[tx] Add cuDNN flash attention

[raulchen]

Summary

• Add dot_product_attention() in tx/layers/attention.py that
  uses cuDNN on GPU for both causal (prefill) and non-causal
  (decode) attention, with XLA fallback for CPU/TPU
• Update Llama3 and Qwen3 models to use the shared attention
  function
• Add GPU CI workflow via Anyscale

Test plan

• GPU tests verify cuDNN output matches XLA mask-based
  attention
• Tests cover: causal with padding, no padding, mixed lengths,
  GQA, decode (non-causal)
• CPU CI excludes GPU tests (--ignore=tests/gpu)

Benchmark results #891 (comment)

[gemini-code-assist]

Code Review

This pull request introduces an excellent optimization by adding support for cuDNN flash attention, which will significantly reduce memory usage during training and prefill on GPUs. The implementation is robust, automatically handling both left- and right-padded sequences by cleverly converting them to the right-padded format required by cuDNN. The fallback to a standard mask-based attention for the decode phase and non-GPU backends is a sound design choice. The new functionality is accompanied by a comprehensive set of unit tests for both CPU and GPU, ensuring numerical correctness across different scenarios. The addition of a dedicated GPU CI workflow is also a great step towards maintaining the stability of this feature. I have one suggestion to further improve the test coverage for an edge case.

[raulchen]

/gemini review

[gemini-code-assist]

Code Review

The pull request introduces cuDNN flash attention support for JAX models, significantly reducing memory usage for causal attention during training and prefill. It includes a new shared attention module that intelligently selects between cuDNN flash attention and a mask-based fallback, handling left-padded sequences by internally converting them to a right-padded format. New GPU CI workflows and comprehensive tests for the new attention mechanism and sequence shifting utility have also been added. The changes are well-structured and integrate smoothly into the existing model architectures.

[pcmoritz]

I'm pretty sure this is not needed any more on modern Ray versions :)

[raulchen]

yeah, copied from other files. removed all

[pcmoritz]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant performance improvement by adding support for cuDNN's flash attention on GPUs. The new dot_product_attention utility is well-designed, providing a clean abstraction with a fallback for other backends. The models are correctly updated to use this new utility, and the changes are accompanied by a comprehensive suite of GPU-specific tests that validate the numerical equivalence with the existing XLA implementation.

My review focuses on the new test code, where I've suggested a couple of improvements for maintainability and adherence to modern JAX practices. Overall, this is a solid contribution.

[gemini-code-assist]

jax.random.key is an alias for jax.random.PRNGKey and its usage is discouraged in new code. Please use jax.random.PRNGKey for clarity and future compatibility.

Suggested change

	q = jax.random.normal(jax.random.key(0), (batch, seq_len, num_heads, head_dim), dtype=dtype)
	k = jax.random.normal(jax.random.key(1), (batch, seq_len, num_kv_heads, head_dim), dtype=dtype)
	v = jax.random.normal(jax.random.key(2), (batch, seq_len, num_kv_heads, head_dim), dtype=dtype)
	q = jax.random.normal(jax.random.PRNGKey(0), (batch, seq_len, num_heads, head_dim), dtype=dtype)
	k = jax.random.normal(jax.random.PRNGKey(1), (batch, seq_len, num_kv_heads, head_dim), dtype=dtype)
	v = jax.random.normal(jax.random.PRNGKey(2), (batch, seq_len, num_kv_heads, head_dim), dtype=dtype)

[gemini-code-assist]

jax.random.key is an alias for jax.random.PRNGKey and its usage is discouraged in new code. Please use jax.random.PRNGKey for clarity and future compatibility.

Additionally, this block duplicates tensor creation logic from make_qkv and test_gqa_decode. Consider refactoring make_qkv to accept separate q_len and kv_len to centralize this logic and improve maintainability.

Suggested change

	q = jax.random.normal(jax.random.key(0), (batch, 1, num_heads, head_dim), dtype=jnp.bfloat16)
	k = jax.random.normal(jax.random.key(1), (batch, kv_len, num_heads, head_dim), dtype=jnp.bfloat16)
	v = jax.random.normal(jax.random.key(2), (batch, kv_len, num_heads, head_dim), dtype=jnp.bfloat16)
	q = jax.random.normal(jax.random.PRNGKey(0), (batch, 1, num_heads, head_dim), dtype=jnp.bfloat16)
	k = jax.random.normal(jax.random.PRNGKey(1), (batch, kv_len, num_heads, head_dim), dtype=jnp.bfloat16)
	v = jax.random.normal(jax.random.PRNGKey(2), (batch, kv_len, num_heads, head_dim), dtype=jnp.bfloat16)

[gemini-code-assist]

jax.random.key is an alias for jax.random.PRNGKey and its usage is discouraged in new code. Please use jax.random.PRNGKey for clarity and future compatibility.

As mentioned in test_decode, this is another instance of duplicated tensor creation logic. Refactoring make_qkv would benefit this test as well.

Suggested change

	q = jax.random.normal(jax.random.key(0), (batch, 1, num_heads, head_dim), dtype=jnp.bfloat16)
	k = jax.random.normal(jax.random.key(1), (batch, kv_len, num_kv_heads, head_dim), dtype=jnp.bfloat16)
	v = jax.random.normal(jax.random.key(2), (batch, kv_len, num_kv_heads, head_dim), dtype=jnp.bfloat16)
	q = jax.random.normal(jax.random.PRNGKey(0), (batch, 1, num_heads, head_dim), dtype=jnp.bfloat16)
	k = jax.random.normal(jax.random.PRNGKey(1), (batch, kv_len, num_kv_heads, head_dim), dtype=jnp.bfloat16)
	v = jax.random.normal(jax.random.PRNGKey(2), (batch, kv_len, num_kv_heads, head_dim), dtype=jnp.bfloat16)