Skip to content

[Attn] Remove head_first & rename offsets to cu_seqlens #345

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 7 commits into from
Apr 10, 2025
Merged

[Attn] Remove head_first & rename offsets to cu_seqlens #345

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
9a46410
[Attn] Remove `head_first` & rename `offsets` to `cu_seqlens`
yzhangcs Apr 10, 2025
ee5e35e
Delete 256 headdim tests
yzhangcs Apr 10, 2025
703a7a9
[DeltaNet] Remove `head_first` & rename `offsets` to `cu_seqlens`
yzhangcs Apr 10, 2025
7b46c6c
Rename `offsets` to `cu_seqlens` across GLA/GSA/RWKV6
yzhangcs Apr 10, 2025
1008294
Fix NSA tests
yzhangcs Apr 10, 2025
c1e7218
Fix TTT checks
yzhangcs Apr 10, 2025
6d4dd41
Rename `offsets` to `cu_seqlens` in DPLR and IPLR implementations, up…
yzhangcs Apr 10, 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
101 changes: 57 additions & 44 deletions fla/ops/attn/parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang

import warnings
from typing import Optional

import torch
Expand All @@ -15,16 +16,16 @@


@triton.heuristics({
'IS_VARLEN': lambda args: args['offsets'] is not None,
'USE_G': lambda args: args['g_cumsum'] is not None
'USE_G': lambda args: args['g_cumsum'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@triton.autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else [])
for num_stages in [2, 3, 4, 5]
],
key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'],
key=['B', 'H', 'HQ', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'IS_VARLEN'],
)
@triton.jit
def parallel_attn_fwd_kernel(
Expand All @@ -35,8 +36,8 @@ def parallel_attn_fwd_kernel(
g_cumsum,
lse,
scale,
offsets,
indices,
cu_seqlens,
chunk_indices,
T,
B: tl.constexpr,
H: tl.constexpr,
Expand All @@ -48,16 +49,16 @@ def parallel_attn_fwd_kernel(
BS: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
IS_VARLEN: tl.constexpr,
USE_G: tl.constexpr
):
i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
i_b, i_hq = i_bh // HQ, i_bh % HQ
i_h = i_hq // G

if IS_VARLEN:
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)
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
i_n = i_b
Expand Down Expand Up @@ -168,16 +169,16 @@ def parallel_attn_bwd_kernel_preprocess(


@triton.heuristics({
'IS_VARLEN': lambda args: args['offsets'] is not None,
'USE_G': lambda args: args['g_cumsum'] is not None
'USE_G': lambda args: args['g_cumsum'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@triton.autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else [])
for num_stages in [2, 3, 4, 5]
],
key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'],
key=['B', 'H', 'HQ', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'IS_VARLEN'],
)
@triton.jit(do_not_specialize=['T'])
def parallel_attn_bwd_kernel_dq(
Expand All @@ -191,8 +192,8 @@ def parallel_attn_bwd_kernel_dq(
dg_cumsum,
g_cumsum,
scale,
offsets,
indices,
cu_seqlens,
chunk_indices,
T,
B: tl.constexpr,
H: tl.constexpr,
Expand All @@ -212,8 +213,8 @@ def parallel_attn_bwd_kernel_dq(
i_h = i_hq // G

if IS_VARLEN:
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)
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
i_n = i_b
Expand Down Expand Up @@ -306,16 +307,16 @@ def parallel_attn_bwd_kernel_dq(


@triton.heuristics({
'IS_VARLEN': lambda args: args['offsets'] is not None,
'USE_G': lambda args: args['g_cumsum'] is not None
'USE_G': lambda args: args['g_cumsum'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@triton.autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else [])
for num_stages in [2, 3, 4, 5]
],
key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G'],
key=['B', 'H', 'HQ', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'IS_VARLEN'],
)
@triton.jit(do_not_specialize=['T'])
def parallel_attn_bwd_kernel_dkv(
Expand All @@ -329,8 +330,8 @@ def parallel_attn_bwd_kernel_dkv(
dk,
dv,
dg_cumsum,
offsets,
indices,
cu_seqlens,
chunk_indices,
scale,
T,
B: tl.constexpr,
Expand All @@ -351,8 +352,8 @@ def parallel_attn_bwd_kernel_dkv(
i_h = i_hq // G

if IS_VARLEN:
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)
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
T = eos - bos
else:
i_n = i_b
Expand Down Expand Up @@ -464,7 +465,7 @@ def parallel_attn_fwd(
g_cumsum: torch.Tensor,
scale: float,
chunk_size: int = 128,
offsets: Optional[torch.LongTensor] = None,
cu_seqlens: Optional[torch.LongTensor] = None,
):
B, T, H, K, V = *k.shape, v.shape[-1]
HQ = q.shape[2]
Expand All @@ -485,8 +486,8 @@ def parallel_attn_fwd(
NK = triton.cdiv(K, BK)
NV = triton.cdiv(V, BV)

indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None
NT = triton.cdiv(T, BT) if offsets is None else len(indices)
chunk_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(chunk_indices)
assert NK == 1, "The key dimension can not be larger than 256"

o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device)
Expand All @@ -500,8 +501,8 @@ def parallel_attn_fwd(
g_cumsum=g_cumsum,
lse=lse,
scale=scale,
offsets=offsets,
indices=indices,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
B=B,
T=T,
H=H,
Expand All @@ -522,7 +523,7 @@ def parallel_attn_bwd_preprocess(
do: torch.Tensor
):
V = o.shape[-1]
delta = torch.empty_like(o[..., 0], dtype=torch.float32)
delta = torch.empty_like(o[..., 0], dtype=torch.float)
parallel_attn_bwd_kernel_preprocess[(delta.numel(),)](
o=o,
do=do,
Expand All @@ -543,7 +544,7 @@ def parallel_attn_bwd(
do: torch.Tensor,
scale: float = None,
chunk_size: int = 128,
offsets: Optional[torch.LongTensor] = None,
cu_seqlens: Optional[torch.LongTensor] = None,
):
B, T, H, K, V = *k.shape, v.shape[-1]
HQ = q.shape[2]
Expand All @@ -554,8 +555,8 @@ def parallel_attn_bwd(
BK = max(16, triton.next_power_of_2(K))
BV = max(16, triton.next_power_of_2(V))

indices = prepare_chunk_indices(offsets, BT) if offsets is not None else None
NT = triton.cdiv(T, BT) if offsets is None else len(indices)
chunk_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(chunk_indices)
NV = triton.cdiv(V, BV)

delta = parallel_attn_bwd_preprocess(o, do)
Expand All @@ -567,8 +568,8 @@ def parallel_attn_bwd(

dg_cumsum, dg_cumsum_k = None, None
if g_cumsum is not None:
dg_cumsum = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device)
dg_cumsum_k = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device)
dg_cumsum = torch.empty(B, T, HQ, dtype=torch.float, device=q.device)
dg_cumsum_k = torch.empty(B, T, HQ, dtype=torch.float, device=q.device)

parallel_attn_bwd_kernel_dq[grid](
q=q,
Expand All @@ -580,8 +581,8 @@ def parallel_attn_bwd(
do=do,
dq=dq,
dg_cumsum=dg_cumsum,
offsets=offsets,
indices=indices,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
scale=scale,
T=T,
B=B,
Expand All @@ -606,8 +607,8 @@ def parallel_attn_bwd(
dk=dk,
dv=dv,
dg_cumsum=dg_cumsum_k,
offsets=offsets,
indices=indices,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
scale=scale,
T=T,
B=B,
Expand Down Expand Up @@ -647,7 +648,7 @@ def forward(ctx, q, k, v, g, scale, cu_seqlens):
g_cumsum=g_cumsum,
scale=scale,
chunk_size=chunk_size,
offsets=cu_seqlens,
cu_seqlens=cu_seqlens,
)
ctx.save_for_backward(q, k, v, o, g_cumsum, lse)
ctx.chunk_size = chunk_size
Expand All @@ -670,12 +671,12 @@ def backward(ctx, do):
do=do,
scale=ctx.scale,
chunk_size=ctx.chunk_size,
offsets=ctx.cu_seqlens,
cu_seqlens=ctx.cu_seqlens,
)
if dg is not None:
dg = chunk_global_cumsum(dg, cu_seqlens=ctx.cu_seqlens, reverse=True)

return dq.to(q), dk.to(k), dv.to(v), dg, None, None, None, None, None, None, None, None
return dq.to(q), dk.to(k), dv.to(v), dg, None, None


def parallel_attn(
Expand Down Expand Up @@ -711,14 +712,26 @@ def parallel_attn(
o (torch.Tensor):
Outputs of shape `[B, T, HQ, V]` if `head_first=False` else `[B, HQ, T, V]`.
"""
if scale is None:
scale = k.shape[-1] ** -0.5
if cu_seqlens is not None:
assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided"
if head_first:
warnings.warn(
"head_first is deprecated and will be removed in a future version. "
"Please use head_first=False for now instead."
)
q, k, v = map(lambda x: rearrange(x, 'b h t ... -> b t h ...'), (q, k, v))
if g is not None:
g = rearrange(g, 'b h t ... -> b t h ...')
if not head_first and q.shape[1] < q.shape[2]:
warnings.warn(
f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
"This may indicate the inputs were passed in head-first format [B, H, T, ...] "
"when head_first=False was specified. "
"Please verify your input tensor format matches the expected shape [B, T, H, ...]."
)
if scale is None:
scale = k.shape[-1] ** -0.5
if cu_seqlens is not None:
assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided"

o = ParallelAttentionFunction.apply(q, k, v, g, scale, cu_seqlens)
if head_first:
o = rearrange(o, 'b t h ... -> b h t ...')
Expand Down
46 changes: 23 additions & 23 deletions tests/ops/test_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,13 +31,13 @@
test_h_list = [2]


@pytest.mark.parametrize("B", test_b_list)
@pytest.mark.parametrize("T", test_t_list)
@pytest.mark.parametrize("H", test_h_list)
@pytest.mark.parametrize("HQ", test_hq_list)
@pytest.mark.parametrize("D", test_d_list)
@pytest.mark.parametrize("scale", [0.1])
@pytest.mark.parametrize("dtype", [torch.float16])
@pytest.mark.parametrize('B', test_b_list)
@pytest.mark.parametrize('T', test_t_list)
@pytest.mark.parametrize('H', test_h_list)
@pytest.mark.parametrize('HQ', test_hq_list)
@pytest.mark.parametrize('D', test_d_list)
@pytest.mark.parametrize('scale', [0.1])
@pytest.mark.parametrize('dtype', [torch.float16])
@pytest.mark.skipif(
not HAS_FLASH,
reason="Skipping test because flash-attn is not installed"
Expand Down Expand Up @@ -79,12 +79,12 @@ def test_parallel(
assert_close("dv", ref_dv, tri_dv, 0.005)


@pytest.mark.parametrize("N", test_b_list)
@pytest.mark.parametrize("T", test_t_varlen_list)
@pytest.mark.parametrize("H", test_h_list)
@pytest.mark.parametrize("HQ", test_hq_list)
@pytest.mark.parametrize("D", test_d_list)
@pytest.mark.parametrize("dtype", [torch.float16])
@pytest.mark.parametrize('N', test_b_list)
@pytest.mark.parametrize('T', test_t_varlen_list)
@pytest.mark.parametrize('H', test_h_list)
@pytest.mark.parametrize('HQ', test_hq_list)
@pytest.mark.parametrize('D', test_d_list)
@pytest.mark.parametrize('dtype', [torch.float16])
@pytest.mark.skipif(
not HAS_FLASH,
reason="Skipping test because flash-attn is not installed"
Expand All @@ -104,7 +104,7 @@ def test_parallel_varlen(

N = min(1, N) if T < 64 else N
# randomly split the sequence into N segments
offsets = torch.cat([
cu_seqlens = torch.cat([
torch.tensor([0], dtype=torch.long),
torch.arange(16, T)[torch.randperm(T - 16)[:N-1]],
torch.tensor([T], dtype=torch.long)
Expand All @@ -119,10 +119,10 @@ def test_parallel_varlen(
q=q.squeeze(0),
k=k.squeeze(0),
v=v.squeeze(0),
cu_seqlens_q=offsets,
cu_seqlens_k=offsets,
max_seqlen_q=prepare_lens(offsets).max(),
max_seqlen_k=prepare_lens(offsets).max(),
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=prepare_lens(cu_seqlens).max(),
max_seqlen_k=prepare_lens(cu_seqlens).max(),
causal=True
)
ref.backward(do.squeeze(0))
Expand All @@ -134,14 +134,14 @@ def test_parallel_varlen(
q=q,
k=k,
v=v,
cu_seqlens=offsets
cu_seqlens=cu_seqlens
)
tri.backward(do)
tri_dq, q.grad = q.grad.clone(), None
tri_dk, k.grad = k.grad.clone(), None
tri_dv, v.grad = v.grad.clone(), None

assert_close(" o", ref, tri, 0.004)
assert_close(" dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005)
assert_close(" dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005)
assert_close(" dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005)
assert_close(" o", ref, tri, 0.004)
assert_close("dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005)
assert_close("dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005)
assert_close("dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005)