Skip to content

Commit 7484d47

Browse files
fix(cute): pass mDynamicCausal to kernel and init is_split_kv in SM80/SM120 FA4 forward (#178)
- Pass mDynamicCausal in the FlashAttentionForwardSm80/Sm120 kernel launch and add it to the kernel signature (was referenced in the kernel body but never passed, causing NameError on first call). - Initialize self.is_split_kv = False in __init__ (read in four places but never assigned, causing AttributeError on first call). Fixes vllm-project/vllm#51776 Signed-off-by: SayHelloToWorld <SayHelloToWorld@users.noreply.github.com> Co-authored-by: SayHelloToWorld <SayHelloToWorld@users.noreply.github.com>
1 parent 617264c commit 7484d47

1 file changed

Lines changed: 3 additions & 0 deletions

File tree

flash_attn/cute/flash_fwd.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -626,6 +626,7 @@ def __init__(self, *args, **kwargs):
626626
assert self.output_quant_key is None, (
627627
f"Fused quant output not implemented for {type(self).__name__}"
628628
)
629+
self.is_split_kv = False
629630

630631
def _get_smem_layout_atom(self):
631632
sQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.tile_hdim)
@@ -793,6 +794,7 @@ def __call__(
793794
aux_data,
794795
fastdiv_mods,
795796
output_scale,
797+
mDynamicCausal,
796798
).launch(
797799
grid=grid_dim,
798800
block=[self.num_threads, 1, 1],
@@ -833,6 +835,7 @@ def kernel(
833835
aux_data: AuxData = AuxData(),
834836
fastdiv_mods=None,
835837
output_scale: Optional[cute.Tensor] = None,
838+
mDynamicCausal: Optional[cute.Tensor] = None,
836839
):
837840
# Thread index, block index
838841
tidx, _, _ = cute.arch.thread_idx()

0 commit comments

Comments
 (0)