chore: import upstream snapshot with attribution

python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.amx_utils import (
+    CPUQuantMethod,
+    _amx_process_weight_after_loading,
+)
+from sglang.srt.layers.moe import MoeRunnerConfig
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.quantization.awq.awq import AWQConfig
+
+__all__ = ["AWQIntelAMXLinearKernel", "AWQIntelAMXMoEKernel"]
+
+
+class AWQIntelAMXLinearKernel:
+    def __init__(self, quant_config: AWQConfig):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        _amx_process_weight_after_loading(
+            layer, ["qweight", "qzeros", "scales"], None, "awq"
+        )
+        layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
+        layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
+        layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return torch.ops.sgl_kernel.int4_scaled_mm_cpu(
+            x,
+            layer.qweight,
+            layer.qzeros,
+            layer.scales,
+            bias,
+        )
+
+
+class AWQIntelAMXMoEKernel:
+    def __init__(self, quant_config: AWQConfig):
+        self.quant_config = quant_config
+        self.moe_runner_config: Optional[MoeRunnerConfig] = None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        _amx_process_weight_after_loading(
+            layer, ["w13_qweight", "w13_qzeros", "w13_scales"], None, "awq"
+        )
+        _amx_process_weight_after_loading(
+            layer, ["w2_qweight", "w2_qzeros", "w2_scales"], None, "awq"
+        )
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        self.moe_runner_config = moe_runner_config
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> torch.Tensor:
+        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
+        assert (
+            self.moe_runner_config.activation == "silu"
+        ), "Only SiLU activation is supported."
+
+        x = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+        topk_weights, topk_ids, _ = topk_output
+        output = torch.ops.sgl_kernel.fused_experts_cpu(
+            x,
+            layer.w13_qweight,
+            layer.w2_qweight,
+            topk_weights,
+            topk_ids,
+            False,  # inplace See [Note] inplace should be False in fused_experts.
+            CPUQuantMethod.INT4_W4A8,
+            layer.w13_scales,  # w1_scale
+            layer.w2_scales,  # w2_scale
+            layer.w13_qzeros,
+            layer.w2_qzeros,
+            None,  # block_size
+            None,  # w1 bias
+            None,  # w3 bias
+            None,  # alpha
+            None,  # limit
+            True,  # is_vnni
+        )
+        return StandardCombineInput(hidden_states=output)

python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.amx_utils import (
+    CPUQuantMethod,
+    _amx_process_weight_after_loading,
+)
+from sglang.srt.layers.moe import MoeRunnerConfig
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig
+
+__all__ = ["GPTQIntelAMXLinearKernel", "GPTQIntelAMXMoEKernel"]
+
+
+class GPTQIntelAMXLinearKernel:
+    def __init__(self, quant_config: GPTQConfig):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        _amx_process_weight_after_loading(
+            layer, ["qweight", "qzeros", "scales"], None, "gptq"
+        )
+        layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
+        layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
+        layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return torch.ops.sgl_kernel.int4_scaled_mm_cpu(
+            x,
+            layer.qweight,
+            layer.qzeros,
+            layer.scales,
+            bias,
+        )
+
+
+class GPTQIntelAMXMoEKernel:
+    def __init__(self, quant_config: GPTQConfig):
+        self.quant_config = quant_config
+        self.moe_runner_config: Optional[MoeRunnerConfig] = None
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        self.moe_runner_config = moe_runner_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        _amx_process_weight_after_loading(
+            layer, ["w13_qweight", "w13_qzeros", "w13_scales"], None, "gptq"
+        )
+        _amx_process_weight_after_loading(
+            layer, ["w2_qweight", "w2_qzeros", "w2_scales"], None, "gptq"
+        )
+
+    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.activation == "silu"
+        ), "Only SiLU activation is supported."
+
+        x = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+        topk_weights, topk_ids, _ = topk_output
+        output = torch.ops.sgl_kernel.fused_experts_cpu(
+            x,
+            layer.w13_qweight,
+            layer.w2_qweight,
+            topk_weights,
+            topk_ids,
+            False,  # inplace See [Note] inplace should be False in fused_experts.
+            CPUQuantMethod.INT4_W4A8,
+            layer.w13_scales,  # w1_scale
+            layer.w2_scales,  # w2_scale
+            layer.w13_qzeros,
+            layer.w2_qzeros,
+            None,  # block_size
+            None,  # w1 bias
+            None,  # w3 bias
+            None,  # alpha
+            None,  # limit
+            True,  # is_vnni
+        )
+        return StandardCombineInput(hidden_states=output)

python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py

@@ -0,0 +1,255 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.moe import MoeRunner
+from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
+from sglang.srt.layers.quantization.marlin_utils import (
+    apply_awq_marlin_linear,
+    awq_to_marlin_zero_points,
+    marlin_make_empty_g_idx,
+    marlin_make_workspace,
+    marlin_moe_permute_scales,
+    marlin_permute_scales,
+    moe_awq_to_marlin_zero_points,
+)
+from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter
+from sglang.srt.utils import is_hip, is_xpu
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        StandardDispatchOutput,
+    )
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+awq_marlin_moe_repack = None
+awq_marlin_repack = None
+
+
+def _unsupported_awq_dequantize(*args, **kwargs):
+    raise RuntimeError("AWQ GPU kernels are unavailable on the current platform.")
+
+
+awq_dequantize = _unsupported_awq_dequantize
+
+if is_xpu():
+    try:
+        from sgl_kernel import awq_dequantize
+    except ImportError:
+        pass
+elif is_hip():
+    try:
+        from sglang.srt.layers.quantization.awq.awq_triton import (
+            awq_dequantize_triton as awq_dequantize,
+        )
+    except ImportError:
+        pass
+else:
+    try:
+        from sglang.jit_kernel.awq_dequantize import awq_dequantize
+        from sglang.jit_kernel.awq_marlin_repack import (
+            awq_marlin_moe_repack,
+            awq_marlin_repack,
+        )
+        from sglang.srt.utils.custom_op import register_custom_op_from_extern
+
+        awq_dequantize = register_custom_op_from_extern(
+            awq_dequantize,
+            fake_impl=lambda qweight, scales, qzeros: qweight.new_empty(
+                qweight.shape[:-1] + (qweight.shape[-1] * 8,), dtype=scales.dtype
+            ),
+        )
+    except ImportError:
+        try:
+            from sglang.srt.layers.quantization.awq.awq_triton import (
+                awq_dequantize_triton as awq_dequantize,
+            )
+        except ImportError:
+            try:
+                from sgl_kernel import awq_dequantize
+            except ImportError:
+                pass
+
+_, scalar_types = get_scalar_types()
+
+
+class AWQLinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
+        layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
+        layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qweight = layer.qweight
+        scales = layer.scales
+        qzeros = layer.qzeros
+        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])
+        out = awq_dequantize(qweight, scales, qzeros)
+        out = torch.matmul(reshaped_x, out)
+
+        if bias is not None:
+            out.add_(bias)
+        return out.reshape(out_shape)
+
+
+class AWQMarlinLinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        device = layer.qweight.device
+        layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
+        layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
+        layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
+
+        layer.workspace = marlin_make_workspace(device)
+
+        marlin_qweight = awq_marlin_repack(
+            layer.qweight,
+            size_k=layer.input_size_per_partition,
+            size_n=layer.output_size_per_partition,
+            num_bits=self.quant_config.quant_type.size_bits,
+        )
+        replace_parameter(layer, "qweight", marlin_qweight)
+
+        marlin_scales = marlin_permute_scales(
+            layer.scales,
+            size_k=layer.input_size_per_partition,
+            size_n=layer.output_size_per_partition,
+            group_size=self.quant_config.group_size,
+        )
+        replace_parameter(layer, "scales", marlin_scales)
+
+        marlin_zp = awq_to_marlin_zero_points(
+            layer.qzeros,
+            size_k=layer.num_groups,
+            size_n=layer.output_size_per_partition,
+            num_bits=self.quant_config.quant_type.size_bits,
+        )
+        replace_parameter(layer, "qzeros", marlin_zp)
+
+        layer.g_idx = marlin_make_empty_g_idx(device)
+        layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return apply_awq_marlin_linear(
+            input=x,
+            weight=layer.qweight,
+            weight_scale=layer.scales,
+            weight_zp=layer.qzeros,
+            g_idx=layer.g_idx,
+            g_idx_sort_indices=layer.g_idx_sort_indices,
+            workspace=layer.workspace,
+            quant_type=self.quant_config.quant_type,
+            output_size_per_partition=layer.output_size_per_partition,
+            input_size_per_partition=layer.input_size_per_partition,
+            bias=bias,
+        )
+
+
+class AWQMoEKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.runner: Optional[MoeRunner] = None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        num_experts = layer.w13_qweight.shape[0]
+        device = layer.w13_qweight.device
+
+        layer.w13_g_idx_sort_indices = torch.nn.Parameter(
+            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+            requires_grad=False,
+        )
+        layer.w2_g_idx_sort_indices = torch.nn.Parameter(
+            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+            requires_grad=False,
+        )
+
+        marlin_w13_qweight = awq_marlin_moe_repack(
+            layer.w13_qweight,
+            layer.w13_g_idx_sort_indices,
+            size_k=layer.w13_qweight.shape[1],
+            size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor,
+            num_bits=self.quant_config.weight_bits,
+        )
+        replace_parameter(layer, "w13_qweight", marlin_w13_qweight)
+
+        marlin_w2_qweight = awq_marlin_moe_repack(
+            layer.w2_qweight,
+            layer.w2_g_idx_sort_indices,
+            size_k=layer.w2_qweight.shape[1],
+            size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor,
+            num_bits=self.quant_config.weight_bits,
+        )
+        replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
+
+        marlin_w13_scales = marlin_moe_permute_scales(
+            s=layer.w13_scales,
+            size_k=layer.intermediate_size_per_partition,
+            size_n=layer.w13_scales.shape[2],
+            group_size=self.quant_config.group_size,
+        )
+        replace_parameter(layer, "w13_scales", marlin_w13_scales)
+
+        marlin_w2_scales = marlin_moe_permute_scales(
+            s=layer.w2_scales,
+            size_k=layer.intermediate_size_per_partition,
+            size_n=layer.w2_scales.shape[2],
+            group_size=self.quant_config.group_size,
+        )
+        replace_parameter(layer, "w2_scales", marlin_w2_scales)
+
+        marlin_w13_zp = moe_awq_to_marlin_zero_points(
+            layer.w13_qzeros,
+            size_k=layer.w13_qzeros.shape[1],
+            size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor,
+            num_bits=self.quant_config.weight_bits,
+        )
+        replace_parameter(layer, "w13_qzeros", marlin_w13_zp)
+
+        marlin_w2_zp = moe_awq_to_marlin_zero_points(
+            layer.w2_qzeros,
+            size_k=layer.w2_qzeros.shape[1],
+            size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor,
+            num_bits=self.quant_config.weight_bits,
+        )
+        replace_parameter(layer, "w2_qzeros", marlin_w2_zp)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        if self.runner is None:
+            raise RuntimeError("moe runner is not initialized")
+
+        quant_info = MarlinMoeQuantInfo(
+            w13_qweight=layer.w13_qweight,
+            w2_qweight=layer.w2_qweight,
+            w13_scales=layer.w13_scales,
+            w2_scales=layer.w2_scales,
+            w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
+            w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
+            w13_qzeros=layer.w13_qzeros,
+            w2_qzeros=layer.w2_qzeros,
+            weight_bits=self.quant_config.weight_bits,
+        )
+        return self.runner.run(dispatch_output, quant_info)

python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py

@@ -0,0 +1,382 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, get_moe_runner_backend
+from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
+from sglang.srt.layers.parameter import BasevLLMParameter, permute_param_layout_
+from sglang.srt.layers.quantization.marlin_utils import (
+    apply_gptq_marlin_linear,
+    check_marlin_supports_shape,
+    marlin_is_k_full,
+    marlin_make_empty_g_idx,
+    marlin_make_workspace,
+    marlin_moe_permute_scales,
+    marlin_permute_scales,
+    marlin_sort_g_idx,
+    marlin_zero_points,
+)
+from sglang.srt.layers.quantization.utils import (
+    get_scalar_types,
+    replace_parameter,
+    unpack_cols,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe import MoeRunnerConfig
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        StandardDispatchOutput,
+    )
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+ScalarType, _ = get_scalar_types()
+
+
+def _unsupported_kernel(*args, **kwargs):
+    raise RuntimeError("GPTQ CUDA kernels are unavailable on the current platform.")
+
+
+gptq_gemm = _unsupported_kernel
+gptq_marlin_repack = _unsupported_kernel
+gptq_shuffle = _unsupported_kernel
+
+try:
+    from sgl_kernel import gptq_gemm, gptq_shuffle
+
+    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+except Exception:
+    pass
+
+
+@dataclass
+class MarlinLinearLayerConfig:
+    full_weight_shape: tuple[int, int]  # [in, out]
+    partition_weight_shape: tuple[int, int]
+    weight_type: ScalarType
+    act_type: torch.dtype
+    group_size: int
+    zero_points: bool
+    has_g_idx: bool
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = gptq_marlin_repack(b_q_weight[e], perm[e], size_k, size_n, num_bits)
+    return output
+
+
+class GPTQLinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.use_shuffle = True
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # for torch.compile
+        layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
+        layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
+        layer.g_idx = torch.nn.Parameter(layer.g_idx.data, requires_grad=False)
+        layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
+
+        # exllama needs to shuffle the weight after the weight is loaded
+        # here we do the shuffle on first forward pass
+        if self.use_shuffle:
+            if self.quant_config.desc_act:
+                layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int)
+            else:
+                layer.g_idx.data = torch.empty(
+                    (0,), dtype=torch.int, device=layer.g_idx.device
+                )
+            gptq_shuffle(layer.qweight, layer.g_idx, self.quant_config.weight_bits)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        out_shape = x.shape[:-1] + (layer.qweight.shape[-1],)
+        reshaped_x = x.reshape(-1, x.shape[-1])
+
+        output = gptq_gemm(
+            reshaped_x,
+            layer.qweight,
+            layer.qzeros,
+            layer.scales,
+            layer.g_idx,
+            self.use_shuffle,
+            self.quant_config.weight_bits,
+        )
+        if bias is not None:
+            output.add_(bias)
+        return output.reshape(out_shape)
+
+
+class GPTQMarlinLinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        device = getattr(layer, "qweight").device
+        c = self.kernel_config
+
+        check_marlin_supports_shape(
+            c.partition_weight_shape[1],  # out_features
+            c.partition_weight_shape[0],  # in_features
+            c.full_weight_shape[0],  # in_features
+            c.group_size,
+        )
+
+        row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0]
+        self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel)
+
+        # Allocate marlin workspace.
+        self.workspace = marlin_make_workspace(device)
+
+        # Default names since marlin requires empty parameters for these,
+        # TODO: remove this requirement from marlin (allow optional tensors)
+        self.w_q_name = "qweight"
+        self.w_s_name = "scales"
+        self.w_zp_name = "qzeros"
+        self.w_gidx_name = "g_idx"
+
+        def _transform_param(
+            layer: torch.nn.Module, name: Optional[str], fn: Callable
+        ) -> None:
+            if name is not None and getattr(layer, name, None) is not None:
+
+                old_param = getattr(layer, name)
+                new_param = fn(old_param)
+                # replace the parameter with torch.nn.Parameter for TorchDynamo
+                # compatibility
+                replace_parameter(
+                    layer, name, torch.nn.Parameter(new_param.data, requires_grad=False)
+                )
+
+        def transform_w_q(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+            x.data = gptq_marlin_repack(
+                x.data.contiguous(),
+                perm=layer.g_idx_sort_indices,
+                size_k=c.partition_weight_shape[0],
+                size_n=c.partition_weight_shape[1],
+                num_bits=c.weight_type.size_bits,
+            )
+            return x
+
+        def transform_w_s(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1)
+            x.data = marlin_permute_scales(
+                x.data.contiguous(),
+                size_k=c.partition_weight_shape[0],
+                size_n=c.partition_weight_shape[1],
+                group_size=c.group_size,
+            )
+            return x
+
+        if c.has_g_idx:
+            g_idx, g_idx_sort_indices = marlin_sort_g_idx(
+                getattr(layer, self.w_gidx_name)
+            )
+            _transform_param(layer, self.w_gidx_name, lambda _: g_idx)
+            layer.g_idx_sort_indices = g_idx_sort_indices
+        else:
+            setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device))
+            layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
+
+        if c.zero_points:
+            grouped_k = (
+                c.partition_weight_shape[0] // c.group_size if c.group_size != -1 else 1
+            )
+            _transform_param(
+                layer,
+                self.w_zp_name,
+                lambda x: marlin_zero_points(
+                    unpack_cols(
+                        x.t(),
+                        c.weight_type.size_bits,
+                        grouped_k,
+                        c.partition_weight_shape[1],
+                    ),
+                    size_k=grouped_k,
+                    size_n=c.partition_weight_shape[1],
+                    num_bits=c.weight_type.size_bits,
+                ),
+            )
+        else:
+            setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device))
+        _transform_param(layer, self.w_q_name, transform_w_q)
+        _transform_param(layer, self.w_s_name, transform_w_s)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        c = self.kernel_config
+
+        def _get_weight_params(
+            layer: torch.nn.Module,
+        ) -> tuple[
+            torch.Tensor,  # w_q
+            torch.Tensor,  # w_s
+            Optional[torch.Tensor],  # w_zp,
+            Optional[torch.Tensor],  # w_gidx
+        ]:
+            return (
+                getattr(layer, self.w_q_name),
+                getattr(layer, self.w_s_name),
+                getattr(layer, self.w_zp_name or "", None),
+                getattr(layer, self.w_gidx_name or "", None),
+            )
+
+        w_q, w_s, w_zp, w_gidx = _get_weight_params(layer)
+
+        # `process_weights_after_loading` will ensure w_zp and w_gidx are not
+        #  None for marlin
+        return apply_gptq_marlin_linear(
+            input=x,
+            weight=w_q,
+            weight_scale=w_s,
+            weight_zp=w_zp,  # type: ignore
+            g_idx=w_gidx,  # type: ignore
+            g_idx_sort_indices=layer.g_idx_sort_indices,
+            workspace=self.workspace,
+            wtype=c.weight_type,
+            input_size_per_partition=c.partition_weight_shape[0],
+            output_size_per_partition=c.partition_weight_shape[1],
+            is_k_full=self.is_k_full,
+            bias=bias,
+        )
+
+
+class GPTQMarlinMoEKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        # Process act_order
+        if self.quant_config.desc_act:
+            # Get sorting based on g_idx
+            num_experts = layer.w13_g_idx.shape[0]
+            w13_g_idx_sort_indices = torch.empty_like(layer.w13_g_idx)
+            w2_g_idx_sort_indices = torch.empty_like(layer.w2_g_idx)
+            w13_sorted_g_idx = torch.empty_like(layer.w13_g_idx)
+            w2_sorted_g_idx = torch.empty_like(layer.w2_g_idx)
+            for e in range(num_experts):
+                w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_g_idx[e]).to(
+                    torch.int32
+                )
+                w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_g_idx[e]).to(
+                    torch.int32
+                )
+                w13_sorted_g_idx[e] = layer.w13_g_idx[e][w13_g_idx_sort_indices[e]]
+                w2_sorted_g_idx[e] = layer.w2_g_idx[e][w2_g_idx_sort_indices[e]]
+            replace_parameter(layer, "w13_g_idx", w13_sorted_g_idx)
+            replace_parameter(layer, "w2_g_idx", w2_sorted_g_idx)
+            replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
+            replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
+        else:
+            # Reset g_idx related tensors
+            num_experts = layer.w13_g_idx.shape[0]
+            device = layer.w13_g_idx.device
+            layer.w13_g_idx = torch.nn.Parameter(
+                torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+                requires_grad=False,
+            )
+            layer.w2_g_idx = torch.nn.Parameter(
+                torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+                requires_grad=False,
+            )
+            layer.w13_g_idx_sort_indices = torch.nn.Parameter(
+                torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+                requires_grad=False,
+            )
+            layer.w2_g_idx_sort_indices = torch.nn.Parameter(
+                torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+                requires_grad=False,
+            )
+        # Repack weights
+        marlin_w13_qweight = gptq_marlin_moe_repack(
+            layer.w13_qweight,
+            layer.w13_g_idx_sort_indices,
+            layer.w13_qweight.shape[1] * self.quant_config.pack_factor,
+            layer.w13_qweight.shape[2],
+            self.quant_config.weight_bits,
+        )
+        replace_parameter(layer, "w13_qweight", marlin_w13_qweight)
+        marlin_w2_qweight = gptq_marlin_moe_repack(
+            layer.w2_qweight,
+            layer.w2_g_idx_sort_indices,
+            layer.w2_qweight.shape[1] * self.quant_config.pack_factor,
+            layer.w2_qweight.shape[2],
+            self.quant_config.weight_bits,
+        )
+        replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
+        # Repack scales
+        marlin_w13_scales = marlin_moe_permute_scales(
+            s=layer.w13_scales,
+            size_k=layer.intermediate_size_per_partition,
+            size_n=layer.w13_scales.shape[2],
+            group_size=self.quant_config.group_size,
+        )
+        replace_parameter(layer, "w13_scales", marlin_w13_scales)
+        marlin_w2_scales = marlin_moe_permute_scales(
+            s=layer.w2_scales,
+            size_k=layer.w2_scales.shape[1]
+            * (
+                self.quant_config.group_size
+                if self.quant_config.group_size != -1
+                else self.quant_config.pack_factor
+            ),
+            size_n=layer.w2_scales.shape[2],
+            group_size=self.quant_config.group_size,
+        )
+        replace_parameter(layer, "w2_scales", marlin_w2_scales)
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        assert get_moe_runner_backend().is_auto()
+        self.moe_runner_config = moe_runner_config
+        self.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        quant_info = MarlinMoeQuantInfo(
+            w13_qweight=layer.w13_qweight,
+            w2_qweight=layer.w2_qweight,
+            w13_scales=layer.w13_scales,
+            w2_scales=layer.w2_scales,
+            w13_g_idx=layer.w13_g_idx,
+            w2_g_idx=layer.w2_g_idx,
+            w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
+            w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
+            weight_bits=self.quant_config.weight_bits,
+            is_k_full=self.is_k_full,
+        )
+
+        return self.runner.run(dispatch_output, quant_info)

python/sglang/srt/hardware_backend/mlx/aot.py

@@ -0,0 +1,240 @@
+"""AOT kernel selection and decode-context helpers for the MLX backend."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import mlx.core as mx
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
+        ContiguousAttentionKVCache,
+    )
+
+
+def _load_metal_rope_pool_fused():
+    try:
+        from sgl_kernel import metal
+    except ImportError as exc:
+        raise ImportError(
+            "sgl_kernel.metal is not importable. Install sgl-kernel in the "
+            "active environment before enabling SGLANG_MLX_USE_CUSTOM_ROPE."
+        ) from exc
+
+    import_error = getattr(metal, "_IMPORT_ERROR", None)
+    if getattr(metal, "_metal", None) is None or import_error is not None:
+        reason = f" Reason: {import_error}." if import_error is not None else ""
+        raise ImportError(
+            "sgl_kernel.metal is importable, but the native Metal extension "
+            f"or metallib is not available.{reason} Install the Metal kernels "
+            "with `uv run sgl-kernel/setup_metal.py install` from the SGLang "
+            "repo root in the active environment."
+        ) from import_error
+    return metal.rope_pool_fused
+
+
+@dataclass
+class MlxAOTRoPEKernel:
+    base: float = 0.0
+    config: dict[str, Any] = field(default_factory=dict)
+    rope_pool_fused: Optional[Any] = None
+
+    @property
+    def enabled(self) -> bool:
+        return (
+            self.base > 0.0 and bool(self.config) and self.rope_pool_fused is not None
+        )
+
+
+@dataclass
+class MlxAOTKernelBuildInputs:
+    sample_attn: Any
+    n_kv_heads: int
+    head_dim: int
+
+
+@dataclass(frozen=True)
+class MlxAOTKernelSpec:
+    name: str
+    kernel_attr: str
+    is_enabled: Callable[[], bool]
+    build: Callable[[MlxAOTKernelBuildInputs], Any]
+
+
+@dataclass
+class MlxAOTKernelSet:
+    rope: MlxAOTRoPEKernel = field(default_factory=MlxAOTRoPEKernel)
+    selected_kernel_names: tuple[str, ...] = ()
+
+
+class MlxAOTKernelRegistry:
+    """Registry for optional MLX AOT kernels.
+
+    Each spec owns one kernel field on ``MlxAOTKernelSet``. The registry is the
+    only place that checks kernel enablement policy and model support.
+    """
+
+    def __init__(self, specs: tuple[MlxAOTKernelSpec, ...]):
+        self._specs = specs
+
+    @property
+    def registered_kernel_names(self) -> tuple[str, ...]:
+        return tuple(spec.name for spec in self._specs)
+
+    def build_kernel_set(
+        self,
+        *,
+        sample_attn: Any,
+        n_kv_heads: int,
+        head_dim: int,
+    ) -> MlxAOTKernelSet:
+        inputs = MlxAOTKernelBuildInputs(
+            sample_attn=sample_attn,
+            n_kv_heads=n_kv_heads,
+            head_dim=head_dim,
+        )
+        kernel_set = MlxAOTKernelSet()
+        selected_kernel_names = []
+        for spec in self._specs:
+            if not spec.is_enabled():
+                continue
+            kernel = spec.build(inputs)
+            if getattr(kernel, "enabled", False):
+                if not hasattr(kernel_set, spec.kernel_attr):
+                    raise ValueError(
+                        f"AOT kernel {spec.name} targets unknown kernel-set "
+                        f"attribute {spec.kernel_attr}"
+                    )
+                setattr(kernel_set, spec.kernel_attr, kernel)
+                selected_kernel_names.append(spec.name)
+        kernel_set.selected_kernel_names = tuple(selected_kernel_names)
+        if kernel_set.selected_kernel_names:
+            logger.info(
+                "MLX AOT kernels selected: %s",
+                ", ".join(kernel_set.selected_kernel_names),
+            )
+        return kernel_set
+
+
+def _build_rope_kernel(inputs: MlxAOTKernelBuildInputs) -> MlxAOTRoPEKernel:
+    from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+        get_num_heads,
+    )
+
+    sample_attn = getattr(inputs.sample_attn, "_inner", inputs.sample_attn)
+    rope = getattr(sample_attn, "rope", None)
+    if rope is None or getattr(rope, "traditional", False):
+        return MlxAOTRoPEKernel()
+
+    rope_dim = int(getattr(rope, "dims", 0))
+    if rope_dim == 0:
+        return MlxAOTRoPEKernel()
+    if rope_dim != inputs.head_dim:
+        # AOT kernel currently requires rope_dim == head_dim.
+        return MlxAOTRoPEKernel()
+
+    base = float(getattr(rope, "base", 10000.0))
+    num_qo_heads = get_num_heads(sample_attn)
+    if num_qo_heads is None:
+        return MlxAOTRoPEKernel()
+    config = {
+        "head_dim": int(inputs.head_dim),
+        "rope_dim": rope_dim,
+        "num_qo_heads": int(num_qo_heads),
+        "num_kv_heads": int(inputs.n_kv_heads),
+    }
+    try:
+        rope_pool_fused = _load_metal_rope_pool_fused()
+    except Exception as exc:  # noqa: BLE001
+        logger.info(
+            "AOT Metal RoPE kernel not available (%s) - falling back to "
+            "mx.fast.rope.",
+            exc,
+        )
+        return MlxAOTRoPEKernel()
+
+    logger.info(
+        f"AOT Metal RoPE kernel ENABLED: head_dim={inputs.head_dim}, "
+        f"n_heads={config['num_qo_heads']}, n_kv={config['num_kv_heads']}, "
+        f"base={base}"
+    )
+    return MlxAOTRoPEKernel(
+        base=base,
+        config=config,
+        rope_pool_fused=rope_pool_fused,
+    )
+
+
+MLX_AOT_KERNEL_REGISTRY = MlxAOTKernelRegistry(
+    specs=(
+        MlxAOTKernelSpec(
+            name="metal_rope_pool_fused",
+            kernel_attr="rope",
+            is_enabled=lambda: envs.SGLANG_MLX_USE_CUSTOM_ROPE.get(),
+            build=_build_rope_kernel,
+        ),
+    )
+)
+
+
+@dataclass
+class MlxAOTRoPEContext:
+    kernel: MlxAOTRoPEKernel
+    kv_pool: Any
+    new_token_slots: Optional[mx.array] = None
+
+
+@dataclass
+class MlxAOTKernelContext:
+    rope: Optional[MlxAOTRoPEContext] = None
+
+    @classmethod
+    def from_decode(
+        cls,
+        *,
+        aot_kernels: MlxAOTKernelSet,
+        kv_pool: Any | None,
+        req_ids: list[str],
+        req_pool_idx: dict[str, int],
+        req_to_token_pool: Any | None,
+        layer_caches: list[list[ContiguousAttentionKVCache]],
+    ) -> MlxAOTKernelContext:
+        """Build optional AOT context for one batched decode step."""
+        if not aot_kernels.rope.enabled or kv_pool is None:
+            return cls()
+
+        new_token_slots = None
+        if req_to_token_pool is not None:
+            try:
+                slot_ids = []
+                for req_idx, req_id in enumerate(req_ids):
+                    pool_idx = req_pool_idx.get(req_id)
+                    if pool_idx is None:
+                        raise KeyError(req_id)
+                    slot = int(
+                        req_to_token_pool.req_to_token[
+                            pool_idx, layer_caches[0][req_idx].offset
+                        ].item()
+                    )
+                    slot_ids.append(slot)
+                new_token_slots = mx.array(slot_ids, dtype=mx.int32)
+            except Exception as exc:  # noqa: BLE001
+                logger.warning(
+                    "AOT RoPE: failed to resolve new-token slots (%s); "
+                    "falling back to RoPE-only for this decode step",
+                    exc,
+                )
+
+        return cls(
+            rope=MlxAOTRoPEContext(
+                kernel=aot_kernels.rope,
+                kv_pool=kv_pool,
+                new_token_slots=new_token_slots,
+            )
+        )

python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py

@@ -0,0 +1,59 @@
+"""Cache components for the MLX backend."""
+
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    get_head_dim,
+    get_num_heads,
+    get_num_kv_heads,
+    is_attention_module,
+    uses_sliding_window_attention,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
+    AttentionOffsetCache,
+    ContiguousAttentionKVCache,
+    PoolBackedAttentionKVCache,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import (
+    MlxAttentionKVPool,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import (
+    BatchedDecodeContext,
+    MLXAttentionWrapper,
+    clear_context,
+    get_context,
+    set_context,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import (
+    MlxAuxiliaryStateComponent,
+    MlxAuxiliaryStatePool,
+    MlxAuxiliaryStateReqToTokenPool,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.layout import MlxModelCacheLayout
+from sglang.srt.hardware_backend.mlx.kv_cache.model_patching import (
+    find_attention_layers,
+    get_num_layers,
+    patch_model_attention,
+)
+
+__all__ = [
+    "BatchedDecodeContext",
+    "clear_context",
+    "AttentionOffsetCache",
+    "ContiguousAttentionKVCache",
+    "find_attention_layers",
+    "get_head_dim",
+    "get_context",
+    "get_num_layers",
+    "get_num_heads",
+    "get_num_kv_heads",
+    "is_attention_module",
+    "MLXAttentionWrapper",
+    "MlxAttentionKVPool",
+    "MlxAuxiliaryStateComponent",
+    "MlxAuxiliaryStatePool",
+    "MlxAuxiliaryStateReqToTokenPool",
+    "MlxModelCacheLayout",
+    "patch_model_attention",
+    "PoolBackedAttentionKVCache",
+    "set_context",
+    "uses_sliding_window_attention",
+]

python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py

@@ -0,0 +1,66 @@
+"""Attention helpers based on duck typing for the MLX backend."""
+
+from __future__ import annotations
+
+from typing import Any, Iterable
+
+# ``rope`` and ``scale`` are required by MLXAttentionWrapper. Keeping them in
+# the contract also prevents recurrent mixers such as DeltaNet from being
+# mistaken for softmax attention just because they expose projection layers.
+ATTENTION_API_ATTRS = ("q_proj", "k_proj", "v_proj", "o_proj", "rope", "scale")
+NUM_HEAD_ATTRS = ("n_heads", "num_heads", "num_attention_heads")
+NUM_KV_HEAD_ATTRS = ("n_kv_heads", "num_k_heads", "num_kv_heads", "num_key_value_heads")
+SLIDING_ATTENTION_ATTRS = (
+    "is_sliding",
+    "use_sliding",
+    "is_sliding_window",
+    "use_sliding_window",
+    "is_swa",
+)
+
+
+def first_present_attr(module: Any, names: Iterable[str]) -> Any | None:
+    """Return the first present attribute value without treating 0 as absent."""
+    for name in names:
+        if hasattr(module, name):
+            return getattr(module, name)
+    return None
+
+
+def get_num_heads(module: Any) -> int | None:
+    return first_present_attr(module, NUM_HEAD_ATTRS)
+
+
+def get_num_kv_heads(module: Any) -> int | None:
+    return first_present_attr(module, NUM_KV_HEAD_ATTRS)
+
+
+def get_head_dim(module: Any) -> int | None:
+    head_dim = first_present_attr(module, ("head_dim",))
+    if head_dim is not None:
+        return head_dim
+
+    n_kv_heads = get_num_kv_heads(module)
+    if n_kv_heads is None:
+        return None
+    if hasattr(module, "hidden_size") and hasattr(module, "num_k_heads"):
+        return module.hidden_size // module.num_k_heads
+    if hasattr(module, "k_proj") and hasattr(module.k_proj, "weight"):
+        return module.k_proj.weight.shape[0] // n_kv_heads
+    return None
+
+
+def is_attention_module(module: Any) -> bool:
+    return (
+        all(hasattr(module, attr) for attr in ATTENTION_API_ATTRS)
+        and get_num_heads(module) is not None
+        and get_num_kv_heads(module) is not None
+    )
+
+
+def uses_sliding_window_attention(*modules: Any) -> bool:
+    return any(
+        bool(getattr(module, attr, False))
+        for module in modules
+        for attr in SLIDING_ATTENTION_ATTRS
+    )

python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py

@@ -0,0 +1,210 @@
+"""Attention KV cache adapters for the MLX backend."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import mlx.core as mx
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import (
+        MlxAttentionKVPool,
+    )
+
+
+class AttentionOffsetCache:
+    """Data-free shim satisfying mlx-lm's cache protocol.
+
+    Provides ``make_mask`` and ``state`` without storing actual K/V.
+    """
+
+    def __init__(self, offset: int = 0):
+        self.offset = offset
+
+    @property
+    def state(self):
+        return ()  # Empty — safe for mx.eval unpacking
+
+    def make_mask(self, N, **kwargs):
+        return None if N == 1 else "causal"
+
+    def update_and_fetch(self, keys, values):
+        raise RuntimeError("AttentionOffsetCache should not store data")
+
+
+_DEFAULT_MAX_SEQ_LEN = 4096
+
+
+class ContiguousAttentionKVCache:
+    """Pre-allocated attention KV buffer for one request and one layer.
+
+    Shape ``(1, n_kv_heads, max_seq_len, head_dim)``.  Slice assignment
+    instead of ``mx.concatenate``.  Lazy-allocated on first write.
+    """
+
+    __slots__ = ("keys", "values", "offset", "max_seq_len")
+
+    def __init__(
+        self,
+        n_kv_heads: int | None = None,
+        head_dim: int | None = None,
+        max_seq_len: int = _DEFAULT_MAX_SEQ_LEN,
+        dtype: mx.Dtype | None = None,
+    ):
+        if n_kv_heads is not None and head_dim is not None and dtype is not None:
+            self.keys = mx.zeros((1, n_kv_heads, max_seq_len, head_dim), dtype=dtype)
+            self.values = mx.zeros((1, n_kv_heads, max_seq_len, head_dim), dtype=dtype)
+        else:
+            self.keys = None
+            self.values = None
+        self.offset = 0
+        self.max_seq_len = max_seq_len
+
+    def _allocate(self, keys: mx.array) -> None:
+        """Allocate buffers matching the first key tensor's shape."""
+        B, n_kv_heads, _, head_dim = keys.shape
+        self.keys = mx.zeros(
+            (B, n_kv_heads, self.max_seq_len, head_dim), dtype=keys.dtype
+        )
+        self.values = mx.zeros(
+            (B, n_kv_heads, self.max_seq_len, head_dim), dtype=keys.dtype
+        )
+
+    @property
+    def state(self):
+        """Arrays for ``mx.eval`` unpacking."""
+        if self.keys is None:
+            return ()
+        return (self.keys, self.values)
+
+    def make_mask(self, N, **kwargs):
+        return None if N == 1 else "causal"
+
+    def _grow(self, required: int) -> None:
+        """Double the buffer until it can hold *required* tokens."""
+        new_max = self.max_seq_len
+        while new_max < required:
+            new_max *= 2
+        B, n_kv_heads, _, head_dim = self.keys.shape
+        new_k = mx.zeros((B, n_kv_heads, new_max, head_dim), dtype=self.keys.dtype)
+        new_v = mx.zeros((B, n_kv_heads, new_max, head_dim), dtype=self.values.dtype)
+        if self.offset > 0:
+            new_k[:, :, : self.offset, :] = self.keys[:, :, : self.offset, :]
+            new_v[:, :, : self.offset, :] = self.values[:, :, : self.offset, :]
+        self.keys = new_k
+        self.values = new_v
+        self.max_seq_len = new_max
+
+    def update_and_fetch(
+        self, keys: mx.array, values: mx.array
+    ) -> tuple[mx.array, mx.array]:
+        """Append K/V and return all valid K/V up to current offset."""
+        if self.keys is None:
+            self._allocate(keys)
+        S = keys.shape[2]
+        end = self.offset + S
+        if end > self.max_seq_len:
+            self._grow(end)
+        self.keys[:, :, self.offset : end, :] = keys
+        self.values[:, :, self.offset : end, :] = values
+        self.offset = end
+        return self.keys[:, :, :end, :], self.values[:, :, :end, :]
+
+    def write_token(self, k: mx.array, v: mx.array) -> None:
+        """Write one token. k, v shape: (1, n_kv_heads, 1, head_dim)."""
+        end = self.offset + 1
+        if end > self.max_seq_len:
+            self._grow(end)
+        self.keys[:, :, self.offset : end, :] = k
+        self.values[:, :, self.offset : end, :] = v
+        self.offset = end
+
+    def get_kv(self) -> tuple[mx.array, mx.array]:
+        """Return valid K/V: (1, n_kv_heads, offset, head_dim)."""
+        return self.keys[:, :, : self.offset, :], self.values[:, :, : self.offset, :]
+
+
+class PoolBackedAttentionKVCache:
+    """Lazily gathers cached attention KV from the shared pool during forward.
+
+    Each ``update_and_fetch`` gathers this layer's prefix from the pool
+    on demand, keeping operations in the lazy compute graph.  Convert to
+    ``ContiguousAttentionKVCache`` via ``to_contiguous`` after the forward pass.
+    """
+
+    __slots__ = (
+        "_pool",
+        "_layer_idx",
+        "_slots",
+        "offset",
+        "_full_keys",
+        "_full_values",
+        "_new_keys",
+        "_new_values",
+    )
+
+    def __init__(
+        self,
+        pool: MlxAttentionKVPool,
+        layer_idx: int,
+        slots: mx.array,
+        prefix_len: int,
+    ):
+        self._pool = pool
+        self._layer_idx = layer_idx
+        self._slots = slots
+        self.offset = prefix_len
+        self._full_keys: mx.array | None = None
+        self._full_values: mx.array | None = None
+        self._new_keys: mx.array | None = None
+        self._new_values: mx.array | None = None
+
+    @property
+    def keys(self) -> mx.array | None:
+        return self._full_keys
+
+    @property
+    def values(self) -> mx.array | None:
+        return self._full_values
+
+    @property
+    def state(self):
+        if self._full_keys is not None:
+            return (self._full_keys, self._full_values)
+        return ()
+
+    def make_mask(self, N, **kwargs):
+        return None if N == 1 else "causal"
+
+    def update_and_fetch(
+        self, keys: mx.array, values: mx.array
+    ) -> tuple[mx.array, mx.array]:
+        """Gather cached prefix from pool, concatenate with new K/V."""
+        S = keys.shape[2]
+
+        if self.offset > 0:
+            k_cached, v_cached = self._pool.get_kv(
+                self._layer_idx, self._slots[: self.offset]
+            )
+            # Pool layout (S, n_kv_heads, head_dim) → cache (1, n_kv_heads, S, head_dim)
+            k_cached = k_cached.transpose(1, 0, 2)[None]
+            v_cached = v_cached.transpose(1, 0, 2)[None]
+            k_all = mx.concatenate([k_cached, keys], axis=2)
+            v_all = mx.concatenate([v_cached, values], axis=2)
+        else:
+            k_all = keys
+            v_all = values
+
+        self.offset += S
+        self._full_keys = k_all
+        self._full_values = v_all
+        self._new_keys = keys
+        self._new_values = values
+        return k_all, v_all
+
+    def to_contiguous(self, max_seq_len: int = 4096) -> ContiguousAttentionKVCache:
+        """Convert to contiguous attention KV reusing forward-pass arrays."""
+        cache = ContiguousAttentionKVCache(max_seq_len=max_seq_len)
+        if self._full_keys is not None:
+            cache.update_and_fetch(self._full_keys, self._full_values)
+        return cache

python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_pool.py

@@ -0,0 +1,86 @@
+"""Flat attention KV pool for the MLX backend.
+
+Each layer buffer has shape ``(pool_size, n_kv_heads, head_dim)``.
+This v1 pool is intentionally uniform: every wrapped softmax-attention
+layer must share the same KV shape and full-context KV semantics.
+Heterogeneous KV shapes and sliding-window KV need per-layer/window-aware
+pools before they can use MLX radix reuse.
+
+Slot 0 is reserved as padding (1-based indexing).
+"""
+
+import logging
+
+import mlx.core as mx
+
+logger = logging.getLogger(__name__)
+
+
+class MlxAttentionKVPool:
+    """Pre-allocated attention KV pool indexed by integer slot IDs."""
+
+    def __init__(
+        self,
+        pool_size: int,
+        num_layers: int,
+        n_kv_heads: int,
+        head_dim: int,
+        dtype: mx.Dtype = mx.float16,
+    ):
+        self.pool_size = pool_size
+        self.num_layers = num_layers
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+        self.dtype = dtype
+
+        # Per-attention-layer buffers: (pool_size, n_kv_heads, head_dim)
+        self.k_buffer: list[mx.array] = [
+            mx.zeros((pool_size, n_kv_heads, head_dim), dtype=dtype)
+            for _ in range(num_layers)
+        ]
+        self.v_buffer: list[mx.array] = [
+            mx.zeros((pool_size, n_kv_heads, head_dim), dtype=dtype)
+            for _ in range(num_layers)
+        ]
+
+        mem_mb = (pool_size * n_kv_heads * head_dim * 2 * num_layers * dtype.size) / (
+            1024 * 1024
+        )
+        logger.info(
+            f"MlxAttentionKVPool: {pool_size} slots x {num_layers} layers "
+            f"x {n_kv_heads} heads x {head_dim} dim, "
+            f"dtype={dtype}, ~{mem_mb:.1f} MB"
+        )
+
+    def set_kv(self, layer_id: int, slots: mx.array, k: mx.array, v: mx.array) -> None:
+        """Scatter K/V into *slots* for one layer."""
+        self.k_buffer[layer_id][slots] = k
+        self.v_buffer[layer_id][slots] = v
+
+    def get_kv(self, layer_id: int, slots: mx.array) -> tuple[mx.array, mx.array]:
+        """Gather K/V from *slots* for one layer."""
+        return self.k_buffer[layer_id][slots], self.v_buffer[layer_id][slots]
+
+    def get_kv_all_layers(self, slots: mx.array) -> tuple[mx.array, mx.array]:
+        """Gather K/V from *slots* across all layers."""
+        k_all = mx.stack([self.k_buffer[i][slots] for i in range(self.num_layers)])
+        v_all = mx.stack([self.v_buffer[i][slots] for i in range(self.num_layers)])
+        return k_all, v_all
+
+    def set_kv_all_layers(
+        self, slots: mx.array, k_all: mx.array, v_all: mx.array
+    ) -> None:
+        """Scatter K/V into *slots* across all layers."""
+        for i in range(self.num_layers):
+            self.set_kv(i, slots, k_all[i], v_all[i])
+
+    def all_buffers(self) -> list[mx.array]:
+        """Return all buffer arrays (for ``mx.eval``)."""
+        return self.k_buffer + self.v_buffer
+
+    def clear(self) -> None:
+        """Zero all buffers."""
+        shape = (self.pool_size, self.n_kv_heads, self.head_dim)
+        for i in range(self.num_layers):
+            self.k_buffer[i] = mx.zeros(shape, dtype=self.dtype)
+            self.v_buffer[i] = mx.zeros(shape, dtype=self.dtype)

python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py

@@ -0,0 +1,298 @@
+"""Batched decode attention wrapper for MLX backend."""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from sglang.srt.hardware_backend.mlx.aot import (
+    MlxAOTKernelContext,
+    MlxAOTKernelSet,
+    MlxAOTRoPEContext,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    get_head_dim,
+    get_num_heads,
+    get_num_kv_heads,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
+    ContiguousAttentionKVCache,
+)
+
+_thread_local = threading.local()
+
+
+# TODO: Move from threading to multiprocessing or asyncio
+@dataclass
+class BatchedDecodeContext:
+    """Context set before batched decode, read by attention wrappers."""
+
+    batch_size: int
+    seq_lens: list[int]  # per-request token count before the new token
+    # attention_layer_caches[attention_pool_idx][req_idx] = ContiguousAttentionKVCache
+    attention_layer_caches: list[list[ContiguousAttentionKVCache]]
+    attention_pool_index_by_layer: dict[int, int] = field(default_factory=dict)
+
+    # Optional AOT kernel state. Keep kernel-specific fields out of the regular
+    # MLX decode path so future AOT kernels can be added without growing this
+    # context one field at a time.
+    aot: MlxAOTKernelContext = field(default_factory=MlxAOTKernelContext)
+
+    # Derived tensors/metadata, shared across all layers in one forward pass.
+    offsets: mx.array = field(init=False)
+    max_len: int = field(init=False)
+    valid_lens: mx.array = field(init=False)
+    needs_padding: bool = field(init=False)
+    pad_sizes: list[int] = field(init=False)
+    positions: Optional[mx.array] = field(init=False)
+
+    def __post_init__(self) -> None:
+        seq_lens = self.seq_lens
+        max_seq_len = max(seq_lens)
+        self.offsets = mx.array(seq_lens, dtype=mx.int32)
+        self.max_len = max_seq_len + 1
+        self.valid_lens = self.offsets + 1
+        self.needs_padding = min(seq_lens) < max_seq_len
+        self.pad_sizes = [max_seq_len - s for s in seq_lens]
+        self.positions = mx.arange(self.max_len) if self.needs_padding else None
+        if not self.attention_pool_index_by_layer:
+            self.attention_pool_index_by_layer = {
+                idx: idx for idx in range(len(self.attention_layer_caches))
+            }
+
+    @classmethod
+    def from_decode(
+        cls,
+        *,
+        caches: list[list[Any]],
+        req_ids: list[str],
+        aot_kernels: MlxAOTKernelSet,
+        kv_pool: Any | None,
+        req_pool_idx: dict[str, int],
+        req_to_token_pool: Any | None,
+        attention_layer_indices: list[int] | None = None,
+        attention_pool_index_by_layer: dict[int, int] | None = None,
+    ) -> BatchedDecodeContext:
+        batch_size = len(req_ids)
+        if attention_layer_indices is None:
+            attention_layer_indices = list(range(len(caches[0])))
+        seq_lens = [
+            caches[i][attention_layer_indices[0]].offset for i in range(batch_size)
+        ]
+        attention_layer_caches = [
+            [caches[i][layer_idx] for i in range(batch_size)]
+            for layer_idx in attention_layer_indices
+        ]
+        return cls(
+            batch_size=batch_size,
+            seq_lens=seq_lens,
+            attention_layer_caches=attention_layer_caches,
+            attention_pool_index_by_layer=attention_pool_index_by_layer or {},
+            aot=MlxAOTKernelContext.from_decode(
+                aot_kernels=aot_kernels,
+                kv_pool=kv_pool,
+                req_ids=req_ids,
+                req_pool_idx=req_pool_idx,
+                req_to_token_pool=req_to_token_pool,
+                layer_caches=attention_layer_caches,
+            ),
+        )
+
+
+def set_context(ctx: Optional[BatchedDecodeContext]) -> None:
+    _thread_local.batched_ctx = ctx
+
+
+def get_context() -> Optional[BatchedDecodeContext]:
+    return getattr(_thread_local, "batched_ctx", None)
+
+
+def clear_context() -> None:
+    _thread_local.batched_ctx = None
+
+
+class MLXAttentionWrapper(nn.Module):
+    """Wraps an mlx-lm Attention for batched decode (BS>1).
+
+    When ``BatchedDecodeContext`` is set, performs per-request RoPE,
+    cache writes, and batched SDPA.  Otherwise delegates to inner module.
+    """
+
+    def __init__(self, inner: nn.Module, layer_idx: int):
+        super().__init__()
+        object.__setattr__(self, "_inner", inner)
+        object.__setattr__(self, "_layer_idx", layer_idx)
+
+    def __call__(self, x: mx.array, mask: Any = None, cache: Any = None) -> mx.array:
+        ctx = get_context()
+        if ctx is None:
+            return self._inner(x, mask=mask, cache=cache)
+        return self._batched_decode(x, ctx)
+
+    def _batched_decode(self, x: mx.array, ctx: BatchedDecodeContext) -> mx.array:
+        inner = self._inner
+        layer_idx = self._layer_idx
+        B = ctx.batch_size
+        n_heads = get_num_heads(inner)
+        n_kv_heads = get_num_kv_heads(inner)
+        if n_heads is None or n_kv_heads is None:
+            raise RuntimeError(
+                f"Cannot determine attention head counts for {type(inner).__name__}"
+            )
+
+        q_proj_output = inner.q_proj(x)
+        keys = inner.k_proj(x)
+        values = inner.v_proj(x)
+
+        head_dim = get_head_dim(inner)
+        if head_dim is None:
+            head_dim = keys.shape[-1] // n_kv_heads
+
+        q_width = n_heads * head_dim
+        gate = None
+        if q_proj_output.shape[-1] == q_width:
+            queries = q_proj_output.reshape(B, 1, n_heads, head_dim)
+        elif q_proj_output.shape[-1] == 2 * q_width:
+            queries, gate = mx.split(
+                q_proj_output.reshape(B, 1, n_heads, 2 * head_dim), 2, axis=-1
+            )
+            gate = gate.reshape(B, 1, q_width)
+        else:
+            raise RuntimeError(
+                f"Unexpected q_proj output shape {q_proj_output.shape} for "
+                f"{type(inner).__name__}"
+            )
+
+        keys = keys.reshape(B, 1, n_kv_heads, head_dim)
+        values = values.reshape(B, 1, n_kv_heads, head_dim)
+
+        if hasattr(inner, "q_norm"):
+            queries = inner.q_norm(queries)
+        if hasattr(inner, "k_norm"):
+            keys = inner.k_norm(keys)
+
+        queries = queries.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.transpose(0, 2, 1, 3)
+
+        # Vectorized RoPE with per-batch offsets (cached on the context).
+        offsets = ctx.offsets
+        attention_pool_idx = ctx.attention_pool_index_by_layer[layer_idx]
+
+        if ctx.aot.rope is not None:
+            # AOT path: real .metallib RoPE + fused KV pool scatter.
+            queries, keys = self._rope_custom_aot(
+                queries,
+                keys,
+                values,
+                offsets,
+                attention_pool_idx,
+                ctx.aot.rope,
+            )
+        else:
+            # Fallback: MLX's built-in mx.fast.rope (used when the AOT kernel
+            # isn't built or the model uses an unsupported RoPE variant).
+            queries = inner.rope(queries, offset=offsets)
+            keys = inner.rope(keys, offset=offsets)
+
+        layer_caches = ctx.attention_layer_caches[attention_pool_idx]
+        pad_sizes = ctx.pad_sizes
+
+        # TODO: replace per-request loop with native batched/ragged
+        # attention once mx.fast.scaled_dot_product_attention supports
+        # variable-length sequences.
+        all_k = []
+        all_v = []
+
+        for i in range(B):
+            layer_caches[i].write_token(keys[i : i + 1], values[i : i + 1])
+
+            k_all, v_all = layer_caches[i].get_kv()
+
+            pad = pad_sizes[i]
+            if pad > 0:
+                k_pad = mx.zeros((1, n_kv_heads, pad, head_dim), dtype=k_all.dtype)
+                v_pad = mx.zeros((1, n_kv_heads, pad, head_dim), dtype=v_all.dtype)
+                k_all = mx.concatenate([k_all, k_pad], axis=2)
+                v_all = mx.concatenate([v_all, v_pad], axis=2)
+
+            all_k.append(k_all)
+            all_v.append(v_all)
+
+        keys_b = mx.concatenate(all_k, axis=0)
+        values_b = mx.concatenate(all_v, axis=0)
+
+        attn_mask = None
+        if ctx.needs_padding:
+            mask_bool = ctx.positions[None, :] >= ctx.valid_lens[:, None]
+            attn_mask = mx.where(
+                mask_bool[:, None, None, :],
+                mx.array(mx.finfo(queries.dtype).min, dtype=queries.dtype),
+                mx.array(0.0, dtype=queries.dtype),
+            )
+
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys_b, values_b, scale=inner.scale, mask=attn_mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, 1, -1)
+        if gate is not None:
+            output = output * mx.sigmoid(gate)
+        return inner.o_proj(output)
+
+    @staticmethod
+    def _rope_custom_aot(
+        queries: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        positions: mx.array,
+        attention_pool_idx: int,
+        rope_ctx: MlxAOTRoPEContext,
+    ) -> tuple[mx.array, mx.array]:
+        """AOT path: rotate Q/K and scatter K/V into the shared pool.
+
+        The kernel call does RoPE on Q/K and scatters
+        rotated K + (untouched) V into ``kv_pool`` at ``new_token_slots``
+        for ``layer_idx``.
+
+        If ``new_token_slots`` is None, slot=-1 sentinel is used (no pool
+        write, RoPE-only mode). Returns rotated (queries, keys) in the
+        original 4-D attention layout. ``values`` is unchanged by RoPE.
+        """
+        # (B, n_heads, 1, head_dim) -> (B, n_heads, head_dim) for kernel
+        q_flat = queries[:, :, 0, :]
+        k_flat = keys[:, :, 0, :]
+        v_flat = values[:, :, 0, :]
+        B = q_flat.shape[0]
+
+        if rope_ctx.new_token_slots is None:
+            slots = mx.full((B,), -1, dtype=mx.int32)
+        else:
+            slots = rope_ctx.new_token_slots.astype(mx.int32)
+
+        k_pool = rope_ctx.kv_pool.k_buffer[attention_pool_idx]
+        v_pool = rope_ctx.kv_pool.v_buffer[attention_pool_idx]
+
+        q_rot, k_rot, k_pool_new, v_pool_new = rope_ctx.kernel.rope_pool_fused(
+            q_flat,
+            k_flat,
+            v_flat,
+            positions,
+            slots,
+            k_pool,
+            v_pool,
+            head_dim=rope_ctx.kernel.config["head_dim"],
+            num_qo_heads=rope_ctx.kernel.config["num_qo_heads"],
+            num_kv_heads=rope_ctx.kernel.config["num_kv_heads"],
+            rope_base=rope_ctx.kernel.base,
+        )
+        # Rebind pool buffers (zero-copy donation result).
+        rope_ctx.kv_pool.k_buffer[attention_pool_idx] = k_pool_new
+        rope_ctx.kv_pool.v_buffer[attention_pool_idx] = v_pool_new
+
+        # (B, n_heads, head_dim) -> (B, n_heads, 1, head_dim) for SDPA path
+        return q_rot[:, :, None, :], k_rot[:, :, None, :]

python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py

@@ -0,0 +1,390 @@
+"""MLX auxiliary-state snapshots for unified radix cache.
+
+Hybrid MLX models may include non-softmax-attention layers whose native
+``mlx-lm`` cache state cannot be reconstructed from the attention KV pool.
+The global scheduler exposes that state through its existing MAMBA component
+contract, so this MLX adapter keeps those scheduler-facing field names while
+storing model-agnostic native cache snapshots.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional
+
+import mlx.core as mx
+import torch
+
+from sglang.srt.mem_cache.base_prefix_cache import EvictParams, InsertResult
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.unified_cache_components.mamba_component import (
+    MambaComponent,
+)
+from sglang.srt.mem_cache.unified_cache_components.tree_component import TreeComponent
+
+_CACHE_ATTRS = ("offset", "lengths", "left_padding")
+_MISSING = object()
+
+
+def _clone_tree(value: Any) -> Any:
+    if isinstance(value, mx.array):
+        return mx.array(value)
+    if isinstance(value, list):
+        return [_clone_tree(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_clone_tree(item) for item in value)
+    if isinstance(value, dict):
+        return {key: _clone_tree(item) for key, item in value.items()}
+    return value
+
+
+def _arrays_in_tree(value: Any) -> list[mx.array]:
+    arrays: list[mx.array] = []
+
+    def collect(item: Any) -> None:
+        if isinstance(item, mx.array):
+            arrays.append(item)
+        elif isinstance(item, (list, tuple)):
+            for child in item:
+                collect(child)
+        elif isinstance(item, dict):
+            for child in item.values():
+                collect(child)
+
+    collect(value)
+    return arrays
+
+
+@dataclass
+class _CacheSnapshot:
+    state: Any
+    meta_state: Any = _MISSING
+    attrs: dict[str, Any] | None = None
+
+
+def _snapshot_cache(cache: Any) -> _CacheSnapshot:
+    state = _clone_tree(getattr(cache, "state", ()))
+    meta_state = (
+        _clone_tree(cache.meta_state) if hasattr(cache, "meta_state") else _MISSING
+    )
+    attrs = {
+        name: _clone_tree(getattr(cache, name))
+        for name in _CACHE_ATTRS
+        if hasattr(cache, name)
+    }
+    arrays = _arrays_in_tree((state, meta_state, attrs))
+    if arrays:
+        mx.eval(*arrays)
+    return _CacheSnapshot(state=state, meta_state=meta_state, attrs=attrs)
+
+
+def _restore_cache(cache: Any, snapshot: _CacheSnapshot) -> None:
+    cache.state = _clone_tree(snapshot.state)
+    if snapshot.meta_state is not _MISSING and hasattr(cache, "meta_state"):
+        cache.meta_state = _clone_tree(snapshot.meta_state)
+    for name, value in (snapshot.attrs or {}).items():
+        setattr(cache, name, _clone_tree(value))
+
+
+class MlxAuxiliaryStatePool:
+    """Index-addressable snapshots of native MLX auxiliary cache state."""
+
+    def __init__(self, size: int, device: str):
+        self.size = size
+        self.device = device
+        self.mamba_cache = None
+        self.mem_usage = 0
+        self._snapshots: dict[int, dict[int, _CacheSnapshot]] = {}
+        self.clear()
+
+    def _tensor(self, indices: Any) -> torch.Tensor:
+        return torch.as_tensor(indices, dtype=torch.int64, device=self.device).view(-1)
+
+    def _index(self, index: Any) -> int:
+        flat = self._tensor(index)
+        assert flat.numel() == 1
+        return int(flat.item())
+
+    def available_size(self) -> int:
+        return int(self.free_slots.numel())
+
+    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
+        if need_size > self.available_size():
+            return None
+        slots = self.free_slots[:need_size].clone()
+        self.free_slots = self.free_slots[need_size:]
+        for slot in slots.tolist():
+            self._snapshots.pop(int(slot), None)
+        return slots
+
+    def free(self, indices: Any) -> None:
+        if indices is None:
+            return
+        indices = self._tensor(indices)
+        if indices.numel() == 0:
+            return
+        for slot in indices.tolist():
+            self._snapshots.pop(int(slot), None)
+        self.free_slots = torch.cat([self.free_slots, indices])
+
+    def clear(self) -> None:
+        self.free_slots = torch.arange(
+            1, self.size + 1, dtype=torch.int64, device=self.device
+        )
+        self._snapshots.clear()
+
+    def copy_from(self, src: Any, dst: Any) -> None:
+        src_indices = self._tensor(src)
+        dst_indices = self._tensor(dst)
+        assert src_indices.numel() == dst_indices.numel()
+        for src_idx, dst_idx in zip(src_indices.tolist(), dst_indices.tolist()):
+            snapshot = self._snapshots.get(int(src_idx))
+            if snapshot is None:
+                self._snapshots.pop(int(dst_idx), None)
+            else:
+                self._snapshots[int(dst_idx)] = {
+                    layer_idx: _CacheSnapshot(
+                        state=_clone_tree(cache_snapshot.state),
+                        meta_state=_clone_tree(cache_snapshot.meta_state),
+                        attrs=_clone_tree(cache_snapshot.attrs),
+                    )
+                    for layer_idx, cache_snapshot in snapshot.items()
+                }
+
+    def fork_from(self, src: Any) -> Optional[torch.Tensor]:
+        src_indices = self._tensor(src)
+        dst = self.alloc(src_indices.numel())
+        if dst is None:
+            return None
+        self.copy_from(src_indices, dst)
+        return dst
+
+    def store_cache(
+        self,
+        index: Any,
+        cache: list[Any],
+        layer_indices: Iterable[int],
+    ) -> None:
+        self._snapshots[self._index(index)] = {
+            layer_idx: _snapshot_cache(cache[layer_idx]) for layer_idx in layer_indices
+        }
+
+    def restore_cache(
+        self,
+        index: Any,
+        cache: list[Any],
+        layer_indices: Iterable[int] | None = None,
+    ) -> bool:
+        snapshot = self._snapshots.get(self._index(index))
+        if snapshot is None:
+            return False
+        selected_layers = set(layer_indices) if layer_indices is not None else None
+        for layer_idx, cache_snapshot in snapshot.items():
+            if selected_layers is not None and layer_idx not in selected_layers:
+                continue
+            _restore_cache(cache[layer_idx], cache_snapshot)
+        return True
+
+    def has_snapshot(self, index: Any) -> bool:
+        return self._index(index) in self._snapshots
+
+
+class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool):
+    """Req-to-token pool with MLX auxiliary-state slot bookkeeping."""
+
+    def __init__(
+        self,
+        *,
+        size: int,
+        max_context_len: int,
+        device: str,
+        enable_memory_saver: bool,
+        auxiliary_state_size: int,
+    ):
+        super().__init__(
+            size=size,
+            max_context_len=max_context_len,
+            device=device,
+            enable_memory_saver=enable_memory_saver,
+        )
+        self.mamba_pool = MlxAuxiliaryStatePool(
+            size=auxiliary_state_size,
+            device=device,
+        )
+        # The unified radix base MAMBA component still reads ``mamba_pool``.
+        # Keep the MLX-owned name beside it so local code can avoid model-
+        # specific terminology.
+        self.auxiliary_state_pool = self.mamba_pool
+        self.enable_mamba_extra_buffer = False
+        self.req_index_to_auxiliary_state_index_mapping = torch.zeros(
+            self._alloc_size, dtype=torch.int32, device=device
+        )
+
+    def alloc(self, reqs):
+        select_index = super().alloc(reqs)
+        if select_index is None:
+            return None
+
+        auxiliary_state_indices = []
+        for req in reqs:
+            if getattr(req, "mamba_pool_idx", None) is not None:
+                mid = req.mamba_pool_idx
+            else:
+                allocated = self.auxiliary_state_pool.alloc(1)
+                assert allocated is not None, "Not enough MLX auxiliary state slots"
+                mid = allocated[0]
+                req.mamba_pool_idx = mid
+            auxiliary_state_indices.append(mid.to(dtype=torch.int32))
+        self.req_index_to_auxiliary_state_index_mapping[select_index] = torch.stack(
+            auxiliary_state_indices
+        )
+        return select_index
+
+    def get_auxiliary_state_indices(self, req_indices) -> torch.Tensor:
+        return self.req_index_to_auxiliary_state_index_mapping[req_indices]
+
+    def get_mamba_indices(self, req_indices) -> torch.Tensor:
+        return self.get_auxiliary_state_indices(req_indices)
+
+    def get_mamba_ping_pong_other_idx(self, mamba_next_track_idx: int) -> int:
+        return 0
+
+    def free_mamba_cache(self, req, mamba_ping_pong_track_buffer_to_keep=None):
+        if getattr(req, "mamba_pool_idx", None) is not None:
+            self.auxiliary_state_pool.free(req.mamba_pool_idx.unsqueeze(0))
+            req.mamba_pool_idx = None
+        track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+        if track_buffer is not None:
+            if mamba_ping_pong_track_buffer_to_keep is None:
+                self.auxiliary_state_pool.free(track_buffer)
+            req.mamba_ping_pong_track_buffer = None
+            req.mamba_next_track_idx = None
+
+    def free_auxiliary_state_cache(self, req, track_buffer_to_keep=None):
+        self.free_mamba_cache(
+            req,
+            mamba_ping_pong_track_buffer_to_keep=track_buffer_to_keep,
+        )
+
+    def free(self, req):
+        super().free(req)
+
+    def clear(self):
+        super().clear()
+        self.auxiliary_state_pool.clear()
+        self.req_index_to_auxiliary_state_index_mapping.zero_()
+
+
+class MlxAuxiliaryStateComponent(MambaComponent):
+    """Unified radix component for MLX native auxiliary-state snapshots."""
+
+    def __init__(self, cache, params):
+        if params.enable_mamba_extra_buffer:
+            raise NotImplementedError(
+                "MLX auxiliary-state radix cache does not support "
+                "enable_mamba_extra_buffer yet."
+            )
+        pool = getattr(cache.req_to_token_pool, "auxiliary_state_pool", None)
+        if not isinstance(pool, MlxAuxiliaryStatePool):
+            raise TypeError(
+                "MlxAuxiliaryStateComponent requires MlxAuxiliaryStatePool, "
+                f"got {type(pool)}"
+            )
+        TreeComponent.__init__(self, cache, params)
+        self.enable_mamba_extra_buffer = False
+        self._mamba_pool_host = None
+
+    @staticmethod
+    def _tracked_value(req) -> tuple[object | None, bool]:
+        track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+        track_len = getattr(req, "mamba_last_track_seqlen", None)
+        if track_buffer is not None and track_len is not None:
+            return track_buffer[0].unsqueeze(-1).clone(), True
+        if getattr(req, "mamba_pool_idx", None) is None:
+            return None, False
+        return req.mamba_pool_idx.unsqueeze(-1).clone(), False
+
+    def prepare_for_caching_req(
+        self,
+        req,
+        insert_params,
+        token_ids_len: int,
+        is_finished: bool,
+    ) -> int | None:
+        cache_len = getattr(req, "mamba_last_track_seqlen", None)
+        auxiliary_value, uses_track_slot = self._tracked_value(req)
+        setattr(insert_params, "mlx_auxiliary_state_uses_track_slot", uses_track_slot)
+
+        if auxiliary_value is None:
+            return 0 if is_finished else None
+
+        if cache_len is None:
+            cache_len = token_ids_len
+        if is_finished:
+            insert_params.mamba_value = auxiliary_value
+        else:
+            source_value = auxiliary_value
+            forked_value = self.cache.req_to_token_pool.auxiliary_state_pool.fork_from(
+                source_value
+            )
+            if forked_value is None:
+                self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
+                forked_value = (
+                    self.cache.req_to_token_pool.auxiliary_state_pool.fork_from(
+                        source_value
+                    )
+                )
+                assert forked_value is not None, "Can not alloc MLX auxiliary cache"
+            insert_params.mamba_value = forked_value
+        return cache_len
+
+    def cleanup_after_caching_req(
+        self,
+        req,
+        is_finished: bool,
+        insert_result: InsertResult | None = None,
+        insert_params=None,
+    ) -> None:
+        if not is_finished:
+            if (
+                insert_params is not None
+                and insert_params.mamba_value is not None
+                and (insert_result is None or insert_result.mamba_exist)
+            ):
+                self.cache.req_to_token_pool.auxiliary_state_pool.free(
+                    insert_params.mamba_value
+                )
+            if bool(
+                getattr(insert_params, "mlx_auxiliary_state_uses_track_slot", False)
+            ):
+                track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+                if track_buffer is not None:
+                    self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer)
+                req.mamba_ping_pong_track_buffer = None
+                req.mamba_next_track_idx = None
+            req.mamba_last_track_seqlen = None
+            return
+
+        auxiliary_value_exists = (
+            insert_result.mamba_exist if insert_result is not None else True
+        )
+        uses_track_slot = bool(
+            getattr(insert_params, "mlx_auxiliary_state_uses_track_slot", False)
+        )
+        if uses_track_slot:
+            keep_track_slot = not auxiliary_value_exists
+            self.cache.req_to_token_pool.free_auxiliary_state_cache(
+                req,
+                track_buffer_to_keep=0 if keep_track_slot else None,
+            )
+        elif auxiliary_value_exists:
+            self.cache.req_to_token_pool.free_auxiliary_state_cache(req)
+        else:
+            # The radix tree now owns the live auxiliary-state slot.
+            track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
+            if track_buffer is not None:
+                self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer)
+                req.mamba_ping_pong_track_buffer = None
+                req.mamba_next_track_idx = None
+            req.mamba_pool_idx = None
+        req.mamba_last_track_seqlen = None

python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py

@@ -0,0 +1,93 @@
+"""Model cache layout helpers for the MLX backend."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Sequence
+
+
+@dataclass(frozen=True)
+class MlxModelCacheLayout:
+    """Map model layers to MLX cache storage components.
+
+    Attention layers store softmax-attention KV in the MLX attention KV pool.
+    Auxiliary layers keep native ``mlx-lm`` cache state and are snapshotted by
+    the MLX auxiliary-state component.
+    """
+
+    layers: tuple[Any, ...]
+    attention_attrs: tuple[str | None, ...]
+    attention_layer_indices: tuple[int, ...]
+    auxiliary_layer_indices: tuple[int, ...]
+    attention_pool_index_by_layer: dict[int, int]
+
+    @classmethod
+    def from_attention_discovery(
+        cls,
+        layers: Sequence[Any],
+        attention_attrs: Sequence[str | None],
+    ) -> MlxModelCacheLayout:
+        if len(layers) != len(attention_attrs):
+            raise ValueError(
+                "Layer count and attention attribute count differ: "
+                f"{len(layers)} != {len(attention_attrs)}"
+            )
+
+        attention_layer_indices = tuple(
+            idx for idx, attr in enumerate(attention_attrs) if attr is not None
+        )
+        auxiliary_layer_indices = tuple(
+            idx for idx, attr in enumerate(attention_attrs) if attr is None
+        )
+        attention_pool_index_by_layer = {
+            layer_idx: pool_idx
+            for pool_idx, layer_idx in enumerate(attention_layer_indices)
+        }
+
+        return cls(
+            layers=tuple(layers),
+            attention_attrs=tuple(attention_attrs),
+            attention_layer_indices=attention_layer_indices,
+            auxiliary_layer_indices=auxiliary_layer_indices,
+            attention_pool_index_by_layer=attention_pool_index_by_layer,
+        )
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.layers)
+
+    @property
+    def num_attention_layers(self) -> int:
+        return len(self.attention_layer_indices)
+
+    @property
+    def has_auxiliary_state(self) -> bool:
+        return bool(self.auxiliary_layer_indices)
+
+    @property
+    def first_attention_layer_index(self) -> int:
+        if not self.attention_layer_indices:
+            raise RuntimeError("MLX model has no supported attention layers")
+        return self.attention_layer_indices[0]
+
+    def attention_pool_index(self, layer_idx: int) -> int:
+        try:
+            return self.attention_pool_index_by_layer[layer_idx]
+        except KeyError as exc:
+            raise KeyError(f"Layer {layer_idx} is not an attention layer") from exc
+
+    def attention_attr(self, layer_idx: int) -> str:
+        attr = self.attention_attrs[layer_idx]
+        if attr is None:
+            raise KeyError(f"Layer {layer_idx} is not an attention layer")
+        return attr
+
+    def attention_layer_caches(
+        self,
+        caches_by_request: list[list[Any]],
+    ) -> list[list[Any]]:
+        """Return layer-major attention caches for batched decode."""
+        return [
+            [request_cache[layer_idx] for request_cache in caches_by_request]
+            for layer_idx in self.attention_layer_indices
+        ]

python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py

@@ -0,0 +1,61 @@
+"""Model introspection and attention patching."""
+
+from typing import Any
+
+import mlx.nn as nn
+
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
+    is_attention_module,
+)
+from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import (
+    MLXAttentionWrapper,
+)
+
+
+def _find_attention_attr(layer: Any) -> str | None:
+    """Return the direct child name that satisfies the attention contract."""
+    if not isinstance(layer, nn.Module):
+        raise TypeError(f"Expected mlx.nn.Module layer, got {type(layer)}")
+    for name, module in layer.children().items():
+        if isinstance(module, MLXAttentionWrapper) or is_attention_module(module):
+            return name
+    return None
+
+
+def find_attention_layers(model: Any) -> tuple[list[Any], list[str | None]]:
+    """Find transformer layers and per-layer attention attribute names."""
+    root = getattr(model, "language_model", model)
+    container = getattr(root, "model", root)
+    layer_list = getattr(container, "layers", None) or getattr(root, "layers", [])
+
+    if layer_list:
+        attn_attrs = [_find_attention_attr(layer) for layer in layer_list]
+        if any(attr is not None for attr in attn_attrs):
+            return layer_list, attn_attrs
+        raise ValueError(f"No attention attribute in layer type {type(layer_list[0])}")
+    return layer_list, []
+
+
+def patch_model_attention(model: Any) -> int:
+    """Install MLXAttentionWrapper on all attention layers (idempotent).
+
+    The wrapper delegates to the inner module when no BatchedDecodeContext
+    is set, so it is always installed and never removed.
+    """
+    layer_list, attn_attrs = find_attention_layers(model)
+    patched = 0
+    for idx, (layer, attn_attr) in enumerate(zip(layer_list, attn_attrs)):
+        if attn_attr is None:
+            continue
+        attn = getattr(layer, attn_attr)
+        if isinstance(attn, MLXAttentionWrapper):
+            continue
+        setattr(layer, attn_attr, MLXAttentionWrapper(attn, idx))
+        patched += 1
+    return patched
+
+
+def get_num_layers(model: Any) -> int:
+    """Return the number of transformer layers."""
+    layer_list, _ = find_attention_layers(model)
+    return len(layer_list)

python/sglang/srt/hardware_backend/mlx/model_runner_stub.py

@@ -0,0 +1,210 @@
+"""Lightweight ModelRunner stub for MLX on Apple Silicon.
+
+Skips PyTorch weight loading.  Creates only the CPU-side bookkeeping
+(req_to_token_pool, token_to_kv_pool_allocator) the scheduler needs.
+"""
+
+import logging
+from typing import Tuple
+
+import torch
+
+from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import (
+    MlxAuxiliaryStateReqToTokenPool,
+)
+from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
+from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
+from sglang.srt.model_executor.model_runner import ModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+class _DummyKVCache(KVCache):
+    """Scheduler-facing KV cache that allocates no GPU memory.
+
+    Satisfies the KVCache interface so that TokenToKVPoolAllocator can be
+    constructed, but every buffer access raises. The MLX backend manages
+    attention KV and auxiliary state internally.
+    """
+
+    def __init__(self, size: int, dtype: torch.dtype, device: str):
+        # Bypass KVCache.__init__ to avoid custom_mem_pool / memory_saver
+        # initialization that may touch CUDA APIs.
+        self.size = size
+        self.page_size = 1
+        self.dtype = dtype
+        self.store_dtype = dtype
+        self.device = device
+        self.layer_num = 0
+        self.start_layer = 0
+        self.end_layer = 0
+        self.mem_usage = 0
+        self.cpu_offloading_chunk_size = 8192
+        self.layer_transfer_counter = None
+        self.enable_custom_mem_pool = False
+        self.custom_mem_pool = None
+
+    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
+        raise RuntimeError("_DummyKVCache has no key buffer (MLX manages cache)")
+
+    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
+        raise RuntimeError("_DummyKVCache has no value buffer (MLX manages cache)")
+
+    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise RuntimeError("_DummyKVCache has no kv buffer (MLX manages cache)")
+
+    def set_kv_buffer(self, layer, loc, cache_k, cache_v) -> None:
+        raise RuntimeError("_DummyKVCache cannot set kv buffer (MLX manages cache)")
+
+    def get_kv_size_bytes(self):
+        return 0, 0
+
+
+class _DummyModel:
+    """Minimal stand-in so that `inspect.signature(model.forward)` and
+    `getattr(model, ...)` calls in ModelRunner.__init__ don't crash."""
+
+    @staticmethod
+    def forward():
+        pass
+
+
+class MlxModelRunnerStub(ModelRunner):
+    """ModelRunner that skips PyTorch weight loading and KV cache allocation.
+
+    Overrides both load_model() and initialize() so that no PyTorch model
+    weights are loaded and no large KV cache tensors are allocated.  Only
+    the minimal bookkeeping pools needed by the scheduler are created.
+    """
+
+    # No KV canary on the MLX path. The base ModelRunner installs it via
+    # install_canary() in its full initialize(), which this lightweight override
+    # skips. Downstream consumers (scheduler, cuda graph runner, speculative
+    # workers) all guard with `canary_manager is not None`, so default to None
+    # as a class attribute to keep those checks working instead of raising
+    # AttributeError.
+    canary_manager = None
+
+    # No prefill-aware SWA on the MLX path. The base ModelRunner derives this in
+    # its full initialize() from `model.is_prefill_aware_swa()`, which this
+    # lightweight override skips (and `_DummyModel` does not implement). The
+    # scheduler reads `model_runner.prefill_aware_swa` unconditionally when
+    # admitting a prefill batch, so default to False as a class attribute to keep
+    # that path working instead of raising AttributeError.
+    prefill_aware_swa = False
+
+    def __init__(self, *args, mlx_pool_size: int | None = None, **kwargs):
+        self._mlx_pool_size = mlx_pool_size
+        super().__init__(*args, **kwargs)
+
+    def load_model(self):
+        """Set only the metadata that downstream code needs, without
+        loading any PyTorch model weights."""
+        logger.info(
+            "MLX stub: skipping PyTorch model weight loading "
+            "(inference runs through MLX)"
+        )
+
+        self.model = _DummyModel()
+
+        self.sliding_window_size = None
+        if (
+            self.model_config.is_hybrid_swa
+            and self.model_config.sliding_window_size is not None
+        ):
+            self.sliding_window_size = self.model_config.sliding_window_size
+        elif self.model_config.attention_chunk_size is not None:
+            self.sliding_window_size = self.model_config.attention_chunk_size
+
+        self.dtype = self.model_config.dtype
+        self.weight_load_mem_usage = 0
+
+    def initialize(self):
+        """Lightweight initialize that skips heavy PyTorch setup.
+
+        Creates minimal req_to_token_pool and token_to_kv_pool_allocator
+        with a dummy KV cache (zero GPU memory) so the scheduler works.
+        """
+        from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=self.server_args.enable_memory_saver
+        )
+
+        # Load model (sets metadata only)
+        self.sampler = None
+        self.load_model()
+
+        # Layer metadata
+        model_num_layers = max(
+            self.model_config.num_hidden_layers,
+            self.model_config.num_attention_layers,
+        )
+        self.start_layer = 0
+        self.end_layer = model_num_layers
+        self.num_effective_layers = model_num_layers
+
+        # KV cache dtype
+        self.kv_cache_dtype = self.dtype
+
+        # Pool sizing — use the MLX runner's auto-sized pool if available,
+        # otherwise fall back to context_len.
+        if self._mlx_pool_size is not None:
+            self.max_total_num_tokens = self._mlx_pool_size
+        else:
+            self.max_total_num_tokens = self.model_config.context_len
+        self.max_running_requests = min(
+            self.max_total_num_tokens // 2,
+            4096,
+        )
+        self.is_hybrid_swa = False
+
+        # Create minimal pools
+        if self.mambaish_config is not None:
+            auxiliary_state_size = self.server_args.max_mamba_cache_size
+            if auxiliary_state_size is None:
+                auxiliary_state_size = self.max_running_requests * 4
+            self.req_to_token_pool = MlxAuxiliaryStateReqToTokenPool(
+                size=self.max_running_requests,
+                max_context_len=self.model_config.context_len,
+                device="cpu",
+                enable_memory_saver=False,
+                auxiliary_state_size=auxiliary_state_size,
+            )
+        else:
+            self.req_to_token_pool = ReqToTokenPool(
+                size=self.max_running_requests,
+                max_context_len=self.model_config.context_len,
+                device="cpu",
+                enable_memory_saver=False,
+            )
+
+        dummy_kv = _DummyKVCache(
+            size=self.max_total_num_tokens,
+            dtype=self.kv_cache_dtype,
+            device="cpu",
+        )
+        self.token_to_kv_pool = dummy_kv
+        self.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
+            size=self.max_total_num_tokens,
+            dtype=self.kv_cache_dtype,
+            device="cpu",
+            kvcache=dummy_kv,
+            need_sort=False,
+        )
+
+        # No CUDA graphs, no attention backend
+        self.decode_cuda_graph_runner = None
+        self.graph_mem_usage = 0
+        self.attn_backend = None
+
+        logger.info(
+            f"MLX stub: initialized minimal pools "
+            f"(max_total_num_tokens={self.max_total_num_tokens}, "
+            f"max_running_requests={self.max_running_requests}, "
+            f"zero GPU KV cache allocation)"
+        )
+
+    def alloc_memory_pool(self, memory_pool_config=None):
+        """No-op: MLX manages its own KV cache."""
+        pass

python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py

@@ -0,0 +1,573 @@
+"""Path B fusion for SwitchGLU: gate gather_qmv with silu(gate) * x_up epilogue.
+
+Why this exists
+---------------
+The existing `FusedSwitchUpGate` (fused_switch_glu.py) concatenates up_proj
+and gate_proj weights along the output dim and runs one gather_qmm. That saves
+one kernel launch per layer but doubles the matmul's output dim, which pushes
+MLX's quantized GEMV into a worse tile/occupancy config. At bs >= 4 on
+Qwen3-30B-A3B-4bit this is a net regression (~2% slower at bs=32).
+
+Path B keeps up_proj and gate_proj separate (matmul kernels see their natural
+N — no tile regression) and instead fuses the *activation* into the gate
+matmul. Concretely:
+
+    Baseline (3 kernels per MoE layer in the swiglu front-half):
+        x_up   = gather_qmm(x, W_up)
+        x_gate = gather_qmm(x, W_gate)
+        out    = silu(x_gate) * x_up         # 1 compiled kernel via mlx_lm.swiglu
+
+    Path B (2 kernels):
+        x_up = gather_qmm(x, W_up)
+        out  = fused_gate_qmv_silu_mul(x, W_gate, ..., x_up)
+                                              # one custom Metal kernel
+
+This removes one kernel launch per MoE layer with no change in matmul kernel
+shapes. The measured end to end decode impact is within run to run noise (bs=1,
+K=12 interleaved trials on Qwen3-30B-A3B-4bit: on minus off 0.4%, a quarter of
+the noise band), so v1 lands off by default as a correct fusion substrate, not
+a measured speedup.
+
+Scope of v1
+-----------
+Targets the configuration shared by Qwen3-30B-A3B-4bit and Qwen1.5-MoE-A2.7B-4bit:
+- bits=4, mode='affine', group_size=64
+- K (input_dim) divisible by 512  (Qwen3: 2048, Qwen1.5: 2048)
+- N (output_dim) divisible by 8   (Qwen3: 768, Qwen1.5: 1408)
+- Scales/biases dtype matches the input dtype (bf16 or fp16)
+
+Anything outside that falls back to the unfused mlx_lm path.
+
+When the fast `gather_qmv` from #22283 lands, the qmv inner loop here can be
+replaced by a call into it; the epilogue (silu * x_up) doesn't change.
+"""
+
+from __future__ import annotations
+
+import logging
+import weakref
+
+import mlx.core as mx
+import mlx.nn as nn
+
+logger = logging.getLogger(__name__)
+
+# Constants matching MLX's affine_qmv_fast for bits=4, group_size=64.
+# Lifted directly from mlx/include/mlx/backend/metal/kernels/quantized.h
+# (qmv_fast_impl), so the inner-loop layout matches MLX's own gather_qmm in
+# that regime.
+_BITS = 4
+_GROUP_SIZE = 64
+_SIMD_SIZE = 32
+_PACK_FACTOR = 8  # 32 / bits
+_BYTES_PER_PACK = 4  # sizeof(uint32_t)
+_PACKS_PER_THREAD = 2
+_NUM_SIMDGROUPS = 2
+_RESULTS_PER_SIMDGROUP = 4
+_VALUES_PER_THREAD = _PACK_FACTOR * _PACKS_PER_THREAD  # 16
+_BLOCK_SIZE = _VALUES_PER_THREAD * _SIMD_SIZE  # 512
+_ROWS_PER_TG = _NUM_SIMDGROUPS * _RESULTS_PER_SIMDGROUP  # 8
+
+
+# Metal source for the fused kernel.
+# Body only — mx.fast.metal_kernel auto-generates the kernel signature
+# based on input_names / output_names and the template params below.
+_KERNEL_SOURCE = r"""
+    // Mirrors qmv_fast_impl<T, group_size=64, bits=4> from MLX's quantized.h
+    // with a silu(result) * x_up write epilogue.
+    //
+    // Inputs:
+    //   x       [M_tok, K]                — pre-gather activations (T)
+    //   w       [E, N, K * 4 / 32]        — packed 4-bit weights (uint32)
+    //   s       [E, N, K / GROUP_SIZE]    — affine scales (T)
+    //   b       [E, N, K / GROUP_SIZE]    — affine biases (T)
+    //   idx     [M_tok * TOPK]            — expert per (token, topk) pair (uint32)
+    //   x_up    [M_tok * TOPK, N]         — precomputed up output (T)
+    // Output:
+    //   y       [M_tok * TOPK, N]         — silu(gate_qmv(x)) * x_up
+
+    constexpr int BITS = 4;
+    constexpr int GROUP_SIZE = 64;
+    constexpr int SIMD_SIZE = 32;
+    constexpr int PACK_FACTOR = 8;
+    constexpr int BYTES_PER_PACK = 4;
+    constexpr int PACKS_PER_THREAD = 2;
+    constexpr int NUM_SIMDGROUPS = 2;
+    constexpr int RESULTS_PER_SIMDGROUP = 4;
+    constexpr int VALUES_PER_THREAD = PACK_FACTOR * PACKS_PER_THREAD;   // 16
+    constexpr int BLOCK_SIZE = VALUES_PER_THREAD * SIMD_SIZE;            // 512
+    constexpr int SCALE_STEP_PER_THREAD = GROUP_SIZE / VALUES_PER_THREAD; // 4
+
+    // Compile-time problem dims (template params)
+    constexpr int K = IN_VEC_SIZE;
+    constexpr int N = OUT_VEC_SIZE;
+    constexpr int TOPK = TOP_K;
+
+    constexpr int in_vec_size_w = K * BYTES_PER_PACK / PACK_FACTOR;    // K/2 bytes per row
+    constexpr int in_vec_size_g = K / GROUP_SIZE;                       // groups per row
+
+    // tid.x = (token, topk) pair index in M_tok*TOPK
+    // tid.y = output row block index (each block writes 8 rows)
+    uint mt = threadgroup_position_in_grid.x;
+    uint out_row_block = threadgroup_position_in_grid.y;
+    uint simd_gid = simdgroup_index_in_threadgroup;
+    uint simd_lid = thread_index_in_simdgroup;
+
+    uint m = mt / TOPK;                      // index into pre-gather x
+    uint e = idx[mt];                        // expert id
+
+    // Base pointers for this (expert) row block.
+    // Weights / scales / biases live in [E, N, ...] tensors.
+    const device uint8_t* ws = (const device uint8_t*)(w
+        + uint64_t(e) * uint64_t(N) * uint64_t(in_vec_size_w / BYTES_PER_PACK));
+    const device T* scales_p = s
+        + uint64_t(e) * uint64_t(N) * uint64_t(in_vec_size_g);
+    const device T* biases_p = b
+        + uint64_t(e) * uint64_t(N) * uint64_t(in_vec_size_g);
+
+    // out_row indexes the first of RESULTS_PER_SIMDGROUP rows this simdgroup handles.
+    int out_row = out_row_block * NUM_SIMDGROUPS * RESULTS_PER_SIMDGROUP
+                + simd_gid * RESULTS_PER_SIMDGROUP;
+
+    ws += out_row * in_vec_size_w + simd_lid * PACKS_PER_THREAD * BYTES_PER_PACK;
+    scales_p += out_row * in_vec_size_g + simd_lid / SCALE_STEP_PER_THREAD;
+    biases_p += out_row * in_vec_size_g + simd_lid / SCALE_STEP_PER_THREAD;
+
+    // Pre-gather x: shape [M_tok, K]
+    const device T* x_p = x + uint64_t(m) * uint64_t(K)
+                            + uint64_t(simd_lid) * uint64_t(VALUES_PER_THREAD);
+
+    float result[RESULTS_PER_SIMDGROUP] = {0, 0, 0, 0};
+    thread float x_thread[VALUES_PER_THREAD];
+
+    // Outer loop over K in BLOCK_SIZE-wide chunks.
+    // K % BLOCK_SIZE == 0 is required (checked in Python wrapper).
+    for (int k = 0; k < K; k += BLOCK_SIZE) {
+        // --- load_vector for bits=4 ---
+        float sum = 0;
+        for (int i = 0; i < VALUES_PER_THREAD; i += 4) {
+            float a0 = float(x_p[i]);
+            float a1 = float(x_p[i + 1]);
+            float a2 = float(x_p[i + 2]);
+            float a3 = float(x_p[i + 3]);
+            sum += a0 + a1 + a2 + a3;
+            x_thread[i]     = a0;
+            x_thread[i + 1] = a1 / 16.0f;
+            x_thread[i + 2] = a2 / 256.0f;
+            x_thread[i + 3] = a3 / 4096.0f;
+        }
+
+        // For each of the 4 output rows this simdgroup is responsible for...
+        for (int row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
+            const device uint16_t* ws_u16 =
+                (const device uint16_t*)(ws + row * in_vec_size_w);
+            float scale_v = float(scales_p[row * in_vec_size_g]);
+            float bias_v  = float(biases_p[row * in_vec_size_g]);
+
+            // --- qdot for bits=4, values_per_thread=16 ---
+            float accum = 0;
+            for (int i = 0; i < VALUES_PER_THREAD / 4; i++) {
+                uint16_t packed = ws_u16[i];
+                accum += (x_thread[4 * i]     * float(packed & 0x000f) +
+                          x_thread[4 * i + 1] * float(packed & 0x00f0) +
+                          x_thread[4 * i + 2] * float(packed & 0x0f00) +
+                          x_thread[4 * i + 3] * float(packed & 0xf000));
+            }
+            result[row] += scale_v * accum + sum * bias_v;
+        }
+
+        ws        += BLOCK_SIZE * BYTES_PER_PACK / PACK_FACTOR;  // += 256 bytes
+        scales_p  += BLOCK_SIZE / GROUP_SIZE;                    // += 8 groups
+        biases_p  += BLOCK_SIZE / GROUP_SIZE;
+        x_p       += BLOCK_SIZE;
+    }
+
+    // Write epilogue: simd-sum across lanes, then silu(gate) * x_up.
+    device T* y_p = y + uint64_t(mt) * uint64_t(N) + uint64_t(out_row);
+    const device T* x_up_p = x_up + uint64_t(mt) * uint64_t(N) + uint64_t(out_row);
+
+    for (int row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
+        float gate_v = simd_sum(result[row]);
+        if (simd_lid == 0) {
+            // silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
+            float silu_v = gate_v / (1.0f + metal::precise::exp(-gate_v));
+            y_p[row] = T(silu_v * float(x_up_p[row]));
+        }
+    }
+"""
+
+
+# Build kernel lazily so import-time on non-MLX systems doesn't fail.
+_kernel_cache: dict = {}
+
+# AOT pre-compile cache: tracks which (dtype, K, N, T) tuples have been warmed.
+# MLX specializes the Metal kernel on template args at first dispatch, so the
+# first call per unique tuple pays a ~3ms compile cost. Pre-warming at patch
+# time moves that cost out of the first forward pass.
+_aot_warmed: set = set()
+
+
+def _get_kernel(dtype: mx.Dtype):
+    """Return a compiled mx.fast.metal_kernel for the given input dtype.
+
+    Kept per-dtype because mx.fast.metal_kernel specializes on template args
+    at first call, and we want clean separation between fp16 / bf16 variants.
+    """
+    if dtype not in _kernel_cache:
+        # Strip the `mlx.core.` prefix and any dots from the dtype repr so the
+        # kernel name is a valid C identifier (Metal's host_name attribute and
+        # function-name slot don't accept '.').
+        dtype_tag = str(dtype).replace("mlx.core.", "").replace(".", "_")
+        _kernel_cache[dtype] = mx.fast.metal_kernel(
+            name=f"affine_gather_qmv_silu_mul_4bit_gs64_{dtype_tag}",
+            input_names=["x", "w", "s", "b", "idx", "x_up"],
+            output_names=["y"],
+            source=_KERNEL_SOURCE,
+        )
+    return _kernel_cache[dtype]
+
+
+def fused_gate_qmv_silu_mul(
+    x: mx.array,
+    gate_w: mx.array,
+    gate_s: mx.array,
+    gate_b: mx.array,
+    indices: mx.array,
+    x_up: mx.array,
+) -> mx.array:
+    """Compute ``silu(gather_qmm(x, W_gate)) * x_up`` in one kernel.
+
+    Shapes:
+        x       : (..., 1, 1, K)               input activations (pre-gather)
+        gate_w  : (E, N, K // PACK_FACTOR)     packed 4-bit weights
+        gate_s  : (E, N, K // GROUP_SIZE)      affine scales
+        gate_b  : (E, N, K // GROUP_SIZE)      affine biases
+        indices : (..., T)                     expert per (token, topk) pair
+        x_up    : (..., T, 1, N)               pre-computed up output
+        y       : (..., T, 1, N)               returned
+
+    Numerical contract: equivalent within floating point tolerance to::
+
+        x_gate = mx.gather_qmm(x, gate_w, gate_s, gate_b, rhs_indices=indices,
+                               transpose=True, group_size=GROUP_SIZE, bits=BITS,
+                               mode='affine')
+        y = nn.silu(x_gate) * x_up
+
+    up to floating-point ordering of accumulations (which matches MLX's own
+    qmv_fast_impl exactly).
+    """
+    # Validate the regime this kernel supports. Outside it, the caller should
+    # fall back to the unfused MLX path.
+    K = gate_s.shape[-1] * _GROUP_SIZE
+    N = gate_w.shape[-2]
+    if K % _BLOCK_SIZE != 0:
+        raise ValueError(
+            f"fused_gate_qmv_silu_mul: K={K} not divisible by {_BLOCK_SIZE}. "
+            f"Use the unfused path."
+        )
+    if N % _ROWS_PER_TG != 0:
+        raise ValueError(
+            f"fused_gate_qmv_silu_mul: N={N} not divisible by {_ROWS_PER_TG}."
+        )
+    # Sanity: scales/biases dtype must match x dtype for in-kernel float() conversion.
+    if gate_s.dtype != x.dtype or gate_b.dtype != x.dtype:
+        raise ValueError(
+            f"fused_gate_qmv_silu_mul: dtype mismatch x={x.dtype} "
+            f"s={gate_s.dtype} b={gate_b.dtype}"
+        )
+
+    # Shape handling: x always has K as its last axis. M_tok is the number of
+    # distinct pre-gather tokens (= x.size // K). T is the top_k axis carried
+    # by `indices`. In the unsorted SwitchGLU path x has shape (B, 1, 1, K) and
+    # indices is (B, T); in the sorted path x has shape (B*T, 1, K) and
+    # indices is (B*T, 1) after our reshape. Either way: M_tok * T == idx.size
+    # and M_tok * K == x.size.
+    assert x.shape[-1] == K, f"x last dim {x.shape[-1]} != K={K}"
+    M_tok = x.size // K
+    T = indices.shape[-1]
+    assert (
+        M_tok * T == indices.size
+    ), f"M_tok({M_tok}) * T({T}) != indices.size({indices.size})"
+    x_flat = x.reshape(M_tok, K)
+    idx_flat = indices.reshape(M_tok * T)
+    if idx_flat.dtype != mx.uint32:
+        idx_flat = idx_flat.astype(mx.uint32)
+
+    # x_up has N as its last axis and total size M_tok * T * N. The singleton
+    # rank dims (1 or 2 of them) get folded away by reshape.
+    assert (
+        x_up.shape[-1] == N and x_up.size == M_tok * T * N
+    ), f"x_up shape {x_up.shape} does not match M_tok({M_tok})*T({T})*N({N})"
+    x_up_flat = x_up.reshape(M_tok * T, N)
+
+    kernel = _get_kernel(x.dtype)
+    (y_flat,) = kernel(
+        inputs=[x_flat, gate_w, gate_s, gate_b, idx_flat, x_up_flat],
+        template=[
+            ("T", x.dtype),
+            ("IN_VEC_SIZE", K),
+            ("OUT_VEC_SIZE", N),
+            ("TOP_K", T),
+        ],
+        # grid is in *threads*, not threadgroups: total threads = product of
+        # (grid_x, grid_y, grid_z). One threadgroup processes one (mt, row_block).
+        # Threadgroup is 64 = 2 simdgroups × 32 lanes.
+        grid=(M_tok * T * 64, N // _ROWS_PER_TG, 1),
+        threadgroup=(64, 1, 1),
+        output_shapes=[(M_tok * T, N)],
+        output_dtypes=[x.dtype],
+    )
+
+    # Reshape to x_up's shape, which is exactly what self.activation(x_up,
+    # x_gate) would have returned (silu*mul is shape-preserving).
+    return y_flat.reshape(x_up.shape)
+
+
+def _aot_warm_kernel(switch_mlp, top_k: int) -> None:
+    """Pre-compile the fused kernel for the (dtype, K, N, T) tuples this
+    layer will dispatch at runtime.
+
+    MLX's mx.fast.metal_kernel specializes Metal source on template args
+    at first dispatch (one Metal compile per unique tuple, ~3ms each).
+    Issuing one dummy dispatch per shape moves that compile out of the
+    first forward pass and into model init. The module-level _aot_warmed
+    set means only the first layer per shape actually compiles; the
+    remaining 47 hit the cache and no-op.
+
+    Warms both the unsorted decode shape (T=top_k) and the sorted
+    large-batch shape (T=1) used by the gather-sort path.
+    """
+    gate = switch_mlp.gate_proj
+    K = gate.scales.shape[-1] * _GROUP_SIZE
+    N = gate.weight.shape[-2]
+    dtype = gate.scales.dtype
+    for T in (top_k, 1):
+        key = (dtype, K, N, T)
+        if key in _aot_warmed:
+            continue
+        M_tok = 1
+        x_dummy = mx.zeros((M_tok, 1, 1, K), dtype=dtype)
+        idx_dummy = mx.zeros((M_tok, T), dtype=mx.uint32)
+        x_up_dummy = mx.zeros((M_tok, T, 1, N), dtype=dtype)
+        out = fused_gate_qmv_silu_mul(
+            x_dummy,
+            gate["weight"],
+            gate["scales"],
+            gate.get("biases"),
+            idx_dummy,
+            x_up_dummy,
+        )
+        mx.eval(out)
+        _aot_warmed.add(key)
+        logger.info(
+            "Path B AOT: warmed fused kernel dtype=%s K=%d N=%d T=%d",
+            dtype,
+            K,
+            N,
+            T,
+        )
+
+
+def can_fuse(switch_mlp) -> bool:
+    """Cheap structural check: does this SwitchGLU match the Path B v1 regime?"""
+    try:
+        from mlx_lm.models.switch_layers import (
+            QuantizedSwitchLinear,
+            SwiGLU,
+            SwitchGLU,
+        )
+    except ImportError:
+        return False
+    up = switch_mlp.up_proj
+    gate = switch_mlp.gate_proj
+    if not isinstance(up, QuantizedSwitchLinear) or not isinstance(
+        gate, QuantizedSwitchLinear
+    ):
+        return False
+    # A model that overrides SwitchGLU.__call__ runs a custom forward, but the
+    # patch installs a subclass __call__ that imposes the stock semantics
+    # fused_forward reimplements, silently bypassing the override. Decline when
+    # the forward is not the stock SwitchGLU.__call__ (evaluated at patch time,
+    # before the class swap, so this sees the model's real class).
+    if type(switch_mlp).__call__ is not SwitchGLU.__call__:
+        return False
+    # fused_forward bakes silu into both the kernel and its fallback, and a
+    # swapped activation= leaves __call__ stock, so the check above cannot see
+    # it. Exact type, fail closed: a SwiGLU subclass may change the math.
+    if type(getattr(switch_mlp, "activation", None)) is not SwiGLU:
+        return False
+    # Learned per-expert bias, added after the matmul in
+    # QuantizedSwitchLinear.__call__ as ``x + bias[indices]`` whenever
+    # ``"bias" in self``. This is the affine learned bias, distinct from the
+    # quant ``biases`` (zero-points) the kernel already consumes. The fused
+    # kernel recomputes the gate matmul in-register and has no slot for the
+    # learned bias, so a gate carrying one would silently drop it. up_proj runs
+    # its normal path (its bias, if any, is already in x_up), so only the gate
+    # is at risk: fall back to the unfused path when the gate has a learned bias.
+    if "bias" in gate:
+        return False
+    if up.bits != 4 or up.group_size != 64 or up.mode != "affine":
+        return False
+    if gate.bits != 4 or gate.group_size != 64 or gate.mode != "affine":
+        return False
+    if up.biases is None or gate.biases is None:
+        return False
+    # The kernel reads gate scales/biases as the activation dtype (in-kernel
+    # float() then writes back as the activation dtype), so the runtime
+    # precondition in fused_gate_qmv_silu_mul requires a single shared float
+    # dtype across the quant params. Check it here so a mismatch falls through
+    # to the unfused path instead of raising at forward. The packed weight is
+    # uint32 and intentionally excluded from the dtype contract.
+    param_dtype = gate.scales.dtype
+    if (
+        gate.biases.dtype != param_dtype
+        or up.scales.dtype != param_dtype
+        or up.biases.dtype != param_dtype
+    ):
+        return False
+    # Validate from the gate projection: fused_gate_qmv_silu_mul recomputes the
+    # gate matmul and keys K/N off gate dims, so the gate is the operand that
+    # must satisfy the tiling constraints (equivalent to up only while
+    # up.shape == gate.shape).
+    K = gate.scales.shape[-1] * _GROUP_SIZE
+    N = gate.weight.shape[-2]
+    if K % _BLOCK_SIZE != 0 or N % _ROWS_PER_TG != 0:
+        return False
+    return True
+
+
+_fallback_warned = False
+
+
+def _fused_gate_or_fallback(gate_proj, x, idx, x_up, sorted_indices=False):
+    """silu(gate_qmv(x)) * x_up via the fused kernel; on ValueError fall back to
+    the unfused gate projection. gather_qmm tolerates the activation dtype the
+    fused kernel rejects, which can_fuse cannot pre-check at patch time. Warns
+    once.
+    """
+    # Kernel layout only: the Metal kernel reads T from indices.shape[-1], and
+    # _gather_sort folded top_k into M_tok, so the sorted path needs an explicit
+    # T=1 axis. The fallback must not see it: gather_qmm broadcasts an (M_tok, 1)
+    # index against sorted x's (M_tok,) batch dim into an M_tok x M_tok cross
+    # product.
+    gate_idx = idx.reshape(-1, 1) if sorted_indices else idx
+    gw = gate_proj["weight"]
+    gs = gate_proj["scales"]
+    gb = gate_proj.get("biases")
+    try:
+        return fused_gate_qmv_silu_mul(x, gw, gs, gb, gate_idx, x_up)
+    except ValueError as e:
+        global _fallback_warned
+        if not _fallback_warned:
+            logger.warning(
+                "Path B: fused gate kernel declined inputs (%s); using the "
+                "unfused gate path for this and matching calls.",
+                e,
+            )
+            _fallback_warned = True
+        # Reference expression by construction: the same projection call with
+        # the same flat idx up_proj/down_proj receive.
+        return nn.silu(gate_proj(x, idx, sorted_indices=sorted_indices)) * x_up
+
+
+class FusedSwitchSwiGLU(nn.Module):
+    """SwitchGLU forward with Path B fusion installed.
+
+    Wraps an existing SwitchGLU instance. Reads up_proj / gate_proj weights
+    directly (no concatenation), runs up_proj as usual, then calls the fused
+    kernel for ``silu(gate_qmv(x)) * x_up`` in one shot.
+
+    Replaces ``switch_mlp.__call__`` via patch_switch_glu_with_fused_swiglu;
+    SwitchGLU.up_proj / gate_proj / activation are *not* replaced and remain
+    available for fallback paths (e.g. sorted-indices large-batch case).
+    """
+
+    def __init__(self, switch_mlp):
+        super().__init__()
+        # Weak ref: sw stores the bound fused_forward, so a strong ref here would
+        # cycle (sw -> method -> self -> sw). sw outlives every call into it.
+        self._switch_mlp = weakref.proxy(switch_mlp)
+
+    def fused_forward(self, x, indices):
+        """Same contract as SwitchGLU.__call__ but with fused activation."""
+        from mlx_lm.models.switch_layers import _gather_sort, _scatter_unsort
+
+        sw = self._switch_mlp
+        x = mx.expand_dims(x, (-2, -3))
+        do_sort = indices.size >= 64
+        idx = indices
+        inv_order = None
+        if do_sort:
+            x, idx, inv_order = _gather_sort(x, indices)
+
+        x_up = sw.up_proj(x, idx, sorted_indices=do_sort)
+
+        swiglu = _fused_gate_or_fallback(
+            sw.gate_proj, x, idx, x_up, sorted_indices=do_sort
+        )
+
+        out = sw.down_proj(swiglu, idx, sorted_indices=do_sort)
+
+        if do_sort:
+            out = _scatter_unsort(out, inv_order, indices.shape)
+        return out.squeeze(-2)
+
+
+def patch_switch_glu_with_fused_swiglu(model) -> int:
+    """Install Path B on every eligible SwitchGLU in the model.
+
+    Replaces ``switch_mlp.__call__`` with FusedSwitchSwiGLU.fused_forward.
+    Leaves ``up_proj``, ``gate_proj``, ``activation`` in place so the original
+    code path is still reachable for sorted-large-batch (handled internally) and
+    for any callers that bypass the patched __call__.
+
+    Returns number of layers patched.
+    """
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    patched = 0
+    for layer in model.model.layers:
+        mlp = getattr(layer, "mlp", None)
+        if mlp is None:
+            continue
+        sw = getattr(mlp, "switch_mlp", None)
+        if not isinstance(sw, SwitchGLU):
+            continue
+        if not can_fuse(sw):
+            continue
+        # Idempotent: skip if already patched.
+        if getattr(sw, "_path_b_installed", False):
+            continue
+        # AOT: pre-compile the fused kernel for this layer's shapes. top_k
+        # lives on the parent MoE block (e.g. Qwen3MoeSparseMoeBlock.top_k).
+        # If absent, skip warming and fall back to lazy JIT on first dispatch.
+        top_k = getattr(mlp, "top_k", None)
+        if top_k is not None:
+            _aot_warm_kernel(sw, int(top_k))
+        # One-off SwitchGLU subclass rather than rewriting up_proj/gate_proj: the
+        # activation fusion folds silu(gate)*x_up into the gate matmul, which has
+        # to intercept the forward (the projection level can't express it).
+        # can_fuse declines a non-stock __call__, so a customized forward falls
+        # back unpatched. Python resolves __call__ on the type, so swap
+        # sw.__class__ to a subclass; cache it on the exact class (cls.__dict__,
+        # not hasattr which walks the MRO) so a SwitchGLU subclass gets its own
+        # entry instead of being downcast to the base.
+        sw._path_b_call = FusedSwitchSwiGLU(sw).fused_forward
+        cls = type(sw)
+        if "_PathBSubclass" not in cls.__dict__:
+            cls._PathBSubclass = type(
+                f"{cls.__name__}_PathB",
+                (cls,),
+                {"__call__": lambda self, *a, **kw: self._path_b_call(*a, **kw)},
+            )
+        sw.__class__ = cls._PathBSubclass
+        sw._path_b_installed = True
+        patched += 1
+
+    if patched == 0:
+        logger.warning(
+            "patch_switch_glu_with_fused_swiglu: no eligible SwitchGLU found"
+        )
+    else:
+        logger.info(f"patch_switch_glu_with_fused_swiglu: patched {patched} layers")
+    return patched

python/sglang/srt/hardware_backend/mlx/moe/tests/test_fused_swiglu.py

@@ -0,0 +1,464 @@
+"""Numerical equivalence and eligibility tests for the Path B fused swiglu kernel.
+
+Two groups:
+  * Model-based equivalence (``@requires_model``): loads a small MoE model, runs
+    the fused gate_qmv + silu + ×x_up kernel against the unfused reference
+    (``mx.gather_qmm`` + ``nn.silu(gate) * x_up``) on both the unsorted and
+    sorted paths. Gated by SGLANG_MLX_TEST_MODEL so CI hosts without a model
+    cache skip them.
+  * Synthetic eligibility (no model, MLX only): the learned-bias fallback. The
+    fused kernel recomputes the gate matmul and has no slot for the per-expert
+    learned bias QuantizedSwitchLinear adds after the matmul, so ``can_fuse``
+    must exclude a gate carrying one, and the patch must leave such a layer
+    unfused. These run whenever MLX is importable.
+"""
+
+import os
+
+import pytest
+
+mx = pytest.importorskip("mlx.core")
+
+
+# Model-based tests need a real checkpoint; synthetic tests below do not.
+requires_model = pytest.mark.skipif(
+    not os.environ.get("SGLANG_MLX_TEST_MODEL"),
+    reason="Set SGLANG_MLX_TEST_MODEL to a HuggingFace model id to enable",
+)
+
+
+def _max_rel_diff(a, b):
+    diff = mx.abs(a.astype(mx.float32) - b.astype(mx.float32))
+    max_abs = diff.max().item()
+    ref_max = mx.abs(a.astype(mx.float32)).max().item()
+    return max_abs, max_abs / max(ref_max, 1e-9)
+
+
+@requires_model
+def test_fused_gate_qmv_silu_mul_matches_unfused():
+    """Kernel output matches ``nn.silu(gate_qmv) * x_up`` within bf16 ULP."""
+    import mlx.nn as nn
+    from mlx_lm import load
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        can_fuse,
+        fused_gate_qmv_silu_mul,
+    )
+
+    model, _ = load(os.environ["SGLANG_MLX_TEST_MODEL"])
+    sw = model.model.layers[0].mlp.switch_mlp
+    assert can_fuse(sw), "layer 0 not eligible for fused swiglu"
+
+    up = sw.up_proj
+    gate = sw.gate_proj
+    in_dim = up.scales.shape[-1] * up.group_size
+    out_dim = up.weight.shape[-2]
+    num_experts = up.weight.shape[0]
+    dtype = up.scales.dtype
+
+    # Two batch sizes both take the unsorted path (indices.size < 64).
+    for B, TOPK in [(1, 8), (4, 8)]:
+        x = mx.random.normal(shape=(B, 1, 1, in_dim)).astype(dtype)
+        indices = mx.random.randint(0, num_experts, shape=(B, TOPK)).astype(mx.uint32)
+
+        x_up = up(x, indices, sorted_indices=False)
+        x_gate = gate(x, indices, sorted_indices=False)
+        y_ref = nn.silu(x_gate) * x_up
+
+        y_fused = fused_gate_qmv_silu_mul(
+            x, gate["weight"], gate["scales"], gate.get("biases"), indices, x_up
+        )
+        mx.eval(y_ref, y_fused)
+
+        assert y_ref.shape == y_fused.shape
+
+        max_abs, rel = _max_rel_diff(y_ref, y_fused)
+        # 2 % relative covers ~2 bf16 ULPs at typical activation magnitudes;
+        # the kernel's fp32 accumulation order matches MLX's qmv_fast_impl so
+        # most elements should land within 1 ULP.
+        assert rel < 2e-2, f"B={B} TOPK={TOPK}: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+@requires_model
+def test_patched_switchglu_matches_unpatched():
+    """Full SwitchGLU forward equivalence on both sorted and unsorted paths."""
+    from mlx_lm import load
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    model, _ = load(os.environ["SGLANG_MLX_TEST_MODEL"])
+    sw = model.model.layers[0].mlp.switch_mlp
+    in_dim = sw.up_proj.scales.shape[-1] * sw.up_proj.group_size
+    num_experts = sw.up_proj.weight.shape[0]
+    dtype = sw.up_proj.scales.dtype
+
+    cases = []
+    # B=2 TOPK=8 -> indices.size=16 < 64 -> unsorted
+    # B=8 TOPK=8 -> indices.size=64 -> sorted
+    for B, TOPK, label in [(2, 8, "unsorted"), (8, 8, "sorted")]:
+        x = mx.random.normal(shape=(B, in_dim)).astype(dtype)
+        indices = mx.random.randint(0, num_experts, shape=(B, TOPK)).astype(mx.uint32)
+        out_ref = sw(x, indices)
+        mx.eval(out_ref)
+        cases.append((label, x, indices, out_ref))
+
+    n_patched = patch_switch_glu_with_fused_swiglu(model)
+    assert n_patched > 0, "no SwitchGLU layers were patched"
+
+    for label, x, indices, out_ref in cases:
+        out_fused = sw(x, indices)
+        mx.eval(out_fused)
+        max_abs, rel = _max_rel_diff(out_ref, out_fused)
+        # 5 % is generous; in practice we see <0.6 % on 48-layer Qwen3-MoE.
+        # The looser bound here absorbs cross-layer ULP propagation through
+        # down_proj's quantized matmul.
+        assert rel < 5e-2, f"full forward {label}: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+# Learned-bias fallback (synthetic, no model): a gate with a learned bias must
+# not fuse, since the kernel has no slot for the bias added after the matmul.
+def _quantized_switch_glu(in_dim, hidden, n_experts, gate_bias):
+    """Small quantized SwitchGLU; gate carries a learned bias iff gate_bias.
+
+    in_dim=512 keeps K%512==0 and hidden%8==0, inside the Path B v1 regime, so
+    the bias-free build is genuinely fusion-eligible (the True control).
+    """
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    sw = SwitchGLU(in_dim, hidden, n_experts, bias=False)
+    sw.up_proj = sw.up_proj.to_quantized(group_size=64, bits=4, mode="affine")
+    sw.down_proj = sw.down_proj.to_quantized(group_size=64, bits=4, mode="affine")
+    gate = sw.gate_proj
+    if gate_bias:
+        # Learned per-expert bias (E, N), nonzero so dropping it would change
+        # the result. to_quantized copies it into the QuantizedSwitchLinear.
+        gate.bias = mx.random.normal((n_experts, hidden)) * 0.1
+    sw.gate_proj = gate.to_quantized(group_size=64, bits=4, mode="affine")
+    return sw
+
+
+def test_can_fuse_excludes_learned_gate_bias():
+    """can_fuse: False for a gate with a learned bias, True when bias-free."""
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import can_fuse
+
+    sw_free = _quantized_switch_glu(512, 64, 8, gate_bias=False)
+    sw_bias = _quantized_switch_glu(512, 64, 8, gate_bias=True)
+    assert "bias" not in sw_free.gate_proj
+    assert "bias" in sw_bias.gate_proj
+    assert can_fuse(sw_free) is True, "bias-free gate in regime should fuse"
+    assert can_fuse(sw_bias) is False, "gate with learned bias must fall back"
+
+
+def test_patch_falls_back_on_gate_bias():
+    """Patching a biased-gate SwitchGLU is a no-op; the forward stays bias-correct."""
+    import types
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 8, 4, 2  # 2*4=8 < 64 -> unsorted
+    sw = _quantized_switch_glu(in_dim, hidden, n_experts, gate_bias=True)
+
+    x = mx.random.normal((B, in_dim))
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+    out_before = sw(x, indices)
+    mx.eval(out_before)
+
+    # Minimal model stand-in: the patch walks model.model.layers[*].mlp.switch_mlp.
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+
+    n_patched = patch_switch_glu_with_fused_swiglu(model)
+    assert n_patched == 0, "biased gate must not be patched"
+
+    out_after = sw(x, indices)
+    mx.eval(out_after)
+    d = mx.abs(out_before.astype(mx.float32) - out_after.astype(mx.float32))
+    diff = d.max().item()
+    assert diff == 0.0, f"forward changed after (no-op) patch: max|delta|={diff:.3e}"
+
+
+# Model-free numerical equivalence + non-stock-forward guard: the central
+# correctness check, runs without a model download (skips where Metal is absent).
+def test_fused_matches_unfused_synthetic():
+    """Synthetic quantized gate weights: fused kernel vs the unfused
+    gather_qmm + silu*x_up path, within the kernel's bf16 bound, plus finiteness."""
+    mx.random.seed(0)
+
+    import mlx.nn as nn
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        fused_gate_qmv_silu_mul,
+    )
+
+    # Gate regime: K%512==0, N%8==0, bits=4, group_size=64, affine.
+    E, N, K, TOPK = 4, 16, 512, 2
+    dtype = mx.bfloat16
+    gate_w = (mx.random.normal((E, N, K)) * 0.02).astype(dtype)
+    gwq, gs, gb = mx.quantize(gate_w, group_size=64, bits=4)
+    mx.eval(gwq, gs, gb)
+
+    # Two routing patterns: spread (hi=E) and collisions (many tokens, few experts).
+    for B, hi in [(2, E), (4, max(1, E // 2))]:
+        x = mx.random.normal((B, 1, 1, K)).astype(dtype)
+        idx = mx.random.randint(0, hi, shape=(B, TOPK)).astype(mx.uint32)
+        x_up = mx.random.normal((B, TOPK, 1, N)).astype(dtype)
+
+        x_gate = mx.gather_qmm(
+            x,
+            gwq,
+            gs,
+            gb,
+            rhs_indices=idx,
+            transpose=True,
+            group_size=64,
+            bits=4,
+            mode="affine",
+        )
+        y_ref = nn.silu(x_gate) * x_up
+        y_fused = fused_gate_qmv_silu_mul(x, gwq, gs, gb, idx, x_up)
+        mx.eval(y_ref, y_fused)
+
+        assert y_ref.shape == y_fused.shape
+        # A broken kernel must not leak NaN/Inf into the downstream down_proj matmul.
+        assert bool(
+            mx.all(mx.isfinite(y_fused.astype(mx.float32))).item()
+        ), f"B={B} hi={hi}: non-finite fused output"
+        # Same bf16 bound as the @requires_model kernel test.
+        max_abs, rel = _max_rel_diff(y_ref, y_fused)
+        assert rel < 2e-2, f"B={B} hi={hi}: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+def test_can_fuse_declines_nonstock_call():
+    """can_fuse: False when SwitchGLU.__call__ is overridden (the fused subclass
+    would impose stock semantics and silently bypass the override), True for stock."""
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import can_fuse
+
+    # hidden=64 keeps down_proj's input dim divisible by the quant group size.
+    sw_stock = _quantized_switch_glu(512, 64, 4, gate_bias=False)
+    assert can_fuse(sw_stock) is True, "stock in-regime SwitchGLU should fuse"
+
+    class _CustomSwitchGLU(SwitchGLU):
+        def __call__(self, x, indices):  # overridden forward
+            return super().__call__(x, indices)
+
+    sw_custom = _quantized_switch_glu(512, 64, 4, gate_bias=False)
+    sw_custom.__class__ = _CustomSwitchGLU  # same swap mechanism the patch uses
+    assert can_fuse(sw_custom) is False, "non-stock __call__ must fall back"
+
+
+def test_can_fuse_declines_non_silu_activation():
+    """can_fuse: False for a non SiLU activation (the kernel and the fallback
+    both bake in silu, which would silently replace the module's formula),
+    True for the stock SwiGLU control."""
+    import types
+
+    import mlx.nn as nn
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        can_fuse,
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    # Same build as _quantized_switch_glu, but the activation kwarg is the
+    # subject under test, so construct directly.
+    sw = SwitchGLU(512, 64, 4, activation=nn.gelu, bias=False)
+    for name in ("up_proj", "gate_proj", "down_proj"):
+        proj = getattr(sw, name)
+        setattr(sw, name, proj.to_quantized(group_size=64, bits=4, mode="affine"))
+    assert can_fuse(sw) is False, "non SiLU activation must fall back"
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=4)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert patch_switch_glu_with_fused_swiglu(model) == 0, "gelu module must not patch"
+
+    sw_stock = _quantized_switch_glu(512, 64, 4, gate_bias=False)
+    assert can_fuse(sw_stock) is True, "stock SwiGLU activation should fuse"
+
+
+def test_fused_forward_falls_back_on_dtype_mismatch():
+    """A runtime activation dtype the fused kernel rejects but gather_qmm
+    tolerates (bf16 gate params, fp16 activations) must fall back, not crash,
+    and match the unfused forward."""
+    import types
+
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
+        patch_switch_glu_with_fused_swiglu,
+    )
+
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 4, 4, 2  # 2*4=8 < 64 -> unsorted
+    sw = SwitchGLU(in_dim, hidden, n_experts, bias=False)
+    for name in ("up_proj", "gate_proj", "down_proj"):
+        lin = getattr(sw, name)
+        lin.weight = lin.weight.astype(mx.bfloat16)  # bf16 weight -> bf16 scales
+        setattr(sw, name, lin.to_quantized(group_size=64, bits=4, mode="affine"))
+    assert sw.gate_proj.scales.dtype == mx.bfloat16
+
+    # fp16 activations mismatch the bf16 gate params: the fused kernel raises,
+    # the unfused gather_qmm tolerates it.
+    x = mx.random.normal((B, in_dim)).astype(mx.float16)
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+
+    out_ref = sw(x, indices)  # stock forward, unpatched
+    mx.eval(out_ref)
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert patch_switch_glu_with_fused_swiglu(model) == 1, "layer should patch"
+
+    out_fb = sw(x, indices)  # patched -> kernel raises -> fallback, no crash
+    mx.eval(out_fb)
+    max_abs, rel = _max_rel_diff(out_ref, out_fb)
+    assert rel < 1e-3, f"fallback != unfused: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+# Fallback index contract (synthetic, no model): the fallback must see the
+# untouched flat indices. The sorted path's (M_tok, 1) kernel reshape once
+# leaked into the fallback and broadcast an M_tok x M_tok cross product
+# (PR #26188 review repro).
+def _bf16_quantized_switch_glu(in_dim, hidden, n_experts):
+    """Quantize from bf16 weights so fp16 activations trip the kernel's runtime
+    dtype check while the unfused path tolerates them."""
+    from mlx_lm.models.switch_layers import SwitchGLU
+
+    sw = SwitchGLU(in_dim, hidden, n_experts, bias=False)
+    for name in ("up_proj", "gate_proj", "down_proj"):
+        lin = getattr(sw, name)
+        lin.weight = lin.weight.astype(mx.bfloat16)
+        setattr(sw, name, lin.to_quantized(group_size=64, bits=4, mode="affine"))
+    return sw
+
+
+def test_sorted_dtype_mismatch_fallback_matches_reference(monkeypatch):
+    """Reviewer repro: B*T == 64 takes the sorted path, the kernel rejects fp16
+    activations on bf16 params, and the fallback must match the reference in
+    shape and value."""
+    import types
+
+    import sglang.srt.hardware_backend.mlx.moe.fused_swiglu as fused_swiglu
+
+    monkeypatch.setattr(fused_swiglu, "_fallback_warned", False)
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 4, 4, 16  # 16*4 = 64 -> sorted
+    sw = _bf16_quantized_switch_glu(in_dim, hidden, n_experts)
+
+    x = mx.random.normal((B, in_dim)).astype(mx.float16)
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+    out_ref = sw(x, indices)
+    mx.eval(out_ref)
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert fused_swiglu.patch_switch_glu_with_fused_swiglu(model) == 1
+
+    out_fb = sw(x, indices)
+    mx.eval(out_fb)
+    assert out_fb.shape == out_ref.shape
+    max_abs, rel = _max_rel_diff(out_ref, out_fb)
+    # Post fix the fallback runs the same MLX ops as the stock forward, so the
+    # bound only absorbs compiled vs eager elementwise ordering (~1 fp16 ULP).
+    assert bool(
+        mx.allclose(
+            out_fb.astype(mx.float32),
+            out_ref.astype(mx.float32),
+            rtol=2e-3,
+            atol=2e-4,
+        ).item()
+    ), f"sorted fallback != reference: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+def test_unsorted_dtype_mismatch_fallback_matches_reference(monkeypatch):
+    """Sibling guard: same dtype mismatch on the unsorted path (B*T < 64)."""
+    import types
+
+    import sglang.srt.hardware_backend.mlx.moe.fused_swiglu as fused_swiglu
+
+    monkeypatch.setattr(fused_swiglu, "_fallback_warned", False)
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k, B = 512, 64, 4, 4, 2  # 2*4 = 8 < 64 -> unsorted
+    sw = _bf16_quantized_switch_glu(in_dim, hidden, n_experts)
+
+    x = mx.random.normal((B, in_dim)).astype(mx.float16)
+    indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+    out_ref = sw(x, indices)
+    mx.eval(out_ref)
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    assert fused_swiglu.patch_switch_glu_with_fused_swiglu(model) == 1
+
+    out_fb = sw(x, indices)
+    mx.eval(out_fb)
+    assert out_fb.shape == out_ref.shape
+    max_abs, rel = _max_rel_diff(out_ref, out_fb)
+    assert bool(
+        mx.allclose(
+            out_fb.astype(mx.float32),
+            out_ref.astype(mx.float32),
+            rtol=2e-3,
+            atol=2e-4,
+        ).item()
+    ), f"unsorted fallback != reference: max_abs={max_abs:.3e} rel={rel:.2%}"
+
+
+def test_forced_kernel_rejection_falls_back_correctly(monkeypatch):
+    """Any ValueError from the fused kernel, not just a dtype mismatch, must
+    take the identical fallback: force one via monkeypatch and check both
+    routing paths against the unpatched module."""
+    import types
+
+    import sglang.srt.hardware_backend.mlx.moe.fused_swiglu as fused_swiglu
+
+    monkeypatch.setattr(fused_swiglu, "_fallback_warned", False)
+    mx.random.seed(0)
+    in_dim, hidden, n_experts, top_k = 512, 64, 4, 4
+    sw = _quantized_switch_glu(in_dim, hidden, n_experts, gate_bias=False)
+
+    cases = []
+    # B=2 -> 8 < 64 -> unsorted; B=16 -> 64 -> sorted.
+    for B, label in [(2, "unsorted"), (16, "sorted")]:
+        x = mx.random.normal((B, in_dim))
+        indices = mx.random.randint(0, n_experts, shape=(B, top_k)).astype(mx.uint32)
+        out_ref = sw(x, indices)
+        mx.eval(out_ref)
+        cases.append((label, x, indices, out_ref))
+
+    mlp = types.SimpleNamespace(switch_mlp=sw, top_k=top_k)
+    layer = types.SimpleNamespace(mlp=mlp)
+    model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[layer]))
+    # Patch before installing the raiser: _aot_warm_kernel dispatches the real
+    # kernel at patch time and does not catch ValueError.
+    assert fused_swiglu.patch_switch_glu_with_fused_swiglu(model) == 1
+
+    def raiser(*args, **kwargs):
+        raise ValueError("forced rejection")
+
+    monkeypatch.setattr(fused_swiglu, "fused_gate_qmv_silu_mul", raiser)
+
+    for label, x, indices, out_ref in cases:
+        out_fb = sw(x, indices)
+        mx.eval(out_fb)
+        assert out_fb.shape == out_ref.shape, label
+        max_abs, rel = _max_rel_diff(out_ref, out_fb)
+        assert bool(
+            mx.allclose(
+                out_fb.astype(mx.float32),
+                out_ref.astype(mx.float32),
+                rtol=1e-5,
+                atol=1e-6,
+            ).item()
+        ), f"forced rejection {label}: max_abs={max_abs:.3e} rel={rel:.2%}"

python/sglang/srt/hardware_backend/mlx/parent_watchdog.py

@@ -0,0 +1,60 @@
+"""Parent-death watchdog for MLX workers on Apple Silicon.
+
+macOS has no ``PR_SET_PDEATHSIG`` equivalent, so the kernel will not signal a
+worker process when its parent dies; the worker would be reparented to PID 1
+and leak (holding GPU/host memory and ports). This module emulates PDEATHSIG
+with a daemon thread that watches the parent PID via kqueue and SIGKILLs the
+current process once it gets orphaned.
+"""
+
+import os
+import select
+import signal
+import threading
+
+
+def start_parent_death_watcher() -> None:
+    """SIGKILL this process once its current parent exits (macOS only).
+
+    kqueue with an ``EVFILT_PROC`` / ``NOTE_EXIT`` filter is the native,
+    event-driven mechanism on macOS (exposed via ``select.kqueue`` /
+    ``select.kevent``), so the watcher thread blocks until the parent actually
+    exits instead of waking up to poll.
+
+    ``SIGKILL`` is sent from this watcher thread and is uncatchable /
+    unblockable, so it works even when the main thread is stuck inside a
+    blocking native call (e.g. an MLX/Metal ``mx.eval`` / ``.tolist()``).
+    """
+    original_ppid = os.getppid()
+
+    def _watch_parent():
+        kq = select.kqueue()
+        kev = select.kevent(
+            original_ppid,
+            filter=select.KQ_FILTER_PROC,
+            flags=select.KQ_EV_ADD,
+            fflags=select.KQ_NOTE_EXIT,
+        )
+        try:
+            # Register the EVFILT_PROC / NOTE_EXIT watch on the parent PID.
+            kq.control([kev], 0, None)
+        except (ProcessLookupError, OSError):
+            # The parent already exited before we could register the watch
+            # (ESRCH); we are already orphaned.
+            os.kill(os.getpid(), signal.SIGKILL)
+            return
+        # Guard against the race where the parent exits between reading
+        # original_ppid and registering the watch above.
+        if os.getppid() != original_ppid:
+            os.kill(os.getpid(), signal.SIGKILL)
+            return
+        # Block until the parent exits, then terminate ourselves.
+        kq.control(None, 1, None)
+        os.kill(os.getpid(), signal.SIGKILL)
+
+    watcher = threading.Thread(
+        target=_watch_parent,
+        name="parent-death-watcher",
+        daemon=True,
+    )
+    watcher.start()

python/sglang/srt/hardware_backend/mlx/profiler.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import gzip
+import json
+import logging
+import os
+import shutil
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Optional
+
+import torch
+
+from sglang.srt.managers.io_struct import ProfileReqOutput
+from sglang.srt.utils.tensor_bridge import use_mlx
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MetalCaptureProfiler:
+    label: str
+    trace_path: Path
+    stop_capture: Callable[[], None]
+    standalone: bool
+
+    @classmethod
+    def start_mlx(cls, trace_path: Path):
+        trace_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            import mlx.core as mx
+
+            mx.metal.start_capture(str(trace_path))
+        except RuntimeError as e:
+            return None, _capture_error("MLX", e)
+
+        return cls._started(
+            label="MLX",
+            trace_path=trace_path,
+            stop_capture=mx.metal.stop_capture,
+            standalone=True,
+        )
+
+    @classmethod
+    def start_mps(cls, trace_path: Path):
+        trace_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            if not hasattr(torch, "mps") or not hasattr(torch.mps, "profiler"):
+                raise RuntimeError("torch.mps.profiler is not available")
+            context = torch.mps.profiler.metal_capture(str(trace_path))
+            context.__enter__()
+        except RuntimeError as e:
+            return None, _capture_error("MPS", e)
+
+        return cls._started(
+            label="MPS",
+            trace_path=trace_path,
+            stop_capture=lambda: context.__exit__(None, None, None),
+            standalone=False,
+        )
+
+    @classmethod
+    def _started(
+        cls,
+        *,
+        label: str,
+        trace_path: Path,
+        stop_capture: Callable[[], None],
+        standalone: bool,
+    ):
+        profiler = cls(
+            label=label,
+            trace_path=trace_path,
+            stop_capture=stop_capture,
+            standalone=standalone,
+        )
+        logger.info("%s Metal capture started, saving to %s", label, trace_path)
+        return profiler, ProfileReqOutput(success=True, message="Succeeded")
+
+    def stop(self) -> str:
+        self.stop_capture()
+
+        logger.info(
+            "%s Metal capture stopped. Trace saved to: %s",
+            self.label,
+            self.trace_path,
+        )
+        return f" Metal trace: {self.trace_path}"
+
+
+def _capture_error(label: str, error: RuntimeError) -> ProfileReqOutput:
+    return ProfileReqOutput(
+        success=False,
+        message=(
+            f"Failed to start {label} Metal capture: {error}. "
+            "Set MTL_CAPTURE_ENABLED=1 in the server's environment "
+            "before launching to enable GPU trace capture."
+        ),
+    )
+
+
+class MetalTorchProfiler:
+    def __init__(
+        self,
+        *,
+        start_metal_capture: Callable[[Path], tuple[Any, ProfileReqOutput]],
+        torch_profiler: Optional[Any] = None,
+    ):
+        self.start_metal_capture = start_metal_capture
+        self.torch_profiler = torch_profiler
+        self.metal_profiler = None
+
+    def start(self):
+        trace_path = _new_temp_gputrace_path()
+        self.metal_profiler, result = self.start_metal_capture(trace_path)
+        if not result.success:
+            raise RuntimeError(result.message)
+        if self.torch_profiler is not None:
+            try:
+                self.torch_profiler.start()
+            except Exception:
+                self.metal_profiler.stop()
+                raise
+
+    def stop(self):
+        try:
+            if self.torch_profiler is not None:
+                self.torch_profiler.stop()
+        finally:
+            if self.metal_profiler is not None:
+                self.metal_profiler.stop()
+
+    def export_chrome_trace(self, path: str):
+        if self.torch_profiler is not None:
+            self.torch_profiler.export_chrome_trace(path)
+        else:
+            _write_empty_chrome_trace(path)
+
+        if self.metal_profiler is None:
+            return
+
+        final_path = _unique_gputrace_path_for_chrome_trace(path)
+        final_path.parent.mkdir(parents=True, exist_ok=True)
+        if self.metal_profiler.trace_path.exists():
+            shutil.move(str(self.metal_profiler.trace_path), str(final_path))
+            logger.info("Metal trace saved to: %s", final_path)
+
+
+def apply_metal_profiler_patches() -> None:
+    if getattr(torch.profiler.profile, "_sglang_metal_patched", False):
+        return
+
+    original_profile = torch.profiler.profile
+
+    def profile(*args, **kwargs):
+        activities = _get_activities(args, kwargs)
+        if not _has_cuda_activity(activities):
+            return original_profile(*args, **kwargs)
+
+        if use_mlx():
+            return MetalTorchProfiler(
+                start_metal_capture=MetalCaptureProfiler.start_mlx
+            )
+
+        torch_activities = [
+            activity for activity in activities if not _is_cuda_activity(activity)
+        ]
+        torch_profiler = None
+        if torch_activities:
+            patched_args, patched_kwargs = _replace_activities(
+                args, kwargs, torch_activities
+            )
+            torch_profiler = original_profile(*patched_args, **patched_kwargs)
+
+        return MetalTorchProfiler(
+            start_metal_capture=MetalCaptureProfiler.start_mps,
+            torch_profiler=torch_profiler,
+        )
+
+    profile._sglang_metal_patched = True
+    profile._sglang_original_profile = original_profile
+    torch.profiler.profile = profile
+
+
+def _get_activities(args, kwargs):
+    if "activities" in kwargs:
+        return kwargs["activities"]
+    if args:
+        return args[0]
+    return None
+
+
+def _replace_activities(args, kwargs, activities):
+    kwargs = dict(kwargs)
+    if "activities" in kwargs:
+        kwargs["activities"] = activities
+        return args, kwargs
+
+    if args:
+        args = list(args)
+        args[0] = activities
+        return tuple(args), kwargs
+
+    kwargs["activities"] = activities
+    return args, kwargs
+
+
+def _has_cuda_activity(activities) -> bool:
+    if activities is None:
+        return False
+    return any(_is_cuda_activity(activity) for activity in activities)
+
+
+def _is_cuda_activity(activity) -> bool:
+    return activity == torch.profiler.ProfilerActivity.CUDA
+
+
+def _new_temp_gputrace_path() -> Path:
+    output_dir = Path(os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")).expanduser()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for i in range(100):
+        candidate = (
+            output_dir / f"sglang-metal-{os.getpid()}-{time.time_ns()}-{i}.gputrace"
+        )
+        if not candidate.exists():
+            return candidate
+    raise RuntimeError(f"Cannot find an unused Metal trace path in {output_dir}")
+
+
+def _unique_gputrace_path_for_chrome_trace(path: str) -> Path:
+    chrome_path = Path(path).expanduser()
+    name = chrome_path.name
+    if name.endswith(".trace.json.gz"):
+        name = name[: -len(".trace.json.gz")] + ".gputrace"
+    else:
+        name = chrome_path.stem + ".gputrace"
+
+    base = chrome_path.with_name(name)
+    if not base.exists():
+        return base
+
+    stem = base.name[: -len(".gputrace")]
+    for i in range(100):
+        candidate = base.with_name(f"{stem}-{time.time_ns()}-{i}.gputrace")
+        if not candidate.exists():
+            return candidate
+    raise RuntimeError(f"Cannot find an unused Metal trace path for {base}")
+
+
+def _write_empty_chrome_trace(path: str):
+    trace = {"traceEvents": []}
+    Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
+    if str(path).endswith(".gz"):
+        with gzip.open(path, "wt") as f:
+            json.dump(trace, f)
+    else:
+        with open(path, "w") as f:
+            json.dump(trace, f)

python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py

@@ -0,0 +1,263 @@
+"""MLX overlap scheduling mixin for the SGLang scheduler.
+
+Provides ``event_loop_overlap_mlx``, which pipelines MLX forward
+passes by keeping two in-flight lazy graphs queued on the GPU while
+the scheduler runs its CPU-side bookkeeping on the tokens of the
+older one.  The lazy-graph primitives live in
+``hardware_backend/mlx/tp_worker.py`` and ``model_runner.py``.
+
+Each request's attention KV lives in per-request, per-layer
+``ContiguousAttentionKVCache`` objects that ``MLXAttentionWrapper`` mutates
+in place during the forward pass. Chained decodes reuse the same cache objects:
+step N+1's graph reads step N's lazy writes via MLX's dependency tracking, so
+the GPU runs both steps back-to-back with no idle gap.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional
+
+import mlx.core as mx
+
+from sglang.srt.environ import envs
+from sglang.srt.managers.overlap_utils import resolve_forward_inputs
+from sglang.srt.utils import DynamicGradMode
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.mlx.model_runner import (
+        MlxPendingDecode,
+        MlxPendingExtend,
+        MlxPendingPrefill,
+    )
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+@dataclass
+class MlxPendingJob:
+    """Unfinished MLX work and graphs queued on the GPU.
+
+    Attributes:
+        lazy_tokens: Lazily evaluated token IDs produced by the forward
+            pass.  Unevaluated; calling ``.tolist()`` / ``.item()`` /
+            ``mx.eval`` on it will block until the Metal kernel finishes.
+            ``None`` for idle batches.
+        prefills: MLX prefill state returned by the model worker — one
+            entry per new request in an extend batch.  Used by
+            ``finalize_mlx_result`` to commit per-request caches.  Empty
+            list for pure-decode steps.
+        extends: Chunked-prefill-continuation state, one entry per
+            already-active request whose extend seq_len > 1.  Also empty
+            for pure-decode steps.
+        decode: Decode state covering full-decode mode AND mixed
+            single-token decodes inside an extend batch.  Used as the
+            chaining root by :meth:`async_chained_decode_mlx`.
+        mode: One of ``"decode"``, ``"extend"``, ``"idle"`` describing
+            which forward pass produced this job.  Drives finalise
+            dispatch and whether chaining is safe.
+        batch_copy: Snapshot of the :class:`ScheduleBatch` at launch
+            time.  Decoupled from the live batch so
+            ``process_batch_result`` can update request state without
+            racing against the next scheduling decision.
+        schedule_batch: The full scheduler batch.  Unlike ``batch_copy``,
+            this keeps allocator/cache fields needed when a prefill batch
+            becomes the next running decode batch.
+        reqs: Snapshot of ``batch.reqs`` at launch time.  The overlap
+            loop uses this to check ``req.finished()`` on the previous
+            step's request list without holding a reference to the
+            mutable batch object.
+    """
+
+    lazy_tokens: Optional[mx.array]
+    prefills: list[MlxPendingPrefill]
+    extends: list[MlxPendingExtend]
+    decode: Optional[MlxPendingDecode]
+    mode: str
+    batch_copy: ScheduleBatch
+    schedule_batch: ScheduleBatch
+    reqs: List[Req]
+
+
+class SchedulerMlxOverlapMixin:
+    """Mixin that adds MLX overlap scheduling to :class:`Scheduler`."""
+
+    def _finalize_mlx_pending_job(self: Scheduler, pending: MlxPendingJob):
+        # Account for this completed forward step. The standard scheduler does
+        # this inside run_batch(), but the MLX overlap loop bypasses run_batch,
+        # so without this forward_ct never advances on MLX. That stalls the
+        # watchdog liveness counter and, more importantly, breaks step-bounded
+        # profiling: _profile_batch_predicate auto-starts/stops based on
+        # forward_ct, so `--profile-steps` (and the server /start_profile
+        # num_steps path) only takes effect once the counter moves here.
+        self.forward_ct += 1
+        self.profiler_manager._profile_batch_predicate(pending.schedule_batch)
+
+        result = self.tp_worker.finalize_mlx_result(
+            pending.prefills,
+            pending.extends,
+            pending.decode,
+            pending.mode,
+            pending.reqs,
+        )
+        if result.next_token_ids is not None:
+            pending.batch_copy.input_ids = result.next_token_ids
+            pending.schedule_batch.input_ids = result.next_token_ids
+        self.last_batch = pending.schedule_batch
+        self.process_batch_result(pending.batch_copy, result)
+
+    @DynamicGradMode()
+    def event_loop_overlap_mlx(self: Scheduler):
+        """MLX-specific overlap loop modelled on ``mlx_lm.generate.generate_step``.
+
+        At steady state we keep TWO in-flight MLX graphs queued on the
+        GPU:
+
+        * ``pending_curr`` — the step whose tokens we are about to block
+          on and feed into the scheduler's bookkeeping.
+        * ``pending_next`` — the step that was built on top of
+          ``pending_curr``'s still-lazy output tokens via
+          ``async_chained_decode_mlx`` and has already been handed to
+          ``mx.async_eval``.  Because MLX tracks the full dependency
+          graph, the GPU will execute ``pending_next`` back-to-back
+          with ``pending_curr`` — there is no scheduling gap on the
+          device.
+
+        Bookkeeping timeline for a steady-state decode loop:
+
+            iter k:
+              build pending_next  (CPU graph build + mx.async_eval; cheap)
+              block on pending_curr via .tolist() (wait only on curr's tokens)
+              process_batch_result(pending_curr)   <-- GPU is running pending_next
+              pending_curr = pending_next
+
+        The chain is broken (we fall back to a "schedule + launch" step)
+        whenever any of the following holds:
+
+        * ``pending_curr`` is not a pure decode (e.g. prefill/extend).
+        * The waiting queue has new requests that need prefill.
+        * Any req in ``pending_curr`` just finished this iteration, so
+          the composition for ``pending_next`` would need to shrink.
+
+        When the chain breaks mid-flight we still finalise the
+        already-launched ``pending_next`` normally (its tokens are
+        valid for all surviving reqs).  With RadixCache-backed caches
+        (#21509) there is no ``extract_cache`` step: per-request caches
+        are the source of truth and are never merged into a shared
+        batched buffer.
+        """
+        pending_curr: Optional[MlxPendingJob] = None
+        pending_next: Optional[MlxPendingJob] = None
+
+        def _launch_fresh(batch: ScheduleBatch) -> MlxPendingJob:
+            # Materialize batch.input_ids from CPU staging (prefill) or the
+            # FutureMap relay (decode) before the forward. With deferred input
+            # materialization, get_next_batch_to_run leaves input_ids unset; the
+            # CUDA paths call resolve_forward_inputs for this, but the MLX overlap
+            # loop must do it too, otherwise async_forward_batch_generation_mlx
+            # dereferences a None input_ids.
+            resolve_forward_inputs(batch, self.future_map)
+            lazy_tokens, prefills, extends, decode, mode = (
+                self.tp_worker.async_forward_batch_generation_mlx(batch)
+            )
+            return MlxPendingJob(
+                lazy_tokens=lazy_tokens,
+                prefills=prefills,
+                extends=extends,
+                decode=decode,
+                mode=mode,
+                batch_copy=batch.copy(),
+                schedule_batch=batch,
+                reqs=list(batch.reqs),
+            )
+
+        def _launch_chained(prev: MlxPendingJob) -> MlxPendingJob:
+            assert prev.decode is not None
+            lazy_tokens, prefills, extends, decode, mode = (
+                self.tp_worker.async_chained_decode_mlx(prev.decode)
+            )
+            # Composition is identical to prev: reuse a fresh batch copy
+            # of the same underlying ScheduleBatch so process_batch_result
+            # updates the same req objects with the new token.
+            return MlxPendingJob(
+                lazy_tokens=lazy_tokens,
+                prefills=prefills,
+                extends=extends,
+                decode=decode,
+                mode=mode,
+                batch_copy=prev.batch_copy.copy(),
+                schedule_batch=prev.schedule_batch,
+                reqs=prev.reqs,
+            )
+
+        while True:
+            recv_reqs = self.request_receiver.recv_requests()
+            self.process_input_requests(recv_reqs)
+            if self._engine_paused:
+                continue
+
+            # 1. If pending_curr is a pure decode AND no new prefill is waiting,
+            #    build pending_next on top of it NOW — before we block on curr.
+            can_chain = (
+                pending_curr is not None
+                and pending_curr.mode == "decode"
+                and pending_curr.decode is not None
+                and not self.waiting_queue
+            )
+            if can_chain and pending_next is None:
+                # Build + launch the chained step BEFORE we block on
+                # pending_curr — this is the "no idle gap" trick.
+                # GPU now has 2 steps queued.
+                pending_next = _launch_chained(pending_curr)
+                self.result_queue.append(pending_next)
+
+            # 2. Finalize/process on pending_curr's tokens.  (GPU is already
+            #    executing pending_next at this point.)
+            if pending_curr is not None:
+                self._finalize_mlx_pending_job(pending_curr)
+                self.result_queue.popleft()
+                pending_curr = None
+
+            # 3. Decide whether pending_next is still valid (if no reqs finished)
+            #    and promote it.
+            finished_any = any(
+                req.finished() for req in (pending_next.reqs if pending_next else [])
+            )
+            new_prefill_waiting = bool(self.waiting_queue)
+            if (
+                pending_next is not None
+                and not finished_any
+                and not new_prefill_waiting
+            ):
+                pending_curr = pending_next
+                pending_next = None
+                self.cur_batch_for_debug = pending_curr.schedule_batch
+                self.last_batch = pending_curr.schedule_batch
+                if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
+                    self.invariant_checker.self_check_during_busy()
+                continue
+
+            # 4. Chain is broken. Finalise pending_next (if any), then
+            #    schedule fresh.
+            if pending_next is not None:
+                self._finalize_mlx_pending_job(pending_next)
+                self.result_queue.popleft()
+                pending_next = None
+            plan = self.get_next_batch_to_run(
+                running_batch=self.running_batch, last_batch=self.last_batch
+            )
+            self.running_batch = plan.running_batch
+            next_batch = plan.batch_to_run
+            self.cur_batch_for_debug = next_batch
+            if next_batch:
+                pending_curr = _launch_fresh(next_batch)
+                self.result_queue.append(pending_curr)
+            else:
+                self.on_idle()
+
+            self.last_batch = next_batch
+            if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
+                self.invariant_checker.self_check_during_busy()

python/sglang/srt/hardware_backend/mlx/tp_worker.py

@@ -0,0 +1,528 @@
+"""MLX-specific TpModelWorker subclass for Apple Silicon.
+
+Routes forward passes through the MLX model runner, bypassing PyTorch
+MPS.  A lightweight stub provides scheduler bookkeeping; the actual
+attention KV data lives in MlxAttentionKVPool.
+
+The worker also exposes an async (lazy-eval) surface used by the MLX
+overlap scheduler: ``async_forward_batch_generation_mlx`` launches a
+batch without blocking on the GPU, ``async_chained_decode_mlx`` builds
+the next decode step on top of a still-lazy previous decode, and
+``finalize_mlx_result`` blocks on the lazy outputs and produces a
+normal ``GenerationBatchResult``.
+"""
+
+import logging
+from typing import Optional, Union
+
+import mlx.core as mx
+import torch
+
+from sglang.srt.hardware_backend.mlx.model_runner import (
+    MlxPendingDecode,
+    MlxPendingExtend,
+    MlxPendingPrefill,
+)
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.managers.utils import GenerationBatchResult
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+
+logger = logging.getLogger(__name__)
+
+
+class MlxTpModelWorker(TpModelWorker):
+    """A tensor parallel model worker that routes inference through MLX.
+
+    Inherits from TpModelWorker for scheduler integration, but replaces
+    the standard ModelRunner with MlxModelRunnerStub (no PyTorch weights,
+    zero-memory KV cache) and delegates all forward passes to a native
+    MlxModelRunner.
+    """
+
+    def _init_model_runner(self):
+        """Create MLX runner first (auto-sizes pool), then stub with matching size."""
+        from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
+        from sglang.srt.hardware_backend.mlx.model_runner_stub import (
+            MlxModelRunnerStub,
+        )
+
+        logger.info("Initializing MlxModelRunner for end-to-end MLX inference")
+        init_kwargs = dict(
+            model_path=self.server_args.model_path,
+            trust_remote_code=self.server_args.trust_remote_code,
+            disable_radix_cache=self.server_args.disable_radix_cache,
+            mem_fraction_static=self.server_args.mem_fraction_static,
+            quantization=self.server_args.quantization,
+        )
+        if self.server_args.max_total_tokens is not None:
+            init_kwargs["pool_size"] = self.server_args.max_total_tokens
+        self._mlx_runner = MlxModelRunner(**init_kwargs)
+
+        self._model_runner = MlxModelRunnerStub(
+            model_config=self.model_config,
+            mem_fraction_static=self.server_args.mem_fraction_static,
+            gpu_id=self.gpu_id,
+            tp_rank=self.tp_rank,
+            tp_size=self.tp_size,
+            moe_ep_rank=self.moe_ep_rank,
+            moe_ep_size=self.ep_size,
+            pp_rank=self.pp_rank,
+            pp_size=self.pp_size,
+            nccl_port=self.nccl_port,
+            dp_rank=self.dp_rank,
+            server_args=self.server_args,
+            is_draft_worker=self.is_draft_worker,
+            req_to_token_pool=self.req_to_token_pool,
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            memory_pool_config=self.memory_pool_config,
+            mlx_pool_size=self._mlx_runner.pool_size,
+        )
+
+        self._mlx_active_rids: set[str] = set()
+        self._mlx_pool_initialized = False
+
+    def get_pad_input_ids_func(self):
+        """Override since the stub ModelRunner has no real model."""
+        return None
+
+    def _ensure_mlx_pool_initialized(self):
+        """Lazily initialize MLX cache pools after the stub pools are ready."""
+        if not self._mlx_pool_initialized:
+            self._mlx_runner.init_cache_pools(self._model_runner.req_to_token_pool)
+            self._mlx_pool_initialized = True
+
+    def forward_batch_generation(
+        self,
+        batch: Optional[ScheduleBatch],
+        forward_batch: Optional[ForwardBatch] = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+        is_verify: bool = False,
+        skip_attn_backend_init: Optional[bool] = None,  # deprecated
+    ) -> GenerationBatchResult:
+        """Override to route through MLX model runner."""
+        if batch is not None:
+            self._ensure_mlx_pool_initialized()
+            return self._forward_batch_generation_mlx(batch)
+
+        # Fallback to standard path for None batches
+        return super().forward_batch_generation(
+            batch,
+            forward_batch,
+            pp_proxy_tensors,
+            is_verify,
+            skip_attn_backend_init,
+        )
+
+    def _cleanup_stale_rids(self, forward_mode, current_rids: set[str]) -> None:
+        """Remove MLX state for decode-mode requests that dropped out of the batch."""
+        if forward_mode.is_decode():
+            stale_rids = self._mlx_active_rids - current_rids
+            for rid in stale_rids:
+                self._mlx_runner.remove_request(rid)
+            self._mlx_active_rids = current_rids
+        else:
+            self._mlx_active_rids |= current_rids
+
+    def prepare_for_kv_cache_release(self, req) -> None:
+        """Snapshot MLX auxiliary state at the scheduler's radix insert point."""
+        if self._mlx_runner.has_request(req.rid):
+            self._mlx_runner.store_auxiliary_state_for_request(req.rid)
+            # Prefer the just-snapshotted live auxiliary state for the final
+            # insert. Any older tracked slot is released during component cleanup.
+            req.mamba_last_track_seqlen = None
+
+    def _route_extend_request(self, rid: str, decoding_rids: set[str]) -> str:
+        """Classify a request within an extend / mixed batch.
+
+        Shared by the sync (:meth:`_forward_batch_generation_mlx`) and async
+        (:meth:`_async_extend_batch`) paths so both route identically.
+
+        Returns one of:
+
+        * ``"prefill"``      -- not seen before; start a fresh prefill.
+        * ``"decode"``       -- a genuine single-token decode step mixed into
+          this batch (present in ``batch.decoding_reqs``).
+        * ``"continuation"`` -- a chunked-prefill continuation.  Routing keys on
+          request state, **not** ``seq_len``: a final continuation chunk can be
+          exactly one token, which must still extend.  Routing it as a decode
+          would drop the real token and feed the model its own previous-chunk
+          prediction, silently corrupting the output.
+        """
+        if not self._mlx_runner.has_request(rid):
+            return "prefill"
+        if rid in decoding_rids:
+            return "decode"
+        return "continuation"
+
+    def _forward_batch_generation_mlx(
+        self, batch: ScheduleBatch
+    ) -> GenerationBatchResult:
+        """Run forward pass through the MLX model runner (greedy only)."""
+        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+
+        forward_mode = batch.forward_mode
+        reqs = batch.reqs
+
+        if forward_mode.is_idle():
+            return GenerationBatchResult(
+                logits_output=LogitsProcessorOutput(next_token_logits=None),
+                can_run_cuda_graph=False,
+            )
+
+        self._cleanup_stale_rids(forward_mode, {req.rid for req in reqs})
+
+        next_token_ids_list: list[int] = []
+
+        if forward_mode.is_extend():
+            # Ensure pool is up-to-date before pool-backed attention reads it
+            # for prefix-cached prefills.  Only runs on extend batches.
+            self._mlx_runner.flush_all_decode_kv()
+            input_ids_cpu = batch.input_ids.cpu().tolist()
+            out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist()
+            extend_seq_lens = batch.extend_lens
+
+            offset = 0  # into input_ids_cpu
+            slot_offset = 0  # into out_cache_loc_cpu
+            prefill_rids: list[tuple[str, int]] = []
+            extend_rids: list[tuple[str, int]] = []
+            decode_rids: list[str] = []
+            # Genuine decode steps mixed into this extend batch; see
+            # _route_extend_request.
+            decoding_rids = {r.rid for r in (batch.decoding_reqs or [])}
+
+            for i, req in enumerate(reqs):
+                seq_len = extend_seq_lens[i]
+                req_token_ids = input_ids_cpu[offset : offset + seq_len]
+                req_new_slots = out_cache_loc_cpu[slot_offset : slot_offset + seq_len]
+                offset += seq_len
+                slot_offset += seq_len
+
+                route = self._route_extend_request(req.rid, decoding_rids)
+                if route == "continuation":
+                    next_token = self._mlx_runner.extend(
+                        req.rid, req_token_ids, req_new_slots
+                    )
+                    extend_rids.append((req.rid, next_token))
+                elif route == "decode":
+                    decode_rids.append(req.rid)
+                else:  # "prefill"
+                    prefix_slot_ids = req.prefix_indices.tolist()
+                    full_token_ids = list(req.get_fill_ids())
+                    next_token = self._mlx_runner.prefill(
+                        req_id=req.rid,
+                        new_token_ids=req_token_ids,
+                        full_token_ids=full_token_ids,
+                        prefix_slot_ids=prefix_slot_ids,
+                        new_slot_ids=req_new_slots,
+                        req_pool_idx=req.req_pool_idx,
+                        req=req,
+                    )
+                    prefill_rids.append((req.rid, next_token))
+
+            # Batch decode all existing requests at once
+            if decode_rids:
+                decode_results = self._mlx_runner.decode_batch(decode_rids)
+                decode_map = dict(zip(decode_rids, decode_results))
+            else:
+                decode_map = {}
+
+            prefill_map = dict(prefill_rids)
+            extend_map = dict(extend_rids)
+
+            for req in reqs:
+                if req.rid in decode_map:
+                    next_token_ids_list.append(decode_map[req.rid])
+                elif req.rid in extend_map:
+                    next_token_ids_list.append(extend_map[req.rid])
+                else:
+                    next_token_ids_list.append(prefill_map[req.rid])
+
+        elif forward_mode.is_decode():
+            req_ids = [req.rid for req in reqs]
+            next_token_ids_list = self._mlx_runner.decode_batch(req_ids)
+
+        else:
+            raise ValueError(
+                f"MLX runner does not support forward mode: {forward_mode}"
+            )
+
+        next_token_ids = torch.tensor(
+            next_token_ids_list, dtype=torch.long, device="cpu"
+        )
+
+        return GenerationBatchResult(
+            logits_output=LogitsProcessorOutput(next_token_logits=None),
+            next_token_ids=next_token_ids,
+            can_run_cuda_graph=False,
+        )
+
+    def async_forward_batch_generation_mlx(self, batch: ScheduleBatch) -> tuple[
+        Union[mx.array, None],
+        list[MlxPendingPrefill],
+        list[MlxPendingExtend],
+        Optional[MlxPendingDecode],
+        str,
+    ]:
+        """Start an async (lazy) forward pass through the MLX model runner.
+
+        Returns ``(lazy_result, prefills, extends, decode, mode)``:
+
+        * ``lazy_result`` — an ``mx.array`` that, when evaluated, forces
+          materialisation of the whole batch's outputs.  ``None`` for
+          idle batches.
+        * ``prefills`` — list of :class:`MlxPendingPrefill` for new
+          requests in an extend batch.
+        * ``extends`` — list of :class:`MlxPendingExtend` for chunked
+          prefill continuations in an extend batch.
+        * ``decode`` — :class:`MlxPendingDecode` for the decode
+          sub-batch (covers full decode mode AND mixed decodes inside
+          an extend batch).
+        * ``mode`` — one of ``"idle"``, ``"decode"``, ``"extend"``.
+
+        The caller must make sure the returned pendings are fed into a
+        subsequent ``mx.async_eval`` or ``.item()`` / ``.tolist()`` call
+        — :meth:`finalize_mlx_result` does that.
+        """
+        self._ensure_mlx_pool_initialized()
+
+        forward_mode = batch.forward_mode
+        reqs = batch.reqs
+
+        if forward_mode.is_idle():
+            return None, [], [], None, "idle"
+
+        self._cleanup_stale_rids(forward_mode, {req.rid for req in reqs})
+
+        if forward_mode.is_decode():
+            req_ids = [req.rid for req in reqs]
+            pending_decode = self._mlx_runner.decode_batch_start(req_ids)
+            mx.async_eval(pending_decode.lazy_tokens)
+            return pending_decode.lazy_tokens, [], [], pending_decode, "decode"
+
+        if forward_mode.is_extend():
+            # TODO (changminbark): Implement per-batch flushing using prefix_slot_ids
+            # Ensure the pool is up-to-date before pool-backed attention
+            # reads it for prefix-cached prefills. Mirror the sync path.
+            self._mlx_runner.flush_all_decode_kv()
+            return self._async_extend_batch(batch)
+
+        raise ValueError(
+            f"MLX async runner does not support forward mode: {forward_mode}"
+        )
+
+    def _async_extend_batch(self, batch: ScheduleBatch) -> tuple[
+        Union[mx.array, None],
+        list[MlxPendingPrefill],
+        list[MlxPendingExtend],
+        Optional[MlxPendingDecode],
+        str,
+    ]:
+        """Launch each request in an EXTEND batch lazily and kick GPU work."""
+        reqs = batch.reqs
+        input_ids_cpu = batch.input_ids.cpu().tolist()
+        out_cache_loc_cpu = batch.out_cache_loc.cpu().tolist()
+        extend_seq_lens = batch.extend_lens
+
+        offset = 0
+        slot_offset = 0
+        pending_prefills: list[MlxPendingPrefill] = []
+        pending_extends: list[MlxPendingExtend] = []
+        mixed_decode_rids: list[str] = []
+        # Genuine decode steps mixed into this extend batch; see
+        # _route_extend_request.
+        decoding_rids = {r.rid for r in (batch.decoding_reqs or [])}
+
+        for i, req in enumerate(reqs):
+            seq_len = extend_seq_lens[i]
+            req_token_ids = input_ids_cpu[offset : offset + seq_len]
+            req_new_slots = out_cache_loc_cpu[slot_offset : slot_offset + seq_len]
+            offset += seq_len
+            slot_offset += seq_len
+
+            route = self._route_extend_request(req.rid, decoding_rids)
+            if route == "continuation":
+                pending_extends.append(
+                    self._mlx_runner.extend_start(
+                        req_id=req.rid,
+                        new_token_ids=req_token_ids,
+                        new_slot_ids=req_new_slots,
+                    )
+                )
+            elif route == "decode":
+                mixed_decode_rids.append(req.rid)
+            else:  # "prefill"
+                prefix_slot_ids = req.prefix_indices.tolist()
+                full_token_ids = list(req.get_fill_ids())
+                pending_prefills.append(
+                    self._mlx_runner.prefill_start(
+                        req_id=req.rid,
+                        new_token_ids=req_token_ids,
+                        full_token_ids=full_token_ids,
+                        prefix_slot_ids=prefix_slot_ids,
+                        new_slot_ids=req_new_slots,
+                        req_pool_idx=req.req_pool_idx,
+                        req=req,
+                    )
+                )
+
+        pending_mixed_decode: Optional[MlxPendingDecode] = None
+        if mixed_decode_rids:
+            pending_mixed_decode = self._mlx_runner.decode_batch_start(
+                mixed_decode_rids
+            )
+
+        # Stack lazy tokens so the caller has a single handle to evaluate
+        # after CPU scheduling work.  We also hand every cache buffer
+        # (and the decode cache arrays) to mx.async_eval so the GPU
+        # kernel-launch stream sees everything the next step depends on
+        # before we actually block on anything.
+        prefill_ext_tokens: list[mx.array] = [p.lazy_token for p in pending_prefills]
+        prefill_ext_tokens.extend(e.lazy_token for e in pending_extends)
+
+        async_args: list[mx.array] = []
+        if prefill_ext_tokens:
+            lazy_stacked = mx.stack(prefill_ext_tokens, axis=0)
+            async_args.append(lazy_stacked)
+        else:
+            lazy_stacked = None
+
+        for p in pending_prefills:
+            async_args.extend(self._cache_state(p.cache))
+        for e in pending_extends:
+            async_args.extend(self._cache_state(self._mlx_runner._req_caches[e.req_id]))
+        if pending_mixed_decode is not None:
+            async_args.append(pending_mixed_decode.lazy_tokens)
+            for c_list in pending_mixed_decode.caches:
+                async_args.extend(self._cache_state(c_list))
+
+        if async_args:
+            mx.async_eval(*async_args)
+
+        return (
+            lazy_stacked,
+            pending_prefills,
+            pending_extends,
+            pending_mixed_decode,
+            "extend",
+        )
+
+    @staticmethod
+    def _cache_state(cache_list) -> list[mx.array]:
+        """Flatten a per-layer cache list to its ``state`` arrays."""
+        arrays: list[mx.array] = []
+
+        def collect(value):
+            if isinstance(value, mx.array):
+                arrays.append(value)
+            elif value is None:
+                return
+            elif isinstance(value, (list, tuple)):
+                for item in value:
+                    collect(item)
+            elif isinstance(value, dict):
+                for item in value.values():
+                    collect(item)
+
+        for cache in cache_list:
+            collect(getattr(cache, "state", ()))
+        return arrays
+
+    def async_chained_decode_mlx(
+        self,
+        prev_pending: MlxPendingDecode,
+    ) -> tuple[mx.array, list, list, MlxPendingDecode, str]:
+        """Launch a decode step that chains off a still-lazy previous decode.
+
+        This is the "no idle gap" pipelining primitive: build the next
+        decode's compute graph using ``prev_pending.lazy_tokens`` (still
+        unevaluated) as its input ids, hand the combined graph to
+        ``mx.async_eval``, and return.  The GPU runs the new step
+        immediately after ``prev_pending`` with no scheduling gap, while
+        the caller is free to block on ``prev_pending`` and run CPU-side
+        bookkeeping.
+
+        Preconditions (caller must ensure):
+
+        * ``prev_pending`` was produced by a previous decode start
+          (either :meth:`async_forward_batch_generation_mlx` in decode
+          mode or a previous :meth:`async_chained_decode_mlx`).
+        * The batch composition for this step is identical to
+          ``prev_pending`` — same requests, same order.  Composition
+          changes (finished reqs, new prefills) must break the chain.
+        * ``prev_pending`` should be finalised BEFORE the returned
+          pending, so per-request token lists are appended in order.
+
+        Returns a 5-tuple matching
+        :meth:`async_forward_batch_generation_mlx` for the decode case:
+        ``(lazy_tokens, [], [], pending_decode, "decode")``.  The empty
+        prefill/extend lists are always absent for chained decodes.
+        """
+        pending = self._mlx_runner.decode_batch_start_chained(prev_pending)
+        mx.async_eval(pending.lazy_tokens)
+        return pending.lazy_tokens, [], [], pending, "decode"
+
+    def finalize_mlx_result(
+        self,
+        prefills: list[MlxPendingPrefill],
+        extends: list[MlxPendingExtend],
+        decode: Optional[MlxPendingDecode],
+        mode: str,
+        reqs: list,
+    ) -> GenerationBatchResult:
+        """Materialise a lazy MLX result into a :class:`GenerationBatchResult`.
+
+        The blocking wait happens inside ``decode_batch_finalize`` /
+        ``prefill_finalize`` / ``extend_finalize`` via ``.tolist()`` /
+        ``.item()`` on the specific lazy outputs.
+        """
+        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+
+        if mode == "idle":
+            return GenerationBatchResult(
+                logits_output=LogitsProcessorOutput(next_token_logits=None),
+                can_run_cuda_graph=False,
+            )
+
+        if mode == "decode":
+            assert decode is not None
+            next_tokens_list = self._mlx_runner.decode_batch_finalize(decode)
+
+        elif mode == "extend":
+            prefill_map: dict[str, int] = {}
+            for pending_p in prefills:
+                prefill_map[pending_p.req_id] = self._mlx_runner.prefill_finalize(
+                    pending_p
+                )
+
+            extend_map: dict[str, int] = {}
+            for pending_e in extends:
+                extend_map[pending_e.req_id] = self._mlx_runner.extend_finalize(
+                    pending_e
+                )
+
+            decode_map: dict[str, int] = {}
+            if decode is not None:
+                mixed_tokens = self._mlx_runner.decode_batch_finalize(decode)
+                decode_map = {
+                    rid: tok for rid, tok in zip(decode.req_ids, mixed_tokens)
+                }
+
+            next_tokens_list = []
+            for req in reqs:
+                if req.rid in decode_map:
+                    next_tokens_list.append(decode_map[req.rid])
+                elif req.rid in extend_map:
+                    next_tokens_list.append(extend_map[req.rid])
+                else:
+                    next_tokens_list.append(prefill_map[req.rid])
+
+        else:
+            raise ValueError(f"Unknown MLX async mode: {mode}")
+
+        next_token_ids = torch.tensor(next_tokens_list, dtype=torch.long, device="cpu")
+        return GenerationBatchResult(
+            logits_output=LogitsProcessorOutput(next_token_logits=None),
+            next_token_ids=next_token_ids,
+            can_run_cuda_graph=False,
+        )

python/sglang/srt/hardware_backend/musa/__init__.py

@@ -0,0 +1 @@
+# MUSA (Moore Threads GPU) hardware backend

python/sglang/srt/hardware_backend/musa/attention/__init__.py

@@ -0,0 +1,3 @@
+from .flashattention_backend import MusaFlashAttentionBackend
+
+__all__ = ["MusaFlashAttentionBackend"]

python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py

@@ -0,0 +1,951 @@
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+
+import torch
+from flash_attn_interface import flash_attn_varlen_func
+from flash_attn_interface import flash_attn_with_kvcache as mate_flash_attn_with_kvcache
+from flash_attn_interface import get_scheduler_metadata
+
+from sglang.srt.distributed import get_pp_group, get_pp_indices
+from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.musa.layers.utils.cp_utils import (
+    musa_cp_attn_forward_extend as cp_attn_forward_extend,
+)
+from sglang.srt.layers.attention.flashattention_backend import (
+    FlashAttentionBackend,
+    FlashAttentionMultiStepBackend,
+    merge_state_v2_wrapper,
+)
+from sglang.srt.layers.radix_attention import AttentionType
+from sglang.srt.layers.utils.cp_utils import (
+    cp_allgather_and_save_kv_cache,
+)
+from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.runtime_context import get_server_args
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+# Global workspace buffer for MLA
+_MATE_MLA_WORKSPACE_SIZE_BYTES = 128 * 1024 * 1024
+
+# Cache for non-MLA scheduler metadata by prefix
+_MATE_NO_MLA_SCHEDULER_METADATA_DICT: dict = {}
+_MATE_NO_MLA_SCHEDULER_METADATA_LOCK = threading.Lock()
+
+# Global reference to the current backend instance (set during __init__)
+_CURRENT_BACKEND: Optional[MusaFlashAttentionBackend] = None
+
+
+def _compute_scheduler_metadata(
+    backend: MusaFlashAttentionBackend,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k_new: Optional[torch.Tensor],
+    cache_seqlens: torch.Tensor,
+    max_seqlen_q: int,
+    page_size: int,
+    causal: bool,
+    window_size: Tuple[int, int],
+    num_splits: int,
+) -> Tuple[torch.Tensor, bool] | torch.Tensor:
+    """Compute scheduler metadata based on backend's current state."""
+    global _MATE_NO_MLA_SCHEDULER_METADATA_DICT
+
+    layer = backend._current_layer
+    current_layer_id = layer.layer_id
+    batch_size = cu_seqlens_q.shape[-1] - 1
+
+    # Determine if scheduler metadata should be updated
+    should_update = True
+    pp_group = get_pp_group()
+    pp_rank = pp_group.rank_in_group
+    start_layer_id, _ = get_pp_indices(
+        backend.num_hidden_layers, pp_group.rank_in_group, pp_group.world_size
+    )
+    if backend._current_can_run_tbo and pp_rank == 0:
+        start_layer_id += (
+            backend.first_k_dense_replace
+            if backend.first_k_dense_replace is not None
+            else 0
+        )
+
+    if backend.full_attention_interval is not None:
+        start_layer_id += backend.full_attention_interval - 1
+
+    if current_layer_id > start_layer_id:
+        should_update = False
+
+    if envs.SGLANG_MUSA_FA3_FORCE_UPDATE_METADATA.get():
+        should_update = True
+
+    if backend.use_mla:
+        from sglang.srt.runtime_context import get_buffer
+
+        workspace = get_buffer(
+            "musa_mate_mla_workspace",
+            lambda: torch.empty(
+                _MATE_MLA_WORKSPACE_SIZE_BYTES, device=backend.device, dtype=torch.uint8
+            ),
+        )
+        return (workspace, not should_update)
+    else:
+        with _MATE_NO_MLA_SCHEDULER_METADATA_LOCK:
+            if (
+                should_update
+                or backend._current_prefix not in _MATE_NO_MLA_SCHEDULER_METADATA_DICT
+            ):
+                _MATE_NO_MLA_SCHEDULER_METADATA_DICT[backend._current_prefix] = (
+                    get_scheduler_metadata(
+                        batch_size=batch_size,
+                        num_heads_q=layer.tp_q_head_num,
+                        num_heads_kv=layer.tp_k_head_num,
+                        headdim=layer.qk_head_dim,
+                        headdim_v=layer.v_head_dim,
+                        cache_seqlens=cache_seqlens,
+                        cu_seqlens_q=cu_seqlens_q,
+                        cu_seqlens_k_new=cu_seqlens_k_new,
+                        max_seqlen_q=max_seqlen_q,
+                        max_seqlen_k=backend._current_max_seqlen_k,
+                        page_size=page_size,
+                        causal=causal,
+                        window_size=window_size,
+                        num_splits=num_splits,
+                    )
+                )
+            return _MATE_NO_MLA_SCHEDULER_METADATA_DICT[backend._current_prefix]
+
+
+def flash_attn_with_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    qv: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_seqlens: Optional[Union[int, torch.Tensor]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    rotary_seqlens: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: Tuple[int, int] = (-1, -1),
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    rotary_interleaved: bool = True,
+    scheduler_metadata: Optional[torch.Tensor] = None,
+    num_splits: int = 0,
+    pack_gqa=None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks=None,
+    score_mod=None,
+    aux_tensors=None,
+    ver=3,
+):
+    """MUSA flash_attn_with_kvcache wrapper that auto-injects scheduler_metadata."""
+    if ver != 3:
+        raise ValueError("Only ver=3 is supported for MUSA FA3.")
+
+    if scheduler_metadata is None and _CURRENT_BACKEND is not None:
+        backend = _CURRENT_BACKEND
+        # Ensure backend has been properly set up for this call
+        if backend._current_layer is not None:
+            page_size = k_cache.shape[1] if k_cache is not None else 1
+            scheduler_metadata = _compute_scheduler_metadata(
+                backend=backend,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k_new=cu_seqlens_k_new,
+                cache_seqlens=cache_seqlens,
+                max_seqlen_q=max_seqlen_q,
+                page_size=page_size,
+                causal=causal,
+                window_size=window_size,
+                num_splits=num_splits,
+            )
+
+    return mate_flash_attn_with_kvcache(
+        q=q,
+        k_cache=k_cache,
+        v_cache=v_cache,
+        k=k,
+        v=v,
+        qv=qv,
+        rotary_cos=rotary_cos,
+        rotary_sin=rotary_sin,
+        cache_seqlens=cache_seqlens,
+        cache_batch_idx=cache_batch_idx,
+        cache_leftpad=cache_leftpad,
+        page_table=page_table,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k_new=cu_seqlens_k_new,
+        max_seqlen_q=max_seqlen_q,
+        rotary_seqlens=rotary_seqlens,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        window_size=window_size,
+        attention_chunk=attention_chunk,
+        softcap=softcap,
+        rotary_interleaved=rotary_interleaved,
+        scheduler_metadata=scheduler_metadata,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        return_softmax_lse=return_softmax_lse,
+        sinks=sinks,
+    )
+
+
+class MusaFlashAttentionBackend(FlashAttentionBackend):
+    def __init__(self, model_runner: ModelRunner, **kwargs):
+        super().__init__(model_runner, **kwargs)
+        self.num_hidden_layers = model_runner.model_config.num_hidden_layers
+        self.first_k_dense_replace = model_runner.model_config.first_k_dense_replace
+        self.full_attention_interval = model_runner.model_config.full_attention_interval
+
+        # State for current attention call (simplified from thread‑local context)
+        self._current_layer: Optional[RadixAttention] = None
+        self._current_prefix: str = ""
+        self._current_max_seqlen_k: int = 0
+        self._current_can_run_tbo: bool = False
+
+        # Disable default scheduler metadata for fa3
+        self._get_scheduler_metadata = None
+
+        # Register this backend as the global current instance for the wrapper
+        global _CURRENT_BACKEND
+        _CURRENT_BACKEND = self
+
+    def _set_current_state(
+        self, layer: RadixAttention, prefix: str, max_seqlen_k: int, can_run_tbo: bool
+    ):
+        """Set the dynamic state for the upcoming flash attention call."""
+        self._current_layer = layer
+        self._current_prefix = prefix
+        self._current_max_seqlen_k = max_seqlen_k
+        self._current_can_run_tbo = can_run_tbo
+
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
+        sinks: Optional[torch.Tensor] = None,
+    ):
+        if k is not None:
+            assert v is not None
+
+            is_cp_mode = (
+                forward_batch.forward_mode.is_context_parallel_extend()
+                and forward_batch.attn_cp_metadata is not None
+                and self.attn_cp_size > 1
+            )
+
+            if save_kv_cache and not is_cp_mode:
+                cache_loc = (
+                    forward_batch.out_cache_loc
+                    if not layer.is_cross_attention
+                    else forward_batch.encoder_out_cache_loc
+                )
+                if not self.use_mla:
+                    self.token_to_kv_pool.set_kv_buffer(
+                        layer,
+                        KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
+                        k,
+                        v,
+                        layer.k_scale,
+                        layer.v_scale,
+                    )
+                else:
+                    self.token_to_kv_pool.set_mla_kv_buffer(
+                        layer,
+                        cache_loc,
+                        k,
+                        k_rope,
+                    )
+            if is_cp_mode:
+                cp_allgather_and_save_kv_cache(
+                    forward_batch,
+                    layer,
+                    k,
+                    v,
+                    self.attn_cp_size,
+                    swa_loc=(
+                        self.forward_metadata.swa_out_cache_loc
+                        if self.use_sliding_window_kv_pool
+                        else None
+                    ),
+                )
+
+        metadata = self.forward_metadata
+
+        is_swa_layer = (
+            layer.sliding_window_size is not None and layer.sliding_window_size > -1
+        )
+        window_size = (layer.sliding_window_size, 0) if is_swa_layer else (-1, -1)
+        k_descale, v_descale = None, None
+        if (
+            self.kv_cache_dtype_str != "auto"
+            and layer.head_dim <= 256
+            and self.fa_impl_ver != 4
+        ):
+            if layer.k_scale is not None:
+                descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
+                k_descale = layer.k_scale.expand(descale_shape)
+                v_descale = layer.v_scale.expand(descale_shape)
+            q = q.to(self.kv_cache_dtype)
+            q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
+            k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None
+        causal = True
+        if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
+            causal = False
+
+        use_local_attn = (
+            self.has_local_attention
+            and self.attention_chunk_size is not None
+            and metadata.local_attn_metadata is not None
+            and (hasattr(layer, "use_irope") and layer.use_irope)
+        )
+
+        use_cascade_attn = (
+            forward_batch.forward_mode.is_target_verify()
+            and self.topk > 1
+            and not is_swa_layer
+        )
+
+        kwargs = {}
+        if sinks is not None:
+            kwargs["sinks"] = sinks
+
+        if use_local_attn:
+            local_metadata = metadata.local_attn_metadata
+            page_table = local_metadata.local_block_table
+            cu_seqlens_q = local_metadata.local_query_start_loc
+            cache_seqlens = local_metadata.local_seqused_k
+            max_seqlen_q = local_metadata.local_max_query_len
+            max_seqlen_k = local_metadata.local_max_seq_len
+        elif is_swa_layer and metadata.swa_spec_metadata is not None:
+            swa_spec_metadata = metadata.swa_spec_metadata
+            page_table = swa_spec_metadata.page_table
+            cu_seqlens_q = swa_spec_metadata.cu_seqlens_q
+            cache_seqlens = swa_spec_metadata.cache_seqlens_int32
+            max_seqlen_q = swa_spec_metadata.max_seq_len_q
+            cu_seqlens_k = swa_spec_metadata.cu_seqlens_k
+            max_seqlen_k = swa_spec_metadata.max_seq_len_k
+        else:
+            page_table = metadata.page_table
+            if is_swa_layer and self.use_sliding_window_kv_pool:
+                if metadata.swa_page_table is not None:
+                    page_table = metadata.swa_page_table
+                else:
+                    page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
+                        metadata.page_table
+                    )
+            cu_seqlens_q = metadata.cu_seqlens_q
+            cache_seqlens = metadata.cache_seqlens_int32
+            max_seqlen_q = metadata.max_seq_len_q
+            cu_seqlens_k = metadata.cu_seqlens_k
+            max_seqlen_k = metadata.max_seq_len_k
+
+        # Set current state for the flash attention call
+        self._set_current_state(
+            layer=layer,
+            prefix="forward_extend",
+            max_seqlen_k=max_seqlen_k,
+            can_run_tbo=forward_batch.can_run_tbo,
+        )
+        if not self.use_mla:
+            key_cache, value_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+
+            key_cache = key_cache.view(
+                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
+            )
+            value_cache = value_cache.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+            )
+            if layer.is_cross_attention:
+                page_table = metadata.encoder_page_table
+                cache_seqlens = metadata.encoder_lens_int32
+                cu_seqlens_k = metadata.encoder_cu_seqlens_k
+                window_size = (-1, -1)
+
+            if (
+                forward_batch.forward_mode.is_context_parallel_extend()
+                and forward_batch.attn_cp_metadata is not None
+                and self.attn_cp_size > 1
+            ):
+
+                def _fa_cp_attn(
+                    q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp
+                ):
+                    return flash_attn_with_kvcache(
+                        q=q_chunk,
+                        k_cache=key_cache,
+                        v_cache=value_cache,
+                        page_table=page_table,
+                        cache_seqlens=cache_seqlens_cp,
+                        cu_seqlens_q=cu_seqlens_q_cp,
+                        cu_seqlens_k_new=(cu_seqlens_k if not use_local_attn else None),
+                        max_seqlen_q=max_seqlen_q_cp,
+                        softmax_scale=layer.scaling,
+                        causal=False if use_cascade_attn else causal,
+                        window_size=window_size,
+                        softcap=layer.logit_cap,
+                        k_descale=k_descale,
+                        v_descale=v_descale,
+                        return_softmax_lse=use_cascade_attn,
+                        num_splits=self.num_splits,
+                        **kwargs,
+                    )
+
+                result = cp_attn_forward_extend(
+                    self,
+                    forward_batch,
+                    q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    self.device,
+                    _fa_cp_attn,
+                )
+            elif (
+                forward_batch.extend_prefix_lens_cpu is not None
+                and any(forward_batch.extend_prefix_lens_cpu)
+            ) or forward_batch.forward_mode.is_target_verify():
+                result = flash_attn_with_kvcache(
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=page_table,
+                    cache_seqlens=cache_seqlens,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
+                    max_seqlen_q=max_seqlen_q,
+                    softmax_scale=layer.scaling,
+                    causal=False if use_cascade_attn else causal,
+                    window_size=window_size,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=use_cascade_attn,
+                    num_splits=self.num_splits,
+                    **kwargs,
+                )
+            else:
+                output = flash_attn_varlen_func(
+                    q=q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k=k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype),
+                    v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim).to(q.dtype),
+                    cu_seqlens_q=metadata.cu_seqlens_q,
+                    cu_seqlens_k=metadata.cu_seqlens_q,
+                    max_seqlen_q=metadata.max_seq_len_q,
+                    max_seqlen_k=metadata.max_seq_len_q,
+                    softmax_scale=layer.scaling,
+                    causal=True,
+                    return_softmax_lse=forward_batch.mha_return_lse,
+                    **kwargs,
+                )
+                if forward_batch.mha_return_lse:
+                    output, lse, *rest = output
+                    lse = torch.transpose(lse, 0, 1).contiguous()
+                    return (
+                        output.view(-1, layer.tp_q_head_num * layer.v_head_dim),
+                        lse,
+                    )
+                return output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+            if use_cascade_attn:
+                # Update state for the second call
+                self._current_prefix = "forward_extend_use_cascade_attn"
+                self._current_max_seqlen_k = (
+                    self.forward_metadata_spec_decode_expand.max_seq_len_k
+                )
+
+                o, softmax_lse, *rest = result
+                o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache(
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k_cache=key_cache.view(-1, 1, layer.tp_k_head_num, layer.head_dim),
+                    v_cache=value_cache.view(
+                        -1, 1, layer.tp_v_head_num, layer.head_dim
+                    ),
+                    page_table=self.forward_metadata_spec_decode_expand.page_table,
+                    cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                    cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                    cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                    max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                    softmax_scale=layer.scaling,
+                    causal=False,
+                    window_size=window_size,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=True,
+                    num_splits=self.num_splits,
+                    **kwargs,
+                )
+                o, _ = merge_state_v2_wrapper(
+                    o,
+                    softmax_lse.T.contiguous(),
+                    o_expand,
+                    softmax_lse_expand.T.contiguous(),
+                )
+            else:
+                o = result
+        else:
+            if (
+                forward_batch.attn_attend_prefix_cache is not None
+                and not forward_batch.forward_mode.is_target_verify()
+                and not forward_batch.forward_mode.is_draft_extend_v2()
+            ):
+                if forward_batch.attn_attend_prefix_cache:
+                    assert not get_server_args().disable_chunked_prefix_cache
+                    assert forward_batch.prefix_chunk_idx is not None
+                    assert forward_batch.prefix_chunk_cu_seq_lens is not None
+                    assert forward_batch.prefix_chunk_max_seq_lens is not None
+
+                    chunk_idx = forward_batch.prefix_chunk_idx
+                    assert chunk_idx >= 0
+
+                    assert forward_batch.mha_return_lse
+                    output = flash_attn_varlen_func(
+                        q=q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                        k=k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype),
+                        v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim).to(q.dtype),
+                        cu_seqlens_q=metadata.cu_seqlens_q,
+                        cu_seqlens_k=forward_batch.prefix_chunk_cu_seq_lens[chunk_idx],
+                        max_seqlen_q=metadata.max_seq_len_q,
+                        max_seqlen_k=forward_batch.prefix_chunk_max_seq_lens[chunk_idx],
+                        softmax_scale=layer.scaling,
+                        causal=False,
+                        return_softmax_lse=True,
+                        **kwargs,
+                    )
+                else:
+                    cu_seqlens_k = (
+                        metadata.cu_seqlens_q
+                        if not forward_batch.mha_one_shot
+                        else metadata.cu_seqlens_k
+                    )
+                    max_seqlen_k = (
+                        metadata.max_seq_len_q
+                        if not forward_batch.mha_one_shot
+                        else metadata.max_seq_len_k
+                    )
+                    output = flash_attn_varlen_func(
+                        q=q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                        k=k.view(-1, layer.tp_k_head_num, layer.head_dim).to(q.dtype),
+                        v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim).to(q.dtype),
+                        cu_seqlens_q=metadata.cu_seqlens_q,
+                        cu_seqlens_k=cu_seqlens_k,
+                        max_seqlen_q=metadata.max_seq_len_q,
+                        max_seqlen_k=max_seqlen_k,
+                        softmax_scale=layer.scaling,
+                        causal=True,
+                        return_softmax_lse=forward_batch.mha_return_lse,
+                        **kwargs,
+                    )
+                if forward_batch.mha_return_lse:
+                    output, lse, *rest = output
+                    lse = torch.transpose(lse, 0, 1).contiguous()
+                    return output, lse
+                return output
+            else:
+                kv_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(
+                    q.dtype
+                )
+                k_rope = kv_cache[:, :, layer.v_head_dim :]
+                c_kv = kv_cache[:, :, : layer.v_head_dim]
+                k_rope_cache = k_rope.view(
+                    -1,
+                    self.page_size,
+                    layer.tp_k_head_num,
+                    layer.head_dim - layer.v_head_dim,
+                )
+                c_kv_cache = c_kv.view(
+                    -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+                )
+                if q_rope is not None:
+                    q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+                    q_rope = q_rope.view(
+                        -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                    )
+                else:
+                    q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+                    q_nope = q_all[:, :, : layer.v_head_dim]
+                    q_rope = q_all[:, :, layer.v_head_dim :]
+
+                result = flash_attn_with_kvcache(
+                    q=q_rope,
+                    k_cache=k_rope_cache,
+                    v_cache=c_kv_cache,
+                    qv=q_nope,
+                    page_table=page_table,
+                    cache_seqlens=cache_seqlens,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
+                    max_seqlen_q=max_seqlen_q,
+                    softmax_scale=layer.scaling,
+                    causal=False if use_cascade_attn else causal,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=use_cascade_attn,
+                    num_splits=self.num_splits,
+                )
+                if use_cascade_attn:
+                    self._current_prefix = "forward_extend_use_cascade_attn"
+                    self._current_max_seqlen_k = (
+                        self.forward_metadata_spec_decode_expand.max_seq_len_k
+                    )
+
+                    o, softmax_lse, *rest = result
+                    o_expand, softmax_lse_expand, *rest_expand = (
+                        flash_attn_with_kvcache(
+                            q=q_rope,
+                            k_cache=k_rope_cache,
+                            v_cache=c_kv_cache,
+                            qv=q_nope,
+                            page_table=self.forward_metadata_spec_decode_expand.page_table,
+                            cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                            cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                            cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                            max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                            softmax_scale=layer.scaling,
+                            causal=False,
+                            window_size=window_size,
+                            softcap=layer.logit_cap,
+                            k_descale=k_descale,
+                            v_descale=v_descale,
+                            return_softmax_lse=True,
+                            num_splits=self.num_splits,
+                        )
+                    )
+                    o, _ = merge_state_v2_wrapper(
+                        o,
+                        softmax_lse.T.contiguous(),
+                        o_expand,
+                        softmax_lse_expand.T.contiguous(),
+                    )
+                else:
+                    o = result
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
+        sinks: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                cache_loc = (
+                    forward_batch.out_cache_loc
+                    if not layer.is_cross_attention
+                    else forward_batch.encoder_out_cache_loc
+                )
+                if not self.use_mla:
+                    self.token_to_kv_pool.set_kv_buffer(
+                        layer,
+                        KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
+                        k,
+                        v,
+                        layer.k_scale,
+                        layer.v_scale,
+                    )
+                else:
+                    self.token_to_kv_pool.set_mla_kv_buffer(
+                        layer,
+                        cache_loc,
+                        k,
+                        k_rope,
+                    )
+
+        metadata = self.forward_metadata
+        local_attn_metadata = getattr(metadata, "local_attn_metadata", None)
+        use_local_attn = (
+            self.has_local_attention
+            and self.attention_chunk_size is not None
+            and local_attn_metadata is not None
+            and (hasattr(layer, "use_irope") and layer.use_irope)
+        )
+
+        use_cascade_attn = forward_batch.spec_info is not None and self.topk > 1
+
+        is_swa_layer = (
+            layer.sliding_window_size is not None and layer.sliding_window_size > -1
+        )
+        window_size = (layer.sliding_window_size, 0) if is_swa_layer else (-1, -1)
+
+        causal = True
+        if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
+            causal = False
+
+        kwargs = {}
+        if sinks is not None:
+            kwargs["sinks"] = sinks
+
+        k_descale, v_descale = None, None
+        if self.kv_cache_dtype_str != "auto" and layer.head_dim <= 256:
+            if layer.k_scale is not None:
+                descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
+                k_descale = layer.k_scale.expand(descale_shape)
+                v_descale = layer.v_scale.expand(descale_shape)
+            q = q.to(self.kv_cache_dtype)
+            q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
+            k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None
+
+        # Set current state for the flash attention call
+        self._set_current_state(
+            layer=layer,
+            prefix="forward_decode",
+            max_seqlen_k=metadata.max_seq_len_k,
+            can_run_tbo=forward_batch.can_run_tbo,
+        )
+        if not self.use_mla:
+            key_cache, value_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+            key_cache = key_cache.view(
+                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
+            )
+            value_cache = value_cache.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+            )
+
+            if layer.is_cross_attention:
+                o = flash_attn_with_kvcache(
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=metadata.encoder_page_table,
+                    cache_seqlens=metadata.encoder_lens_int32,
+                    cu_seqlens_q=metadata.cu_seqlens_q,
+                    cu_seqlens_k_new=metadata.encoder_cu_seqlens_k,
+                    max_seqlen_q=1,
+                    softmax_scale=layer.scaling,
+                    causal=False,
+                    window_size=(-1, -1),
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    num_splits=self.num_splits,
+                    **kwargs,
+                )
+            elif use_local_attn:
+                o = flash_attn_with_kvcache(
+                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=local_attn_metadata.local_block_table,
+                    cache_seqlens=local_attn_metadata.local_seqused_k,
+                    cu_seqlens_q=local_attn_metadata.local_query_start_loc,
+                    cu_seqlens_k_new=None,
+                    max_seqlen_q=local_attn_metadata.local_max_query_len,
+                    softmax_scale=layer.scaling,
+                    causal=True,
+                    window_size=(-1, -1),
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    num_splits=self.num_splits,
+                    **kwargs,
+                )
+            else:
+                page_table = metadata.page_table
+                if is_swa_layer and self.use_sliding_window_kv_pool:
+                    if metadata.swa_page_table is not None:
+                        page_table = metadata.swa_page_table
+                    else:
+                        page_table = (
+                            self.token_to_kv_pool.translate_loc_from_full_to_swa(
+                                metadata.page_table
+                            )
+                        )
+                cache_seqlens = metadata.cache_seqlens_int32
+                cu_seqlens_k = metadata.cu_seqlens_k
+                max_seqlen_q = metadata.max_seq_len_q
+                q_reshaped = q.contiguous().view(
+                    -1, layer.tp_q_head_num, layer.head_dim
+                )
+
+                result = flash_attn_with_kvcache(
+                    q=q_reshaped,
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    page_table=page_table,
+                    cache_seqlens=cache_seqlens,
+                    cu_seqlens_q=metadata.cu_seqlens_q,
+                    max_seqlen_q=max_seqlen_q,
+                    softmax_scale=layer.scaling,
+                    causal=False if use_cascade_attn else causal,
+                    window_size=window_size,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=use_cascade_attn,
+                    num_splits=self.num_splits,
+                    **kwargs,
+                )
+                if use_cascade_attn:
+                    self._current_prefix = "forward_decode_use_cascade_attn"
+                    self._current_max_seqlen_k = (
+                        self.forward_metadata_spec_decode_expand.max_seq_len_k
+                    )
+
+                    o, softmax_lse, *rest = result
+                    o_expand, softmax_lse_expand, *rest_expand = (
+                        flash_attn_with_kvcache(
+                            q=q_reshaped,
+                            k_cache=key_cache,
+                            v_cache=value_cache,
+                            page_table=self.forward_metadata_spec_decode_expand.page_table,
+                            cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                            cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                            cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                            max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                            softmax_scale=layer.scaling,
+                            causal=False,
+                            window_size=window_size,
+                            softcap=layer.logit_cap,
+                            k_descale=k_descale,
+                            v_descale=v_descale,
+                            return_softmax_lse=True,
+                            num_splits=self.num_splits,
+                            **kwargs,
+                        )
+                    )
+                    o, _ = merge_state_v2_wrapper(
+                        o,
+                        softmax_lse.T.contiguous(),
+                        o_expand,
+                        softmax_lse_expand.T.contiguous(),
+                    )
+                else:
+                    o = result
+        else:
+            kv_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype)
+            k_rope = kv_cache[:, :, layer.v_head_dim :]
+            c_kv = kv_cache[:, :, : layer.v_head_dim]
+            k_rope_cache = k_rope.view(
+                -1,
+                self.page_size,
+                layer.tp_k_head_num,
+                layer.head_dim - layer.v_head_dim,
+            )
+            c_kv_cache = c_kv.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+            )
+
+            if q_rope is not None:
+                q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+                q_rope = q_rope.view(
+                    -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                )
+            else:
+                q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+                q_nope = q_all[:, :, : layer.v_head_dim]
+                q_rope = q_all[:, :, layer.v_head_dim :]
+            max_seqlen_q = metadata.max_seq_len_q
+
+            result = flash_attn_with_kvcache(
+                q=q_rope,
+                k_cache=k_rope_cache,
+                v_cache=c_kv_cache,
+                qv=q_nope,
+                page_table=metadata.page_table,
+                cache_seqlens=metadata.cache_seqlens_int32,
+                cu_seqlens_q=metadata.cu_seqlens_q,
+                cu_seqlens_k_new=metadata.cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                softmax_scale=layer.scaling,
+                causal=False if use_cascade_attn else causal,
+                softcap=layer.logit_cap,
+                k_descale=k_descale,
+                v_descale=v_descale,
+                return_softmax_lse=use_cascade_attn,
+                num_splits=self.num_splits,
+            )
+            if use_cascade_attn:
+                self._current_prefix = "forward_decode_use_cascade_attn"
+                self._current_max_seqlen_k = (
+                    self.forward_metadata_spec_decode_expand.max_seq_len_k
+                )
+
+                o, softmax_lse, *rest = result
+                o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache(
+                    q=q_rope,
+                    k_cache=k_rope_cache,
+                    v_cache=c_kv_cache,
+                    qv=q_nope,
+                    page_table=self.forward_metadata_spec_decode_expand.page_table,
+                    cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
+                    cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
+                    cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
+                    max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
+                    softmax_scale=layer.scaling,
+                    causal=False,
+                    window_size=window_size,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                    return_softmax_lse=True,
+                    num_splits=self.num_splits,
+                )
+                o, _ = merge_state_v2_wrapper(
+                    o,
+                    softmax_lse.T.contiguous(),
+                    o_expand,
+                    softmax_lse_expand.T.contiguous(),
+                )
+            else:
+                o = result
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+
+class MusaFlashAttentionMultiStepBackend(FlashAttentionMultiStepBackend):
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+        fa_impl_ver: int = 3,
+    ):
+        self.model_runner = model_runner
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.attn_backends = []
+        for i in range(self.speculative_num_steps - 1):
+            self.attn_backends.append(
+                MusaFlashAttentionBackend(
+                    model_runner,
+                    speculative_step_id=i,
+                    topk=self.topk,
+                    speculative_num_steps=self.speculative_num_steps,
+                    fa_impl_ver=fa_impl_ver,
+                )
+            )

python/sglang/srt/hardware_backend/musa/kernels/topk.py

@@ -0,0 +1,300 @@
+from typing import (
+    Optional,
+)
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1, num_stages=1),
+        triton.Config({}, num_warps=1, num_stages=2),
+        triton.Config({}, num_warps=2, num_stages=1),
+        triton.Config({}, num_warps=2, num_stages=2),
+        triton.Config({}, num_warps=4, num_stages=1),
+        triton.Config({}, num_warps=4, num_stages=2),
+        triton.Config({}, num_warps=4, num_stages=3),
+        triton.Config({}, num_warps=8, num_stages=1),
+        triton.Config({}, num_warps=8, num_stages=2),
+        triton.Config({}, num_warps=8, num_stages=3),
+        triton.Config({}, num_warps=16, num_stages=1),
+        triton.Config({}, num_warps=16, num_stages=2),
+        triton.Config({}, num_warps=16, num_stages=3),
+        triton.Config({}, num_warps=32, num_stages=1),
+        triton.Config({}, num_warps=32, num_stages=2),
+    ],
+    key=["num_tokens", "num_experts", "has_correction_bias"],
+)
+@triton.jit
+def topk_softmax_triton_kernel(
+    gating_output_ptr,
+    selected_expert_ptr,
+    moe_weights_ptr,
+    renormalize_flag,
+    num_experts,
+    num_tokens,  # for autotune key
+    moe_softcapping,
+    correction_bias_ptr,
+    has_correction_bias: tl.constexpr,
+    K: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_WIDTH_SIZE_UP: tl.constexpr,
+):
+    curr_row_idx = tl.program_id(0)
+
+    FLOAT_MINIMUM = -10000.0
+    LOG2E = 1.4426950408889634
+
+    weights_local_final = tl.zeros((BLOCK_K,), dtype=tl.float32)
+    selected_local_final = tl.zeros((BLOCK_K,), dtype=tl.int32)
+
+    offset = tl.arange(0, BLOCK_WIDTH_SIZE_UP)
+    k_offset = tl.arange(0, BLOCK_K)
+    mask_expert = offset < num_experts
+    mask_topk = k_offset < K
+
+    row_offset = curr_row_idx * num_experts
+
+    logits = tl.load(
+        gating_output_ptr + row_offset + offset, mask=mask_expert, other=FLOAT_MINIMUM
+    )
+    logits = tl.cast(logits, tl.float32)
+
+    if has_correction_bias:
+        bias = tl.load(correction_bias_ptr + offset, mask=mask_expert, other=0.0)
+        logits = logits + bias
+
+    if moe_softcapping > 0.0:
+        logits = moe_softcapping * tanh(logits / moe_softcapping)
+
+    row_max = tl.max(logits, axis=0)
+    probs = tl.exp2((logits - row_max) * LOG2E)
+    row_sum = tl.sum(probs, axis=0)
+    inv_row_sum = 1.0 / row_sum
+    probs = probs * inv_row_sum
+    probs = tl.where(mask_expert, probs, FLOAT_MINIMUM)
+
+    weights_selected_sum = 0.0
+    for k_idx in range(K):
+        top_k_index = tl.argmax(probs, axis=0)
+        mask = offset == top_k_index
+        top_k_value = tl.sum(tl.where(mask, probs, 0.0))
+
+        weights_local_final = tl.where(
+            k_offset == k_idx, top_k_value, weights_local_final
+        )
+        selected_local_final = tl.where(
+            k_offset == k_idx, top_k_index, selected_local_final
+        )
+        weights_selected_sum += top_k_value
+
+        probs = tl.where(offset == top_k_index, FLOAT_MINIMUM, probs)
+
+    if renormalize_flag:
+        weights_local_final = weights_local_final / weights_selected_sum
+
+    tl.store(
+        moe_weights_ptr + curr_row_idx * K + k_offset,
+        weights_local_final,
+        mask=mask_topk,
+    )
+    tl.store(
+        selected_expert_ptr + curr_row_idx * K + k_offset,
+        selected_local_final,
+        mask=mask_topk,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    moe_softcapping: float = 0.0,
+    correction_bias: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Compute top-k softmax for MoE routing.
+
+    Args:
+        topk_weights: Output tensor for top-k weights [num_tokens, topk]
+        topk_ids: Output tensor for top-k expert indices [num_tokens, topk]
+        gating_output: Gating logits [num_tokens, num_experts]
+        renormalize: Whether to renormalize the top-k weights
+        moe_softcapping: Tanh softcapping value (0.0 to disable)
+        correction_bias: Per-expert bias correction [num_experts], must be float32 if provided
+    """
+
+    num_tokens, num_experts = gating_output.shape
+    topk = topk_weights.shape[-1]
+    has_correction_bias = correction_bias is not None
+
+    block_width_up = triton.next_power_of_2(num_experts)
+    grid = (num_tokens,)
+
+    topk_softmax_triton_kernel[grid](
+        gating_output,
+        topk_ids,
+        topk_weights,
+        renormalize,
+        num_experts,
+        num_tokens,
+        moe_softcapping,
+        correction_bias,
+        has_correction_bias,
+        K=topk,
+        BLOCK_K=triton.next_power_of_2(topk),
+        BLOCK_WIDTH_SIZE_UP=block_width_up,
+    )
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1, num_stages=1),
+        triton.Config({}, num_warps=1, num_stages=2),
+        triton.Config({}, num_warps=2, num_stages=1),
+        triton.Config({}, num_warps=2, num_stages=2),
+        triton.Config({}, num_warps=4, num_stages=1),
+        triton.Config({}, num_warps=4, num_stages=2),
+        triton.Config({}, num_warps=4, num_stages=3),
+        triton.Config({}, num_warps=8, num_stages=1),
+        triton.Config({}, num_warps=8, num_stages=2),
+        triton.Config({}, num_warps=8, num_stages=3),
+        triton.Config({}, num_warps=16, num_stages=1),
+        triton.Config({}, num_warps=16, num_stages=2),
+        triton.Config({}, num_warps=16, num_stages=3),
+        triton.Config({}, num_warps=32, num_stages=1),
+        triton.Config({}, num_warps=32, num_stages=2),
+    ],
+    key=["num_tokens", "num_experts"],
+)
+@triton.jit
+def topk_sigmoid_triton_kernel(
+    gating_output_ptr,
+    selected_expert_ptr,
+    moe_weights_ptr,
+    renormalize_flag,
+    correction_bias_ptr,
+    has_correction_bias: tl.constexpr,
+    num_experts,
+    num_tokens,  # for autotune key
+    K: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_WIDTH_SIZE_UP: tl.constexpr,
+):
+    curr_row_idx = tl.program_id(0)
+
+    FLOAT_MINIMUM = -10000.0
+    LOG2E = 1.4426950408889634
+
+    weights_local_final = tl.zeros((BLOCK_K,), dtype=tl.float32)
+    selected_local_final = tl.zeros((BLOCK_K,), dtype=tl.int32)
+
+    offset = tl.arange(0, BLOCK_WIDTH_SIZE_UP)
+    k_offset = tl.arange(0, BLOCK_K)
+    mask_expert = offset < num_experts
+    mask_topk = k_offset < K
+
+    row_offset = curr_row_idx * num_experts
+
+    x = tl.load(
+        gating_output_ptr + row_offset + offset, mask=mask_expert, other=FLOAT_MINIMUM
+    )
+    x = tl.cast(x, tl.float32)
+
+    # Compute sigmoid(x)
+    is_positive = x >= 0
+    neg_x = tl.where(is_positive, -x, x)
+    exp_neg_x = tl.exp2(neg_x * LOG2E)
+    probs = tl.where(
+        is_positive,
+        1.0 / (1.0 + exp_neg_x),
+        exp_neg_x / (1.0 + exp_neg_x),
+    )
+
+    if has_correction_bias:
+        bias = tl.load(correction_bias_ptr + offset, mask=mask_expert, other=0.0)
+        probs_for_choice = probs + bias
+    else:
+        probs_for_choice = probs
+
+    probs_for_choice = tl.where(mask_expert, probs_for_choice, FLOAT_MINIMUM)
+
+    weights_selected_sum = 0.0
+    for k_idx in range(K):
+        top_k_index = tl.argmax(probs_for_choice, axis=0)
+        mask = offset == top_k_index
+        top_k_value = tl.sum(tl.where(mask, probs, 0.0))
+
+        weights_local_final = tl.where(
+            k_offset == k_idx, top_k_value, weights_local_final
+        )
+        selected_local_final = tl.where(
+            k_offset == k_idx, top_k_index, selected_local_final
+        )
+        weights_selected_sum += top_k_value
+
+        probs_for_choice = tl.where(
+            offset == top_k_index, FLOAT_MINIMUM, probs_for_choice
+        )
+
+    if renormalize_flag:
+        weights_local_final = weights_local_final / weights_selected_sum
+
+    tl.store(
+        moe_weights_ptr + curr_row_idx * K + k_offset,
+        weights_local_final,
+        mask=mask_topk,
+    )
+    tl.store(
+        selected_expert_ptr + curr_row_idx * K + k_offset,
+        selected_local_final,
+        mask=mask_topk,
+    )
+
+
+def topk_sigmoid(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    correction_bias: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Compute top-k sigmoid for MoE routing.
+
+    Args:
+        topk_weights: Output tensor for top-k weights [num_tokens, topk]
+        topk_ids: Output tensor for top-k expert indices [num_tokens, topk]
+        gating_output: Gating logits [num_tokens, num_experts]
+        renormalize: Whether to renormalize the top-k weights
+        correction_bias: Per-expert bias correction [num_experts], must be float32 if provided
+    """
+    num_tokens, num_experts = gating_output.shape
+    topk = topk_weights.shape[-1]
+    has_correction_bias = correction_bias is not None
+
+    block_width_up = triton.next_power_of_2(num_experts)
+    grid = (num_tokens,)
+
+    topk_sigmoid_triton_kernel[grid](
+        gating_output,
+        topk_ids,
+        topk_weights,
+        renormalize,
+        correction_bias,
+        has_correction_bias,
+        num_experts,
+        num_tokens,
+        K=topk,
+        BLOCK_K=triton.next_power_of_2(topk),
+        BLOCK_WIDTH_SIZE_UP=block_width_up,
+    )

python/sglang/srt/hardware_backend/musa/layers/utils/cp_utils.py

@@ -0,0 +1,57 @@
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.musa.attention.flashattention_backend import (
+        MusaFlashAttentionBackend,
+    )
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def musa_cp_attn_forward_extend(
+    musa_fa_backend: "MusaFlashAttentionBackend",
+    forward_batch: "ForwardBatch",
+    q: torch.Tensor,
+    device: torch.device,
+    attn_fn: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int], torch.Tensor],
+) -> torch.Tensor:
+    """
+    Split q into prev/next zigzag halves based on CP metadata, call the
+    backend-specific attention function twice with appropriate per-half
+    metadata, and concatenate the results.
+
+    attn_fn signature:
+        attn_fn(q, cu_seqlens_q, cache_seqlens, max_seqlen_q) -> result
+    where only these four CP-varying parameters differ between halves.
+    All other backend-specific args should be captured in the closure.
+    """
+    cp_meta = forward_batch.attn_cp_metadata
+
+    q_prev, q_next = torch.chunk(q, 2, dim=0)
+
+    cu_seqlens_q_prev = torch.tensor(
+        [0, cp_meta.actual_seq_q_prev], device=device, dtype=torch.int32
+    )
+    if hasattr(musa_fa_backend, "_current_prefix"):
+        musa_fa_backend._current_prefix = "forward_extend_cp_prev"
+    result_prev = attn_fn(
+        q_prev,
+        cu_seqlens_q_prev,
+        cp_meta.kv_len_prev_tensor,
+        cp_meta.actual_seq_q_prev,
+    )
+
+    cu_seqlens_q_next = torch.tensor(
+        [0, cp_meta.actual_seq_q_next], device=device, dtype=torch.int32
+    )
+    if hasattr(musa_fa_backend, "_current_prefix"):
+        musa_fa_backend._current_prefix = "forward_extend_cp_next"
+    result_next = attn_fn(
+        q_next,
+        cu_seqlens_q_next,
+        cp_meta.kv_len_next_tensor,
+        cp_meta.actual_seq_q_next,
+    )
+
+    return torch.concat([result_prev, result_next], dim=0)

python/sglang/srt/hardware_backend/musa/utils/patch_torch.py

@@ -0,0 +1,63 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import re
+from dataclasses import replace as _dataclass_replace
+
+import torch
+import torch.fx.graph as fx_graph
+
+_DEVICE_REPR_RE = re.compile(r"\bdevice\(type='([^']+)'(?:,\s*index=(\d+))?\)")
+
+
+def _replace_device_repr(m: re.Match) -> str:
+    dev_type = m.group(1)
+    dev_index = m.group(2)
+    if dev_index is not None:
+        return f"torch.device('{dev_type}:{dev_index}')"
+    return f"torch.device('{dev_type}')"
+
+
+def patch_fx_custom_device() -> None:
+    """
+    Fix FX codegen serialization for non-standard devices (e.g. torch_musa).
+
+    Root cause:
+    torch.device is registered as a custom builtin named 'device', imported
+    via 'from torch import device'. repr(torch.device('musa', 0)) produces
+    "device(type='musa', index=0)", which is syntactically valid but fails
+    at runtime because torch.device does not recognize 'musa' as a type when
+    invoked through the standard import path.
+
+    Fix:
+    Post-process the generated src string, replacing all occurrences of
+    device(type='x', index=N) with torch.device('x:N'), and ensure 'torch'
+    is present in the graph globals.
+
+    Note:
+    _get_repr is a closure inside _gen_python_code and cannot be patched
+    directly, so we wrap _gen_python_code and rewrite its output instead.
+    """
+    original = fx_graph.CodeGen._gen_python_code
+
+    def patched(self, nodes, root_module, namespace, **kwargs):
+        result = original(self, nodes, root_module, namespace, **kwargs)
+        new_src = _DEVICE_REPR_RE.sub(_replace_device_repr, result.src)
+        if new_src is result.src:
+            return result
+        result.globals.setdefault("torch", torch)
+        if hasattr(result, "_replace"):
+            return result._replace(src=new_src)
+        return _dataclass_replace(result, src=new_src)
+
+    fx_graph.CodeGen._gen_python_code = patched

python/sglang/srt/hardware_backend/npu/allocator_npu.py

@@ -0,0 +1,155 @@
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.mem_cache.allocator import (
+    PagedTokenToKVPoolAllocator,
+    alloc_extend_naive,
+)
+from sglang.srt.utils import get_num_new_pages, next_power_of_2
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import KVCache
+
+
+class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        device: str,
+        kvcache: "KVCache",
+        need_sort: bool,
+    ):
+        super().__init__(size, page_size, dtype, device, kvcache, need_sort)
+        self.roundup = page_size - 1
+
+    def alloc_extend(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+        num_new_pages: int = None,
+    ):
+        if self.debug_mode:
+            assert torch.all(
+                (last_loc + 1) % self.page_size == prefix_lens % self.page_size
+            )
+
+        if num_new_pages is None:
+            num_new_pages_tensor = (
+                (seq_lens + self.roundup) // self.page_size
+                - (prefix_lens + self.roundup) // self.page_size
+            ).sum()
+            num_new_pages_item = num_new_pages_tensor.item()
+        else:
+            num_new_pages_item = num_new_pages
+        if self.need_sort and num_new_pages_item > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        if num_new_pages_item > len(self.free_pages):
+            return None
+
+        if num_new_pages_item < 200:
+            from sgl_kernel_npu.mem_cache.allocator import alloc_extend_kernel
+
+            out_indices = torch.empty(
+                (extend_num_tokens,),
+                dtype=torch.int64,
+                device=self.device,
+            )
+            max_num_extend_tokens = next_power_of_2(extend_num_tokens)
+            bs = prefix_lens.shape[0]
+            alloc_extend_kernel[(bs,)](
+                prefix_lens,
+                seq_lens,
+                last_loc,
+                self.free_pages,
+                out_indices,
+                next_power_of_2(bs),
+                self.page_size,
+                max_num_extend_tokens,
+            )
+
+        else:
+            out_indices = torch.empty(
+                (extend_num_tokens,),
+                dtype=torch.int32,
+                device=self.device,
+            )
+            alloc_extend_naive(
+                prefix_lens,
+                seq_lens,
+                last_loc,
+                self.free_pages,
+                out_indices,
+                self.page_size,
+                self.device,
+            )
+
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
+        self.free_pages = self.free_pages[num_new_pages_item:]
+        return out_indices.int()
+
+    def alloc_decode(
+        self,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+    ):
+        if self.debug_mode:
+            assert torch.all(
+                (last_loc + 2) % self.page_size == seq_lens % self.page_size
+            )
+
+        num_new_pages = get_num_new_pages(
+            seq_lens=seq_lens_cpu,
+            page_size=self.page_size,
+            decode=True,
+        )
+
+        if num_new_pages > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        if num_new_pages > len(self.free_pages):
+            return None
+
+        need_new_pages = (seq_lens % self.page_size == 1).int()
+        end_new_pages = torch.cumsum(need_new_pages, 0)
+        start_new_pages = end_new_pages - need_new_pages
+        if num_new_pages == 0:
+            out_indices = last_loc + 1
+        else:
+            out_indices = (last_loc + 1) * (1 - need_new_pages) + self.free_pages[
+                start_new_pages
+            ] * self.page_size * need_new_pages
+
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
+        self.free_pages = self.free_pages[num_new_pages:]
+        return out_indices.int()
+
+    def free(self, free_index: torch.Tensor):
+        if free_index.numel() == 0:
+            return
+
+        if self.is_not_in_free_group:
+            device = free_index.device
+            free_page_indices = torch.unique(free_index.cpu() // self.page_size)
+            free_page_indices = free_page_indices.to(device)
+            if self.need_sort:
+                self.release_pages = torch.cat((free_page_indices, self.release_pages))
+            else:
+                self.free_pages = torch.cat((free_page_indices, self.free_pages))
+        else:
+            self.free_group.append(free_index)
+
+        if self.debug_mode:
+            assert len(torch.unique(self.free_pages)) == len(self.free_pages)

python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py

@@ -0,0 +1,391 @@
+from typing import Optional, Tuple, Union
+
+import torch
+from sgl_kernel_npu.fla.fused_gdn_gating import (
+    fused_gdn_gating_kernel_without_sigmoid,
+    fused_gdn_gating_npu,
+)
+from sgl_kernel_npu.mamba.causal_conv1d import (
+    causal_conv1d_fn_npu,
+    causal_conv1d_update_npu,
+    causal_conv1d_update_v2,
+)
+
+from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import (
+    AscendMambaAttnBackendBase,
+)
+from sglang.srt.layers.attention.linear.gdn_backend import GDNKernelDispatcher
+from sglang.srt.layers.attention.linear.utils import (
+    get_linear_attn_decode_backend,
+    get_linear_attn_prefill_backend,
+)
+from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
+from sglang.srt.mem_cache.memory_pool import MambaPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+
+fused_gdn_gating = fused_gdn_gating_npu
+causal_conv1d_fn = causal_conv1d_fn_npu
+causal_conv1d_update = causal_conv1d_update_npu
+
+
+class AscendGDNAttnBackend(AscendMambaAttnBackendBase):
+
+    def __init__(self, model_runner: ModelRunner):
+        super().__init__(model_runner)
+        self.conv_states_shape = torch.Size(
+            (
+                *model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[
+                    :-2
+                ],
+                model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[-1],
+                model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[-2],
+            )
+        )
+        decode_backend = get_linear_attn_decode_backend()
+        prefill_backend = get_linear_attn_prefill_backend()
+        self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend)
+
+    def _prepare_mamba_track_metadata(self, forward_batch: ForwardBatch):
+        if self.forward_metadata.has_mamba_track_mask:
+            mamba_track_mask_indices = forward_batch.mamba_track_mask.nonzero(
+                as_tuple=True
+            )[0]
+            self.forward_metadata.conv_states_mask_indices = (
+                forward_batch.mamba_track_indices[mamba_track_mask_indices]
+            )
+
+    def prepare_gdn_inputs(
+        self,
+        bs: int,
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        cache_indices = self.forward_metadata.mamba_cache_indices
+        self.num_accept_tokens = torch.ones(
+            [bs], dtype=torch.int32, device=cache_indices.device
+        )
+        self.actual_seq_lengths = torch.ones(
+            [bs], dtype=torch.int32, device=cache_indices.device
+        )
+        if forward_mode.is_target_verify():
+            seq_len = spec_info.draft_token_num
+            self.actual_seq_lengths = self.actual_seq_lengths * seq_len
+            # indices
+            self.ssm_state_indices = torch.arange(
+                cache_indices.shape[0] * seq_len,
+                dtype=torch.int32,
+                device=cache_indices.device,
+            )
+        else:
+            self.ssm_state_indices = cache_indices
+
+    def init_forward_metadata_out_graph(
+        self,
+        forward_batch: ForwardBatch,
+        in_capture: bool = False,
+    ):
+        if forward_batch.forward_mode.is_draft_extend_v2():
+            return
+        super().init_forward_metadata_out_graph(forward_batch, in_capture=in_capture)
+        self.prepare_gdn_inputs(
+            forward_batch.batch_size,
+            forward_batch.forward_mode,
+            forward_batch.spec_info,
+        )
+        self._prepare_mamba_track_metadata(forward_batch)
+        self.graph_mode = True
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        if forward_batch.forward_mode.is_draft_extend_v2():
+            return
+        super().init_forward_metadata(forward_batch)
+        self.prepare_gdn_inputs(
+            forward_batch.batch_size,
+            forward_batch.forward_mode,
+            forward_batch.spec_info,
+        )
+        self._prepare_mamba_track_metadata(forward_batch)
+        self.graph_mode = False
+
+    def forward_decode(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,
+    ):
+        layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        conv_states = layer_cache.conv[0]
+        ssm_states = layer_cache.temporal
+        query_start_loc = self.forward_metadata.query_start_loc
+        cache_indices = self.forward_metadata.mamba_cache_indices
+
+        assert isinstance(mixed_qkv, torch.Tensor)
+        conv_states_tmp = conv_states.transpose(1, 2).clone()
+        mixed_qkv = causal_conv1d_update(
+            mixed_qkv,
+            conv_states_tmp,
+            layer.conv_weights,
+            layer.bias,
+            layer.activation,
+            conv_state_indices=cache_indices,
+        )
+        conv_states[:] = conv_states_tmp.transpose(1, 2)
+
+        query, key, value = torch.split(
+            mixed_qkv,
+            [layer.q_dim, layer.k_dim, layer.v_dim],
+            dim=-1,
+        )
+        bs = forward_batch.batch_size
+        query = query.view(1, bs, layer.num_q_heads, layer.head_q_dim)
+        key = key.view(1, bs, layer.num_k_heads, layer.head_k_dim)
+        value = value.view(1, bs, layer.num_v_heads, layer.head_v_dim)
+
+        core_attn_out = self.kernel_dispatcher.decode(
+            q=query,
+            k=key,
+            v=value,
+            a=a,
+            b=b,
+            A_log=layer.A_log,
+            dt_bias=layer.dt_bias,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+        )
+
+        self._track_mamba_state_decode(
+            forward_batch, conv_states, ssm_states, cache_indices
+        )
+        return core_attn_out
+
+    def forward_extend(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,
+    ):
+        assert isinstance(mixed_qkv, torch.Tensor)
+        seq_len = mixed_qkv.shape[0]
+        is_target_verify = forward_batch.forward_mode.is_target_verify()
+        forward_metadata = self.forward_metadata
+
+        query_start_loc = forward_metadata.query_start_loc
+        cache_indices = forward_metadata.mamba_cache_indices
+        retrieve_next_token = forward_metadata.retrieve_next_token
+        retrieve_next_sibling = forward_metadata.retrieve_next_sibling
+        retrieve_parent_token = forward_metadata.retrieve_parent_token
+
+        mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        conv_states = mamba_cache_params.conv[0]
+        ssm_states = mamba_cache_params.temporal
+        if is_target_verify:
+            assert isinstance(mamba_cache_params, MambaPool.SpeculativeState)
+            intermediate_state_cache = mamba_cache_params.intermediate_ssm
+            intermediate_conv_window_cache = (
+                mamba_cache_params.intermediate_conv_window[0]
+            )
+            has_initial_states = torch.ones(
+                seq_len // forward_batch.spec_info.draft_token_num,
+                dtype=torch.bool,
+                device=forward_batch.input_ids.device,
+            )
+        else:
+            has_initial_states = forward_batch.extend_prefix_lens > 0
+        if is_target_verify:
+            num_token_padding = mixed_qkv.shape[0]
+            if (
+                not self.graph_mode
+                and forward_batch.num_token_non_padded_cpu != num_token_padding
+            ):
+                mixed_qkv = mixed_qkv[: forward_batch.num_token_non_padded_cpu]
+                a = a[: forward_batch.num_token_non_padded_cpu]
+                b = b[: forward_batch.num_token_non_padded_cpu]
+                seq_len = forward_batch.num_token_non_padded_cpu
+
+            batch_size = cache_indices.shape[0]
+            draft_token_num = forward_batch.spec_info.draft_token_num
+            num_accepted_tokens = torch.full(
+                (batch_size,),
+                draft_token_num,
+                dtype=torch.int32,
+                device=mixed_qkv.device,
+            )
+            mixed_qkv = causal_conv1d_update_v2(
+                x=mixed_qkv.view(batch_size, draft_token_num, -1).contiguous(),
+                conv_state=conv_states.contiguous(),
+                weight=layer.conv_weights.transpose(0, 1).contiguous(),
+                bias=layer.bias,
+                activation=layer.activation,
+                conv_state_indices=cache_indices,
+                num_accepted_tokens=num_accepted_tokens,
+                pad_slot_id=-1,
+                validate_data=False,
+            ).view(seq_len, -1)
+        else:
+            mixed_qkv = mixed_qkv.transpose(0, 1)
+            if forward_metadata.has_mamba_track_mask:
+                mixed_qkv_to_track = mixed_qkv[
+                    :, forward_metadata.track_conv_indices
+                ].transpose(0, 1)
+                conv_states.transpose(1, 2)[
+                    forward_metadata.conv_states_mask_indices
+                ] = mixed_qkv_to_track
+            kernel_size = layer.conv_weights.shape[-1]
+            conv_states_for_prefill = conv_states[:, -(kernel_size - 1) :, :]
+            conv_states_tmp = conv_states_for_prefill.transpose(1, 2).contiguous()
+
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                layer.conv_weights,
+                layer.bias,
+                activation=layer.activation,
+                conv_states=conv_states_tmp,
+                has_initial_state=has_initial_states,
+                cache_indices=cache_indices,
+                query_start_loc=query_start_loc,
+                seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+            ).transpose(0, 1)[:seq_len]
+            conv_states[:, -(kernel_size - 1) :, :] = conv_states_tmp.transpose(
+                1, 2
+            ).contiguous()
+        if is_target_verify:
+            g, beta = fused_gdn_gating_kernel_without_sigmoid(
+                layer.A_log, a, b, layer.dt_bias
+            )
+            beta = beta.unsqueeze(0)
+            num_heads, head_k_dim = layer.num_q_heads, layer.head_q_dim
+            num_value_heads, head_v_dim = layer.num_v_heads, layer.head_v_dim
+
+            mixed_qkv_last_dim = mixed_qkv.shape[-1]
+
+            mixed_qkv = mixed_qkv.view(batch_size, -1, mixed_qkv_last_dim)
+            beta = beta.view(batch_size, -1, num_value_heads)
+            g = g.view(batch_size, -1, num_value_heads)
+
+            core_attn_out = self.fused_recurrent_gated_delta_rule_update(
+                mixed_qkv,
+                num_heads,
+                num_value_heads,
+                head_k_dim,
+                head_v_dim,
+                recurrent_state=ssm_states,
+                beta=beta,
+                g=g,
+                cache_indices=cache_indices,
+                intermediate_state=intermediate_state_cache,
+            )
+            core_attn_out = core_attn_out.view(-1, num_value_heads, head_v_dim)
+            if (not self.graph_mode) and core_attn_out.shape[0] < num_token_padding:
+                core_attn_out = torch.cat(
+                    [
+                        core_attn_out,
+                        core_attn_out.new_zeros(
+                            num_token_padding - core_attn_out.shape[0],
+                            *core_attn_out.shape[1:],
+                        ),
+                    ],
+                    dim=0,
+                )
+        else:
+            query, key, value = torch.split(
+                mixed_qkv,
+                [layer.q_dim, layer.k_dim, layer.v_dim],
+                dim=-1,
+            )
+
+            actual_seq_len = query.shape[0]
+            query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim)
+            key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim)
+            value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim)
+
+            g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
+            core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
+                q=query,
+                k=key,
+                v=value,
+                g=g,
+                beta=beta,
+                ssm_states=ssm_states,
+                cache_indices=cache_indices,
+                query_start_loc=query_start_loc,
+            )
+            if last_recurrent_state is not None:
+                last_recurrent_state = last_recurrent_state.to(
+                    ssm_states.dtype, copy=False
+                )
+                ssm_states[cache_indices] = last_recurrent_state
+            last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False)
+            ssm_states[cache_indices] = last_recurrent_state
+            if h is not None:
+                self._track_mamba_state_extend(
+                    forward_batch, h, ssm_states, forward_metadata
+                )
+
+        return core_attn_out
+
+    def fused_recurrent_gated_delta_rule_update(
+        self,
+        mix_qkv: torch.Tensor,
+        num_heads,
+        num_value_heads,
+        head_k_dim,
+        head_v_dim,
+        recurrent_state: torch.Tensor,
+        beta: torch.Tensor,
+        g: torch.Tensor,
+        cache_indices: torch.Tensor,
+        intermediate_state: Optional[torch.Tensor] = None,
+    ):
+        beta = beta.to(torch.bfloat16)
+        g = g.to(torch.float32)
+        batch_size = mix_qkv.shape[0]
+        seq_len = mix_qkv.shape[1]
+        scale = 1 / (head_k_dim**0.5)
+
+        if intermediate_state is not None:
+            intermediate_state = intermediate_state.view(
+                -1, num_value_heads, head_k_dim, head_v_dim
+            )
+
+        if self.graph_mode:
+            num_accept_tokens = torch.full(
+                [batch_size], 1, dtype=torch.int32, device=cache_indices.device
+            )
+            actual_seq_lengths = torch.full(
+                [batch_size], seq_len, dtype=torch.int32, device=cache_indices.device
+            )
+            ssm_state_indices = self.forward_metadata.mamba_cache_indices_gdn
+        else:
+            num_accept_tokens = self.num_accept_tokens
+            actual_seq_lengths = self.actual_seq_lengths
+            ssm_state_indices = self.ssm_state_indices
+
+        attn_core_out = torch.ops.npu.recurrent_gated_delta_rule(
+            mix_qkv,
+            recurrent_state,
+            beta=beta,
+            scale=scale,
+            actual_seq_lengths=actual_seq_lengths,
+            ssm_state_indices=ssm_state_indices.view(batch_size, seq_len),
+            nk=num_heads,
+            nv=num_value_heads,
+            intermediate_state=intermediate_state,
+            cache_indices=cache_indices,
+            num_accepted_tokens=num_accept_tokens,
+            g=g,
+        )
+
+        if intermediate_state is not None:
+            intermediate_state = intermediate_state.view(
+                -1, seq_len, num_value_heads, head_k_dim, head_v_dim
+            )
+        return attn_core_out

python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py

@@ -0,0 +1,320 @@
+import logging
+from typing import Optional, Union
+
+import torch
+from sgl_kernel_npu.mamba.mamba_state_update_triton import (
+    conv_state_rollback,
+    move_intermediate_cache,
+)
+
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    HybridLinearAttnBackend,
+    MambaAttnBackendBase,
+)
+from sglang.srt.layers.attention.mamba.mamba2_metadata import (
+    ForwardMetadata,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+from sglang.srt.speculative.spec_info import SpecInput
+
+logger = logging.getLogger(__name__)
+
+
+class AscendMambaAttnBackendBase(MambaAttnBackendBase):
+    def __init__(self, model_runner: ModelRunner):
+        super().__init__(model_runner)
+        self.state_indices_list_gdn = []
+
+    def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
+        assert (
+            max_num_tokens % max_bs == 0
+        ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}"
+        draft_token_num = max_num_tokens // max_bs
+        for i in range(max_bs):
+            self.state_indices_list.append(
+                torch.full(
+                    (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device
+                )
+            )
+            self.state_indices_list_gdn.append(
+                torch.full(
+                    ((i + 1) * draft_token_num,),
+                    self.pad_slot_id,
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+            )
+            self.query_start_loc_list.append(
+                torch.zeros((i + 2,), dtype=torch.int32, device=self.device)
+            )
+            self.retrieve_next_token_list.append(
+                torch.zeros(
+                    (i + 1, draft_token_num), dtype=torch.int32, device=self.device
+                )
+            )
+            self.retrieve_next_sibling_list.append(
+                torch.zeros(
+                    (i + 1, draft_token_num), dtype=torch.int32, device=self.device
+                )
+            )
+            self.retrieve_parent_token_list.append(
+                torch.zeros(
+                    (i + 1, draft_token_num), dtype=torch.int32, device=self.device
+                )
+            )
+        self.cached_cuda_graph_decode_query_start_loc = torch.arange(
+            0, max_bs + 1, dtype=torch.int32, device=self.device
+        )
+        self.cached_cuda_graph_verify_query_start_loc = torch.arange(
+            0,
+            max_bs * draft_token_num + 1,
+            step=draft_token_num,
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+    def _capture_metadata(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
+        self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
+        if forward_mode.is_decode_or_idle():
+            self.query_start_loc_list[bs - 1].copy_(
+                self.cached_cuda_graph_decode_query_start_loc[: bs + 1]
+            )
+        elif forward_mode.is_target_verify():
+            self.query_start_loc_list[bs - 1].copy_(
+                self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
+            )
+            ssm_state_indices = torch.arange(
+                mamba_indices.shape[0] * spec_info.draft_token_num,
+                dtype=torch.int32,
+                device=mamba_indices.device,
+            )
+            self.state_indices_list_gdn[bs - 1][
+                : len(mamba_indices) * spec_info.draft_token_num
+            ].copy_(ssm_state_indices)
+        else:
+            raise ValueError(f"Invalid forward mode: {forward_mode=}")
+
+        # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
+        if forward_mode.is_target_verify() and spec_info.topk > 1:
+            # They are None during cuda graph capture so skip the copy_...
+            # self.retrieve_next_token_list[bs - 1].copy_(spec_info.retrive_next_token)
+            # self.retrieve_next_sibling_list[bs - 1].copy_(spec_info.retrive_next_sibling)
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                retrieve_next_token=self.retrieve_next_token_list[bs - 1],
+                retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1],
+                retrieve_parent_token=self.retrieve_parent_token_list[bs - 1],
+            )
+        else:
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1],
+            )
+
+    def _replay_metadata(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInput],
+        seq_lens_cpu: Optional[torch.Tensor],
+        num_padding: Optional[int] = None,
+        in_capture: bool = False,
+        mamba_track_indices: Optional[torch.Tensor] = None,
+        *args,
+        **kwargs,
+    ):
+        # out_graph passes seq_lens_cpu=None at capture; mirror the base guard.
+        if seq_lens_cpu is None:
+            num_padding = 0
+        else:
+            num_padding = torch.count_nonzero(
+                seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value()
+            )
+        # Make sure forward metadata is correctly handled for padding reqs
+        req_pool_indices[bs - num_padding :] = 0
+        mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
+        mamba_indices[bs - num_padding :] = 0
+        self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
+        track_buf = None
+        if mamba_track_indices is not None:
+            track_buf = mamba_track_indices
+        if forward_mode.is_decode_or_idle():
+            if num_padding == 0:
+                self.query_start_loc_list[bs - 1].copy_(
+                    self.cached_cuda_graph_decode_query_start_loc[: bs + 1]
+                )
+            else:
+                self.query_start_loc_list[bs - 1][: bs - num_padding].copy_(
+                    self.cached_cuda_graph_decode_query_start_loc[: bs - num_padding]
+                )
+                self.query_start_loc_list[bs - 1][bs - num_padding :].fill_(
+                    bs - num_padding
+                )
+        elif forward_mode.is_target_verify():
+            ssm_state_indices = torch.arange(
+                bs * spec_info.draft_token_num,
+                dtype=torch.int32,
+                device=mamba_indices.device,
+            )
+            self.state_indices_list_gdn[bs - 1].copy_(ssm_state_indices)
+            if num_padding == 0:
+                self.query_start_loc_list[bs - 1].copy_(
+                    self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
+                )
+            else:
+                self.query_start_loc_list[bs - 1][: bs - num_padding].copy_(
+                    self.cached_cuda_graph_verify_query_start_loc[: bs - num_padding]
+                )
+                self.query_start_loc_list[bs - 1][bs - num_padding :].fill_(
+                    (bs - num_padding) * spec_info.draft_token_num
+                )
+        else:
+            raise ValueError(f"Invalid forward mode: {forward_mode=}")
+
+        # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
+        if forward_mode.is_target_verify() and spec_info.topk > 1:
+            bs_without_pad = spec_info.retrive_next_token.shape[0]
+            self.retrieve_next_token_list[bs - 1][:bs_without_pad].copy_(
+                spec_info.retrive_next_token
+            )
+            self.retrieve_next_sibling_list[bs - 1][:bs_without_pad].copy_(
+                spec_info.retrive_next_sibling
+            )
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_track_indices=track_buf,
+                retrieve_next_token=self.retrieve_next_token_list[bs - 1],
+                retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1],
+                retrieve_parent_token=self.retrieve_parent_token_list[bs - 1],
+            )
+        else:
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1],
+                mamba_track_indices=track_buf,
+            )
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 0  # Mamba attn does not use seq lens to index kv cache
+
+
+class AscendMamba2AttnBackend(AscendMambaAttnBackendBase):
+    pass
+
+
+class AscendHybridLinearAttnBackend(HybridLinearAttnBackend):
+    def __init__(
+        self,
+        full_attn_backend: AttentionBackend,
+        linear_attn_backend: AscendMambaAttnBackendBase,
+        full_attn_layers: list[int],
+    ):
+        super().__init__(full_attn_backend, linear_attn_backend, full_attn_layers)
+
+    def update_mamba_state_after_mtp_verify(
+        self,
+        last_correct_step_indices: torch.Tensor,
+        mamba_track_indices: Optional[torch.Tensor],
+        mamba_steps_to_track: Optional[torch.Tensor],
+        model,
+    ):
+        """
+        Update mamba states after MTP verify using fully fused Triton kernel.
+
+        This replaces the original advanced indexing operations with a single fused
+        gather-scatter kernel that also handles masking internally, avoiding:
+        - index_elementwise_kernel from tensor[bool_mask]
+        - index_select kernel launches
+        - nonzero kernel launches
+        """
+        request_number = last_correct_step_indices.shape[0]
+
+        state_indices_tensor = (
+            self.linear_attn_backend.forward_metadata.mamba_cache_indices[
+                :request_number
+            ]
+        )
+
+        mamba_caches = (
+            self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers()
+        )
+
+        conv_states = mamba_caches.conv[0]
+        ssm_states = mamba_caches.temporal
+        intermediate_state_cache = mamba_caches.intermediate_ssm
+        dst_indices_tensor = state_indices_tensor.to(torch.int64)  # [N]
+        src_indices_tensor = torch.arange(
+            dst_indices_tensor.shape[0],
+            device=dst_indices_tensor.device,
+            dtype=torch.int64,
+        )
+        last_steps = last_correct_step_indices.to(torch.int64)  # [N]
+
+        move_intermediate_cache(
+            ssm_states,
+            intermediate_state_cache,
+            dst_indices_tensor,
+            src_indices_tensor,
+            last_steps,
+        )
+
+        draft_token_num = intermediate_state_cache.shape[2]
+        if mamba_track_indices is not None:
+            assert mamba_steps_to_track is not None
+            mamba_track_indices = mamba_track_indices.to(torch.int64)
+            mamba_steps_to_track = mamba_steps_to_track.to(torch.int64)
+
+            move_intermediate_cache(
+                ssm_states,
+                intermediate_state_cache,
+                mamba_track_indices,
+                src_indices_tensor,
+                mamba_steps_to_track,
+            )
+
+            track_mask = mamba_steps_to_track >= 0
+            # Track conv state from the verify-time window before rolling back
+            # the working slot; NPU does not keep per-step conv intermediates.
+            track_indices = mamba_track_indices[track_mask]
+            if track_indices.numel() > 0:
+                conv_states[:, track_indices] = conv_states[
+                    :, dst_indices_tensor[track_mask]
+                ]
+
+        if dst_indices_tensor.numel() > 0:
+            conv_state_rollback(
+                conv_states,
+                dst_indices_tensor,
+                last_steps,
+                draft_token_num,
+            )
+
+        if mamba_track_indices is not None and mamba_track_indices.numel() > 0:
+            conv_state_rollback(
+                conv_states,
+                mamba_track_indices,
+                mamba_steps_to_track,
+                draft_token_num,
+            )
+
+        return
+
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
+    ):
+        pass

python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+from torch.nn.functional import scaled_dot_product_attention
+
+
+class AscendTorchNativeAttnBackend:
+    def __init__(self):
+        pass
+
+    def scaled_dot_product_attention_with_softcapping(
+        self,
+        query,
+        key,
+        value,
+        attn_mask=None,
+        is_causal=False,
+        scale=None,
+        enable_gqa=False,
+        logit_cap=0.0,
+        logit_capping_method="tanh",
+    ) -> torch.Tensor:
+        L, S = query.size(-2), key.size(-2)
+        scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+        attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+        if is_causal:
+            assert attn_mask is None
+            temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(
+                diagonal=0
+            )
+            attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+            attn_bias.to(query.dtype)
+
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+            else:
+                attn_bias = attn_mask + attn_bias
+
+        if enable_gqa:
+            key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
+            value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
+
+        attn_weight = query @ key.transpose(-2, -1) * scale_factor
+
+        if logit_cap > 0:
+            if logit_capping_method == "tanh":
+                attn_weight = logit_cap * torch.tanh(attn_weight / logit_cap)
+
+        attn_weight += attn_bias
+        attn_weight = torch.softmax(attn_weight, dim=-1)
+        return attn_weight @ value
+
+    def run_sdpa_forward_extend(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_prefix_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        encoder_lens: torch.Tensor = None,
+        is_cross_attention: bool = False,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+        sliding_window_size: int = -1,
+        full_to_swa_mapping: Optional[torch.Tensor] = None,
+        logit_cap: float = 0.0,
+        logit_capping_method: str = "tanh",
+    ):
+        """Run the extend forward by using torch native sdpa op.
+
+        Args:
+            query: [num_tokens, num_heads, head_size]
+            output: [num_tokens, num_heads, head_size]
+            k_cache: [max_total_num_tokens, num_heads, head_size]
+            v_cache: [max_total_num_tokens, num_heads, head_size]
+            req_to_token: [max_num_reqs, max_context_len]
+            req_pool_indices: [num_seqs]
+            seq_lens: [num_seqs]
+            extend_prefix_lens: [num_seqs]
+            extend_seq_lens: [num_seqs]
+            encoder_lens: [num_seqs]
+            is_cross_attention: [bool]
+            scaling: float or None
+            enable_gqa: bool
+            causal: bool
+            sliding_window_size: int, -1 means no sliding window
+            full_to_swa_mapping: mapping from full pool index to SWA pool index,
+                required for SWA layers to translate req_to_token indices
+
+        Returns:
+            output: [num_tokens, num_heads, head_size]
+        """
+
+        assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
+        assert seq_lens.shape[0] == extend_seq_lens.shape[0]
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # Need optimize the performance later.
+
+            extend_seq_len_q = int(extend_seq_lens[seq_idx].item())
+            prefill_seq_len_q = int(extend_prefix_lens[seq_idx].item())
+
+            seq_len_kv = int(seq_lens[seq_idx].item())
+            end_q = start_q + extend_seq_len_q
+            end_kv = start_kv + seq_len_kv
+            atten_start_kv = 0
+            atten_end_kv = seq_len_kv
+            # support cross attention
+            if encoder_lens is not None:
+                if is_cross_attention:
+                    atten_end_kv = int(encoder_lens[seq_idx].item())
+                else:
+                    atten_start_kv = int(encoder_lens[seq_idx].item())
+                    atten_end_kv = atten_start_kv + extend_seq_len_q
+
+            is_swa_self_attn = (
+                sliding_window_size is not None
+                and sliding_window_size > -1
+                and encoder_lens is None
+            )
+            if is_swa_self_attn:
+                # For extend, the sliding window must be anchored at the first
+                # query token in this chunk rather than the final sequence
+                # length. Otherwise a large extend chunk can no longer fit in
+                # the cropped query suffix, which breaks the native fallback.
+                atten_start_kv = max(
+                    prefill_seq_len_q - sliding_window_size, atten_start_kv
+                )
+
+            per_req_query = query[:, start_q:end_q, :]
+
+            # SWA crops the front of the KV window, so the redundant query
+            # tensor must match the cropped window to keep Q.len == K.len for
+            # the causal mask. In cross-attention (and non-SWA self-attention)
+            # the original sizing — text seq len — must be preserved, since Q
+            # (text) and KV (encoder) lengths legitimately differ there.
+            if is_swa_self_attn:
+                redundant_len = atten_end_kv - atten_start_kv
+                query_start_idx = max(prefill_seq_len_q - atten_start_kv, 0)
+            else:
+                redundant_len = int(seq_lens[seq_idx].item())
+                query_start_idx = prefill_seq_len_q
+
+            per_req_query_redundant = torch.zeros(
+                (per_req_query.shape[0], redundant_len, per_req_query.shape[2]),
+                dtype=per_req_query.dtype,
+                device=per_req_query.device,
+            )
+
+            per_req_query_redundant[:, query_start_idx:, :] = per_req_query
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, atten_start_kv:atten_end_kv]
+            # For SWA layers, k_cache/v_cache are from the SWA pool but
+            # req_to_token stores full pool indices. Translate before indexing.
+            if full_to_swa_mapping is not None:
+                per_req_tokens = full_to_swa_mapping[per_req_tokens]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype):
+                # scaled_dot_product_attention() expects query, key, and value to have the same dtype
+                per_req_key = per_req_key.to(per_req_query.dtype)
+                per_req_value = per_req_value.to(per_req_query.dtype)
+
+            if logit_cap > 0:
+                per_req_out_redundant = (
+                    self.scaled_dot_product_attention_with_softcapping(
+                        per_req_query_redundant.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                        logit_cap=logit_cap,
+                        logit_capping_method=logit_capping_method,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            else:
+                per_req_out_redundant = (
+                    scaled_dot_product_attention(
+                        per_req_query_redundant.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            output[start_q:end_q, :, :] = per_req_out_redundant[
+                query_start_idx : query_start_idx + extend_seq_len_q, :, :
+            ]
+            start_q, start_kv = end_q, end_kv
+        return output
+
+    def run_sdpa_forward_decode(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: torch.Tensor = None,
+        is_cross_attention: bool = False,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+        sliding_window_size: int = -1,
+        full_to_swa_mapping: Optional[torch.Tensor] = None,
+        logit_cap: float = 0.0,
+        logit_capping_method: str = "tanh",
+    ):
+        """Run the decode forward by using torch native sdpa op.
+
+        Args:
+            query: [num_tokens, num_heads, head_size]
+            output: [num_tokens, num_heads, head_size]
+            k_cache: [max_total_num_tokens, num_heads, head_size]
+            v_cache: [max_total_num_tokens, num_heads, head_size]
+            req_to_token: [max_num_reqs, max_context_len]
+            req_pool_indices: [num_seqs]
+            seq_lens: [num_seqs]
+            encoder_lens: [num_seqs]
+            is_cross_attention: [bool]
+            scaling: float or None
+            enable_gqa: bool
+            causal: bool
+
+        Returns:
+            output: [num_tokens, num_heads, head_size]
+        """
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # Need optimize the performance later.
+
+            seq_len_q = 1
+            seq_len_kv = int(seq_lens[seq_idx].item())
+            end_q = start_q + seq_len_q
+            end_kv = start_kv + seq_len_kv
+            atten_start_kv = 0
+            atten_end_kv = seq_len_kv
+            # support cross attention
+            if encoder_lens is not None:
+                if is_cross_attention:
+                    atten_end_kv = int(encoder_lens[seq_idx].item())
+                else:
+                    atten_start_kv = int(encoder_lens[seq_idx].item())
+                    atten_end_kv = atten_start_kv + seq_len_kv
+
+            if (
+                sliding_window_size is not None
+                and sliding_window_size > -1
+                and encoder_lens is None
+            ):
+                atten_start_kv = max(
+                    atten_end_kv - (sliding_window_size + 1), atten_start_kv
+                )
+
+            per_req_query = query[:, start_q:end_q, :]
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, atten_start_kv:atten_end_kv]
+            # For SWA layers, k_cache/v_cache are from the SWA pool but
+            # req_to_token stores full pool indices. Translate before indexing.
+            if full_to_swa_mapping is not None:
+                per_req_tokens = full_to_swa_mapping[per_req_tokens]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype):
+                # scaled_dot_product_attention() expects query, key, and value to have the same dtype
+                per_req_key = per_req_key.to(per_req_query.dtype)
+                per_req_value = per_req_value.to(per_req_query.dtype)
+
+            if logit_cap > 0:
+                per_req_out = (
+                    self.scaled_dot_product_attention_with_softcapping(
+                        per_req_query.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                        logit_cap=logit_cap,
+                        logit_capping_method=logit_capping_method,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            else:
+                per_req_out = (
+                    scaled_dot_product_attention(
+                        per_req_query.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            output[start_q:end_q, :, :] = per_req_out
+            start_q, start_kv = end_q, end_kv
+
+        return output
+
+    def support_triton(self):
+        return False

python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py

@@ -0,0 +1,493 @@
+import re
+from functools import lru_cache
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.hardware_backend.npu.utils import npu_format_cast
+from sglang.srt.model_executor.forward_context import (
+    get_attn_backend,
+    get_token_to_kv_pool,
+)
+from sglang.srt.utils import get_bool_env_var
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+
+@lru_cache(maxsize=1)
+def is_mla_preprocess_enabled() -> bool:
+    return get_bool_env_var("SGLANG_NPU_USE_MLAPO")
+
+
+@lru_cache(maxsize=1)
+def is_fia_nz() -> bool:
+    is_fia_nz_ = get_bool_env_var("SGLANG_USE_FIA_NZ")
+    if is_fia_nz_:
+        assert (
+            is_mla_preprocess_enabled()
+        ), "SGLANG_USE_FIA_NZ must be enable with SGLANG_NPU_USE_MLAPO"
+    return is_fia_nz_
+
+
+def round_up(val: int, align: int) -> int:
+    if align == 0:
+        return 0
+    return -(val // -align) * align
+
+
+def transdata(nd_mat, block_size: tuple = (16, 16)):
+    r = round_up(nd_mat.shape[0], block_size[0])
+    c = round_up(nd_mat.shape[1], block_size[1])
+    r_pad = r - nd_mat.shape[0]
+    c_pad = c - nd_mat.shape[1]
+    nd_mat = F.pad(nd_mat, ((0, r_pad, 0, c_pad)))
+    nz_mat = torch.permute(
+        torch.reshape(
+            nd_mat,
+            (r // block_size[0], block_size[0], c // block_size[1], block_size[1]),
+        ),
+        [2, 0, 1, 3],
+    )
+    nz_mat = torch.reshape(
+        nz_mat, (nz_mat.shape[0], nz_mat.shape[1] * nz_mat.shape[2], nz_mat.shape[3])
+    )
+    return nz_mat
+
+
+def trans_rope_weight(weight, rope_dim):
+    weight_1 = weight[..., -rope_dim::2, :].contiguous()
+    weight_2 = weight[..., -rope_dim + 1 :: 2, :].contiguous()
+    weight[..., -rope_dim:, :] = torch.cat([weight_1, weight_2], dim=-2)
+
+    return weight.contiguous()
+
+
+class NPUFusedMLAPreprocess(torch.nn.Module):
+    def __init__(
+        self,
+        fused_qkv_a_proj_with_mqa,
+        q_a_layernorm,
+        kv_a_layernorm,
+        q_b_proj,
+        w_kc,
+        rotary_emb,
+        layer_id,
+        num_local_heads,
+        qk_nope_head_dim,
+        qk_rope_head_dim,
+        v_head_dim,
+        quant_config: Optional["QuantizationConfig"] = None,
+    ):
+        super().__init__()
+        self.qkv_a_proj = fused_qkv_a_proj_with_mqa
+        self.q_a_layernorm = q_a_layernorm
+        self.kv_a_layernorm = kv_a_layernorm
+        self.q_b_proj = q_b_proj
+        self.w_kc = w_kc.contiguous()
+        self.rotary_emb = rotary_emb
+        self.layer_id = layer_id
+        self.quant_config = quant_config
+        self.has_preprocess_weights = False
+        self.dtype = None
+
+        self.q_lora_rank = self.q_b_proj.input_size  # 1536
+        self.kv_lora_rank = self.kv_a_layernorm.hidden_size  # 512
+        self.num_local_heads = num_local_heads  # tp
+        self.qk_nope_head_dim = qk_nope_head_dim  # 128
+        self.qk_rope_head_dim = qk_rope_head_dim  # 64
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.q_b_proj_weight_scale = self.q_b_proj.weight_scale.view(1, -1).to(
+            torch.float
+        )
+
+    def preprocess_weights(self, hidden_states):
+        self.dummy = torch.zeros(
+            (hidden_states.shape[-1]),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        self.qkv_a_proj_input_offset = self.qkv_a_proj.input_offset.to(dtype=torch.int8)
+        self.q_b_proj_input_offset = self.q_b_proj.input_offset.to(dtype=torch.int8)
+
+        # matmul_0 weight [7168, 2112]
+        fused_qkv_a_proj_with_mqa_weight_q = self.qkv_a_proj.weight.data[
+            :, : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_weight_kv = self.qkv_a_proj.weight.data[
+            :, self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_weight_kv_t = (
+            fused_qkv_a_proj_with_mqa_weight_kv.t().contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_weight_kv_t = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_weight_kv_t, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_weight_kv = (
+            fused_qkv_a_proj_with_mqa_weight_kv_t.t().contiguous()
+        )
+        # cat nz
+        fused_qkv_a_proj_with_mqa_weight_new = torch.cat(
+            (fused_qkv_a_proj_with_mqa_weight_kv, fused_qkv_a_proj_with_mqa_weight_q),
+            dim=-1,
+        )
+        fused_qkv_a_proj_with_mqa_weight = (
+            fused_qkv_a_proj_with_mqa_weight_new.t().contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_weight_nz = (
+            transdata(fused_qkv_a_proj_with_mqa_weight, block_size=(16, 32))
+            .unsqueeze(0)
+            .contiguous()
+        )
+        self.qkv_a_proj_weight_nz = npu_format_cast(fused_qkv_a_proj_with_mqa_weight_nz)
+
+        # matmul_0 deq_scale [2112]
+        fused_qkv_a_proj_with_mqa_deq_scale_q = self.qkv_a_proj.deq_scale.data[
+            : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = self.qkv_a_proj.deq_scale.data[
+            self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = (
+            fused_qkv_a_proj_with_mqa_deq_scale_kv.reshape(
+                self.kv_lora_rank + self.qk_rope_head_dim, -1
+            ).contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_deq_scale_kv, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = (
+            fused_qkv_a_proj_with_mqa_deq_scale_kv.view(
+                self.kv_lora_rank + self.qk_rope_head_dim
+            ).contiguous()
+        )
+        self.qkv_a_proj_deq_scale_kvq = torch.cat(
+            (
+                fused_qkv_a_proj_with_mqa_deq_scale_kv,
+                fused_qkv_a_proj_with_mqa_deq_scale_q,
+            ),
+            dim=-1,
+        )
+
+        # matmul_0 quant_bias [2112]
+        fused_qkv_a_proj_with_mqa_quant_bias_q = self.qkv_a_proj.quant_bias.data[
+            : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = self.qkv_a_proj.quant_bias.data[
+            self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = (
+            fused_qkv_a_proj_with_mqa_quant_bias_kv.reshape(
+                self.kv_lora_rank + self.qk_rope_head_dim, -1
+            ).contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_quant_bias_kv, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = (
+            fused_qkv_a_proj_with_mqa_quant_bias_kv.view(
+                self.kv_lora_rank + self.qk_rope_head_dim
+            ).contiguous()
+        )
+        self.qkv_a_proj_quant_bias_kvq = torch.cat(
+            (
+                fused_qkv_a_proj_with_mqa_quant_bias_kv,
+                fused_qkv_a_proj_with_mqa_quant_bias_q,
+            ),
+            dim=-1,
+        )
+
+        # matmul_1 weight [1536, num_head * 192]
+        q_b_proj_weight = self.q_b_proj.weight.data.clone()
+        q_b_proj_weight = q_b_proj_weight.t().reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_weight = trans_rope_weight(q_b_proj_weight, self.qk_rope_head_dim)
+        q_b_proj_weight = q_b_proj_weight.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim), -1
+        )
+        q_b_proj_weight_nz = (
+            transdata(q_b_proj_weight, block_size=(16, 32)).unsqueeze(0).contiguous()
+        )
+        self.q_b_proj_weight_nz = npu_format_cast(q_b_proj_weight_nz)
+
+        # matmul_1 deq_scale [num_head * 192]
+        q_b_proj_deq_scale = self.q_b_proj.deq_scale.data.clone()
+        q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_deq_scale = trans_rope_weight(
+            q_b_proj_deq_scale, self.qk_rope_head_dim
+        )
+        self.q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
+        )
+
+        # matmul_1 quant_bias [num_head * 192]
+        q_b_proj_quant_bias = self.q_b_proj.quant_bias.data.clone()
+        q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_quant_bias = trans_rope_weight(
+            q_b_proj_quant_bias, self.qk_rope_head_dim
+        )
+        self.q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
+        )
+
+    def mlaprolog_preprocess_weight(self):
+        self.qkv_a_proj.weight.data = self.qkv_a_proj.weight.data.transpose(0, 1)
+        qkv_a_proj_weight_q = self.qkv_a_proj.weight.data[:, : self.q_lora_rank].clone()
+        qkv_a_proj_weight_kv = self.qkv_a_proj.weight.data[
+            :, self.q_lora_rank :
+        ].clone()
+        self.q_a_proj_weight = npu_format_cast(qkv_a_proj_weight_q)
+        self.kv_a_proj_weight = npu_format_cast(qkv_a_proj_weight_kv)
+
+    def get_sin_cos(self, positions):
+        cos_sin = self.rotary_emb.cos_sin_cache[positions]
+        cos, sin = cos_sin.chunk(2, dim=-1)
+        cos = cos.repeat(1, 2)
+        sin = sin.repeat(1, 2)
+        return cos, sin
+
+    def get_kv_cache_and_cache_idx(self, forward_batch):
+        k_cache, v_cache = get_token_to_kv_pool().get_kv_buffer(self.layer_id)
+        slot_mapping = forward_batch.out_cache_loc.to(dtype=torch.int32)
+        return k_cache, v_cache, slot_mapping
+
+    def forward_absorb_prepare_npu_rms_norm_cache(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch,
+        zero_allocator,
+    ):
+        bsz, _ = hidden_states.view(-1, hidden_states.shape[-1]).shape
+        self.dtype = hidden_states.dtype
+        if self.layer_id == 0:
+            self.cos, self.sin = self.get_sin_cos(positions)
+            self.rotary_emb.cos_cached, self.rotary_emb.sin_cache = self.cos, self.sin
+        else:
+            self.cos, self.sin = self.rotary_emb.cos_cached, self.rotary_emb.sin_cache
+
+        self.kvCache, self.kvCacheRope, self.slotmapping = (
+            self.get_kv_cache_and_cache_idx(forward_batch)
+        )
+
+        if not self.has_preprocess_weights:
+            self.has_preprocess_weights = True
+
+        cos, sin = self.cos, self.sin
+
+        if self.q_lora_rank is not None:
+            fused_qkv_a_proj_out = self.qkv_a_proj(hidden_states)[0]
+            q_lowrank, latent_cache = fused_qkv_a_proj_out.split(
+                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
+            )
+            q = self.q_a_layernorm(q_lowrank)
+            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+        else:
+            q = self.q_proj(hidden_states)[0].view(
+                -1, self.num_local_heads, self.qk_head_dim
+            )
+            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+
+        q_nope, q_pe = torch.split(
+            q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+        )  # b*s,n,d
+
+        q_nope = q_nope.view(-1, self.num_local_heads, self.qk_nope_head_dim)
+        q_nope = torch.matmul(q_nope.transpose(0, 1), self.w_kc).transpose(0, 1)
+
+        q_pe = q_pe.view(-1, self.num_local_heads, 1, self.qk_rope_head_dim)
+        cos = cos.view(-1, 1, 1, self.qk_rope_head_dim)
+        sin = sin.view(-1, 1, 1, self.qk_rope_head_dim)
+        q_pe = torch.ops.npu.npu_interleave_rope(q_pe, cos, sin)  # (B,N,S,D)
+        q_pe = q_pe.view(cos.shape[0], self.num_local_heads, self.qk_rope_head_dim)
+
+        latent_cache = latent_cache.view(
+            -1, 1, 1, self.kv_lora_rank + self.qk_rope_head_dim
+        )  # (B*S,N,1,D)
+
+        cache_mode = "PA_NZ" if is_fia_nz() else "PA_BNSD"
+        self.kvCache = self.kvCache.view(
+            -1,
+            get_attn_backend().page_size,
+            1,
+            get_attn_backend().kv_lora_rank,
+        )
+        self.kvCacheRope = self.kvCacheRope.view(
+            -1,
+            get_attn_backend().page_size,
+            1,
+            get_attn_backend().qk_rope_head_dim,
+        )
+        k_rope, k_nope, _, _ = torch.ops.npu.npu_kv_rmsnorm_rope_cache(
+            latent_cache,
+            self.kv_a_layernorm.weight,
+            cos,
+            sin,
+            self.slotmapping.to(torch.int64),
+            self.kvCacheRope,
+            self.kvCache,
+            epsilon=self.kv_a_layernorm.variance_epsilon,
+            cache_mode=cache_mode,
+        )
+
+        return (q_pe, k_rope, q_nope, k_nope, forward_batch, zero_allocator, positions)
+
+    def forward_mlapo(self, positions, hidden_states, forward_batch, zero_allocator):
+        input_dtype = hidden_states.dtype
+        if not self.has_preprocess_weights:
+            self.preprocess_weights(hidden_states)
+            self.has_preprocess_weights = True
+            self.dtype = hidden_states.dtype
+
+        if self.layer_id == 0:
+            cos, sin = self.get_sin_cos(positions)
+            self.rotary_emb.cos_cached, self.rotary_emb.sin_cache = cos, sin
+        else:
+            cos, sin = self.rotary_emb.cos_cached, self.rotary_emb.sin_cache
+
+        k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch)
+
+        q_nope_out = torch.empty(
+            (hidden_states.shape[0], self.w_kc.shape[0], k_cache.shape[-1]),
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+        q_rope_out = torch.empty(
+            (hidden_states.shape[0], self.w_kc.shape[0], v_cache.shape[-1]),
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+        if is_fia_nz():
+            kv_shape, kv_rope_shape = k_cache.shape, v_cache.shape
+            num_blocks, block_size, num_heads, _ = kv_shape
+            k_cache = k_cache.view(
+                num_blocks, num_heads * self.kv_lora_rank // 16, block_size, 16
+            )
+            v_cache = v_cache.view(
+                num_blocks, num_heads * self.qk_rope_head_dim // 16, block_size, 16
+            )
+        # TODO: dummy inputs to be removed
+        # https://github.com/sgl-project/sgl-kernel-npu/issues/78
+        if hasattr(self.q_a_layernorm, "bias"):
+            q_a_layernorm_bias = self.q_a_layernorm.bias
+        else:
+            q_a_layernorm_bias = self.dummy
+
+        torch.ops.npu.mla_preprocess(
+            hidden_states,
+            self.dummy,
+            self.dummy,
+            self.qkv_a_proj_weight_nz,
+            self.qkv_a_proj_deq_scale_kvq,
+            self.q_a_layernorm.weight,
+            q_a_layernorm_bias,
+            self.q_b_proj_weight_nz,
+            self.q_b_proj_deq_scale,
+            self.kv_a_layernorm.weight,
+            cos,
+            sin,
+            self.w_kc,
+            k_cache,
+            v_cache,
+            slot_mapping,
+            quant_scale0=self.qkv_a_proj.input_scale,
+            quant_offset0=self.qkv_a_proj_input_offset,
+            bias0=self.qkv_a_proj_quant_bias_kvq,
+            quant_scale1=self.q_b_proj.input_scale,
+            quant_offset1=self.q_b_proj_input_offset,
+            bias1=self.q_b_proj_quant_bias,
+            cache_mode="nzcache" if is_fia_nz() else "krope_ctkv",
+            quant_mode="per_tensor_quant_asymm",
+            q_out0=q_nope_out,
+            kv_cache_out0=k_cache,
+            q_out1=q_rope_out,
+            kv_cache_out1=v_cache,
+        )
+
+        if is_fia_nz():
+            k_cache = k_cache.view(kv_shape)
+            v_cache = v_cache.view(kv_rope_shape)
+
+        return (
+            q_rope_out,
+            v_cache,
+            q_nope_out,
+            k_cache,
+            forward_batch,
+            zero_allocator,
+            positions,
+        )
+
+    def forward_mlaprolog(self, positions, hidden_states, forward_batch):
+        if not self.has_preprocess_weights:
+            self.mlaprolog_preprocess_weight()
+            self.has_preprocess_weights = True
+        self.cos, self.sin = self.get_sin_cos(positions)
+        k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch)
+        mla_prolog_input_args = {
+            "token_x": hidden_states,
+            "weight_dq": self.q_a_proj_weight,
+            "weight_uq_qr": self.q_b_proj.weight,
+            "weight_uk": self.w_kc,
+            "weight_dkv_kr": self.kv_a_proj_weight,
+            "rmsnorm_gamma_cq": self.q_a_layernorm.weight,
+            "rmsnorm_gamma_ckv": self.kv_a_layernorm.weight,
+            "rope_sin": self.sin,
+            "rope_cos": self.cos,
+            "kv_cache": k_cache,
+            "kr_cache": v_cache,
+            "cache_index": slot_mapping.to(dtype=torch.int64),
+            "dequant_scale_w_uq_qr": self.q_b_proj_weight_scale,
+            "rmsnorm_epsilon_cq": self.q_a_layernorm.variance_epsilon,
+            "rmsnorm_epsilon_ckv": self.kv_a_layernorm.variance_epsilon,
+            "cache_mode": "PA_BSND",
+            "query_norm_flag": True,
+            "weight_quant_mode": 1,  # 0:no quant; 1:uq_qr: quant; 2: weight_dq,weight_uq_qr,weight_dkv_kr: quant
+        }
+        q_nope, q_pe, dequant_scale_q_nope, qr, dequant_q_norm = (
+            torch.ops.custom.npu_mla_prolog_v3(**mla_prolog_input_args)
+        )
+        dequant_q_norm = dequant_q_norm.view(hidden_states.shape[0])
+        return (
+            q_pe,
+            v_cache,
+            q_nope,
+            k_cache,
+            qr,
+            forward_batch,
+            positions,
+            dequant_q_norm,
+        )
+
+    def forward(self, positions, hidden_states, forward_batch, zero_allocator):
+        # assert self.quant_config and self.quant_config.get_name() == "modelslim"
+        # route by `qkv_a_proj` quant type as MTP layers can be unquantized
+        _is_w8a8 = (
+            hasattr(self.qkv_a_proj.quant_method, "quantization_config")
+            and self.qkv_a_proj.quant_method.quantization_config.get_name()
+            == "modelslim"
+        )
+        # with the mlaprolog enabled, the kv_b_proj layers are unquantized
+        _is_mlaprolog = hasattr(self.quant_config, "ignore") and any(
+            re.fullmatch(r".*kv_b_proj", l) for l in self.quant_config.ignore
+        )
+        if _is_w8a8:
+            return self.forward_mlapo(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+        elif _is_mlaprolog:
+            return self.forward_mlaprolog(positions, hidden_states, forward_batch)
+        else:
+            return self.forward_absorb_prepare_npu_rms_norm_cache(
+                positions, hidden_states, forward_batch, zero_allocator
+            )

python/sglang/srt/hardware_backend/npu/batch_invariant_ops/npu_batch_invariant_ops.py

@@ -0,0 +1,62 @@
+# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/batch_invariant_ops.py
+
+import batch_invariant_ops  # noqa: F401
+import torch
+import torch_npu
+
+
+def npu_mm_batch_invariant(a, b):
+    return torch.ops.batch_invariant_ops.npu_mm_batch_invariant(a, b)
+
+
+def npu_matmul_batch_invariant(a, b):
+    return torch.ops.batch_invariant_ops.npu_matmul_batch_invariant(a, b)
+
+
+def npu_mean_batch_invariant(
+    input, dim, keepdim=False, dtype: torch.dtype | None = None
+):
+    assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}"
+    if len(dim) == 1:
+        return torch.ops.batch_invariant_ops.npu_reduce_mean_batch_invariant(
+            input, dim[0], keepdim=keepdim
+        )
+    else:
+        assert input.dtype in {
+            torch.float16,
+            torch.bfloat16,
+            torch.float32,
+        }, "only float types supported for now"
+        n_elems = 1
+        for d in dim:
+            n_elems *= input.shape[d]
+        return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems
+
+
+def npu_log_softmax_batch_invariant(input, dim, _half_to_float):
+    assert not _half_to_float, "not implemented"
+    return torch.ops.batch_invariant_ops.npu_log_softmax_batch_invariant(input, dim=dim)
+
+
+def npu_fused_infer_attention_score_batch_invariant(*args, **kwargs):
+    return (
+        torch.ops.batch_invariant_ops.npu_fused_infer_attention_score_batch_invariant(
+            *args, **kwargs
+        )
+    )
+
+
+def npu_add_rms_norm_batch_invariant(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+):
+    """
+    AclnnAddRmsNorm can't ensure batch invariant,
+    so we need to split it into add and rms_norm.
+    """
+    x_ = x + residual
+    residual_ = x_
+    x_, _ = torch_npu.npu_rms_norm(x_, weight, eps)
+    return x_, None, residual_

python/sglang/srt/hardware_backend/npu/cmo.py

@@ -0,0 +1,83 @@
+import torch
+
+cmo_stream = None
+share_stream = None
+
+
+def get_cmo_stream():
+    """
+    Cache Management Operation(CMO).
+    Launch a new stream to prefetch the weight of matmul when running other
+    AIV or communication kernels, aiming to overlap the memory access time.
+    """
+    global cmo_stream
+    return cmo_stream
+
+
+def set_cmo_stream(stream):
+    global cmo_stream
+    cmo_stream = stream
+
+
+def prepare_weight_cache(handle, cache, PREFETCH_MAX_SIZE=1000000000):
+    """
+    PREFETCH_MAX_SIZE: maximum size (bytes) for each prefetch operation.
+    This affects the time spent in prefetch:
+        time ≈ PREFETCH_MAX_SIZE / system_bandwidth
+    """
+    import torch_npu
+
+    stream = get_cmo_stream()
+    if stream is None:
+        stream = torch.npu.Stream()
+        set_cmo_stream(stream)
+    stream.wait_stream(torch.npu.current_stream())
+    with torch.npu.stream(stream):
+        if isinstance(cache, list):
+            for weight in cache:
+                torch_npu.npu_prefetch(
+                    weight,
+                    handle,
+                    PREFETCH_MAX_SIZE,
+                )
+        else:
+            torch_npu.npu_prefetch(
+                cache,
+                handle,
+                PREFETCH_MAX_SIZE,
+            )
+
+
+def wait_cmo_stream():
+    stream = get_cmo_stream()
+    if stream is not None:
+        cur_stream = torch.npu.current_stream()
+        cur_stream.wait_stream(stream)
+
+
+def get_share_stream():
+    global share_stream
+    return share_stream
+
+
+def set_share_stream(stream):
+    global share_stream
+    share_stream = stream
+
+
+def wait_share_stream():
+    stream = get_share_stream()
+    if stream is not None:
+        cur_stream = torch.npu.current_stream()
+        cur_stream.wait_stream(stream)
+
+
+def shared_expert_on_independent_stream(hidden_states, forward_func):
+    stream = get_share_stream()
+    if stream is None:
+        stream = torch.npu.Stream()
+        set_share_stream(stream)
+    stream.wait_stream(torch.npu.current_stream())
+    with torch.npu.stream(stream):
+        shared_output = forward_func(hidden_states)
+        return shared_output

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py

@@ -0,0 +1,749 @@
+"""DSV4-NPU SWA + c4/c128 paged allocator.
+
+Subclasses :class:`SWATokenToKVPoolAllocator` and adds paged allocation for the
+c4/c128 compressed-KV pools and their tail-only compress-state pools, alongside
+the parent's full + SWA pools.
+
+Per ``alloc_extend`` / ``alloc_decode``:
+  1. super() allocates the full + SWA slots (``out_full_loc``).
+  2. Allocate c4/c128 KV slots — one compressed token per ``ratio`` raw tokens
+     (``seq_len // ratio - prefix_len // ratio``) — via the standard
+     :class:`NPUPagedTokenToKVPoolAllocator` over the pool's c4/c128 KV buffers.
+  3. Allocate the c4/c128 compress-state slots the same way, tail-only per req,
+     using the per-req lens the scheduler packed into ``DSV4StateLens``.
+  4. Return a :class:`DSV4OutCacheLoc` bundling all five slot families.
+
+State slots are paged because the NPU fused compressor runs ``cache_mode=1``; the
+base class' ``translate_kv_loc_to_compress_state_loc`` ring-hash is the CUDA-only
+path and is unused on NPU. The bundle is the explicit return value:
+mem_cache/common.py unpacks ``out_full_loc`` and stashes the bundle on
+``batch.out_cache_loc_dsv4``; ``DSV4NPUReqToTokenPool`` writes the per-req
+``req_to_token_c{4,128}[_state]`` tables that :meth:`free` and the last_loc
+lookups read back.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.configs.model_config import is_deepseek_v4
+from sglang.srt.hardware_backend.npu.allocator_npu import NPUPagedTokenToKVPoolAllocator
+from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
+    maybe_write_dsv4_extend,
+)
+from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+from sglang.srt.mem_cache.common import alloc_paged_token_slots_extend
+from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, DSV4StateLens
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+
+def get_last_loc(
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> torch.Tensor:
+    """Slot id of each req's last already-allocated token, or -1 when
+    ``prefix_lens[i] == 0`` (fresh req).
+
+    Looks up ``req_to_token[req, prefix_lens - 1]`` to anchor the paged
+    allocator's ``alloc_extend`` on the real previous tail slot, preserving the
+    intra-page slot continuity the kernel's ``cmp_block_table`` relies on (the
+    allocator debug-asserts ``(last_loc + 1) % page_size == prefix_lens %
+    page_size``). Result dtype matches ``prefix_lens``.
+    """
+    req_pool_indices = req_pool_indices.to(torch.int64)
+    safe_idx = (prefix_lens.to(torch.int64) - 1).clamp(min=0)
+    looked_up = req_to_token[req_pool_indices, safe_idx].to(prefix_lens.dtype)
+    return torch.where(
+        prefix_lens > 0,
+        looked_up,
+        torch.full_like(prefix_lens, -1),
+    )
+
+
+def alloc_paged_token_slots_extend_npu(*args, batch=None, **kwargs):
+    if batch is not None and is_deepseek_v4(batch.model_config.hf_config):
+        return alloc_paged_token_slots_reserve_extend(*args, batch=batch, **kwargs)
+    return alloc_paged_token_slots_extend(*args, batch=batch, **kwargs)
+
+
+def alloc_paged_token_slots_reserve_extend(
+    tree_cache,
+    prefix_lens: torch.Tensor,
+    prefix_lens_cpu: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seq_lens_cpu: torch.Tensor,
+    last_loc: torch.Tensor,
+    extend_num_tokens: int,
+    *,
+    req_pool_indices: Optional[torch.Tensor] = None,
+    dsv4_state_lens: Optional[DSV4StateLens] = None,
+    batch=None,
+):
+    """Allocate reserved draft slots and update DSV4 per-request tables."""
+    if dsv4_state_lens is None and batch is not None:
+        allocator = batch.token_to_kv_pool_allocator
+        dsv4_state_lens = (
+            allocator.compute_dsv4_state_lens_reserve(
+                batch.reqs, prefix_lens_cpu, seq_lens_cpu
+            )
+            if hasattr(allocator, "compute_dsv4_state_lens_reserve")
+            else None
+        )
+
+    out_cache_loc = alloc_paged_token_slots_extend(
+        tree_cache,
+        prefix_lens,
+        prefix_lens_cpu,
+        seq_lens,
+        seq_lens_cpu,
+        last_loc,
+        extend_num_tokens,
+        req_pool_indices=req_pool_indices,
+        dsv4_state_lens=dsv4_state_lens,
+        batch=batch,
+    )
+    if batch is not None:
+        maybe_write_dsv4_extend(
+            batch,
+            batch.req_pool_indices_cpu,
+            prefix_lens_cpu,
+            seq_lens_cpu,
+            c4_state_alloc_offsets=prefix_lens_cpu,
+            c128_state_alloc_offsets=prefix_lens_cpu,
+        )
+    return out_cache_loc
+
+
+class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
+    """SWA allocator + c4/c128 KV and compress-state paged allocators for DSV4 on NPU."""
+
+    def __init__(
+        self,
+        size: int,
+        size_swa: int,
+        page_size: int,
+        dtype: torch.dtype,
+        device: str,
+        kvcache,
+        need_sort: bool,
+    ):
+        super().__init__(
+            size=size,
+            size_swa=size_swa,
+            page_size=page_size,
+            dtype=dtype,
+            device=device,
+            kvcache=kvcache,
+            need_sort=need_sort,
+        )
+
+        def mk(pool_size, pool):
+            # c4/c128 KV and state sub-pools implement KVCache, so they drop into
+            # the standard paged allocator. pool_size is in compressed-token units.
+            return NPUPagedTokenToKVPoolAllocator(
+                pool_size,
+                page_size=page_size,
+                dtype=dtype,
+                device=device,
+                kvcache=pool,
+                need_sort=need_sort,
+            )
+
+        self.c4_attn_allocator = mk(kvcache.c4_size, kvcache.c4_kv_pool)
+        self.c128_attn_allocator = mk(kvcache.c128_size, kvcache.c128_kv_pool)
+
+        # State allocators (paged, NPU-only). Any layer's pool works as KVCache
+        # pointer (slot alloc is layer-agnostic); None when no c{ratio} layers or
+        # zero budget.
+        self.c4_state_attn_allocator: Optional[NPUPagedTokenToKVPoolAllocator] = None
+        self.c128_state_attn_allocator: Optional[NPUPagedTokenToKVPoolAllocator] = None
+        state_pools = getattr(kvcache, "compress_state_pools", None)
+        if state_pools:
+
+            def first_state_pool(want_ratio):
+                return next(
+                    (
+                        p
+                        for r, p in zip(kvcache.compression_ratios, state_pools)
+                        if r == want_ratio and p is not None
+                    ),
+                    None,
+                )
+
+            c4_state_pool = first_state_pool(4)
+            c128_state_pool = first_state_pool(128)
+            if c4_state_pool is not None and kvcache.c4_state_pool_size > 0:
+                self.c4_state_attn_allocator = mk(
+                    kvcache.c4_state_pool_size, c4_state_pool
+                )
+            if c128_state_pool is not None and kvcache.c128_state_pool_size > 0:
+                self.c128_state_attn_allocator = mk(
+                    kvcache.c128_state_pool_size, c128_state_pool
+                )
+
+        # Returned by the c-pool helpers when a step adds no compressed tokens.
+        self._empty_loc = torch.empty((0,), dtype=torch.int64, device=device)
+
+        # Per-call handle to the DSV4NPUReqToTokenPool, stashed by alloc_extend/
+        # alloc_decode for last_loc lookups; avoids a permanent allocator->pool ref.
+        self._cur_req_to_token_pool = None
+
+    @staticmethod
+    def _compute_c_extend_counts(
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        ratio: int,
+    ) -> int:
+        """New compressed-K tokens this extend produces across the batch:
+        ``sum_i (seq_lens[i] // ratio - prefix_lens[i] // ratio)``."""
+        if prefix_lens_cpu is None or seq_lens_cpu is None:
+            return 0
+        diff = ((seq_lens_cpu // ratio) - (prefix_lens_cpu // ratio)).clamp(min=0)
+        return int(diff.sum().item())
+
+    @staticmethod
+    def _pool_exhausted(
+        ratio: int, kind: str, need: int, available: int
+    ) -> RuntimeError:
+        return RuntimeError(
+            f"DSV4 c{ratio} {kind} pool exhausted: need {need} new slots, "
+            f"available={available}. Raise --mem-fraction-static, lower "
+            f"--max-running-requests, or check that "
+            f"DSV4NPUTokenToKVPoolAllocator.free(req=...) releases {kind} slots "
+            f"on req finish."
+        )
+
+    def _alloc_state_extend(
+        self,
+        allocator: Optional[NPUPagedTokenToKVPoolAllocator],
+        raw_prefix_lens: torch.Tensor,
+        state_prefix_lens: torch.Tensor,
+        state_prefix_lens_cpu: torch.Tensor,
+        state_seq_lens: torch.Tensor,
+        state_seq_lens_cpu: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        last_loc_dtype: torch.dtype,
+        state_extend_num_tokens: int,
+        ratio: int,
+    ) -> torch.Tensor:
+        """Allocate tail-only state-pool slots for an extend at ``ratio``.
+
+        The state pool is a separate paged slot space; each req allocates only
+        its trailing window (cumulative lens precomputed by
+        ``ScheduleBatch._compute_dsv4_state_lens_*`` and passed via
+        ``DSV4StateLens``). ``state_last_loc`` is looked up from
+        ``req_to_token_c{ratio}_state`` at the RAW position
+        ``raw_prefix_lens - 1`` (the last position the previous extend/decode
+        populated). Returns ``_empty_loc`` when the allocator is absent (no
+        c{ratio} layers) or there is nothing to add.
+        """
+        if allocator is None or state_extend_num_tokens == 0:
+            return self._empty_loc
+
+        assert self._cur_req_to_token_pool is not None, (
+            "alloc_extend/alloc_decode must be called with req_to_token_pool= "
+            "for the state-pool last_loc lookup."
+        )
+        state_table = (
+            self._cur_req_to_token_pool.req_to_token_c4_state
+            if ratio == 4
+            else self._cur_req_to_token_pool.req_to_token_c128_state
+        )
+        state_last_loc = get_last_loc(
+            state_table, req_pool_indices, raw_prefix_lens
+        ).to(last_loc_dtype)
+
+        result = allocator.alloc_extend(
+            state_prefix_lens,
+            state_prefix_lens_cpu,
+            state_seq_lens,
+            state_seq_lens_cpu,
+            state_last_loc,
+            state_extend_num_tokens,
+        )
+        if result is None:
+            raise self._pool_exhausted(
+                ratio, "state", state_extend_num_tokens, allocator.available_size()
+            )
+        return result
+
+    def _alloc_c_extend(
+        self,
+        allocator: NPUPagedTokenToKVPoolAllocator,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        last_loc_dtype: torch.dtype,
+        ratio: int,
+    ) -> torch.Tensor:
+        """Allocate compressed-KV slots for an extend at ``ratio``.
+
+        Prefix/seq lens are translated to compressed units (``// ratio``); the
+        c-pool last_loc comes from ``req_to_token_c{ratio}`` via
+        :func:`get_last_loc` so the paged allocator continues in-page (or opens
+        a fresh page at a ratio boundary), keeping the intra-page continuity the
+        ``cmp_block_table`` reader relies on. Returns ``_empty_loc`` when this
+        step closes no compressed token.
+        """
+        c_extend = self._compute_c_extend_counts(prefix_lens_cpu, seq_lens_cpu, ratio)
+        if c_extend == 0:
+            return self._empty_loc
+
+        assert self._cur_req_to_token_pool is not None, (
+            "alloc_extend/alloc_decode must be called with req_to_token_pool= "
+            "for the c-pool last_loc lookup."
+        )
+        c_table = (
+            self._cur_req_to_token_pool.req_to_token_c4
+            if ratio == 4
+            else self._cur_req_to_token_pool.req_to_token_c128
+        )
+        c_prefix = (prefix_lens // ratio).to(prefix_lens.dtype)
+        c_seq = (seq_lens // ratio).to(seq_lens.dtype)
+        c_last_loc = get_last_loc(c_table, req_pool_indices, c_prefix).to(
+            last_loc_dtype
+        )
+
+        result = allocator.alloc_extend(
+            c_prefix,
+            prefix_lens_cpu // ratio,
+            c_seq,
+            seq_lens_cpu // ratio,
+            c_last_loc,
+            c_extend,
+        )
+        if result is None:
+            raise self._pool_exhausted(
+                ratio, "KV", c_extend, allocator.available_size()
+            )
+        return result
+
+    def _alloc_c_and_state(
+        self,
+        out_full_loc: torch.Tensor,
+        out_swa_loc: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc_dtype: torch.dtype,
+        req_pool_indices: Optional[torch.Tensor],
+        dsv4_state_lens: Optional[DSV4StateLens],
+    ) -> DSV4OutCacheLoc:
+        """Allocate c4/c128 KV + state slots and bundle them with full/swa loc.
+
+        Shared by alloc_extend / alloc_decode (which differ only in how
+        prefix_lens is derived). State lens are tail-only, precomputed by
+        ScheduleBatch._compute_dsv4_state_lens_*; raw prefix_lens drives the
+        state last_loc lookup.
+        """
+        assert req_pool_indices is not None, (
+            "DSV4NPUTokenToKVPoolAllocator requires req_pool_indices "
+            "(forwarded from batch.req_pool_indices)."
+        )
+        if dsv4_state_lens is not None:
+            out_c4_state_loc = self._alloc_state_extend(
+                self.c4_state_attn_allocator,
+                prefix_lens,
+                dsv4_state_lens.c4_prefix_lens,
+                dsv4_state_lens.c4_prefix_lens_cpu,
+                dsv4_state_lens.c4_seq_lens,
+                dsv4_state_lens.c4_seq_lens_cpu,
+                req_pool_indices,
+                last_loc_dtype,
+                dsv4_state_lens.c4_extend_num_tokens,
+                ratio=4,
+            )
+            out_c128_state_loc = self._alloc_state_extend(
+                self.c128_state_attn_allocator,
+                prefix_lens,
+                dsv4_state_lens.c128_prefix_lens,
+                dsv4_state_lens.c128_prefix_lens_cpu,
+                dsv4_state_lens.c128_seq_lens,
+                dsv4_state_lens.c128_seq_lens_cpu,
+                req_pool_indices,
+                last_loc_dtype,
+                dsv4_state_lens.c128_extend_num_tokens,
+                ratio=128,
+            )
+        else:
+            out_c4_state_loc = self._empty_loc
+            out_c128_state_loc = self._empty_loc
+        out_c4_loc = self._alloc_c_extend(
+            self.c4_attn_allocator,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            req_pool_indices,
+            last_loc_dtype,
+            ratio=4,
+        )
+        out_c128_loc = self._alloc_c_extend(
+            self.c128_attn_allocator,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            req_pool_indices,
+            last_loc_dtype,
+            ratio=128,
+        )
+        return DSV4OutCacheLoc(
+            out_full_loc=out_full_loc,
+            out_swa_loc=out_swa_loc,
+            out_c4_loc=out_c4_loc,
+            out_c128_loc=out_c128_loc,
+            out_c4_state_loc=out_c4_state_loc,
+            out_c128_state_loc=out_c128_state_loc,
+        )
+
+    def compute_dsv4_state_lens_extend(
+        self, reqs: List[Req], seq_lens: List[int]
+    ) -> Optional[DSV4StateLens]:
+        """Per-req c{4,128}_state pool alloc lens for extend (tail-only).
+
+        State pool stores only the trailing portion of each sequence (the c{N}
+        compressor's read/write window); the tail length depends on raw
+        seq_len's alignment to the SWA page boundary (128)::
+
+            c4_alloc_len  = tail + 128 if (tail <= 3 and seq_len >= 128) else tail
+            c128_alloc_len = tail                  where tail = seq_len % 128
+
+        Long prefills allocate only the trailing partial window, not slots for
+        already-compressed positions, so the small paged state pool (~256
+        slots/req) stays sufficient even for 28k-token prompts.
+
+        Mutates per-req cumulative state via getattr/setattr so the community
+        ``Req`` needs no DSV4 field declarations:
+          * ``req.c{4,128}_state_kv_len`` — cumulative slot count (prefix for
+            the paged allocator; never decreases on eviction).
+          * ``req.c{4,128}_state_alloc_offset`` — low-water raw-position mark
+            for eviction (see ``dsv4_common_hooks.maybe_evict_dsv4_state``).
+
+        Returns None when this model has no paged state pools (CUDA / non-V4 /
+        zero budget) — callers pass that straight through as ``dsv4_state_lens``.
+        """
+        if self.c4_state_attn_allocator is None:
+            return None
+        c4_prefix: List[int] = []
+        c4_seq: List[int] = []
+        c128_prefix: List[int] = []
+        c128_seq: List[int] = []
+        for req, seq_len in zip(reqs, seq_lens):
+            tail = seq_len % 128
+            c4_alloc_len = tail + 128 if (tail <= 3 and seq_len >= 128) else tail
+            c128_alloc_len = tail
+
+            prev_c4 = getattr(req, "c4_state_kv_len", 0)
+            prev_c128 = getattr(req, "c128_state_kv_len", 0)
+            new_c4 = prev_c4 + c4_alloc_len
+            new_c128 = prev_c128 + c128_alloc_len
+
+            c4_prefix.append(prev_c4)
+            c4_seq.append(new_c4)
+            c128_prefix.append(prev_c128)
+            c128_seq.append(new_c128)
+
+            req.c4_state_kv_len = new_c4
+            req.c128_state_kv_len = new_c128
+            req.c4_state_alloc_offset = seq_len - c4_alloc_len
+            req.c128_state_alloc_offset = seq_len - c128_alloc_len
+
+        return self._pack_state_lens(
+            c4_prefix,
+            c4_seq,
+            c128_prefix,
+            c128_seq,
+            c4_extend_num_tokens=int(sum(s - p for s, p in zip(c4_seq, c4_prefix))),
+            c128_extend_num_tokens=int(
+                sum(s - p for s, p in zip(c128_seq, c128_prefix))
+            ),
+        )
+
+    def compute_dsv4_state_lens_decode(
+        self, reqs: List[Req]
+    ) -> Optional[DSV4StateLens]:
+        """Per-req c{4,128}_state pool alloc lens for decode: exactly 1 new
+        state slot per req per pool. ``c{N}_state_alloc_offset`` does NOT
+        advance here (only eviction advances it). Returns None when there are
+        no paged state pools."""
+        if self.c4_state_attn_allocator is None:
+            return None
+        c4_prefix: List[int] = []
+        c4_seq: List[int] = []
+        c128_prefix: List[int] = []
+        c128_seq: List[int] = []
+        for req in reqs:
+            prev_c4 = getattr(req, "c4_state_kv_len", 0)
+            prev_c128 = getattr(req, "c128_state_kv_len", 0)
+            c4_prefix.append(prev_c4)
+            c4_seq.append(prev_c4 + 1)
+            c128_prefix.append(prev_c128)
+            c128_seq.append(prev_c128 + 1)
+            req.c4_state_kv_len = prev_c4 + 1
+            req.c128_state_kv_len = prev_c128 + 1
+
+        bs = len(reqs)
+        return self._pack_state_lens(
+            c4_prefix,
+            c4_seq,
+            c128_prefix,
+            c128_seq,
+            c4_extend_num_tokens=bs,
+            c128_extend_num_tokens=bs,
+        )
+
+    def compute_dsv4_state_lens_reserve(
+        self, reqs: List[Req], prefix_lens: List[int], seq_lens: List[int]
+    ) -> Optional[DSV4StateLens]:
+        """Allocate state slots for a speculative pre-reserved raw interval."""
+        if self.c4_state_attn_allocator is None:
+            return None
+
+        c4_prefix: List[int] = []
+        c4_seq: List[int] = []
+        c128_prefix: List[int] = []
+        c128_seq: List[int] = []
+        for req, prefix_len, seq_len in zip(reqs, prefix_lens, seq_lens):
+            reserve = max(0, int(seq_len) - int(prefix_len))
+            prev_c4 = getattr(req, "c4_state_kv_len", 0)
+            prev_c128 = getattr(req, "c128_state_kv_len", 0)
+            c4_prefix.append(prev_c4)
+            c4_seq.append(prev_c4 + reserve)
+            c128_prefix.append(prev_c128)
+            c128_seq.append(prev_c128 + reserve)
+            req.c4_state_kv_len = prev_c4 + reserve
+            req.c128_state_kv_len = prev_c128 + reserve
+
+        total = sum(max(0, int(s) - int(p)) for p, s in zip(prefix_lens, seq_lens))
+        return self._pack_state_lens(
+            c4_prefix,
+            c4_seq,
+            c128_prefix,
+            c128_seq,
+            c4_extend_num_tokens=total,
+            c128_extend_num_tokens=total,
+        )
+
+    def _pack_state_lens(
+        self,
+        c4_prefix: List[int],
+        c4_seq: List[int],
+        c128_prefix: List[int],
+        c128_seq: List[int],
+        *,
+        c4_extend_num_tokens: int,
+        c128_extend_num_tokens: int,
+    ) -> DSV4StateLens:
+        c4_prefix_cpu = torch.tensor(c4_prefix, dtype=torch.int64)
+        c4_seq_cpu = torch.tensor(c4_seq, dtype=torch.int64)
+        c128_prefix_cpu = torch.tensor(c128_prefix, dtype=torch.int64)
+        c128_seq_cpu = torch.tensor(c128_seq, dtype=torch.int64)
+        return DSV4StateLens(
+            c4_prefix_lens=c4_prefix_cpu.to(self.device, non_blocking=True),
+            c4_prefix_lens_cpu=c4_prefix_cpu,
+            c4_seq_lens=c4_seq_cpu.to(self.device, non_blocking=True),
+            c4_seq_lens_cpu=c4_seq_cpu,
+            c4_extend_num_tokens=c4_extend_num_tokens,
+            c128_prefix_lens=c128_prefix_cpu.to(self.device, non_blocking=True),
+            c128_prefix_lens_cpu=c128_prefix_cpu,
+            c128_seq_lens=c128_seq_cpu.to(self.device, non_blocking=True),
+            c128_seq_lens_cpu=c128_seq_cpu,
+            c128_extend_num_tokens=c128_extend_num_tokens,
+        )
+
+    def alloc_extend(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+        *,
+        req_pool_indices: Optional[torch.Tensor] = None,
+        dsv4_state_lens: Optional[DSV4StateLens] = None,
+        req_to_token_pool=None,
+    ) -> Optional[DSV4OutCacheLoc]:
+        # Stash per-req tables for this call's last_loc lookups (read by
+        # _alloc_c_extend / _alloc_state_extend); no permanent allocator->pool ref.
+        self._cur_req_to_token_pool = req_to_token_pool
+        out_full_loc = super().alloc_extend(
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc,
+            extend_num_tokens,
+        )
+        if out_full_loc is None:
+            return None
+
+        out_swa_loc = self.translate_loc_from_full_to_swa(out_full_loc)
+        assert out_swa_loc is not None, (
+            "translate_loc_from_full_to_swa returned None — "
+            "full_to_swa_index_mapping not initialized?"
+        )
+        return self._alloc_c_and_state(
+            out_full_loc,
+            out_swa_loc,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc.dtype,
+            req_pool_indices,
+            dsv4_state_lens,
+        )
+
+    def alloc_decode(
+        self,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        *,
+        req_pool_indices: Optional[torch.Tensor] = None,
+        dsv4_state_lens: Optional[DSV4StateLens] = None,
+        req_to_token_pool=None,
+    ) -> Optional[DSV4OutCacheLoc]:
+        self._cur_req_to_token_pool = req_to_token_pool
+        out_full_loc = super().alloc_decode(seq_lens, seq_lens_cpu, last_loc)
+        if out_full_loc is None:
+            return None
+
+        out_swa_loc = self.translate_loc_from_full_to_swa(out_full_loc)
+        # One new token per req. Model as an extend from (seq_len-1)//ratio to
+        # seq_len//ratio so _alloc_c_extend anchors on the real c-pool last_loc.
+        prefix_lens = (seq_lens - 1).clamp(min=0)
+        prefix_lens_cpu = (seq_lens_cpu - 1).clamp(min=0)
+        return self._alloc_c_and_state(
+            out_full_loc,
+            out_swa_loc,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc.dtype,
+            req_pool_indices,
+            dsv4_state_lens,
+        )
+
+    def free(
+        self,
+        free_index: Optional[torch.Tensor] = None,
+        *,
+        req=None,
+        req_to_token_pool=None,
+    ):
+        """Unified free for full/swa/c4/c128 pools. Two forms (may co-fire):
+
+          * ``free(free_index)`` — full + SWA only (tail/radix eviction; no req
+            identity, so c-pool free can't run).
+          * ``free(req=, req_to_token_pool=)`` — from DSV4NPUReqToTokenPool.free
+            on req finish: reads the per-req slot lists from
+            ``req_to_token_c{4,128}[_state]`` and returns them to the c-pools
+            (the paged allocator dedupes by page).
+
+        KV pools free ``[0, kv_len // ratio)``. State pools are 1-per-raw-token
+        and free only the tail ``[c{N}_state_alloc_offset, kv_len)`` — the prefix
+        was already returned by ScheduleBatch._evict_swa (state rides SWA
+        eviction); freeing it again would double-free (caught by the paged
+        allocator's debug_mode assert, corrupts the free list otherwise).
+        """
+        if free_index is not None:
+            super().free(free_index)
+
+        if req is None or req_to_token_pool is None:
+            return
+        kv_len = max(req.kv_committed_len, req.kv_allocated_len)
+        req_pool_idx = req.req_pool_idx
+        if kv_len <= 0 or req_pool_idx is None:
+            return
+
+        # KV pools: free the leading [0, kv_len // ratio) compressed slots.
+        for ratio, allocator, table_attr in (
+            (4, self.c4_attn_allocator, "req_to_token_c4"),
+            (128, self.c128_attn_allocator, "req_to_token_c128"),
+        ):
+            n = kv_len // ratio
+            if n > 0 and hasattr(req_to_token_pool, table_attr):
+                slots = getattr(req_to_token_pool, table_attr)[req_pool_idx, :n]
+                # to int64 — paged allocator's free does cpu()//page_size on it.
+                allocator.free(slots.to(torch.int64))
+
+        # State pools: free only the tail [c{N}_state_alloc_offset, kv_len).
+        for ratio, allocator, table_attr, off_attr in (
+            (
+                4,
+                self.c4_state_attn_allocator,
+                "req_to_token_c4_state",
+                "c4_state_alloc_offset",
+            ),
+            (
+                128,
+                self.c128_state_attn_allocator,
+                "req_to_token_c128_state",
+                "c128_state_alloc_offset",
+            ),
+        ):
+            if allocator is None or not hasattr(req_to_token_pool, table_attr):
+                continue
+            off = getattr(req, off_attr, 0)
+            if kv_len > off:
+                slots = getattr(req_to_token_pool, table_attr)[req_pool_idx, off:kv_len]
+                allocator.free(slots.to(torch.int64))
+
+    def backup_state(self):
+        # EAGLE/NEXTN draft preprocess allocates speculative c{4,128} KV via
+        # alloc_extend(backup_state=True) and rolls it back with restore_state.
+        # The base SWATokenToKVPoolAllocator only snapshots the full + SWA pools,
+        # so without this override the draft's c{4,128} (+ state) slots are never
+        # rolled back -> they leak every draft step until the c4 pool exhausts.
+        # Snapshot the sub-allocators alongside the base pools.
+        return (
+            super().backup_state(),
+            self.c4_attn_allocator.backup_state(),
+            self.c128_attn_allocator.backup_state(),
+            (
+                self.c4_state_attn_allocator.backup_state()
+                if self.c4_state_attn_allocator is not None
+                else None
+            ),
+            (
+                self.c128_state_attn_allocator.backup_state()
+                if self.c128_state_attn_allocator is not None
+                else None
+            ),
+        )
+
+    def restore_state(self, state):
+        base, c4, c128, c4_state, c128_state = state
+        super().restore_state(base)
+        self.c4_attn_allocator.restore_state(c4)
+        self.c128_attn_allocator.restore_state(c128)
+        if self.c4_state_attn_allocator is not None and c4_state is not None:
+            self.c4_state_attn_allocator.restore_state(c4_state)
+        if self.c128_state_attn_allocator is not None and c128_state is not None:
+            self.c128_state_attn_allocator.restore_state(c128_state)
+
+    def clear(self):
+        super().clear()
+        # super().__init__ calls clear() before our sub-allocators exist;
+        # getattr(..., None) tolerates that and the always-None state allocators.
+        for attr in (
+            "c4_attn_allocator",
+            "c128_attn_allocator",
+            "c4_state_attn_allocator",
+            "c128_state_attn_allocator",
+        ):
+            allocator = getattr(self, attr, None)
+            if allocator is not None:
+                allocator.clear()

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py

@@ -0,0 +1,404 @@
+"""Helpers used by mem_cache/common.py to wire DSV4-NPU per-req tables.
+
+mem_cache/common.py runs platform-agnostic alloc flow. When the model is
+DSV4 on NPU, ``alloc_paged_token_slots_{extend,decode}`` already stashed the
+:class:`DSV4OutCacheLoc` the allocator returned onto
+``batch.out_cache_loc_dsv4``. After each ``alloc_extend`` / ``alloc_decode``
+these hooks then:
+
+  1. Read the bundle from ``batch.out_cache_loc_dsv4``.
+  2. Write the per-pool slot ids into the per-req tables on the
+     :class:`DSV4NPUReqToTokenPool`.
+
+Non-DSV4 paths leave ``batch.out_cache_loc_dsv4`` None, so this module is a
+no-op for them.
+
+TODO: the disagg DSV4 path bypasses these hooks — it calls
+``allocator.alloc_extend`` directly then ``req_to_token_pool.write`` without
+going through ``mem_cache/common.py`` (see ``disaggregation/decode.py``). The
+DSV4OutCacheLoc bundle is still produced but never written into the per-req
+tables, so disagg + DSV4 is unsupported here (c-pages leak). Fixing requires
+calling these hooks from disagg's per-req alloc loop, or moving the write
+into the allocator itself.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+
+
+def maybe_write_dsv4_extend(
+    batch: ScheduleBatch,
+    req_pool_indices_cpu: torch.Tensor,
+    prefix_lens_cpu: torch.Tensor,
+    seq_lens_cpu: torch.Tensor,
+    *,
+    c4_state_alloc_offsets: Sequence[int] | torch.Tensor | None = None,
+    c128_state_alloc_offsets: Sequence[int] | torch.Tensor | None = None,
+) -> None:
+    """Post-alloc_extend hook for DSV4. No-op when allocator/pool is not DSV4.
+
+    For each compressed pool (c4 / c128), spreads the flat
+    ``out_c{4,128}_loc`` tensor across requests using per-req extend
+    counts (``seq_lens[i] // ratio - prefix_lens[i] // ratio``) and writes
+    the resulting slot ids into ``req_to_token_c{4,128}[req, prefix:seq]``.
+
+    Also writes ``req_to_token_swa[req, prefix:seq]`` with the swa slots
+    derived from out_full_loc via the SWA index mapping.
+    """
+    # Bundle stashed on batch.out_cache_loc_dsv4 by mem_cache/common.py;
+    # None on CUDA / non-V4 paths → no-op.
+    bundle = batch.out_cache_loc_dsv4
+    if bundle is None:
+        return
+
+    req_to_token_pool = batch.req_to_token_pool
+    if not hasattr(req_to_token_pool, "write_c4"):
+        return  # non-DSV4 pool; skip defensively (shouldn't happen)
+
+    # SWA writes: prefix..seq token positions, one slot per raw token.
+    _write_per_req_slice(
+        req_to_token_pool.write_swa,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_swa_loc,
+        ratio=1,
+    )
+
+    # c4 / c128 writes: prefix//ratio .. seq//ratio compressed positions.
+    _write_per_req_slice(
+        req_to_token_pool.write_c4,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c4_loc,
+        ratio=4,
+    )
+    _write_per_req_slice(
+        req_to_token_pool.write_c128,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c128_loc,
+        ratio=128,
+    )
+
+    # c4_state / c128_state writes: tail-only. Bundle length is
+    # sum(c{N}_state_alloc_len_i), NOT total raw extend tokens. Normal extend
+    # uses the per-Req low-water marks; reserve callers can pass explicit raw
+    # offsets for the pre-reserved interval.
+    if c4_state_alloc_offsets is None:
+        c4_state_alloc_offsets = [
+            getattr(r, "c4_state_alloc_offset", 0) for r in batch.reqs
+        ]
+    if c128_state_alloc_offsets is None:
+        c128_state_alloc_offsets = [
+            getattr(r, "c128_state_alloc_offset", 0) for r in batch.reqs
+        ]
+    if bundle.out_c4_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c4_state"
+    ):
+        _write_state_tail_per_req(
+            req_to_token_pool.write_c4_state,
+            req_pool_indices_cpu,
+            c4_state_alloc_offsets,
+            seq_lens_cpu,
+            bundle.out_c4_state_loc,
+        )
+    if bundle.out_c128_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c128_state"
+    ):
+        _write_state_tail_per_req(
+            req_to_token_pool.write_c128_state,
+            req_pool_indices_cpu,
+            c128_state_alloc_offsets,
+            seq_lens_cpu,
+            bundle.out_c128_state_loc,
+        )
+
+
+def maybe_write_dsv4_decode(
+    batch: ScheduleBatch,
+    seq_lens_cpu: torch.Tensor,
+    token_per_req: int,
+) -> None:
+    """Post-alloc_decode hook for DSV4. Spreads the new token slot ids
+    (one per req for swa, gated by ratio boundary for c4/c128) into the
+    per-req tables on DSV4NPUReqToTokenPool.
+
+    ``seq_lens_cpu`` is the POST-decode seq len (already incremented by
+    ``token_per_req``); the new compressed tokens go at positions
+    ``[(old_seq) // ratio, (new_seq) // ratio)``.
+    """
+    # Bundle stashed on batch.out_cache_loc_dsv4 by mem_cache/common.py;
+    # None on CUDA / non-V4 paths → no-op.
+    bundle = batch.out_cache_loc_dsv4
+    if bundle is None:
+        return
+
+    req_to_token_pool = batch.req_to_token_pool
+    if not hasattr(req_to_token_pool, "write_c4"):
+        return
+
+    prefix_lens_cpu = (seq_lens_cpu - token_per_req).clamp(min=0)
+    req_pool_indices_cpu = batch.req_pool_indices.cpu()
+
+    _write_per_req_slice(
+        req_to_token_pool.write_swa,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_swa_loc,
+        ratio=1,
+    )
+    _write_per_req_slice(
+        req_to_token_pool.write_c4,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c4_loc,
+        ratio=4,
+    )
+    _write_per_req_slice(
+        req_to_token_pool.write_c128,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c128_loc,
+        ratio=128,
+    )
+
+    # State table decode writes: one slot per raw decode token (ratio=1).
+    if bundle.out_c4_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c4_state"
+    ):
+        _write_per_req_slice(
+            req_to_token_pool.write_c4_state,
+            req_pool_indices_cpu,
+            prefix_lens_cpu,
+            seq_lens_cpu,
+            bundle.out_c4_state_loc,
+            ratio=1,
+        )
+    if bundle.out_c128_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c128_state"
+    ):
+        _write_per_req_slice(
+            req_to_token_pool.write_c128_state,
+            req_pool_indices_cpu,
+            prefix_lens_cpu,
+            seq_lens_cpu,
+            bundle.out_c128_state_loc,
+            ratio=1,
+        )
+
+
+def maybe_build_dsv4_verify_bundle(batch: ScheduleBatch, draft_token_num: int):
+    """Build the DSV4 cache-location view for one target-verify pass.
+
+    Spec-v2 reserves cache ahead of time, so target verify must select only the
+    current draft interval from the per-request DSV4 tables instead of reusing
+    the larger allocation bundle produced during decode preparation.
+    """
+    pool = batch.req_to_token_pool
+    if not hasattr(pool, "req_to_token_c4"):
+        return None
+    reserve_bundle = batch.out_cache_loc_dsv4
+    if reserve_bundle is None:
+        return None
+
+    req_indices = batch.req_pool_indices_cpu.tolist()
+    seq_lens = batch.seq_lens_cpu.tolist()
+
+    def flatten_interval(table: torch.Tensor, ratio: int) -> torch.Tensor:
+        chunks = []
+        for req_idx, seq_len in zip(req_indices, seq_lens):
+            start = int(seq_len) // ratio
+            end = (int(seq_len) + draft_token_num) // ratio
+            if end > start:
+                chunks.append(table[int(req_idx), start:end])
+        return torch.cat(chunks) if chunks else table.new_empty((0,))
+
+    return type(reserve_bundle)(
+        out_full_loc=batch.out_cache_loc,
+        out_swa_loc=flatten_interval(pool.req_to_token_swa, 1),
+        out_c4_loc=flatten_interval(pool.req_to_token_c4, 4),
+        out_c128_loc=flatten_interval(pool.req_to_token_c128, 128),
+        out_c4_state_loc=flatten_interval(pool.req_to_token_c4_state, 1),
+        out_c128_state_loc=flatten_interval(pool.req_to_token_c128_state, 1),
+    )
+
+
+def _write_per_req(
+    write_fn,
+    req_pool_indices_cpu: torch.Tensor,
+    flat_loc: torch.Tensor,
+    bounds_fn,
+) -> None:
+    """Distribute a flat ``[total_alloc]`` slot tensor across reqs.
+
+    ``bounds_fn(i) -> (lo, hi)`` gives req i's write window; the matching
+    ``hi - lo`` slots are sliced off ``flat_loc`` in order and written via
+    ``write_fn((req_idx, slice(lo, hi)), values)``. flat_loc may be None /
+    empty when the alloc path bypassed DSV4NPUTokenToKVPoolAllocator (e.g.
+    page_size=1 or HiSparse wrapper); skip then.
+    """
+    if flat_loc is None or flat_loc.numel() == 0:
+        return
+    pt = 0
+    for i in range(req_pool_indices_cpu.shape[0]):
+        lo, hi = bounds_fn(i)
+        alloc_len = max(0, hi - lo)
+        if alloc_len == 0:
+            continue
+        req_idx = int(req_pool_indices_cpu[i].item())
+        chunk = flat_loc[pt : pt + alloc_len].to(torch.int32)
+        write_fn((req_idx, slice(lo, hi)), chunk)
+        pt += alloc_len
+
+
+def _write_state_tail_per_req(
+    write_fn,
+    req_pool_indices_cpu: torch.Tensor,
+    state_alloc_offsets: list,
+    seq_lens_cpu: torch.Tensor,
+    flat_loc: torch.Tensor,
+) -> None:
+    """Tail-only state write: req i's slots go at ``[state_alloc_offsets[i],
+    seq_lens[i])`` in ``req_to_token_c{N}_state``."""
+    _write_per_req(
+        write_fn,
+        req_pool_indices_cpu,
+        flat_loc,
+        lambda i: (int(state_alloc_offsets[i]), int(seq_lens_cpu[i].item())),
+    )
+
+
+def _write_per_req_slice(
+    write_fn,
+    req_pool_indices_cpu: torch.Tensor,
+    prefix_lens_cpu: torch.Tensor,
+    seq_lens_cpu: torch.Tensor,
+    flat_loc: torch.Tensor,
+    ratio: int,
+) -> None:
+    """Compressed-position write: req i's slots go at
+    ``[prefix_lens[i] // ratio, seq_lens[i] // ratio)``."""
+    _write_per_req(
+        write_fn,
+        req_pool_indices_cpu,
+        flat_loc,
+        lambda i: (
+            int(prefix_lens_cpu[i].item()) // ratio,
+            int(seq_lens_cpu[i].item()) // ratio,
+        ),
+    )
+
+
+def maybe_evict_dsv4_state(batch: ScheduleBatch, req: Req, pre_len: int) -> None:
+    """Per-decode evict for the DSV4-NPU compress-state pools, independent of
+    SWA evict cadence. Called every decode step from ``ScheduleBatch``.
+
+    The state pool is small (~2 pages c4 / ~3 pages c128 of raw positions per
+    req) — with a large sliding_window (SWA evict fires every
+    ``eviction_interval`` and needs ``pre_len > sliding_window + page_size`` to
+    free anything) the pool exhausts before the first SWA frontier advance, so
+    we drain it here on its own cadence.
+
+    Retention windows (kernel read window + decode lookahead margin):
+    c4 = 8 + 16, c128 = 128 + 64 raw positions — intentionally smaller than one
+    SWA page so the first eviction fires before the small pool fills. Watermarks
+    are page-aligned so freed slots are whole pages reclaimable by the paged
+    allocator. ``req.c{4,128}_state_alloc_offset`` (read/written via getattr/
+    setattr) is the low-water mark. No-op on non-DSV4-NPU paths.
+    """
+    allocator = batch.token_to_kv_pool_allocator
+    pool = batch.req_to_token_pool
+    if not hasattr(allocator, "c4_state_attn_allocator") or (
+        allocator.c4_state_attn_allocator is None
+        and allocator.c128_state_attn_allocator is None
+    ):
+        return
+
+    page_size = batch.tree_cache.page_size
+    c4_watermark = ((max(0, pre_len - (8 + 16))) // page_size) * page_size
+    c128_watermark = ((max(0, pre_len - (128 + 64))) // page_size) * page_size
+
+    _free_state_range(
+        allocator.c4_state_attn_allocator,
+        pool,
+        "req_to_token_c4_state",
+        req,
+        "c4_state_alloc_offset",
+        c4_watermark,
+    )
+    _free_state_range(
+        allocator.c128_state_attn_allocator,
+        pool,
+        "req_to_token_c128_state",
+        req,
+        "c128_state_alloc_offset",
+        c128_watermark,
+    )
+
+
+def maybe_evict_dsv4_state_on_swa(
+    allocator, pool, req: Req, new_swa_evicted_seqlen: int
+) -> None:
+    """Free compress-state slots that ride along with SWA eviction.
+
+    State at raw positions < ``swa_evicted_seqlen`` is no longer readable (the
+    compressor only reads the trailing ``2*ratio`` window) and is returned to
+    its paged allocator to keep the small state pool from exhausting on long
+    generations. No-op when the DSV4-NPU state allocators are absent.
+
+    This path is needed for small-sliding-window models where
+    ``sliding_window < retention`` (e.g. c128 retention 192 > window 128):
+    in that case the watermark-based eviction alone may not free slots
+    fast enough, and the SWA-ride eviction is the primary reclaim mechanism.
+    For typical large-window models (DS-V4 with window >> 192), the
+    watermark eviction always runs first, making this path a no-op.
+    """
+    if not hasattr(allocator, "c4_state_attn_allocator"):
+        return
+    _free_state_range(
+        allocator.c4_state_attn_allocator,
+        pool,
+        "req_to_token_c4_state",
+        req,
+        "c4_state_alloc_offset",
+        new_swa_evicted_seqlen,
+    )
+    _free_state_range(
+        allocator.c128_state_attn_allocator,
+        pool,
+        "req_to_token_c128_state",
+        req,
+        "c128_state_alloc_offset",
+        new_swa_evicted_seqlen,
+    )
+
+
+def _free_state_range(
+    state_allocator,
+    pool,
+    table_attr: str,
+    req: Req,
+    offset_attr: str,
+    watermark: int,
+) -> None:
+    """Free ``[alloc_offset, watermark)`` raw-position state slots for ``req``
+    and advance its low-water mark. No-op when the allocator/table is absent or
+    the watermark hasn't advanced past the current offset."""
+    offset = getattr(req, offset_attr, 0)
+    if state_allocator is None or not hasattr(pool, table_attr) or watermark <= offset:
+        return
+    free_slots = getattr(pool, table_attr)[req.req_pool_idx, offset:watermark]
+    state_allocator.free(free_slots.to(torch.int64))
+    setattr(req, offset_attr, watermark)

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py

@@ -0,0 +1,593 @@
+"""NPU-only KV pool variant for DeepSeek-V4.
+
+Subclasses :class:`DeepSeekV4TokenToKVPool` to swap the ring-buffered
+:class:`CompressStatePool` for the paged :class:`NPUCompressStatePool` that
+the on-NPU fused compressor kernel (``torch.ops.custom.compressor`` with
+``cache_mode=1``) requires. Atlas A3 rejects ``cache_mode=2`` (ring) entirely,
+so this is the only valid layout on that hardware.
+
+Selected at pool construction time by
+:meth:`ModelRunnerKVCacheMixin._init_pools` when the model is DSV4 AND the
+device is NPU. CUDA continues to use the unchanged base class.
+
+The subclass overrides only:
+
+  * ``_make_attn_state_pool`` / ``_make_indexer_state_pool`` — the per-ratio
+    state-pool factories the base ``_init_paged_compress_states`` loop calls.
+    Both return :class:`NPUCompressStatePool` (paged, ``cache_mode=1``)
+    instead of the base's ring-buffered :class:`CompressStatePool`.
+  * ``translate_kv_loc_to_compress_state_loc`` — raise loudly. The ring
+    hash this method implements is meaningless on the paged kernel; callers
+    must consume ``out_cache_loc_dsv4.out_c{4,128}_state_loc`` from the
+    allocator bundle instead. Currently the only NPU caller that still
+    invokes translate is the unfused Python compressor decode path
+    (``layers/attention/dsv4/compressor.py``); with USE_FUSED_COMPRESSOR=1
+    that path is dead. If someone disables the fused compressor, they hit
+    the raise with a clear message.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch_npu
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+    ONLINE_C128,
+    DeepSeekV4IndexerPool,
+    DeepSeekV4SingleKVPool,
+    DeepSeekV4TokenToKVPool,
+)
+
+
+class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool):
+    """NPU bf16 variant of the full / SWA / c4 / c128 single-KV pool.
+
+    ``npu_sparse_attn_sharedkv`` reads KV in PA_ND layout
+    ``(num_pages, kernel_page_size, num_kv_heads=1, dim)`` with ``dim`` packing
+    K_nope + K_rope as bf16, and requires ``cmp_kv.shape[1] == ori_kv.shape[1]``.
+    So the c4/c128 pools (whose token-level page_size is ``page_size // ratio``)
+    are allocated at the GLOBAL ``kernel_page_size`` rather than their own
+    per-ratio page_size; the SWA pool uses ``kernel_page_size == page_size``.
+    The CUDA fp8-packed-bytes layout (the base ``create_buffer``) is untouched.
+    """
+
+    def __init__(self, *args, kernel_page_size: int, **kwargs):
+        # Set before super().__init__ — it calls _create_buffers() ->
+        # create_buffer(), which reads self.kernel_page_size.
+        self.kernel_page_size = kernel_page_size
+        super().__init__(*args, **kwargs)
+
+    def create_buffer(self, *, num_pages: int):
+        # Non-bf16 store dtype (shouldn't happen here) falls back to base layout.
+        if self.store_dtype != torch.bfloat16:
+            return super().create_buffer(num_pages=num_pages)
+        kv_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.kv_cache_total_dim = kv_dim
+        # GLOBAL kernel_page_size keeps cmp_kv.shape[1] == ori_kv.shape[1]; writes
+        # are flat-indexed by loc, so page granularity affects shape not location.
+        npu_num_pages = (self.size + self.kernel_page_size + 1) // self.kernel_page_size
+        return torch.zeros(
+            npu_num_pages,
+            self.kernel_page_size,
+            1,
+            kv_dim,
+            dtype=torch.bfloat16,
+            device=self.device,
+        )
+
+
+def npu_state_pool_size(
+    *,
+    ratio: int,
+    page_size: int,
+    max_num_reqs: int,
+) -> int:
+    """Per-pool state slot count for the NPU paged state pool's
+    :class:`NPUPagedTokenToKVPoolAllocator`.
+
+    Sizing formula::
+
+        max(2, ceil(1.8 * ratio / page_size) + 1) * max_num_reqs * page_size
+
+    Sized for steady-state during decode: each req keeps roughly the trailing
+    ``sliding_window_size`` worth of state slots live at any one time (SWA
+    eviction in :meth:`ScheduleBatch._evict_swa` frees state slots as it
+    advances), and the 1.8x factor adds headroom for the tail-only allocation
+    pattern across page boundaries.
+
+    Prefill no longer drives sizing because allocation is tail-only — long
+    prompts only allocate ``c{ratio}_alloc_len`` slots (``≤ tail + 128`` for
+    c4, ``≤ tail`` for c128, where ``tail = seq_len % 128``), not the full raw
+    seqlen. See :meth:`ScheduleBatch._compute_dsv4_state_lens_extend` for the
+    per-req formula.
+
+    Result is in TOKEN units (matches the SGLang allocator
+    ``PagedTokenToKVPoolAllocator(size, ...)`` convention where
+    ``num_pages = size // page_size`` is the count of USABLE pages handed out
+    by ``free_pages = arange(1, num_pages+1)``). The BUFFER allocates one extra
+    page (see :class:`NPUCompressStatePool`, sized ``(num_pages + 1) *
+    page_size`` — page 0 is the kernel's skip-sentinel).
+    """
+    blocks_per_req = max(2, math.ceil(1.8 * ratio / page_size) + 1)
+    num_usable_pages = blocks_per_req * max_num_reqs
+    return num_usable_pages * page_size
+
+
+class NPUCompressStatePool(CompressStatePool):
+    """Paged compress-state pool for the NPU fused compressor kernel.
+
+    ``torch.ops.custom.compressor`` (cache_mode=1) reads/writes the compress
+    state via ``state_cache`` shape ``(block_num, page_size, 2*coff*head_dim)``
+    indexed by a paged ``state_block_table`` (block ids from 1; value 0 means
+    "skip this slot"). The CUDA :class:`CompressStatePool` sizes itself
+    ring-style, which misaddresses slots under cache_mode=1 (ring is also
+    unsupported on Atlas A3). This subclass keeps the parent's buffer layout
+    (``(self._size, 2*coff*head_dim)`` flat; ``state_cache_3d`` reshapes to
+    ``(num_blocks, page_size, 2*coff*head_dim)``) but replaces the size formula
+    with a paged one derived from ``max_num_reqs``. Block 0 is reserved as the
+    kernel's skip-sentinel (zero kv / -inf score) so any ``state_block_table``
+    entry defaulting to 0 lands in a deterministic, attention-neutral place.
+
+    NPU-only; CUDA keeps using the unchanged :class:`CompressStatePool`.
+    """
+
+    def __init__(
+        self,
+        *,
+        size: int,
+        overlap: bool,
+        head_dim: int,
+        dtype: torch.dtype,
+        device: str,
+        enable_memory_saver: bool,
+        ratio: int,
+        page_size: int,
+    ):
+        # Bypass parent __init__ — its ring-based sizing is incompatible with the
+        # kernel's paged block-id contract. We redo buffer alloc and set the same
+        # fields so the parent API (state_cache_3d, kv_score_buffer) stays intact.
+        assert ratio in (
+            4,
+            128,
+        ), f"NPUCompressStatePool only supports ratio in (4, 128); got {ratio}"
+        assert page_size > 1, (
+            "NPUCompressStatePool requires page_size>1 (kernel's "
+            "state_cache_3d view is (block_num, page_size, slot_dim)). "
+            "Got page_size=%d." % page_size
+        )
+
+        # ``size`` is the ALLOCATOR's size (npu_state_pool_size output). Buffer
+        # needs one EXTRA page so the free list arange(1, num_pages+1) indexes it
+        # without OOB (page 0 = skip sentinel; pages 1..num_pages handed out).
+        num_usable_pages = (size + page_size - 1) // page_size
+        num_buffer_pages = num_usable_pages + 1
+        self._size = num_buffer_pages * page_size
+        self.page_size = page_size
+        # ring_size=0 marks "not ring-buffered" (paged allocator replaces the
+        # parent's ring hashing); kept so downstream hasattr probes don't break.
+        self.ring_size = 0
+        # online compress is a CUDA-only opt with no NPU fused-compressor support;
+        # force off so layout matches kernel expectations.
+        self.online = False
+
+        # Slot dim = 2 * coff * head_dim = [kv | score]; coff = 1 (no overlap) or
+        # 2 (overlap). Matches CompressStatePool non-online layout.
+        self.last_dim = 2 * (1 + int(overlap)) * head_dim
+
+        # Reuse parent's buffer-alloc helper; only self._size differs from the
+        # ring-based parent path.
+        self._alloc_kv_score_buffer(
+            dtype=dtype, device=device, enable_memory_saver=enable_memory_saver
+        )
+
+        # Block 0 = kernel skip-sentinel: kv zeroed, score -inf (softmax → 0).
+        # The free list excludes it; only stale state_block_table entries land here.
+        self.kv_score_buffer.kv[:page_size].zero_()
+        self.kv_score_buffer.score[:page_size].fill_(float("-inf"))
+
+
+class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool):
+    """NPU c4-indexer pool. Keeps the base packed CUDA buffer (read by
+    get_contiguous_buf_infos / NSA) and ADDS dedicated int8 K + float16 scale
+    buffers in PA_ND layout at the global ``kernel_page_size``, written by
+    ``torch_npu.npu_scatter_nd_update_`` and read by
+    ``torch.ops.custom.npu_quant_lightning_indexer``.
+    """
+
+    def __init__(self, *args, kernel_page_size: int, **kwargs):
+        # Set before super().__init__ — it calls _create_buffer().
+        self._kernel_page_size = kernel_page_size
+        super().__init__(*args, **kwargs)
+
+    def _create_buffer(self):
+        # Base allocates the packed CUDA index_k_with_scale_buffer (kept for
+        # get_contiguous_buf_infos / NSA compat); then add the NPU buffers.
+        super()._create_buffer()
+        kp = self._kernel_page_size
+        npu_num_pages = (self.size + kp + 1) // kp
+        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            self.index_k_buffer = [
+                torch.zeros(
+                    npu_num_pages,
+                    kp,
+                    1,
+                    self.index_head_dim,
+                    dtype=torch.int8,
+                    device=self.device,
+                )
+                for _ in range(self.layer_num)
+            ]
+            self.index_scale_buffer = [
+                torch.zeros(
+                    npu_num_pages,
+                    kp,
+                    1,
+                    1,
+                    dtype=torch.float16,
+                    device=self.device,
+                )
+                for _ in range(self.layer_num)
+            ]
+
+    @property
+    def has_npu_storage(self) -> bool:
+        return True
+
+    def get_index_k(self, layer_id: int) -> torch.Tensor:
+        return self.index_k_buffer[layer_id]
+
+    def get_index_scale(self, layer_id: int) -> torch.Tensor:
+        return self.index_scale_buffer[layer_id]
+
+    def set_index_k_scale(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        index_k: torch.Tensor,
+        index_k_scale: Optional[torch.Tensor],
+    ) -> None:
+        # int8 K + fp16 scale come from _compressor_epilog_npu's npu_dynamic_quant
+        # output (index_k: int8 [T, D], index_k_scale: fp16 [T, 1]).
+        d = self.index_head_dim
+        loc_long = loc.view(-1, 1).long()
+        torch_npu.npu_scatter_nd_update_(
+            self.index_k_buffer[layer_id].view(-1, 1, d),
+            loc_long,
+            index_k.to(torch.int8).view(-1, 1, d),
+        )
+        if index_k_scale is not None:
+            torch_npu.npu_scatter_nd_update_(
+                self.index_scale_buffer[layer_id].view(-1, 1, 1),
+                loc_long,
+                index_k_scale.to(torch.float16).view(-1, 1, 1),
+            )
+
+
+class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
+    """NPU-only DSV4 KV pool with paged compress-state buffers.
+
+    The full / SWA / c4 / c128 KV pools use the NPU bf16 PA_ND layout
+    (:class:`NPUDeepSeekV4SingleKVPool`); the compress-state pool is paged
+    (:class:`NPUCompressStatePool`) rather than ring-buffered; and the indexer
+    pool adds dedicated int8 K + fp16 scale buffers
+    (:class:`NPUDeepSeekV4IndexerPool`). The generic-accessor / port-hook
+    methods at the bottom of this class are the NPU equivalents of the CUDA
+    DSV4 store-cache chain — kept here, not in the community base, which raises
+    ``NotImplementedError`` for them (CUDA goes through the radix / store_cache
+    accessors instead).
+    """
+
+    def _make_kv_pool(
+        self,
+        *,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+        global_page_size: int,
+        cls: type = DeepSeekV4SingleKVPool,
+    ) -> NPUDeepSeekV4SingleKVPool:
+        # NPU does not use the HiSparse c4 device pool; fail loud if someone
+        # enables it so the silent layout mismatch surfaces at init.
+        assert cls is DeepSeekV4SingleKVPool, (
+            "enable_hisparse is not supported on the NPU DSV4 KV pool "
+            f"(got c4 pool class {cls.__name__})."
+        )
+        return NPUDeepSeekV4SingleKVPool(
+            size,
+            page_size,
+            dtype,
+            self.qk_nope_head_dim,
+            self.qk_rope_head_dim,
+            layer_num,
+            device,
+            enable_memory_saver,
+            kernel_page_size=global_page_size,
+        )
+
+    def _get_state_pool(self, layer_id: int, from_indexer: bool) -> CompressStatePool:
+        """Select this layer's attention vs c4-indexer compress-state pool.
+        Wraps the community getters so the NPU port hooks below don't index the
+        pool lists directly."""
+        if from_indexer:
+            return self.get_indexer_compress_states(layer_id)
+        return self.get_attention_compress_states(layer_id)
+
+    def _make_attn_state_pool(
+        self, ratio: int, enable_memory_saver: bool
+    ) -> NPUCompressStatePool:
+        # ONLINE_C128 (CUDA-only) collapses the c128 ring to size 1; the NPU fused
+        # compressor has no online mode, so assert the config mismatch early.
+        assert not (ratio == 128 and ONLINE_C128), (
+            "SGLANG_OPT_USE_ONLINE_COMPRESS is incompatible with the "
+            "NPU fused compressor (no online mode in the kernel)."
+        )
+        return NPUCompressStatePool(
+            size=self._state_pool_size(ratio),
+            overlap=ratio == 4,
+            head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
+            dtype=self.c4_state_dtype if ratio == 4 else self.c128_state_dtype,
+            device=self.device,
+            enable_memory_saver=enable_memory_saver,
+            ratio=ratio,
+            page_size=self.swa_page_size,
+        )
+
+    def _make_indexer_state_pool(
+        self, ratio: int, enable_memory_saver: bool
+    ) -> NPUCompressStatePool:
+        # c4 indexer shares the c4 state pool size budget but has its own
+        # slot_dim (indexer_head_dim vs attention head_dim).
+        return NPUCompressStatePool(
+            size=self.c4_state_pool_size,
+            overlap=ratio == 4,
+            head_dim=self.indexer_head_dim,
+            device=self.device,
+            dtype=self.c4_state_dtype,
+            enable_memory_saver=enable_memory_saver,
+            ratio=ratio,
+            page_size=self.swa_page_size,
+        )
+
+    def clear_unaccepted_c128_draft_states(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        accept_lens: torch.Tensor,
+        num_draft_tokens: int,
+    ) -> None:
+        pass
+
+    def _make_indexer_pool(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        index_head_dim: int,
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+    ) -> NPUDeepSeekV4IndexerPool:
+        # NPU dedicated int8 K + fp16 scale buffers use the GLOBAL page_size
+        # (= self.page_size) as kernel_page_size, matching ori_kv for the kernel.
+        return NPUDeepSeekV4IndexerPool(
+            size,
+            page_size,
+            dtype,
+            index_head_dim,
+            layer_num,
+            device,
+            enable_memory_saver,
+            kernel_page_size=self.page_size,
+        )
+
+    def get_state_cache(self, layer_id: int, from_indexer: bool) -> torch.Tensor:
+        """fp32 ``[block_num, page_size, 2*coff*D]`` view of this layer's
+        kv+score buffer — the fused compressor op
+        (``torch.ops.custom.compressor``)'s ``state_cache`` argument."""
+        return self._get_state_pool(layer_id, from_indexer).state_cache_3d
+
+    # ------------------------------------------------------------------
+    # Generic KV accessors (community base raises NotImplementedError; CUDA uses
+    # store_cache). AscendAttnBackend reads KV through these, routed to the right
+    # sub-pool by compression ratio.
+    # ------------------------------------------------------------------
+
+    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
+        item = self.layer_mapping[layer_id]
+        ratio = item.compress_ratio
+        if ratio == 0:
+            return self.swa_kv_pool.kv_buffer[item.compress_layer_id]
+        if ratio == 4:
+            return self.c4_kv_pool.kv_buffer[item.compress_layer_id]
+        if ratio == 128:
+            return self.c128_kv_pool.kv_buffer[item.compress_layer_id]
+        raise ValueError(f"unsupported compress_ratio={ratio} for get_key_buffer")
+
+    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
+        # V4 uses MQA / latent attention — the K buffer doubles as V.
+        return self.get_key_buffer(layer_id)
+
+    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        buf = self.get_key_buffer(layer_id)
+        return buf, buf
+
+    def get_swa_buffer(
+        self, layer_id: int, loc: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Return the SWA layer's KV cache in PA_ND layout
+        (num_pages, page_size, num_kv_heads=1, dim). When ``loc`` is given,
+        flatten across (num_pages, page_size) and gather the matching tokens —
+        shape becomes (num_tokens, 1, dim).
+        """
+        # Index by RAW layer_id, not compress_layer_id (a per-bucket counter that
+        # would collide across ratios). swa_kv_pool is sized layer_num=total_layers.
+        kv = self.swa_kv_pool.kv_buffer[layer_id]
+        if loc is not None:
+            kv = kv.flatten(0, 1)[loc]
+        return kv
+
+    def get_compress_buffer(
+        self,
+        layer_id: int,
+        from_indexer: bool = False,
+        loc: Optional[torch.Tensor] = None,
+    ) -> Optional[torch.Tensor]:
+        """Return the compressed KV buffer for a c4 / c128 layer.
+
+        Routes to c4 / c128 kv_pool by layer compression ratio. Returns
+        ``None`` for ratio == 0 (no compress KV exists). The
+        from_indexer=True branch returns the dedicated int8 K buffer that
+        ``torch.ops.custom.npu_quant_lightning_indexer`` consumes.
+        """
+        item = self.layer_mapping[layer_id]
+        if item.compress_ratio == 4:
+            if from_indexer:
+                kv = self.c4_indexer_kv_pool.get_index_k(item.compress_layer_id)
+            else:
+                kv = self.c4_kv_pool.kv_buffer[item.compress_layer_id]
+        elif item.compress_ratio == 128:
+            assert not from_indexer, "c128 has no indexer pool"
+            kv = self.c128_kv_pool.kv_buffer[item.compress_layer_id]
+        else:
+            return None
+        if loc is not None:
+            kv = kv.flatten(0, 1)[loc]
+        return kv
+
+    def set_swa_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        cache: torch.Tensor,
+    ) -> None:
+        """Write ``cache`` into the SWA pool at flat token positions ``loc``.
+
+        ``cache`` shape: (num_tokens, num_kv_heads=1, dim). The buffer view is
+        (num_pages, page_size, 1, dim) so we flatten the first two dims and
+        index_put.
+        """
+        # Index by raw layer_id (see get_swa_buffer) to avoid bucket collision.
+        buf = self.swa_kv_pool.kv_buffer[layer_id]
+        buf_flat = buf.flatten(0, 1)  # (num_pages * page_size, 1, dim)
+        # Caller (V4 MQALayer) may hand us cache shaped (T, dim); the buffer has
+        # an explicit num_kv_heads=1 axis, so insert it.
+        if cache.ndim == buf_flat.ndim - 1:
+            cache = cache.unsqueeze(1)
+        buf_flat[loc] = cache.to(buf_flat.dtype)
+
+    # ------------------------------------------------------------------
+    # NPU port hooks — used by dsv4/{compressor,indexer}.py forward_npu.
+    # CompressStatePool stores a fused [kv | score] tensor; split is a last-dim slice.
+    # ------------------------------------------------------------------
+
+    def set_state_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        kv: torch.Tensor,
+        score: torch.Tensor,
+        from_indexer: bool,
+    ) -> None:
+        # KVAndScore.kv_score is [..., 2*coff*head_dim] = [kv | score].
+        kv_score = self._get_state_pool(layer_id, from_indexer).kv_score_buffer.kv_score
+        last_dim = kv_score.shape[-1]
+        half = last_dim // 2
+        kv_view = kv.reshape(-1, half).to(kv_score.dtype)
+        score_view = score.reshape(-1, half).to(kv_score.dtype)
+        kv_score[loc, :half] = kv_view
+        kv_score[loc, half:] = score_view
+
+    def get_state_buffer(
+        self,
+        layer_id: int,
+        from_indexer: bool,
+        kv_indices: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        kv_score = self._get_state_pool(layer_id, from_indexer).kv_score_buffer.kv_score
+        if kv_indices is not None:
+            kv_score = kv_score[kv_indices]
+        last_dim = kv_score.shape[-1]
+        half = last_dim // 2
+        kv = kv_score[..., :half].unsqueeze(-2)  # add num_kv_heads=1 axis
+        score = kv_score[..., half:].unsqueeze(-2)
+        return kv, score
+
+    def set_compress_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        kv: torch.Tensor,
+        kv_scale: Optional[torch.Tensor],
+        from_indexer: bool,
+    ) -> None:
+        # Routes to c4_indexer (from_indexer) / c4_kv (ratio 4) / c128_kv (ratio
+        # 128). NPU bypasses CUDA fused_store_cache with direct bf16 writes.
+        ratio, compress_layer_id, _ = self.layer_mapping[layer_id]
+        device_type = kv.device.type
+        if from_indexer:
+            assert ratio == 4, f"indexer only on c4 layers, got ratio={ratio}"
+            if device_type == "npu":
+                assert (
+                    self.c4_indexer_kv_pool.has_npu_storage
+                ), "NPU index buffers not allocated — pool was init'd on CUDA?"
+                self.c4_indexer_kv_pool.set_index_k_scale(
+                    compress_layer_id, loc, kv, kv_scale
+                )
+                return
+            if kv_scale is None:
+                self.c4_indexer_kv_pool.set_index_fused(compress_layer_id, loc, kv)
+                return
+            self.c4_indexer_kv_pool.set_index_k_scale_buffer(
+                compress_layer_id, loc, kv, kv_scale
+            )
+            return
+        compress_pool = self.c4_kv_pool if ratio == 4 else self.c128_kv_pool
+        if device_type == "npu":
+            # PA_ND layout: kv_buffer[layer_id] shape = (num_pages, page_size,
+            # 1, kv_dim). Flatten (num_pages, page_size) and index by `loc`.
+            buf = compress_pool.kv_buffer[compress_layer_id]
+            buf_flat = buf.flatten(0, 1)
+            kv_view = kv.to(buf_flat.dtype)
+            if kv_view.ndim == buf_flat.ndim - 1:
+                kv_view = kv_view.unsqueeze(1)
+            buf_flat[loc] = kv_view
+            return
+        compress_pool.set_key_buffer_fused(compress_layer_id, loc, kv)
+
+    def get_compress_dequant_scale_buffer(
+        self,
+        layer_id: int,
+        from_indexer: bool,
+    ) -> torch.Tensor:
+        # Returns the float16 dequant scale buffer (NPU indexer pool's dedicated
+        # scale buffer alongside the int8 K buffer).
+        assert from_indexer, "only indexer compress pool has dequant scale"
+        compress_layer_id = self.layer_mapping[layer_id].compress_layer_id
+        return self.c4_indexer_kv_pool.get_index_scale(compress_layer_id)
+
+    def translate_kv_loc_to_compress_state_loc(
+        self,
+        kv_loc: torch.Tensor,
+        compress_ratio: int,
+    ) -> torch.Tensor:
+        # Parent's ring-buffer hash is meaningless under the paged cache_mode=1
+        # contract; returning a stale value would silently corrupt state. Fail loud.
+        raise RuntimeError(
+            "DSV4NPUTokenToKVPool.translate_kv_loc_to_compress_state_loc was "
+            "called, but the NPU fused compressor kernel uses a paged state "
+            "pool (cache_mode=1) and does not support ring-buffer state "
+            "addressing (cache_mode=2 is explicitly unsupported on Atlas A3). "
+            "Callers must consume out_cache_loc_dsv4.out_c{4,128}_state_loc "
+            "from the allocator bundle (set during alloc_extend/alloc_decode) "
+            "and read state_page_table from req_to_token_c{4,128}_state on "
+            "the DSV4NPUReqToTokenPool instead. See "
+            "hardware_backend/npu/dsv4_memory_pool.py for the rationale."
+        )

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py

@@ -0,0 +1,128 @@
+"""DSV4-NPU per-request mapping pool.
+
+Subclass of ``ReqToTokenPool`` that adds five auxiliary per-request tables
+needed by the DSV4 attention backend:
+
+  * ``req_to_token_swa``        — slot ids in the SWA full-pool view
+  * ``req_to_token_c4``         — slot ids in the c4 compressed-KV pool
+  * ``req_to_token_c128``       — slot ids in the c128 compressed-KV pool
+  * ``req_to_token_c4_state``   — c4 state-pool slot ids, 1 per raw token
+  * ``req_to_token_c128_state`` — c128 state-pool slot ids, 1 per raw token
+
+Compressed KV pools store 1 slot per ``ratio`` raw tokens, so their per-req
+table column count is ``max_context_len // ratio``. swa mirrors the raw
+token count. Elements are token-level slot ids; the attention backend
+converts to page ids via ``// page_size`` when constructing PA_ND block
+tables.
+
+The c4/c128 STATE pools also have per-req tables here: the NPU fused
+compressor uses a paged state pool (``cache_mode=1``), so each raw token's
+state slot id is recorded (1 column per raw token) and the backend builds
+``state_block_table = req_to_token_c{N}_state[req, ::page_size] // page_size``
+to feed the kernel. (The base class' ``translate_kv_loc_to_compress_state_loc``
+ring-hash is the CUDA-only path; it is disabled on NPU.)
+
+Memory cost example (size=64, max_context_len=32K): swa 8MB + c4 2MB +
+c128 64KB ≈ 10MB extra on top of the base req_to_token (8MB).
+
+The tables are populated by the ``dsv4_common_hooks`` writers (driven from
+``mem_cache/common.py``) immediately after a successful alloc_extend /
+alloc_decode, using the per-pool slot indices returned in ``DSV4OutCacheLoc``.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+
+class DSV4NPUReqToTokenPool(ReqToTokenPool):
+    """ReqToTokenPool extended with DSV4 SWA + c4/c128 per-req tables.
+
+    Drop-in replacement for ReqToTokenPool when the model is DeepSeek-V4 on
+    NPU. Selected by ``model_runner_kv_cache_mixin`` based on model arch +
+    device. Non-DSV4 and non-NPU paths continue to use the base class.
+
+    The auxiliary tables are intentionally NOT zeroed on ``clear()``: they are
+    indexed only by active rows (via req_pool_idx) and only each row's
+    ``[:seq_len]`` prefix is read, so stale entries past kv_committed_len are
+    unreachable by the attention metadata builder.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        max_context_len: int,
+        device: str,
+        enable_memory_saver: bool,
+    ):
+        super().__init__(size, max_context_len, device, enable_memory_saver)
+
+        memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+        )
+
+        # Back-ref to DSV4NPUTokenToKVPoolAllocator, wired via
+        # register_dsv4_allocator after both exist, so free(req) can release
+        # c4/c128 pages. None at construction so base clear() runs safely.
+        self._dsv4_allocator = None
+
+        # (name, columns). swa + state tables: 1 slot per raw token; c4/c128:
+        # 1 slot per `ratio` raw tokens. Init zero so unallocated columns map to
+        # block 0 (kernel skip sentinel cleared by NPUCompressStatePool).
+        with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            for name, cols in (
+                ("req_to_token_swa", max_context_len),
+                ("req_to_token_c4", max(1, max_context_len // 4)),
+                ("req_to_token_c128", max(1, max_context_len // 128)),
+                ("req_to_token_c4_state", max_context_len),
+                ("req_to_token_c128_state", max_context_len),
+            ):
+                setattr(
+                    self,
+                    name,
+                    torch.zeros(
+                        (self._alloc_size, cols),
+                        dtype=torch.int32,
+                        device=device,
+                    ),
+                )
+
+    # ------------------------------------------------------------------
+    # Per-pool write helpers, called by mem_cache/common.py after alloc, using
+    # slot indices from DSV4OutCacheLoc. Args: (req_pool_idx, token_offset), slot.
+    # ------------------------------------------------------------------
+
+    def write_swa(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_swa[indices] = values
+
+    def write_c4(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c4[indices] = values
+
+    def write_c128(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c128[indices] = values
+
+    def write_c4_state(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c4_state[indices] = values
+
+    def write_c128_state(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c128_state[indices] = values
+
+    def register_dsv4_allocator(self, allocator) -> None:
+        """Wire the DSV4NPUTokenToKVPoolAllocator ref so ``free(req)`` can
+        release c4/c128 pool pages alongside the req_pool_idx slot. This is a
+        one-way ref (pool -> allocator). The reverse direction (the allocator
+        reading these per-req tables for its c-pool / state last_loc lookup) is
+        no longer a stored back-ref: mem_cache/common.py passes this pool into
+        ``alloc_extend`` / ``alloc_decode`` per call instead."""
+        self._dsv4_allocator = allocator
+
+    def free(self, req):
+        # Trigger c4/c128 free via the allocator's unified free path. May be None
+        # between __init__ and register_dsv4_allocator — defensive None check.
+        if self._dsv4_allocator is not None:
+            self._dsv4_allocator.free(req=req, req_to_token_pool=self)
+        super().free(req)

python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py

@@ -0,0 +1,49 @@
+# Copyright 2024-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.configs.model_config import is_deepseek_dsa
+from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
+    EAGLEDraftExtendCudaGraphRunner,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
+
+
+class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
+    def __init__(self, eagle_worker: EagleDraftWorker):
+        super().__init__(eagle_worker)
+
+    def _cache_loc_dtype(self):
+        return torch.int32
+
+    def _replay_graph(self, shape_key, forward_batch):
+        if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
+            seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
+                self.bs - self.raw_bs
+            )
+            return self.backend.replay_with_input_update(
+                shape_key,
+                seq_lens=seq_lens,
+                attr_name="actual_seq_lengths_kv",
+                attr_type=[],
+            )
+        else:
+            return self.backend.replay(shape_key, forward_batch)

python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py

@@ -0,0 +1,69 @@
+# Copyright 2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Union
+
+import torch
+
+from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
+from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
+    EAGLEDraftCudaGraphRunner,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
+
+
+class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
+    def __init__(self, eagle_worker: EagleDraftWorker):
+        self._init_arch_map()
+        super().__init__(eagle_worker)
+
+    def _init_arch_map(self):
+        self.attr_name: Dict[str, str] = {
+            AttentionArch.MLA: "actual_seq_lengths_kv",
+            AttentionArch.MHA: "context_lens",
+        }
+        self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
+            AttentionArch.MLA: [],
+            AttentionArch.MHA: torch.Tensor(),
+        }
+
+    def _cache_loc_dtype(self):
+        return torch.int32
+
+    def _get_update_attr_name(self):
+        return self.attr_name[AttentionArch.MLA]
+
+    def _get_update_attr_type(self):
+        return self.attr_type[AttentionArch.MLA]
+
+    def _replay_graph(self, shape_key, forward_batch):
+        if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
+            seq_lens_for_each_draft_step = []
+            for speculative_step_id in range(self.speculative_num_steps - 1):
+                seq_lens_cpu = (
+                    forward_batch.seq_lens_cpu[: self.raw_bs] + speculative_step_id + 1
+                )
+                seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
+                seq_lens_for_each_draft_step.append(seq_lens)
+            attr_name = self._get_update_attr_name()
+            cpu_update_input = [{attr_name: sl} for sl in seq_lens_for_each_draft_step]
+            return self.backend.replay_with_input_update(
+                shape_key, seq_lens=None, cpu_update_input=cpu_update_input
+            )
+        else:
+            return self.backend.replay(shape_key, forward_batch)

python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py

@@ -0,0 +1,54 @@
+# Copyright 2024-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
+from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
+    MultiLayerEagleDraftExtendCudaGraphRunner,
+    MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
+)
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    pass
+
+
+class MultiLayerEagleDraftExtendNpuGraphRunner(
+    MultiLayerEagleDraftExtendCudaGraphRunner
+):
+    def _replay_graph(self, shape_key, forward_batch):
+        seq_lens = self.buffers.seq_lens_cpu[: self.raw_bs].tolist() + [0] * (
+            self.bs - self.raw_bs
+        )
+        return self.backend.replay_with_input_update(
+            shape_key,
+            seq_lens=seq_lens,
+            attr_name="actual_seq_kvlen",
+            attr_type=[],
+        )
+
+
+class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
+    MultiLayerEagleMultiStepDraftExtendCudaGraphRunner
+):
+    def _create_runner(self, step: int) -> MultiLayerEagleDraftExtendNpuGraphRunner:
+        return MultiLayerEagleDraftExtendNpuGraphRunner(self.eagle_worker, step)
+
+    def _cuda_graph_disabled(self) -> bool:
+        return cuda_graph_fully_disabled()

python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py

@@ -0,0 +1,177 @@
+"""NPUCudaGraphBackend — Ascend NPU full-graph capture (torch.npu.NPUGraph).
+
+Mirrors FullCudaGraphBackend with two differences:
+  - Captures via torch.npu.graph(...) into torch.npu.NPUGraph.
+  - replay_with_input_update(shape_key, seq_lens, attr_name) rebinds
+    the recorded graph's input bindings for variable seq_lens at replay
+    time (NPU's NPUGraph.update(...) API).
+
+torch.npu is imported lazily inside methods so the module loads on
+non-NPU hosts.
+"""
+
+from __future__ import annotations
+
+import threading
+from contextlib import AbstractContextManager, contextmanager
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.model_executor.runner.shape_key import ShapeKey
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.utils import empty_context, get_bool_env_var
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+
+
+class NPUCudaGraphBackend(BaseCudaGraphBackend):
+    """One torch.npu.NPUGraph per shape; attention metadata captured
+    inside the graph. replay_with_input_update substitutes fresh
+    seq_lens without re-recording."""
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+        *,
+        enable_memory_saver: bool = False,
+    ) -> None:
+        self._graphs: Dict[Any, Any] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = cuda_graph_runner.model_runner.tp_group
+        self._capture_stream = None
+        self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+        self._enable_torch_compile = getattr(
+            cuda_graph_runner, "enable_torch_compile", False
+        )
+
+    @contextmanager
+    def capture_session(self, stream):
+        if self._pool is None:
+            self._pool = self._device_module.graph_pool_handle()
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        import torch_npu  # noqa: F401  (verifies NPU availability)
+
+        # Two warmups so kernels are loaded and one-time setup is paid before capture.
+        # post_warmup_hook lets the attention backend reset state that warmup mutated.
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = torch.npu.NPUGraph()
+
+        if self._enable_torch_compile:
+            skip_guard_context = torch.compiler.set_stance(skip_guard_eval_unsafe=True)
+        else:
+            skip_guard_context = empty_context()
+
+        graph_ctx: Callable[..., AbstractContextManager]
+        if (
+            self._memory_saver_adapter is not None
+            and self._memory_saver_adapter.enabled
+        ):
+            graph_ctx = partial(
+                self._memory_saver_adapter.cuda_graph,
+                tag=GPU_MEMORY_TYPE_CUDA_GRAPH,
+            )
+        else:
+            graph_ctx = torch.npu.graph
+
+        with skip_guard_context, graph_ctx(
+            graph,
+            pool=self._pool,
+            stream=self._capture_stream,
+            auto_dispatch_capture=True,
+        ):
+            out = forward_fn()
+
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def replay_with_input_update(
+        self,
+        shape_key: ShapeKey,
+        seq_lens: Any,
+        attr_name: str = None,
+        attr_type: Any = None,
+        cpu_update_input: list = None,
+    ) -> Any:
+        """Rebind seq_lens on the recorded NPU graph in a background
+        thread, then replay. Used when the model is not deepseek-nsa.
+
+        Two calling conventions:
+        1. (legacy) seq_lens + attr_name + attr_type:
+           Constructs cpu_update_input=[{attr_name: seq_lens}] internally.
+        2. cpu_update_input: A list of {attr_name: seq_lens} dicts,
+           one per speculative step.  Used by EAGLE draft runners.
+        """
+        if cpu_update_input is None:
+            if isinstance(attr_type, torch.Tensor):
+                seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
+            cpu_update_input = [{attr_name: seq_lens}]
+
+        graph = self._graphs[shape_key]
+
+        def _update():
+            graph.update(cpu_update_input=cpu_update_input)
+
+        thread = threading.Thread(target=_update)
+        thread.start()
+        graph.replay()
+        thread.join()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -0,0 +1,284 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Run the model with NPU graph and torch.compile.
+
+NPUGraphRunner is a thin subclass of DecodeCudaGraphRunner: the
+factory returns NPUCudaGraphBackend for NPU devices, so all
+capture/replay mechanics live in the backend. This class adds:
+  - NPU-specific patch_model monkey-patch for the decode-Full +
+    torch.compile path.
+  - Profile context override (NPU profiler emits to disk, not in-mem).
+  - Replay override that issues an async NPUGraph.update for
+    seq_lens before replay (skipped for deepseek-nsa).
+  - Smaller cache_loc dtype (int32 instead of int64).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, Optional, Union
+
+import numpy as np
+import torch
+
+from sglang.srt.configs.model_config import (
+    AttentionArch,
+    is_deepseek_dsa,
+    is_deepseek_v4,
+)
+from sglang.srt.distributed.parallel_state import GroupCoordinator
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+from sglang.srt.utils import (
+    empty_context,
+    get_bool_env_var,
+    get_compiler_backend,
+    is_npu,
+)
+
+is_npu = is_npu()
+
+if is_npu:
+    import torch_npu
+    from torch_npu.profiler import ProfilerActivity, profile
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+
+
+@contextmanager
+def patch_model_npu(
+    model: torch.nn.Module,
+    enable_compile: bool,
+    num_tokens: int,
+    tp_group: GroupCoordinator,
+):
+    if enable_compile:
+        backend = get_compiler_backend("npugraph_ex")
+        yield torch.compile(
+            torch.no_grad()(model.forward),
+            fullgraph=True,
+            dynamic=False,
+            backend=backend,
+        )
+    else:
+        yield model.forward
+
+
+class NPUGraphRunner(DecodeCudaGraphRunner):
+    """A NPUGraphRunner runs the forward pass of a model with NPU graph and torch.compile."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        *,
+        attn_backend=None,
+        speculative_num_steps: Optional[int] = None,
+        speculative_num_draft_tokens: Optional[int] = None,
+    ):
+        # NPU patch_model override: monkey-patch torch_compile_decoration's
+        # patch_model with the NPU-specific version.
+        from sglang.srt.compilation import torch_compile_decoration
+
+        torch_compile_decoration.patch_model = patch_model_npu
+        super().__init__(
+            model_runner,
+            attn_backend=attn_backend,
+            speculative_num_steps=speculative_num_steps,
+            speculative_num_draft_tokens=speculative_num_draft_tokens,
+        )
+        self.update_attr_name = None
+        self.update_attr_type = None
+        self.model_runner = model_runner
+        self._init_arch_map()
+        self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
+        self.if_use_v2 = any(
+            arch
+            in ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM", "Step3p5ForCausalLM")
+            for arch in (model_runner.model_config.hf_config.architectures or [])
+        )
+
+    def _init_arch_map(self):
+        if self.is_dllm:
+            self.attr_name: Dict[str, str] = {
+                AttentionArch.MLA: "actual_seq_lengths_kv",
+                AttentionArch.MHA: "actual_seq_lengths_kv",
+                "TARGET_VERIFY": "actual_seq_kvlen",
+            }
+        else:
+            self.attr_name: Dict[str, str] = {
+                AttentionArch.MLA: "actual_seq_lengths_kv",
+                AttentionArch.MHA: "context_lens",
+                "TARGET_VERIFY": "actual_seq_kvlen",
+            }
+        self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
+            AttentionArch.MLA: [],
+            AttentionArch.MHA: torch.Tensor(),
+            "TARGET_VERIFY": [],
+        }
+
+    def _create_device_graph(self):
+        return torch.npu.NPUGraph()
+
+    def _capture_graph(self, graph, pool, stream, run_once_fn):
+        if self.enable_torch_compile:
+            skip_guard_context = torch.compiler.set_stance(skip_guard_eval_unsafe=True)
+        else:
+            skip_guard_context = empty_context()
+
+        with (
+            skip_guard_context,
+            torch.npu.graph(
+                graph,
+                pool=pool,
+                stream=stream,
+                auto_dispatch_capture=True,
+            ),
+        ):
+            out = run_once_fn()
+        return out
+
+    def _get_update_attr_name(self):
+        if self.if_use_v2:
+            return self.attr_name["TARGET_VERIFY"]
+        return self.attr_name[AttentionArch.MLA]
+
+    def _get_update_attr_type(self):
+        if self.if_use_v2:
+            return self.attr_type["TARGET_VERIFY"]
+        return self.attr_type[AttentionArch.MLA]
+
+    def _update_inputs(self, seq_lens):
+        if isinstance(self.update_attr_type, torch.Tensor):
+            seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
+
+        self.graphs[self.bs].update(
+            cpu_update_input=[{self.update_attr_name: seq_lens}]
+        )
+
+    def _cache_loc_dtype(self):
+        return torch.int32
+
+    def _init_profile_context_and_memory_record(self):
+        output_dir = os.path.join(
+            os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp"), "graph_capture_profile"
+        )
+        if not Path(output_dir).exists():
+            Path(output_dir).mkdir(parents=True, exist_ok=True)
+        logger.info(
+            f"Profiling starts for graph capture for NPU. Traces will be saved to: {output_dir}"
+        )
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            export_type=[torch_npu.profiler.ExportType.Text],
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+        )
+        profile_context = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
+            record_shapes=True,
+            profile_memory=True,
+            on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
+                output_dir, async_mode=True
+            ),
+            experimental_config=experimental_config,
+        )
+        return profile_context
+
+    def _post_process_after_profile(self, prof_context):
+        # for NPU, profile data will be saved to disk for further analysis.
+        pass
+
+    def execute(
+        self,
+        forward_batch: ForwardBatch,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+    ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
+        if forward_batch.needs_forward_metadata_init():
+            self.load_batch(forward_batch, pp_proxy_tensors)
+        else:
+            # In speculative decoding, these two fields are still needed.
+            self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
+            self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
+            if (
+                self.model_runner.spec_algorithm.is_dflash()
+                and self.model_runner.is_draft_worker
+                and forward_batch.input_embeds is not None
+            ):
+                self.buffers.input_embeds[: self.raw_num_token].copy_(
+                    forward_batch.input_embeds
+                )
+            if (
+                envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get()
+                and forward_batch.mrope_positions is not None
+            ):
+                self.buffers.mrope_positions[:, : self.raw_num_token].copy_(
+                    forward_batch.mrope_positions
+                )
+
+        graph_key = self._make_graph_key(self.bs)
+
+        if not (
+            is_deepseek_dsa(self.model_runner.model_config.hf_config)
+            or is_deepseek_v4(self.model_runner.model_config.hf_config)
+        ):
+            if forward_batch.forward_mode.is_target_verify():
+                seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
+                seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
+            else:
+                seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
+                    self.bs - self.raw_bs
+                )
+            output = self.backend.replay_with_input_update(
+                graph_key,
+                seq_lens=seq_lens,
+                attr_name=self._get_update_attr_name(),
+                attr_type=self._get_update_attr_type(),
+            )
+        else:
+            output = self.backend.replay(graph_key, forward_batch)
+
+        if isinstance(output, LogitsProcessorOutput):
+            if self.is_dllm:
+                next_token_logits = None
+                full_logits = (
+                    output.full_logits[: self.raw_num_token]
+                    if output.full_logits is not None
+                    else None
+                )
+            else:
+                full_logits = None
+                next_token_logits = (
+                    output.next_token_logits[: self.raw_num_token]
+                    if output.next_token_logits is not None
+                    else None
+                )
+            return LogitsProcessorOutput(
+                next_token_logits=next_token_logits,
+                full_logits=full_logits,
+                hidden_states=(
+                    output.hidden_states[: self.raw_num_token]
+                    if output.hidden_states is not None
+                    else None
+                ),
+            )
+        else:
+            assert isinstance(output, PPProxyTensors)
+            return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()})

python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py

@@ -0,0 +1,229 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""ViT NPU Graph Runner class."""
+
+from __future__ import annotations
+
+from typing import Dict, Hashable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch_npu
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
+from sglang.srt.runtime_context import get_server_args
+
+
+class ViTNpuGraphRunner(ViTCudaGraphRunner):
+    """Generic ViT NPU Graph Runner.
+
+    This runner captures the "blocks + merger + deepstack merger (optional)" part
+    of a vision transformer into a NPU graph and replays it for identical shapes.
+
+    Optional for Qwen3 deepstack:
+      - vit.deepstack_vision_indexes: Sequence[int]
+      - vit.deepstack_merger_list: nn.ModuleList (same length as deepstack_vision_indexes)
+    """
+
+    _graph_memory_pool = None
+
+    def __init__(
+        self,
+        vit: nn.Module,
+    ) -> None:
+        super().__init__(vit)
+        self.device_module = torch.get_device_module(self.device)
+        self.cu_seq_lens: Dict[Hashable, torch.Tensor] = {}
+
+        # rotary position buffers shared across graphs
+        self.sin_cos_ws: Dict[Hashable, Tuple[torch.Tensor, torch.Tensor]] = {}
+
+    @property
+    def device(self) -> torch.device:
+        return self.vit.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.vit.dtype
+
+    def _create_graph(
+        self,
+        graph_key: int,
+    ):
+
+        graph = torch_npu.npu.NPUGraph()
+        vit = self.vit
+
+        override_backend = get_server_args().mm_attention_backend
+        with torch_npu.npu.graph(graph, pool=ViTNpuGraphRunner._graph_memory_pool):
+            y = None
+            deepstack_outs: List[torch.Tensor] = []
+            deepstack_capture_idx = 0
+
+            for layer_num, blk in enumerate(vit.blocks):
+                if override_backend == "ascend_attn":
+                    cu_seq_lens = self.cu_seq_lens[graph_key]
+                else:
+                    raise RuntimeError("Not supported ViT attention backend")
+
+                if layer_num == 0:
+                    y = blk(
+                        self.block_input[graph_key],
+                        cu_seqlens=cu_seq_lens,
+                        rotary_pos_emb_cos=self.sin_cos_ws[graph_key][0],
+                        rotary_pos_emb_sin=self.sin_cos_ws[graph_key][1],
+                        output_ws=self.block_ws[graph_key],
+                    )
+                else:
+                    y = blk(
+                        y,
+                        cu_seqlens=cu_seq_lens,
+                        rotary_pos_emb_cos=self.sin_cos_ws[graph_key][0],
+                        rotary_pos_emb_sin=self.sin_cos_ws[graph_key][1],
+                        output_ws=self.block_ws[graph_key],
+                    )
+
+                # Optional deepstack support (Qwen3-VL)
+                if (
+                    self._deepstack_visual_indexes
+                    and layer_num in self._deepstack_visual_indexes
+                ):
+                    if self._deepstack_merger_list is None:
+                        raise RuntimeError(
+                            "deepstack_visual_indexes exists but deepstack_merger_list is missing."
+                        )
+                    deepstack_out = self._deepstack_merger_list[deepstack_capture_idx](
+                        y
+                    )
+                    deepstack_outs.append(deepstack_out)
+                    deepstack_capture_idx += 1
+
+            main_out = vit.merger(y)
+
+            if deepstack_outs:
+                self.block_output[graph_key] = torch.cat(
+                    [main_out] + deepstack_outs, dim=1
+                )
+            else:
+                self.block_output[graph_key] = main_out
+
+        self.block_graphs[graph_key] = graph
+
+    def create_graph(
+        self,
+        x_3d: torch.Tensor,  # [S, 1, H]
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+    ) -> int:
+        vit = self.vit
+        graph_key = self._get_graph_key(x_3d)
+
+        if graph_key in self.block_graphs:
+            return graph_key
+
+        if ViTNpuGraphRunner._graph_memory_pool is None:
+            ViTNpuGraphRunner._graph_memory_pool = (
+                self.device_module.graph_pool_handle()
+            )
+        # Set graph pool id globally to be able to use symmetric memory
+        set_graph_pool_id(ViTNpuGraphRunner._graph_memory_pool)
+
+        # pre-allocate workspace
+        attn_module: VisionAttention = vit.blocks[0].attn
+        num_heads = attn_module.num_attention_heads_per_partition
+        attn_head_dim = attn_module.head_size
+
+        if graph_key not in self.block_output:
+            self.block_output[graph_key] = x_3d
+            self.block_input[graph_key] = x_3d
+            self.block_ws[graph_key] = torch.empty(
+                graph_key,
+                num_heads,
+                attn_head_dim,
+                device=self.device,
+                dtype=self.dtype,
+            )
+            if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+                self.sin_cos_ws[graph_key] = (rotary_pos_emb_cos, rotary_pos_emb_sin)
+
+        if graph_key not in self.cu_seq_lens:
+            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            self.cu_seq_lens[graph_key] = seq_lens.to("cpu").to(torch.int32)
+
+        if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+            self._create_graph(
+                graph_key=graph_key,
+            )
+
+        return graph_key
+
+    def replay(
+        self,
+        graph_key: int,
+        x_3d: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+        output_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+            # update rotary workspace content
+            self.sin_cos_ws[graph_key][0].copy_(rotary_pos_emb_cos)
+            self.sin_cos_ws[graph_key][1].copy_(rotary_pos_emb_sin)
+
+        # copy input
+        self.block_input[graph_key].copy_(x_3d)
+
+        # replay
+        self.block_graphs[graph_key].replay()
+
+        out = self.block_output[graph_key]
+
+        # Optional output reordering (Qwen2.5-VL window permutation inverse)
+        if output_indices is not None:
+            out = out.index_select(0, output_indices)
+
+        return out
+
+    def run(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+        output_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # x: [seq_len, hidden] -> [S, B=1, H]
+        x_3d = x.unsqueeze(1)
+        graph_key = self._get_graph_key(x_3d)
+        if graph_key not in self.block_graphs:
+            self.create_graph(
+                x_3d=x_3d,
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb_cos=rotary_pos_emb_cos,
+                rotary_pos_emb_sin=rotary_pos_emb_sin,
+            )
+
+        return self.replay(
+            graph_key=graph_key,
+            x_3d=x_3d,
+            rotary_pos_emb_cos=rotary_pos_emb_cos,
+            rotary_pos_emb_sin=rotary_pos_emb_sin,
+            output_indices=output_indices,
+        )

python/sglang/srt/hardware_backend/npu/memory_pool_npu.py

@@ -0,0 +1,553 @@
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.mem_cache.memory_pool import (
+    MHATokenToKVPool,
+    MLATokenToKVPool,
+    get_tensor_size_bytes,
+    unwrap_write_loc,
+)
+from sglang.srt.utils import get_bool_env_var
+from sglang.srt.utils.common import is_npu
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+
+if is_npu():
+    import torch_npu
+
+
+def _init_npu_conv_state(
+    conv_state_in, conv_state_shape, speculative_num_draft_tokens: Optional[int] = None
+):
+    extra_conv_len = 0
+    if speculative_num_draft_tokens is not None:
+        extra_conv_len = speculative_num_draft_tokens - 1
+
+    # conv_state shape (layers, pool_size, conv_wind + draft_step, dim) for conv1d ascendc ops require dim as last dim
+    conv_state = [
+        torch.zeros(
+            size=(
+                conv_state_in.shape[0],
+                conv_state_in.shape[1],
+                conv_shape[1] + extra_conv_len,
+                conv_shape[0],
+            ),
+            dtype=conv_state_in.dtype,
+            device=conv_state_in.device,
+        )
+        for conv_shape in conv_state_shape
+    ]
+    return conv_state
+
+
+class NPUMHATokenToKVPool(MHATokenToKVPool):
+
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        head_num: int,
+        head_dim: int,
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+        v_head_dim: Optional[int] = None,
+        swa_head_num: Optional[int] = None,
+        swa_head_dim: Optional[int] = None,
+        swa_v_head_dim: Optional[int] = None,
+        start_layer: Optional[int] = None,
+        end_layer: Optional[int] = None,
+        enable_alt_stream: bool = True,
+        enable_kv_cache_copy: bool = False,
+        **kwargs,
+    ):
+        self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
+        super().__init__(
+            size=size,
+            page_size=page_size,
+            dtype=dtype,
+            head_num=head_num,
+            head_dim=head_dim,
+            layer_num=layer_num,
+            device=device,
+            enable_memory_saver=enable_memory_saver,
+            v_head_dim=v_head_dim,
+            swa_head_num=swa_head_num,
+            swa_head_dim=swa_head_dim,
+            swa_v_head_dim=swa_v_head_dim,
+            start_layer=start_layer,
+            end_layer=end_layer,
+            enable_alt_stream=enable_alt_stream,
+            enable_kv_cache_copy=enable_kv_cache_copy,
+            **kwargs,
+        )
+
+    def _create_buffers(self):
+        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            # [size, head_num, head_dim] for each layer
+            # The padded slot 0 is used for writing dummy outputs from padded tokens.
+            # Continuous memory improves the efficiency of Ascend`s transmission backend,
+            # while other backends remain unchanged.
+            self.k_buffer = torch.zeros(
+                (
+                    self.layer_num,
+                    self.size // self.page_size + 1,
+                    self.page_size,
+                    self.head_num,
+                    self.head_dim,
+                ),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+            self.v_buffer = torch.zeros(
+                (
+                    self.layer_num,
+                    self.size // self.page_size + 1,
+                    self.page_size,
+                    self.head_num,
+                    self.v_head_dim,
+                ),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+
+            if self.use_fia:
+                # Use per-layer Python lists to avoid torch.compile capturing
+                # the entire multi-layer tensor (OOM during graph capture).
+                # Each layer view: [P*ps, 1, H, D], sharing the contiguous
+                # storage allocated above.
+                self.k_buffer = [
+                    self.k_buffer[i].view(-1, 1, self.head_num, self.head_dim)
+                    for i in range(self.layer_num)
+                ]
+                self.v_buffer = [
+                    self.v_buffer[i].view(-1, 1, self.head_num, self.v_head_dim)
+                    for i in range(self.layer_num)
+                ]
+
+    def _init_kv_copy_and_warmup(self):
+        # implementation relies on self.data_strides / self.data_ptrs, which the
+        # NPU paged buffer layout never builds.
+        self._kv_copy_config = None
+
+    # for disagg
+    def get_contiguous_buf_infos(self):
+        # layer_num x [seq_len, head_num, head_dim]
+        # layer_num x [page_num, page_size, head_num, head_dim]
+        kv_data_ptrs = [
+            self.get_key_buffer(i).data_ptr()
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ] + [
+            self.get_value_buffer(i).data_ptr()
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ]
+        kv_data_lens = [
+            self.get_key_buffer(i).nbytes
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ] + [
+            self.get_value_buffer(i).nbytes
+            for i in range(self.start_layer, self.start_layer + self.layer_num)
+        ]
+        if self.use_fia:
+            kv_item_lens = [
+                self.get_key_buffer(i)[0].nbytes * self.page_size
+                for i in range(self.start_layer, self.start_layer + self.layer_num)
+            ] + [
+                self.get_value_buffer(i)[0].nbytes * self.page_size
+                for i in range(self.start_layer, self.start_layer + self.layer_num)
+            ]
+        else:
+            kv_item_lens = [
+                self.get_key_buffer(i)[0].nbytes
+                for i in range(self.start_layer, self.start_layer + self.layer_num)
+            ] + [
+                self.get_value_buffer(i)[0].nbytes
+                for i in range(self.start_layer, self.start_layer + self.layer_num)
+            ]
+        return kv_data_ptrs, kv_data_lens, kv_item_lens
+
+    def set_kv_buffer(
+        self,
+        layer: "RadixAttention",
+        loc_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        layer_id_override: Optional[int] = None,
+        dcp_kv_mask: Optional[torch.Tensor] = None,
+    ):
+        loc, _, _ = unwrap_write_loc(loc_info)
+        if layer_id_override is not None:
+            layer_id = layer_id_override
+        else:
+            layer_id = layer.layer_id
+        if cache_k.dtype != self.dtype:
+            if k_scale is not None:
+                cache_k.div_(k_scale)
+            if v_scale is not None:
+                cache_v.div_(v_scale)
+            cache_k = cache_k.to(self.dtype)
+            cache_v = cache_v.to(self.dtype)
+
+        if self.store_dtype != self.dtype:
+            cache_k = cache_k.view(self.store_dtype)
+            cache_v = cache_v.view(self.store_dtype)
+
+        if self.use_fia:
+            k_buffer_layer = self.k_buffer[layer_id - self.start_layer]
+            v_buffer_layer = self.v_buffer[layer_id - self.start_layer]
+
+            torch_npu.npu_scatter_nd_update_(
+                k_buffer_layer,
+                loc.view(-1, 1),
+                cache_k.view(-1, 1, self.head_num, self.head_dim),
+            )
+            torch_npu.npu_scatter_nd_update_(
+                v_buffer_layer,
+                loc.view(-1, 1),
+                cache_v.view(-1, 1, self.head_num, self.v_head_dim),
+            )
+        else:
+            loc = loc.to(torch.int32)
+            torch_npu._npu_reshape_and_cache(
+                key=cache_k,
+                value=cache_v,
+                key_cache=self.k_buffer[layer_id - self.start_layer].view(
+                    -1, self.page_size, self.head_num, self.head_dim
+                ),
+                value_cache=self.v_buffer[layer_id - self.start_layer].view(
+                    -1, self.page_size, self.head_num, self.v_head_dim
+                ),
+                slot_indices=loc,
+            )
+
+    def _chunk_copy_npu_to_cpu(self, buf_of_layers, indices):
+        chunk_size = self.cpu_offloading_chunk_size
+        out = []
+        for tensors_per_layer in buf_of_layers:  # [k_buf, v_buf]
+            layer_chunks = []
+            for i in range(0, len(indices), chunk_size):
+                ci = indices[i : i + chunk_size]
+                layer_chunks.append(
+                    [
+                        t[ci].to("cpu", non_blocking=True)
+                        for t in tensors_per_layer
+                        if t is not None
+                    ]
+                )
+            out.append(layer_chunks)
+        return out
+
+    # Parent MHATokenToKVPool.get_cpu_copy / load_cpu_copy use
+    # `self.k_buffer[layer_id][chunk_indices]` which indexes the first dim.
+    # NPUMHATokenToKVPool stores buffers as
+    #   (num_pages, page_size, head_num, head_dim)            # use_fia=False
+    #   (num_pages*page_size, 1, head_num, head_dim)          # use_fia=True
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        buf_of_layers = []
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            v_layer = self.v_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            buf_of_layers.append([k_layer, v_layer])
+        kv_cache_cpu = self._chunk_copy_npu_to_cpu(buf_of_layers, indices)
+        torch.npu.synchronize()
+        return kv_cache_cpu
+
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        chunk_size = self.cpu_offloading_chunk_size
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            v_layer = self.v_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            for i in range(0, len(indices), chunk_size):
+                chunk_indices = indices[i : i + chunk_size]
+                k_cpu, v_cpu = (
+                    kv_cache_cpu[local_layer_id][i // chunk_size][0],
+                    kv_cache_cpu[local_layer_id][i // chunk_size][1],
+                )
+                assert k_cpu.shape[0] == v_cpu.shape[0] == len(chunk_indices)
+                k_layer[chunk_indices] = k_cpu.to(k_layer.device, non_blocking=True)
+                v_layer[chunk_indices] = v_cpu.to(v_layer.device, non_blocking=True)
+        torch.npu.synchronize()
+
+
+class NPUMLATokenToKVPool(MLATokenToKVPool):
+
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        kv_lora_rank: int,
+        qk_rope_head_dim: int,
+        index_head_dim: Optional[int],
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+        start_layer: Optional[int] = None,
+        end_layer: Optional[int] = None,
+    ):
+        super(MLATokenToKVPool, self).__init__(
+            size=size,
+            page_size=page_size,
+            dtype=dtype,
+            layer_num=layer_num,
+            device=device,
+            enable_memory_saver=enable_memory_saver,
+            start_layer=start_layer,
+            end_layer=end_layer,
+        )
+
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.index_head_dim = index_head_dim
+
+        self.custom_mem_pool = None
+
+        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            # The padded slot 0 is used for writing dummy outputs from padded tokens.
+            self.k_buffer = torch.zeros(
+                (
+                    layer_num,
+                    self.size // self.page_size + 1,
+                    self.page_size,
+                    1,
+                    self.kv_lora_rank,
+                ),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+            self.v_buffer = torch.zeros(
+                (
+                    layer_num,
+                    self.size // self.page_size + 1,
+                    self.page_size,
+                    1,
+                    self.qk_rope_head_dim,
+                ),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+            self.index_k_buffer = None
+            if self.index_head_dim is not None:
+                self.index_k_buffer = torch.zeros(
+                    (
+                        layer_num,
+                        self.size // self.page_size + 1,
+                        self.page_size,
+                        1,
+                        self.index_head_dim,
+                    ),
+                    dtype=self.store_dtype,
+                    device=self.device,
+                )
+
+        self._finalize_allocation_log(size)
+
+    def get_kv_size_bytes(self):
+        assert hasattr(self, "k_buffer")
+        assert hasattr(self, "v_buffer")
+        kv_size_bytes = 0
+        for k_cache in self.k_buffer:
+            kv_size_bytes += get_tensor_size_bytes(k_cache)
+        for v_cache in self.v_buffer:
+            kv_size_bytes += get_tensor_size_bytes(v_cache)
+        if self.index_head_dim is not None:
+            assert hasattr(self, "index_k_buffer")
+            for index_k_cache in self.index_k_buffer:
+                kv_size_bytes += get_tensor_size_bytes(index_k_cache)
+        return kv_size_bytes
+
+    def get_kv_buffer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+        return (
+            self.k_buffer[layer_id - self.start_layer],
+            self.v_buffer[layer_id - self.start_layer],
+        )
+
+    def get_state_buf_infos(self):
+        if self.index_head_dim is None:
+            return [], [], []
+        data_ptrs = [self.index_k_buffer[i].data_ptr() for i in range(self.layer_num)]
+        data_lens = [self.index_k_buffer[i].nbytes for i in range(self.layer_num)]
+        item_lens = [self.index_k_buffer[i][0].nbytes for i in range(self.layer_num)]
+        return data_ptrs, data_lens, item_lens
+
+    def get_key_buffer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+
+        if self.store_dtype != self.dtype:
+            return self.k_buffer[layer_id - self.start_layer].view(self.dtype)
+        return self.k_buffer[layer_id - self.start_layer]
+
+    def get_value_buffer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+
+        if self.store_dtype != self.dtype:
+            return self.v_buffer[layer_id - self.start_layer].view(self.dtype)
+        return self.v_buffer[layer_id - self.start_layer]
+
+    def get_index_k_buffer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+
+        if self.store_dtype != self.dtype:
+            return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype)
+        return self.index_k_buffer[layer_id - self.start_layer]
+
+    # for disagg
+    def get_contiguous_buf_infos(self):
+        # MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
+        kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)] + [
+            self.v_buffer[i].data_ptr() for i in range(self.layer_num)
+        ]
+        kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)] + [
+            self.v_buffer[i].nbytes for i in range(self.layer_num)
+        ]
+        kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)] + [
+            self.v_buffer[i][0].nbytes for i in range(self.layer_num)
+        ]
+        if self.index_head_dim is not None:
+            kv_data_ptrs += [
+                self.index_k_buffer[i].data_ptr() for i in range(self.layer_num)
+            ]
+            kv_data_lens += [
+                self.index_k_buffer[i].nbytes for i in range(self.layer_num)
+            ]
+            kv_item_lens += [
+                self.index_k_buffer[i][0].nbytes for i in range(self.layer_num)
+            ]
+        return kv_data_ptrs, kv_data_lens, kv_item_lens
+
+    def set_kv_buffer(
+        self,
+        layer: "RadixAttention",
+        loc_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+    ):
+        loc, _, _ = unwrap_write_loc(loc_info)
+        layer_id = layer.layer_id
+        if cache_k.dtype != self.dtype:
+            cache_k = cache_k.to(self.dtype)
+            cache_v = cache_v.to(self.dtype)
+
+        if self.store_dtype != self.dtype:
+            cache_k = cache_k.view(self.store_dtype)
+            cache_v = cache_v.view(self.store_dtype)
+
+        if cache_v is None:
+            cache_k, cache_v = cache_k.split(
+                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+            )
+
+        torch_npu.npu_scatter_nd_update_(
+            self.k_buffer[layer_id - self.start_layer].view(-1, 1, self.kv_lora_rank),
+            loc.view(-1, 1),
+            cache_k.view(-1, 1, self.kv_lora_rank),
+        )
+        torch_npu.npu_scatter_nd_update_(
+            self.v_buffer[layer_id - self.start_layer].view(
+                -1, 1, self.qk_rope_head_dim
+            ),
+            loc.view(-1, 1),
+            cache_v.view(-1, 1, self.qk_rope_head_dim),
+        )
+
+    def set_index_k_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        index_k: torch.Tensor,
+    ):
+        if index_k.dtype != self.dtype:
+            index_k = index_k.to(self.dtype)
+
+        if self.store_dtype != self.dtype:
+            index_k = index_k.view(self.store_dtype)
+
+        torch_npu.npu_scatter_nd_update_(
+            self.index_k_buffer[layer_id - self.start_layer].view(
+                -1, 1, self.index_head_dim
+            ),
+            loc.view(-1, 1),
+            index_k.view(-1, 1, self.index_head_dim),
+        )
+
+    def _chunk_copy_npu_to_cpu(self, buf_of_layers, indices):
+        chunk_size = self.cpu_offloading_chunk_size
+        out = []
+        for tensors_per_layer in buf_of_layers:  # [k_buf, v_buf, ik_buf/None]
+            layer_chunks = []
+            for i in range(0, len(indices), chunk_size):
+                ci = indices[i : i + chunk_size]
+                layer_chunks.append(
+                    [
+                        t[ci].to("cpu", non_blocking=True)
+                        for t in tensors_per_layer
+                        if t is not None
+                    ]
+                )
+            out.append(layer_chunks)
+        return out
+
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        buf_of_layers = []
+        has_ik = self.index_head_dim is not None
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(-1, 1, self.kv_lora_rank)
+            v_layer = self.v_buffer[local_layer_id].view(-1, 1, self.qk_rope_head_dim)
+            ik_layer = (
+                self.index_k_buffer[local_layer_id].view(-1, 1, self.index_head_dim)
+                if has_ik
+                else None
+            )
+            buf_of_layers.append([k_layer, v_layer, ik_layer])
+
+        kv_cache_cpu = self._chunk_copy_npu_to_cpu(buf_of_layers, indices)
+        torch.npu.synchronize()
+        return kv_cache_cpu
+
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        chunk_size = self.cpu_offloading_chunk_size
+        has_ik = self.index_head_dim is not None
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(-1, 1, self.kv_lora_rank)
+            v_layer = self.v_buffer[local_layer_id].view(-1, 1, self.qk_rope_head_dim)
+            ik_layer = (
+                self.index_k_buffer[local_layer_id].view(-1, 1, self.index_head_dim)
+                if has_ik
+                else None
+            )
+            for i in range(0, len(indices), chunk_size):
+                chunk_indices = indices[i : i + chunk_size]
+                chunk = kv_cache_cpu[local_layer_id][i // chunk_size]
+                k_cpu, v_cpu = chunk[0], chunk[1]
+                assert k_cpu.shape[0] == len(chunk_indices)
+                k_layer[chunk_indices] = k_cpu.to(k_layer.device, non_blocking=True)
+                v_layer[chunk_indices] = v_cpu.to(v_layer.device, non_blocking=True)
+                if has_ik:
+                    ik_cpu = chunk[2]
+                    ik_layer[chunk_indices] = ik_cpu.to(
+                        ik_layer.device, non_blocking=True
+                    )
+        torch.npu.synchronize()

python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py

@@ -0,0 +1,604 @@
+import re
+from typing import TYPE_CHECKING
+
+import torch
+import torch_npu
+from sgl_kernel_npu.norm.fused_split_qk_norm import fused_split_qk_norm
+
+from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
+    NPUFusedMLAPreprocess,
+    is_fia_nz,
+    is_mla_preprocess_enabled,
+)
+from sglang.srt.layers.attention.dsa.dsa_indexer import scattered_to_tp_attn_full
+from sglang.srt.layers.attention.dsa.utils import (
+    dsa_use_prefill_cp,
+)
+from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+    from sglang.srt.utils import BumpAllocator
+_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
+
+
+# region MHA
+def forward_mha_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+):
+    if m.q_lora_rank is not None:
+        q, latent_cache = (
+            get_attn_tp_context()
+            .fetch_qkv_latent()
+            .split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
+                dim=-1,
+            )
+        )
+
+        # DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
+
+        if m.use_dsa:
+            q_lora = m.q_a_layernorm(q)
+            q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+            _ = m.indexer(
+                x=hidden_states,
+                q_lora=q_lora,
+                positions=positions,
+                forward_batch=forward_batch,
+                layer_id=m.layer_id,
+                return_indices=False,
+            )
+
+        else:
+            q = m.q_a_layernorm(q)
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+            q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+
+    else:
+        q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+        latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
+
+    _, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
+    kv_a, _ = latent_cache.split([m.kv_lora_rank, m.qk_rope_head_dim], dim=-1)
+    latent_cache = latent_cache.unsqueeze(1)
+
+    if m.use_deepseek_yarn_rope:
+        B, S = q.shape[0], 1
+        cos, sin = m.rotary_emb.get_cos_sin_cache(
+            positions, hidden_states.dtype, offsets=None
+        )
+        q_pe = torch_npu.npu_interleave_rope(
+            q_pe.reshape(B, -1, S, m.qk_rope_head_dim),
+            cos,
+            sin,
+        )
+        q_pe = q_pe.reshape(B, -1, m.qk_rope_head_dim)
+
+        ckv_cache, k_rope_cache = get_token_to_kv_pool().get_kv_buffer(m.layer_id)
+        _, _, k_pe, kv_a = torch_npu.npu_kv_rmsnorm_rope_cache(
+            latent_cache.view(-1, 1, 1, m.kv_lora_rank + m.qk_rope_head_dim),  # bnsd
+            m.kv_a_layernorm.weight,
+            cos,
+            sin,
+            forward_batch.out_cache_loc.to(torch.int64),
+            k_rope_cache,
+            ckv_cache,
+            k_rope_scale=None,
+            c_kv_scale=None,
+            k_rope_offset=None,
+            c_kv_offset=None,
+            epsilon=m.kv_a_layernorm.variance_epsilon,
+            cache_mode="PA_NZ" if is_fia_nz() else "PA_BNSD",
+            is_output_kv=True,
+        )  # adapter NZ
+
+        k_pe = k_pe.reshape(B, -1, m.qk_rope_head_dim)
+    else:
+        kv_a = m.kv_a_layernorm(kv_a)
+        k_pe = latent_cache[:, :, m.kv_lora_rank :]
+        if m.rotary_emb is not None:
+            q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+        # this is for model kimi-vl-a3B-instruct
+        get_token_to_kv_pool().set_kv_buffer(
+            m, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
+        )
+
+    q[..., m.qk_nope_head_dim :] = q_pe
+
+    kv = m.kv_b_proj(kv_a)[0]
+    kv = kv.view(-1, m.num_local_heads, m.qk_nope_head_dim + m.v_head_dim)
+    k_nope = kv[..., : m.qk_nope_head_dim]
+    v = kv[..., m.qk_nope_head_dim :]
+
+    k = m._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
+    return q, k, v, forward_batch
+
+
+def forward_mha_core_npu(
+    m: "DeepseekV2AttentionMLA",
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    forward_batch: "ForwardBatch",
+) -> torch.Tensor:
+    attn_output = m.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
+    attn_output = attn_output.reshape(-1, m.num_local_heads * m.v_head_dim)
+    output, _ = m.o_proj(attn_output)
+    return output
+
+
+# endregion
+
+
+# region MLA
+def forward_mla_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+):
+    if is_mla_preprocess_enabled():
+        if not hasattr(m, "mla_preprocess"):
+            m.mla_preprocess = NPUFusedMLAPreprocess(
+                m.fused_qkv_a_proj_with_mqa,
+                m.q_a_layernorm,
+                m.kv_a_layernorm,
+                m.q_b_proj,
+                m.w_kc,
+                m.rotary_emb,
+                m.layer_id,
+                m.num_local_heads,
+                m.qk_nope_head_dim,
+                m.qk_rope_head_dim,
+                m.quant_config,
+            )
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            forward_batch,
+            zero_allocator,
+            positions,
+        ) = m.mla_preprocess.forward(
+            positions, hidden_states, forward_batch, zero_allocator
+        )
+        topk_indices = None
+    else:
+        q_lora = None
+        if m.q_lora_rank is not None:
+            qkv_latent = get_attn_tp_context().fetch_qkv_latent()
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q, latent_cache = qkv_latent.split(
+                    [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
+                    dim=-1,
+                )
+                k_nope = latent_cache[..., : m.kv_lora_rank]
+
+                q = m.q_a_layernorm(q)
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+
+                k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+                k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+            else:
+                if qkv_latent.shape[0] < 65536 and not dsa_use_prefill_cp(
+                    forward_batch
+                ):
+                    q, k_nope, k_pe = fused_split_qk_norm(
+                        qkv_latent,
+                        m.q_a_layernorm,
+                        m.kv_a_layernorm,
+                        m.q_lora_rank,
+                        m.kv_lora_rank,
+                        m.qk_rope_head_dim,
+                        eps=m.q_a_layernorm.variance_epsilon,
+                    )
+                else:
+                    q, latent_cache = qkv_latent.split(
+                        [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
+                        dim=-1,
+                    )
+                    k_nope = latent_cache[..., : m.kv_lora_rank]
+
+                    q = m.q_a_layernorm(q)
+
+                    k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+                    k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+
+            # q_lora needed by indexer
+            if m.use_dsa:
+                q_lora = q
+
+            q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+        else:
+            q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+            latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
+            k_nope = latent_cache[..., : m.kv_lora_rank]
+            k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+            k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+
+        q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
+
+        q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
+
+        q_nope_out = q_nope_out.transpose(0, 1)
+
+        q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+
+        if dsa_use_prefill_cp(forward_batch):
+            # support allgather+rerrange
+            k_nope, k_pe = m.rebuild_cp_kv_cache(
+                latent_cache, forward_batch, k_nope, k_pe
+            )
+        topk_indices = None
+        if q_lora is not None:
+            topk_indices = m.indexer(
+                x=hidden_states,
+                q_lora=q_lora,
+                positions=positions,
+                forward_batch=forward_batch,
+                layer_id=m.layer_id,
+            )
+
+    return (
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        forward_batch,
+        zero_allocator,
+        positions,
+        topk_indices,
+    )
+
+
+def forward_mla_core_npu(
+    m: "DeepseekV2AttentionMLA",
+    q_pe: torch.Tensor,
+    k_pe: torch.Tensor,
+    q_nope_out: torch.Tensor,
+    k_nope: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    positions: torch.Tensor,
+    topk_indices: torch.Tensor,
+) -> torch.Tensor:
+    attn_output = m.attn_mqa(
+        q_nope_out,
+        k_nope,
+        k_nope,
+        forward_batch,
+        q_rope=q_pe,
+        k_rope=k_pe,
+        **(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
+    )
+
+    attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
+
+    attn_bmm_output = torch.empty(
+        (attn_output.shape[0], m.num_local_heads, m.v_head_dim),
+        dtype=attn_output.dtype,
+        device=attn_output.device,
+    )
+
+    attn_output = attn_output.contiguous()
+    torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output)
+
+    attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
+    output, _ = m.o_proj(attn_bmm_output)
+
+    return output
+
+
+# endregion
+
+
+# region DSA
+def forward_dsa_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+    prev_topk_indices: torch.Tensor = None,
+):
+    dynamic_scale = None
+    if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode():
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            q_lora,
+            forward_batch,
+            zero_allocator,
+            positions,
+            dynamic_scale,
+        ) = npu_mla_preprocess(
+            m,
+            hidden_states,
+            positions,
+            forward_batch,
+            zero_allocator,
+        )
+    else:
+        fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
+        if m.rotary_emb.is_neox_style:
+            q, latent_cache = fused_qkv_a_proj_out.split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+            )
+            # overlap qk norm
+            q = m.q_a_layernorm(q)
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+            q_lora = q.clone()  # required for topk_indices
+
+            q_event = None
+            if m.alt_stream is not None:
+                m.alt_stream.wait_stream(torch.npu.current_stream())
+                with torch.npu.stream(m.alt_stream):
+                    q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+                    # record q to ensure memory space will not be released
+                    q.record_stream(m.alt_stream)
+                    q_event = m.alt_stream.record_event()
+            else:
+                q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+
+            k_nope, k_pe = latent_cache.unsqueeze(1).split(
+                [m.kv_lora_rank, m.qk_rope_head_dim], dim=-1
+            )
+            k_nope = m.kv_a_layernorm(k_nope)
+            # main stream waits for the completion of the event on the alt stream to ensure data dependency is complete
+            if q_event is not None:
+                torch.npu.current_stream().wait_event(q_event)
+        else:
+            if fused_qkv_a_proj_out.shape[0] < 65535 and not dsa_use_prefill_cp(
+                forward_batch
+            ):
+                q_lora, k_nope, k_pe = fused_split_qk_norm(
+                    fused_qkv_a_proj_out,
+                    m.q_a_layernorm,
+                    m.kv_a_layernorm,
+                    m.q_lora_rank,
+                    m.kv_lora_rank,
+                    m.qk_rope_head_dim,
+                    eps=m.q_a_layernorm.variance_epsilon,
+                )
+            else:
+                q, latent_cache = fused_qkv_a_proj_out.split(
+                    [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+                )
+                # overlap qk norm
+                q = m.q_a_layernorm(q)
+
+                q_lora = q.clone()  # required for topk_indices
+                k_nope, k_pe = latent_cache.unsqueeze(1).split(
+                    [m.kv_lora_rank, m.qk_rope_head_dim], dim=-1
+                )
+                k_nope = m.kv_a_layernorm(k_nope)
+            q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+
+        q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
+
+        q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
+
+        q_nope_out = q_nope_out.transpose(0, 1)
+
+        if m.layer_id == 0:
+            m.rotary_emb.sin_cos_cache = m.rotary_emb.cos_sin_cache.index_select(
+                0, positions
+            )
+
+        q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+
+        if dsa_use_prefill_cp(forward_batch):
+            # support allgather+rerrange
+            k_nope, k_pe = m.rebuild_cp_kv_cache(
+                latent_cache, forward_batch, k_nope, k_pe
+            )
+
+    if not m.skip_topk or (m.is_nextn and prev_topk_indices is None):
+        topk_indices = m.indexer(
+            hidden_states,
+            q_lora,
+            positions,
+            forward_batch,
+            m.layer_id,
+            layer_scatter_modes,
+            dynamic_scale,
+        )
+    else:
+        topk_indices = prev_topk_indices
+
+    return (
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        topk_indices,
+        forward_batch,
+        zero_allocator,
+        positions,
+    )
+
+
+def forward_dsa_core_npu(
+    m: "DeepseekV2AttentionMLA",
+    q_pe: torch.Tensor,
+    k_pe: torch.Tensor,
+    q_nope_out: torch.Tensor,
+    k_nope: torch.Tensor,
+    topk_indices: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    positions: torch.Tensor,
+) -> torch.Tensor:
+    attn_output = m.attn_mqa(
+        q_nope_out.contiguous(),
+        k_nope.contiguous(),
+        k_nope.contiguous(),
+        forward_batch,
+        save_kv_cache=True,  # False if forward_batch.forward_mode.is_extend() else True,
+        q_rope=q_pe.contiguous(),
+        k_rope=k_pe.contiguous(),
+        topk_indices=topk_indices,
+    )
+    attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
+
+    attn_bmm_output = torch.empty(
+        (attn_output.shape[0], m.num_local_heads, m.v_head_dim),
+        dtype=attn_output.dtype,
+        device=attn_output.device,
+    )
+
+    if (
+        forward_batch.forward_mode.is_extend()
+        and not forward_batch.forward_mode.is_draft_extend_v2()
+        and not forward_batch.forward_mode.is_target_verify()
+    ):
+        attn_output = attn_output.transpose(0, 1)
+        torch.bmm(
+            attn_output,
+            m.w_vc,
+            out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(
+                0, 1
+            ),
+        )
+    else:
+        attn_output = attn_output.contiguous()
+        torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output)
+
+    attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
+
+    output, _ = m.o_proj(attn_bmm_output)
+    if not m.next_skip_topk:
+        return output, None
+    else:
+        return output, topk_indices
+
+
+def npu_mla_preprocess(
+    m: "DeepseekV2AttentionMLA",
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+):
+    dynamic_scale = None
+    if not hasattr(m, "mla_preprocess"):
+        m.mla_preprocess = NPUFusedMLAPreprocess(
+            m.fused_qkv_a_proj_with_mqa,
+            m.q_a_layernorm,
+            m.kv_a_layernorm,
+            m.q_b_proj,
+            m.w_kc,
+            m.rotary_emb,
+            m.layer_id,
+            m.num_local_heads,
+            m.qk_nope_head_dim,
+            m.qk_rope_head_dim,
+            m.v_head_dim,
+            m.quant_config,
+        )
+    # mlaprolog does not require additional calculation of q_lora
+    _is_mlaprolog = hasattr(m.quant_config, "ignore") and any(
+        re.fullmatch(r".*kv_b_proj", l) for l in m.quant_config.ignore
+    )
+    if _is_mlaprolog:
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            q_lora,
+            forward_batch,
+            positions,
+            dynamic_scale,
+        ) = m.mla_preprocess.forward(
+            positions, hidden_states, forward_batch, zero_allocator
+        )
+    else:
+        if m.alt_stream is not None:
+            mla_event = torch.npu.Event()
+            mla_event.record()
+            with torch.npu.stream(m.alt_stream):
+                # alt stream waits for the completion of the event on the main stream to ensure data dependency is complete
+                torch.npu.current_stream().wait_event(mla_event)
+                (
+                    q_pe,
+                    k_pe,
+                    q_nope_out,
+                    k_nope,
+                    forward_batch,
+                    zero_allocator,
+                    positions,
+                ) = m.mla_preprocess.forward(
+                    positions, hidden_states, forward_batch, zero_allocator
+                )
+
+            fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
+            q, _ = fused_qkv_a_proj_out.split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+            )
+            q_lora = m.q_a_layernorm(q)
+            torch.npu.current_stream().wait_event(m.alt_stream)
+        else:
+            (
+                q_pe,
+                k_pe,
+                q_nope_out,
+                k_nope,
+                forward_batch,
+                zero_allocator,
+                positions,
+            ) = m.mla_preprocess.forward(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+            fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
+            q, _ = fused_qkv_a_proj_out.split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+            )
+            q_lora = m.q_a_layernorm(q)
+
+    return (
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        q_lora,
+        forward_batch,
+        zero_allocator,
+        positions,
+        dynamic_scale,
+    )
+
+
+# endregion

python/sglang/srt/hardware_backend/npu/modules/glm46v_processor.py

@@ -0,0 +1,285 @@
+"""NPU patch for GLM-4.6V image and video preprocessing.
+
+The GLM-4.6V image processor (Glm46VImageProcessorFast) and video processor
+(Glm46VVideoProcessor) create 10-dimensional tensors during patch extraction,
+which exceeds Ascend NPU's 8-dimension limit.
+
+This patch restructures the computation to stay within 8 dimensions, following
+the same pattern as the Qwen VL NPU patch.
+"""
+
+import torch
+import torchvision.transforms.v2.functional as tvF
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_processing_utils_fast import (
+    group_images_by_shape,
+    reorder_images,
+)
+from transformers.image_utils import (
+    ChannelDimension,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+)
+from transformers.models.glm46v.image_processing_glm46v import smart_resize
+from transformers.utils import TensorType
+from transformers.video_utils import group_videos_by_shape, reorder_videos
+
+from sglang.srt.hardware_backend.npu.modules.qwen_vl_processor import (
+    transform_patches_to_flatten,
+)
+from sglang.srt.utils import apply_module_patch
+
+
+# Func refers to transformers.models.glm46v.image_processing_glm46v_fast.py
+# Glm46VImageProcessorFast._preprocess
+def npu_wrapper_glm46v_preprocess(func):
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int,
+        temporal_patch_size: int,
+        merge_size: int,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        **kwargs,
+    ):
+        grouped_images, grouped_images_index = group_images_by_shape(
+            images, disable_grouping=disable_grouping
+        )
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            height, width = stacked_images.shape[-2:]
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    num_frames=temporal_patch_size,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_images = self.resize(
+                    stacked_images,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+            resized_images_grouped[shape] = stacked_images
+
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        grouped_images, grouped_images_index = group_images_by_shape(
+            resized_images, disable_grouping=disable_grouping
+        )
+        processed_images_grouped = {}
+        processed_grids = {}
+
+        for shape, stacked_images in grouped_images.items():
+            resized_height, resized_width = stacked_images.shape[-2:]
+
+            patches = self.rescale_and_normalize(
+                stacked_images,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            if patches.ndim == 4:
+                patches = patches.unsqueeze(1)
+
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(
+                    1,
+                    temporal_patch_size - (patches.shape[1] % temporal_patch_size),
+                    1,
+                    1,
+                    1,
+                )
+                patches = torch.cat([patches, repeats], dim=1)
+
+            batch_size, t_len, channel = patches.shape[:3]
+            grid_t = t_len // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_images_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_images = reorder_images(
+            processed_images_grouped, grouped_images_index
+        )
+        processed_grids = reorder_images(processed_grids, grouped_images_index)
+
+        pixel_values = torch.cat(processed_images, dim=0)
+        image_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    return _preprocess
+
+
+# Func refers to transformers.models.glm46v.video_processing_glm46v.py
+# Glm46VVideoProcessor._preprocess
+def npu_wrapper_glm46v_video_preprocess(func):
+
+    def _preprocess(
+        self,
+        videos: list[torch.Tensor],
+        do_convert_rgb: bool = True,
+        do_resize: bool = True,
+        size: SizeDict | None = None,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None" = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        image_mean: float | list[float] | None = None,
+        image_std: float | list[float] | None = None,
+        patch_size: int | None = None,
+        temporal_patch_size: int | None = None,
+        merge_size: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ):
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+
+        for shape, stacked_videos in grouped_videos.items():
+            B, T, C, H, W = stacked_videos.shape
+            num_frames, height, width = T, H, W
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    num_frames=num_frames,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_videos = stacked_videos.view(B * T, C, H, W)
+                stacked_videos = self.resize(
+                    stacked_videos,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+                stacked_videos = stacked_videos.view(
+                    B, T, C, resized_height, resized_width
+                )
+            resized_videos_grouped[shape] = stacked_videos
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+
+        # Group videos by size for further processing
+        # Needed in case do_resize is False, or resize returns videos with different sizes
+        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
+        processed_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            resized_height, resized_width = get_image_size(
+                stacked_videos[0], channel_dim=ChannelDimension.FIRST
+            )
+
+            # Fused rescale and normalize
+            stacked_videos = self.rescale_and_normalize(
+                stacked_videos,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            patches = stacked_videos
+
+            # Check that videos have `num_frames` divisible by `temporal_patch_size`
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)
+                patches = torch.cat([patches, repeats], dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_videos_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_videos = reorder_videos(
+            processed_videos_grouped, grouped_videos_index
+        )
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(processed_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+        data = {
+            "pixel_values_videos": pixel_values_videos,
+            "video_grid_thw": video_grid_thw,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    return _preprocess
+
+
+_npu_glm46v_preprocess_patched = False
+
+
+def npu_apply_glm46v_image_preprocess_patch():
+    global _npu_glm46v_preprocess_patched
+    if _npu_glm46v_preprocess_patched:
+        return
+    apply_module_patch(
+        "transformers.models.glm46v.image_processing_glm46v_fast.Glm46VImageProcessorFast",
+        "_preprocess",
+        [npu_wrapper_glm46v_preprocess],
+    )
+    apply_module_patch(
+        "transformers.models.glm46v.video_processing_glm46v.Glm46VVideoProcessor",
+        "_preprocess",
+        [npu_wrapper_glm46v_video_preprocess],
+    )
+    _npu_glm46v_preprocess_patched = True

python/sglang/srt/hardware_backend/npu/modules/qwen_vl_processor.py

@@ -0,0 +1,304 @@
+import torch
+import torchvision.transforms.v2.functional as tvF
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_transforms import group_images_by_shape, reorder_images
+from transformers.image_utils import (
+    ChannelDimension,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+)
+from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
+from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
+    smart_resize as smart_resize_video,
+)
+from transformers.utils import TensorType
+from transformers.video_utils import group_videos_by_shape, reorder_videos
+
+from sglang.srt.utils import apply_module_patch
+
+
+def transform_patches_to_flatten(
+    patches: torch.Tensor,
+    batch_size: int,
+    grid_t: int,
+    temporal_patch_size: int,
+    channel: int,
+    grid_h: int,
+    grid_w: int,
+    patch_size: int,
+    merge_size: int,
+) -> torch.Tensor:
+    patches = patches.view(
+        batch_size * grid_t,
+        temporal_patch_size * channel,
+        grid_h // merge_size,
+        merge_size,
+        patch_size,
+        grid_w // merge_size,
+        merge_size,
+        patch_size,
+    )
+    patches = patches.permute(0, 1, 2, 5, 3, 6, 4, 7)
+    patches = patches.reshape(
+        batch_size,
+        grid_t,
+        temporal_patch_size,
+        channel,
+        grid_h * grid_w,
+        patch_size,
+        patch_size,
+    )
+    patches = patches.permute(0, 1, 4, 3, 2, 5, 6)
+    flatten_patches = patches.reshape(
+        batch_size,
+        grid_t * grid_h * grid_w,
+        -1,
+    )
+    return flatten_patches
+
+
+# Func refers to transformers.models.qwen2_vl.image_processing_qwen2_vl.py
+# Qwen2VLImageProcessor._preprocess
+def npu_wrapper_preprocess(func):
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int,
+        temporal_patch_size: int,
+        merge_size: int,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        **kwargs,
+    ):
+        # Group images by size for batched resizing
+        grouped_images, grouped_images_index = group_images_by_shape(
+            images, disable_grouping=disable_grouping
+        )
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            height, width = stacked_images.shape[-2:]
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    height,
+                    width,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_images = self.resize(
+                    image=stacked_images,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+            resized_images_grouped[shape] = stacked_images
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        # Group images by size for further processing
+        # Needed in case do_resize is False, or resize returns images with different sizes
+        grouped_images, grouped_images_index = group_images_by_shape(
+            resized_images, disable_grouping=disable_grouping
+        )
+        processed_images_grouped = {}
+        processed_grids = {}
+        for shape, stacked_images in grouped_images.items():
+            resized_height, resized_width = stacked_images.shape[-2:]
+            # Fused rescale and normalize
+            patches = self.rescale_and_normalize(
+                stacked_images,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            if patches.ndim == 4:
+                # add a temporal dimension if we have images
+                patches = patches.unsqueeze(1)
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)
+                patches = torch.cat([patches, repeats], dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_images_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_images = reorder_images(
+            processed_images_grouped, grouped_images_index
+        )
+        processed_grids = reorder_images(processed_grids, grouped_images_index)
+        pixel_values = torch.cat(processed_images, dim=0)
+        image_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    return _preprocess
+
+
+# Func refers to transformers.models.qwen3_vl.video_processing_qwen3_vl.py
+# Qwen3VLVideoProcessor._preprocess
+def npu_wrapper_video_preprocess(func):
+
+    def _preprocess(
+        self,
+        videos: list[torch.Tensor],
+        do_convert_rgb: bool = True,
+        do_resize: bool = True,
+        size: SizeDict | None = None,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None" = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        image_mean: float | list[float] | None = None,
+        image_std: float | list[float] | None = None,
+        patch_size: int | None = None,
+        temporal_patch_size: int | None = None,
+        merge_size: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ):
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+
+        for shape, stacked_videos in grouped_videos.items():
+            B, T, C, H, W = stacked_videos.shape
+            num_frames, height, width = T, H, W
+            if do_resize:
+                resized_height, resized_width = smart_resize_video(
+                    num_frames=num_frames,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_videos = stacked_videos.view(B * T, C, H, W)
+                stacked_videos = self.resize(
+                    stacked_videos,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+                stacked_videos = stacked_videos.view(
+                    B, T, C, resized_height, resized_width
+                )
+            resized_videos_grouped[shape] = stacked_videos
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+
+        # Group videos by size for further processing
+        # Needed in case do_resize is False, or resize returns videos with different sizes
+        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
+        processed_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            resized_height, resized_width = get_image_size(
+                stacked_videos[0], channel_dim=ChannelDimension.FIRST
+            )
+
+            # Fused rescale and normalize
+            stacked_videos = self.rescale_and_normalize(
+                stacked_videos,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            patches = stacked_videos
+
+            # Check that videos have `num_frames` divisible by `temporal_patch_size`
+            T = patches.shape[1]
+            if pad := -T % temporal_patch_size:
+                repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1)
+                patches = torch.cat((patches, repeats), dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_videos_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_videos = reorder_videos(
+            processed_videos_grouped, grouped_videos_index
+        )
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(processed_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+        data = {
+            "pixel_values_videos": pixel_values_videos,
+            "video_grid_thw": video_grid_thw,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    return _preprocess
+
+
+_npu_preprocess_patched = False
+
+
+def npu_apply_qwen_image_preprocess_patch():
+    global _npu_preprocess_patched
+    if _npu_preprocess_patched:
+        return
+    apply_module_patch(
+        "transformers.models.qwen2_vl.image_processing_qwen2_vl.Qwen2VLImageProcessor",
+        "_preprocess",
+        [npu_wrapper_preprocess],
+    )
+    apply_module_patch(
+        "transformers.models.qwen3_vl.video_processing_qwen3_vl.Qwen3VLVideoProcessor",
+        "_preprocess",
+        [npu_wrapper_video_preprocess],
+    )
+    _npu_preprocess_patched = True

python/sglang/srt/hardware_backend/npu/moe/fuseep.py

@@ -0,0 +1,171 @@
+"""Ascend FuseEP fused dispatch+GEMM+combine forward path.
+
+Follows the mega_moe shape: a free-function bypass invoked from
+``FusedMoE.forward`` when ``--moe-a2a-backend ascend_fuseep`` is set, plus a
+weight-postprocess helper that NPU quant_methods call from their
+``process_weights_after_loading`` when the same backend is selected.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+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.layers.moe.token_dispatcher.deepep import DeepEPBuffer
+from sglang.srt.layers.moe.utils import DeepEPMode
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+    from sglang.srt.layers.moe.topk import TopKOutput
+
+
+_PARAMS_BYTES = 2  # bf16 — Ascend's Dispatch & Combine does not support fp16
+
+
+def _get_fuseep_buffer(layer: FusedMoE):
+    DeepEPBuffer.set_dispatch_mode_as_low_latency()
+    return DeepEPBuffer.get_deepep_buffer(
+        get_tp_group().device_group,
+        layer.hidden_size,
+        _PARAMS_BYTES,
+        DeepEPMode.LOW_LATENCY,
+        envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get(),
+        layer.num_experts,
+    )
+
+
+def forward_fuseep(
+    layer: FusedMoE,
+    hidden_states: torch.Tensor,
+    topk_output: TopKOutput,
+) -> torch.Tensor:
+    buf = _get_fuseep_buffer(layer)
+    hidden_states, _ = buf.fused_deep_moe(
+        hidden_states,
+        topk_idx=topk_output.topk_ids,
+        topk_weights=topk_output.topk_weights,
+        gmm1_permuted_weight=layer.w13_weight,
+        gmm1_permuted_weight_scale=layer.w13_weight_scale,
+        gmm2_weight=layer.w2_weight,
+        gmm2_weight_scale=layer.w2_weight_scale,
+        num_max_dispatch_tokens_per_rank=(
+            envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
+        ),
+        num_experts=layer.num_experts,
+        fuse_mode=envs.SGLANG_NPU_FUSED_MOE_MODE.get(),
+    )
+    return hidden_states
+
+
+def _permute_w13_weight_scale(w: torch.Tensor, tile_n: int) -> torch.Tensor:
+    if tile_n % 2 != 0:
+        raise ValueError(f"tile_n must be even, got {tile_n}")
+
+    *dims, n = w.shape
+    if n % tile_n != 0:
+        raise ValueError(f"Last dimension {n} must be divisible by tile_n {tile_n}")
+
+    w_reshaped = w.reshape(*dims, 2, n // tile_n, tile_n // 2)
+    perm_order = list(range(len(dims))) + [-2, -3, -1]
+    return w_reshaped.permute(perm_order).reshape(*dims, n)
+
+
+def _reshape_w13_weight(
+    weight: torch.Tensor, dim: int, chunk_size: int = 64
+) -> torch.Tensor:
+    # Achieving greater computing power through reshape on Ascend.
+    original_shape = weight.shape
+    if dim < 0:
+        dim += len(original_shape)
+
+    if original_shape[dim] % (2 * chunk_size) != 0:
+        raise ValueError(
+            f"Dimension {dim} size {original_shape[dim]} must be divisible by "
+            f"{2 * chunk_size}"
+        )
+
+    new_shape = (
+        *original_shape[:dim],
+        2,
+        original_shape[dim] // (2 * chunk_size),
+        chunk_size,
+        *original_shape[dim + 1 :],
+    )
+
+    weight = weight.view(new_shape)
+    weight = weight.transpose(dim, dim + 1).contiguous()
+    return weight.view(*original_shape[:dim], -1, *original_shape[dim + 1 :])
+
+
+def _release_weight_cache(weight: torch.Tensor) -> torch.Tensor:
+    # .contiguous() introduces additional memory overhead; release with resize_(0)
+    origin_weight = weight.data.transpose(1, 2)
+    new_weight = origin_weight.contiguous()
+    origin_weight.untyped_storage().resize_(0)
+    return new_weight
+
+
+def _scale_from_float_to_int64(scale: torch.Tensor) -> torch.nn.Parameter:
+    import numpy as np
+
+    converted = torch.from_numpy(
+        np.frombuffer(
+            scale.cpu().to(torch.float32).numpy().tobytes(), dtype=np.int32
+        ).astype(np.int64)
+    ).to(scale.device)
+    return torch.nn.Parameter(converted, requires_grad=False)
+
+
+def process_fuseep_weights(layer: torch.nn.Module) -> None:
+    """Apply the Ascend FuseEP-specific weight layout.
+
+    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.
+    """
+    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)
+
+        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,
+        )

python/sglang/srt/hardware_backend/npu/moe/topk.py

@@ -0,0 +1,113 @@
+from typing import TYPE_CHECKING, Optional
+
+import torch
+from sgl_kernel_npu.norm.l1_norm import l1_norm
+
+from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
+from sglang.srt.eplb.expert_location_dispatch import topk_ids_logical_to_physical
+from sglang.srt.layers.moe.topk import (
+    StandardTopKOutput,
+    capture_routed_experts_if_allowed,
+    select_experts,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
+    from sglang.srt.layers.moe.topk import TopKConfig, TopKOutput
+
+
+def fused_topk_npu(
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    topk_config: "TopKConfig",
+    num_token_non_padded: Optional[torch.Tensor] = None,
+    expert_location_dispatch_info: Optional["ExpertLocationDispatchInfo"] = None,
+    layer_id: Optional[int] = None,
+) -> "TopKOutput":
+
+    use_grouped_topk = topk_config.use_grouped_topk
+    renormalize = topk_config.renormalize
+    correction_bias = topk_config.correction_bias
+
+    # Fast path: simple top-k without grouped routing and bias
+    if not use_grouped_topk and correction_bias is None:
+        topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax(
+            router_logits,
+            k=topk_config.top_k,
+        )
+
+        if renormalize:
+            topk_weights = l1_norm(
+                topk_weights
+                if topk_config.num_fused_shared_experts == 0
+                else topk_weights[:, :-1]
+            )
+        topk_weights = topk_weights.to(torch.float32)
+
+    # sqrtsoftplus (DSV4 noaux_tc): the NPU op only scores sigmoid/softmax, so use
+    # a torch path. top-k over (scores + bias); weights from un-biased scores.
+    elif topk_config.scoring_func == "sqrtsoftplus":
+        scores = torch.nn.functional.softplus(router_logits.float()).sqrt()
+        scores_for_choice = (
+            scores + correction_bias.unsqueeze(0).float()
+            if correction_bias is not None
+            else scores
+        )
+        _, topk_ids = torch.topk(
+            scores_for_choice, k=topk_config.top_k, dim=-1, sorted=False
+        )
+        topk_ids = topk_ids.to(torch.int32)
+        topk_weights = scores.gather(1, topk_ids)
+        if renormalize:
+            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        else:
+            topk_weights = topk_weights * topk_config.routed_scaling_factor
+        topk_weights = topk_weights.to(torch.float32)
+
+    # Support grouped top-k or correction bias or sigmoid or routed_scaling_factor
+    elif (
+        correction_bias is not None
+        or topk_config.scoring_func == "sigmoid"
+        or num_token_non_padded is not None
+    ):
+        topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k(
+            router_logits.to(torch.float32),
+            k=topk_config.top_k,
+            bias=(
+                correction_bias.to(torch.float32)
+                if correction_bias is not None
+                else None
+            ),
+            # num_expert_group and topk_group in some topk_config without group is None, (not supported by this ops)
+            k_group=topk_config.topk_group if use_grouped_topk else 1,
+            group_count=topk_config.num_expert_group if use_grouped_topk else 1,
+            group_select_mode=(1 if use_grouped_topk else 0),
+            renorm=0,
+            # 1 for sigmoid, 0 for softmax
+            norm_type=1,
+            routed_scaling_factor=(
+                1 if renormalize else topk_config.routed_scaling_factor
+            ),
+            eps=float(1e-20),
+        )
+        topk_weights = topk_weights.to(torch.float32)
+
+    # torch native is not yet supported num_token_non_padded
+    # Fallback to torch native implementation
+    else:
+        topk_config.torch_native = True
+        return select_experts(
+            hidden_states=hidden_states,
+            layer_id=layer_id,
+            router_logits=router_logits,
+            topk_config=topk_config,
+            num_token_non_padded=num_token_non_padded,
+            expert_location_dispatch_info=expert_location_dispatch_info,
+        )
+
+    if expert_location_dispatch_info is not None:
+        topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
+    get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
+    capture_routed_experts_if_allowed(topk_config, layer_id, topk_ids)
+
+    return StandardTopKOutput(topk_weights, topk_ids, router_logits)

python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
+    NPUW4A16Int4DynamicMoEMethod,
+)
+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:
+        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 = []
+        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)
+            )
+
+        qweight_tmp.bitwise_xor_(0x88888888)
+
+        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)
+
+        layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
+        layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        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])
+
+        if bias is not None and bias.dtype == torch.bfloat16:
+            bias = bias.float()
+
+        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,
+        )
+
+        return out.reshape(out_shape)
+
+
+class AWQAscendMoEKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.kernel = NPUW4A16Int4DynamicMoEMethod()
+
+    @staticmethod
+    def _register_or_replace_parameter(
+        layer: torch.nn.Module, name: str, tensor: torch.Tensor
+    ) -> None:
+        if hasattr(layer, name):
+            replace_parameter(layer, name, tensor)
+        else:
+            layer.register_parameter(
+                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
+            ),
+        )
+
+        self.kernel.process_weights_after_loading(layer)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> torch.Tensor:
+        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,
+        )

python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py

@@ -0,0 +1,315 @@
+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
+
+
+def unpack_from_int32(
+    weight: torch.Tensor,
+    num_bits: int,
+    packed_dim: int = 1,
+) -> torch.Tensor:
+    """
+    Unpacks quantized weights from int32 format back to original bits.
+
+    :param weight: The packed int32 tensor containing quantized weights
+    :param num_bits: The number of bits used for quantization (<= 8)
+    :param packed_dim: Dimension along which weights are packed (0 or 1), defaults to 1
+    :return: Unpacked tensor with int8 dtype after applying offset correction
+    """
+    assert (
+        weight.dtype == torch.int32
+    ), f"Expecting `weight.dtype` is torch.int32 but got {weight.dtype}."
+    assert (
+        num_bits <= 8
+    ), f"Expecting `num_bits` should not be larger than 8 but got {num_bits}."
+
+    pack_factor = 32 // num_bits
+    mask = (1 << num_bits) - 1
+
+    if packed_dim == 1:
+        unpacked_weight = torch.zeros(
+            (weight.shape[0], weight.shape[1] * pack_factor),
+            device=weight.device,
+            dtype=torch.int32,
+        )
+        for i in range(pack_factor):
+            unpacked_weight[:, i::pack_factor] = (weight >> (num_bits * i)) & mask
+    else:
+        unpacked_weight = torch.zeros(
+            (weight.shape[0] * pack_factor, weight.shape[1]),
+            device=weight.device,
+            dtype=torch.int32,
+        )
+        for i in range(pack_factor):
+            unpacked_weight[i::pack_factor, :] = (weight >> (num_bits * i)) & mask
+    offset = pow(2, num_bits) // 2
+    unpacked_weight = (unpacked_weight - offset).to(torch.int8)
+    return unpacked_weight
+
+
+class GPTQLinearAscendKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        layer.qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                layer.qzeros.data.contiguous(),
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            ).to(layer.scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.qzeros += 1
+
+        qweight_tmp = unpack_from_int32(
+            layer.qweight.data.contiguous(), self.quant_config.weight_bits, packed_dim=0
+        )
+        # use int8 to store weight by default
+        if self.quant_config.weight_bits != 4:
+            layer.qweight = torch.nn.Parameter(
+                qweight_tmp,
+                requires_grad=False,
+            )
+            return
+
+        # 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)),
+            requires_grad=False,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qweight = layer.qweight
+        scales = layer.scales
+        qzeros = layer.qzeros
+
+        reshaped_x = x.reshape(-1, x.shape[-1])
+
+        if bias is not None and bias.dtype == torch.bfloat16:
+            bias = bias.float()
+
+        # 4bit weight is packed to int32(8 x int4)
+        if self.quant_config.weight_bits == 4:
+            out_shape = x.shape[:-1] + (qweight.shape[-1] * 8,)
+        else:
+            out_shape = x.shape[:-1] + (qweight.shape[-1],)
+
+        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,
+        )
+
+        return out.reshape(out_shape)
+
+
+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:
+        w13_qzeros_2d = layer.w13_qzeros.data.contiguous().reshape(
+            -1, layer.w13_qzeros.shape[-1]
+        )
+        layer.w13_qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                w13_qzeros_2d,
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            )
+            .reshape(layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1)
+            .to(layer.w13_scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.w13_qzeros += 1
+
+        w2_qzeros_2d = layer.w2_qzeros.data.contiguous().reshape(
+            -1, layer.w2_qzeros.shape[-1]
+        )
+        layer.w2_qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                w2_qzeros_2d,
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            )
+            .reshape(layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1)
+            .to(layer.w2_scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.w2_qzeros += 1
+
+        w13_qweight_2d = (
+            layer.w13_qweight.data.transpose(-1, -2)
+            .contiguous()
+            .reshape(-1, layer.w13_qweight.shape[-2])
+        )
+        w13_qweight_tmp = unpack_from_int32(
+            w13_qweight_2d, self.quant_config.weight_bits, packed_dim=1
+        )
+
+        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)
+
+            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_()
+
+            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
+                    )
+                    .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,
+            )
+        # use int8 to store weight by default
+        else:
+            layer.w13_qweight = torch.nn.Parameter(
+                w13_qweight_tmp.reshape(
+                    layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
+                )
+                .transpose(-1, -2)
+                .contiguous(),
+                requires_grad=False,
+            )
+
+        w2_qweight_2d = (
+            layer.w2_qweight.data.transpose(-1, -2)
+            .contiguous()
+            .reshape(-1, layer.w2_qweight.shape[-2])
+        )
+        w2_qweight_tmp = unpack_from_int32(
+            w2_qweight_2d, self.quant_config.weight_bits, packed_dim=1
+        )
+
+        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)
+
+            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_()
+
+            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
+                    )
+                    .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,
+            )
+        # use int8 to store weight by default
+        else:
+            layer.w2_qweight = torch.nn.Parameter(
+                w2_qweight_tmp.reshape(
+                    layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
+                )
+                .transpose(-1, -2)
+                .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)

python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py

@@ -0,0 +1,622 @@
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.srt.hardware_backend.npu.utils import NPUACLFormat, npu_format_cast
+from sglang.srt.layers.quantization.base_config import LinearMethodBase
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+logger = logging.getLogger(__name__)
+
+MXFP8_BLOCK_SIZE = 32
+# W4A8_MXFP block (group) size — fixed at 32 by the msmodelslim export format.
+MXFP4_BLOCK_SIZE = 32
+
+
+# NPU ops are reached via torch.ops.npu.* (registered when torch_npu is imported
+# by the runtime), so this module needs no top-level `import torch_npu` and stays
+# importable on CUDA/CPU/AMD/XPU CI.
+def _get_float8_e8m0fnu_dtype():
+    # Resolve lazily rather than as a module-level constant: this module is
+    # imported early (during quant-scheme registration), so reading the dtype at
+    # call time keeps it correct regardless of import order / platform.
+    return getattr(torch, "float8_e8m0fnu", None)
+
+
+def _get_float4_e2m1fn_x2_dtype():
+    # The packed-FP4 dtype MUST come from torch_npu (an int enum, e.g. 296), not
+    # from torch. The NPU ops that consume it -- npu_dynamic_mx_quant(dst_type=),
+    # npu_quant_matmul(x2_dtype=), npu_format_cast(input_dtype=) -- REJECT the
+    # torch dtype object torch.float4_e2m1fn_x2 in op-plugin on recent torch_npu
+    # builds (it raises, or with None gives "output y must be same shape as input
+    # x"), even though torch.float4_e2m1fn_x2 exists. This is fp4-specific: fp8 /
+    # float8_e8m0fnu is accepted from torch either way. Verified on A5 /
+    # torch_npu 2.10.0.post2.dev20260704 (see llm/probe_fp4_w4a8_chain.py: dst=296
+    # passes the full quant->format_cast->matmul chain, dst=torch dtype fails).
+    #
+    # Lazy import so this NPU-only path keeps the module importable on
+    # CUDA/CPU/AMD/XPU CI (no top-level torch_npu; see AGENTS.md known pitfalls).
+    from sglang.srt.utils import is_npu
+
+    if is_npu():
+        import torch_npu
+
+        npu_dtype = getattr(torch_npu, "float4_e2m1fn_x2", None)
+        if npu_dtype is not None:
+            return npu_dtype
+    return getattr(torch, "float4_e2m1fn_x2", None)
+
+
+class _NPULinearMethodBase(LinearMethodBase):
+
+    def __init__(
+        self,
+        quant_config: Optional["QuantizationConfig"] = None,
+    ):
+        self.quant_config = quant_config
+
+
+class NPUW8A8Int8LinearMethod(_NPULinearMethodBase):
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight.data = npu_format_cast(layer.weight.data)
+
+        layer.weight_scale.data = layer.weight_scale.data.flatten()
+        # Compressed-tensors format doesn't have this field
+        if hasattr(layer, "weight_offset"):
+            layer.weight_offset.data = layer.weight_offset.data.flatten()
+
+        expanding_factor = layer.weight.data.shape[0]
+        layer.aclnn_input_scale = torch.nn.Parameter(
+            layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
+            requires_grad=False,
+        )
+        layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter(
+            layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
+            requires_grad=False,
+        )
+        layer.aclnn_input_offset = torch.nn.Parameter(
+            layer.input_offset.data.repeat(expanding_factor).to(device="npu"),
+            requires_grad=False,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        from sglang.srt.layers.linear import RowParallelLinear
+
+        original_dtype = x.dtype
+        if original_dtype != torch.int8:
+            x = torch.ops.npu.npu_quantize(
+                x,
+                layer.aclnn_input_scale_reciprocal,
+                layer.aclnn_input_offset,
+                torch.qint8,
+                -1,
+                False,
+            )
+        # Only fuse bias add into GEMM for rank 0 (this ensures that
+        # bias will not get added more than once in Attention TP>1 case)
+        if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0:
+            quant_bias = None
+        else:
+            quant_bias = layer.quant_bias
+        return torch.ops.npu.npu_quant_matmul(
+            x,
+            layer.weight,
+            layer.deq_scale,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight.data = npu_format_cast(layer.weight.data)
+
+        layer.weight_scale.data = layer.weight_scale.data.flatten()
+        # Compressed-tensors format doesn't have this field
+        if hasattr(layer, "weight_offset"):
+            layer.weight_offset.data = layer.weight_offset.data.flatten()
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if isinstance(x, tuple):
+            """dynamic_scale is calculated in malprolog kernel"""
+            original_dtype = torch.bfloat16
+            quant_out, dynamic_scale = x
+        else:
+            original_dtype = x.dtype
+            quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(x)
+        return torch.ops.npu.npu_quant_matmul(
+            quant_out,
+            layer.weight,
+            layer.weight_scale,
+            pertoken_scale=dynamic_scale.flatten(),
+            bias=bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUMXFP8LinearMethod(_NPULinearMethodBase):
+    """Ascend NPU MXFP8 linear method for LLM (SRT) models.
+
+    Shared kernel for both the online config path (``--quantization mxfp8``) and
+    the offline ModelSlimMXFP8Scheme (which delegates to this as ``self.kernel``).
+    process_weights_after_loading branches on weight dtype: FP16/BF16 weights are
+    quantised to MXFP8 at load time (online); pre-quantised float8_e4m3fn weights
+    are only re-laid-out (offline). Inference: dynamic MXFP8 activation quant +
+    MXFP8 matmul (block_size=32).
+    """
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes,
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        from sglang.srt.layers.parameter import ModelWeightParameter
+
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # Load weights in original dtype; quantise later in process_weights_after_loading
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data
+        if weight.dtype == torch.float8_e4m3fn:
+            # Offline (ModelSlim) path: weight is already MXFP8-quantised and
+            # layer.weight_scale holds the uint8 block scales [out, in/32]. Only
+            # re-layout to [in, out] / [in//64, out, 2] strided views below.
+            n_dim, k_dim = layer.weight_scale.data.shape
+            scale = layer.weight_scale.data.reshape(n_dim, k_dim // 2, 2)
+            layer.weight = Parameter(weight.transpose(0, 1), requires_grad=False)
+            layer.weight_scale_inv = Parameter(
+                scale.transpose(0, 1), requires_grad=False
+            )
+            # weight_scale is now folded into weight_scale_inv (which keeps the
+            # underlying storage alive via its view); drop the stale parameter so
+            # it doesn't linger in named_parameters() / state_dict().
+            del layer.weight_scale
+        else:
+            # Online path: quantise FP16/BF16 weights to MXFP8 at load time.
+            if weight.dtype not in (torch.float16, torch.bfloat16):
+                logger.warning(
+                    "NPUMXFP8LinearMethod: weight dtype %s is not float16/bfloat16; "
+                    "casting to bfloat16 before MXFP8 quantisation.",
+                    weight.dtype,
+                )
+                weight = weight.to(torch.bfloat16)
+            # Move weight to NPU if needed (cpu offload may move it back to CPU).
+            if not weight.is_npu:
+                weight = weight.to(f"npu:{torch.npu.current_device()}")
+            # Online MXFP8 quantisation of weights (block_size=32).
+            # qw: [out, in] float8_e4m3fn, w_scale: [out, in//64, 2] uint8.
+            qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant(
+                weight, dst_type=torch.float8_e4m3fn
+            )
+            layer.weight = Parameter(qw.transpose(0, 1), requires_grad=False)
+            layer.weight_scale_inv = Parameter(
+                w_scale.transpose(0, 1), requires_grad=False
+            )
+
+        # Both paths produce weight [in, out] and weight_scale_inv [in//64, out,
+        # 2] as strided transpose views — DO NOT call .contiguous(). The matmul
+        # reduction loop scans the in-dim per output column; the [out, in]
+        # row-major source gives stride-1 access for that scan via the transpose
+        # view (matches msmodelslim's offline layout and vllm-ascend's
+        # AscendW8A8MXFP8DynamicLinearMethod). Calling .contiguous() physically
+        # reorders to [in, out] row-major, making the inner-loop stride = out and
+        # tanking HBM bandwidth.
+
+        # Cache FP32 bias once to avoid a per-forward dtype conversion + alloc.
+        if (
+            getattr(layer, "bias", None) is not None
+            and layer.bias.dtype != torch.float32
+        ):
+            layer.bias_fp32 = Parameter(
+                layer.bias.data.to(torch.float32), requires_grad=False
+            )
+        else:
+            layer.bias_fp32 = None
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation
+        qx, input_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        # MXFP8 matmul (weight & scale already transposed at load time)
+        # Use the cached FP32 bias from process_weights_after_loading; fall back
+        # to per-call conversion if the cache was bypassed (e.g. dynamic bias).
+        if bias is None:
+            quant_bias = None
+        elif (
+            bias is getattr(layer, "bias", None)
+            and getattr(layer, "bias_fp32", None) is not None
+        ):
+            quant_bias = layer.bias_fp32
+        else:
+            quant_bias = bias.to(torch.float32)
+
+        e8m0_dtype = _get_float8_e8m0fnu_dtype()
+        output = torch.ops.npu.npu_quant_matmul(
+            qx,
+            layer.weight,
+            layer.weight_scale_inv,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=input_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+            group_sizes=[1, 1, MXFP8_BLOCK_SIZE],
+        )
+
+        # Restore original shape (replace last dim with output features)
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)
+
+
+class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):
+
+    def process_weights_after_loading(self, layer):
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        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)
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        tp_rank: Optional[int] = 0,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(
+            x, dst_type=torch.quint4x2
+        )
+        return torch.ops.npu.npu_quant_matmul(
+            quant_out,
+            layer.weight,
+            layer.weight_scale,
+            pertoken_scale=dynamic_scale.flatten(),
+            bias=bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
+    """Ascend NPU W4A8 online quantization: MXFP4 weights + MXFP8 activations.
+
+    This is a *true* W4(weight) A8(activation) path: it mirrors the offline
+    ``W4A8_MXFP`` kernel (``NPUMXFP4W4A8OfflineLinearMethod``) exactly — the only
+    difference is that the FP4 weights are produced online from BF16/FP16
+    (round-to-nearest, no calibration) instead of being loaded from a msmodelslim
+    checkpoint. An earlier version of this method ran a *dual-level* scheme that
+    also compressed the activation to FP4 (W4A4 compute via
+    ``npu_dual_level_quant_matmul``); that was a large accuracy regression — 4-bit
+    activations — so it was replaced with the single-level FP8-activation path
+    below, aligned with the offline W4A8 implementation.
+
+    Weight quantization (process_weights_after_loading):
+        BF16/FP16 weight → npu_dynamic_mx_quant(dst=float4_e2m1fn_x2) → packed FP4
+        + UE8M0 block scale → npu_format_cast to FRACTAL_NZ → transpose [in//2, out]
+
+    Inference (apply):
+        BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn)  (A8, FP8)
+        → npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block])
+
+    Hardware: Ascend 950 (A5) + a recent torch_npu with the FP4 npu_quant_matmul
+    (same requirement as the offline W4A8 path — see that class's docstring).
+    """
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes,
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """Register an unquantized (``params_dtype``) weight placeholder.
+
+        Online quantization needs its own ``create_weights`` because the
+        checkpoint still holds full-precision BF16/FP16 weights: the loader
+        fills this buffer, then ``process_weights_after_loading`` quantizes it to
+        MXFP4 in place. This differs from the offline/int8 methods, whose weights
+        are created by the scheme's own ``create_weights`` to match the
+        already-quantized (FP8 / uint8-packed) layout the checkpoint provides.
+        """
+        from sglang.srt.layers.parameter import ModelWeightParameter
+
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # Load weights in original dtype; quantise to MXFP4 in
+        # process_weights_after_loading.
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Online single-level MXFP4 weight quant, then lay the weight out exactly
+        # like the offline W4A8 path so the same npu_quant_matmul(x2_dtype=fp4)
+        # kernel accepts it. All NPU ops go through torch.ops.npu.* (no torch_npu).
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        weight_fp = layer.weight.data
+        if weight_fp.dtype not in (torch.float16, torch.bfloat16):
+            weight_fp = weight_fp.to(torch.bfloat16)
+        # Move to NPU if needed (cpu offload may have put it on CPU).
+        if not weight_fp.is_npu:
+            weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
+
+        # BF16 -> packed FP4 (float4_e2m1fn_x2, [out, in//2]) + UE8M0 block scale.
+        # npu_dynamic_mx_quant returns the scale as [out, in//64, 2] (3D); older
+        # builds may return [out, in//32] (2D) — handle both before the transpose.
+        qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            weight_fp, dst_type=fp4_dtype, round_mode="round"
+        )
+
+        # weight: packed FP4 -> FRACTAL_NZ (float8_e4m3fn view) -> transpose
+        # [in//2, out]. Mirror the offline path (no .contiguous() on the NZ view);
+        # view as uint8 first because npu_format_cast only accepts int-dtype tensors.
+        qw_nz = npu_format_cast(
+            qw.view(torch.uint8),
+            NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+            customize_dtype=torch.float8_e4m3fn,
+            input_dtype=fp4_dtype,
+        )
+        layer.weight = Parameter(qw_nz.transpose(-1, -2), requires_grad=False)
+
+        # weight_scale -> [in//64, out, 2] to match npu_quant_matmul.
+        if w_scale.dim() == 2:
+            n, k = w_scale.shape
+            w_scale = w_scale.reshape(n, k // 2, 2)
+        layer.weight_scale = Parameter(w_scale.transpose(-3, -2), requires_grad=False)
+
+        # Cache FP32 bias once to avoid a per-forward dtype conversion + alloc.
+        if (
+            getattr(layer, "bias", None) is not None
+            and layer.bias.dtype != torch.float32
+        ):
+            layer.bias_fp32 = Parameter(
+                layer.bias.data.to(torch.float32), requires_grad=False
+            )
+        else:
+            layer.bias_fp32 = None
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        e8m0_dtype = _get_float8_e8m0fnu_dtype()
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant.
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation (A8 — FP8, not FP4).
+        quantized_x, dynamic_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        # Use the cached FP32 bias from process_weights_after_loading; fall back
+        # to per-call conversion if the cache was bypassed (e.g. dynamic bias).
+        if bias is None:
+            quant_bias = None
+        elif (
+            bias is getattr(layer, "bias", None)
+            and getattr(layer, "bias_fp32", None) is not None
+        ):
+            quant_bias = layer.bias_fp32
+        else:
+            quant_bias = bias.to(torch.float32)
+
+        # True W4(weight)A8(activation) matmul, identical to the offline path.
+        output = torch.ops.npu.npu_quant_matmul(
+            quantized_x,
+            layer.weight,
+            layer.weight_scale,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=dynamic_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+            x2_dtype=fp4_dtype,
+            group_sizes=[0, 0, MXFP4_BLOCK_SIZE],
+        )
+
+        # Restore original shape (replace last dim with output features).
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)
+
+
+class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):
+    """Ascend NPU offline W4A8 (ModelSlim ``W4A8_MXFP``): packed-FP4 weights + MXFP8 activations.
+
+    Kernel for the offline ModelSlimMXFP4W4A8Scheme (delegated as ``self.kernel``).
+    The msmodelslim ``W4A8_MXFP`` checkpoint stores weights as *packed FP4*
+    (``pack_fp4_to_uint8`` → ``uint8`` shape ``[out, in//2]``) plus UE8M0 block
+    scales (``uint8`` shape ``[out, in//group_size]``):
+
+      process_weights_after_loading:
+        weight (uint8 packed FP4 [out, in//2]) → npu_format_cast(29,
+            customize_dtype=float8_e4m3fn, input_dtype=float4_e2m1fn_x2) → FRACTAL_NZ
+            → transpose [in//2, out]
+        weight_scale [out, in/32] → reshape [out, in/64, 2] → transpose → [in/64, out, 2]
+
+      apply:
+        BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn)  (A8, MXFP8)
+        → npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block])
+
+    Mirrors vllm-ascend ``AscendW4A8MXFPDynamicLinearMethod`` exactly (Ascend 950/A5).
+    The weight is cast to FRACTAL_NZ then transposed; ``npu_dynamic_mx_quant`` already
+    returns a 3D ``[tokens, in//64, 2]`` block scale so the matmul needs no extra
+    scale-layout normalization.
+
+    ⚠️ REQUIRES a recent torch_npu build for the FP4 ``npu_quant_matmul``. On the
+    A5 this device forces ``allow_internal_format=False`` (the NZ cast still produces
+    a ``FRACTAL_NZ_C0_16`` tensor, which is fine). Older torch_npu (e.g.
+    ``2.10.0.dev20260320``) had a broken FP4 matmul that rejected the NZ weight in
+    *prefill* with ``x2 should be in ... nz format, but it is 2``;
+    ``2.10.0.post1.dev20260624`` (and later) runs the vllm-aligned NZ path
+    correctly. If you hit ``it is 2``, update torch_npu — do NOT "fix" it by
+    switching the weight to ND.
+
+    ⚠️ A ``atb::OperationSetup`` *segfault during decode* (not prefill) is a
+    DIFFERENT, unrelated issue: it is the eager-decode ``ascend`` attention
+    backend, NOT this matmul (verified by stage-sync bisection — qkv's matmul
+    syncs clean, the fault surfaces at the entry-sync of the next layer, i.e. the
+    decode attention between qkv and o_proj). Run with the NPU decode graph (do
+    NOT pass ``--disable-cuda-graph``); graph mode is the NPU default and what
+    vllm uses. This attention issue is model-agnostic and out of scope for W4A8.
+
+    This is a true W4(weight) A8(activation) single-level matmul. The *online*
+    ``NPUMXFP4W4A8LinearMethod`` now uses this exact apply path — the only
+    difference is that it quantizes BF16/FP16 weights to FP4 at load time instead
+    of loading them from a msmodelslim checkpoint. ``group_size`` is fixed at 32
+    by the ``W4A8_MXFP`` export format.
+    """
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Mirror vllm-ascend AscendW4A8MXFPDynamicLinearMethod: cast the packed-FP4
+        # weight to FRACTAL_NZ then transpose. All NPU ops go through
+        # torch.ops.npu.* (no torch_npu). Requires a recent torch_npu build (see
+        # class docstring): older builds reject the NZ weight ("x2 ... it is 2").
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        # weight: packed-FP4 uint8 [out, in//2] -> FRACTAL_NZ (float8_e4m3fn view)
+        # -> transpose to [in//2, out].
+        layer.weight.data = npu_format_cast(
+            layer.weight.data,
+            NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+            customize_dtype=torch.float8_e4m3fn,
+            input_dtype=fp4_dtype,
+        )
+        layer.weight.data = layer.weight.data.transpose(-1, -2)
+        # weight_scale: [out, in/32] uint8 -> [in/64, out, 2].
+        n, k = layer.weight_scale.data.shape
+        layer.weight_scale.data = layer.weight_scale.data.reshape(
+            n, k // 2, 2
+        ).transpose(-3, -2)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        e8m0_dtype = _get_float8_e8m0fnu_dtype()
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant.
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation (A8).
+        quantized_x, dynamic_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        if bias is not None and bias.dtype != torch.float32:
+            bias = bias.to(torch.float32)
+
+        # W4(weight)A8(activation) matmul, mirroring vllm-ascend exactly.
+        output = torch.ops.npu.npu_quant_matmul(
+            quantized_x,
+            layer.weight,
+            layer.weight_scale,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=dynamic_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=bias,
+            output_dtype=original_dtype,
+            x2_dtype=fp4_dtype,
+            group_sizes=[0, 0, MXFP4_BLOCK_SIZE],
+        )
+
+        # Restore original shape (replace last dim with output features).
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)

python/sglang/srt/hardware_backend/npu/utils.py

@@ -0,0 +1,366 @@
+import functools
+import logging
+import sys
+from enum import IntEnum
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_npu_memory_capacity, is_npu
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+_is_npu = is_npu()
+indexer_weight_stream = None
+gva_is_inited = False
+
+
+class NPUACLFormat(IntEnum):
+    ACL_FORMAT_UNDEFINED = -1
+    ACL_FORMAT_ND = 2
+    ACL_FORMAT_FRACTAL_NZ = 29
+
+
+class FusedMoEMode(IntEnum):
+    FUSED_DEEP_MOE = 1
+    DISPATCH_FFN_COMBINE = 2
+
+
+def _call_once(fn: Callable):
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        if getattr(fn, "_has_been_called", False):
+            logger.debug("Function {} has already been called.", fn.__name__)
+            return
+
+        fn._has_been_called = True
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+def set_default_server_args(args: "ServerArgs"):
+    """
+    Set default server arguments for NPU backend.
+    """
+
+    # NPU only works with "ascend" attention backend for now
+    args.attention_backend = "ascend"
+    args.prefill_attention_backend = "ascend"
+    args.decode_attention_backend = "ascend"
+    if args.page_size is None:
+        args.page_size = 128
+
+    # NPU memory settings
+    decode = args.cuda_graph_config.decode
+    npu_mem = get_npu_memory_capacity()
+    if npu_mem <= 32 * 1024:
+        # Ascend 910B4,910B4_1
+        # (chunked_prefill_size 4k, max_bs 16 if tp < 4 else 64)
+        if args.chunked_prefill_size is None:
+            args.chunked_prefill_size = 4 * 1024
+        if decode.max_bs is None:
+            if args.tp_size < 4:
+                decode.max_bs = 16
+            else:
+                decode.max_bs = 64
+    elif npu_mem <= 64 * 1024:
+        # Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362
+        # (chunked_prefill_size 8k, max_bs 64 if tp < 4 else 256)
+        if args.chunked_prefill_size is None:
+            args.chunked_prefill_size = 8 * 1024
+        if decode.max_bs is None:
+            if args.tp_size < 4:
+                decode.max_bs = 64
+            else:
+                decode.max_bs = 256
+
+    # NPU does not support CustomAllReduce
+    args.disable_custom_all_reduce = True
+
+    # handles hierarchical cache configs
+    if args.enable_hierarchical_cache:
+        args.hicache_io_backend = "kernel_ascend"
+        if args.use_mla_backend():
+            args.hicache_mem_layout = "page_first_kv_split"
+        else:
+            args.hicache_mem_layout = "page_first_direct"
+
+
+@_call_once
+def init_npu_backend():
+    """
+    Initialize NPU backend. This function should be called only once.
+    """
+
+    assert _is_npu, "NPU backend initialization called on non-NPU device."
+
+    try:
+        import custom_ops  # noqa: F401
+        import sgl_kernel_npu  # noqa: F401
+    except ImportError as e:
+        logger.warning("NPU custom kernel packages unavailable: %s", e)
+
+    import torch_npu
+    from torch_npu.contrib import transfer_to_npu  # noqa: F401
+
+    # Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
+    torch.cuda.is_available = lambda: False
+
+    torch_npu.npu.config.allow_internal_format = True
+    torch_npu.npu.set_compile_mode(jit_compile=False)
+
+
+def _is_nz_aligned(tensor: torch.Tensor) -> bool:
+    """Check whether the last two dims satisfy FRACTAL_NZ alignment rules.
+
+    Ascend FRACTAL_NZ requires:
+      BF16 / FP16 : both dims divisible by 16
+      INT8         : k % 16 == 0  and  n % 32 == 0
+      INT4         : k % 16 == 0  and  n % 64 == 0
+      FP4          : both dims divisible by 64
+    """
+    if tensor.dim() < 2:
+        return False
+    k, n = tensor.shape[-2], tensor.shape[-1]
+    if tensor.dtype in (torch.bfloat16, torch.float16):
+        return k % 16 == 0 and n % 16 == 0
+    if tensor.dtype == torch.int8:
+        return k % 16 == 0 and n % 32 == 0
+    if tensor.dtype in (torch.uint8, torch.int32):
+        # INT4 is typically packed into uint8/int32; be conservative
+        return k % 16 == 0 and n % 64 == 0
+    return True
+
+
+def npu_format_cast(
+    tensor: torch.Tensor,
+    acl_format: NPUACLFormat = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+    *,
+    customize_dtype=None,
+    input_dtype=None,
+) -> torch.Tensor:
+    """
+    Cast a tensor to a specific NPU ACL format.
+
+    Args:
+        tensor (torch.Tensor): The input tensor.
+        acl_format (NPUACLFormat): The target NPU ACL format.
+        customize_dtype / input_dtype: packed-FP4 unpack kwargs (e.g.
+            ``customize_dtype=torch.float8_e4m3fn``,
+            ``input_dtype=torch.float4_e2m1fn_x2``). When either is set the unpack
+            kwargs are forwarded to the op and the ``_is_nz_aligned`` ND fallback
+            is skipped: the FP4 matmul strictly requires FRACTAL_NZ, so a silent
+            ND fallback would corrupt results.
+
+    Returns:
+        torch.Tensor: The tensor cast to the specified NPU ACL format.
+    """
+
+    if not _is_npu:
+        return tensor
+
+    if envs.SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT.get():
+        return tensor
+
+    if tensor.device == torch.device("cpu"):
+        logger.warning_once(
+            "Warning: The conversion from 'ND' to 'NZ' does not work on the CPU. "
+            "Please disable offloading, otherwise the performance will be "
+            "significantly reduced. --dit-cpu-offload false"
+        )
+        return tensor
+
+    # Skip format cast for meta tensors (used in offloader)
+    if tensor.device.type == "meta":
+        return tensor
+
+    # Packed-FP4 → FRACTAL_NZ: forward the unpack kwargs to the op, and skip the
+    # _is_nz_aligned ND fallback — the FP4 matmul strictly requires NZ, so a
+    # silent ND fallback would corrupt results.
+    if customize_dtype is not None or input_dtype is not None:
+        return torch.ops.npu.npu_format_cast(
+            tensor,
+            int(acl_format),
+            customize_dtype=customize_dtype,
+            input_dtype=input_dtype,
+        )
+
+    if acl_format == NPUACLFormat.ACL_FORMAT_FRACTAL_NZ and not _is_nz_aligned(tensor):
+        k, n = tensor.shape[-2], tensor.shape[-1]
+        logger.warning_once(
+            "Skipping FRACTAL_NZ format cast: tensor shape (%d, %d) dtype %s "
+            "is not aligned to NZ requirements. Falling back to 'ND' format, "
+            "which may reduce NPU performance.",
+            k,
+            n,
+            tensor.dtype,
+        )
+        return tensor
+
+    return torch.ops.npu.npu_format_cast(tensor, acl_format.value)
+
+
+def get_indexer_weight_stream():
+    global indexer_weight_stream
+    if indexer_weight_stream is None:
+        indexer_weight_stream = torch.npu.Stream()
+    return indexer_weight_stream
+
+
+def init_zbal(world_size, gpu_id, world_rank, do_check=True):
+    """
+    init zbal, if is mix alloc mode, only register for sma & comm
+    """
+    zbal_mem_size = envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get()
+    if not zbal_mem_size > 0:
+        return 1
+
+    global gva_is_inited
+    from zbal import is_mix_alloc, switch_to_allocator, zbal_init
+
+    if is_mix_alloc():
+        switch_to_allocator()
+        # use lazy init for mix alloc
+        return 1
+    else:
+        if envs.SGLANG_ZBAL_BOOTSTRAP_URL.get():
+            ret = zbal_init(
+                world_size,
+                gpu_id,
+                world_rank,
+                zbal_mem_size * (1024**2),
+                ip_port=envs.SGLANG_ZBAL_BOOTSTRAP_URL.get(),
+            )
+        else:
+            ret = zbal_init(world_size, gpu_id, world_rank, zbal_mem_size * (1024**2))
+
+        gva_is_inited = True
+
+        if do_check and not ret:
+            logger.error("[ZBAL] zbal init failed!")
+            sys.exit(-1)
+
+        return ret
+
+
+def lazy_init_zbal_gva_mem(
+    device, gpu_id, world_rank, world_size, cpu_group=None, do_check=True
+):
+    """
+    lazy init zbal gva mem, keep weights and kv remains alloc by dma vmm to avoid memory fragment
+    """
+    from zbal import is_mix_alloc, zbal_init
+
+    if not is_mix_alloc():
+        logger.info(
+            "lazy init is supported only in mix alloc mode, this action will be passed"
+        )
+        return 1
+
+    global gva_is_inited
+    from sglang.srt.utils.common import get_available_gpu_memory
+
+    # TODO need to use allgather if you want use total_memory stats from mem_get_info as unbalance os
+    total_memory = 61.2  # 2.5GB for other (workspace & os) outside torch
+    free_gpu_memory = get_available_gpu_memory(
+        device,
+        gpu_id,
+        distributed=world_size > 1,
+        cpu_group=cpu_group,
+        empty_cache=True,
+    )
+
+    used_memory = total_memory - free_gpu_memory
+
+    used_memory_in_mb = int(used_memory * 1024)
+    gva_in_mb = envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() - used_memory_in_mb
+    gva_in_mb = gva_in_mb - gva_in_mb % 128  # align to 128MB
+    print(f"[ZBAL] rank {world_rank} allocated {gva_in_mb} MB gva space.")
+
+    assert not gva_is_inited, "zbal gva should be inited only once"
+    # zbal_set_logger_level(0)
+    if envs.SGLANG_ZBAL_BOOTSTRAP_URL.get():
+        res = zbal_init(
+            world_size,
+            gpu_id,
+            world_rank,
+            gva_in_mb * (1024**2),
+            ip_port=envs.SGLANG_ZBAL_BOOTSTRAP_URL.get(),
+        )
+    else:
+        res = zbal_init(world_size, gpu_id, world_rank, gva_in_mb * (1024**2))
+
+    gva_is_inited = True
+    if do_check and not res:
+        logger.error("[ZBAL] zbal lazy init failed!")
+        sys.exit(-1)
+    return res
+
+
+share_stream = None
+routed_stream = None
+
+
+def get_share_stream():
+    global share_stream
+    return share_stream
+
+
+def set_share_stream(stream):
+    global share_stream
+    share_stream = stream
+    # TODO LKL: set stream limit has impact on precision
+    # torch.npu.set_stream_limit(share_stream, 8, 16)
+
+
+def get_routed_stream():
+    global routed_stream
+    return routed_stream
+
+
+def set_routed_stream(stream):
+    global routed_stream
+    routed_stream = stream
+    # TODO LKL: set stream limit has impact on precision
+    # torch.npu.set_stream_limit(routed_stream, 16, 32)
+
+
+def wait_share_stream():
+    stream = get_share_stream()
+    if stream is not None:
+        cur_stream = torch.get_device_module().current_stream()
+        cur_stream.wait_stream(stream)
+
+
+def wait_routed_stream():
+    stream = get_routed_stream()
+    if stream is not None:
+        cur_stream = torch.get_device_module().current_stream()
+        cur_stream.wait_stream(stream)
+
+
+def process_shared_expert(hidden_states, forward_func):
+    stream = get_share_stream()
+    if stream is None:
+        stream = torch.get_device_module().Stream()
+        set_share_stream(stream)
+    stream.wait_stream(torch.get_device_module().current_stream())
+    with torch.get_device_module().stream(stream):
+        shared_output = forward_func(hidden_states)
+    return shared_output
+
+
+def process_routed_expert(hidden_states, topk_output, forward_func):
+    stream = get_routed_stream()
+    if stream is None:
+        stream = torch.get_device_module().Stream()
+        set_routed_stream(stream)
+    stream.wait_stream(torch.get_device_module().current_stream())
+    with torch.get_device_module().stream(stream):
+        shared_output = forward_func(hidden_states, topk_output)
+    return shared_output

python/sglang/srt/hardware_backend/xpu/__init__.py

@@ -0,0 +1 @@
+# XPU (Intel GPU) hardware backend

python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py

@@ -0,0 +1,100 @@
+"""FullXPUGraphBackend — Intel XPU full-graph capture (torch.xpu.XPUGraph).
+
+Mirrors FullCudaGraphBackend with XPU-specific differences:
+  - Captures via torch.xpu.graph(xpu_graph=...) into torch.xpu.XPUGraph.
+  - Shares the global graph memory pool with the prefill backend so that
+    decode + prefill graphs occupy max(decode, prefill) rather than their sum.
+  - No set_graph_pool_id: SymmetricMemoryContext is never triggered on XPU
+    (oneCCL has no ncclMemAlloc equivalent; enable_symm_mem defaults False).
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import torch
+
+from sglang.srt.model_executor.runner.shape_key import ShapeKey
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+
+
+class FullXPUGraphBackend(BaseCudaGraphBackend):
+    """One torch.xpu.XPUGraph per shape for Intel XPU devices."""
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+    ) -> None:
+        self._graphs: Dict[Any, torch.xpu.XPUGraph] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = cuda_graph_runner.model_runner.tp_group
+        self._capture_stream: Optional[torch.xpu.Stream] = None
+
+    @contextmanager
+    def capture_session(self, stream: torch.xpu.Stream):
+        if self._pool is None:
+            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = torch.xpu.XPUGraph()
+
+        with self._device_module.graph(
+            xpu_graph=graph, pool=self._pool, stream=self._capture_stream
+        ):
+            out = forward_fn()
+
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None

python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py

@@ -0,0 +1,169 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Run the model with xpu graph and torch.compile."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+from torch.profiler import ProfilerActivity, profile
+
+from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+from sglang.srt.utils import register_xpu_device_properties_for_dynamo
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+_fake_ops_registered = False
+
+
+def register_fake_ops():
+    """Register fake/abstract implementations for XPU sgl_kernel ops so that
+    torch.compile (Dynamo) can trace through them using FakeTensors for shape
+    and dtype propagation, without executing the real GPU kernels.
+    """
+    global _fake_ops_registered
+    if _fake_ops_registered:
+        return
+    _fake_ops_registered = True
+
+    @torch.library.register_fake("sgl_kernel::fwd")
+    def _(
+        q,
+        k,
+        v,
+        q_v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        page_table,
+        kv_batch_idx,
+        leftpad_k,
+        rotary_cos,
+        rotary_sin,
+        seqlens_rotary,
+        q_descale,
+        k_descale,
+        v_descale,
+        softmax_scale,
+        sinks,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        is_rotary_interleaved,
+        scheduler_metadata,
+        num_kv_splits,
+        pack_gqa,
+        sm_margin,
+        out=None,
+    ):
+        total_q = q.shape[0]
+        num_heads_q = q.shape[1]
+        head_size_v = v.shape[-1]
+        if out is None:
+            out = q.new_empty(total_q, num_heads_q, head_size_v)
+        softmax_lse = q.new_empty(num_heads_q, total_q, dtype=torch.float32)
+        # out_accum and softmax_lse_accum are intermediate split-kv buffers;
+        # they are only read when num_kv_splits > 1, which is determined at
+        # runtime.  Return empty tensors with correct rank so downstream ops
+        # that index into the list do not fail shape propagation.
+        out_accum = q.new_empty(0)
+        softmax_lse_accum = q.new_empty(0, dtype=torch.float32)
+        return (out, softmax_lse, out_accum, softmax_lse_accum)
+
+    @torch.library.register_fake("sgl_kernel::flash_mla_decode")
+    def _(
+        out,
+        q_nope,
+        q_pe,
+        kv_c_and_k_pe_cache,
+        seq_lens,
+        page_table,
+        workspace,
+        sm_scale,
+        num_kv_splits,
+    ):
+        return
+
+
+class XPUGraphRunner(DecodeCudaGraphRunner):
+    """A XPUGraphRunner runs the forward pass of a model with xpu graph and torch.compile."""
+
+    @staticmethod
+    def _apply_xpu_compile_config() -> None:
+        """Apply XPU-specific torch.compile / dynamo settings.
+
+        Called unconditionally before super().__init__() so that the settings
+        are in place regardless of whether --enable-torch-compile is passed.
+        The critical flag is suppress_errors: when the Intel IGC compiler
+        crashes with SIGFPE on certain reduction kernels (ocloc -device bmg
+        returns exit code 245), dynamo falls back to eager for that subgraph
+        instead of propagating the crash.
+        """
+        import torch._dynamo.config
+
+        torch._dynamo.config.suppress_errors = True
+
+    def __init__(self, model_runner: ModelRunner):
+        assert (
+            not model_runner.server_args.enable_memory_saver
+        ), "XPUGraphRunner does not support Torch Memory Saver yet."
+        register_fake_ops()
+        self._apply_xpu_compile_config()
+        register_xpu_device_properties_for_dynamo()
+        super().__init__(model_runner)
+
+        assert (
+            not self.enable_two_batch_overlap
+        ), "XPUGraphRunner does not support two batch overlap yet."
+        assert (
+            not self.require_mlp_tp_gather
+        ), "XPUGraphRunner does not support MLP TP gather yet."
+        assert (
+            not self.require_mlp_sync
+        ), "XPUGraphRunner does not support MLP sync yet."
+        assert (
+            not self.require_gathered_buffer
+        ), "XPUGraphRunner does not support gathered buffer yet."
+
+    def _init_profile_context_and_memory_record(self):
+        profile_context = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.XPU],
+            record_shapes=True,
+        )
+        torch.xpu.memory._record_memory_history()
+        return profile_context
+
+    def _post_process_after_profile(self, prof_context):
+        torch.xpu.memory._dump_snapshot("xpu_graph_runner_memory_usage.pickle")
+        torch.xpu.memory._record_memory_history(enabled=None)
+        log_message = (
+            "Sorted by XPU Time:\n"
+            + prof_context.key_averages(group_by_input_shape=True).table(
+                sort_by="self_xpu_time_total"
+            )
+            + "\n\nSorted by CPU Time:\n"
+            + prof_context.key_averages(group_by_input_shape=True).table(
+                sort_by="self_cpu_time_total"
+            )
+            + "\n\nMemory Usage is saved to xpu_graph_runner_memory_usage.pickle\n"
+        )
+        logger.info(log_message)

python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py

@@ -0,0 +1,280 @@
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.fla.index import (
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+)
+from sglang.srt.layers.attention.fla.op import exp, make_tensor_descriptor, safe_exp
+from sglang.srt.layers.attention.fla.utils import (
+    autotune_cache_kwargs,
+)
+
+CHUNK_SIZE = 64
+
+
+# This kernel handles K blocks in a for loop to minimize register spills.
+# Time is the OUTER loop; K blocks are processed in two inner phases per step:
+#   Phase 1: store h to output, accumulate v_correction = sum_k(w_k @ h_k^T)
+#   Phase 2: update h = gate * h + k^T @ v_gated, save to scratch (initial_state)
+@triton.autotune(
+    configs=[triton.Config({"BV": 64}, num_warps=8, num_stages=2)],
+    key=["H", "K", "V", "BT", "USE_GK", "USE_INITIAL_STATE", "NT_BUCKET"],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop(
+    k,
+    v,
+    w,
+    v_new,
+    g,
+    gk,
+    h,
+    initial_state,
+    initial_state_indices,
+    cu_seqlens,
+    chunk_offsets,
+    T,
+    H: tl.constexpr,
+    Hg: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_GK: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    INPLACE_UPDATE: tl.constexpr,
+    SAVE_NEW_VALUE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    NT_BUCKET: tl.constexpr,  # this arg is kept to align with the triton kernel for CUDA
+):
+    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+            cu_seqlens + i_n + 1
+        ).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+        boh = tl.load(chunk_offsets + i_n).to(tl.int32)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+        boh = i_n * NT
+
+    # calculate offset
+    h += ((boh * H + i_h) * V * K).to(tl.int64)
+    v += ((bos * H + i_h) * V).to(tl.int64)
+    k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64)
+    w += ((bos * H + i_h) * K).to(tl.int64)
+    if SAVE_NEW_VALUE:
+        v_new += ((bos * H + i_h) * V).to(tl.int64)
+    stride_v = H * V
+    stride_h = H * V * K
+    stride_k = Hg * K
+    stride_w = H * K
+
+    w_desc = make_tensor_descriptor(
+        base=w,
+        shape=(T, K),
+        strides=(stride_w, 1),
+        block_shape=(BT, 64),
+    )
+    v_desc = make_tensor_descriptor(
+        base=v,
+        shape=(T, V),
+        strides=(stride_v, 1),
+        block_shape=(BT, BV),
+    )
+    k_desc = make_tensor_descriptor(
+        base=k,
+        shape=(T, K),
+        strides=(stride_k, 1),
+        block_shape=(BT, 64),
+    )
+    if SAVE_NEW_VALUE:
+        v_new_desc = make_tensor_descriptor(
+            base=v_new,
+            shape=(T, V),
+            strides=(stride_v, 1),
+            block_shape=(BT, BV),
+        )
+
+    index = tl.load(initial_state_indices + i_n).to(tl.int32)
+    h0 = initial_state + index * stride_h
+    ht = initial_state + index * stride_h
+    if USE_INITIAL_STATE:
+        h0 = h0 + i_h * V * K
+    if INPLACE_UPDATE:
+        ht = ht + i_h * V * K
+
+    # main recurrence — time is the outer loop
+    for i_t in range(NT):
+        ########################################################################
+        # Phase 1: store h to output, compute v_new = u - sum_k(w_k @ h_k^T)
+        ########################################################################
+        b_v_corr = tl.zeros([BT, BV], dtype=tl.float32)
+        for k_blk in range(0, K, 64):
+            # Load h: from initial_state (i_t==0) or scratch (i_t>0)
+            if i_t == 0:
+                if USE_INITIAL_STATE:
+                    p_hs = tl.make_block_ptr(
+                        h0, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
+                    )
+                    b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
+                else:
+                    b_h = tl.zeros([BV, 64], dtype=tl.float32)
+            else:
+                p_hs = tl.make_block_ptr(
+                    ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
+                )
+                b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
+
+            # Store pre-update h to output
+            p_ho = tl.make_block_ptr(
+                h + i_t * stride_h,
+                (V, K),
+                (K, 1),
+                (i_v * BV, k_blk),
+                (BV, 64),
+                (1, 0),
+            )
+            tl.store(p_ho, b_h.to(p_ho.dtype.element_ty), boundary_check=(0, 1))
+
+            # Accumulate correction: w_k @ h_k^T
+            b_w = w_desc.load([i_t * BT, k_blk])
+            b_v_corr += tl.dot(b_w, tl.trans(b_h).to(b_w.dtype))
+
+        # v_new = u - correction
+        b_v = v_desc.load([i_t * BT, i_v * BV]) - b_v_corr
+
+        if SAVE_NEW_VALUE:
+            v_new_desc.store([i_t * BT, i_v * BV], b_v.to(v_new.dtype.element_ty))
+
+        # Apply gate to v
+        last_idx = min((i_t + 1) * BT, T) - 1
+        if USE_G:
+            b_g_last = tl.load(g + bos * H + last_idx * H + i_h)
+            p_g = tl.make_block_ptr(
+                g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)
+            )
+            b_g = tl.load(p_g, boundary_check=(0,))
+            b_v = b_v * tl.expand_dims(safe_exp(b_g_last - b_g), 1)
+            b_g_last = exp(b_g_last)
+
+        b_v = b_v.to(k.dtype.element_ty)
+
+        ########################################################################
+        # Phase 2: reload h, apply gate, update h += k^T @ v, save to scratch
+        ########################################################################
+        for k_blk in range(0, K, 64):
+            # Reload h (same source as Phase 1)
+            if i_t == 0:
+                if USE_INITIAL_STATE:
+                    p_hs = tl.make_block_ptr(
+                        h0, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
+                    )
+                    b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
+                else:
+                    b_h = tl.zeros([BV, 64], dtype=tl.float32)
+            else:
+                p_hs = tl.make_block_ptr(
+                    ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
+                )
+                b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32)
+
+            # Gate decay on h
+            if USE_G:
+                b_h = b_h * b_g_last
+
+            if USE_GK:
+                o_k1 = tl.arange(0, 64) + k_blk
+                b_gk_last1 = tl.load(
+                    gk + (bos + last_idx) * H * K + i_h * K + o_k1,
+                    mask=(o_k1 < K),
+                    other=0.0,
+                )
+                b_h *= tl.expand_dims(exp(b_gk_last1), 0)
+
+            # Delta update: h += k^T @ v
+            b_k = tl.trans(k_desc.load([i_t * BT, k_blk]))
+            b_h += tl.trans(tl.dot(b_k, b_v))
+
+            # Save updated h to scratch (initial_state) for next time step
+            if INPLACE_UPDATE:
+                p_hs = tl.make_block_ptr(
+                    ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0)
+                )
+                tl.store(p_hs, b_h.to(p_hs.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gated_delta_rule_fwd_h(
+    k: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    gk: Optional[torch.Tensor] = None,
+    initial_state: Optional[torch.Tensor] = None,
+    initial_state_indices: Optional[torch.Tensor] = None,
+    save_new_value: bool = True,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_indices: Optional[torch.LongTensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, T, Hg, K, V = *k.shape, u.shape[-1]
+    H = u.shape[-2]
+    BT = CHUNK_SIZE
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE)
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    if cu_seqlens is None:
+        N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
+    else:
+        N, NT, chunk_offsets = (
+            len(cu_seqlens) - 1,
+            len(chunk_indices),
+            prepare_chunk_offsets(cu_seqlens, BT),
+        )
+    assert K <= 256, "current kernel does not support head dimension larger than 256."
+
+    h = k.new_empty(B, NT, H, V, K)
+
+    v_new = torch.empty_like(u) if save_new_value else None
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), N * H)
+
+    kernel = chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop
+
+    kernel[grid](
+        k=k,
+        v=u,
+        w=w,
+        v_new=v_new,
+        g=g,
+        gk=gk,
+        h=h,
+        initial_state=initial_state,
+        initial_state_indices=initial_state_indices,
+        cu_seqlens=cu_seqlens,
+        chunk_offsets=chunk_offsets,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        V=V,
+        BT=BT,
+        USE_G=g is not None,
+        USE_GK=gk is not None,
+        USE_INITIAL_STATE=initial_state is not None,
+        INPLACE_UPDATE=True,
+        SAVE_NEW_VALUE=v_new is not None,
+        IS_VARLEN=cu_seqlens is not None,
+        NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)),
+    )
+    return h, v_new

python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py

@@ -0,0 +1,325 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
+from sglang.srt.layers.attention.fla.op import safe_exp
+from sglang.srt.layers.attention.fla.utils import (
+    autotune_cache_kwargs,
+)
+from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd
+
+_MERGE_DOT_PRECISION = tl.constexpr("ieee")
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({"BK": BK}, num_warps=num_warps)
+        for BK in [16, 32, 64]
+        for num_warps in [2, 4, 8, 16, 32]
+    ],
+    key=["H", "Hg", "K", "BC", "BK", "USE_G", "IS_VARLEN"],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gated_delta_rule_fwd_kkt_solve_kernel_low_reg(
+    k,
+    g,
+    beta,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    Hg: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    USE_G: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    """
+    Low-reg version: one [BC,BC] accumulator at a time to minimise register pressure.
+
+    Pass 1: loop over 4 diagonal blocks (tl.static_range unrolls to 4 K-loops).
+    Pass 2: nested loop over off-diagonal distance d=1,2,3 and column j.
+      d=1 (nearest): Ai_{ij} = -Ai_ii @ A_ij_raw @ Ai_jj
+      d>1 (farther):  Ai_{ij} = -(Ai_ii @ A_ij_raw
+                                  + sum_{m=j+1}^{i-1} Ai_im @ A_mj_raw) @ Ai_jj
+    Each K-loop holds exactly one [BC,BC] accumulator.  Raw blocks needed by
+    later correction terms are spilled to upper-triangular scratch slots in A
+    (see _KKT_SCRATCH_COL for the layout; boundary_check makes out-of-bounds
+    stores/loads safe so no runtime `if i_tcX < T` guards are needed).
+    """
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+            chunk_indices + i_t * 2 + 1
+        ).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+            cu_seqlens + i_n + 1
+        ).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    if i_t * BT >= T:
+        return
+
+    i_tc0 = i_t * BT
+    k += (bos * Hg + i_h // (H // Hg)) * K
+    A += (bos * H + i_h) * BT
+
+    o_i = tl.arange(0, BC)
+    m_d = o_i[:, None] > o_i[None, :]
+    m_I = o_i[:, None] == o_i[None, :]
+
+    ############################################################################
+    # Pass 1: diagonal blocks — one K-loop per sub-chunk (tl.static_range → 4)
+    ############################################################################
+    for i_b in tl.static_range(4):
+        i_tci = i_tc0 + i_b * BC
+        m_tci = (i_tci + o_i) < T
+        p_bi = tl.make_block_ptr(
+            beta + bos * H + i_h, (T,), (H,), (i_tci,), (BC,), (0,)
+        )
+        b_bi = tl.load(p_bi, boundary_check=(0,)).to(tl.float32)
+        if USE_G:
+            p_gi = tl.make_block_ptr(
+                g + bos * H + i_h, (T,), (H,), (i_tci,), (BC,), (0,)
+            )
+            b_gi = tl.load(p_gi, boundary_check=(0,)).to(tl.float32)
+        b_A = tl.zeros([BC, BC], dtype=tl.float32)
+        for i_k in range(tl.cdiv(K, BK)):
+            p_k = tl.make_block_ptr(
+                k, (T, K), (Hg * K, 1), (i_tci, i_k * BK), (BC, BK), (1, 0)
+            )
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_A += tl.dot(b_k, tl.trans(b_k))
+        if USE_G:
+            b_A *= safe_exp(b_gi[:, None] - b_gi[None, :])
+        b_A = (
+            tl.where(m_d & (m_tci[:, None] & m_tci[None, :]), b_A, 0.0) * b_bi[:, None]
+        )
+        # Forward substitution: solve (I + A_diag) x = I column by column.
+        # Extra iterations for out-of-bounds rows are no-ops (b_A rows are zero).
+        b_Ai = -b_A
+        for i in range(2, BC):
+            b_a = tl.sum(tl.where((o_i == i)[:, None], -b_A, 0.0), 0)
+            b_a = tl.where(o_i < i, b_a, 0.0)
+            b_a = b_a + tl.sum(b_a[:, None] * b_Ai, 0)
+            b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai)
+        b_Ai += m_I
+        p_Aii = tl.make_block_ptr(
+            A, (T, BT), (H * BT, 1), (i_tci, i_b * BC), (BC, BC), (1, 0)
+        )
+        tl.store(p_Aii, b_Ai.to(A.dtype.element_ty), boundary_check=(0, 1))
+
+    ############################################################################
+    # Pass 2: off-diagonal blocks — one K-loop per (i,j) pair.
+    # Outer loop: d = i-j (diagonal distance) from 1 to 3.
+    # Inner loop: j from 0 to 3-d  (so i = j+d).
+    # Processing order ensures Ai_im (m in j+1..i-1) is already in A when needed.
+    # Formula: Ai_ij = -(Ai_ii @ A_ij_raw
+    #                    + sum_{m=j+1}^{i-1} Ai_im @ A_mj_raw) @ Ai_jj
+    ############################################################################
+    for d in tl.static_range(1, 4):
+        for j in tl.static_range(0, 4 - d):
+            i = j + d  # compile-time ints from static_range
+            i_tci = i_tc0 + i * BC
+            i_tcj = i_tc0 + j * BC
+            p_bi = tl.make_block_ptr(
+                beta + bos * H + i_h, (T,), (H,), (i_tci,), (BC,), (0,)
+            )
+            b_bi = tl.load(p_bi, boundary_check=(0,)).to(tl.float32)
+            if USE_G:
+                p_gi = tl.make_block_ptr(
+                    g + bos * H + i_h, (T,), (H,), (i_tci,), (BC,), (0,)
+                )
+                p_gj = tl.make_block_ptr(
+                    g + bos * H + i_h, (T,), (H,), (i_tcj,), (BC,), (0,)
+                )
+                b_gi = tl.load(p_gi, boundary_check=(0,)).to(tl.float32)
+                b_gj = tl.load(p_gj, boundary_check=(0,)).to(tl.float32)
+            # K-loop: accumulate k_i @ k_j^T
+            b_A = tl.zeros([BC, BC], dtype=tl.float32)
+            for i_k in range(tl.cdiv(K, BK)):
+                p_ki = tl.make_block_ptr(
+                    k, (T, K), (Hg * K, 1), (i_tci, i_k * BK), (BC, BK), (1, 0)
+                )
+                p_kj = tl.make_block_ptr(
+                    k, (T, K), (Hg * K, 1), (i_tcj, i_k * BK), (BC, BK), (1, 0)
+                )
+                b_A += tl.dot(
+                    tl.load(p_ki, boundary_check=(0, 1)),
+                    tl.trans(tl.load(p_kj, boundary_check=(0, 1))),
+                )
+            if USE_G:
+                b_A *= safe_exp(b_gi[:, None] - b_gj[None, :])
+            b_A *= b_bi[:, None]
+            # Scratch this raw block if a later row in this column needs it as
+            # a correction term.  i < 3 means rows i+1..3 exist and will use it.
+            # Scratch column layout (compile-time ternary, evaluated at trace time):
+            #   (i=1,j=0)->col1  (i=2,j=1)->col2  (i=2,j=0)->col3
+            if i < 3:
+                sc = 1 if (i == 1 and j == 0) else (2 if (i == 2 and j == 1) else 3)
+                p_s = tl.make_block_ptr(
+                    A, (T, BT), (H * BT, 1), (i_tc0, sc * BC), (BC, BC), (1, 0)
+                )
+                tl.store(p_s, b_A.to(A.dtype.element_ty), boundary_check=(0, 1))
+            # Correction sum: sum_{m=j+1}^{i-1} Ai_im @ A_mj_raw
+            # Unrolled manually (d is a compile-time Python int from tl.static_range):
+            #   d=1: no corrections; d=2: m=j+1; d=3: m=j+1 then m=j+2
+            b_corr = tl.zeros([BC, BC], dtype=tl.float32)
+            if d >= 2:
+                m1 = j + 1
+                sc_m1j = (
+                    1 if (m1 == 1 and j == 0) else (2 if (m1 == 2 and j == 1) else 3)
+                )
+                p_s_m1j = tl.make_block_ptr(
+                    A, (T, BT), (H * BT, 1), (i_tc0, sc_m1j * BC), (BC, BC), (1, 0)
+                )
+                b_A_m1j = tl.load(p_s_m1j, boundary_check=(0, 1)).to(tl.float32)
+                p_Ai_im1 = tl.make_block_ptr(
+                    A, (T, BT), (H * BT, 1), (i_tci, m1 * BC), (BC, BC), (1, 0)
+                )
+                b_corr += tl.dot(
+                    tl.load(p_Ai_im1, boundary_check=(0, 1)).to(tl.float32),
+                    b_A_m1j,
+                    input_precision=_MERGE_DOT_PRECISION,
+                )
+            if d >= 3:
+                m2 = j + 2
+                sc_m2j = (
+                    1 if (m2 == 1 and j == 0) else (2 if (m2 == 2 and j == 1) else 3)
+                )
+                p_s_m2j = tl.make_block_ptr(
+                    A, (T, BT), (H * BT, 1), (i_tc0, sc_m2j * BC), (BC, BC), (1, 0)
+                )
+                b_A_m2j = tl.load(p_s_m2j, boundary_check=(0, 1)).to(tl.float32)
+                p_Ai_im2 = tl.make_block_ptr(
+                    A, (T, BT), (H * BT, 1), (i_tci, m2 * BC), (BC, BC), (1, 0)
+                )
+                b_corr += tl.dot(
+                    tl.load(p_Ai_im2, boundary_check=(0, 1)).to(tl.float32),
+                    b_A_m2j,
+                    input_precision=_MERGE_DOT_PRECISION,
+                )
+            # Compute and store Ai_ij
+            p_Ai_ii = tl.make_block_ptr(
+                A, (T, BT), (H * BT, 1), (i_tci, i * BC), (BC, BC), (1, 0)
+            )
+            p_Ai_jj = tl.make_block_ptr(
+                A, (T, BT), (H * BT, 1), (i_tcj, j * BC), (BC, BC), (1, 0)
+            )
+            b_Ai_ii = tl.load(p_Ai_ii, boundary_check=(0, 1)).to(tl.float32)
+            b_Ai_jj = tl.load(p_Ai_jj, boundary_check=(0, 1)).to(tl.float32)
+            b_Ai_ij = -tl.dot(
+                tl.dot(b_Ai_ii, b_A, input_precision=_MERGE_DOT_PRECISION) + b_corr,
+                b_Ai_jj,
+                input_precision=_MERGE_DOT_PRECISION,
+            )
+            p_Ai_ij = tl.make_block_ptr(
+                A, (T, BT), (H * BT, 1), (i_tci, j * BC), (BC, BC), (1, 0)
+            )
+            tl.store(p_Ai_ij, b_Ai_ij.to(A.dtype.element_ty), boundary_check=(0, 1))
+
+    # Clean up scratch slots: Pass 2 stored raw A_ij blocks in the upper-triangular
+    # part of row i_tc0 (cols BC..3*BC). These must be zeroed because
+    # recompute_w_u_fwd reads the full BT×BT block.
+    b_zero = tl.zeros([BC, BC], dtype=tl.float32)
+    for sc in tl.static_range(1, BT // BC):
+        p_scratch = tl.make_block_ptr(
+            A, (T, BT), (H * BT, 1), (i_tc0, sc * BC), (BC, BC), (1, 0)
+        )
+        tl.store(p_scratch, b_zero.to(A.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gated_delta_rule_fwd_intra(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor | None = None,
+    beta: torch.Tensor | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    chunk_indices: torch.LongTensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    r"""
+    GDN intra-chunk forward: fused kkt + solve_tril + recompute_w_u.
+
+    Equivalent to:
+        A = chunk_scaled_dot_kkt_fwd(k, g, beta, ...)       # kernel 1
+        A = solve_tril(A, ...)                                # kernel 2
+        w, u = recompute_w_u_fwd(k, v, beta, A, g, ...)      # kernel 3
+
+    Fuses kernels 1+2 into a single kernel, reducing from 3 to 2 kernel launches
+    and eliminating the HBM round-trip for the intermediate A matrix.
+
+    Args:
+        k (torch.Tensor):
+            The key tensor of shape `[B, T, H, K]`.
+        v (torch.Tensor):
+            The value tensor of shape `[B, T, H, V]`.
+        g (torch.Tensor):
+            The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.
+        beta (torch.Tensor):
+            The beta tensor of shape `[B, T, H]`.
+        cu_seqlens (torch.LongTensor):
+            The cumulative sequence lengths. Default: `None`.
+        chunk_size (int):
+            The chunk size. Default: 64.
+        chunk_indices (torch.LongTensor):
+            Precomputed chunk indices. Default: `None`.
+
+    Returns:
+        w (torch.Tensor): shape `[B, T, H, K]`
+        u (torch.Tensor): shape `[B, T, H, V]`
+        A (torch.Tensor): shape `[B, T, H, BT]`, the solved (I+A)^{-1} matrix
+    """
+    B, T, Hg, K = k.shape
+    H = beta.shape[-1]
+    BT = chunk_size
+    BC = 16
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    # Step 1: fused kkt + solve_tril
+    A = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype)
+    kernel = chunk_gated_delta_rule_fwd_kkt_solve_kernel_low_reg
+    kernel[(NT, B * H)](
+        k=k,
+        g=g,
+        beta=beta,
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        BT=BT,
+        BC=BC,
+    )
+
+    # Step 2: recompute_w_u
+    w, u = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        g_cumsum=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    return w, u, A

python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py

@@ -0,0 +1,128 @@
+from typing import Optional
+
+import torch
+import triton
+
+from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+    fused_sigmoid_gating_delta_rule_update_kernel,
+)
+
+
+def fused_sigmoid_gating_delta_rule_update(
+    A_log: torch.Tensor,
+    a: torch.Tensor,
+    dt_bias: torch.Tensor,
+    softplus_beta: float,
+    softplus_threshold: float,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    b: torch.Tensor,
+    initial_state_source: torch.Tensor,
+    initial_state_indices: torch.Tensor,
+    scale: Optional[float] = None,
+    use_qk_l2norm_in_kernel: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    is_kda: bool = False,
+    # Optional parameters for target_verify support
+    disable_state_update: bool = False,
+    intermediate_states_buffer: Optional[torch.Tensor] = None,
+    intermediate_state_indices: Optional[torch.Tensor] = None,
+    cache_steps: Optional[int] = None,
+    retrieve_parent_token: Optional[torch.Tensor] = None,
+):
+    """
+    Fused triton implementation of sigmoid gating delta rule update.
+    This function uses a single fused kernel that combines both sigmoid gating computation
+    and the recurrent delta rule update for better performance.
+
+    Supports both decode and target_verify modes:
+    - decode: standard single-step update with state write-back
+    - target_verify: multi-step with intermediate state caching, optional tree attention,
+                     and optional state update disable
+    """
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    stride_q = q.stride()[1]
+    stride_k = k.stride()[1]
+    stride_v = v.stride()[1]
+    stride_b = b.stride()[-2]
+    # Both paths (KDA/GDN) advance p_a once per token, so use the token-axis stride.
+    # For 2D a ([T, ...]) this is stride(0); for 3D a ([B, T, ...]) this is stride(1).
+    # Using stride()[-2] covers GDN [T, HV] and KDA layouts ([T, HV*K] / [B, T, HV*K]).
+    stride_a = a.stride()[-2]
+    HV = v.shape[2]
+    N = B if cu_seqlens is None else len(cu_seqlens) - 1
+    BK, BV = triton.next_power_of_2(K), min(
+        triton.next_power_of_2(V), 16
+    )  # use 16 here to reduce register pressure
+    NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+    assert NK == 1, "NK > 1 is not supported yet"
+    num_stages = 3
+    num_warps = 1
+
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    else:
+        assert scale > 0, "scale must be positive"
+
+    o = q.new_empty(NK, *v.shape)
+
+    # Prepare retrieve_parent_token strides
+    if retrieve_parent_token is not None:
+        stride_retrieve_parent_token_seq = retrieve_parent_token.stride(0)
+        stride_retrieve_parent_token_token = retrieve_parent_token.stride(1)
+    else:
+        stride_retrieve_parent_token_seq = 0
+        stride_retrieve_parent_token_token = 0
+
+    NP2_T = triton.next_power_of_2(T)
+
+    grid = (NK, NV, N * HV)
+
+    fused_sigmoid_gating_delta_rule_update_kernel[grid](
+        A_log=A_log,
+        a=a,
+        dt_bias=dt_bias,
+        softplus_beta=softplus_beta,
+        softplus_threshold=softplus_threshold,
+        q=q,
+        k=k,
+        v=v,
+        b=b,
+        o=o,
+        h0_source=initial_state_source,
+        h0_indices=initial_state_indices,
+        cu_seqlens=cu_seqlens,
+        intermediate_states_buffer=intermediate_states_buffer,
+        intermediate_state_indices=intermediate_state_indices,
+        cache_steps=0 if cache_steps is None else cache_steps,
+        retrieve_parent_token_ptr=retrieve_parent_token,
+        stride_retrieve_parent_token_seq=stride_retrieve_parent_token_seq,
+        stride_retrieve_parent_token_token=stride_retrieve_parent_token_token,
+        scale=scale,
+        T=T,
+        stride_a=stride_a,
+        stride_q=stride_q,
+        stride_k=stride_k,
+        stride_v=stride_v,
+        stride_b=stride_b,
+        NP2_T=NP2_T,
+        B=B,
+        H=H,
+        HV=HV,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        USE_INITIAL_STATE=initial_state_source is not None,
+        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        IS_VARLEN=cu_seqlens is not None,
+        IS_KDA=is_kda,
+        DISABLE_STATE_UPDATE=disable_state_update,
+        CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
+        HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    o = o.squeeze(0)
+    return o