[CuTe] Fix lint failures (#2625)

stack-info: PR: #2625, branch: drisspg/stack/42

Author: drisspg

flash_attn/cute/flash_bwd.py

@@ -11,7 +11,7 @@
 import cutlass
 import cutlass.cute as cute
 from cutlass.cute.nvgpu import cpasync, warp
-from cutlass import Float32, Int32
+from cutlass import Int32
 import cutlass.utils as utils_basic
 
 from quack import layout_utils
@@ -167,13 +167,13 @@ def _check_type(
         else:
             if cutlass.const_expr(not (mdK_type == mdV_type == cutlass.Float32)):
                 raise TypeError("mdKaccum and mdVaccum tensors must have the data type Float32")
-        if cutlass.const_expr(not mQ_type in [cutlass.Float16, cutlass.BFloat16]):
+        if cutlass.const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]):
             raise TypeError("Only Float16 or BFloat16 is supported")
-        if cutlass.const_expr(not mLSE_type in [cutlass.Float32]):
+        if cutlass.const_expr(mLSE_type not in [cutlass.Float32]):
             raise TypeError("LSE tensor must be Float32")
-        if cutlass.const_expr(not mdPsum_type in [cutlass.Float32]):
+        if cutlass.const_expr(mdPsum_type not in [cutlass.Float32]):
             raise TypeError("dPsum tensor must be Float32")
-        if cutlass.const_expr(not mdQaccum_type in [cutlass.Float32]):
+        if cutlass.const_expr(mdQaccum_type not in [cutlass.Float32]):
             raise TypeError("dQaccum tensor must be Float32")
         if cutlass.const_expr(mCuSeqlensQ_type not in [None, cutlass.Int32]):
             raise TypeError("cuSeqlensQ tensor must be Int32")

flash_attn/cute/flash_fwd.py

@@ -7,14 +7,14 @@
 
 import math
 from types import SimpleNamespace
-from typing import Type, Callable, Optional, List
+from typing import Type, Callable, Optional
 from functools import partial
 
 import cuda.bindings.driver as cuda
 
 import cutlass
 import cutlass.cute as cute
-from cutlass import Constexpr, Float32, Int32, const_expr, Boolean
+from cutlass import Float32, Int32, const_expr
 from cutlass.cute.nvgpu import cpasync, warp
 import cutlass.utils as utils_basic
 from cutlass.base_dsl.arch import Arch

flash_attn/cute/interface.py

@@ -10,7 +10,6 @@
 import torch
 
 
-import cuda.bindings.driver as cuda
 
 import cutlass
 import cutlass.cute as cute
@@ -55,7 +54,6 @@
     to_cute_block_sparse_tensors,
     normalize_block_sparse_config,
     normalize_block_sparse_config_bwd,
-    get_block_sparse_broadcast_pattern,
 )
 
 def _parse_arch_str(arch_str):