perf: optimize causal conv1d and mega chunk gdn for qwen3.5.

test/python_test/RegisterOps.cpp

@@ -173,6 +173,115 @@ at::Tensor causal_conv1d(
     return output;
 }
 
+at::Tensor causal_conv1d_packed_qkv_general_impl(
+    const at::Tensor& x,
+    const at::Tensor& weight,
+    const at::Tensor& conv_state,
+    at::IntArrayRef query_start_loc_opt,
+    at::IntArrayRef cache_indices_opt,
+    at::IntArrayRef initial_state_mode_opt,
+    int64_t q_dim,
+    int64_t k_dim,
+    int64_t v_dim,
+    int64_t head_dim,
+    at::ScalarType output_dtype)
+{
+    constexpr int64_t kPackedQkvActivationMode = 2;
+    constexpr int64_t kPadSlotId = -1;
+    constexpr int64_t kForwardRunMode = 0;
+    at::Tensor output = at::empty(x.sizes(), x.options().dtype(output_dtype));
+    c10::optional<at::Tensor> bias_opt = c10::nullopt;
+    at::IntArrayRef num_accepted_tokens_opt;
+    EXEC_NPU_CMD(aclnnCausalConv1dQkv,
+                 x,
+                 weight,
+                 bias_opt,
+                 conv_state,
+                 query_start_loc_opt,
+                 cache_indices_opt,
+                 initial_state_mode_opt,
+                 num_accepted_tokens_opt,
+                 kPackedQkvActivationMode,
+                 kPadSlotId,
+                 kForwardRunMode,
+                 q_dim,
+                 k_dim,
+                 v_dim,
+                 head_dim,
+                 output);
+    return output;
+}
+
+at::Tensor causal_conv1d_packed_qkv_general(
+    const at::Tensor& x,
+    const at::Tensor& weight,
+    const at::Tensor& conv_state,
+    at::IntArrayRef query_start_loc_opt,
+    at::IntArrayRef cache_indices_opt,
+    at::IntArrayRef initial_state_mode_opt,
+    int64_t q_dim,
+    int64_t k_dim,
+    int64_t v_dim,
+    int64_t head_dim)
+{
+    return causal_conv1d_packed_qkv_general_impl(x,
+                                                 weight,
+                                                 conv_state,
+                                                 query_start_loc_opt,
+                                                 cache_indices_opt,
+                                                 initial_state_mode_opt,
+                                                 q_dim,
+                                                 k_dim,
+                                                 v_dim,
+                                                 head_dim,
+                                                 at::kHalf);
+}
+
+at::Tensor causal_conv1d_packed_qkv_general_bf16(
+    const at::Tensor& x,
+    const at::Tensor& weight,
+    const at::Tensor& conv_state,
+    at::IntArrayRef query_start_loc_opt,
+    at::IntArrayRef cache_indices_opt,
+    at::IntArrayRef initial_state_mode_opt,
+    int64_t q_dim,
+    int64_t k_dim,
+    int64_t v_dim,
+    int64_t head_dim)
+{
+    return causal_conv1d_packed_qkv_general_impl(x,
+                                                 weight,
+                                                 conv_state,
+                                                 query_start_loc_opt,
+                                                 cache_indices_opt,
+                                                 initial_state_mode_opt,
+                                                 q_dim,
+                                                 k_dim,
+                                                 v_dim,
+                                                 head_dim,
+                                                 at::kBFloat16);
+}
+
+at::Tensor causal_conv1d_packed_qkv(
+    const at::Tensor& x,
+    const at::Tensor& weight,
+    const at::Tensor& conv_state,
+    at::IntArrayRef query_start_loc_opt,
+    at::IntArrayRef cache_indices_opt,
+    at::IntArrayRef initial_state_mode_opt)
+{
+    return causal_conv1d_packed_qkv_general(x,
+                                            weight,
+                                            conv_state,
+                                            query_start_loc_opt,
+                                            cache_indices_opt,
+                                            initial_state_mode_opt,
+                                            1024,
+                                            1024,
+                                            3072,
+                                            128);
+}
+
 at::Tensor recurrent_gated_delta_rule(
     const at::Tensor &query,
     const at::Tensor &key,
@@ -1026,6 +1135,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         &beam_search_rec_final_select_impl_npu,
         "beam_search_rec_final_select");
   m.def("causal_conv1d", &causal_conv1d, "causal_conv1d");
