feat: support Qwen3.5 eager-mode on NPU with Python executor.

tests/python/test_registry.py

@@ -21,10 +21,10 @@
 
 def test_unsupported_model_fails_before_import(monkeypatch: pytest.MonkeyPatch) -> None:
     import_model = Mock()
-    monkeypatch.setattr(registry.current_platform, "device_type", lambda: "npu")
+    monkeypatch.setattr(registry.current_platform, "device_type", lambda: "cuda")
     monkeypatch.setattr(registry, "import_module", import_model)
 
-    with pytest.raises(NotImplementedError, match="qwen3_5.*npu"):
-        registry.get_model_class("qwen3_5")
+    with pytest.raises(NotImplementedError, match="qwen3_vl.*cuda"):
+        registry.get_model_class("qwen3_vl")
 
     import_model.assert_not_called()

xllm/core/kernels/npu/npu_ops_library.cpp

@@ -27,6 +27,8 @@ limitations under the License.
 
 #include "kernels/npu/xllm_ops/xllm_ops_api.h"
 #include "npu_ops_api.h"
+#include "tilelang/tilelang_ops_api.h"
+#include "triton_npu/torch_api/triton_ops_api.h"
 
 namespace xllm {
 
@@ -38,6 +40,194 @@ torch::Tensor rms_norm_npu(const torch::Tensor& input,
   return xllm::kernel::npu::rms_norm(input, weight, eps, "rmsnorm");
 }
 
+torch::Tensor rms_norm_gated_npu(const torch::Tensor& input,
+                                 const torch::Tensor& gate,
+                                 const torch::Tensor& weight,
+                                 double eps) {
+  return xllm::kernel::npu::layer_norm_fwd_aclnn(input,
+                                                 weight,
+                                                 /*bias=*/torch::Tensor(),
+                                                 eps,
+                                                 /*z=*/gate,
+                                                 /*group_size=*/input.size(-1),
+                                                 /*norm_before_gate=*/true,
+                                                 /*is_rms_norm=*/true);
+}
+
+torch::Tensor l2_norm_npu(torch::Tensor input, double eps) {
+  return xllm::kernel::npu::npu_l2norm_last_dim(input, eps);
+}
+
+torch::Tensor causal_conv1d_update_npu(torch::Tensor x,
+                                       torch::Tensor conv_state,
+                                       torch::Tensor weight,
+                                       torch::Tensor state_indices) {
+  // Python layer stores weight as [dim, kernel_width]; tilelang expects
+  // [kernel_width, dim]. Transpose if the first dim is larger.
+  if (weight.size(0) > weight.size(1)) {
+    weight = weight.t().contiguous();
+  }
+  return xllm::kernel::npu::tilelang::causal_conv1d_update(
+      x,
+      conv_state,
+      weight,
+      /*bias=*/std::nullopt,
+      /*conv_state_indices=*/state_indices,
+      /*query_start_loc=*/std::nullopt,
+      /*max_query_len=*/1,
+      /*activation=*/true);
+}
+
+torch::Tensor causal_conv1d_prefill_npu(torch::Tensor x,
+                                        torch::Tensor weight,
+                                        torch::Tensor conv_state,
+                                        torch::Tensor state_indices,
+                                        torch::Tensor has_initial_state,
+                                        torch::Tensor query_start_loc) {
+  // Python layer stores weight as [dim, kernel_width]; CANN expects
+  // [kernel_width, dim].
+  if (weight.size(0) > weight.size(1)) {
+    weight = weight.t().contiguous();
+  }
+
+  // Convert device tensors to host vectors for IntArrayRef parameters.
+  auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous();
+  auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous();
+  auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous();
+
+  std::vector<int64_t> qsl_vec(qsl_cpu.data_ptr<int64_t>(),
+                               qsl_cpu.data_ptr<int64_t>() + qsl_cpu.numel());
+  std::vector<int64_t> si_vec(si_cpu.data_ptr<int64_t>(),
+                              si_cpu.data_ptr<int64_t>() + si_cpu.numel());
+  std::vector<int64_t> ism_vec(ism_cpu.data_ptr<int64_t>(),
+                               ism_cpu.data_ptr<int64_t>() + ism_cpu.numel());
+
+  constexpr int64_t kActivationSilu = 1;
+  constexpr int64_t kPadSlotId = -1;
+  constexpr int64_t kRunModeForward = 0;
+
+  return xllm::kernel::npu::causal_conv1d(
+      x,
+      weight,
+      conv_state,
+      /*bias_opt=*/std::nullopt,
+      torch::IntArrayRef(qsl_vec),
+      torch::IntArrayRef(si_vec),
+      torch::IntArrayRef(ism_vec),
+      /*num_accepted_tokens_opt=*/torch::IntArrayRef{},
+      kActivationSilu,
+      kPadSlotId,
+      kRunModeForward);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+causal_conv1d_qkv_prefill_npu(torch::Tensor x,
+                              torch::Tensor weight,
+                              torch::Tensor conv_state,
+                              torch::Tensor state_indices,
+                              torch::Tensor has_initial_state,
+                              torch::Tensor query_start_loc,
+                              int64_t num_qk_heads,
+                              int64_t num_v_heads,
+                              int64_t head_k_dim,
+                              int64_t head_v_dim) {
+  // Python layer stores weight as [dim, kernel_width]; CANN expects
+  // [kernel_width, dim].
+  if (weight.size(0) > weight.size(1)) {
+    weight = weight.t().contiguous();
+  }
+
+  auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous();
+  auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous();
+  auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous();
+
+  std::vector<int64_t> qsl_vec(qsl_cpu.data_ptr<int64_t>(),
+                               qsl_cpu.data_ptr<int64_t>() + qsl_cpu.numel());
+  std::vector<int64_t> si_vec(si_cpu.data_ptr<int64_t>(),
+                              si_cpu.data_ptr<int64_t>() + si_cpu.numel());
+  std::vector<int64_t> ism_vec(ism_cpu.data_ptr<int64_t>(),
+                               ism_cpu.data_ptr<int64_t>() + ism_cpu.numel());
+
+  return xllm::kernel::npu::causal_conv1d_qkv(x,
+                                              weight,
+                                              conv_state,
+                                              torch::IntArrayRef(qsl_vec),
+                                              torch::IntArrayRef(si_vec),
+                                              torch::IntArrayRef(ism_vec),
+                                              num_qk_heads,
+                                              num_v_heads,
+                                              head_k_dim,
+                                              head_v_dim);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> chunk_gated_delta_rule_npu(
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    torch::Tensor g,
+    torch::Tensor beta,
+    torch::Tensor initial_state,
+    torch::Tensor cu_seqlens) {
+  return xllm::kernel::npu::npu_mega_chunk_gdn(
+      q,
+      k,
+      v,
+      g,
+      beta,
+      /*scale=*/std::nullopt,
+      /*initial_state=*/initial_state,
+      /*output_final_state=*/true,
+      /*cu_seqlens=*/cu_seqlens,
+      /*q_seq_lens=*/{},
+      /*use_qk_l2norm_in_kernel=*/true);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> fused_gdn_gating_npu(
+    torch::Tensor a_log,
+    torch::Tensor a,
+    torch::Tensor b,
+    torch::Tensor dt_bias) {
+  auto [g, beta] = xllm::kernel::npu::tilelang::fused_gdn_gating(
+      a_log.to(torch::kFloat32),
+      a,
+      b,
+      dt_bias.to(torch::kFloat32),
+      /*softplus_beta=*/1.0f,
+      /*softplus_threshold=*/20.0f);
+  return std::make_tuple(g, beta);
+}
+
+torch::Tensor fused_sigmoid_gating_delta_rule_decode_npu(
+    torch::Tensor a_log,
+    torch::Tensor a,
+    torch::Tensor dt_bias,
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    torch::Tensor b,
+    torch::Tensor ssm_state,
+    torch::Tensor state_indices,
+    torch::Tensor cu_seqlens,
+    double scale) {
+  auto a_log_f32 = a_log.to(torch::kFloat32);
+  auto dt_bias_f32 = dt_bias.to(torch::kFloat32);
+  return xllm::kernel::npu::npu_fused_sigmoid_gating_delta_rule_update(
+      a_log_f32,
+      a,
+      dt_bias_f32,
+      q,
+      k,
+      v,
+      b,
+      ssm_state,
+      state_indices,
+      cu_seqlens,
+      /*scale=*/static_cast<float>(scale),
+      /*use_qk_l2norm_in_kernel=*/true,
+      /*softplus_beta=*/1.0f,
+      /*softplus_threshold=*/20.0f);
+}
+
 std::tuple<torch::Tensor, torch::Tensor> fused_add_rms_norm_npu(
     torch::Tensor& input,
     torch::Tensor& residual,
@@ -279,6 +469,35 @@ void ensure_xllm_ops_registered() {
 // compiled only under USE_NPU (mutually exclusive with USE_CUDA).
 TORCH_LIBRARY(xllm_ops, m) {
   m.def("rms_norm(Tensor input, Tensor weight, float eps) -> Tensor");
+  m.def(
+      "rms_norm_gated(Tensor input, Tensor gate, Tensor weight, float eps) -> "
+      "Tensor");
+  m.def("l2_norm(Tensor input, float eps) -> Tensor");
+  m.def(
+      "causal_conv1d_update(Tensor x, Tensor(a!) conv_state, Tensor weight, "
+      "Tensor state_indices) -> Tensor");
+  m.def(
+      "chunk_gated_delta_rule(Tensor q, Tensor k, Tensor v, Tensor g, "
+      "Tensor beta, Tensor initial_state, Tensor cu_seqlens) -> "
+      "(Tensor, Tensor)");
+  m.def(
+      "causal_conv1d_prefill(Tensor x, Tensor weight, Tensor(a!) conv_state, "
+      "Tensor state_indices, Tensor has_initial_state, "
+      "Tensor query_start_loc) -> Tensor");
+  m.def(
+      "causal_conv1d_qkv_prefill(Tensor x, Tensor weight, "
+      "Tensor(a!) conv_state, Tensor state_indices, "
+      "Tensor has_initial_state, Tensor query_start_loc, "
+      "int num_qk_heads, int num_v_heads, "
+      "int head_k_dim, int head_v_dim) -> (Tensor, Tensor, Tensor)");
+  m.def(
+      "fused_gdn_gating(Tensor a_log, Tensor a, Tensor b, Tensor dt_bias) -> "
+      "(Tensor, Tensor)");
+  m.def(
+      "fused_sigmoid_gating_delta_rule_decode(Tensor a_log, Tensor a, "
+      "Tensor dt_bias, Tensor q, Tensor k, Tensor v, Tensor b, "
+      "Tensor(a!) ssm_state, Tensor state_indices, Tensor cu_seqlens, "
+      "float scale) -> Tensor");
   m.def(
       "fused_add_rms_norm(Tensor(a!) input, Tensor(b!) residual, Tensor "
       "weight, "
@@ -353,6 +572,16 @@ TORCH_LIBRARY(xllm_ops, m) {
 
 TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) {
   m.impl("rms_norm", TORCH_FN(xllm::rms_norm_npu));
+  m.impl("rms_norm_gated", TORCH_FN(xllm::rms_norm_gated_npu));
+  m.impl("l2_norm", TORCH_FN(xllm::l2_norm_npu));
+  m.impl("causal_conv1d_update", TORCH_FN(xllm::causal_conv1d_update_npu));
+  m.impl("chunk_gated_delta_rule", TORCH_FN(xllm::chunk_gated_delta_rule_npu));
+  m.impl("causal_conv1d_prefill", TORCH_FN(xllm::causal_conv1d_prefill_npu));
+  m.impl("causal_conv1d_qkv_prefill",
+         TORCH_FN(xllm::causal_conv1d_qkv_prefill_npu));
+  m.impl("fused_gdn_gating", TORCH_FN(xllm::fused_gdn_gating_npu));
+  m.impl("fused_sigmoid_gating_delta_rule_decode",
+         TORCH_FN(xllm::fused_sigmoid_gating_delta_rule_decode_npu));
   m.impl("fused_add_rms_norm", TORCH_FN(xllm::fused_add_rms_norm_npu));
   m.impl("silu_and_mul", TORCH_FN(xllm::silu_and_mul_npu));
   m.impl("reshape_paged_cache", TORCH_FN(xllm::reshape_paged_cache_npu));

xllm/core/kernels/npu/tilelang/CMakeLists.txt

@@ -180,6 +180,7 @@ cc_library(
     tilelang_ops_api.h
   SRCS
     ${TILELANG_KERNEL_SRCS}
+    causal_conv1d_update_wrapper.cpp
   DEPS
     torch
     torch_npu
@@ -200,6 +201,7 @@ target_link_libraries(tilelang_kernels
     "$ENV{NPU_HOME_PATH}/lib64/libascend_dump.so"
     "$ENV{NPU_HOME_PATH}/lib64/libprofapi.so"
     "$ENV{NPU_HOME_PATH}/lib64/libmmpa.so"
+    "$ENV{NPU_HOME_PATH}/lib64/libunified_dlog.so"
     m
     dl
 )