perf: optimize qwen3.5 gdn decode super operator.

test/python_test/RegisterOps.cpp

@@ -1430,6 +1430,160 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> moe_gating_top_k_hash_impl_npu(
   return std::make_tuple(y, expert_idx, out);
 }
 
+std::tuple<at::Tensor, at::Tensor&, at::Tensor&, at::Tensor>
+qwen35_gdn_decode_super_op(
+    const at::Tensor& qkv,
+    const at::Tensor& z,
+    const at::Tensor& b,
+    const at::Tensor& a,
+    const at::Tensor& conv_weight,
+    at::Tensor& conv_state,
+    const at::Tensor& a_log,
+    const at::Tensor& dt_bias,
+    at::Tensor& ssm_state,
+    const at::Tensor& state_indices,
+    const at::Tensor& norm_weight) {
+  at::Tensor conv_out = at::empty_like(qkv);
+  at::Tensor out = at::empty_like(z);
+  EXEC_NPU_CMD(aclnnQwen35GdnDecodeSuperOp,
+               qkv,
+               z,
+               b,
+               a,
+               conv_weight,
+               conv_state,
+               a_log,
+               dt_bias,
+               ssm_state,
+               state_indices,
+               norm_weight,
+               conv_out,
+               conv_state,
+               ssm_state,
+               out);
+  return std::tuple<at::Tensor, at::Tensor&, at::Tensor&, at::Tensor>(
+      conv_out, conv_state, ssm_state, out);
+}
+
+std::tuple<at::Tensor,
+           at::Tensor,
+           at::Tensor,
+           at::Tensor,
+           at::Tensor,
+           at::Tensor,
+           at::Tensor&,
+           at::Tensor&,
+           at::Tensor>
+qwen35_gdn_prefill_super_op(
+    const at::Tensor& mixed_qkv,
+    const at::Tensor& z,
+    const at::Tensor& b,
+    const at::Tensor& a,
+    const at::Tensor& conv_weight,
+    at::Tensor& conv_state,
+    const at::Tensor& a_log,
+    const at::Tensor& dt_bias,
+    at::Tensor& ssm_state,
+    const at::Tensor& norm_weight,
+    const at::Tensor& mask_lower,
+    const at::Tensor& mask_full,
+    const at::Tensor& minus_identity,
+    const at::Tensor& cu_seqlens,
+    int64_t conv_state_index,
+    int64_t ssm_state_index) {
+  constexpr int64_t kChunkSize = 128;
+  constexpr int64_t kHeadDim = 128;
+  constexpr int64_t kValueHeads = 24;
+  const int64_t total_tokens = mixed_qkv.size(0);
+  const int64_t num_matrices =
+      ((total_tokens + kChunkSize - 1) / kChunkSize) * kValueHeads;
+  auto opts_fp16 = mixed_qkv.options().dtype(at::kHalf);
+  auto opts_fp32 = mixed_qkv.options().dtype(at::kFloat);
+
+  at::Tensor packed_qkv = at::empty({total_tokens, 5120}, opts_fp16);
+  at::Tensor g = at::empty({1, total_tokens, kValueHeads}, opts_fp32);
+  at::Tensor beta = at::empty({1, total_tokens, kValueHeads}, opts_fp16);
+  at::Tensor initial_state =
+      at::empty({1, kValueHeads, kHeadDim, kHeadDim}, opts_fp16);
+  at::Tensor mega_out =
+      at::empty({1, total_tokens, kValueHeads, kHeadDim}, opts_fp16);
+  at::Tensor g_sum =
+      at::empty({1, total_tokens, kValueHeads}, opts_fp32);
+  at::Tensor g_t = at::empty({kValueHeads, total_tokens}, opts_fp32);
+  at::Tensor beta_t = at::empty({kValueHeads, total_tokens}, opts_fp16);
+  at::Tensor mega_a =
+      at::zeros({1, total_tokens, kValueHeads, kChunkSize}, opts_fp16);
+  at::Tensor a_inv_f32 =
+      at::zeros({1, total_tokens, kValueHeads, kChunkSize}, opts_fp32);
+  at::Tensor a_inv =
+      at::zeros({1, total_tokens, kValueHeads, kChunkSize}, opts_fp16);
+  at::Tensor w =
+      at::empty({1, total_tokens, kValueHeads, kHeadDim}, opts_fp16);
+  at::Tensor u = at::empty_like(w);
+  at::Tensor h =
+      at::zeros({num_matrices, kHeadDim, kHeadDim}, opts_fp16);
+  at::Tensor v_new = at::empty_like(w);
+  at::Tensor final_state =
+      at::zeros({kValueHeads, kHeadDim, kHeadDim}, opts_fp16);
+  at::Tensor out =
+      at::empty({total_tokens, kValueHeads, kHeadDim}, mixed_qkv.options());
+
+  EXEC_NPU_CMD(aclnnQwen35GdnPrefillSuperOp,
+               mixed_qkv,
+               z,
+               b,
+               a,
+               conv_weight,
+               conv_state,
+               a_log,
+               dt_bias,
+               ssm_state,
+               norm_weight,
+               mask_lower,
+               mask_full,
+               minus_identity,
+               cu_seqlens,
+               num_matrices,
+               conv_state_index,
+               ssm_state_index,
+               packed_qkv,
+               g,
+               beta,
+               initial_state,
+               mega_out,
+               g_sum,
+               g_t,
+               beta_t,
+               mega_a,
+               a_inv_f32,
+               a_inv,
+               w,
+               u,
+               h,
+               v_new,
+               final_state,
+               conv_state,
+               ssm_state,
+               out);
+  return std::tuple<at::Tensor,
+                    at::Tensor,
+                    at::Tensor,
+                    at::Tensor,
+                    at::Tensor,
+                    at::Tensor,
+                    at::Tensor&,
+                    at::Tensor&,
+                    at::Tensor>(packed_qkv,
+                                g,
+                                beta,
+                                initial_state,
+                                mega_out,
+                                final_state,
+                                conv_state,
+                                ssm_state,
+                                out);
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("select_unshared_kv", &select_unshared_kv_impl_npu, "select_unshared_kv");
   m.def("cache_unshared_kv", &cache_unshared_kv_impl_npu, "cache_unshared_kv");
@@ -1452,6 +1606,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("recurrent_gated_delta_rule", &recurrent_gated_delta_rule, "recurrent_gated_delta_rule");
   m.def("rec_constrained_topk", &rec_constrained_topk_impl_npu, "rec_constrained_topk");
   m.def("mega_chunk_gdn", &mega_chunk_gdn, "mega_chunk_gdn");
+  m.def("qwen35_gdn_decode_super_op",
+        &qwen35_gdn_decode_super_op,
+        "qwen35_gdn_decode_super_op");
+  m.def("qwen35_gdn_prefill_super_op",
+        &qwen35_gdn_prefill_super_op,
+        "qwen35_gdn_prefill_super_op");
   m.def("layer_norm_fwd", &layer_norm_fwd_impl_npu, "layer_norm_fwd");
   m.def("moe_fused_add_topk", &moe_fused_add_topk_impl_npu, "moe_fused_add_topk");
   m.def("moe_fused_reducesum_div", &moe_fused_reducesum_div_impl_npu, "moe_fused_reducesum_div");

test/python_test/test_qwen35_gdn_decode_super_op.py

@@ -0,0 +1,89 @@
+# Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+import pytest
+import torch
+
+torch_npu = pytest.importorskip("torch_npu")
+custom_ops = pytest.importorskip("custom_ops_lib")
+
+
+HEAD_DIM = 128
+LOCAL_HEAD_SHAPES = [
+    (num_k_heads, num_k_heads * ratio)
+    for ratio in range(1, 5)
+    for num_k_heads in (1, 2, 4, 8, 16)
+]
+
+
+def _has_npu():
+    return hasattr(torch, "npu") and torch.npu.is_available()
+
+
+pytestmark = pytest.mark.skipif(not _has_npu(), reason="NPU is required")
+
+
+def _run_case(num_k_heads, num_v_heads, batch_size):
+    conv_dim = (2 * num_k_heads + num_v_heads) * HEAD_DIM
+    bf16_options = {"device": "npu:0", "dtype": torch.bfloat16}
+    float_options = {"device": "npu:0", "dtype": torch.float32}
+
+    qkv = torch.ones((batch_size, conv_dim), **bf16_options)
+    z = torch.zeros((batch_size, num_v_heads, HEAD_DIM), **bf16_options)
+    b = torch.zeros((batch_size, num_v_heads), **bf16_options)
+    a = torch.zeros((batch_size, num_v_heads), **bf16_options)
+    conv_weight = torch.zeros((4, conv_dim), **bf16_options)
+    conv_state = torch.zeros((batch_size, 3, conv_dim), **bf16_options)
+    a_log = torch.zeros((num_v_heads,), **float_options)
+    dt_bias = torch.zeros((num_v_heads,), **float_options)
+    ssm_state = torch.ones(
+        (batch_size, num_v_heads, HEAD_DIM, HEAD_DIM), **float_options
+    )
+    state_indices = torch.arange(batch_size, device="npu:0", dtype=torch.int32)
+    norm_weight = torch.ones((HEAD_DIM,), **bf16_options)
+
+    outputs = custom_ops.qwen35_gdn_decode_super_op(
+        qkv,
+        z,
+        b,
+        a,
+        conv_weight,
+        conv_state,
+        a_log,
+        dt_bias,
+        ssm_state,
+        state_indices,
+        norm_weight,
+    )
+    torch.npu.synchronize()
+
+    assert torch.equal(outputs[3], torch.zeros_like(outputs[3]))
+    assert torch.equal(conv_state[:, :2], torch.zeros_like(conv_state[:, :2]))
+    assert torch.equal(conv_state[:, 2], qkv)
+    assert torch.allclose(
+        ssm_state,
+        torch.full_like(ssm_state, 0.5),
+        rtol=0.0,
+        atol=1e-6,
+    )
+
+
+@pytest.mark.parametrize(
+    ("num_k_heads", "num_v_heads"),
+    LOCAL_HEAD_SHAPES,
+)
+def test_all_qwen35_local_head_shapes(num_k_heads, num_v_heads):
+    _run_case(num_k_heads, num_v_heads, batch_size=1)
+
+
+@pytest.mark.parametrize(
+    ("num_k_heads", "num_v_heads", "batch_size"),
+    [
+        (8, 16, 2),
+        (8, 24, 3),
+        (8, 24, 4),
+        (8, 24, 32),
+        (16, 64, 2),
+    ],
+)
+def test_representative_batch_shapes(num_k_heads, num_v_heads, batch_size):
+    _run_case(num_k_heads, num_v_heads, batch_size)

xllm_ops/build_aclnn.sh

@@ -133,6 +133,8 @@ elif [[ "$SOC_VERSION" =~ ^(ascend)?910b ]]; then
         "cache_unshared_kv"
         "causal_conv1d"
         "causal_conv1d_qkv"
+        "qwen35_gdn_decode_super_op"
+        "qwen35_gdn_prefill_super_op"
         "convert_kv_cache_format"
         "beam_search"
         "index_group_matmul"
@@ -243,6 +245,8 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then
         "cache_unshared_kv"
         "causal_conv1d"
         "causal_conv1d_qkv"
+        "qwen35_gdn_decode_super_op"
+        "qwen35_gdn_prefill_super_op"
         "convert_kv_cache_format"
         "beam_search"
         "index_group_matmul"

xllm_ops/mega_chunk_gdn/op_kernel/mega_chunk_gdn.cpp

@@ -374,8 +374,6 @@ AICORE inline void mega_kernel_impl(GM_ADDR q_ptr, GM_ADDR k_ptr, GM_ADDR v_ptr,
     return;
 #endif
 
-    SyncAllImpl<false>();
-
 #ifdef MEGA_STOP_AFTER_SYNC_BEFORE_WY
     return;
 #endif

xllm_ops/qwen35_gdn_decode_super_op/CMakeLists.txt

@@ -0,0 +1,7 @@
+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+foreach(SUB_DIR ${CURRENT_DIRS})
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
+        add_subdirectory(${SUB_DIR})
+    endif()
+endforeach()
+add_op_to_compiled_list()

xllm_ops/qwen35_gdn_decode_super_op/op_host/CMakeLists.txt

@@ -0,0 +1,17 @@
+if (BUILD_OPEN_PROJECT)
+    target_sources(op_host_aclnn PRIVATE
+        qwen35_gdn_decode_super_op_def.cpp
+    )
+endif()
+
+add_ops_compile_options(
+    OP_NAME Qwen35GdnDecodeSuperOp
+    OPTIONS --cce-auto-sync=on
+            -Wno-deprecated-declarations
+            -Wno-error
+            -I${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/pto-isa/include
+)
+
+if (NOT BUILD_OPS_RTY_KERNEL)
+    add_modules_sources(OPTYPE qwen35_gdn_decode_super_op ACLNNTYPE aclnn)
+endif()

xllm_ops/qwen35_gdn_decode_super_op/op_host/qwen35_gdn_decode_super_op_def.cpp

@@ -0,0 +1,81 @@
+#include "register/op_def_registry.h"
+
+namespace ops {
+
+class Qwen35GdnDecodeSuperOp : public OpDef {
+public:
+    explicit Qwen35GdnDecodeSuperOp(const char* name) : OpDef(name)
+    {
+        this->Input("qkv").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous();
+        this->Input("z").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous();
+        this->Input("b").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous();
+        this->Input("a").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous();
+        this->Input("convWeight")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_BF16})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Input("convState")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_BF16})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Input("aLog")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_FLOAT})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Input("dtBias")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_FLOAT})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Input("ssmState")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_FLOAT})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Input("stateIndices")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_INT32})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Input("normWeight")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_BF16})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+
+        this->Output("convOut")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_BF16})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Output("convStateOut")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_BF16})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Output("ssmStateOut")
+            .ParamType(REQUIRED)
+            .DataType({ge::DT_FLOAT})
+            .FormatList({ge::FORMAT_ND})
+            .AutoContiguous();
+        this->Output("out").ParamType(REQUIRED).DataType({ge::DT_BF16}).FormatList({ge::FORMAT_ND}).AutoContiguous();
+
+        OpAICoreConfig config;
+        config.DynamicCompileStaticFlag(true)
+            .DynamicFormatFlag(false)
+            .DynamicRankSupportFlag(false)
+            .DynamicShapeSupportFlag(true)
+            .NeedCheckSupportFlag(false)
+            .PrecisionReduceFlag(false)
+            .ExtendCfgInfo("coreType.value", "AiCore");
+        this->AICore().AddConfig("ascend910b", config);
+        this->AICore().AddConfig("ascend910_93", config);
+    }
+};
+
+OP_ADD(Qwen35GdnDecodeSuperOp);
+
+} // namespace ops