EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

167
tests/e2e/nightly/310p/single_node/ops/singlecard_ops/test_chunk_fwd_o_310.py Normal file
View File

@@ -0,0 +1,167 @@
"""
chunk_fwd_o correctness tests on Ascend 310P via torch.ops._C_ascend binding.
"""
import pytest
import torch
import torch_npu # noqa: F401
from vllm_ascend.utils import enable_custom_op
CHUNK_SIZE = 64
def npu_chunk_fwd_o(q, k, v, h, g, scale):
enable_custom_op()
return torch.ops._C_ascend.chunk_fwd_o(
q,
k,
v,
h,
scale,
g=g,
g_gamma=None,
cu_seqlens=None,
chunk_indices=None,
chunk_size=CHUNK_SIZE,
transpose_state_layout=False,
)
def golden_chunk_fwd_o(q, k, v, h_state, g, scale):
"""CPU fp32 reference.
Per chunk c (CS tokens starting at t0):
attn = q[c] @ k[c].T [CS, CS]
gate[i,j] = exp(min(0, g[j] - g[i])) * (j<=i) [CS, CS]
attn_masked = attn * gate
h_work = q[c] @ h_state[c] [CS, Dv]
v_work = attn_masked @ v[c] [CS, Dv]
o[c] = scale * (v_work + exp(g[c]) * h_work)
"""
q, k, v, g = q.float(), k.float(), v.float(), g.float()
h_state = h_state.float()
B, H_k, L, D_k = q.shape
H_v, D_v = v.shape[1], v.shape[3]
CS = CHUNK_SIZE
NT = L // CS
head_groups = H_v // H_k
o = torch.zeros(B, H_v, L, D_v)
for b in range(B):
for hv in range(H_v):
hk = hv // head_groups
for c in range(NT):
t0 = c * CS
q_c = q[b, hk, t0 : t0 + CS]
k_c = k[b, hk, t0 : t0 + CS]
v_c = v[b, hv, t0 : t0 + CS]
g_c = g[b, hv, t0 : t0 + CS]
h_c = h_state[b, hv, c * D_k : (c + 1) * D_k]
attn = q_c @ k_c.T
g_row = g_c.unsqueeze(1)
g_col = g_c.unsqueeze(0)
gate = torch.exp(torch.clamp(g_col - g_row, max=0.0))
causal = torch.tril(torch.ones(CS, CS))
attn_masked = attn * gate * causal
h_work = q_c @ h_c
v_work = attn_masked @ v_c
g_exp = torch.exp(g_c).unsqueeze(1)
o[b, hv, t0 : t0 + CS] = scale * (v_work + g_exp * h_work)
return o
class TestChunkFwdO310:
"""chunk_fwd_o kernel correctness on Ascend 310P."""
@pytest.mark.parametrize(
"B,Hk,Hv,L,Dk,Dv",
[
(1, 2, 2, 128, 128, 128),
(1, 4, 4, 256, 128, 128),
],
)
def test_constant_inputs(self, B, Hk, Hv, L, Dk, Dv):
"""Constant q=k=v, h=0, g=0 => analytically verifiable output."""
scale = 1.0 / (Dk**0.5)
NC = L // CHUNK_SIZE
c = 0.01
q = torch.full((B, Hk, L, Dk), c, dtype=torch.float16).npu()
k = torch.full((B, Hk, L, Dk), c, dtype=torch.float16).npu()
v = torch.full((B, Hv, L, Dv), c, dtype=torch.float16).npu()
h = torch.zeros(B, Hv, NC * Dk, Dv, dtype=torch.float16).npu()
g = torch.zeros(B, Hv, L, dtype=torch.float32).npu()
o = npu_chunk_fwd_o(q, k, v, h, g, scale)
oc = o.cpu().float()
assert torch.isnan(oc).sum() == 0, "output has NaN"
assert torch.isinf(oc).sum() == 0, "output has Inf"
attn_val = c * c * Dk
for i in range(min(CHUNK_SIZE, 8)):
expected = scale * (i + 1) * attn_val * c
actual = oc[0, 0, i, 0].item()
rel_err = abs(actual - expected) / max(abs(expected), 1e-10)
assert rel_err < 0.10, f"row {i}: actual={actual:.8f} expected={expected:.8f} rel_err={rel_err:.2f}"
@pytest.mark.parametrize(
"B,Hk,Hv,L,Dk,Dv",
[
(1, 2, 2, 128, 128, 128),
(1, 4, 4, 256, 128, 128),
],
)
def test_random_inputs_no_nan(self, B, Hk, Hv, L, Dk, Dv):
"""Random small inputs: no NaN/Inf in output."""
torch.manual_seed(42)
scale = 1.0 / (Dk**0.5)
NC = L // CHUNK_SIZE
q = (torch.randn(B, Hk, L, Dk) * 0.01).half().npu()
k = (torch.randn(B, Hk, L, Dk) * 0.01).half().npu()
v = (torch.randn(B, Hv, L, Dv) * 0.01).half().npu()
h = (torch.randn(B, Hv, NC * Dk, Dv) * 0.01).half().npu()
g = torch.randn(B, Hv, L, dtype=torch.float32).npu() * 0.001
o = npu_chunk_fwd_o(q, k, v, h, g, scale)
oc = o.cpu().float()
assert torch.isnan(oc).sum() == 0, "output has NaN"
assert torch.isinf(oc).sum() == 0, "output has Inf"
assert oc.abs().max() > 0, "output is all zeros"
def test_g_zero_reduces_to_standard_attention(self):
"""g=0 => gate=1, so kernel = scale*(causal_attn@v + q@h)."""
torch.manual_seed(123)
B, Hk, Hv, L, Dk, Dv = 1, 2, 2, 128, 128, 128
scale = 1.0 / (Dk**0.5)
NC = L // CHUNK_SIZE
q = (torch.randn(B, Hk, L, Dk) * 0.01).half()
k = (torch.randn(B, Hk, L, Dk) * 0.01).half()
v = (torch.randn(B, Hv, L, Dv) * 0.01).half()
h = torch.zeros(B, Hv, NC * Dk, Dv, dtype=torch.float16)
g = torch.zeros(B, Hv, L, dtype=torch.float32)
o_npu = npu_chunk_fwd_o(q.npu(), k.npu(), v.npu(), h.npu(), g.npu(), scale)
o_ref = golden_chunk_fwd_o(q, k, v, h, g, scale)
cos = torch.nn.functional.cosine_similarity(o_npu.cpu().float().flatten(), o_ref.flatten(), dim=0).item()
assert cos > 0.999, f"cosine {cos:.4f} too low for g=0 h=0 case"
def test_chunk_boundary_independence(self):
"""Each chunk should produce the same output for identical data."""
B, Hk, Hv, L, Dk, Dv = 1, 2, 2, 128, 128, 128
scale = 1.0 / (Dk**0.5)
NC = L // CHUNK_SIZE
c = 0.02
q = torch.full((B, Hk, L, Dk), c, dtype=torch.float16).npu()
k = torch.full((B, Hk, L, Dk), c, dtype=torch.float16).npu()
v = torch.full((B, Hv, L, Dv), c, dtype=torch.float16).npu()
h = torch.zeros(B, Hv, NC * Dk, Dv, dtype=torch.float16).npu()
g = torch.zeros(B, Hv, L, dtype=torch.float32).npu()
o = npu_chunk_fwd_o(q, k, v, h, g, scale).cpu().float()
chunk0 = o[0, 0, :CHUNK_SIZE, :]
chunk1 = o[0, 0, CHUNK_SIZE:, :]
cos = torch.nn.functional.cosine_similarity(chunk0.flatten(), chunk1.flatten(), dim=0).item()
assert cos > 0.999, f"chunks differ: cosine={cos:.6f}"

