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
This commit is contained in:
project6-dev
29
qwen3_6_scripts/build_corex_gdn_chunk_recurrent.sh
Normal file
29
qwen3_6_scripts/build_corex_gdn_chunk_recurrent.sh
Normal file
@@ -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}"
|
||||
276
qwen3_6_scripts/corex_gdn_chunk_recurrent.cu
Normal file
276
qwen3_6_scripts/corex_gdn_chunk_recurrent.cu
Normal file
@@ -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)");
|
||||
}
|
||||
242
verify_gdn_cpp.py
Normal file
242
verify_gdn_cpp.py
Normal file
@@ -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())
|
||||
Reference in New Issue
Block a user