Causal mask in Mosaic GPU

[southfreebird]

Hi team,

I'm exploring Mosaic GPU DSL and it looks really promising, great work! I'm currently going through the attention_mgpu example and wanted to ask: is there a good way to add a causal mask there?

For example, I’d like to reproduce the behaviour of causal_mask from attention on pallas:

span_q = q_step * block_q + jnp.arange(block_q)
span_k = kv_step * block_kv + jnp.arange(block_kv)
causal_mask = span_q[:, None] >= span_k[None, :]
qk = jnp.where(causal_mask, qk, DEFAULT_MASK_VALUE)

I tried this approach, but encountered an unimplemented error due to the use of iota (invoked via jnp.arange):

jax._src.source_info_util.JaxStackTraceBeforeTransformation: NotImplementedError: Unimplemented primitive in Pallas Mosaic GPU lowering: iota. Please file an issue on https://github.com/jax-ml/jax/issues.

Is there an alternative or recommended way to implement causal masking in Mosaic GPU?

Thanks in advance!

[southfreebird]

Hi @apaszke,
Apologies for the mention, I believe this might be a quick question for you to answer, and I’d really appreciate your insight. :)

[apaszke]

Sorry I missed this entirely! iota is a bit annoying here, because Mosaic GPU does "forward-only" layout inference, meaning that there's no way for a use of a value to influence the format in which it is produced. This is a problem for those "array factory" methods since we can't really know how you'll use them when the array is created. As a workaround we have plgpu.broadcasted_iota. You should be able to use that!

[southfreebird]

Thank you very much for your answer. It was very helpful!