Skip to content

Commit e6e3993

Browse files
fix linter error from PR 2557
1 parent 0727938 commit e6e3993

1 file changed

Lines changed: 2 additions & 4 deletions

File tree

flash_attn/cute/flash_bwd_sm100.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2756,14 +2756,12 @@ def apply_score_mod(
27562756
):
27572757
"""Apply forward score modification for SM100 backward pass."""
27582758
# In bwd, S is computed as K @ Q.T so dimensions are (tile_n, tile_m).
2759-
# With 2CTA, partition_C must see the full cluster tile so each CTA
2759+
# With 2CTA, partition_C must see the full cluster tile so each CTA
27602760
# gets its own half of the tile.
27612761
cluster_tile_n = self.tile_n * self.cta_group_size
27622762
cluster_n_block = n_block // self.cta_group_size
27632763
cS = cute.make_identity_tensor((cluster_tile_n, self.tile_m))
2764-
cS = cute.domain_offset(
2765-
(cluster_n_block * cluster_tile_n, m_block * self.tile_m), cS
2766-
)
2764+
cS = cute.domain_offset((cluster_n_block * cluster_tile_n, m_block * self.tile_m), cS)
27672765
tScS = thr_mma_S.partition_C(cS)
27682766
tScS_idx = thr_copy_t2r.partition_D(tScS)
27692767

0 commit comments

Comments
 (0)