+  m.def("causal_conv1d_packed_qkv",
+        &causal_conv1d_packed_qkv,
+        "causal_conv1d_packed_qkv");
+  m.def("causal_conv1d_packed_qkv_general",
+        &causal_conv1d_packed_qkv_general,
+        "causal_conv1d_packed_qkv_general");
+  m.def("causal_conv1d_packed_qkv_general_bf16",
+        &causal_conv1d_packed_qkv_general_bf16,
+        "causal_conv1d_packed_qkv_general_bf16");
   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");

test/python_test/custom_ops.py

@@ -114,10 +114,10 @@ def _mega_get_masks(device):
     return _MEGA_MASK_CACHE[key]
 
 
-def _mega_get_minus_identity(device):
-    key = _mega_device_key(device)
+def _mega_get_minus_identity(device, dtype):
+    key = (*_mega_device_key(device), dtype)
     if key not in _MEGA_MINUS_IDENTITY_CACHE:
-        minus_identity = torch.zeros(CHUNK_SIZE, CHUNK_SIZE, device=device, dtype=torch.float16)
+        minus_identity = torch.zeros(CHUNK_SIZE, CHUNK_SIZE, device=device, dtype=dtype)
         minus_identity.fill_diagonal_(-1)
         _MEGA_MINUS_IDENTITY_CACHE[key] = minus_identity
     return _MEGA_MINUS_IDENTITY_CACHE[key]
