Calling a Pallas kernel

There is no torch-side Pallas. trace_pallas takes the kernel and the tensors you would call it with, replaces every tensor with a jax.ShapeDtypeStruct carrying only shape and dtype, and jits the kernel over those meta values. Nothing executes, because meta values have no storage; the point is the module that lowering produces.

In that module the kernel appears as a single stablehlo.custom_call named tpu_custom_call, and its backend_config holds the compiled MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → body as a base64 blob. _extract_backend_config walks two levels of the module looking for exactly that operation and returns its config string. The docstring in the source shows a full example module, which is worth reading once because it makes the shape of the thing concrete.

So the object that crosses from jax to torch is text. Not a function pointer, not a compiled binary the torch runtime loads, and not anything torch can inspect. A string, handed onward.

  # Here we ignore the kwargs for execution as most of the time, the kwargs is only used in traced code.
  ir = jax.jit(
      kernel, static_argnums=static_argnums,
      static_argnames=static_argnames).lower(*jax_args, **kwargs).compiler_ir()
  payload = _extract_backend_config(ir)

  if use_cache:
    # if we reach here it means we have a cache miss.
    trace_pallas_arg_to_payload[hash_key] = payload

  return payload, tensor_args

That lowering step is the one piece of the path a machine without a TPU can attempt, so it was attempted here. jax 0.4.38 on CPU imports the MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → registration happily and then refuses the lowering, in one sentence: only interpret mode is supported on the CPU backend.

Run the same kernel with interpret=True and it lowers fine, into ordinary StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → with no custom call in it at all. Feed that module to the extractor from the pinned source and it returns None, which is the correct answer rather than a bug: interpret mode inlines the kernel, so there is no payload to extract and nothing for torch to carry.

The runnable half of this lesson is therefore Colab-pending, in the same sense LAB·P5 uses the term. The mechanism below is read from source at the pinned commit. Producing a real payload, printing its first eighty characters, and finding tpu_custom_call in a dumped HLO module all need a TPU runtime, and they belong in the same lab pass that closes chapter 12.

$ python pallas_lower.py           # the trace_pallas lowering, unmodified
jax 0.4.38 [CpuDevice(id=0)]
mosaic registration: imported
ValueError: Only interpret mode is supported on CPU backend.

$ python pallas_interp.py          # same kernel, interpret=True
tpu_custom_call in module: False
custom_call in module: False
module lines: 45

$ python extract_none.py           # _extract_backend_config, verbatim from 41398bf
payload: None

On the torch side the payload goes to _xla_tpu_custom_call, which lands in tensor_methods::tpu_custom_call and builds one TpuCustomCall IR node. Look at what that node's constructor does with the payload: it passes it as the hash seed. The bridge arc's first lesson enumerated the ingredients of a graph hash; this is the kernel body entering that number.

The consequence is worth stating plainly, because it surprises people who think of a kernel as a library call. Edit two lines inside the kernel body, and the payload changes, so the node hash changes, so the graph hash changes, so the step compiles again. A kernel is part of the program's identity, not a dependency the program links against.

Lowering the node emits an XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → custom call whose target name is the string tpu_custom_call, with layouts forced to torch's own major-to-minor order on both the inputs and the outputs, because MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → requires that ordering. A single-output kernel skips the tuple wrapper entirely, with a comment saying Mosaic rejects a tuple of one.

// torch_xla/csrc/ops/tpu_custom_call.cpp: the payload is the node hash seed
TpuCustomCall::TpuCustomCall(torch::lazy::OpList inputs,
                             xla::Shape output_shape,
                             const std::string& payload)
    : XlaNode(xla_tpu_custom_call, inputs, output_shape,
              /*num_outputs=*/output_shape.tuple_shapes_size(),
              torch::lazy::MHash(payload)),
      payload_(payload) {}

// torch_xla/csrc/xla_lower_util.cpp: and the lowering names the call target
  // Mosaic has some weird checks that disallow using a tuple output for single
  // element.
  if (output_shapes.size() == 1) {
    return {xla::CustomCallWithLayout(inputs[0].builder(),
                                      /*call_target_name=*/"tpu_custom_call",
                                      inputs, output_shapes[0], input_shapes,
                                      payload)};
  }

Lowering a kernel through jax is not free, and doing it on every call would put a jit trace inside your training step. There is a cache, and it is a module-level dictionary rather than anything the runtime knows about.

Read its key and you learn what counts as the same kernel: the default matmul precision, the kernel function object itself, the static argnums and argnames, the meta arguments with their shapes and dtypes, and the keyword arguments repr'd and sorted. A hit bumps a trace_pallas_cache_hit counter you can read in the same metrics report as the seam names from LAB·P5.

