feat: add DeepSeek V4 NPU kernel Python bindings.

tests/python/test_grouped_moe.py

@@ -0,0 +1,123 @@
+# Copyright 2026 The xLLM Authors.
+#
+# 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
+#
+#     https://github.com/xLLM-AI/xllm/blob/main/LICENSE
+#
+# 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.
+
+"""Contracts for the NPU pre-selected grouped MoE path."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+import torch
+
+_REPO_ROOT = Path(__file__).parents[2]
+
+
+def _load_npu_moe_module():
+    path = _REPO_ROOT / "xllm/python/kernels_npu/moe.py"
+    spec = importlib.util.spec_from_file_location("pr5_npu_moe", path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_selected_expert_moe_matches_native_call_contract(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from xllm.python import kernels
+
+    moe = _load_npu_moe_module()
+
+    hidden = torch.empty(3, 16, dtype=torch.bfloat16)
+    topk_weights = torch.ones(3, 2, dtype=torch.bfloat16)
+    topk_ids = torch.zeros(3, 2, dtype=torch.int32)
+    expanded = torch.empty(6, 16, dtype=torch.bfloat16)
+    row_ids = torch.arange(6, dtype=torch.int32)
+    group_list = torch.tensor([1, 3, 5, 6], dtype=torch.int64)
+    quantized = torch.empty(6, 16, dtype=torch.int8)
+    input_scale = torch.empty(6, dtype=torch.float32)
+    gemm1 = torch.empty(6, 32, dtype=torch.int32)
+    activated = torch.empty(6, 16, dtype=torch.int8)
+    activation_scale = torch.empty(6, dtype=torch.float32)
+    gemm2 = torch.empty(6, 16, dtype=torch.bfloat16)
+    calls: list[tuple[str, object]] = []
+
+    def init_routing(*args, **kwargs):
+        calls.append(("routing", kwargs))
+        return expanded, row_ids, group_list, torch.empty(0)
+
+    def dynamic_quant(value):
+        assert value is expanded
+        calls.append(("dynamic_quant", value))
+        return quantized, input_scale
+
+    def dequant_swiglu_quant(**kwargs):
+        calls.append(("dequant_swiglu_quant", kwargs))
+        return activated, activation_scale
+
+    gemm_calls: list[dict[str, object]] = []
+
+    def group_gemm(**kwargs):
+        gemm_calls.append(kwargs)
+        return gemm1 if len(gemm_calls) == 1 else gemm2
+
+    def token_unpermute(**kwargs):
+        calls.append(("unpermute", kwargs))
+        return hidden
+
+    monkeypatch.setattr(moe, "_group_gemm", group_gemm)
+    monkeypatch.setattr(moe.torch_npu, "npu_moe_init_routing_v2", init_routing)
+    monkeypatch.setattr(moe.torch_npu, "npu_moe_token_unpermute", token_unpermute)
+    monkeypatch.setattr(kernels, "dynamic_quant", dynamic_quant, raising=False)
+    monkeypatch.setattr(kernels, "dequant_swiglu_quant", dequant_swiglu_quant, raising=False)
+
+    result = moe._grouped_moe_with_selected_experts_impl(
+        hidden,
+        topk_weights,
+        topk_ids,
+        torch.empty(4, 16, 32, dtype=torch.int8),
+        torch.empty(4, 16, 16, dtype=torch.int8),
+        torch.empty(4, 32),
+        torch.empty(4, 16),
+        num_total_experts=16,
+        start_expert_id=4,
+        num_experts_per_rank=4,
+        swiglu_limit=7.0,
+    )
+
+    assert result is hidden
+    routing = dict(calls)["routing"]
+    assert isinstance(routing, dict)
+    assert routing["active_expert_range"] == [4, 8]
+    assert routing["expert_num"] == 16
+    assert routing["quant_mode"] == -1
+
+    assert len(gemm_calls) == 2
+    assert gemm_calls[0]["scale"] is None
+    assert gemm_calls[0]["per_token_scale"] is None
+    assert gemm_calls[0]["output_dtype"] == torch.int32
+    assert gemm_calls[1]["scale"].dtype == torch.bfloat16
+    assert gemm_calls[1]["per_token_scale"] is activation_scale
+    assert gemm_calls[1]["output_dtype"] == torch.bfloat16
+    assert all(call["group_list"] is group_list for call in gemm_calls)
+    assert all(call["group_list_type"] == 1 for call in gemm_calls)
+
+    dequant = dict(calls)["dequant_swiglu_quant"]
+    assert isinstance(dequant, dict)
+    assert dequant["x"] is gemm1
+    assert dequant["activation_scale"] is input_scale
+    assert dequant["group_index"] is group_list
+    assert dequant["clamp_limit"] == 7.0

tests/python/test_kernels_import.py

@@ -57,6 +57,9 @@
     "quantize_per_tensor(Tensor self, Tensor scales, Tensor zero_points, ScalarType dtype, int axis) -> Tensor",
     "dynamic_quant(Tensor input, Tensor? smooth_scales, Tensor? group_index, "
     "ScalarType? dst_type) -> (Tensor, Tensor?)",
+    "group_gemm(Tensor x, Tensor weight, Tensor? scale, Tensor? "
+    "per_token_scale, Tensor group_list, int split_item, int group_type, int "
+    "group_list_type, ScalarType? output_dtype) -> Tensor",
     "lightning_indexer(Tensor query, Tensor key, Tensor weights, Tensor? "
     "query_seq_lengths, Tensor? key_seq_lengths, Tensor? block_table, str "
     "layout_query, str layout_key, int selected_count, int sparse_mode, int "
@@ -77,6 +80,17 @@
     "Tensor? actual_seq_lengths_kv, Tensor? query_rope, Tensor? key_rope, "
     "float scale_value, int sparse_block_size, str layout_query, str layout_kv, "
     "int sparse_mode, Tensor(a!) output) -> Tensor",
+    "rms_norm_dynamic_quant(Tensor input, Tensor weight, float eps) -> (Tensor, Tensor)",
+    "npu_inplace_partial_rotary_mul(Tensor(a!) x, Tensor r1, Tensor r2, str rotary_mode, int[] partial_slice) -> ()",
+    "moe_gating_top_k_hash(Tensor x, int k, Tensor? bias, Tensor? input_ids, Tensor? tid2eid, int k_group, int group_count, float routed_scaling_factor, float eps, int group_select_mode, int renorm, int norm_type, bool out_flag) -> (Tensor, Tensor, Tensor)",
+    "dequant_swiglu_quant(Tensor x, Tensor? weight_scale, Tensor? activation_scale, Tensor? bias, Tensor? quant_scale, Tensor? quant_offset, Tensor? group_index, bool activate_left, int quant_mode, int swiglu_mode, float clamp_limit, float glu_alpha, float glu_bias) -> (Tensor, Tensor)",
+    "hc_pre(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, int hc_mult, int hc_sinkhorn_iters, float norm_eps, float hc_eps) -> (Tensor, Tensor, Tensor)",
+    "hc_post(Tensor x, Tensor residual, Tensor post, Tensor comb) -> Tensor",
+    "compressor(Tensor x, Tensor wkv, Tensor wgate, Tensor(a!) kv_state, Tensor(b!) score_state, Tensor ape, Tensor norm_weight, Tensor rope_sin, Tensor rope_cos, Tensor? kv_block_table, Tensor? score_block_table, Tensor? cu_seqlens, Tensor? seqused, Tensor? start_pos, int rope_head_dim, int cmp_ratio, int coff, float norm_eps, int rotary_mode, bool enable_grad) -> (Tensor, Tensor, Tensor, Tensor, Tensor)",
+    "sparse_attn_sharedkv(Tensor q, Tensor? ori_kv, Tensor? cmp_kv, Tensor? ori_sparse_indices, Tensor? cmp_sparse_indices, Tensor? ori_block_table, Tensor? cmp_block_table, Tensor? cu_seqlens_q, Tensor? cu_seqlens_ori_kv, Tensor? cu_seqlens_cmp_kv, Tensor? seqused_q, Tensor? seqused_kv, Tensor? sinks, Tensor? metadata, float softmax_scale, int cmp_ratio, int ori_mask_mode, int cmp_mask_mode, int ori_win_left, int ori_win_right, str layout_q, str layout_kv, bool return_softmax_lse) -> (Tensor, Tensor)",
+    "sparse_attn_sharedkv_metadata(int num_heads_q, int num_heads_kv, int head_dim, Tensor? cu_seqlens_q, Tensor? cu_seqlens_ori_kv, Tensor? cu_seqlens_cmp_kv, Tensor? seqused_q, Tensor? seqused_kv, int batch_size, int max_seqlen_q, int max_seqlen_kv, int ori_topk, int cmp_topk, int cmp_ratio, int ori_mask_mode, int cmp_mask_mode, int ori_win_left, int ori_win_right, str layout_q, str layout_kv, bool has_ori_kv, bool has_cmp_kv) -> Tensor",
+    "quant_lightning_indexer(Tensor query, Tensor key, Tensor weights, Tensor query_dequant_scale, Tensor key_dequant_scale, int query_quant_mode, int key_quant_mode, Tensor? actual_seq_lengths_query, Tensor? actual_seq_lengths_key, Tensor? block_table, Tensor? metadata, str layout_query, str layout_key, int sparse_count, int sparse_mode, int pre_tokens, int next_tokens, int cmp_ratio, bool return_value) -> (Tensor, Tensor)",
+    "quant_lightning_indexer_metadata(int num_heads_q, int num_heads_k, int head_dim, int query_quant_mode, int key_quant_mode, Tensor? actual_seq_lengths_query, Tensor? actual_seq_lengths_key, int batch_size, int max_seqlen_q, int max_seqlen_k, str layout_query, str layout_key, int sparse_count, int sparse_mode, int pre_tokens, int next_tokens, int cmp_ratio, str device) -> Tensor",
 )
 
 _PLATFORM_REQUIRED = pytest.mark.skipif(
@@ -193,11 +207,12 @@ def test_registry_does_not_preload_model_modules() -> None:
 
 
 def test_npu_fake_tensor_and_mutation_contracts() -> None:
-    """Quantization and sparse attention shapes traced without an NPU."""
+    """NPU wrapper shape and mutation contracts traced without an NPU."""
     _run_isolated_python(
         """
         import xllm.python.kernels_npu._custom_op  # noqa: F401