@@ -265,7 +265,13 @@ def mega_chunk_gdn_npu(q, k, v, g, beta, scale=None, initial_state=None,
         scale = k.shape[-1] ** -0.5
 
     q_dtype, k_dtype, v_dtype = q.dtype, k.dtype, v.dtype
-    q, k, v, beta = (t.half() for t in (q, k, v, beta))
+    supported_compute_dtypes = (torch.float16, torch.bfloat16)
+    compute_dtype = (
+        q.dtype
+        if q.dtype in supported_compute_dtypes and k.dtype == q.dtype and v.dtype == q.dtype
+        else torch.float16
+    )
+    q, k, v, beta = (t.to(compute_dtype) for t in (q, k, v, beta))
 
     _, total_tokens, _, _ = q.shape
     num_value_heads = v.shape[-2]
@@ -279,19 +285,19 @@ def mega_chunk_gdn_npu(q, k, v, g, beta, scale=None, initial_state=None,
     num_matrices = num_chunks * num_value_heads
     has_initial_state = initial_state is not None
     if has_initial_state:
-        initial_state_arg = initial_state.half()
+        initial_state_arg = initial_state.to(compute_dtype)
     else:
         initial_state_arg = torch.zeros(
             num_sequences,
             num_value_heads,
             q.shape[-1],
             q.shape[-1],
             device=q.device,
-            dtype=torch.float16,
+            dtype=compute_dtype,
         )
 
     mask_lower, mask_full = _mega_get_masks(q.device)
-    minus_identity = _mega_get_minus_identity(q.device)
+    minus_identity = _mega_get_minus_identity(q.device, compute_dtype)
 
     (out, g_sum, _, _, _, _, a_inv, w, _, h, v_new, final_state) = custom_ops_lib.mega_chunk_gdn(
         q,

test/python_test/test_causal_conv1d_qkv.py

@@ -0,0 +1,150 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+
+torch_npu = pytest.importorskip("torch_npu")
+custom_ops_lib = pytest.importorskip("custom_ops_lib")
+
+HEAD_DIM = 128
+WIDTH = 4
+
+
+def _cumulative_lengths(lengths):
+    result = [0]
+    for length in lengths:
+        result.append(result[-1] + length)
+    return tuple(result)
+
+
+def _conv_reference(
+    x,
+    weight,
+    conv_state,
+    lengths,
+    cache_indices,
+    initial_state_mode,
+):
+    outputs = []
+    token_offset = 0
+    weight_fp32 = weight.float()
+    for batch_idx, length in enumerate(lengths):
+        slot = cache_indices[batch_idx]
+        sequence = x[token_offset : token_offset + length]
+        history = (
+            conv_state[slot].clone()
+            if initial_state_mode[batch_idx]
+            else torch.zeros((WIDTH - 1, x.size(1)), dtype=x.dtype)
+        )
+        padded = torch.cat((history, sequence), dim=0)
+        accumulator = torch.zeros_like(sequence, dtype=torch.float32)
+        for tap in range(WIDTH):
+            accumulator.add_(
+                padded[tap : tap + length].float() * weight_fp32[tap]
+            )
+        outputs.append(F.silu(accumulator).to(x.dtype))
+        conv_state[slot].copy_(padded[-(WIDTH - 1) :])
+        token_offset += length
+    return torch.cat(outputs, dim=0)
+
+
+def _packed_reference(conv_output, q_dim, k_dim, v_dim, output_dtype):
+    q, k, v = torch.split(conv_output, (q_dim, k_dim, v_dim), dim=-1)
+
+    def normalize_qk(value):
+        heads = value.view(value.size(0), -1, HEAD_DIM).float()
+        norm = torch.sqrt((heads * heads).sum(dim=-1, keepdim=True) + 1.0e-6)
+        return (heads / norm).to(torch.bfloat16).to(output_dtype)
+
+    return torch.cat(
+        (
+            normalize_qk(q).reshape(-1),
+            normalize_qk(k).reshape(-1),
+            v.to(output_dtype).reshape(-1),
+        )
+    ).view_as(conv_output)
+
+
+@pytest.mark.parametrize(
+    ("q_heads", "v_heads", "lengths", "initial_state_mode"),
+    [
+        pytest.param(16, 48, (64,), (1,), id="tp1"),
+        pytest.param(8, 24, (127, 129), (0, 1), id="tp2-ragged"),
+        pytest.param(4, 12, (256,), (0,), id="tp4"),
+        pytest.param(2, 6, (63, 65), (1, 0), id="tp8-ragged"),
+        pytest.param(1, 3, (129,), (1,), id="tp16"),
+    ],
+)
+@pytest.mark.parametrize(
+    "output_dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]
+)
+def test_causal_conv1d_qkv_general(
+    q_heads,
+    v_heads,
+    lengths,
+    initial_state_mode,
+    output_dtype,
+):
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(20260716 + q_heads + v_heads)
+    q_dim = q_heads * HEAD_DIM
+    k_dim = q_dim
+    v_dim = v_heads * HEAD_DIM
+    conv_dim = q_dim + k_dim + v_dim
+    total_tokens = sum(lengths)
+    num_slots = len(lengths) + 2
+    cache_indices = tuple(range(1, len(lengths) + 1))
+
+    x = (torch.randn((total_tokens, conv_dim), generator=generator) * 0.25).to(
+        torch.bfloat16
+    )
+    weight = (torch.randn((WIDTH, conv_dim), generator=generator) * 0.25).to(
+        torch.bfloat16
+    )
+    conv_state = (
+        torch.randn((num_slots, WIDTH - 1, conv_dim), generator=generator) * 0.1
+    ).to(torch.bfloat16)
+    conv_state_ref = conv_state.clone()
+
+    conv_output = _conv_reference(
+        x,
+        weight,
+        conv_state_ref,
+        lengths,
+        cache_indices,
+        initial_state_mode,
+    )
+    expected = _packed_reference(
+        conv_output, q_dim, k_dim, v_dim, output_dtype
+    )
+
+    device = torch.device("npu")
+    x_npu = x.to(device)
+    weight_npu = weight.to(device)
+    conv_state_npu = conv_state.to(device)
+    op = (
+        custom_ops_lib.causal_conv1d_packed_qkv_general_bf16
+        if output_dtype == torch.bfloat16
+        else custom_ops_lib.causal_conv1d_packed_qkv_general
+    )
+    actual = op(
+        x_npu,
+        weight_npu,
+        conv_state_npu,
+        _cumulative_lengths(lengths),
+        cache_indices,
+        initial_state_mode,
+        q_dim,
+        k_dim,
+        v_dim,
+        HEAD_DIM,
+    )
+    torch.npu.synchronize()
+
+    assert actual.dtype == output_dtype
+    atol = 1.0e-2 if output_dtype == torch.bfloat16 else 1.0e-3
+    rtol = atol
+    torch.testing.assert_close(actual.cpu(), expected, atol=atol, rtol=rtol)
+    torch.testing.assert_close(
+        conv_state_npu.cpu(), conv_state_ref, atol=0.0, rtol=0.0
+    )