Feat: internal 1d and 2d ring attention kernel

[eltsai]

No description provided.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[gemini-code-assist]

Code Review

This pull request introduces an internal ring attention kernel for MaxDiffusion, which performs KV permutation directly within the Pallas kernel using remote DMA to improve performance. The changes include the implementation of the new kernel, updates to kernel registration and configuration, and a refined VMEM ceiling model for tile size grid search to better optimize tile selection. The review identified a critical bug in the DMA descriptor handling within the new kernel and noted redundant entries in the kernel registration list, both of which require attention.

[gemini-code-assist]

In _prime(), _local_load(0, 0) is called twice, which constructs two separate sets of DMA descriptors. The first set is started but never waited on, and the second set is waited on but never started. This can lead to compilation issues, deadlocks, or undefined behavior in Pallas. Store the descriptors in a local variable and reuse them for both .start() and .wait().

Suggested change

	@pl.when(is_first_hop)
	def _prime():
	for dma in _local_load(0, 0):
	dma.start()
	for dma in _local_load(0, 0):
	dma.wait()
	@pl.when(is_first_hop)
	def _prime():
	dmas = _local_load(0, 0)
	for dma in dmas:
	dma.start()
	for dma in dmas:
	dma.wait()

[gemini-code-assist]

There are duplicate and redundant attention kernel names in the in tuple. Specifically, "ulysses_ring_custom_iperm_fixed_m_nocond" is repeated multiple times, and "ulysses_ring_custom_iperm_fixed_m" is also duplicated. These redundant entries should be removed to keep the code clean and maintainable.

Suggested change

	"ulysses_ring_custom_iperm",
	"ulysses_ring_custom_iperm_fixed_m",
	"ulysses_ring_custom_iperm_fixed_m_nocond",
	"ulysses_ring_custom_iperm_fixed_m_nocond",
	"ulysses_ring_custom_iperm_fixed_m",
	"ulysses_ring_custom_iperm_fixed_m_nocond",
	"ulysses_ring_custom_iperm",
	"ulysses_ring_custom_iperm_fixed_m",
	"ulysses_ring_custom_iperm_fixed_m_nocond",