-        from xllm.python.kernels_npu import quantization, sparse_attention
+        from xllm.python.kernels_npu import dsa, normalization, quantization
+        from xllm.python.kernels_npu import rotary_embedding, sparse_attention
 
         mode = torch._subclasses.fake_tensor.FakeTensorMode()
         with mode:
@@ -213,6 +228,37 @@ def test_npu_fake_tensor_and_mutation_contracts() -> None:
             assert quantized.shape == (2, 2) and quantized.dtype == torch.int32
             assert scale.shape == (2,) and scale.dtype == torch.float32
 
+            grouped = torch.ops.xllm_ops.group_gemm(
+                torch.empty(6, 16, dtype=torch.int8),
+                torch.empty(4, 16, 32, dtype=torch.int8),
+                None,
+                None,
+                torch.empty(4, dtype=torch.int64),
+                2,
+                0,
+                1,
+                torch.int32,
+            )
+            assert grouped.shape == (6, 32)
+            assert grouped.dtype == torch.int32
+
+            from xllm.python.kernels_npu import moe
+
+            selected_moe = moe.grouped_moe_with_selected_experts(
+                torch.empty(3, 16, dtype=torch.bfloat16),
+                torch.empty(3, 2, dtype=torch.bfloat16),
+                torch.empty(3, 2, dtype=torch.int32),
+                torch.empty(4, 16, 32, dtype=torch.int8),
+                torch.empty(4, 16, 16, dtype=torch.int8),
+                torch.empty(4, 32),
+                torch.empty(4, 16),
+                num_total_experts=16,
+                start_expert_id=4,
+                num_experts_per_rank=4,
+            )
+            assert selected_moe.shape == (3, 16)
+            assert selected_moe.dtype == torch.bfloat16
+
             query = torch.empty(8, 4, 16)
             key = torch.empty(8, 2, 16)
             indices = sparse_attention.lightning_indexer(
@@ -229,6 +275,107 @@ def test_npu_fake_tensor_and_mutation_contracts() -> None:
                 torch.empty(2, 2, 16),
             ) is None
 
