minor clean up

Author: reubenconducts

flash_attn/cute/flash_fwd_combine.py

@@ -365,7 +365,6 @@ def kernel(
 
         # Handle semaphore reset — wait for dependent grids first
         if const_expr(semaphore_to_reset is not None):
-            # maybe handle on first CTA?
             if (
                 tidx == 0
                 and m_block == cute.arch.grid_dim()[0] - 1

tests/cute/test_flash_attn.py

@@ -34,7 +34,6 @@
     _flash_attn_fwd,
     _flash_attn_bwd,
 )
-from flash_attn.cute.prepare_scheduler import SchedulerMetadataTensorsTorch
 
 def retry_on_oom(func):
     @wraps(func)
@@ -1001,7 +1000,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device):
 @pytest.mark.parametrize("has_leftpad", [False])
 # @pytest.mark.parametrize("has_batch_idx", [False, True])
 @pytest.mark.parametrize("has_batch_idx", [False])
-@pytest.mark.parametrize("varlen_q", [True])
+@pytest.mark.parametrize("varlen_q", [False, True])
 # @pytest.mark.parametrize("varlen_q", [False])
 # @pytest.mark.parametrize("d", [32, 59, 64, 80, 128, 256])
 # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256])