JAX kernel fusion

When jit-ting a JAX function, the XLA compiler tries to fuse its operations into fewer kernel launches. When the fused result is still too slow, we can either use Pallas, JAX’s API for writing custom GPU/TPU kernels, or write custom CUDA code.

Here we test Pallas and JAX’s foreign function interface to bridge to a custom CUDA kernel using a simple QK-Norm+RoPE operation (which one typically finds in the attention mechanisms of modern LLMs).

We’ll evaluate the implementations by examining the HLO, the intermediate representation (IR) of the XLA compiler, which shows the fusions, primitives and operations the GPU has to execute. For the Pallas and CUDA implementations the kernel appears as a single custom call, so the HLO shows where our code takes over rather than what it does.

The code to reproduce all results is on GitHub at dirmeier/jax-kernel-fusion.

QK-Norm+RoPE

Let $q\in {\mathbb{R}}^{B\times S\times H\times D}$ be the query vector of a self-attention mechanism. If $x=q[b,s,h,:]$ is a single attention head vector, QK-Norm+RoPE first normalises over the head dimension

where $\gamma$ are trainable parameters. It then rotates each channel $i$ against its partner $i+m$, where $m=D/2$, by an angle that depends on the sequence position $s$,

where ${\theta }_i={\Theta }^{-2i/D}$ and $\Theta$ is some constant.

From jaxpr to HLO

To understand what XLA is doing, let’s examine the entire path from Python to machine code through a Python implementation of QK-Norm+RoPE. In plain JAX, the operation looks like this:

def rope_tables(seq_len, head_dim, dtype=jnp.float32):
  inv_freq = 10_000 ** (-jnp.arange(0, head_dim, 2, dtype=dtype) / head_dim)
  angles = jnp.arange(seq_len, dtype=dtype)[:, None] * inv_freq[None, :]
  return jnp.cos(angles), jnp.sin(angles)

def qk_norm_rope(query, gamma, cos, sin, EPS=1e-6):
  mean_square = jnp.mean(jnp.square(query), axis=-1, keepdims=True)
  qn = query * jax.lax.rsqrt(mean_square + EPS) * gamma

  half = query.shape[-1] // 2
  lo, hi = qn[..., :half], qn[..., half:]
  c = cos[None, :, None, :]
  s = sin[None, :, None, :]
  return jnp.concatenate([lo * c - hi * s, hi * c + lo * s], axis=-1)

Jitting this function traces it into a jaxpr, JAX internal IR of a program:

def make_inputs(shape, seed=0):
  _, seq_len, _, head_dim = shape
  queries = jax.random.normal(jax.random.key(seed), shape, jnp.float32)
  gamma = jnp.linspace(0.5, 1.5, head_dim, dtype=jnp.float32)
  cos, sin = rope_tables(seq_len, head_dim)
  return queries, gamma, cos, sin

args = make_inputs((2, 32, 4, 128))
print(jax.make_jaxpr(qk_norm_rope)(*args))

{ lambda ; a:f32[2,32,4,128] b:f32[128] c:f32[32,64] d:f32[32,64]. let
    e:f32[2,32,4,128] = square a
    f:f32[2,32,4] = reduce_sum[axes=(3,) out_sharding=None] e
    g:f32[2,32,4,1] = broadcast_in_dim[broadcast_dimensions=(0, 1, 2)] f
    h:f32[2,32,4,1] = div g 128.0:f32[]
    i:f32[2,32,4,1] = add h 9.999999974752427e-07:f32[]
    j:f32[2,32,4,1] = rsqrt i
    k:f32[2,32,4,128] = mul a j
    l:f32[1,1,1,128] = broadcast_in_dim[broadcast_dimensions=(3,)] b
    m:f32[2,32,4,128] = mul k l
    n:f32[2,32,4,64] = slice[
      limit_indices=(2, 32, 4, 64)
      start_indices=(0, 0, 0, 0)
      strides=None
    ] m
    o:f32[2,32,4,64] = slice[
      limit_indices=(2, 32, 4, 128)
      start_indices=(0, 0, 0, 64)
      strides=None
    ] m
    p:f32[1,32,1,64] = broadcast_in_dim[broadcast_dimensions=(1, 3)] c
    q:f32[1,32,1,64] = broadcast_in_dim[broadcast_dimensions=(1, 3)] d
    r:f32[2,32,4,64] = mul n p
    s:f32[2,32,4,64] = mul o q
    t:f32[2,32,4,64] = sub r s
    u:f32[2,32,4,64] = mul o p
    v:f32[2,32,4,64] = mul n q
    w:f32[2,32,4,64] = add u v
    x:f32[2,32,4,128] = concatenate[dimension=3] t w
  in (x,) }

