Skip to content

[WY representation] Faster lower triangle inverse #289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yzhangcs merged 14 commits into from
Apr 3, 2025
Merged

[WY representation] Faster lower triangle inverse #289

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
027d951
[WY representation] Faster lower triangle inverse
sustcsonglin Apr 3, 2025
0442fdc
[WY] Fix when `cu_seqlens` is None
zhiyuan1i Apr 3, 2025
566a503
run pre-commit
zhiyuan1i Apr 3, 2025
912c740
Fix code formats
yzhangcs Apr 3, 2025
c2674c1
Minor update
yzhangcs Apr 3, 2025
d55608e
WY repr by fast tril solve
yzhangcs Apr 3, 2025
b00e556
Minor fix
yzhangcs Apr 3, 2025
415ceb0
Tests with fixed seed
yzhangcs Apr 3, 2025
b08bb88
Merge branch 'main' into faster-deltanet
yzhangcs Apr 3, 2025
527d742
Tested under stable version
yzhangcs Apr 3, 2025
e578a23
Fix index errs in varlen tests
yzhangcs Apr 3, 2025
5f272f3
Merge remote-tracking branch 'origin/main' into faster-deltanet
zhiyuan1i Apr 3, 2025
1446e9e
[Test] Skip test on Intel GPU
zhiyuan1i Apr 3, 2025
67d86a1
skip
zhiyuan1i Apr 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 126 additions & 0 deletions fla/ops/common/chunk_scaled_dot_kkt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

from typing import Optional

import torch
import triton
import triton.language as tl

from fla.ops.common.utils import prepare_chunk_indices


@triton.heuristics({
'USE_OFFSETS': lambda args: args['offsets'] is not None
})
@triton.autotune(
configs=[
triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages)
for BK in [32, 64, 128]
for num_warps in [2, 4, 8]
for num_stages in [2, 3, 4]
],
key=['H', 'K', 'BT', 'USE_OFFSETS'],
)
@triton.jit(do_not_specialize=['T'])
def chunk_scaled_dot_kkt_fwd_kernel(
k,
beta,
A,
offsets,
indices,
T,
H: tl.constexpr,
K: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
HEAD_FIRST: tl.constexpr,
USE_OFFSETS: tl.constexpr,
):
i_t, i_bh = tl.program_id(0), tl.program_id(1)
i_b, i_h = i_bh // H, i_bh % H
if USE_OFFSETS:
i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_t = tl.arange(0, BT)

Comment on lines +39 to +49

@coderabbitai coderabbitai Bot Apr 3, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Consider renaming i_t to avoid overshadowing.
Here, i_t is read from tl.program_id(0), but then overwritten via the offsets logic:

i_n, i_t = tl.load(indices + i_t * 2), tl.load(indices + i_t * 2 + 1)

This can create confusion and obscure the kernel’s indexing logic. A clearer naming convention for the chunk index (e.g., chunk_id) would improve maintainability.

Below is a sample diff showing how to rename the overwritten variable:

- i_t, i_bh = tl.program_id(0), tl.program_id(1)
- if USE_OFFSETS:
-     i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32)
+ chunk_id, i_bh = tl.program_id(0), tl.program_id(1)
+ if USE_OFFSETS:
+     i_n, local_t = tl.load(indices + chunk_id * 2).to(tl.int32), tl.load(indices + chunk_id * 2 + 1).to(tl.int32)

if HEAD_FIRST:
p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
else:
p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_beta = tl.load(p_beta, boundary_check=(0,))

b_A = tl.zeros([BT, BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
if HEAD_FIRST:
p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
else:
p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
b_kb = b_k * b_beta[:, None]
b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k))

b_A = tl.where(o_t[:, None] > o_t[None, :], b_A, 0)
if HEAD_FIRST:
p_A = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
else:
p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0))
tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))


def chunk_scaled_dot_kkt_fwd(
k: torch.Tensor,
beta: torch.Tensor,
cu_seqlens: Optional[torch.LongTensor],
head_first: bool = False,
chunk_size: int = 64,
output_dtype: torch.dtype = torch.float32
) -> torch.Tensor:
r"""
Compute beta * K * K^T.

Args:
k (torch.Tensor):
The key tensor of shape `[B, T, H, K]` if not `head_first` else `[B, H, T, K]`.
beta (torch.Tensor):
The beta tensor of shape `[B, T, H]` if not `head_first` else `[B, H, T]`.
cu_seqlens (torch.LongTensor):
The cumulative sequence lengths of the input tensor.
Default: None
head_first (bool):
If False, the input/output tensor is in the shape of `[B, T, H, K]`.
If True, the input/output tensor is in the shape of `[B, H, T, K]`.
Default: False
chunk_size (int):
The chunk size. Default: 64.
output_dtype (torch.dtype):
The dtype of the output tensor. Default: `torch.float32`

Returns:
beta * K * K^T of shape `[B, T, H, BT]` if not `head_first` else `[B, H, T, BT]`,
where `BT` is the chunk size.
"""
if head_first:
B, H, T, K = k.shape
else:
B, T, H, K = k.shape
BT = chunk_size
indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices)
A = torch.empty(B, *((H, T) if head_first else (T, H)), BT, device=k.device, dtype=output_dtype)
chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)](
k=k,
beta=beta,
A=A,
offsets=cu_seqlens,
indices=indices,
T=T,
H=H,
K=K,
BT=BT,
HEAD_FIRST=head_first
)
return A
241 changes: 17 additions & 224 deletions fla/ops/delta_rule/wy_fast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,192 +7,13 @@
import triton
import triton.language as tl

