Causal Mask HLO for jnp.tril(jnp.ones()) can be simplified

[apivovarov]

Quite often we see the following JAX code to create and apply causal mask in the self-attention layer.

import jax
import jax.numpy as jnp
from jax import Array, random

key = random.PRNGKey(42)
qk = random.uniform(key, shape=(4,4))

def apply_causal_mask(qk: Array):
  seq_len = qk.shape[-1]
  mask = jnp.tril(jnp.ones((seq_len, seq_len))).astype('bool')
  return jnp.where(mask, qk, -jnp.inf)

If we jit + lower + compile + as_text this function

print(jax.jit(apply_causal_mask).lower(qk).compile().as_text())

then we will get the following hlo:

HloModule jit_apply_causal_mask, entry_computation_layout={(f32[4,4]{1,0})->f32[4,4]{1,0}}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation (param_0.1: f32[4,4]) -> f32[4,4] {
  %iota.5 = s32[4,4]{1,0} iota(), iota_dimension=0, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/iota[dtype=int32 shape=(4, 4) dimension=0]" source_file="<stdin>" source_line=5}
  %iota.4 = s32[4,4]{1,0} iota(), iota_dimension=1, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/iota[dtype=int32 shape=(4, 4) dimension=1]" source_file="<stdin>" source_line=5}
  %compare.2 = pred[4,4]{1,0} compare(s32[4,4]{1,0} %iota.5, s32[4,4]{1,0} %iota.4), direction=GE, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/ge" source_file="<stdin>" source_line=5}

  %constant.5 = f32[] constant(1)
  %broadcast.8 = f32[4,4]{1,0} broadcast(f32[] %constant.5), dimensions={}
  %constant.3 = f32[] constant(0)
  %broadcast.7 = f32[4,4]{1,0} broadcast(f32[] %constant.3), dimensions={}
  %select.3 = f32[4,4]{1,0} select(pred[4,4]{1,0} %compare.2, f32[4,4]{1,0} %broadcast.8, f32[4,4]{1,0} %broadcast.7), metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/select_n" source_file="<stdin>" source_line=5}
  %compare.1 = pred[4,4]{1,0} compare(f32[4,4]{1,0} %select.3, f32[4,4]{1,0} %broadcast.7), direction=NE, metadata={op_name="jit(apply_causal_mask)/jit(main)/convert_element_type[new_dtype=bool weak_type=False]" source_file="<stdin>" source_line=5}

  %param_0.1 = f32[4,4]{1,0} parameter(0)
  %constant.1 = f32[] constant(-inf)
  %broadcast.6 = f32[4,4]{1,0} broadcast(f32[] %constant.1), dimensions={}, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/broadcast_in_dim[shape=(4, 4) broadcast_dimensions=()]" source_file="<stdin>" source_line=8}
  ROOT %select.2 = f32[4,4]{1,0} select(pred[4,4]{1,0} %compare.1, f32[4,4]{1,0} %param_0.1, f32[4,4]{1,0} %broadcast.6), metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/select_n" source_file="<stdin>" source_line=8}
}

ENTRY %main.26 (Arg_0.1: f32[4,4]) -> f32[4,4] {
  %Arg_0.1 = f32[4,4]{1,0} parameter(0), sharding={replicated}
  ROOT %fusion = f32[4,4]{1,0} fusion(f32[4,4]{1,0} %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/select_n" source_file="<stdin>" source_line=8}
}

We noticed that fused_computation can be simplified.

In particular, Mask related code can be simplified to just iota + iota + compare

%fused_computation (param_0.1: f32[4,4]) -> f32[4,4] {
  %iota.5 = s32[4,4]{1,0} iota(), iota_dimension=0, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/iota[dtype=int32 shape=(4, 4) dimension=0]" source_file="<stdin>" source_line=5}
  %iota.4 = s32[4,4]{1,0} iota(), iota_dimension=1, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/iota[dtype=int32 shape=(4, 4) dimension=1]" source_file="<stdin>" source_line=5}
  %compare.2 = pred[4,4]{1,0} compare(s32[4,4]{1,0} %iota.5, s32[4,4]{1,0} %iota.4), direction=GE, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/ge" source_file="<stdin>" source_line=5}

  %param_0.1 = f32[4,4]{1,0} parameter(0)
  %constant.1 = f32[] constant(-inf)
  %broadcast.6 = f32[4,4]{1,0} broadcast(f32[] %constant.1), dimensions={}, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/broadcast_in_dim[shape=(4, 4) broadcast_dimensions=()]" source_file="<stdin>" source_line=8}
  ROOT %select.2 = f32[4,4]{1,0} select(pred[4,4]{1,0} %compare.2, f32[4,4]{1,0} %param_0.1, f32[4,4]{1,0} %broadcast.6), metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/select_n" source_file="<stdin>" source_line=8}
}

As a result the following several ops can be removed

  %constant.5 = f32[] constant(1)
  %broadcast.8 = f32[4,4]{1,0} broadcast(f32[] %constant.5), dimensions={}
  %constant.3 = f32[] constant(0)
  %broadcast.7 = f32[4,4]{1,0} broadcast(f32[] %constant.3), dimensions={}
  %select.3 = f32[4,4]{1,0} select(pred[4,4]{1,0} %compare.2, f32[4,4]{1,0} %broadcast.8, f32[4,4]{1,0} %broadcast.7), metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(tril)/select_n" source_file="<stdin>" source_line=5}
  %compare.1 = pred[4,4]{1,0} compare(f32[4,4]{1,0} %select.3, f32[4,4]{1,0} %broadcast.7), direction=NE, metadata={op_name="jit(apply_causal_mask)/jit(main)/convert_element_type[new_dtype=bool weak_type=False]" source_file="<stdin>" source_line=5}

What do you think about recognizing such a pattern and applying the described simplification to it?

Seems like a very common use-case in LLM models.

Our team will be happy to work on it.

[akuegel]

This looks like a useful optimisation to me. However I think it could also be possible to pre-optimize this on the jax side, maybe you can try

mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))

