feat: C++ GDN chunk+recurrent from xllm upstream + verification script

Extracted torch_chunk_gated_delta_rule and torch_recurrent_gated_delta_rule
from xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp.

Pure PyTorch C++ — no NPU/ACL deps, no custom CUDA kernels.
Same algorithm as our Python _torch_chunk_gated_delta_rule but
avoids Python interpreter overhead in the chunk loop.

Verify on real BI-V100: python3 verify_gdn_cpp.py

qwen3_6_scripts/build_corex_gdn_chunk_recurrent.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_gdn_chunk_recurrent.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_gdn_chunk_recurrent.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_gdn_chunk_recurrent \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    -I"${COREX_ROOT}/include" \
+    -I"${SCRIPT_DIR}" \
+    "${SCRIPT_DIR}/corex_gdn_chunk_recurrent.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX GDN chunk+recurrent C++ extension %s\n' "${OUTPUT}"

qwen3_6_scripts/corex_gdn_chunk_recurrent.cu

@@ -0,0 +1,276 @@
+// corex_gdn_chunk_recurrent.cu — C++ GDN chunk + recurrent algorithms
+//
+// Extracted from: xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp
+// These are pure PyTorch C++ implementations — no NPU/ACL/CUDA custom kernels.
+// Benefit: avoids Python loop overhead in _torch_chunk_gated_delta_rule.
+//
+// Functions:
+//   torch_chunk_gated_delta_rule(q,k,v,g,beta, chunk_size, initial_state,
+//                                output_final_state, use_qk_l2norm)
+//     → (core_attn_out, last_recurrent_state)
+//
+//   torch_recurrent_gated_delta_rule(q,k,v,g,beta, initial_state,
+//                                    output_final_state, use_qk_l2norm)
+//     → (core_attn_out, last_recurrent_state)
+
+#include <torch/extension.h>
+#include <optional>
+#include <tuple>
+#include <vector>
+
+namespace {
+
+torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) {
+  auto norm = torch::sqrt(torch::sum(torch::square(x), dim, true) + eps);
+  return x / norm;
+}
+
+torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor,
+                                  int64_t target_heads,
+                                  int64_t head_dim) {
+  const int64_t current_heads = tensor.size(head_dim);
+  if (current_heads == target_heads) {
+    return tensor;
+  }
+  const int64_t repeats = target_heads / current_heads;
+  std::vector<int64_t> view_shape = tensor.sizes().vec();
+  view_shape.insert(view_shape.begin() + head_dim + 1, 1);
+  std::vector<int64_t> expand_shape = view_shape;
+  expand_shape[head_dim + 1] = repeats;
+  std::vector<int64_t> output_shape = tensor.sizes().vec();
+  output_shape[head_dim] = target_heads;
+  return tensor.unsqueeze(head_dim + 1)
+      .expand(expand_shape)
+      .reshape(output_shape)
+      .contiguous();
+}
+
+}  // namespace
+
+std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor g,
+    torch::Tensor beta,
+    c10::optional<torch::Tensor> initial_state,
+    bool output_final_state,
+    bool use_qk_l2norm_in_kernel) {
+  auto initial_dtype = query.dtype();
+
+  if (use_qk_l2norm_in_kernel) {
+    query = l2norm(query, -1, 1e-6);
+    key = l2norm(key, -1, 1e-6);
+  }
+
+  auto to_float32_and_transpose = [](torch::Tensor x) {
+    return x.transpose(1, 2).contiguous().to(torch::kFloat32);
+  };
+  query = to_float32_and_transpose(query);
+  key = to_float32_and_transpose(key);
+  value = to_float32_and_transpose(value);
+  beta = to_float32_and_transpose(beta);
+  g = to_float32_and_transpose(g);
+  const int64_t value_num_heads = value.size(1);
+  query = repeat_tensor_heads(query, value_num_heads, 1);
+  key = repeat_tensor_heads(key, value_num_heads, 1);
+
+  int64_t batch_size = key.size(0);
+  int64_t num_heads = key.size(1);
+  int64_t sequence_length = key.size(2);
+  int64_t k_head_dim = key.size(3);
+  int64_t v_head_dim = value.size(3);
+
+  float scale_val = 1.0f / std::sqrt(static_cast<float>(query.size(-1)));
+  query = query * scale_val;
+
+  torch::Tensor core_attn_out = torch::zeros(
+      {batch_size, num_heads, sequence_length, v_head_dim},
+      torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
+  torch::Tensor last_recurrent_state;
+  if (!initial_state.has_value()) {
+    last_recurrent_state = torch::zeros(
+        {batch_size, num_heads, k_head_dim, v_head_dim},
+        torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
+  } else {
+    last_recurrent_state =
+        initial_state.value().to(value.device(), torch::kFloat32);
+  }
+
+  for (int64_t i = 0; i < sequence_length; ++i) {
+    torch::Tensor q_t = query.select(2, i);
+    torch::Tensor k_t = key.select(2, i);
+    torch::Tensor v_t = value.select(2, i);
+    torch::Tensor g_t = g.select(2, i).exp().unsqueeze(-1).unsqueeze(-1);
+    torch::Tensor beta_t = beta.select(2, i).unsqueeze(-1);
+    last_recurrent_state = last_recurrent_state * g_t;
+    torch::Tensor kv_mem =
+        torch::sum(last_recurrent_state * k_t.unsqueeze(-1), -2);
+    torch::Tensor delta = (v_t - kv_mem) * beta_t;
+    last_recurrent_state =
+        last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2);
+    core_attn_out.select(2, i) =
+        torch::sum(last_recurrent_state * q_t.unsqueeze(-1), -2);
+  }
+
+  core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
+  return std::make_tuple(core_attn_out, last_recurrent_state);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor g,
+    torch::Tensor beta,
+    int64_t chunk_size,
+    c10::optional<torch::Tensor> initial_state,
+    bool output_final_state,
+    bool use_qk_l2norm_in_kernel) {
+  auto initial_dtype = query.dtype();
+  if (use_qk_l2norm_in_kernel) {
+    query = l2norm(query, -1, 1e-6);
+    key = l2norm(key, -1, 1e-6);
+  }
+  auto to_float32 = [](torch::Tensor x) {
+    return x.transpose(1, 2).contiguous().to(torch::kFloat32);
+  };
+
+  query = to_float32(query);
+  key = to_float32(key);
+  value = to_float32(value);
+  beta = to_float32(beta);
+  g = to_float32(g);
+  const int64_t value_num_heads = value.size(1);
+  query = repeat_tensor_heads(query, value_num_heads, 1);
+  key = repeat_tensor_heads(key, value_num_heads, 1);
+
+  int64_t batch_size = query.size(0);
+  int64_t num_heads = query.size(1);
+  int64_t sequence_length = query.size(2);
+  int64_t k_head_dim = key.size(-1);
+  int64_t v_head_dim = value.size(-1);
+
+  int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
+  query = torch::nn::functional::pad(
+      query, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  key = torch::nn::functional::pad(
+      key, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  value = torch::nn::functional::pad(
+      value, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  beta = torch::nn::functional::pad(
+      beta, torch::nn::functional::PadFuncOptions({0, pad_size}));
+  g = torch::nn::functional::pad(
+      g, torch::nn::functional::PadFuncOptions({0, pad_size}));
+
+  int64_t total_sequence_length = sequence_length + pad_size;
+  float scale = 1.0f / std::sqrt(static_cast<float>(query.size(-1)));
+  query = query * scale;
+  auto v_beta = value * beta.unsqueeze(-1);
+  auto k_beta = key * beta.unsqueeze(-1);
+  auto reshape_to_chunks = [chunk_size](torch::Tensor x) {
+    auto shape = x.sizes();
+    std::vector<int64_t> new_shape = {
+        shape[0], shape[1], shape[2] / chunk_size, chunk_size, shape[3]};
+    return x.reshape(new_shape);
+  };
+
+  query = reshape_to_chunks(query);
+  key = reshape_to_chunks(key);
+  value = reshape_to_chunks(value);
+  k_beta = reshape_to_chunks(k_beta);
+  v_beta = reshape_to_chunks(v_beta);
+
+  auto g_shape = g.sizes();
+  std::vector<int64_t> g_new_shape = {
+      g_shape[0], g_shape[1], g_shape[2] / chunk_size, chunk_size};
+  g = g.reshape(g_new_shape);
+  auto mask = torch::triu(
+      torch::ones(
+          {chunk_size, chunk_size},
+          torch::TensorOptions().dtype(torch::kBool).device(query.device())),
+      0);
+
+  g = g.cumsum(-1);
+  auto g_diff = g.unsqueeze(-1) - g.unsqueeze(-2);
+  auto decay_mask = g_diff.tril().exp().to(torch::kFloat32);
+  decay_mask = decay_mask.tril();
+  auto attn = -(torch::matmul(k_beta, key.transpose(-1, -2)) * decay_mask)
+                   .masked_fill(mask, 0.0);
+  for (int64_t i = 1; i < chunk_size; ++i) {
+    if (!attn.is_contiguous()) {
+      attn = attn.contiguous();
+    }
+    auto row = attn.slice(-2, i, i + 1)
+                   .slice(-1, 0, i)
+                   .squeeze(-2)
+                   .clone()
+                   .contiguous();
+    auto sub = attn.slice(-2, 0, i).slice(-1, 0, i).clone().contiguous();
+    auto row_unsq = row.unsqueeze(-1).contiguous();
+    auto row_sub_mul = (row_unsq * sub).contiguous();
+    auto row_sub_sum = row_sub_mul.sum(-2).contiguous();
+    auto row_final = (row + row_sub_sum).contiguous();
+    attn.index_put_({torch::indexing::Ellipsis,
+                     torch::indexing::Slice(i, i + 1),
+                     torch::indexing::Slice(0, i)},
+                    row_final.unsqueeze(-2));
+  }
+
+  attn = attn +
+         torch::eye(
+             chunk_size,
+             torch::TensorOptions().dtype(attn.dtype()).device(attn.device()));
+  value = torch::matmul(attn, v_beta);
+  auto k_cumdecay = torch::matmul(attn, (k_beta * g.exp().unsqueeze(-1)));
+  torch::Tensor last_recurrent_state;
+  if (!initial_state.has_value()) {
+    last_recurrent_state = torch::zeros(
+        {batch_size, num_heads, k_head_dim, v_head_dim},
+        torch::TensorOptions().dtype(value.dtype()).device(value.device()));
+  } else {
+    last_recurrent_state = initial_state.value().to(value);
+  }
+  auto core_attn_out = torch::zeros_like(value);
+  mask = torch::triu(
+      torch::ones(
+          {chunk_size, chunk_size},
+          torch::TensorOptions().dtype(torch::kBool).device(query.device())),
+      1);
+  int64_t num_chunks = total_sequence_length / chunk_size;
+  for (int64_t i = 0; i < num_chunks; ++i) {
+    auto q_i = query.select(2, i);
+    auto k_i = key.select(2, i);
+    auto v_i = value.select(2, i);
+    auto attn_i =
+        (torch::matmul(q_i, k_i.transpose(-1, -2)) * decay_mask.select(2, i))
+            .masked_fill_(mask, 0.0);
+    auto v_prime = torch::matmul(k_cumdecay.select(2, i), last_recurrent_state);
+    auto v_new = v_i - v_prime;
+    auto attn_inter = torch::matmul(q_i * g.select(2, i).unsqueeze(-1).exp(),
+                                    last_recurrent_state);
+    core_attn_out.select(2, i) = attn_inter + torch::matmul(attn_i, v_new);
+    auto g_i_last = g.select(2, i).select(-1, -1).unsqueeze(-1);
+    auto g_exp_term = (g_i_last - g.select(2, i)).exp().unsqueeze(-1);
+    auto k_g_exp = (k_i * g_exp_term).transpose(-1, -2).contiguous();
+    last_recurrent_state = last_recurrent_state * g_i_last.unsqueeze(-1).exp() +
+                           torch::matmul(k_g_exp, v_new);
+  }
+  auto core_attn_out_shape = core_attn_out.sizes();
+  std::vector<int64_t> reshape_shape = {
+      core_attn_out_shape[0],
+      core_attn_out_shape[1],
+      core_attn_out_shape[2] * core_attn_out_shape[3],
+      core_attn_out_shape[4]};
+  core_attn_out = core_attn_out.reshape(reshape_shape);
+  core_attn_out = core_attn_out.slice(2, 0, sequence_length);
+  core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
+  return std::make_tuple(core_attn_out, last_recurrent_state);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("torch_chunk_gated_delta_rule", &torch_chunk_gated_delta_rule,
+        "C++ chunked gated delta rule (from xllm upstream)");
+  m.def("torch_recurrent_gated_delta_rule", &torch_recurrent_gated_delta_rule,
+        "C++ recurrent gated delta rule (from xllm upstream)");
+}

verify_gdn_cpp.py

@@ -0,0 +1,242 @@
+#!/usr/bin/env python3
+"""Verify C++ GDN chunk+recurrent on real BI-V100.
+
+Compiles corex_gdn_chunk_recurrent.cu, then tests:
+1. torch_chunk_gated_delta_rule: C++ vs Python output match
+2. torch_recurrent_gated_delta_rule: C++ vs Python output match
+3. Performance comparison
+
+Qwen3.5 GDN dimensions (TP=4):
+  num_k_heads=4, num_v_heads=8, head_k_dim=128, head_v_dim=128
+  Input: (1, seq_len, 8, 128) for v, (1, seq_len, 4, 128) for q/k
+
+Run: python3 verify_gdn_cpp.py
+"""
+
+import sys
+import os
+import time
+import importlib.util
+import torch
+import torch.nn.functional as F
+
+
+def compile_gdn():
+    script_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                              "qwen3_6_scripts")
+    build_sh = os.path.join(script_dir, "build_corex_gdn_chunk_recurrent.sh")
+    tmp_root = "/tmp/gdn_test"
+    os.makedirs(tmp_root, exist_ok=True)
+    ret = os.system(f"bash {build_sh} {tmp_root} 2>&1")
+    so_path = os.path.join(tmp_root, "corex_gdn_chunk_recurrent.so")
+    if ret != 0 or not os.path.exists(so_path):
+        print(f"[FAIL] Compilation failed (exit={ret})")
+        return None
+    print(f"[OK] Compiled: {so_path}")
+    spec = importlib.util.spec_from_file_location(
+        "corex_gdn_chunk_recurrent", so_path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def python_chunk_gated_delta_rule(q, k, v, g, beta,
+                                   chunk_size=64,
+                                   initial_state=None,
+                                   output_final_state=True,
+                                   use_qk_l2norm_in_kernel=True):
+    """Python reference — same as qwen3_5.py _torch_chunk_gated_delta_rule."""
+    def _l2norm(x, dim=-1, eps=1e-6):
+        norm = torch.sqrt(torch.sum(x ** 2, dim=dim, keepdim=True) + eps)
+        return x / norm
+
+    initial_dtype = q.dtype
+    if use_qk_l2norm_in_kernel:
+        q = _l2norm(q, dim=-1)
+        k = _l2norm(k, dim=-1)
+
+    q = q.transpose(1, 2).contiguous().float()
+    k = k.transpose(1, 2).contiguous().float()
+    v = v.transpose(1, 2).contiguous().float()
+    beta = beta.transpose(1, 2).contiguous().float()
+    g = g.transpose(1, 2).contiguous().float()
+
+    vnh = v.size(1)
+    q = q.repeat_interleave(vnh // q.size(1), dim=1) if q.size(1) != vnh else q
+    k = k.repeat_interleave(vnh // k.size(1), dim=1) if k.size(1) != vnh else k
+
+    B, H, T, Dk = q.shape
+    Dv = v.size(-1)
+    scale = Dk ** -0.5
+    q = q * scale
+
+    pad = (chunk_size - T % chunk_size) % chunk_size
+    if pad > 0:
+        q = F.pad(q, (0, 0, 0, pad))
+        k = F.pad(k, (0, 0, 0, pad))
+        v = F.pad(v, (0, 0, 0, pad))
+        beta = F.pad(beta, (0, pad))
+        g = F.pad(g, (0, pad))
+
+    Tp = T + pad
+    v_beta = v * beta.unsqueeze(-1)
+    k_beta = k * beta.unsqueeze(-1)
+
+    q = q.reshape(B, H, Tp // chunk_size, chunk_size, Dk)
+    k = k.reshape(B, H, Tp // chunk_size, chunk_size, Dk)
+    v = v.reshape(B, H, Tp // chunk_size, chunk_size, Dv)
+    k_beta = k_beta.reshape(B, H, Tp // chunk_size, chunk_size, Dk)
+    v_beta = v_beta.reshape(B, H, Tp // chunk_size, chunk_size, Dv)
+    g = g.reshape(B, H, Tp // chunk_size, chunk_size)
+
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), 0)
+    g = g.cumsum(-1)
+    g_diff = g.unsqueeze(-1) - g.unsqueeze(-2)
+    decay_mask = g_diff.tril().exp().float().tril()
+
+    attn = -(torch.matmul(k_beta, k.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0.0)
+    for i in range(1, chunk_size):
+        row = attn[..., i:i+1, :i].squeeze(-2).clone()
+        sub = attn[..., :i, :i].clone()
+        row_final = row + (row.unsqueeze(-1) * sub).sum(-2)
+        attn[..., i:i+1, :i] = row_final.unsqueeze(-2)
+
+    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+    v = torch.matmul(attn, v_beta)
+    k_cumdecay = torch.matmul(attn, k_beta * g.exp().unsqueeze(-1))
+
+    if initial_state is None:
+        state = torch.zeros(B, H, Dk, Dv, dtype=v.dtype, device=v.device)
+    else:
+        state = initial_state.to(v)
+
+    out = torch.zeros_like(v)
+    mask2 = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), 1)
+    nc = Tp // chunk_size
+    for i in range(nc):
+        qi = q[:, :, i]
+        ki = k[:, :, i]
+        vi = v[:, :, i]
+        ai = (torch.matmul(qi, ki.transpose(-1, -2)) * decay_mask[:, :, i]).masked_fill_(mask2, 0.0)
+        vp = torch.matmul(k_cumdecay[:, :, i], state)
+        vn = vi - vp
+        inter = torch.matmul(qi * g[:, :, i].unsqueeze(-1).exp(), state)
+        out[:, :, i] = inter + torch.matmul(ai, vn)
+        gl = g[:, :, i, -1].unsqueeze(-1)
+        ge = (gl - g[:, :, i]).exp().unsqueeze(-1)
+        kg = (ki * ge).transpose(-1, -2).contiguous()
+        state = state * gl.unsqueeze(-1).exp() + torch.matmul(kg, vn)
+
+    out = out.reshape(B, H, Tp, Dv)[:, :, :T]
+    out = out.transpose(1, 2).contiguous().to(initial_dtype)
+    return out, state
+
+
+def main():
+    print("=" * 60)
+    print("BI-V100 C++ GDN chunk+recurrent verification")
+    print("=" * 60)
+
+    if not torch.cuda.is_available():
+        print("FATAL: No CUDA device")
+        return 1
+
+    mod = compile_gdn()
+    if mod is None:
+        return 1
+
+    # Qwen3.5 GDN dimensions (TP=4)
+    B, T = 1, 128
+    num_k_heads, num_v_heads = 4, 8
+    head_dim = 128
+    chunk_size = 64
+
+    torch.manual_seed(42)
+    q = torch.randn(B, T, num_k_heads, head_dim, device="cuda", dtype=torch.float16)
+    k = torch.randn(B, T, num_k_heads, head_dim, device="cuda", dtype=torch.float16)
+    v = torch.randn(B, T, num_v_heads, head_dim, device="cuda", dtype=torch.float16)
+    g = torch.randn(B, T, num_v_heads, device="cuda", dtype=torch.float16)
+    beta = torch.randn(B, T, num_v_heads, device="cuda", dtype=torch.float16)
+
+    # --- Test 1: chunk ---
+    print(f"\n--- Test 1: torch_chunk_gated_delta_rule (B={B}, T={T}, chunk={chunk_size}) ---")
+    ref_out, ref_state = python_chunk_gated_delta_rule(
+        q.clone(), k.clone(), v.clone(), g.clone(), beta.clone(),
+        chunk_size=chunk_size)
+
+    cpp_out, cpp_state = mod.torch_chunk_gated_delta_rule(
+        q.clone(), k.clone(), v.clone(), g.clone(), beta.clone(),
+        chunk_size, None, True, True)
+
+    diff_out = (ref_out.float() - cpp_out.float()).abs().max().item()
+    diff_state = (ref_state.float() - cpp_state.float()).abs().max().item()
+    print(f"  Output max diff: {diff_out:.8f}")
+    print(f"  State max diff: {diff_state:.8f}")
+    print(f"  Match (tol=1e-2): {diff_out < 1e-2 and diff_state < 1e-2}")
+
+    # --- Test 2: recurrent (decode, T=1) ---
+    print(f"\n--- Test 2: torch_recurrent_gated_delta_rule (B=1, T=1) ---")
+    q1 = torch.randn(1, 1, num_k_heads, head_dim, device="cuda", dtype=torch.float16)
+    k1 = torch.randn(1, 1, num_k_heads, head_dim, device="cuda", dtype=torch.float16)
+    v1 = torch.randn(1, 1, num_v_heads, head_dim, device="cuda", dtype=torch.float16)
+    g1 = torch.randn(1, 1, num_v_heads, device="cuda", dtype=torch.float16)
+    beta1 = torch.randn(1, 1, num_v_heads, device="cuda", dtype=torch.float16)
+    state0 = torch.randn(1, num_v_heads, head_dim, head_dim,
+                          device="cuda", dtype=torch.float32)
+
+    cpp_out1, cpp_state1 = mod.torch_recurrent_gated_delta_rule(
+        q1.clone(), k1.clone(), v1.clone(), g1.clone(), beta1.clone(),
+        state0.clone(), True, True)
+    print(f"  Output shape: {cpp_out1.shape}")
+    print(f"  State shape: {cpp_state1.shape}")
+    print(f"  Output has NaN: {cpp_out1.isnan().any().item()}")
+    print(f"  State has NaN: {cpp_state1.isnan().any().item()}")
+
+    # --- Test 3: Performance ---
+    print(f"\n--- Performance: chunk (B=1, T=512, chunk=64) ---")
+    T_perf = 512
+    q_p = torch.randn(1, T_perf, num_k_heads, head_dim, device="cuda", dtype=torch.float16)
+    k_p = torch.randn(1, T_perf, num_k_heads, head_dim, device="cuda", dtype=torch.float16)
+    v_p = torch.randn(1, T_perf, num_v_heads, head_dim, device="cuda", dtype=torch.float16)
+    g_p = torch.randn(1, T_perf, num_v_heads, device="cuda", dtype=torch.float16)
+    beta_p = torch.randn(1, T_perf, num_v_heads, device="cuda", dtype=torch.float16)
+
+    # Warmup
+    for _ in range(3):
+        mod.torch_chunk_gated_delta_rule(
+            q_p.clone(), k_p.clone(), v_p.clone(), g_p.clone(), beta_p.clone(),
+            64, None, True, True)
+        python_chunk_gated_delta_rule(
+            q_p.clone(), k_p.clone(), v_p.clone(), g_p.clone(), beta_p.clone(),
+            chunk_size=64)
+    torch.cuda.synchronize()
+
+    N = 5
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N):
+        mod.torch_chunk_gated_delta_rule(
+            q_p.clone(), k_p.clone(), v_p.clone(), g_p.clone(), beta_p.clone(),
+            64, None, True, True)
+    torch.cuda.synchronize()
+    cpp_ms = (time.perf_counter() - t0) / N * 1000
+
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N):
+        python_chunk_gated_delta_rule(
+            q_p.clone(), k_p.clone(), v_p.clone(), g_p.clone(), beta_p.clone(),
+            chunk_size=64)
+    torch.cuda.synchronize()
+    py_ms = (time.perf_counter() - t0) / N * 1000
+
+    print(f"  C++: {cpp_ms:.1f} ms")
+    print(f"  Python: {py_ms:.1f} ms")
+    print(f"  Speedup: {py_ms/cpp_ms:.2f}x")
+
+    print("\n" + "=" * 60)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())