160
tests/e2e/nightly/310p/single_node/ops/singlecard_ops/test_chunk_gated_delta_rule_fwd_h_aclnn.py Normal file
View File

@@ -0,0 +1,160 @@
"""
chunk_gated_delta_rule_fwd_h correctness tests on Ascend 310P
via torch.ops._C_ascend binding.
"""
import pytest
import torch
import torch_npu # noqa: F401
from vllm_ascend.utils import enable_custom_op
CHUNK_SIZE = 64
def npu_chunk_gdr_fwd_h(k, w, u, g, initial_state=None, chunk_size=64):
enable_custom_op()
return torch.ops._C_ascend.chunk_gated_delta_rule_fwd_h(
k,
w,
u,
g=g,
initial_state=initial_state,
output_final_state=False,
chunk_size=chunk_size,
save_new_value=True,
)
def cpu_reference(k, w, u, g, initial_state=None, chunk_size=64):
"""CPU fp32 reference matching kernel semantics."""
k, w, u, g = k.float(), w.float(), u.float(), g.float()
B, Hg, T, K = k.shape
HV, V = u.shape[1], u.shape[3]
NT = T // chunk_size
h = initial_state.float().clone() if initial_state is not None else torch.zeros(B, HV, K, V)
h_chunks = [h.clone()]
v_new = torch.zeros_like(u)
for c in range(NT):
t0 = c * chunk_size
W_chunk = w[:, :, t0 : t0 + chunk_size, :]
ws = torch.einsum("bhik,bhkv->bhiv", W_chunk, h)
g_chunk = g[:, :, t0 : t0 + chunk_size]
v_update = torch.zeros(B, HV, chunk_size, V)
for i in range(chunk_size):
gi_cum = g_chunk[:, :, -1] - g_chunk[:, :, i]
vn = u[:, :, t0 + i, :] - ws[:, :, i, :]
v_new[:, :, t0 + i, :] = vn
v_update[:, :, i, :] = gi_cum.unsqueeze(-1).exp() * vn
K_chunk = k[:, :, t0 : t0 + chunk_size, :]
h_work = torch.einsum("bhik,bhiv->bhkv", K_chunk, v_update)
h = h * g_chunk[:, :, -1:].unsqueeze(-1).exp() + h_work
h_chunks.append(h.clone())
return h_chunks, v_new
def cosine(a, b):
a, b = a.flatten().double(), b.flatten().double()
if a.norm() == 0 and b.norm() == 0:
return 1.0
if a.norm() == 0 or b.norm() == 0:
return 0.0
return torch.nn.functional.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item()
class TestChunkGatedDeltaRuleFwdH310:
"""chunk_gated_delta_rule_fwd_h kernel correctness on Ascend 310P."""
@pytest.mark.parametrize(
"B,Hg,HV,T,K,V",
[
(1, 1, 1, 128, 128, 128),
(1, 2, 2, 128, 128, 128),
],
)
def test_h_state_correctness(self, B, Hg, HV, T, K, V):
torch.manual_seed(42)
DTYPE = torch.float16
k = torch.randn(B, Hg, T, K, dtype=DTYPE) * 0.1
w = torch.randn(B, Hg, T, K, dtype=DTYPE) * 0.1
u = torch.randn(B, HV, T, V, dtype=DTYPE) * 0.1
g = (-torch.rand(B, HV, T) * 0.1).float()
init = torch.randn(B, HV, K, V, dtype=DTYPE) * 0.01
h_ref, _ = cpu_reference(k, w, u, g, init, CHUNK_SIZE)
h_out, _, _ = npu_chunk_gdr_fwd_h(
k.npu(),
w.npu(),
u.npu(),
g.npu(),
initial_state=init.npu(),
chunk_size=CHUNK_SIZE,
)
h_npu = h_out.cpu().float()
NT = T // CHUNK_SIZE
for c in range(min(NT + 1, h_npu.shape[2])):
ref = h_ref[c].flatten()
npu = h_npu[0, :, c].flatten()
cos = cosine(npu, ref)
assert cos >= 0.99, f"h[{c}] cos={cos:.6f} too low"
@pytest.mark.parametrize(
"B,Hg,HV,T,K,V",
[
(1, 1, 1, 128, 128, 128),
(1, 2, 2, 128, 128, 128),
],
)
def test_v_new_correctness(self, B, Hg, HV, T, K, V):
torch.manual_seed(42)
DTYPE = torch.float16
k = torch.randn(B, Hg, T, K, dtype=DTYPE) * 0.1
w = torch.randn(B, Hg, T, K, dtype=DTYPE) * 0.1
u = torch.randn(B, HV, T, V, dtype=DTYPE) * 0.1
g = (-torch.rand(B, HV, T) * 0.1).float()
init = torch.randn(B, HV, K, V, dtype=DTYPE) * 0.01
_, vn_ref = cpu_reference(k, w, u, g, init, CHUNK_SIZE)
_, vn_out, _ = npu_chunk_gdr_fwd_h(
k.npu(),
w.npu(),
u.npu(),
g.npu(),
initial_state=init.npu(),
chunk_size=CHUNK_SIZE,
)
vn_npu = vn_out.cpu().float()
NT = T // CHUNK_SIZE
for c in range(NT):
t0, t1 = c * CHUNK_SIZE, (c + 1) * CHUNK_SIZE
ref = vn_ref[:, :, t0:t1].flatten()
npu = vn_npu[:, :, t0:t1].flatten()
cos = cosine(npu, ref)
assert cos >= 0.99, f"v_new chunk {c} cos={cos:.6f} too low"
def test_no_nan(self):
torch.manual_seed(42)
B, Hg, HV, T, K, V = 1, 1, 1, 128, 128, 128
DTYPE = torch.float16
k = (torch.randn(B, Hg, T, K, dtype=DTYPE) * 0.1).npu()
w = (torch.randn(B, Hg, T, K, dtype=DTYPE) * 0.1).npu()
u = (torch.randn(B, HV, T, V, dtype=DTYPE) * 0.1).npu()
g = (-torch.rand(B, HV, T).float() * 0.1).npu()
init = (torch.randn(B, HV, K, V, dtype=DTYPE) * 0.01).npu()
h_out, vn_out, _ = npu_chunk_gdr_fwd_h(
k,
w,
u,
g,
initial_state=init,
chunk_size=CHUNK_SIZE,
)
assert torch.isnan(h_out.cpu()).sum() == 0, "h_out has NaN"
assert torch.isnan(vn_out.cpu()).sum() == 0, "vn_out has NaN"
assert torch.isinf(h_out.cpu()).sum() == 0, "h_out has Inf"
assert torch.isinf(vn_out.cpu()).sum() == 0, "vn_out has Inf"

