[Pallas TPU] Support SparseCore-to-TensorCore Remote DMA on Cloud TPU
Author: ayaka14732Created Jul 31, 2026Updated Sep 17, 2026
Labelsenhancement
Description
I am trying to use the functionality demonstrated by MpmdAsyncTest.test_async_sc_tc_prefetch_vmem, but found that it is skipped on Cloud TPU.
This functionality would be extremely useful for writing my Pallas kernels. I would greatly appreciate Cloud TPU support as soon as possible.
System info
jax: 0.11.1.dev20260730+5b77e49
jaxlib: 0.11.0
numpy: 2.5.1
python: 3.14.6 (main, Jul 23 2026, 14:45:24) [Clang 22.1.3 ]
device info: TPU v6 lite-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='instance-20260718-100907', release='6.8.0-1063-gcp', version='#69~22.04.1-Ubuntu SMP Tue Jun 30 15:12:15 UTC 2026', machine='x86_64')
Source: jax-ml/jax