Clean up imports in flash.py, flash_sharding.py, and ring_attention.py.

Author: nshepperd

scripts/install-cuda-linux.sh

@@ -1,4 +1,5 @@
 #!/bin/bash
+# Install CUDA on manylinux docker image.
 set -eux
 
 VER=${1:-12.8}

src/flash_attn_jax/flash.py

@@ -8,16 +8,8 @@
 from jax.interpreters import batching
 from jax.interpreters import mlir
 from jax.interpreters import xla
-from jax.interpreters.mlir import ir
-from jax.lib import xla_client
-from jaxlib.hlo_helpers import custom_call
-from jax.experimental.custom_partitioning import custom_partitioning
 from jax.extend.core import Primitive
 
-from jax.sharding import PartitionSpec as P
-from jax.sharding import Mesh
-from jax.sharding import NamedSharding
-
 from einops import rearrange
 import einops
 import math
@@ -206,11 +198,14 @@ def custom_vjp(cls, nondiff_argnums=()):
 # bwd.
 @partial(custom_vjp, nondiff_argnums=(3,))
 class _flash_mha_vjp:
+    @staticmethod
     def base(q,k,v,config):
         return _flash_mha_fwd(q,k,v, **config)[0]
+    @staticmethod
     def fwd(q,k,v,config):
         out, lse = _flash_mha_fwd(q,k,v, **config)
         return out, (q,k,v,out,lse)
+    @staticmethod
     def bwd(config, pack, dout):
         (q,k,v,out,lse) = pack
         dq, dk, dv = _flash_mha_bwd(dout, q, k, v, out, lse, **config)

src/flash_attn_jax/flash_sharding.py

@@ -5,14 +5,7 @@
 import numpy as np
 import jax
 import jax.numpy as jnp
-from jax import core, dtypes
-from jax.core import ShapedArray
-from jax.interpreters import batching
-from jax.interpreters import mlir
-from jax.interpreters import xla
 from jax.interpreters.mlir import ir
-from jax.lib import xla_client
-from jaxlib.hlo_helpers import custom_call
 from jax.experimental.custom_partitioning import custom_partitioning, SdyShardingRule, ArrayMapping, CompoundFactor
 
 from jax.sharding import PartitionSpec as P
@@ -35,9 +28,9 @@ def is_replicated(sharding):
         return sharding.is_fully_replicated
     raise ValueError(f"Unsupported sharding type: {type(sharding)}")
 
-def partition_fwd(softmax_scale, is_causal, window_size, 
-                  mesh: Mesh, 
-                  arg_shapes: List[jax.ShapeDtypeStruct], 
+def partition_fwd(softmax_scale, is_causal, window_size,
+                  mesh: Mesh,
+                  arg_shapes: List[jax.ShapeDtypeStruct],
                   result_shape: List[jax.ShapeDtypeStruct]):
     result_shardings = [x.sharding for x in result_shape],
     arg_shardings = [x.sharding for x in arg_shapes]

src/flash_attn_jax/ring_attention.py

@@ -3,15 +3,6 @@
 import numpy as np
 import jax
 import jax.numpy as jnp
-from jax import core, dtypes
-from jax.core import ShapedArray
-from jax.interpreters import batching
-from jax.interpreters import mlir
-from jax.interpreters import xla
-from jax.interpreters.mlir import ir
-from jax.lib import xla_client
-from jaxlib.hlo_helpers import custom_call
-from jax.experimental.custom_partitioning import custom_partitioning
 
 from einops import rearrange
 import math