+            normed, norm_scale = normalization.rms_norm_dynamic_quant(
+                torch.empty(8, 16), torch.empty(16), 1e-6
+            )
+            assert normed.shape == (8, 16) and normed.dtype == torch.int8
+            assert norm_scale.shape == (8,) and norm_scale.dtype == torch.float32
+
+            rotary_input = torch.empty(8, 2, 128)
+            rotary_ptr = rotary_input.data_ptr()
+            assert rotary_embedding.npu_inplace_partial_rotary_mul(
+                rotary_input, torch.empty(8, 64), torch.empty(8, 64), 64, 64
+            ).data_ptr() == rotary_ptr
+
+            compressor_out = dsa.compressor(
+                torch.empty(8, 16),
+                torch.empty(8, 16),
+                torch.empty(8, 16),
+                torch.empty(1, 128, 8),
+                torch.empty(1, 128, 8),
+                torch.empty(4, 8),
+                torch.empty(8),
+                torch.empty(2, 4),
+                torch.empty(2, 4),
+                None, None, None, None, None,
+                4, 4, 1, 1e-6, 1, False,
+            )
+            assert compressor_out[0].shape == (2, 8)
+            assert all(tensor.numel() == 0 for tensor in compressor_out[1:])
+
+            seq_lens = torch.empty(1, dtype=torch.int32)
+            sparse_metadata = dsa.sparse_attn_sharedkv_metadata(
+                64, 1, 512, None, None, None, seq_lens, seq_lens,
+                1, 4, 16, 0, 0, 1, 4, 3, 127, 0,
+                "BSND", "PA_ND", True, False,
+            )
+            assert sparse_metadata.shape == (1024,)
+            assert sparse_metadata.dtype == torch.int32
+
+            dsa_query = torch.empty(1, 4, 64, 512)
+            sparse_out, sparse_lse = dsa.sparse_attn_sharedkv(
+                dsa_query, None, None, None, None, None, None,
+                None, None, None, None, None, None, sparse_metadata,
+                1.0, 1, 4, 3, 127, 0, "BSND", "PA_ND", False,
+            )
+            assert sparse_out.shape == dsa_query.shape
+            assert sparse_lse.numel() == 0
+
+            qli_metadata = dsa.quant_lightning_indexer_metadata(
+                64, 1, 128, 0, 0, seq_lens, seq_lens,
+                1, 8, 8, "TND", "PA_BSND", 512, 3,
+                2**63 - 1, 2**63 - 1, 4, "cpu",
+            )
+            assert qli_metadata.shape == (1024,)
+            assert qli_metadata.dtype == torch.int32
+
+            qli_indices, qli_values = dsa.quant_lightning_indexer(
+                torch.empty(8, 64, 128, dtype=torch.int8),
+                torch.empty(1, 128, 1, 128, dtype=torch.int8),
+                torch.empty(8, 64),
+                torch.empty(8, 64),
+                torch.empty(1, 128, 1),
+                0, 0, seq_lens, seq_lens, torch.empty(1, 1), qli_metadata,
+                "TND", "PA_BSND", 512, 3,
+                2**63 - 1, 2**63 - 1, 4, False,
+            )
+            assert qli_indices.shape == (8, 1, 512)
+            assert qli_indices.dtype == torch.int32
+            assert qli_values.numel() == 0
+
+            hc_input = torch.empty(8, 4, 16)
+            hc_attn, hc_post, hc_comb = dsa.hc_pre(
+                hc_input,
+                torch.empty(24, 64),
+                torch.empty(3),
+                torch.empty(24),
+                4, 20, 1e-6, 1e-6,
+            )
+            assert hc_attn.shape == (8, 16)
+            assert hc_post.shape == (8, 4)
+            assert hc_comb.shape == (8, 4, 4)
+            assert dsa.hc_post(
+                hc_attn, hc_input, hc_post, hc_comb
+            ).shape == hc_input.shape
+
+            gate_weights, expert_idx, gate_out = dsa.moe_gating_top_k_hash(
+                torch.empty(8, 256), 6, None, None, None,
+                1, 1, 1.0, 1e-20, 1, 0, 2, False,
+            )
+            assert gate_weights.shape == (8, 6)
+            assert expert_idx.shape == (8, 6)
+            assert expert_idx.dtype == torch.int32
+            assert gate_out.shape == (8, 256)
+            assert gate_out.dtype == torch.float32
+
+            swiglu_out, swiglu_scale = dsa.dequant_swiglu_quant(
+                torch.empty(8, 32, dtype=torch.int32), None, None
+            )
+            assert swiglu_out.shape == (8, 16)
+            assert swiglu_out.dtype == torch.int8
+            assert swiglu_scale.shape == (8,)
+            assert swiglu_scale.dtype == torch.float32
+
             try:
                 quantization.dynamic_quant(
                     torch.empty(2, 15), dst_type=torch.quint4x2

xllm/core/kernels/npu/CMakeLists.txt

@@ -37,6 +37,7 @@ cc_library(
     w4a8_dynamic_moe_preprocess.cpp
     rec_constrained_topk.cpp
     npu_ops_library.cpp
+    group_gemm_wrapper.cpp
   DEPS
     :torch_npu_kernels
     :tilelang_kernels

xllm/core/kernels/npu/group_gemm_wrapper.cpp

@@ -0,0 +1,65 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "npu_ops_api.h"
+
+namespace xllm::kernel::npu {
+
+torch::Tensor group_gemm(const torch::Tensor& x,
+                         const torch::Tensor& weight,
+                         const std::optional<torch::Tensor>& scale,
+                         const std::optional<torch::Tensor>& per_token_scale,
+                         const torch::Tensor& group_list,
+                         int64_t split_item,
+                         int64_t group_type,
+                         int64_t group_list_type,
+                         std::optional<at::ScalarType> output_dtype) {
+  std::vector<torch::Tensor> x_list = {x};
+  std::vector<torch::Tensor> weight_list = {weight};
+  std::vector<torch::Tensor> scale_storage;
+  std::vector<torch::Tensor> per_token_scale_storage;
+  std::optional<torch::TensorList> scale_list = std::nullopt;
+  if (scale.has_value()) {
+    scale_storage.push_back(scale.value());
+    scale_list = torch::TensorList(scale_storage);
+  }
+  std::optional<torch::TensorList> per_token_scale_list = std::nullopt;
+  if (per_token_scale.has_value()) {
+    per_token_scale_storage.push_back(per_token_scale.value());
+    per_token_scale_list = torch::TensorList(per_token_scale_storage);
+  }
+  auto outputs =
+      apply_npu_grouped_matmul(torch::TensorList(x_list),
+                               torch::TensorList(weight_list),
+                               /*bias=*/std::nullopt,
+                               scale_list,
+                               /*offset=*/std::nullopt,
+                               /*antiquant_scale=*/std::nullopt,
+                               /*antiquant_offset=*/std::nullopt,
+                               per_token_scale_list,
+                               group_list,
+                               /*activation_input=*/std::nullopt,
+                               /*activation_quant_scale=*/std::nullopt,
+                               /*activation_quant_offset=*/std::nullopt,
+                               split_item,
+                               group_type,
+                               group_list_type,
+                               /*act_type=*/std::nullopt,
+                               /*tuning_config=*/c10::nullopt,
+                               output_dtype);
+  return outputs.back();
+}
+
+}  // namespace xllm::kernel::npu