Is Pallas SparseCore supported for TPUv4?

[wtedw]

Hi,

I noticed the Pallas and OpenXLA docs say the TPU v4 has SparseCores, but the current JAX code doesn't support accessing them.

from jax.experimental.pallas import tpu as pltpu

print(pltpu.get_tpu_info())
assert pltpu.get_tpu_info().sparse_core is not None, "No SparseCore found"

TpuInfo(chip_version=<ChipVersion.TPU_V4: 'v4'>, generation=4, num_cores=2, num_lanes=128, num_sublanes=8, mxu_column_size=128, vmem_capacity_bytes=16777216, cmem_capacity_bytes=67000000, smem_capacity_bytes=1048576, hbm_capacity_bytes=17200000000, mem_bw_bytes_per_second=615000000000, bf16_ops_per_second=137500000000000, int8_ops_per_second=0, fp8_ops_per_second=0, int4_ops_per_second=0, sparse_core=None)
Traceback (most recent call last):
  File "/tmp/sparse2.py", line 4, in <module>
    assert pltpu.get_tpu_info().sparse_core is not None, "No SparseCore found"
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: No SparseCore found

References

• jax/docs/pallas/tpu/sparsecore.md

  Lines 57 to 66 in 231b735

  	| Attribute | TPU v4 | TPU v5p | TPU v6e | TPU 7x (Ironwood) |
  	: : : : (Trillium) : :
  	| :---------------- | :----- | :------ | :------------ | :---------------- |
  	| SparseCores / | 4 | 4 | 2 | 2 (4 physical |
  	: Chip : : : : cores) :
  	| Vector subcores / | 16 | 16 | 16 | 16 |
  	: SparseCore : : : : :
  	| SIMD Width | 8 | 8 | 8 (F32)<br>16 | 16 |
  	: : : : (BF16) : (F32)<br>32(BF16) :
  	| HBM Capacity | 32 GiB | 96 GiB | 32 GiB | 192 GB |

• https://openxla.org/xla/sparsecore#specifications_at_a_glance

Furthermore, from the tests, it seems like Pallas SparseCore support is explicitly restricted to v5p+:

jax/tests/pallas/tpu_sparsecore_pallas_test.py

Lines 49 to 54 in 231b735

	class PallasSCMeshTest(jtu.JaxTestCase):
	def setUp(self):
	if not jtu.is_device_tpu(5, "p") and not jtu.is_device_tpu_at_least(6):
	self.skipTest("SparseCore only supported on TPU v5p+")
	super().setUp()

My questions:

1. Is it currently expected that Pallas cannot target TPU v4 SparseCores?
2. Is there a roadmap to add v4 support, or will Pallas SparseCore kernels remain a v5p+ exclusive feature?

Thank you!