deciax: Running a Decision Model on TPU with JAX and Pallas

Routing a support ticket, checking whether it asks for a refund, and scoring its urgency all call for structured decisions. A model can return choices and probabilities directly, without generating a textual answer token by token.

Laya is a Decision Model that handles choice, score, and yes/no questions in a single forward pass. I ported its inference path to JAX and built deciax to run these models on TPU, then optimized the computation with Pallas.

In this benchmark¹, a single TPU v6e chip reached a single-question p50 of 2.18 ms for the English model and 1.55 ms for the multilingual model. The fastest measured results on A100 SXM4, H100 PCIe, and H200 NVL were 6.4 / 4.0 ms, 4.8 / 3.2 ms, and 5.8 / 4.1 ms, respectively (English / multilingual).

At 50 questions per call, v6e Pallas reached 33.69 / 17.83 ms, roughly 2.3× faster than plain JAX. The fastest measured A100, H100, and H200 results were 68.8 / 36.2 ms, 44.0 / 23.2 ms, and 39.8 / 21.2 ms, respectively.

# One input, several decisions

deciax takes an input and a set of questions, then returns structured results. Here is a duplicate-charge ticket:

import deciax

decider = deciax.load("laya")
result = decider.decide(
    {"user": "I was charged twice this month. Please refund the duplicate."},
    {
        "department": {
            "type": "choice",
            "instructions": "Which department should handle this?",
            "criteria": {
                "billing": "payments, invoices, refunds",
                "technical": "bugs, outages, system errors",
                "other": "everything else",
            },
        },
        "urgency": {
            "type": "score",
            "instructions": "How urgent is this?",
            "criteria": ["not urgent", "soon", "blocking"],
        },
        "refund_requested": {
            "type": "noul",
            "instructions": "Does the user ask for a refund?",
        },
    },
)
print(result)

# Move the inference path to JAX

deciax preserves Laya’s input construction, ModernBERT encoder, and decision heads while moving inference to JAX. Its runtime does not require PyTorch, Transformers, or Laya; the checkpoints and tokenizer still come from the original model. It also offers a /v1/systemone HTTP endpoint. The project documentation lists its current scope.

# Start with the baseline: PyTorch eager and JAX

Measurements ran on Oct 12, 2026, across RTX 3090, A100 SXM4, RTX A6000, H100 PCIe, and H200 NVL hosts, plus single-chip TPU v5e and v6e runs.

Latency figures show warmed in-process p50 per call², including input encoding, inference, and result handling. Solid bars show p50; dashed markers show p95.

First, compare PyTorch eager on GPU with ordinary JAX on TPU:

For English single-question calls, GPU eager p50 ranged from 13.3 to 25.8 ms, while v6e JAX took 2.27 ms. At 50 questions, v6e JAX took 76.79 ms, compared with the lowest GPU eager result of 57.5 ms on H200 NVL. Multilingual followed a similar pattern: v6e JAX grew from 1.61 ms to 41.24 ms, while H200 eager took 34.0 ms at 50 questions.

# Where time goes: dispatch overhead and model computation

Next, break down baseline calls using the English model³:

Single-question GPU eager calls issue many small kernels: the English model used 1,090–1,213 cudaLaunchKernel calls per inference. The large remainder outside device execution points toward both kernel fusion and reducing dispatch overhead.

At 50 questions, computation accounts for most of the call. English v6e JAX took 76.79 ms end to end, with separately measured forward-only time of 71.27 ms and encoding near 3 ms. In XProf, dot_general and sub accounted for about 38% and 21% of this device time, directing attention to matrix computation and data movement in elementwise operations.

After the port, JAX operations serve as the baseline for exploring how to fuse computation on TPU. At the measured shapes, Pallas fuses multiple computations within each layer into one kernel⁴.

# Optimized: the fastest measured implementation on each platform

GPU testing covers TileLang, CUDA Graph, and torch.compile. We select the fastest measured implementation for each host, model, and question count to make the comparison as fair as possible. The TPU path uses v6e Pallas.

Single-question v6e Pallas p50 was 2.18 ms for English and 1.55 ms for multilingual. The lowest English GPU result was 4.8 ms on H100 PCIe using torch.compile with cudagraph; the lowest multilingual result was 3.2 ms. Pallas makes little difference relative to plain JAX at this shape.

At 50 questions, English v6e latency fell from 76.79 ms with JAX to 33.69 ms with Pallas; multilingual fell from 41.24 ms to 17.83 ms. These are speedups of 2.28× and 2.31×. English forward-only time fell from 71.27 ms to 28.25 ms, and multilingual from 36.54 ms to 13.36 ms.

# Latency as the question count grows

Connect the measurements at 1, 5, 10, and 50 questions to compare latency on H100, H200, and v6e Pallas:

# HTTP serving and installation

We also measured deciax.serve. For single-question calls over loopback with keep-alive, v6e Pallas HTTP p50 was 3.0 ms for English and 2.4 ms for multilingual. This includes the local HTTP round trip, but no external network latency or concurrent load test.

Use Python 3.14 or newer. Install it on a TPU host with:

python -m pip install "deciax[tpu]==0.1.1"

This release supports the Python interface example above. To inspect the Pallas implementation measured here, check out the corresponding code from the repository:

git clone https://github.com/Junyi-99/deciax.git
cd deciax
git checkout 6f3a930
uv sync --extra tpu

# Cost: Million model input tokens per US dollar

Cost compares only optimized English-model implementations: the fastest measured GPU variant and v6e Pallas on TPU. We calculate both on-demand and preemptible / spot scenarios.

Performance and price data were collected on Oct 12, 2026.

Prices are USD per GPU-hour or TPU chip-hour.

The calculation is⁶:

Million model input tokens per US dollar
  = model input tokens per second × 3600 / hourly price / 1,000,000

Model input tokens include the ticket text repeated for each question.

# Conclusion

deciax brings Laya’s structured decisions to TPU. In this experiment, plain JAX already handled single-question calls in about 2 ms, with little additional benefit from Pallas at that shape. At 50 questions per call, the fused Pallas path was roughly 2.3× faster than plain JAX. The question count changes both the dominant overhead and the benefit of optimization.

For the inputs, models, and implementations tested here, v6e Pallas had lower latency than the compared GPU implementations at both one and 50 questions.

You can try deciax in Colab or explore the code in the GitHub repository.