The jaxpr shows that JAX records 20 primitive operations. Conveniently, it doesn’t matter if we use vanilla JAX, Pallas or CUDA FFI calls: all three are traced to a jaxpr before they are lowered to HLO. However, while vanilla JAX records 20 primitives, Pallas and the FFI record a single one that stands for the whole kernel. Additionally, the 3 implementations take 3 different “routes” in order to produce machine code (shown below with the help of Gemini):

  route 1                 route 2                 route 3
  XLA, compiler-fused     Pallas                  CUDA via FFI

  jnp ops in Python       kernel fn in Python     qk_norm_rope.cu in C/CUDA
     |                       |                       |
  jaxpr                   pallas_call             nvcc, at BUILD time
     |                       |                       |
  StableHLO               StableHLO custom_call   StableHLO custom_call
                          @mosaic_gpu_v2          @qk_norm_rope_cuda
     |                       :                       :
     |                       :                       :
  HLO passes              HLO                     HLO
  (fusion decided)        custom-call, target=    custom-call, target=
                          mosaic_gpu_v2           qk_norm_rope_cuda
     |                       |                       |
  Triton IR               Mosaic GPU                 |
  (on GH200)              MLIR -> NVVM               |
     |                       |                       |
  LLVM IR                    |                       |
     +-----------------------+-----------------------+
                             |
                            PTX             virtual ISA
                             |
                           ptxas            optimises PTX for one architecture
                             |
                            SASS            machine code for sm_90

For vanilla JAX Python code, after tracing, the jaxpr is lowered to StableHLO, a HLO dialect, from where it is to lowered HLO. The HLO (at least with some caveats) shows how many kernels each implementation launches per call which can be used as a diagnostic tool. Importantly, a fusion represents exactly one kernel launch, while a custom-call represents an unknown number of them.

When we jit the qk_norm_rope function on an GH200, XLA produces the following HLO fusing everything into a single kernel operation:

args = make_inputs((128, 1024, 16, 128))
print(jax.jit(qk_norm_rope).lower(*args).compile().as_text())

  ENTRY %main.2 (queries.1: f32[128,1024,16,128], gamma.1: f32[128], cos.1: f32[1024,64], sin.1: f32[1024,64]) -> f32[128,1024,16,128] {
    %sin.1 = f32[1024,64]{1,0} parameter(3)
    %cos.1 = f32[1024,64]{1,0} parameter(2)
    %gamma.1 = f32[128]{0} parameter(1)
    %queries.1 = f32[128,1024,16,128]{3,2,1,0} parameter(0)
+   ROOT %fusion.8 = f32[128,1024,16,128]{3,2,1,0} fusion(%cos.1, %sin.1, %gamma.1, %queries.1),
+       kind=kCustom, calls=%fused_computation.6,
        backend_config={"fusion_backend_config":{
          "kind":"__triton",
          "block_level_fusion_config":{
            "num_warps":"8","output_tiles":[{"sizes":["8","1","16","64"]}]}}}
  }

The HLO shows that XLA, interestingly, uses OpenAI’s Triton utilizing 8 warps (8 x 32 threads) and an output tile shape of [8, 1, 16, 64].