instead of converting to bool at the end. Still, ideally the compiler should be able to figure out that this can be simplified.

[apivovarov]

Tried this line to generate the mask - fused_computation hlo code is the same as before.

[akuegel]

Thanks for checking. It is strange that with jnp.ones it would still create a broadcast of a f32 1.0 constant value when dtype is set to bool. But I guess this is then because of jnp.tril op.

[apivovarov]

I double checked the solution again - it gives optimized HLO. Sorry for confusion.

def apply_causal_mask(qk: Array):
  seq_len = qk.shape[-1]
  mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=bool))
  return jnp.where(mask, qk, -jnp.inf)

[jakevdp]

A more direct way to generate this mask is

mask = jnp.tri(seq_len, dtype=bool)

That results in more-or-less the HLO you mentioned:

>>> print(jax.jit(apply_causal_mask).lower(qk).compile().as_text())
HloModule jit_apply_causal_mask, entry_computation_layout={(f32[4,4]{1,0})->f32[4,4]{1,0}}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation (param_0.1: f32[4,4]) -> f32[4,4] {
  %iota.4 = s32[4,4]{1,0} iota(), iota_dimension=0, metadata={op_name="jit(apply_causal_mask)/jit(main)/iota[dtype=int32 shape=(4, 4) dimension=0]" source_file="<ipython-input-4-5c7489ff92a8>" source_line=10}
  %iota.2 = s32[4,4]{1,0} iota(), iota_dimension=1, metadata={op_name="jit(apply_causal_mask)/jit(main)/iota[dtype=int32 shape=(4, 4) dimension=1]" source_file="<ipython-input-4-5c7489ff92a8>" source_line=10}
  %compare.0 = pred[4,4]{1,0} compare(s32[4,4]{1,0} %iota.4, s32[4,4]{1,0} %iota.2), direction=GE, metadata={op_name="jit(apply_causal_mask)/jit(main)/ge" source_file="<ipython-input-4-5c7489ff92a8>" source_line=10}
  %param_0.1 = f32[4,4]{1,0} parameter(0)
  %constant.0 = f32[] constant(-inf)
  %broadcast.1 = f32[4,4]{1,0} broadcast(f32[] %constant.0), dimensions={}, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/broadcast_in_dim[shape=(4, 4) broadcast_dimensions=()]" source_file="<ipython-input-4-5c7489ff92a8>" source_line=11}
  ROOT %select.1 = f32[4,4]{1,0} select(pred[4,4]{1,0} %compare.0, f32[4,4]{1,0} %param_0.1, f32[4,4]{1,0} %broadcast.1), metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/select_n" source_file="<ipython-input-4-5c7489ff92a8>" source_line=11}
}

ENTRY %main.15 (Arg_0.1: f32[4,4]) -> f32[4,4] {
  %Arg_0.1 = f32[4,4]{1,0} parameter(0), sharding={replicated}
  ROOT %fusion = f32[4,4]{1,0} fusion(f32[4,4]{1,0} %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(apply_causal_mask)/jit(main)/jit(_where)/select_n" source_file="<ipython-input-4-5c7489ff92a8>" source_line=11}
}

[apivovarov]

That works! Thank you!

[apivovarov]

I think we identified common JAX code issues encountered when generating a Causal mask. The user either neglected to specify dtype=bool in ones() or mistakenly applied .astype(bool) to the result of tril() instead of to ones(). Consequently, the mask will be converted from f32 to bool, resulting in suboptimal HLO.

I opened PR-9867 (Merged) to simplify potential Causal mask suboptimal HLO