Refactor flash HLO primitives to use jax.ffi.ffi_call and remove custom call implementation

Author: nshepperd

src/flash_attn_jax/flash_hlo.py

@@ -5,12 +5,9 @@
 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 jax.experimental.custom_partitioning import custom_partitioning
 
 from jax.extend.core import Primitive
 
@@ -20,6 +17,8 @@
 
 import flash_attn_jax_lib.flash_api as flash_api
 
+# jax.ffi.ffi_call()
+
 # ==== Register primitives ====
 
 _flash_mha_fwd_hlo_p = Primitive("flash_mha_fwd_hlo")
@@ -30,10 +29,6 @@
 _flash_mha_bwd_hlo_p.multiple_results = True
 _flash_mha_bwd_hlo_p.def_impl(partial(xla.apply_primitive, _flash_mha_bwd_hlo_p))
 
-_custom_call_p = Primitive("custom_call")
-_custom_call_p.multiple_results = True
-_custom_call_p.def_impl(partial(xla.apply_primitive, _custom_call_p))
-
 # ==== Primitive wrapper ====
 
 def _flash_mha_fwd_hlo(q, k, v, softmax_scale, is_causal, window_size):
@@ -44,18 +39,10 @@ def _flash_mha_bwd_hlo(dout, q, k, v, out, lse, softmax_scale, is_causal, window
     dq, dk, dv = _flash_mha_bwd_hlo_p.bind(dout, q, k, v, out, lse, softmax_scale=softmax_scale, is_causal=is_causal, window_size=window_size)
     return dq, dk, dv
 
-def custom_call(*args, call_target_name, result_types, backend_config, operand_layouts, result_layouts):
-    return _custom_call_p.bind(*args, call_target_name=call_target_name, 
-                               result_types=tuple(result_types),
-                               backend_config=backend_config,
-                               operand_layouts=tuple(operand_layouts), 
-                               result_layouts=tuple(result_layouts))
-
 # ==== HLO lowerings ====
 
 # Register functions defined in gpu_ops as custom call target for GPUs
 for _name, _value in flash_api.get_registrations().items():
-    # xla_client.register_custom_call_target(_name, _value, platform="gpu")
     jax.ffi.register_ffi_target(_name, _value, platform="gpu", api_version=0)
 
 def default_layouts(*shapes):
@@ -109,26 +96,24 @@ def fwd(q, k, v):
             k = jnp.pad(k, ((0,0),(0,0),(0,0),(0,dpad)), 'constant')
             v = jnp.pad(v, ((0,0),(0,0),(0,0),(0,dpad)), 'constant')
 
-        q_shape = [n, l, h, d+dpad]
-        k_shape = [n, lk, hk, d+dpad]
-        v_shape = [n, lk, hk, d+dpad]
+        # q_shape = [n, l, h, d+dpad]
+        # k_shape = [n, lk, hk, d+dpad]
+        # v_shape = [n, lk, hk, d+dpad]
         o_shape = [n, l, h, d+dpad]
         lse_shape = [n, h, l]
-
 
-        lse_type = ir.RankedTensorType.get([n, h, l], mlir.dtype_to_ir_type(jnp.float32.dtype))
-        out_types = [ir.RankedTensorType.get(o_shape, element_type), lse_type]
-        operand_layouts = default_layouts(q_shape, k_shape, v_shape)
-        result_layouts = default_layouts(o_shape, lse_shape)
-
-        o, lse = custom_call(
-            q, k, v,
-            call_target_name = b"flash_mha_fwd",
-            result_types=out_types,
-            backend_config=opaque,
-            operand_layouts=operand_layouts,
-            result_layouts=result_layouts,
-        )
+        jax_dtype = jnp.bfloat16 if type(element_type) == ir.BF16Type else jnp.float16
+        out_types = [jax.ShapeDtypeStruct(o_shape, jax_dtype), jax.ShapeDtypeStruct(lse_shape, jnp.float32)]
+
+        o, lse = jax.ffi.ffi_call(
+            "flash_mha_fwd", 
+            result_shape_dtypes=out_types,
+            has_side_effect=False,
+            legacy_backend_config=opaque,
+            input_layouts=[None, None, None], # default row major
+            output_layouts=[None, None],
+            custom_call_api_version=1
+            )(q, k, v)
 
         if dpad > 0:
             o = o[:,:,:,:d]
@@ -197,19 +182,20 @@ def fwd(dout, q, k, v, out, lse):
             dout = jnp.pad(dout, ((0,0),(0,0),(0,0),(0,dpad)), 'constant')
 
         # For MQA/GQA, hq != hk, but we pass a hq sized output tensor to the kernel and sum over it afterwards to reduce the size.
-        out_types = [ir.RankedTensorType.get([n, lq, hq, d+dpad], dtype),
-                    ir.RankedTensorType.get([n, lk, hq, d+dpad], dtype),
-                    ir.RankedTensorType.get([n, lk, hq, d+dpad], dtype)]
-        out_layouts = default_layouts([n, lq, hq, d+dpad], [n, lk, hq, d+dpad], [n, lk, hq, d+dpad])
-
-        dq, dk, dv = custom_call(
-            dout, q, k, v, out, lse,
-            call_target_name=b"flash_mha_bwd",
-            operand_layouts=default_layouts(dout.shape, q.shape, k.shape, v.shape, out.shape, lse.shape),
-            backend_config=opaque,
-            result_types=out_types,
-            result_layouts=out_layouts,
-        )
+        jax_dtype = jnp.bfloat16 if type(dtype) == ir.BF16Type else jnp.float16
+        out_types = [jax.ShapeDtypeStruct((n, lq, hq, d+dpad), jax_dtype),
+                    jax.ShapeDtypeStruct((n, lk, hq, d+dpad), jax_dtype),
+                    jax.ShapeDtypeStruct((n, lk, hq, d+dpad), jax_dtype)]
+
+        dq, dk, dv = jax.ffi.ffi_call(
+            "flash_mha_bwd", 
+            result_shape_dtypes=out_types,
+            has_side_effect=False,
+            legacy_backend_config=opaque,
+            input_layouts=[None]*6, # default row major
+            output_layouts=[None]*3,
+            custom_call_api_version=1
+            )(dout, q, k, v, out, lse)
 
         if hq != hk:
             assert hq > hk and hq % hk == 0
@@ -264,32 +250,3 @@ def _flash_mha_bwd_abstract(dout, q, k, v, out, lse, softmax_scale=None, is_caus
         ShapedArray(v.shape, v_dtype),
     )
 _flash_mha_bwd_hlo_p.def_abstract_eval(_flash_mha_bwd_abstract)
-
-# ==== Custom Call ====
-
-def _custom_call_abstract_eval(*args, call_target_name, result_types, backend_config, operand_layouts, result_layouts):
-    def convert(ty):
-        ty = ir.RankedTensorType(ty)
-        shape = tuple(ty.shape)
-        dtype = ir_type_to_dtype(ty.element_type)
-        return ShapedArray(shape, dtype)
-    out_types = [convert(o) for o in result_types]
-    return tuple(out_types)
-
-_custom_call_p.def_abstract_eval(_custom_call_abstract_eval)
-
-def _custom_call_hlo_lowering(ctx, *args, call_target_name, result_types, backend_config, operand_layouts, result_layouts):
-    out = mlir.custom_call(
-            call_target_name,
-            operands=args,
-            result_types=list(result_types),
-            backend_config=backend_config,
-            operand_layouts=list(operand_layouts),
-            result_layouts=list(result_layouts),
-        ).results
-    return out
-
-mlir.register_lowering(
-    _custom_call_p,
-    _custom_call_hlo_lowering
-)