Skip to content

SM100 dynamic causal #144

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

Open
MatthewBonanni wants to merge 1 commit into
base: main
Choose a base branch
from
+132 −5
Open

SM100 dynamic causal #144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
129 changes: 124 additions & 5 deletions flash_attn/cute/flash_fwd_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -379,6 +379,10 @@ def __call__(
mCuSeqlensK: Optional[cute.Tensor] = None,
mSeqUsedQ: Optional[cute.Tensor] = None,
mSeqUsedK: Optional[cute.Tensor] = None,
# Per-request causal flag (int32, one entry per batch: 1=causal,
# 0=bidirectional). Same positional slot as the SM90 kernel so the
# shared interface positional arg list aligns across arches.
mDynamicCausal: Optional[cute.Tensor] = None,
mPageTable: Optional[cute.Tensor] = None, # (b_k, max_num_pages_per_seq)
window_size_left: Int32 | int | None = None,
window_size_right: Int32 | int | None = None,
Expand Down Expand Up @@ -751,6 +755,7 @@ class SharedStorage:
mCuSeqlensK,
mSeqUsedQ,
mSeqUsedK,
mDynamicCausal,
mPageTable,
tma_atom_Q,
tma_atom_K,
Expand Down Expand Up @@ -811,6 +816,7 @@ def kernel(
mCuSeqlensK: Optional[cute.Tensor],
mSeqUsedQ: Optional[cute.Tensor],
mSeqUsedK: Optional[cute.Tensor],
mDynamicCausal: Optional[cute.Tensor],
mPageTable: Optional[cute.Tensor],
tma_atom_Q: Optional[cute.CopyAtom],
tma_atom_K: Optional[cute.CopyAtom],
Expand Down Expand Up @@ -1090,6 +1096,9 @@ def kernel(
AttentionMaskCls = self._generate_attention_mask_cls(
window_size_left, window_size_right
)
# Stash the per-request causal flag so every warp method (load, mma,
# softmax, correction, epilogue) can read it via self.
self._mDynamicCausal = mDynamicCausal
# Cluster wait before tensor memory alloc
pipeline_init_wait(cluster_shape_mn=cta_layout_vmnk)

Expand Down Expand Up @@ -1475,9 +1484,35 @@ def load(
)

if const_expr(not self.use_block_sparsity):
psc = (
self._mDynamicCausal[batch_idx]
if const_expr(self._mDynamicCausal is not None)
else None
)
n_block_min, n_block_max = block_info.get_n_block_min_max(
seqlen, m_block, split_idx, num_splits
)
if const_expr(self._mDynamicCausal is not None):
# Per-sequence causal: psc == 0 means this sequence is
# processed bidirectionally. get_n_block_min_max may have
# applied a causal upper bound (the kernel is compiled
# causal) and, for split-KV, partitioned that range. For a
# bidirectional sequence each split must instead own a
# DISJOINT slice of the FULL key range. Recompute it here --
# IDENTICALLY in every warp (load/mma/softmax/correction/
# epilogue) -- so the pipeline block counts agree; a
# mismatch deadlocks the kernel.
if not psc:
n_block_max_full = cute.ceil_div(seqlen.seqlen_k, self.n_block_size)
if const_expr(self.is_split_kv):
num_n_blocks_per_split = cute.ceil_div(n_block_max_full, num_splits)
n_block_min = split_idx * num_n_blocks_per_split
n_block_max = cutlass.min(
n_block_min + num_n_blocks_per_split, n_block_max_full
)
else:
n_block_min = Int32(0)
n_block_max = n_block_max_full
if const_expr(not self.is_split_kv) or n_block_min < n_block_max:
n_block_first = n_block_max - 1 if n_block_max > 0 else 0
page_idx = (
Expand Down Expand Up @@ -1680,7 +1715,28 @@ def mma(
)
process_tile = block_iter_count > Int32(0)
else:
n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
psc = (
self._mDynamicCausal[batch_idx]
if const_expr(self._mDynamicCausal is not None)
else None
)
n_block_min, n_block_max = block_info.get_n_block_min_max(
seqlen, m_block, split_idx, num_splits
)
if const_expr(self._mDynamicCausal is not None):
# Bidirectional (psc == 0) full-range widening; must match
# the load warp exactly (see there) or the pipeline hangs.
if not psc:
n_block_max_full = cute.ceil_div(seqlen.seqlen_k, self.n_block_size)
if const_expr(self.is_split_kv):
num_n_blocks_per_split = cute.ceil_div(n_block_max_full, num_splits)
n_block_min = split_idx * num_n_blocks_per_split
n_block_max = cutlass.min(
n_block_min + num_n_blocks_per_split, n_block_max_full
)
else:
n_block_min = Int32(0)
n_block_max = n_block_max_full
block_iter_count = n_block_max - n_block_min
if const_expr(not self.is_split_kv):
process_tile = True
Expand Down Expand Up @@ -1959,9 +2015,30 @@ def softmax_loop(
m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx
kv_head_idx = self._kv_head_idx(head_idx)
seqlen = SeqlenInfoCls(batch_idx)
n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
psc = (
self._mDynamicCausal[batch_idx]
if const_expr(self._mDynamicCausal is not None)
else None
)
n_block_min, n_block_max = block_info.get_n_block_min_max(
seqlen, m_block, split_idx, num_splits
)
if const_expr(self._mDynamicCausal is not None):
# Bidirectional (psc == 0) full-range widening; must match
# the load warp exactly (see there) or the pipeline hangs.
if not psc:
n_block_max_full = cute.ceil_div(seqlen.seqlen_k, self.n_block_size)
if const_expr(self.is_split_kv):
num_n_blocks_per_split = cute.ceil_div(n_block_max_full, num_splits)
n_block_min = split_idx * num_n_blocks_per_split
n_block_max = cutlass.min(
n_block_min + num_n_blocks_per_split, n_block_max_full
)
else:
n_block_min = Int32(0)
n_block_max = n_block_max_full

mask = AttentionMaskCls(seqlen)
mask = AttentionMaskCls(seqlen, dynamic_causal=psc)
shared_mask_kwargs = dict(
m_block=(self.q_stage * m_block + stage) * self.cta_group_size,
thr_mma=thr_mma_qk,
Expand Down Expand Up @@ -2441,7 +2518,28 @@ def correction_loop(
Float32(256.0) if cutlass.const_expr(self.q_dtype.width == 8) else Float32(1.0)
)
seqlen = SeqlenInfoCls(batch_idx)
n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
psc = (
self._mDynamicCausal[batch_idx]
if const_expr(self._mDynamicCausal is not None)
else None
)
n_block_min, n_block_max = block_info.get_n_block_min_max(
seqlen, m_block, split_idx, num_splits
)
if const_expr(self._mDynamicCausal is not None):
# Bidirectional (psc == 0) full-range widening; must match
# the load warp exactly (see there) or the pipeline hangs.
if not psc:
n_block_max_full = cute.ceil_div(seqlen.seqlen_k, self.n_block_size)
if const_expr(self.is_split_kv):
num_n_blocks_per_split = cute.ceil_div(n_block_max_full, num_splits)
n_block_min = split_idx * num_n_blocks_per_split
n_block_max = cutlass.min(
n_block_min + num_n_blocks_per_split, n_block_max_full
)
else:
n_block_min = Int32(0)
n_block_max = n_block_max_full

if const_expr(self.is_split_kv):
mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx, split_idx]
Expand Down Expand Up @@ -2892,7 +2990,28 @@ def epilogue_s2g(
while work_tile.is_valid_tile:
m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx
seqlen = SeqlenInfoCls(batch_idx)
n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits)
psc = (
self._mDynamicCausal[batch_idx]
if const_expr(self._mDynamicCausal is not None)
else None
)
n_block_min, n_block_max = block_info.get_n_block_min_max(
seqlen, m_block, split_idx, num_splits
)
if const_expr(self._mDynamicCausal is not None):
# Bidirectional (psc == 0) full-range widening; must match
# the load warp exactly (see there) or the pipeline hangs.
if not psc:
n_block_max_full = cute.ceil_div(seqlen.seqlen_k, self.n_block_size)
if const_expr(self.is_split_kv):
num_n_blocks_per_split = cute.ceil_div(n_block_max_full, num_splits)
n_block_min = split_idx * num_n_blocks_per_split
n_block_max = cutlass.min(
n_block_min + num_n_blocks_per_split, n_block_max_full
)
else:
n_block_min = Int32(0)
n_block_max = n_block_max_full
has_work = const_expr(self.use_block_sparsity or not self.is_split_kv) or n_block_min < n_block_max

if has_work:
Expand Down
1 change: 1 addition & 0 deletions flash_attn/cute/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -712,6 +712,7 @@ def _flash_attn_fwd(
head_dim_v,
qhead_per_kvhead,
causal,
dynamic_causal is not None,
score_mod_hash,
mask_mod_hash,
use_block_sparsity,
Expand Down
7 changes: 7 additions & 0 deletions flash_attn/cute/mask.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -687,6 +687,13 @@ def apply_mask_sm100(
row_idx = row_idx // self.qhead_per_kvhead_packgqa
if const_expr(mask_causal):
col_limit_right = row_idx + causal_row_offset + 1
# Per-request dynamic causal: for bidirectional sequences
# (psc == 0), drop the causal column clip so the row attends
# to the full (seqlen-bounded) KV. Mirrors the SM90 path.
if const_expr(self.dynamic_causal is not None):
col_limit_right = (
col_limit_right if self.dynamic_causal else seqlenk_col_limit
)
if const_expr(mask_seqlen):
col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit)
# if cute.arch.thread_idx()[0] % 32 == 0:
Expand Down