perf(deltanet): CCCL block_scan_raking pattern — replace Python loop with solve_triangular

CCCL block_scan_raking.cuh: parallel prefix scan over C elements using
GPU-native raking threads, not sequential host-driven loops.

Our _forward_sub_lower was a Python for-loop over chunk_size=64 rows,
each launching a separate matmul kernel. This is 64 sequential kernel
launches per DeltaNet layer per chunk.

Fix: Use torch.linalg.solve_triangular (cuBLAS trsm) which solves
the entire (I-A)@X=RHS system in ONE kernel launch. Falls back to
the Python loop if cuSOLVER is unavailable on BI-V100.

CCCL source: cub/cub/block/specializations/block_scan_raking.cuh
Maps to: qwen3_6_scripts/qwen3_5.py (_forward_sub_lower)
This commit is contained in:
project6
2026-08-07 08:57:51 +00:00
parent a1558b6e50
commit 86ca125b47

View File

@@ -140,14 +140,28 @@ def _torch_chunk_gated_delta_rule(
A_lower: (..., C, C) strictly lower-triangular
rhs: (..., C, D)
Returns X: (..., C, D)
CCCL block_scan_raking.cuh insight: sequential scan over C elements
is the bottleneck. torch.linalg.solve_triangular delegates to
cuBLAS trsm which is O(C²) but fully GPU-parallel, vs our Python
loop which is O(C²) but with C sequential kernel launches.
"""
C = rhs.shape[-2]
x = torch.zeros_like(rhs)
x[..., 0, :] = rhs[..., 0, :]
for i in range(1, C):
# x[i] = rhs[i] + A[i, :i] @ x[:i]
x[..., i, :] = rhs[..., i, :] + (A_lower[..., i, :i].unsqueeze(-2) @ x[..., :i, :]).squeeze(-2)
return x
# Build (I - A_lower) which is unit lower-triangular
eye = torch.eye(C, dtype=A_lower.dtype, device=A_lower.device)
IminusA = eye - A_lower
try:
# cuBLAS trsm: solve IminusA @ X = rhs for X
# unitriangular=True tells solver diagonal is all 1s (skip division)
return torch.linalg.solve_triangular(
IminusA, rhs, upper=False, unitriangular=True)
except RuntimeError:
# BI-V100 may lack cuSOLVER — fall back to row-by-row
x = torch.zeros_like(rhs)
x[..., 0, :] = rhs[..., 0, :]
for i in range(1, C):
x[..., i, :] = rhs[..., i, :] + (A_lower[..., i, :i].unsqueeze(-2) @ x[..., :i, :]).squeeze(-2)
return x
value = _forward_sub_lower(A, v_beta)