from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
from fla.ops.utils.solve_tril import solve_tril
from fla.utils import check_shared_mem, is_nvidia_hopper

NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8]


@triton.heuristics({
'USE_OFFSETS': lambda args: args['offsets'] is not None
})
@triton.autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in [2, 4, 8]
for num_stages in [2, 3, 4]
],
key=['H', 'K', 'BT', 'BK', 'BC', 'HEAD_FIRST', 'USE_OFFSETS'],
)
@triton.jit(do_not_specialize=['T'])
def fwd_prepare_wy_repr_kernel_chunk32(
k,
beta,
A,
offsets,
indices,
T,
H: tl.constexpr,
K: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BC: tl.constexpr,
HEAD_FIRST: tl.constexpr,
USE_OFFSETS: tl.constexpr,
):
i_t, i_bh = tl.program_id(0), tl.program_id(1)
i_b, i_h = i_bh // H, i_bh % H
if USE_OFFSETS:
i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T

if HEAD_FIRST:
p_beta = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BT,), (0,))
else:
p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
b_beta = tl.load(p_beta, boundary_check=(0,))

b_A = tl.zeros([BT, BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
if HEAD_FIRST:
p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
else:
p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
b_k = tl.load(p_k, boundary_check=(0, 1))
b_kb = (b_k * b_beta[:, None]).to(b_k.dtype)
b_A += tl.dot(b_kb, tl.trans(b_k), allow_tf32=False)

b_A = -tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A, 0)
for i in range(1, BT):
mask = tl.arange(0, BT) == i
b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0)
b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) * (tl.arange(0, BT) < i)
b_A = tl.where(mask[:, None], b_a, b_A)
b_A += tl.arange(0, BT)[:, None] == tl.arange(0, BT)[None, :]

if HEAD_FIRST:
p_A = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
else:
p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
tl.store(p_A, (b_A).to(p_A.dtype.element_ty), boundary_check=(0, 1))
b_A = b_A.to(k.dtype.element_ty)


@triton.heuristics({
'USE_OFFSETS': lambda args: args['offsets'] is not None
})
@triton.autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in NUM_WARPS
for num_stages in [2, 3, 4]
],
key=['H', 'K', 'BT', 'BK', 'BC', 'HEAD_FIRST', 'USE_OFFSETS'],
)
@triton.jit(do_not_specialize=['T'])
def fwd_prepare_wy_repr_kernel_chunk64(
k,
beta,
A,
At,
offsets,
indices,
T,
H: tl.constexpr,
K: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BC: tl.constexpr,
HEAD_FIRST: tl.constexpr,
USE_OFFSETS: tl.constexpr
):
i_t, i_bh = tl.program_id(0), tl.program_id(1)
i_b, i_h = i_bh // H, i_bh % H
if USE_OFFSETS:
i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_c = tl.arange(0, BC)

if HEAD_FIRST:
p_beta1 = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT,), (BC,), (0,))
p_beta2 = tl.make_block_ptr(beta + i_bh * T, (T,), (1,), (i_t * BT + BC,), (BC,), (0,))
p_A1 = tl.make_block_ptr(At + i_bh * T*BC, (T, BC), (BC, 1), (i_t * BT, 0), (BC, BC), (1, 0))
p_A2 = tl.make_block_ptr(At + i_bh * T*BC, (T, BC), (BC, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0))
else:
p_beta1 = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BC,), (0,))
p_beta2 = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT + BC,), (BC,), (0,))
p_A1 = tl.make_block_ptr(At + (bos*H + i_h) * BC, (T, BC), (H*BC, 1), (i_t * BT, 0), (BC, BC), (1, 0))
p_A2 = tl.make_block_ptr(At + (bos*H + i_h) * BC, (T, BC), (H*BC, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0))
b_beta1 = tl.load(p_beta1, boundary_check=(0,))
b_beta2 = tl.load(p_beta2, boundary_check=(0,))

b_A1 = tl.zeros([BC, BC], dtype=tl.float32)
b_A2 = tl.zeros([BC, BC], dtype=tl.float32)
b_A3 = tl.zeros([BC, BC], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
if HEAD_FIRST:
p_k1 = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_k2 = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT + BC, i_k * BK), (BC, BK), (1, 0))
else:
p_k1 = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0))
p_k2 = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + BC, i_k * BK), (BC, BK), (1, 0))
b_k1 = tl.load(p_k1, boundary_check=(0, 1))
b_k2 = tl.load(p_k2, boundary_check=(0, 1))
b_kb1 = (b_k1 * b_beta1[:, None]).to(b_k1.dtype)
b_kb2 = (b_k2 * b_beta2[:, None]).to(b_k2.dtype)
b_A1 += tl.dot(b_kb1, tl.trans(b_k1), allow_tf32=False)
b_A2 += tl.dot(b_kb2, tl.trans(b_k2), allow_tf32=False)
b_A3 += tl.dot(b_kb2, tl.trans(b_k1), allow_tf32=False)