As a comparison, let’s also have a look at the CUDA FFI HLO:

  ENTRY %main.1 (queries.1: f32[128,1024,16,128], gamma.1: f32[128], cos.1: f32[1024,64], sin.1: f32[1024,64]) -> f32[128,1024,16,128] {
    %sin.1 = f32[1024,64]{1,0} parameter(3)
    %cos.1 = f32[1024,64]{1,0} parameter(2)
    %gamma.1 = f32[128]{0} parameter(1)
    %queries.1 = f32[128,1024,16,128]{3,2,1,0} parameter(0)
+   ROOT %ffi_call.1 = f32[128,1024,16,128]{3,2,1,0} custom-call(%queries.1, %gamma.1, %cos.1, %sin.1), custom_call_target="qk_norm_rope_cuda", operand_layout_constraints={f32[128,1024,16,128]{3,2,1,0}, f32[128]{0}, f32[1024,64]{1,0}, f32[1024,64]{1,0}}, api_version=API_VERSION_TYPED_FFI
}

As expected (since we wrote only a single CUDA kernel), there is only a single primitive.

Now, more interestingly, the HLO of the tiled Pallas kernel:

  ENTRY %main.1 (queries.1: f32[128,1024,16,128], gamma.1: f32[128], cos.1: f32[1024,64], sin.1: f32[1024,64]) -> f32[128,1024,16,128] {
    %sin.1 = f32[1024,64]{1,0} parameter(3)
    %cos.1 = f32[1024,64]{1,0} parameter(2)
    %gamma.1 = f32[128]{0} parameter(1)
    %queries.1 = f32[128,1024,16,128]{3,2,1,0} parameter(0)
    %bitcast.2 = f32[131072,2048]{1,0} bitcast(%queries.1)
+   %wrapped_broadcast = f32[64,128]{1,0} fusion(%gamma.1), kind=kLoop, calls=%wrapped_broadcast_computation
+   %pallas_call.1 = f32[131072,2048]{1,0} custom-call(%bitcast.2, %wrapped_broadcast, %cos.1, %sin.1), custom_call_target="mosaic_gpu_v2", operand_layout_constraints={f32[131072,2048]{1,0}, f32[64,128]{1,0}, f32[1024,64]{1,0}, f32[1024,64]{1,0}}, api_version=API_VERSION_TYPED_FFI
  ROOT %bitcast.1.0 = f32[128,1024,16,128]{3,2,1,0} bitcast(%pallas_call.1)
}

Even though there is a custom-call within the pallas_call primitive, it (pallas_call) only uses a single kernel launch, since Pallas compiles every kernel through Mosaic’s _lower_as_gpu_kernel which corresponds to exactly one CUDA kernel.

Runtimes

I ran the three implementations on a GH200 (which uses a H100 Hopper GPU) using the query dimensionality (128, 1024, 16, 128). Results are shown below:

Conclusion

Here, we developed QK-Norm+RoPE implementations using vanilla JAX, Pallas and CUDA FFI, and evaluated how XLA lowers them to HLO. Surprisingly (and counter-intuitively), I found the FFI path significantly easier thatn the Pallas approach. The results of the runtime measurements were expectedly unconclusive, given the simplicity of the operation and the consequential kernel fusion that was achieved by XLA. In addition to gaining a better understanding of HLO and Pallas, there’s also some insights I’ve gained:

• A fusion is exactly one kernel launch while a custom-call launches an unknown number.
• copy, transpose or concatenate outside a fusion are real kernels, and usually mark where fusion was blocked.
• Compare pre- and post-optimisation HLO with --xla_dump_to. If they look similar, XLA fused little or nothing, which is worth investigating.
• HLO for the same code differs per architecture: on my M1 it yielded three fusions while it was a single fusion on a GH200.
• Consider the actual hardware architecture that you work with. If the data fits into the cache, using SMEM does not increase throughput.

Hope reading this was informative to some 🙂🦉.