178
tests/e2e/nightly/310p/single_node/ops/singlecard_ops/test_recurrent_gated_delta_rule_v310.py Normal file
View File

@@ -0,0 +1,178 @@
"""Test V310 kernel via ctypes API against golden CPU reference."""
import ctypes
import os
import pytest
import torch
import torch_npu
torch_npu.npu.set_compile_mode(jit_compile=False)
_CANN = os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest")
_CUST = f"{_CANN}/opp/vendors/custom_transformer/op_api/lib"
_LIB_PATHS = [_CUST, f"{_CANN}/lib64", f"{_CANN}/aarch64-linux/lib64"]
def _find_lib(name, paths):
for p in paths:
full = os.path.join(p, name)
if os.path.exists(full):
return full
return name
_acl = ctypes.CDLL(_find_lib("libnnopbase.so", _LIB_PATHS))
_opapi = ctypes.CDLL(_find_lib("libcust_opapi.so", _LIB_PATHS))
_acl.aclCreateTensor.restype = ctypes.c_void_p
_acl.aclCreateTensor.argtypes = [
ctypes.POINTER(ctypes.c_int64),
ctypes.c_uint64,
ctypes.c_int,
ctypes.POINTER(ctypes.c_int64),
ctypes.c_int64,
ctypes.c_int,
ctypes.POINTER(ctypes.c_int64),
ctypes.c_uint64,
ctypes.c_void_p,
]
_acl.aclDestroyTensor.argtypes = [ctypes.c_void_p]
_DTYPE_MAP = {torch.float16: 1, torch.float32: 0, torch.int32: 3}
def mk(t):
if t is None:
return None
shape, strides, ndim = list(t.shape), list(t.stride()), len(t.shape)
return ctypes.c_void_p(
_acl.aclCreateTensor(
(ctypes.c_int64 * ndim)(*shape),
ndim,
_DTYPE_MAP[t.dtype],
(ctypes.c_int64 * ndim)(*strides),
ctypes.c_int64(0),
2,
(ctypes.c_int64 * ndim)(*shape),
ndim,
ctypes.c_void_p(t.data_ptr()),
)
)
def call_v310(query, key, value, beta, state, seq_lens, indices, g, nat, scale):
out = torch.empty_like(value)
ws_size = ctypes.c_uint64(0)
executor = ctypes.c_void_p(0)
ret = _opapi.aclnnRecurrentGatedDeltaRuleV310GetWorkspaceSize(
mk(query),
mk(key),
mk(value),
mk(beta),
mk(state),
mk(seq_lens),
mk(indices),
mk(g),
None,
mk(nat),
ctypes.c_float(scale),
mk(out),
ctypes.byref(ws_size),
ctypes.byref(executor),
)
assert ret == 0, f"GetWorkspaceSize failed: {ret}"
ws_ptr = ctypes.c_void_p(0)
if ws_size.value > 0:
ws = torch.empty(ws_size.value, dtype=torch.uint8, device=query.device)
ws_ptr = ctypes.c_void_p(ws.data_ptr())
stream = torch.npu.current_stream().npu_stream
ret = _opapi.aclnnRecurrentGatedDeltaRuleV310(ws_ptr, ws_size, executor, ctypes.c_void_p(stream))
assert ret == 0, f"Execute failed: {ret}"
torch.npu.synchronize()
return out
def golden(query, key, value, state, beta, scale, seq_lens, indices, g, nat):
k = key.float()
q = query.float()
v = value.float()
S = state.clone().float()
T, nv, Dv = v.shape
nk = q.shape[1]
g_f = torch.ones(T, nv) if g is None else g.float().exp()
beta_f = beta.float()
o = torch.empty_like(v, dtype=torch.float32)
q = q * scale
seq_start = 0
for i in range(len(seq_lens)):
init_idx = indices[seq_start + nat[i] - 1] if nat is not None else indices[seq_start]
for head in range(nv):
s = S[init_idx][head].clone()
for t in range(seq_start, seq_start + seq_lens[i]):
qi = q[t][head // (nv // nk)]
ki = k[t][head // (nv // nk)]
vi = v[t][head]
s = s * g_f[t][head]
x = (s * ki.unsqueeze(-2)).sum(dim=-1)
y = (vi - x) * beta_f[t][head]
s = s + y[:, None] * ki[None, :]
S[indices[t]][head] = s
o[t][head] = (s * qi.unsqueeze(-2)).sum(dim=-1)
seq_start += seq_lens[i]
return o, S
@pytest.mark.parametrize(
"batch_size,mtp,nk,nv,dk,dv,num_slots",
[
(1, 1, 8, 16, 128, 128, 444),
(2, 2, 8, 16, 128, 128, 444),
(4, 2, 4, 4, 64, 64, 32),
],
)
def test_recurrent_gated_delta_rule_v310(batch_size, mtp, nk, nv, dk, dv, num_slots):
torch.manual_seed(42)
scale = dk**-0.5
seq_lens = torch.ones(batch_size, dtype=torch.int32) * mtp
T = int(seq_lens.sum())
state = torch.rand(num_slots, nv, dv, dk, dtype=torch.float16)
indices = torch.randperm(num_slots, dtype=torch.int32)[:T]
nat = torch.ones(batch_size, dtype=torch.int32)
query = torch.nn.functional.normalize(torch.randn(T, nk, dk), dim=-1).to(torch.float16)
key = torch.nn.functional.normalize(torch.randn(T, nk, dk), dim=-1).to(torch.float16)
value = torch.randn(T, nv, dv, dtype=torch.float16)
beta = torch.rand(T, nv, dtype=torch.float16)
g = torch.rand(T, nv, dtype=torch.float32)
out_gold, state_gold = golden(query, key, value, state, beta, scale, seq_lens, indices, g, nat)
state_npu = state.clone().npu()
out_npu = call_v310(
query.npu(),
key.npu(),
value.npu(),
beta.npu(),
state_npu,
seq_lens.npu(),
indices.npu(),
g.npu(),
nat.npu(),
scale,
)
touched = indices.long()
torch.testing.assert_close(
out_npu.float().cpu(),
out_gold,
rtol=3e-3,
atol=2e-3,
equal_nan=True,
)
torch.testing.assert_close(
state_npu.float().cpu()[touched],
state_gold.float()[touched],
rtol=3e-3,
atol=2e-3,
equal_nan=True,
)