What is the default multicore behaviour on TPUs?

[oluwatimilehin]

Hello, I'm benchmarking some simple kernels like jax.lax.add, jax.lax.subtract, etc., on the TPUv4, and it looks like the kernel gets executed on both cores:

The above is the trace from running jax.lax.abs on a 1x4096x256 tensor. It looks like both cores execute the kernel at the same time.

What I find confusing, though, is that the tensor does not seem to be split between the tensor cores. For example, selecting one of the blocks gives the same HLO for both cores:

%abs.2 = s32[1,4096,256]{2,1,0:T(8,128)} abs(s32[1,4096,256]{2,1,0:T(8,128)} %Arg_0.1)

I'm not using any Jax primitives to split the work, and I don't understand why it's running the operation on both cores. From the profile, it also looks like CMEM is not being used at all, as the read and write bandwidth utilization for both is zero. Could you please explain the behaviour?

[ayaka14732]

On TPU v4, a JAX logical device is normally configured as a Megacore containing two physical TensorCores. You can confirm the active configuration with jax.devices()[0].num_cores.

This is intra-device parallelism, so it does not require an explicit JAX sharding primitive. HLO describes the computation and shapes for the whole logical device. Consequently, XProf can associate both TensorCore executions with the same full-shaped HLO instruction. That does not mean that both cores process the entire tensor.

In a reproduction of abs(s32[1,4096,256]), the final low-level program reads the Megacore core ID and uses it to compute different HBM offsets. The full tiled input window is (1,512,2), while each core handles (1,256,2), corresponding to s32[1,2048,256] per core. Thus, both cores execute the same SPMD program concurrently, but on disjoint halves of the tensor.

Zero CMEM bandwidth is expected for this kernel. CMEM is shared scratch memory, not a cache that every Megacore operation must traverse. This elementwise operation has no cross-core communication or shared-data reuse, so XLA uses HBM → per-core VMEM/registers → HBM and does not use CMEM.

If explicit control is required, Pallas exposes the TensorCore mesh within the logical device. A kernel can use pltpu.TensorCoreMesh(axis_name="tc_core", num_cores=2) together with jax.lax.axis_index("tc_core") and pl.ds(...) to assign an exact, disjoint HBM slice to each TensorCore and explicitly stage it through per-core pltpu.VMEM.