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:
@@ -140,14 +140,28 @@ def _torch_chunk_gated_delta_rule(
|
|||||||
A_lower: (..., C, C) strictly lower-triangular
|
A_lower: (..., C, C) strictly lower-triangular
|
||||||
rhs: (..., C, D)
|
rhs: (..., C, D)
|
||||||
Returns X: (..., 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]
|
C = rhs.shape[-2]
|
||||||
x = torch.zeros_like(rhs)
|
# Build (I - A_lower) which is unit lower-triangular
|
||||||
x[..., 0, :] = rhs[..., 0, :]
|
eye = torch.eye(C, dtype=A_lower.dtype, device=A_lower.device)
|
||||||
for i in range(1, C):
|
IminusA = eye - A_lower
|
||||||
# x[i] = rhs[i] + A[i, :i] @ x[:i]
|
try:
|
||||||
x[..., i, :] = rhs[..., i, :] + (A_lower[..., i, :i].unsqueeze(-2) @ x[..., :i, :]).squeeze(-2)
|
# cuBLAS trsm: solve IminusA @ X = rhs for X
|
||||||
return 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)
|
value = _forward_sub_lower(A, v_beta)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user