The cache is opt-in per call site, and the comment says why: the key assumes every keyword argument is hashable and not a tensor, which holds for the grouped matmul kernels that use it and is not promised in general. A kernel wrapper of your own that turns use_cache on has just taken on that assumption.

  hash_key = ()
  if use_cache:
    global trace_pallas_arg_to_payload
    # implcit assumption here that everything in kwargs is hashable and not a tensor,
    # which is true for the gmm and tgmm.
    hash_key = (jax.config.jax_default_matmul_precision, kernel, static_argnums,
                tuple(static_argnames)
                if static_argnames is not None else static_argnames,
                tuple(jax_args), repr(sorted(kwargs.items())).encode())
    if hash_key in trace_pallas_arg_to_payload:
      torch_xla._XLAC._xla_increment_counter('trace_pallas_cache_hit', 1)
      return trace_pallas_arg_to_payload[hash_key], tensor_args

The kernels the repo ships are not left as bare Python functions. Flash attention, paged attention, grouped matmul and the rest are each defined as a torch library op with a schema, then given two implementations. The XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → one calls the real thing. The CompositeExplicitAutograd one returns an empty tensor shaped like the query, and warns when the tensors are not on a meta device.

That second implementation is what makes the kernel usable under torch.compile. Dynamo builds fake outputs with meta tensors while it is capturing, and a fake tensor cannot run a MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → kernel; what it needs is a shape rule. The empty tensor is that rule, written as code.

The same registration is why a model containing one of these kernels still imports and traces on a CPU box, with a warning rather than an import error. It will not compute anything correct there, and the warning says so, but the graph is capturable, which is usually what you were doing on the CPU box anyway.

def non_xla_attetion(q, k, v, attention_type):
  # This will be called when dynamo use fake tensor to construct the fake output.
  # We need to make sure output tensor's shape is correct.
  if k.device != torch.device("meta"):
    warnings.warn(
        f'XLA {attention_type} attention should only be applied to tensors on XLA device'
    )

  # Return orignal shape of q.
  return torch.empty_like(q)


XLA_LIB.define(
    "flash_attention(Tensor q, Tensor k, Tensor v, bool casual=False) -> Tensor",
)


@impl(XLA_LIB, "flash_attention", "XLA")
def flash_attention_xla(q: torch.Tensor,
                        k: torch.Tensor,
                        v: torch.Tensor,
                        causal: bool = False):
  return flash_attention(q, k, v, causal=causal)


@impl(XLA_LIB, "flash_attention", "CompositeExplicitAutograd")
def flash_attention_non_xla(q: torch.Tensor,
                            k: torch.Tensor,
                            v: torch.Tensor,
                            causal: bool = False):
  return non_xla_attetion(q, k, v, "flash")

A kernel runs on one device's worth of data, so a sharded model cannot hand it a global tensor and hope. The wrapper for that is _shard_map, which calls enable_manual_sharding on each input, runs the single-device function, and calls disable_manual_sharding on each output after computing what the full shape should be from the mesh axis sizes.

Flash attention uses exactly that when a mesh is passed: nine input specs, three output specs, and the same single-device forward underneath either way. The specs for the log-sum-exp outputs differ from the specs for the queries, which is the kind of detail that only shows up when the kernel has more outputs than the obvious one.

Set this next to lesson one and the difference is the interesting part. mark_sharding leaves an annotation for the partitioner to propagate, and the compiler decides what the shard shapes are. Manual sharding ends propagation at a boundary and hands you the shard shapes as your problem. A kernel is the case where you wanted the second one.

The chapter above says vLLM's TPU backend runs PyTorch model definitions through a jax lowering path, and the repository states it in its own words: tpu-inference is described as a hardware plugin unifying JAX and PyTorch under a single lowering path within the vLLM project, whose aim includes running PyTorch model definitions performantly on TPU without any additional code changes.

The wrapper that does it is one file, and its imports are the argument. It imports torchax, pulls jax_view and torch_view out of torchax.interop, calls the torch model through torch.func.functional_call with jax-backed parameters viewed as torch tensors, and shards its inputs with jax.sharding.NamedSharding over a PartitionSpec. Those are the same two objects the jax path's chapter 10 built, applied to a torch module, which is the convergence chapter 10 of this path described made into imports.

Read that as a map rather than a recipe. The kernels sit in their own directory, the sharding is jax's, the model definition is torch's, and each of those layers has a chapter on this site that teaches it. Placing them is chapter 11's mastery item, and doing it yourself is worth more than reading someone else's answer.

Provenance for this section: read on 2026-08-15 from tpu-inference at commit 878eb5e, not run. Nothing here was measured.