b_A1 = -tl.where(o_c[:, None] > o_c[None, :], b_A1, 0)
b_A2 = -tl.where(o_c[:, None] > o_c[None, :], b_A2, 0)
tl.store(p_A1, b_A1.to(p_A1.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_A2, b_A2.to(p_A2.dtype.element_ty), boundary_check=(0, 1))
tl.debug_barrier()

for i in range(1, BC):
mask = o_c == i
if HEAD_FIRST:
p_a1 = At + (i_bh * T + i_t * BT + i) * BC + o_c
p_a2 = At + (i_bh * T + i_t * BT + BC + i) * BC + o_c
else:
p_a1 = At + ((bos + i_t * BT + i)*H + i_h) * BC + o_c
p_a2 = At + ((bos + i_t * BT + BC + i)*H + i_h) * BC + o_c
b_a1 = tl.load(p_a1, mask=(i_t * BT + i < T), other=0)
b_a2 = tl.load(p_a2, mask=(i_t * BT + BC + i < T), other=0)
b_a1 = b_a1 + tl.sum(b_a1[:, None] * b_A1, 0) * (o_c < i)
b_a2 = b_a2 + tl.sum(b_a2[:, None] * b_A2, 0) * (o_c < i)
b_A1 = tl.where(mask[:, None], b_a1, b_A1)
b_A2 = tl.where(mask[:, None], b_a2, b_A2)

# blockwise computation of lower triangular matrix's inverse
# i.e., [A11, 0; A21, A22]^-1 = [A11^-1, 0; -A22^-1 A21 A11^-1, A22^-1]
b_A1 += o_c[:, None] == o_c[None, :]
b_A2 += o_c[:, None] == o_c[None, :]
b_A3 = -tl.dot(tl.dot(b_A2, b_A3, allow_tf32=False), b_A1, allow_tf32=False)

if HEAD_FIRST:
p_A1 = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT, 0), (BC, BC), (1, 0))
p_A2 = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0))
p_A3 = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0))
p_A4 = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT, BC), (BC, BC), (1, 0))
else:
p_A1 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0))
p_A2 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0))
p_A3 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0))
p_A4 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, BC), (BC, BC), (1, 0))
tl.store(p_A1, b_A1.to(p_A1.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_A2, b_A2.to(p_A2.dtype.element_ty), boundary_check=(0, 1))
tl.store(p_A3, b_A3.to(p_A3.dtype.element_ty), boundary_check=(0, 1))
# causal mask
tl.store(p_A4, tl.zeros([BC, BC], dtype=tl.float32).to(p_A4.dtype.element_ty), boundary_check=(0, 1))


@triton.heuristics({
'USE_OFFSETS': lambda args: args['offsets'] is not None
})
Expand Down Expand Up @@ -396,51 +217,23 @@ def fwd_prepare_wy_repr(
beta: torch.Tensor,
offsets: Optional[torch.LongTensor],
indices: Optional[torch.LongTensor],
head_first: bool = True,
head_first: bool = False,
chunk_size: int = 64
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
if head_first:
B, H, T, K = k.shape
else:
B, T, H, K = k.shape
BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
BC = min(BT, 32)
BK = min(triton.next_power_of_2(K), 64)
NT = triton.cdiv(T, BT) if offsets is None else len(indices)

A = torch.empty(B, *((H, T) if head_first else (T, H)), BT, device=k.device, dtype=k.dtype)
if BT == 64:
At = torch.empty(B, *((H, T) if head_first else (T, H)), BC, device=k.device, dtype=torch.float)
fwd_prepare_wy_repr_kernel_chunk64[(NT, B * H)](
k=k,
beta=beta,
A=A,
At=At,
offsets=offsets,
indices=indices,
T=T,
H=H,
K=K,
BT=BT,
BK=BK,
BC=BC,
HEAD_FIRST=head_first
)
else:
fwd_prepare_wy_repr_kernel_chunk32[(NT, B * H)](
k=k,
beta=beta,
A=A,
offsets=offsets,
indices=indices,
T=T,
H=H,
K=K,
BT=BT,
BK=BK,
BC=BC,
HEAD_FIRST=head_first
)
A = chunk_scaled_dot_kkt_fwd(
k=k,
beta=beta,
cu_seqlens=offsets,
head_first=head_first,
chunk_size=chunk_size,
output_dtype=torch.float32
)
A = solve_tril(
A=A,
cu_seqlens=offsets,
head_first=head_first,
output_dtype=k.dtype
)

w, u = fwd_recompute_w_u(
k=k,
Expand Down
Loading