diff --git a/CHANGELOG.md b/CHANGELOG.md
index ca8ca44..7104d65 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,39 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 -->
 
+## [Unreleased]
+
+### Added
+
+## [Unreleased]
+
+### Added
+
+* PyTorch backend: `tabfm.src.pytorch.seqpar` -- sequence-parallel (row-
+  sharded) multi-GPU inference under `torch.distributed`. Shards one ensemble
+  member's in-context rows across ranks with exact cross-rank attention
+  (log-sum-exp-combined induced attention; all-gathered context K/V for the
+  ICL blocks), enabling contexts that exceed a single device's memory
+  (e.g. a 1M-row context on 4x80GB at ~35GB/GPU). `seqpar.predict` /
+  `seqpar.predict_proba` mirror the estimators' own prediction paths for any
+  `n_estimators`.
+* JAX backend: `tabfm.src.jax.seqpar` -- the same sequence-parallel
+  inference for the JAX backend: a single process shards the rows across all
+  local devices with `jax.shard_map` (log-sum-exp-combined induced attention
+  scanned in key chunks; bias-masked K/V gathers for the ICL blocks), with
+  the same `predict` / `predict_proba` API.
+* JAX backend: `AttentionImplementation.SPLASH` ('splash') -- fused Pallas
+  splash-attention kernel for TPU inference (the TPU analogue of 'cudnn'):
+  fp32 softmax accumulation, key-prefix masks expressed via segment ids.
+  CPU-testable via `set_splash_interpret(True)` (Pallas interpret mode).
+* JAX backend: `AttentionImplementation.CUDNN` ('cudnn') -- fused cuDNN flash
+  attention for GPU inference. Boolean prefix masks are translated to cuDNN's
+  variable sequence-length support, so no `[T, T_src]` mask materializes.
+  On an H100 at a 135k-row context this takes a single-member
+  `predict_proba` from ~630s to ~7s with unchanged predictions
+  (bf16-noise-level differences).
+
+
 ## [1.0.1] - 2026-07-09
 
 ### Fixed
diff --git a/README.md b/README.md
index 0fcffd5..a0d79e2 100644
--- a/README.md
+++ b/README.md
@@ -136,6 +136,37 @@ print("Predicted Prices:", predictions)
 
 ---
 
+## Multi-GPU inference (PyTorch)
+
+TabFM reads the whole training fold as one in-context sequence, so very large
+contexts (roughly >450k rows on an 80GB GPU) exceed a single device's memory.
+`tabfm.src.pytorch.seqpar` shards the rows of the sequence across the ranks of
+a `torch.distributed` process group with mathematically exact attention (no
+approximation; results match the single-device path up to bf16 summation
+order). Launch one process per GPU, e.g. with `torchrun`:
+
+```python
+import torch.distributed as dist
+from tabfm import TabFMRegressor, tabfm_v1_0_0_pytorch
+from tabfm.src.pytorch import seqpar
+
+dist.init_process_group("nccl")
+model = tabfm_v1_0_0_pytorch.load(model_type="regression", device=f"cuda:{rank}")
+reg = TabFMRegressor(model=model, n_estimators=1)
+reg.fit(X_train, y_train)          # cheap: no GPU forward
+preds = seqpar.predict(reg, X_test)  # collective call; every rank returns preds
+```
+
+`seqpar.predict_proba` is the classification equivalent. A runnable script is
+provided in [examples/seqpar_regression_example.py](examples/seqpar_regression_example.py).
+
+The JAX backend has an equivalent, `tabfm.src.jax.seqpar`: a single process
+shards the rows across all local devices with `jax.shard_map`, with the same
+`seqpar.predict(reg, X_test)` / `seqpar.predict_proba(clf, X_test)` API (no
+`torchrun` needed). A 1M-row context fits
+in ~35GB/GPU on 4 devices (a single 80GB device cannot run it at all), and at
+single-device-feasible sizes the sharded path is ~5x faster on 4 GPUs.
+
 ## Examples Directory
 
 You can find runnable scripts for both classification and regression under the [examples/](examples/) folder:
diff --git a/examples/seqpar_regression_example.py b/examples/seqpar_regression_example.py
new file mode 100644
index 0000000..e9a6f05
--- /dev/null
+++ b/examples/seqpar_regression_example.py
@@ -0,0 +1,73 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Multi-GPU regression with TabFM v1.0.0 via sequence-parallel inference.
+
+Shards the in-context rows of each ensemble member across all GPUs of a
+``torch.distributed`` process group, so training folds that exceed a single
+device's memory can be used as context. Launch one process per GPU:
+
+    torchrun --standalone --nproc_per_node=4 examples/seqpar_regression_example.py
+
+The script also runs on a single GPU (``--nproc_per_node=1``).
+"""
+
+import os
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+import tabfm
+from tabfm.src.pytorch import seqpar
+
+
+def make_data(n_train=20_000, n_test=1_000, n_features=20, seed=0):
+  """Synthetic regression data: linear signal plus noise."""
+  rng = np.random.default_rng(seed)
+  x = rng.standard_normal((n_train + n_test, n_features)).astype(np.float32)
+  w = np.random.default_rng(1).standard_normal(n_features)
+  y = x @ w + 0.1 * rng.standard_normal(n_train + n_test)
+  return x[:n_train], y[:n_train], x[n_train:], y[n_train:]
+
+
+def main():
+  rank = int(os.environ["RANK"])
+  local_rank = int(os.environ["LOCAL_RANK"])
+  torch.cuda.set_device(local_rank)
+  dist.init_process_group("nccl")
+
+  x_train, y_train, x_test, y_test = make_data()
+
+  model = tabfm.tabfm_v1_0_0_pytorch.load(
+      model_type="regression", device=f"cuda:{local_rank}"
+  )
+  reg = tabfm.TabFMRegressor(model=model, n_estimators=4, random_state=0)
+  reg.fit(x_train, y_train)  # cheap: preprocessing only, no GPU forward
+
+  # Collective call: every rank participates and returns the full predictions.
+  preds = seqpar.predict(reg, x_test)
+
+  if rank == 0:
+    rmse = float(np.sqrt(np.mean((y_test - preds) ** 2)))
+    r2 = 1 - np.sum((y_test - preds) ** 2) / np.sum(
+        (y_test - y_test.mean()) ** 2
+    )
+    print(f"world_size={dist.get_world_size()}  RMSE={rmse:.4f}  R2={r2:.4f}")
+
+  dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+  main()
diff --git a/tabfm/src/jax/model.py b/tabfm/src/jax/model.py
index 42f3019..6612ea1 100644
--- a/tabfm/src/jax/model.py
+++ b/tabfm/src/jax/model.py
@@ -41,13 +41,18 @@
 import typing
 
 import einops
+
 rearrange = einops.rearrange
 repeat = einops.repeat
 from flax import nnx
 from flax.nnx import Module
 import jax
 import jax.numpy as jnp
-import jaxtyping as jt; import typeguard; import numpy as np; jt.typed = jt.jaxtyped(typechecker=typeguard.typechecked)
+import jaxtyping as jt
+import typeguard
+import numpy as np
+
+jt.typed = jt.jaxtyped(typechecker=typeguard.typechecked)
 
 Array = jax.Array | np.ndarray
 DType = Any
@@ -57,12 +62,32 @@
 from . import memory_efficient_attention
 
 
+# Run the SPLASH attention kernel in Pallas interpret mode (CPU-executable,
+# slow; for tests only). Set before the first predict.
+SPLASH_INTERPRET = False
+
+
+def set_splash_interpret(value):
+  """Enables/disables Pallas interpret mode for SPLASH attention."""
+  global SPLASH_INTERPRET
+  SPLASH_INTERPRET = value
+
 
 class AttentionImplementation(str, enum.Enum):
   JAX = 'jax'
   # vmaps jax.nn.dot_product_attention over the head dimension to save memory
   JAX_VMAP_ON_HEAD_DIM = 'jax_vmap_on_head_dim'
   FLASH = 'flash'
+  # Fused cuDNN flash attention (GPU only, fp16/bf16). Much faster than FLASH
+  # on long sequences; masks must be key-prefix (padding) masks, which is the
+  # only mask shape TabFM uses. Head dims > 128 need Hopper (sm90) or later --
+  # TabFM v1.0.0's ICL stage uses head dim 256, so 'cudnn' for the ICL layers
+  # requires an H100-class GPU (cuDNN raises NotImplementedError otherwise).
+  CUDNN = 'cudnn'
+  # Fused Pallas flash attention (TPU only, the analogue of CUDNN on GPU).
+  # fp32 softmax accumulation; masks must be key-prefix (padding) masks,
+  # which are expressed via splash segment ids.
+  SPLASH = 'splash'
   NONE = 'none'
 
 
@@ -82,7 +107,9 @@ def default(val: Any, d: Any) -> Any:
 
 
 @jt.typed
-def rotate_half(x: jt.Float[jax.Array | np.ndarray, '*B D_r']) -> jt.Float[jax.Array | np.ndarray, '*B D_r']:
+def rotate_half(
+    x: jt.Float[jax.Array | np.ndarray, '*B D_r'],
+) -> jt.Float[jax.Array | np.ndarray, '*B D_r']:
   """Rotates half of the input tensor's dimensions for rotary positional embeddings.
 
   This operation splits the last dimension into two halves and swaps them,
@@ -178,7 +205,9 @@ def __init__(self, num_dims: int, *, rngs: Any):
     self.per_dim_scale = nnx.Param(jnp.zeros(shape=(num_dims,)))
 
   @jt.typed
-  def __call__(self, x: jt.Float[jax.Array | np.ndarray, '*B L']) -> jt.Float[jax.Array | np.ndarray, '*B L']:
+  def __call__(
+      self, x: jt.Float[jax.Array | np.ndarray, '*B L']
+  ) -> jt.Float[jax.Array | np.ndarray, '*B L']:
     """Applies per-dimension scaling.
 
     Args:
@@ -247,7 +276,7 @@ def __init__(
     self.interpolate_factor = interpolate_factor
     self.dtype = dtype
 
-    assert dim >=2, f'dim must be at least 2. Got {dim}'
+    assert dim >= 2, f'dim must be at least 2. Got {dim}'
     # Apply theta rescaling based on NTK-aware scaling for longer sequence lengths
     theta *= theta_rescale_factor ** (dim / (dim - 2))
 
@@ -257,8 +286,7 @@ def __init__(
     elif freqs_for == 'lang':
       # Language-based frequencies (standard RoPE formulation)
       freqs_init = 1.0 / (
-          theta
-          ** (jnp.arange(0, dim, 2)[: (dim // 2)].astype(dtype) / dim)
+          theta ** (jnp.arange(0, dim, 2)[: (dim // 2)].astype(dtype) / dim)
       )
     elif freqs_for == 'pixel':
       # Pixel-based frequencies (often used for images)
@@ -432,7 +460,9 @@ def __init__(
     self.dtype = dtype
 
   @jt.typed
-  def __call__(self, src: jt.Shaped[jax.Array | np.ndarray, '... T']) -> jt.Float[jax.Array | np.ndarray, '... T E']:
+  def __call__(
+      self, src: jt.Shaped[jax.Array | np.ndarray, '... T']
+  ) -> jt.Float[jax.Array | np.ndarray, '... T E']:
     """Transforms integer indices to dense embeddings.
 
     Args:
@@ -489,7 +519,11 @@ def __init__(
     for hidden_dim in hidden_dims:
       self.layers.append(
           nnx.Linear(
-              prev_dim, hidden_dim, use_bias=use_bias, rngs=rngs, dtype=self.dtype
+              prev_dim,
+              hidden_dim,
+              use_bias=use_bias,
+              rngs=rngs,
+              dtype=self.dtype,
           )
       )
       self.layers.append(act_fn)
@@ -503,7 +537,9 @@ def __init__(
       )
 
   @jt.typed
-  def __call__(self, x: jt.Float[jax.Array | np.ndarray, '... L_in']) -> jt.Float[jax.Array | np.ndarray, '... L_out']:
+  def __call__(
+      self, x: jt.Float[jax.Array | np.ndarray, '... L_in']
+  ) -> jt.Float[jax.Array | np.ndarray, '... L_out']:
     """Forward pass through the MLP.
 
     Args:
@@ -525,10 +561,14 @@ def _logn(n: int, dtype: jnp.dtype) -> jax.Array | np.ndarray:
 
 @jt.typed
 def _extract_kv_from_cache(
-    cached_kv: Tuple[jt.Float[jax.Array | np.ndarray, 'B T N D'],
-                     jt.Float[jax.Array | np.ndarray, 'B T N D']]
-) -> Tuple[jt.Float[jax.Array | np.ndarray, 'B T N D'],
-           jt.Float[jax.Array | np.ndarray, 'B T N D']]:
+    cached_kv: Tuple[
+        jt.Float[jax.Array | np.ndarray, 'B T N D'],
+        jt.Float[jax.Array | np.ndarray, 'B T N D'],
+    ],
+) -> Tuple[
+    jt.Float[jax.Array | np.ndarray, 'B T N D'],
+    jt.Float[jax.Array | np.ndarray, 'B T N D'],
+]:
   """Extracts key and value tensors from the cache.
 
   Args:
@@ -547,9 +587,11 @@ def _extract_kv_from_cache(
 @jt.typed
 def _encode_kv_into_cache(
     k: jt.Float[jax.Array | np.ndarray, 'B T N D'],
-    v: jt.Float[jax.Array | np.ndarray, 'B T N D']
-) -> Tuple[jt.Float[jax.Array | np.ndarray, 'B T N D'],
-           jt.Float[jax.Array | np.ndarray, 'B T N D']]:
+    v: jt.Float[jax.Array | np.ndarray, 'B T N D'],
+) -> Tuple[
+    jt.Float[jax.Array | np.ndarray, 'B T N D'],
+    jt.Float[jax.Array | np.ndarray, 'B T N D'],
+]:
   """Encodes key and value tensors into the cache.
 
   Args:
@@ -600,7 +642,6 @@ def __init__(
     self.attention_impl = attention_impl
     self.dtype = dtype
 
-
     assert (
         self.head_dim * num_heads == self.embed_dim
     ), 'embed_dim must be divisible by num_heads'
@@ -623,13 +664,25 @@ def __init__(
     # if self.out_proj.bias is not None:
     #     self.out_proj.bias[...] = jnp.zeros_like(self.out_proj.bias[...])
     self.q_proj = nnx.Linear(
-        self.embed_dim, self.embed_dim, use_bias=use_bias, rngs=rngs, dtype=self.dtype
+        self.embed_dim,
+        self.embed_dim,
+        use_bias=use_bias,
+        rngs=rngs,
+        dtype=self.dtype,
     )
     self.k_proj = nnx.Linear(
-        self.embed_dim, self.embed_dim, use_bias=use_bias, rngs=rngs, dtype=self.dtype
+        self.embed_dim,
+        self.embed_dim,
+        use_bias=use_bias,
+        rngs=rngs,
+        dtype=self.dtype,
     )
     self.v_proj = nnx.Linear(
-        self.embed_dim, self.embed_dim, use_bias=use_bias, rngs=rngs, dtype=self.dtype
+        self.embed_dim,
+        self.embed_dim,
+        use_bias=use_bias,
+        rngs=rngs,
+        dtype=self.dtype,
     )
     self.query_ln = nnx.RMSNorm(self.head_dim, rngs=rngs, dtype=self.dtype)
     self.key_ln = nnx.RMSNorm(self.head_dim, rngs=rngs, dtype=self.dtype)
@@ -641,7 +694,9 @@ def __call__(
       query: jt.Float[jax.Array | np.ndarray, 'B T E'],
       key: jt.Float[jax.Array | np.ndarray, 'B T_src E'] | None,
       value: jt.Float[jax.Array | np.ndarray, 'B T_src E'] | None,
-      attn_mask: Optional[jt.Bool[jax.Array | np.ndarray, 'B #N #T T_src']] = None,
+      attn_mask: Optional[
+          jt.Bool[jax.Array | np.ndarray, 'B #N #T T_src']
+      ] = None,
       rope: Optional[RotaryEmbedding] = None,
       cached_kv: Optional[
           Tuple[
@@ -656,7 +711,9 @@ def __call__(
           jt.Float[jax.Array | np.ndarray, 'B T E'],  # Output tensor.
           Tuple[
               jt.Float[jax.Array | np.ndarray, 'B T_src N D'],  # New key cache
-              jt.Float[jax.Array | np.ndarray, 'B T_src N D'],  # New value cache
+              jt.Float[
+                  jax.Array | np.ndarray, 'B T_src N D'
+              ],  # New value cache
           ],
       ],
   ]:
@@ -689,7 +746,8 @@ def __call__(
 
     if attn_mask is not None:
       assert attn_mask.shape[0] == batch_size, (
-          f'attn_mask batch size must match query batch size: {attn_mask.shape[0]} != {batch_size}'
+          'attn_mask batch size must match query batch size:'
+          f' {attn_mask.shape[0]} != {batch_size}'
       )
     if self.attention_impl == AttentionImplementation.NONE:
       # For None attention, kv cache is empty.
@@ -705,7 +763,9 @@ def __call__(
     # 2. Handle K, V
     if cached_kv is not None:
       assert key is None, f'key must be None if cached_kv is not None {key=}'
-      assert value is None, f'value must be None if cached_kv is not None {value=}.'
+      assert (
+          value is None
+      ), f'value must be None if cached_kv is not None {value=}.'
       k, v = _extract_kv_from_cache(cached_kv)
       src_len = k.shape[-3]  # (Batch, Seq, Head, Dim)
     else:
@@ -731,7 +791,6 @@ def __call__(
       k = self.key_ln(k)
     q = self.per_dim_scale(q)
 
-
     new_kv = _encode_kv_into_cache(k, v)
     if attn_mask is not None:
       assert attn_mask.ndim == 4, 'attn_mask must be 4D'
@@ -754,6 +813,77 @@ def __call__(
           query_chunk_size=128 if tgt_len >= 128 else tgt_len,
           key_chunk_size=128 if src_len >= 128 else src_len,
       )
+    elif self.attention_impl == AttentionImplementation.CUDNN:
+      # Fused cuDNN flash attention: never materializes the [T, T_src] score
+      # matrix and runs orders of magnitude faster than the chunked FLASH
+      # path on long sequences (e.g. 24 ICL blocks at 135k-row context:
+      # ~137s -> ~1.6s on an H100).
+      #
+      # cuDNN cannot take TabFM's broadcastable boolean masks (it requires a
+      # full [B, N, T, T_src] mask, which would materialize T*T_src bools).
+      # Every mask in this model is a key-prefix (padding) mask -- "attend to
+      # the first n keys" -- so it maps exactly onto cuDNN's variable
+      # sequence-length support instead.
+      kv_seq_lens = None
+      if attn_mask is not None:
+        kv_seq_lens = (
+            attn_mask.reshape(batch_size, -1)[:, -src_len:]
+            .sum(-1)
+            .astype(jnp.int32)
+        )
+      attn_output = jax.nn.dot_product_attention(
+          query=q,
+          key=k,
+          value=v,
+          scale=1.0,
+          key_value_seq_lengths=kv_seq_lens,
+          implementation='cudnn',
+      )
+    elif self.attention_impl == AttentionImplementation.SPLASH:
+      # Fused Pallas splash-attention kernel (TPU; CPU only via the interpret
+      # knob below, for tests). The kernel applies no internal scaling, which
+      # matches this module's convention (scaling is folded into q). TabFM's
+      # boolean masks are always key-prefix (padding) masks; they are
+      # expressed via segment ids: queries carry segment 1, valid keys 1,
+      # padded keys 0, and splash only attends within equal segments.
+      from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel as _sak  # pylint: disable=g-import-not-at-top
+      from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as _sam  # pylint: disable=g-import-not-at-top
+
+      if attn_mask is not None:
+        kv_valid = attn_mask.reshape(batch_size, -1)[:, -src_len:]
+      else:
+        kv_valid = jnp.ones((batch_size, src_len), dtype=bool)
+      seg_q = jnp.ones((batch_size, tgt_len), jnp.int32)
+      seg_kv = kv_valid.astype(jnp.int32)
+
+      # Sequence lengths are 128-multiples in this codebase (inputs are padded
+      # to the chunked-attention granularity), so fixed 128 blocks divide.
+      blocks = _sak.BlockSizes(
+          block_q=min(128, tgt_len), block_kv=min(128, src_len)
+      )
+      kernel = _sak.make_splash_mha(
+          _sam.MultiHeadMask(
+              [_sam.FullMask((tgt_len, src_len))] * self.num_heads
+          ),
+          block_sizes=blocks,
+          head_shards=1,
+          q_seq_shards=1,
+          interpret=SPLASH_INTERPRET,
+      )
+
+      def _one(q_i, k_i, v_i, sq_i, skv_i):
+        return kernel(
+            q_i, k_i, v_i, segment_ids=_sak.SegmentIds(q=sq_i, kv=skv_i)
+        )
+
+      # splash takes [num_heads, seq, head_dim]; vmap over the batch axis.
+      attn_output = jax.vmap(_one)(
+          q.transpose(0, 2, 1, 3),
+          k.transpose(0, 2, 1, 3),
+          v.transpose(0, 2, 1, 3),
+          seg_q,
+          seg_kv,
+      ).transpose(0, 2, 1, 3)
     elif self.attention_impl == AttentionImplementation.JAX:
       attn_output = jax.nn.dot_product_attention(
           query=q,
@@ -788,8 +918,7 @@ def _attention_fn(inputs):
       attn_output = jnp.swapaxes(attn_output_h, 0, -2)
     else:
       raise ValueError(
-          'Unsupported attention implementation: %s'
-          % self.attention_impl
+          'Unsupported attention implementation: %s' % self.attention_impl
       )
 
     # 6. Reshape and final projection
@@ -909,17 +1038,28 @@ def __call__(
       q: jt.Float[jax.Array | np.ndarray, 'B T E'],
       k: Optional[jt.Float[jax.Array | np.ndarray, 'B T_src E']] = None,
       v: Optional[jt.Float[jax.Array | np.ndarray, 'B T_src E']] = None,
-      attn_mask: Optional[jt.Bool[jax.Array | np.ndarray, 'B #N #T T_src']] = None,
+      attn_mask: Optional[
+          jt.Bool[jax.Array | np.ndarray, 'B #N #T T_src']
+      ] = None,
       rope: Optional[RotaryEmbedding] = None,
       *,
-      cached_kv: Optional[Tuple[jt.Float[jax.Array | np.ndarray, 'B T_src N D'],
-                                jt.Float[jax.Array | np.ndarray, 'B T_src N D']]] = None,
+      cached_kv: Optional[
+          Tuple[
+              jt.Float[jax.Array | np.ndarray, 'B T_src N D'],
+              jt.Float[jax.Array | np.ndarray, 'B T_src N D'],
+          ]
+      ] = None,
       return_kv: bool = False,
-  ) -> Union[jt.Float[jax.Array | np.ndarray, 'B T E'],   # output
-             Tuple[jt.Float[jax.Array | np.ndarray, 'B T E'],  # output
-                   Tuple[jt.Float[jax.Array | np.ndarray, 'B T_src N D'],  # key cache
-                         jt.Float[jax.Array | np.ndarray, 'B T_src N D']]]  # value cache
-             ]:
+  ) -> Union[
+      jt.Float[jax.Array | np.ndarray, 'B T E'],  # output
+      Tuple[
+          jt.Float[jax.Array | np.ndarray, 'B T E'],  # output
+          Tuple[
+              jt.Float[jax.Array | np.ndarray, 'B T_src N D'],  # key cache
+              jt.Float[jax.Array | np.ndarray, 'B T_src N D'],
+          ],
+      ],  # value cache
+  ]:
     """Forward pass through the attention block.
 
     Args:
@@ -1036,7 +1176,9 @@ def _induced_attention(
       src: jt.Float[jax.Array | np.ndarray, 'B T E'],
       train_size: Optional[jt.Int[jax.Array | np.ndarray, 'B']] = None,
       *,
-      cached_inducing_repr: Optional[jt.Float[jax.Array | np.ndarray, 'B I E']] = None,
+      cached_inducing_repr: Optional[
+          jt.Float[jax.Array | np.ndarray, 'B I E']
+      ] = None,
       return_inducing_repr: bool = False,
   ) -> Union[Array, Tuple[Array, Array]]:
     """Helper to run the two-stage attention with static masking."""
@@ -1084,13 +1226,17 @@ def __call__(
       src: jt.Float[jax.Array | np.ndarray, 'B T E'],
       train_size: Optional[jt.Int[jax.Array | np.ndarray, 'B']] = None,
       *,
-      cached_inducing_repr: Optional[jt.Float[jax.Array | np.ndarray, 'B I E']] = None,
+      cached_inducing_repr: Optional[
+          jt.Float[jax.Array | np.ndarray, 'B I E']
+      ] = None,
       return_inducing_repr: bool = False,
-  ) -> Union[jt.Float[jax.Array | np.ndarray, 'B T E'],
-             Tuple[jt.Float[jax.Array | np.ndarray, 'B T E'],   # output
-                   jt.Float[jax.Array | np.ndarray, 'B I E']    # cached_inducing_repr
-                  ]
-            ]:
+  ) -> Union[
+      jt.Float[jax.Array | np.ndarray, 'B T E'],
+      Tuple[
+          jt.Float[jax.Array | np.ndarray, 'B T E'],  # output
+          jt.Float[jax.Array | np.ndarray, 'B I E'],  # cached_inducing_repr
+      ],
+  ]:
     """Apply induced self-attention.
 
     Args:
@@ -1186,22 +1332,38 @@ def create_block(rngs):
   def __call__(
       self,
       src: jt.Float[jax.Array | np.ndarray, 'B T E'],
-      attn_mask: Optional[jt.Bool[jax.Array | np.ndarray, 'B #N #T T_prefill']] = None,
+      attn_mask: Optional[
+          jt.Bool[jax.Array | np.ndarray, 'B #N #T T_prefill']
+      ] = None,
       *,
       cached_kv: (
-          jt.Float[jax.Array | np.ndarray, 'Y B T_prefill E'] |         # For cache_icl_input_only
-          Tuple[jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D'],    # Key cache
-                jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D']] |  # Value cache
-          None
-      )=None,
+          jt.Float[
+              jax.Array | np.ndarray, 'Y B T_prefill E'
+          ]  # For cache_icl_input_only
+          | Tuple[
+              jt.Float[
+                  jax.Array | np.ndarray, 'Y B T_prefill N D'
+              ],  # Key cache
+              jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D'],
+          ]  # Value cache
+          | None
+      ) = None,
       return_kv: bool = False,
-  ) -> (jt.Float[jax.Array | np.ndarray, 'B T E'] |
-        Tuple[
-          jt.Float[jax.Array | np.ndarray, 'B T E'],            # Output
-          jt.Float[jax.Array | np.ndarray, 'Y B T_prefill E'] |         # For cache_icl_input_only
-          Tuple[jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D'],    # Key cache
-                jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D']]    # Value cache
-        ]):
+  ) -> (
+      jt.Float[jax.Array | np.ndarray, 'B T E']
+      | Tuple[
+          jt.Float[jax.Array | np.ndarray, 'B T E'],  # Output
+          jt.Float[
+              jax.Array | np.ndarray, 'Y B T_prefill E'
+          ]  # For cache_icl_input_only
+          | Tuple[
+              jt.Float[
+                  jax.Array | np.ndarray, 'Y B T_prefill N D'
+              ],  # Key cache
+              jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D'],
+          ],  # Value cache
+      ]
+  ):
     """Forward pass through the stacked encoder blocks.
 
     Args:
@@ -1250,7 +1412,9 @@ def scan_fn_cached_input(
         )
         @nnx.remat
         def scan_fn_cached(
-            block: MultiheadAttentionBlock, carry: jax.Array | np.ndarray, layer_kv
+            block: MultiheadAttentionBlock,
+            carry: jax.Array | np.ndarray,
+            layer_kv,
         ):
           out = block(
               q=carry,
@@ -1291,12 +1455,15 @@ def scan_fn_return_input(
         final_out, kvs = scan_fn_return_input(self.blocks, src)
         return final_out, kvs
       else:
+
         @nnx.scan(
             in_axes=(0, nnx.Carry),  # block, carry
             out_axes=(nnx.Carry, 0),  # carry, layer_kv
         )
         @nnx.remat
-        def scan_fn_return_kv(block: MultiheadAttentionBlock, carry: jax.Array | np.ndarray):
+        def scan_fn_return_kv(
+            block: MultiheadAttentionBlock, carry: jax.Array | np.ndarray
+        ):
           out, new_kv = block(
               q=carry,
               k=carry,
@@ -1369,6 +1536,7 @@ def __init__(
       raise ValueError(f'Activation must be one of {list(activations.keys())}')
 
     act_fn = activations[activation]
+
     @nnx.split_rngs(splits=num_blocks)
     @nnx.vmap(axis_size=num_blocks)
     def create_block(rngs):
@@ -1393,15 +1561,20 @@ def __call__(
       src: jt.Float[jax.Array | np.ndarray, 'B T E'],
       train_size: Optional[jt.Int[jax.Array | np.ndarray, 'B']] = None,
       *,
-      cached_inducing_repr: Optional[jt.Float[jax.Array | np.ndarray, 'Y B I E']] = None,
+      cached_inducing_repr: Optional[
+          jt.Float[jax.Array | np.ndarray, 'Y B I E']
+      ] = None,
       return_inducing_repr: bool = False,
-  ) -> Union[jt.Float[jax.Array | np.ndarray, 'B T E'],    # Output
-             # For cached_inducing_repr:
-             Tuple[
-                 # Output
-                 jt.Float[jax.Array | np.ndarray, 'B T E'],
-                 # Cache
-                 jt.Float[jax.Array | np.ndarray, 'Y B I E']]]:
+  ) -> Union[
+      jt.Float[jax.Array | np.ndarray, 'B T E'],  # Output
+      # For cached_inducing_repr:
+      Tuple[
+          # Output
+          jt.Float[jax.Array | np.ndarray, 'B T E'],
+          # Cache
+          jt.Float[jax.Array | np.ndarray, 'Y B I E'],
+      ],
+  ]:
     """Applies the Set Transformer to the input.
 
     Args:
@@ -1426,7 +1599,9 @@ def __call__(
       )
       @nnx.remat
       def scan_fn_cached(
-          block: InducedSelfAttentionBlock, carry: jax.Array | np.ndarray, layer_repr
+          block: InducedSelfAttentionBlock,
+          carry: jax.Array | np.ndarray,
+          layer_repr,
       ):
         out = block(
             carry,
@@ -1445,7 +1620,9 @@ def scan_fn_cached(
           out_axes=(nnx.Carry, 0),
       )
       @nnx.remat
-      def scan_fn_return(block: InducedSelfAttentionBlock, carry: jax.Array | np.ndarray):
+      def scan_fn_return(
+          block: InducedSelfAttentionBlock, carry: jax.Array | np.ndarray
+      ):
         out, repr = block(
             carry,
             train_size=train_size,
@@ -1463,10 +1640,10 @@ def scan_fn_return(block: InducedSelfAttentionBlock, carry: jax.Array | np.ndarr
           out_axes=nnx.Carry,
       )
       @nnx.remat
-      def scan_fn(block: InducedSelfAttentionBlock, carry: jax.Array | np.ndarray):
-        out = block(
-            carry, train_size=train_size
-        )
+      def scan_fn(
+          block: InducedSelfAttentionBlock, carry: jax.Array | np.ndarray
+      ):
+        out = block(carry, train_size=train_size)
         return out
 
       final_out = scan_fn(self.blocks, src)
@@ -1635,7 +1812,7 @@ def __init__(
   def feature_grouping(
       self,
       X: jt.Shaped[jax.Array | np.ndarray, 'B T H'],
-      d: Optional[jt.Int[jax.Array | np.ndarray, 'B']] = None
+      d: Optional[jt.Int[jax.Array | np.ndarray, 'B']] = None,
   ) -> jt.Shaped[jax.Array | np.ndarray, 'B T H G']:
     """Groups features with overlap using shifts.
 
@@ -1703,7 +1880,9 @@ def __call__(
       # features_expanded shape: (B, T, HC, in_dim)
       # fourier_frequencies shape: (in_dim, num_frequencies)
       x_proj = jnp.einsum(
-          '...i,if->...if', features_expanded_raw, self.fourier_frequencies.value
+          '...i,if->...if',
+          features_expanded_raw,
+          self.fourier_frequencies.value,
       )
       # Always keep fourier_feats in per-slot shape: (B, T, HC, G, num_freq*2).
       # When feature_group=False, feature_grouping returns (B, T, HC, 1), so G=1
@@ -1719,7 +1898,9 @@ def __call__(
       if cat_mask is not None:
         # Group cat_mask exactly like features (X) were grouped.
         # Expand to (B, 1, HC) so feature_grouping sees shape (B, T, HC).
-        cat_mask_grouped = self.feature_grouping(cat_mask[:, None, :], d=d)  # (B, 1, HC, G)
+        cat_mask_grouped = self.feature_grouping(
+            cat_mask[:, None, :], d=d
+        )  # (B, 1, HC, G)
         # Broadcast to (B, 1, HC, G, 1) to match (B, T, HC, G, E)
         cat_mask_per_slot = cat_mask_grouped[:, :, :, :, None]
 
@@ -1757,7 +1938,9 @@ def __call__(
         # TODO: Try using Fourier features for the continuous y embedding (as
         # done for X in feature_grouping) instead of a plain MLP, so the model
         # gets the same frequency-rich representation for target values.
-        y_embedded: Float[Array, '... T E'] = self.y_embedder_lookup(y[..., None])
+        y_embedded: Float[Array, '... T E'] = self.y_embedder_lookup(
+            y[..., None]
+        )
 
       if train_size is not None:
         # Create a mask for the training samples.
@@ -1867,12 +2050,17 @@ def __call__(
       train_size: Optional[jt.Int[jax.Array | np.ndarray, 'B']] = None,
       feature_shuffles: Optional[List[List[int]]] = None,
       *,
-      cached_repr: Optional[jt.Float[jax.Array | np.ndarray, 'Y B*H I E']] = None,
+      cached_repr: Optional[
+          jt.Float[jax.Array | np.ndarray, 'Y B*H I E']
+      ] = None,
       return_repr: bool = False,
-  ) -> Union[jt.Float[jax.Array | np.ndarray, 'B T H E'],  # Output
-             Tuple[jt.Float[jax.Array | np.ndarray, 'B T H E'], # Output
-                   jt.Float[jax.Array | np.ndarray, 'Y B*H I E']  # cached_repr
-                  ]]:
+  ) -> Union[
+      jt.Float[jax.Array | np.ndarray, 'B T H E'],  # Output
+      Tuple[
+          jt.Float[jax.Array | np.ndarray, 'B T H E'],  # Output
+          jt.Float[jax.Array | np.ndarray, 'Y B*H I E'],  # cached_repr
+      ],
+  ]:
     """Transform input table into embeddings.
 
     Args:
@@ -1885,15 +2073,15 @@ def __call__(
     Returns:
       Embeddings or (embeddings, new_repr).
     """
-    assert not (cached_repr is not None and return_repr), (
-        'Cannot have both cached_repr not None and return_repr set to True.'
-    )
+    assert not (
+        cached_repr is not None and return_repr
+    ), 'Cannot have both cached_repr not None and return_repr set to True.'
     B, T, HC, E = X.shape
     padded_train_size = T
     if cached_repr is not None:
-      assert train_size is None, (
-          'train_size must be None when cached_repr is not None.'
-      )
+      assert (
+          train_size is None
+      ), 'train_size must be None when cached_repr is not None.'
 
     assert padded_train_size is not None
     train_size_expanded = None
@@ -1913,7 +2101,7 @@ def __call__(
     new_repr: Optional[Array] = None
 
     tf_col_st = typing.cast(SetTransformer, self.tf_col)
-    if cached_repr is not None:    # Decode.
+    if cached_repr is not None:  # Decode.
       representations = typing.cast(
           Array,
           tf_col_st(
@@ -1923,13 +2111,13 @@ def __call__(
               cached_inducing_repr=cached_repr,
           ),
       )
-    elif return_repr:             # Prefill.
+    elif return_repr:  # Prefill.
       representations, new_repr = tf_col_st(
           src,
           train_size=train_size_expanded,
           return_inducing_repr=True,
       )
-    else:                        # Train.
+    else:  # Train.
       representations = typing.cast(
           Array,
           tf_col_st(
@@ -1952,8 +2140,10 @@ def __call__(
       return final_embeddings, typing.cast(Array, new_repr)
     return final_embeddings
 
+
 # Row-wise embedding
 
+
 class RowInteraction(nnx.Module):
   """Context-aware row-wise interaction, rewritten in Flax NNX.
 
@@ -2008,8 +2198,10 @@ def __call__(
       self,
       embeddings: jt.Float[jax.Array | np.ndarray, 'B T H E'],
       d: Optional[jt.Int[jax.Array | np.ndarray, 'B']] = None,
-  ) -> (jt.Float[jax.Array | np.ndarray, 'B T H E']
-        | jt.Float[jax.Array | np.ndarray, 'B T C_TIMES_E']): # if output_full_sequence is False.
+  ) -> (
+      jt.Float[jax.Array | np.ndarray, 'B T H E']
+      | jt.Float[jax.Array | np.ndarray, 'B T C_TIMES_E']
+  ):  # if output_full_sequence is False.
     """Captures interactions between features within each row.
 
     Args:
@@ -2061,7 +2253,9 @@ def __call__(
 @nnx.dataclass
 class ICLearningCache:
   layer_caches: (
-      jt.Float[jax.Array | np.ndarray, 'Y B T_prefill E']  # For cache_icl_input_only
+      jt.Float[
+          jax.Array | np.ndarray, 'Y B T_prefill E'
+      ]  # For cache_icl_input_only
       | Tuple[
           jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D'],  # Key cache
           jt.Float[jax.Array | np.ndarray, 'Y B T_prefill N D'],  # Value cache
@@ -2095,6 +2289,7 @@ class ICLearning(nnx.Module):
   rngs : nnx.Rngs, optional
       RNGs for initialization.
   """
+
   @jt.typed
   def __init__(
       self,
@@ -2168,9 +2363,7 @@ def __init__(
       )
 
   @jt.typed
-  def _prefill_sequence_length_from_cache(
-      self, cache: ICLearningCache
-  ) -> Any:
+  def _prefill_sequence_length_from_cache(self, cache: ICLearningCache) -> Any:
     """Returns the sequence length of the prefill cache."""
     assert cache.layer_caches is not None, 'Layer caches must be non-empty.'
     if self.cache_icl_input_only:
@@ -2190,11 +2383,15 @@ def __call__(
       *,
       cache: Optional[ICLearningCache] = None,
       return_cache: bool = False,
-  ) -> (jt.Float[jax.Array | np.ndarray, 'B T K'] | jt.Float[jax.Array | np.ndarray, 'B T E']
-        | Tuple[
-            jt.Float[jax.Array | np.ndarray, 'B T K'] | jt.Float[jax.Array | np.ndarray, 'B T E'],
-            ICLearningCache
-          ]):
+  ) -> (
+      jt.Float[jax.Array | np.ndarray, 'B T K']
+      | jt.Float[jax.Array | np.ndarray, 'B T E']
+      | Tuple[
+          jt.Float[jax.Array | np.ndarray, 'B T K']
+          | jt.Float[jax.Array | np.ndarray, 'B T E'],
+          ICLearningCache,
+      ]
+  ):
     """Forward pass for ICLearning."""
     is_prefill = return_cache
     is_decode = cache is not None
@@ -2216,9 +2413,7 @@ def __call__(
       assert train_size is not None
       if train_size.ndim == 2:
         train_size = jnp.squeeze(train_size, axis=-1)
-      train_mask = (
-          jnp.arange(sequence_length)[None, :] < train_size[:, None]
-      )
+      train_mask = jnp.arange(sequence_length)[None, :] < train_size[:, None]
       full_attn_mask = train_mask[:, None, None, :]
       y_encoded = y_encoded * train_mask.reshape(B, sequence_length, 1)
       R = R + y_encoded
@@ -2227,7 +2422,8 @@ def __call__(
       prefill_train_size = cache.prefill_train_size
       prefill_sequence_length = self._prefill_sequence_length_from_cache(cache)
       train_mask = (
-          jnp.arange(prefill_sequence_length)[None, :] < prefill_train_size[:, None]
+          jnp.arange(prefill_sequence_length)[None, :]
+          < prefill_train_size[:, None]
       )
       full_attn_mask = train_mask[:, None, None, :]
 
@@ -2239,8 +2435,7 @@ def __call__(
           R, attn_mask=full_attn_mask, return_kv=True
       )
       new_cache = ICLearningCache(
-          layer_caches=new_layer_caches,
-          prefill_train_size=train_size
+          layer_caches=new_layer_caches, prefill_train_size=train_size
       )
     else:
       assert is_decode
@@ -2249,8 +2444,6 @@ def __call__(
           R, attn_mask=full_attn_mask, cached_kv=cache.layer_caches
       )
 
-
-
     result = self.ln(result)
     result = self.decoder(result)
     if return_cache:
@@ -2258,8 +2451,10 @@ def __call__(
       return result, new_cache
     return result
 
+
 # TabFM
 
+
 class TabFM(nnx.Module):
   """A Tabular Foundation Model (TabFM), rewritten in Flax NNX.
 
@@ -2484,13 +2679,22 @@ def __call__(
       self,
       X: jt.Float[jax.Array | np.ndarray, 'B T H'],
       y: jt.Shaped[Array, 'B T'],
-      train_size: jt.Int[jax.Array | np.ndarray, 'B'] | jt.Int[jax.Array | np.ndarray, 'B 1'],
-      d: jt.Int[jax.Array | np.ndarray, 'B'] | jt.Int[jax.Array | np.ndarray, 'B 1'] | None = None,
+      train_size: (
+          jt.Int[jax.Array | np.ndarray, 'B']
+          | jt.Int[jax.Array | np.ndarray, 'B 1']
+      ),
+      d: (
+          jt.Int[jax.Array | np.ndarray, 'B']
+          | jt.Int[jax.Array | np.ndarray, 'B 1']
+          | None
+      ) = None,
       cat_mask: Optional[jt.Bool[jax.Array | np.ndarray, 'B H']] = None,
       softmax_temperature: float = 0.9,
       num_classes: Optional[int] = None,
-  ) -> (jt.Float[jax.Array | np.ndarray, 'B T E']      # Regression
-        | jt.Float[jax.Array | np.ndarray, 'B T K']):  # Classification
+  ) -> (
+      jt.Float[jax.Array | np.ndarray, 'B T E']  # Regression
+      | jt.Float[jax.Array | np.ndarray, 'B T K']
+  ):  # Classification
     """Processes tabular data through nested encoders and ICL predictor.
 
     Args:
@@ -2523,7 +2727,6 @@ def __call__(
         X, y, train_size=train_size, d=d, cat_mask=cat_mask
     )  # (B, T, HC, E)
 
-
     # Column-wise embedding
     embeddings = typing.cast(
         Array,
@@ -2538,16 +2741,13 @@ def __call__(
         self.cls_tokens[...], (B1, T1, self.row_num_cls, self.embed_dim)
     )
 
-    embeddings = jnp.concatenate(
-        [cls_tokens_expanded, embeddings], axis=-2
-    )
+    embeddings = jnp.concatenate([cls_tokens_expanded, embeddings], axis=-2)
 
     embeddings = self.row_interactor(
         embeddings,
         d=d,
     )
 
-
     embeddings = typing.cast(
         Array,
         self.col_embedder_2(
@@ -2586,10 +2786,13 @@ def prefill(
       y: jt.Shaped[jax.Array | np.ndarray, 'B T'],
       d: jt.Int[jax.Array | np.ndarray, 'B'] = None,
       cat_mask: jt.Bool[jax.Array | np.ndarray, 'B H'] | None = None,
-  ) -> (
-      Tuple[(jt.Float[jax.Array | np.ndarray, 'B T K'] | jt.Float[jax.Array | np.ndarray, 'B T E']),
-             Dict[str, Any] # cache
-      ]):
+  ) -> Tuple[
+      (
+          jt.Float[jax.Array | np.ndarray, 'B T K']
+          | jt.Float[jax.Array | np.ndarray, 'B T E']
+      ),
+      Dict[str, Any],  # cache
+  ]:
     """Prefills the model with training data and returns the cache.
 
     Args:
@@ -2629,7 +2832,6 @@ def prefill(
         X, y, train_size=train_size, d=d, cat_mask=cat_mask
     )
 
-
     # Stage 1: Column-wise embedding
     res = typing.cast(
         Tuple[jnp.ndarray, Array],
@@ -2641,7 +2843,6 @@ def prefill(
     )
     embeddings, cache_col1 = res
 
-
     # Prepend CLS tokens to embeddings before row interactor
     B1, T1, _, _ = embeddings.shape
     cls_tokens_expanded = jnp.broadcast_to(
@@ -2649,9 +2850,7 @@ def prefill(
     )
     embeddings = jnp.concatenate([cls_tokens_expanded, embeddings], axis=-2)
 
-    embeddings = self.row_interactor(
-        embeddings, d=d
-    )
+    embeddings = self.row_interactor(embeddings, d=d)
 
     res2 = typing.cast(
         Tuple[jnp.ndarray, Array],
@@ -2663,9 +2862,7 @@ def prefill(
     )
     embeddings, cache_col2 = res2
 
-    representations = self.row_interactor_2(
-        embeddings, d=d
-    )
+    representations = self.row_interactor_2(embeddings, d=d)
 
     # Stage 3: ICL
     logits_icl, cache_icl = self.icl_predictor(
@@ -2692,7 +2889,10 @@ def decode(
       cat_mask: Optional[jt.Bool[jax.Array | np.ndarray, 'B H']] = None,
       softmax_temperature: float = 0.9,
       num_classes: Optional[int] = None,
-  ) -> (jt.Float[jax.Array | np.ndarray, 'B T K'] | jt.Float[jax.Array | np.ndarray, 'B T E']):
+  ) -> (
+      jt.Float[jax.Array | np.ndarray, 'B T K']
+      | jt.Float[jax.Array | np.ndarray, 'B T E']
+  ):
     """Generates predictions for test data using cached KVs.
 
     Args:
@@ -2722,14 +2922,14 @@ def decode(
 
     # Stage 0: Cell-wise embedding for test rows
     cell_embeddings = self.cell_embedder(
-        X, y,
+        X,
+        y,
         # train_size is 0 for the current batch (all are test rows relative to the cache)
-        train_size=jnp.zeros((B,),
-                             dtype=jnp.int32),
-        d=d, cat_mask=cat_mask
+        train_size=jnp.zeros((B,), dtype=jnp.int32),
+        d=d,
+        cat_mask=cat_mask,
     )
 
-
     # Stage 1: Column-wise embedding with cache
     embeddings = typing.cast(
         Array,
@@ -2740,7 +2940,6 @@ def decode(
         ),
     )
 
-
     # Prepend CLS tokens
     B1, T1, _, _ = embeddings.shape
     cls_tokens_expanded = jnp.broadcast_to(
@@ -2748,9 +2947,7 @@ def decode(
     )
     embeddings = jnp.concatenate([cls_tokens_expanded, embeddings], axis=-2)
 
-    embeddings = self.row_interactor(
-        embeddings, d=d
-    )
+    embeddings = self.row_interactor(embeddings, d=d)
 
     embeddings = typing.cast(
         Array,
@@ -2762,9 +2959,7 @@ def decode(
     )
 
     # Stage 2: Row-wise interaction
-    representations = self.row_interactor_2(
-        embeddings, d=d
-    )
+    representations = self.row_interactor_2(embeddings, d=d)
 
     # Stage 3: ICL with cache
     out = self.icl_predictor(
diff --git a/tabfm/src/jax/seqpar.py b/tabfm/src/jax/seqpar.py
new file mode 100644
index 0000000..c3ff1db
--- /dev/null
+++ b/tabfm/src/jax/seqpar.py
@@ -0,0 +1,375 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sequence-parallel (row-sharded) inference for the JAX TabFM backend.
+
+Mirrors ``tabfm.src.pytorch.seqpar``: the in-context rows of each ensemble
+member are sharded across the devices of a mesh with ``jax.shard_map``, so
+training folds that exceed a single device's memory can be used as context.
+Weights are replicated; attention is exact (no approximation):
+
+  * Cell embedding, row interaction (attention over features), feed-forward
+    blocks and the decoder are row-independent and run locally per shard,
+    reusing the stock nnx modules.
+  * The column set-transformer's inducing attention runs an fp32 online
+    softmax over each device's context shard (scanned in key chunks so no
+    full-length score matrix materializes) and combines the partial results
+    exactly across devices via log-sum-exp weights.
+  * ICL self-attention all-gathers each device's projected context K/V (test
+    rows are never keys) and runs the chunked memory-efficient attention
+    locally per device, with padded key slots masked via an additive bias.
+
+Typical use (single process driving all local devices)::
+
+    model = tabfm_v1_0_0.load(model_type="regression")
+    reg = TabFMRegressor(model=model, n_estimators=4)
+    reg.fit(X_train, y_train)              # cheap; no device forward
+    preds = seqpar.predict(reg, X_test)    # sharded over jax.devices()
+
+``predict`` / ``predict_proba`` run every ensemble member through the sharded
+forward and apply the same ensemble combination as the estimators' own
+prediction paths. Device blocks are padded to a common 128-multiple length
+(the chunked-attention granularity), so unequal shards are supported; the
+padded rows are masked out of every attention and trimmed from the output.
+
+Multi-process (e.g. multi-host TPU slice) runs are supported: initialize the
+runtime with ``jax.distributed.initialize()`` first, then run the same
+program on every process -- ``fit`` (deterministic given ``random_state``)
+and ``predict`` are called on all processes, and every process returns the
+full predictions. Input shards are assembled per process with
+``jax.make_array_from_callback`` and the (tiny) output is re-replicated
+across devices so it is readable everywhere.
+"""
+
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+from jax.sharding import NamedSharding
+from jax.sharding import PartitionSpec as P
+
+from flax import nnx
+
+from tabfm.src.jax import memory_efficient_attention as mea
+
+_AXIS = "seqpar_rows"
+_PAD = 128  # chunked-attention length granularity
+_MAB1_KEY_CHUNK = 2048  # keys per online-softmax step (multiple of _PAD)
+_ROW_CHUNK = 32768  # rows per slice for the local mab2 / stage-tail calls
+
+
+def _round_up(n, k):
+  return ((n + k - 1) // k) * k
+
+
+def _shard_bounds(n, world, rank):
+  """Contiguous near-equal split of ``n`` items; returns (start, stop)."""
+  base, rem = divmod(n, world)
+  start = rank * base + min(rank, rem)
+  return start, start + base + (1 if rank < rem else 0)
+
+
+def _iter_blocks(stacked):
+  """Yields per-layer modules from an ``nnx.vmap``-stacked block module."""
+  if hasattr(stacked, "__iter__"):
+    yield from stacked
+    return
+  graphdef, state = nnx.split(stacked)
+  n = jax.tree.leaves(state)[0].shape[0]
+  for i in range(n):
+    yield nnx.merge(graphdef, jax.tree.map(lambda a, i=i: a[i], state))
+
+
+def _row_chunked(fn, src):
+  """Applies a row-independent fn over row slices of ``[B, T, E]``."""
+  outs = [
+      fn(src[:, s : s + _ROW_CHUNK]) for s in range(0, src.shape[1], _ROW_CHUNK)
+  ]
+  return outs[0] if len(outs) == 1 else jnp.concatenate(outs, axis=1)
+
+
+def _mab1_combined(mab, ind, k_src, ts_local):
+  """mab1 (inducing queries over all context rows) with sharded keys.
+
+  ``k_src``: [B, c_blk, E] local context block; rows at positions >=
+  ``ts_local`` are padding and masked out. Online softmax over key chunks
+  (fp32 accumulators), then an exact cross-device log-sum-exp combine. Only
+  softmax statistics (whose shapes do not depend on the key count) cross
+  devices, so unequal shards work natively.
+  """
+  attn = mab.attn
+  b, c_blk, e = k_src.shape
+  ni = ind.shape[0]
+  nh, hd = attn.num_heads, attn.head_dim
+  f32 = jnp.float32
+
+  q0 = jnp.broadcast_to(ind, (b, ni, e))
+  qn = mab.pre_attn_ln(q0)
+  q = attn.q_proj(qn).reshape(b, ni, nh, hd)
+  q = attn.per_dim_scale(attn.query_ln(q))
+  qf = jnp.einsum("bind->bnid", q.astype(f32))  # [b, nh, I, hd]
+
+  n_chunks = c_blk // _MAB1_KEY_CHUNK
+  k_chunks = k_src.reshape(b, n_chunks, _MAB1_KEY_CHUNK, e).transpose(
+      (1, 0, 2, 3)
+  )
+
+  def step(carry, kc_and_idx):
+    m, den, num = carry
+    kc, idx = kc_and_idx
+    kn = mab.pre_attn_ln(kc)  # [b, KC, e]
+    k = attn.key_ln(attn.k_proj(kn).reshape(b, -1, nh, hd))
+    v = attn.v_proj(kn).reshape(b, -1, nh, hd)
+    s = jnp.einsum("bnid,bjnd->bnij", qf, k.astype(f32))  # [b,nh,I,KC]
+    pos = idx * _MAB1_KEY_CHUNK + jnp.arange(_MAB1_KEY_CHUNK)
+    s = jnp.where(pos[None, None, None, :] < ts_local, s, -jnp.inf)
+    m_new = jnp.maximum(m, s.max(-1))
+    alpha = jnp.exp(m - m_new)
+    p = jnp.exp(s - m_new[..., None])
+    den = den * alpha + p.sum(-1)
+    num = num * alpha[..., None] + jnp.einsum(
+        "bnij,bjnd->bnid", p, v.astype(f32)
+    )
+    return (m_new, den, num), None
+
+  init = (
+      jnp.full((b, nh, ni), -jnp.inf),
+      jnp.zeros((b, nh, ni)),
+      jnp.zeros((b, nh, ni, hd)),
+  )
+  (m, den, num), _ = jax.lax.scan(step, init, (k_chunks, jnp.arange(n_chunks)))
+
+  # Exact cross-device combine of the online-softmax statistics.
+  m_glob = jax.lax.pmax(m, _AXIS)
+  scale = jnp.exp(m - m_glob)
+  den = jax.lax.psum(den * scale, _AXIS)
+  num = jax.lax.psum(num * scale[..., None], _AXIS)
+  out = (num / den[..., None]).astype(q0.dtype)  # [b, nh, I, hd]
+
+  out = jnp.einsum("bnid->bind", out).reshape(b, ni, nh * hd)
+  x = q0 + mab.post_attn_ln(attn.out_proj(out))
+  return x + mab._ff_block(x)  # pylint: disable=protected-access
+
+
+def _col_sharded(col, emb, c_blk, ts_local):
+  """ColEmbedding with the row axis sharded. ``emb``: [B, T_local, HC, E]."""
+  b, t, hc, e = emb.shape
+  src = emb.transpose((0, 2, 1, 3)).reshape(b * hc, t, e)
+  for blk in _iter_blocks(col.tf_col.blocks):
+    hidden = _mab1_combined(
+        blk.mab1, blk.ind_vectors[...], src[:, :c_blk], ts_local
+    )
+    # mab2 rows attend only to the replicated inducing outputs: local.
+    src = _row_chunked(lambda s: blk.mab2(q=s, k=hidden, v=hidden), src)  # pylint: disable=cell-var-from-loop
+  out = _row_chunked(lambda s: col.ln_w(col.out_w(s)), src)
+  return out.reshape(b, hc, t, e).transpose((0, 2, 1, 3))
+
+
+def _icl_block_sharded(blk, r, c_blk, key_bias):
+  """One ICL self-attention block; keys = all devices' context blocks."""
+  attn = blk.attn
+  nh, hd = attn.num_heads, attn.head_dim
+  xn = blk.pre_attn_ln(r)
+  q = attn.q_proj(xn).reshape(1, -1, nh, hd)
+  kn = xn[:, :c_blk]
+  k = attn.key_ln(attn.k_proj(kn).reshape(1, c_blk, nh, hd))
+  v = attn.v_proj(kn).reshape(1, c_blk, nh, hd)
+  q = attn.per_dim_scale(attn.query_ln(q))
+  key = jax.lax.all_gather(k, _AXIS, axis=1, tiled=True)
+  val = jax.lax.all_gather(v, _AXIS, axis=1, tiled=True)
+  tgt = q.shape[1]
+  ao = mea.dot_product_attention_multihead(
+      query=q,
+      key=key,
+      value=val,
+      bias=key_bias,
+      dtype=np.dtype(r.dtype.name),
+      enable_dropout=False,
+      query_chunk_size=_PAD if tgt >= _PAD else tgt,
+      key_chunk_size=_PAD,
+  )
+  o = attn.out_proj(ao.reshape(1, -1, nh * hd))
+  x = r + blk.post_attn_ln(o)
+  return x + blk._ff_block(x)  # pylint: disable=protected-access
+
+
+def _make_forward(graphdef, c_blk, t_blk, has_cat, has_d):
+  """Builds the per-device shard_map body for one ensemble member."""
+
+  def forward(state, x, y, ts, cat_mask, d):
+    m = nnx.merge(graphdef, state)
+    dtype = m.dtype
+    x = jnp.nan_to_num(x, nan=-100.0).astype(dtype)
+    y = y.astype(dtype)
+    cm = cat_mask if has_cat else None
+    dd = d if has_d else None
+    emb = m.cell_embedder(x, y, train_size=ts, d=dd, cat_mask=cm)
+    emb = _col_sharded(m.col_embedder, emb, c_blk, ts[0])
+    b1, t1 = emb.shape[:2]
+    cls = jnp.broadcast_to(
+        m.cls_tokens[...], (b1, t1, m.row_num_cls, m.embed_dim)
+    )
+    emb = jnp.concatenate([cls, emb], axis=-2)
+    emb = m.row_interactor(emb, d=dd)
+    emb = _col_sharded(m.col_embedder_2, emb, c_blk, ts[0])
+    reps = m.row_interactor_2(emb, d=dd)
+
+    icl = m.icl_predictor
+    if icl.is_classifier:
+      y_enc = icl.y_encoder(y.astype(jnp.int32))
+    else:
+      y_enc = icl.y_encoder(y[..., None])
+    tmask = jnp.arange(reps.shape[1])[None, :] < ts[:, None]
+    r = reps + y_enc * tmask[..., None]
+
+    # Key-validity bias for the gathered context blocks of every device.
+    ts_all = jax.lax.all_gather(ts[0], _AXIS)  # [world]
+    valid = (jnp.arange(c_blk)[None, :] < ts_all[:, None]).reshape(-1)
+    key_bias = jnp.where(valid, 0.0, -1e30)[None, None, None, :]
+    for blk in _iter_blocks(icl.tf_icl.blocks):
+      r = _icl_block_sharded(blk, r, c_blk, key_bias)
+    out = icl.decoder(icl.ln(r))  # [1, T_local, L]
+    return out[:, c_blk:, :]  # [1, t_blk, L]
+
+  return forward
+
+
+def _member_outputs(estimator, X, mesh):
+  """Runs every ensemble member through the sharded forward.
+
+  Returns ``[n_members, n_test, L_out]`` float32 outputs.
+  """
+  x_enc = estimator.X_encoder_.transform(X)
+  data = estimator.ensemble_generator_.transform(x_enc)
+  xs, ys, cat_masks, ds, _ = (
+      estimator.ensemble_generator_.prepare_ensemble_tensors(data)
+  )
+  n_members = xs.shape[0]
+  n_train = ys.shape[1]
+  n_test = xs.shape[1] - n_train
+  h = xs.shape[-1]
+  world = mesh.devices.size
+  if n_train < world:
+    raise ValueError(
+        f"n_train ({n_train}) must be >= the device count ({world})."
+    )
+
+  bounds_c = [_shard_bounds(n_train, world, r) for r in range(world)]
+  bounds_t = [_shard_bounds(n_test, world, r) for r in range(world)]
+  c_blk = _round_up(max(c1 - c0 for c0, c1 in bounds_c), _MAB1_KEY_CHUNK)
+  t_blk = max(_round_up(max(t1 - t0 for t0, t1 in bounds_t), _PAD), _PAD)
+
+  graphdef, state = nnx.split(estimator.model)
+  state = jax.device_put(state, NamedSharding(mesh, P()))
+  has_cat = cat_masks is not None
+  has_d = ds is not None
+  fwd = _make_forward(graphdef, c_blk, t_blk, has_cat, has_d)
+  # The output is re-replicated across all devices so that every process of a
+  # multi-process (e.g. multi-host TPU) run can read it back directly; the
+  # gathered tensor is tiny ([1, W * t_blk, L]).
+  sharded = jax.jit(
+      jax.shard_map(
+          fwd,
+          mesh=mesh,
+          in_specs=(
+              P(),
+              P(None, _AXIS, None),
+              P(None, _AXIS),
+              P(_AXIS),
+              P(),
+              P(),
+          ),
+          out_specs=P(None, _AXIS, None),
+          check_vma=False,
+      ),
+      out_shardings=NamedSharding(mesh, P()),
+  )
+
+  x_shard = NamedSharding(mesh, P(None, _AXIS, None))
+  y_shard = NamedSharding(mesh, P(None, _AXIS))
+  ts_shard = NamedSharding(mesh, P(_AXIS))
+  repl = NamedSharding(mesh, P())
+
+  def to_global(arr, sharding):
+    # Every process holds the full host copy (fit is deterministic, so all
+    # processes computed identical tensors); each contributes the slices its
+    # addressable devices need. Works identically in single-process runs.
+    return jax.make_array_from_callback(
+        arr.shape, sharding, lambda idx: arr[idx]
+    )
+
+  outs = []
+  for mi in range(n_members):
+    xg = np.zeros((1, world * (c_blk + t_blk), h), np.float32)
+    yg = np.full((1, world * (c_blk + t_blk)), -100.0, np.float32)
+    ts_g = np.zeros((world,), np.int32)
+    for r, ((c0, c1), (t0, t1)) in enumerate(zip(bounds_c, bounds_t)):
+      base = r * (c_blk + t_blk)
+      xg[0, base : base + c1 - c0] = xs[mi, c0:c1]
+      yg[0, base : base + c1 - c0] = ys[mi, c0:c1]
+      xg[0, base + c_blk : base + c_blk + t1 - t0] = xs[
+          mi, n_train + t0 : n_train + t1
+      ]
+      ts_g[r] = c1 - c0
+    args = (
+        state,
+        to_global(xg, x_shard),
+        to_global(yg, y_shard),
+        to_global(ts_g, ts_shard),
+        to_global(
+            np.asarray(cat_masks[mi : mi + 1])
+            if has_cat
+            else np.zeros((1, h), bool),
+            repl,
+        ),
+        to_global(
+            np.asarray(ds[mi : mi + 1], np.int32)
+            if has_d
+            else np.zeros((1,), np.int32),
+            repl,
+        ),
+    )
+    out = np.asarray(
+        jax.block_until_ready(sharded(*args)), np.float32
+    )  # [1, W*t_blk, L]; fully replicated, so readable from any process
+    parts = [
+        out[0, r * t_blk : r * t_blk + (t1 - t0)]
+        for r, (t0, t1) in enumerate(bounds_t)
+    ]
+    outs.append(np.concatenate(parts, axis=0))
+  return np.stack(outs, axis=0)
+
+
+def _default_mesh():
+  return jax.make_mesh((len(jax.devices()),), (_AXIS,))
+
+
+def predict(estimator, X, mesh=None):
+  """Sharded equivalent of ``TabFMRegressor.predict``."""
+  outputs = _member_outputs(estimator, X, mesh or _default_mesh())
+  return estimator._combine_predictions(outputs.squeeze(-1))  # pylint: disable=protected-access
+
+
+def predict_proba(estimator, X, mesh=None):
+  """Sharded equivalent of ``TabFMClassifier.predict_proba``."""
+  outputs = _member_outputs(estimator, X, mesh or _default_mesh())
+  outputs = outputs[..., : estimator.n_classes_]
+  offsets = []
+  for offs in estimator.ensemble_generator_.class_shift_offsets_.values():
+    offsets.extend(offs)
+  logits = np.stack([
+      np.concatenate([out[..., off:], out[..., :off]], axis=-1)
+      for out, off in zip(outputs, offsets)
+  ])
+  return estimator._process_logits(logits)  # pylint: disable=protected-access
diff --git a/tabfm/src/jax/seqpar_multiprocess_test.py b/tabfm/src/jax/seqpar_multiprocess_test.py
new file mode 100644
index 0000000..f248982
--- /dev/null
+++ b/tabfm/src/jax/seqpar_multiprocess_test.py
@@ -0,0 +1,145 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Multi-process tests for sequence-parallel JAX inference.
+
+Rehearses the multi-host execution model (e.g. a multi-host TPU slice) on
+CPU: two OS processes are spawned, each with two simulated host devices
+(four global devices), joined via ``jax.distributed.initialize``. Every
+process runs the same fit + sharded predict, mirroring how every host of a
+TPU slice runs the same program; the results are compared against the
+estimator's plain single-process prediction path.
+"""
+
+import multiprocessing
+import os
+
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=2")
+
+from absl.testing import absltest
+import numpy as np
+
+try:
+  import jax
+
+  # Full-precision fp32 matmuls: the parent-process reference may run on GPU
+  # dev machines, where TF32 would exceed the comparison tolerance.
+  jax.config.update("jax_default_matmul_precision", "highest")
+  HAS_JAX = True
+except ImportError:
+  HAS_JAX = False
+
+_NPROCS = 2
+
+
+def _make_fitted_estimator(is_classifier):
+  """Builds a tiny fitted estimator + test rows. Deterministic everywhere.
+
+  Imports are lazy so spawned workers can configure JAX (platform, device
+  count, distributed runtime) before its first use.
+  """
+  from flax import nnx  # pylint: disable=g-import-not-at-top
+  import jax.numpy as jnp  # pylint: disable=g-import-not-at-top
+  from tabfm.src.classifier_and_regressor import TabFMClassifier  # pylint: disable=g-import-not-at-top
+  from tabfm.src.classifier_and_regressor import TabFMRegressor  # pylint: disable=g-import-not-at-top
+  from tabfm.src.jax import model as tabfm_model  # pylint: disable=g-import-not-at-top
+
+  model = tabfm_model.TabFM(
+      loss="cross_entropy" if is_classifier else "mse",
+      max_classes=2,
+      embed_dim=8,
+      col_num_blocks=1,
+      col_nhead=2,
+      col_num_inds=8,
+      row_num_blocks=1,
+      row_nhead=2,
+      row_num_cls=1,
+      icl_num_blocks=1,
+      icl_nhead=2,
+      dtype=jnp.float32,
+      rngs=nnx.Rngs(0),
+  )
+  n_train, n_test = 21, 9
+  rng = np.random.RandomState(0)
+  X = rng.rand(n_train + n_test, 4)
+  if is_classifier:
+    y = rng.randint(0, 2, size=n_train)
+    est = TabFMClassifier(model=model, n_estimators=2, random_state=0)
+  else:
+    y = X[:n_train] @ rng.rand(4)
+    est = TabFMRegressor(model=model, n_estimators=2, random_state=0)
+  est.fit(X[:n_train], y)
+  return est, X[n_train:]
+
+
+def _worker(pid, port, is_classifier, q):
+  """One simulated host: joins the distributed runtime, runs the predict."""
+  os.environ["JAX_PLATFORMS"] = "cpu"
+  os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=2"
+  import jax  # pylint: disable=g-import-not-at-top
+
+  jax.distributed.initialize(
+      coordinator_address=f"localhost:{port}",
+      num_processes=_NPROCS,
+      process_id=pid,
+  )
+  assert jax.process_count() == _NPROCS
+  assert len(jax.devices()) == 2 * _NPROCS
+
+  from tabfm.src.jax import seqpar  # pylint: disable=g-import-not-at-top
+
+  est, x_test = _make_fitted_estimator(is_classifier)
+  if is_classifier:
+    out = seqpar.predict_proba(est, x_test)
+  else:
+    out = seqpar.predict(est, x_test)
+  q.put((pid, out))  # every process returns the full predictions
+
+
+@absltest.skipIf(not HAS_JAX, "JAX not installed")
+class SeqparMultiprocessTest(absltest.TestCase):
+
+  def _run(self, is_classifier):
+    ctx = multiprocessing.get_context("spawn")
+    q = ctx.Queue()
+    port = 21212 + int(is_classifier)
+    procs = [
+        ctx.Process(target=_worker, args=(pid, port, is_classifier, q))
+        for pid in range(_NPROCS)
+    ]
+    for p in procs:
+      p.start()
+    results = dict(q.get(timeout=300) for _ in range(_NPROCS))
+    for p in procs:
+      p.join(timeout=300)
+      self.assertEqual(p.exitcode, 0)
+
+    # Reference: the estimator's plain (non-distributed) prediction path,
+    # computed in this parent process.
+    est, x_test = _make_fitted_estimator(is_classifier)
+    ref = est.predict_proba(x_test) if is_classifier else est.predict(x_test)
+
+    for pid in range(_NPROCS):
+      np.testing.assert_allclose(results[pid], ref, rtol=1e-4, atol=1e-4)
+
+  def test_regressor(self):
+    self._run(is_classifier=False)
+
+  def test_classifier(self):
+    self._run(is_classifier=True)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tabfm/src/jax/seqpar_test.py b/tabfm/src/jax/seqpar_test.py
new file mode 100644
index 0000000..b8ae11a
--- /dev/null
+++ b/tabfm/src/jax/seqpar_test.py
@@ -0,0 +1,112 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for sequence-parallel JAX inference.
+
+These run on simulated CPU devices (two host devices requested before JAX
+initializes, which is why this lives in its own module), so they exercise the
+sharded forward -- including the cross-device softmax combine and the padded,
+bias-masked K/V gathers -- in CI without GPUs. Predictions are compared
+against the estimators' plain single-device prediction paths.
+"""
+
+import os
+
+os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=2")
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import numpy as np
+
+try:
+  import jax
+  import jax.numpy as jnp
+  from flax import nnx
+  from tabfm.src.jax import model as tabfm_model
+  from tabfm.src.jax import seqpar
+
+  # Full-precision fp32 matmuls so the comparison against the single-device
+  # path stays within tolerance on GPU dev machines too (XLA otherwise uses
+  # TF32 for float32 matmuls on Ampere+).
+  jax.config.update("jax_default_matmul_precision", "highest")
+
+  HAS_JAX = True
+except ImportError:
+  HAS_JAX = False
+
+from tabfm.src.classifier_and_regressor import TabFMClassifier
+from tabfm.src.classifier_and_regressor import TabFMRegressor
+
+# pylint: disable=invalid-name
+
+
+def _tiny_model(loss):
+  return tabfm_model.TabFM(
+      loss=loss,
+      max_classes=2,
+      embed_dim=8,
+      col_num_blocks=1,
+      col_nhead=2,
+      col_num_inds=8,
+      row_num_blocks=1,
+      row_nhead=2,
+      row_num_cls=1,
+      icl_num_blocks=1,
+      icl_nhead=2,
+      dtype=jnp.float32,
+      rngs=nnx.Rngs(0),
+  )
+
+
+@absltest.skipIf(not HAS_JAX, "JAX not installed")
+class SeqparPredictTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ("regressor_even", False, 20, 10, 1),
+      ("regressor_ragged", False, 21, 9, 1),
+      ("regressor_two_members", False, 20, 10, 2),
+      ("classifier_even", True, 20, 10, 1),
+      ("classifier_ragged", True, 23, 7, 2),
+  )
+  def test_matches_single_device(
+      self, is_classifier, n_train, n_test, n_estimators
+  ):
+    rng = np.random.RandomState(0)
+    X = rng.rand(n_train + n_test, 4)
+    if is_classifier:
+      y = rng.randint(0, 2, size=n_train)
+      est = TabFMClassifier(
+          model=_tiny_model("cross_entropy"),
+          n_estimators=n_estimators,
+          random_state=0,
+      )
+    else:
+      y = X[:n_train] @ rng.rand(4)
+      est = TabFMRegressor(
+          model=_tiny_model("mse"), n_estimators=n_estimators, random_state=0
+      )
+    est.fit(X[:n_train], y)
+
+    if is_classifier:
+      ref = est.predict_proba(X[n_train:])
+      out = seqpar.predict_proba(est, X[n_train:])
+    else:
+      ref = est.predict(X[n_train:])
+      out = seqpar.predict(est, X[n_train:])
+
+    np.testing.assert_allclose(out, ref, rtol=1e-4, atol=1e-4)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tabfm/src/jax/splash_attention_test.py b/tabfm/src/jax/splash_attention_test.py
new file mode 100644
index 0000000..744cef8
--- /dev/null
+++ b/tabfm/src/jax/splash_attention_test.py
@@ -0,0 +1,83 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the SPLASH attention implementation.
+
+The splash kernel is TPU-only; these tests run it in Pallas interpret mode
+(CPU-executable, exact kernel semantics) and compare against the stock JAX
+attention implementation on identical weights and inputs, with and without a
+key-prefix mask. This validates the segment-id mask translation and layout
+handling in CI; kernel performance is only measurable on real TPUs.
+"""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import numpy as np
+
+try:
+  import jax
+  import jax.numpy as jnp
+  from flax import nnx
+  from tabfm.src.jax import model as tabfm_model
+
+  jax.config.update("jax_default_matmul_precision", "highest")
+  HAS_JAX = True
+except ImportError:
+  HAS_JAX = False
+
+
+def _mha(impl):
+  return tabfm_model.MultiheadAttention(
+      embed_dim=16,
+      num_heads=2,
+      attention_impl=impl,
+      dtype=jnp.float32,
+      rngs=nnx.Rngs(0),
+  )
+
+
+@absltest.skipIf(not HAS_JAX, "JAX not installed")
+class SplashAttentionTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    tabfm_model.set_splash_interpret(True)
+    self.addCleanup(tabfm_model.set_splash_interpret, False)
+
+  @parameterized.named_parameters(
+      ("full", None),
+      ("prefix_mask", 96),
+  )
+  def test_matches_jax_attention(self, prefix):
+    rng = np.random.RandomState(0)
+    tgt_len, src_len = 128, 256
+    query = jnp.asarray(rng.randn(1, tgt_len, 16), jnp.float32)
+    key = jnp.asarray(rng.randn(1, src_len, 16), jnp.float32)
+    mask = None
+    if prefix is not None:
+      mask = (jnp.arange(src_len)[None, :] < prefix)[:, None, None, :]
+
+    ref = _mha(tabfm_model.AttentionImplementation.JAX)(
+        query, key, key, attn_mask=mask
+    )
+    out = _mha(tabfm_model.AttentionImplementation.SPLASH)(
+        query, key, key, attn_mask=mask
+    )
+    np.testing.assert_allclose(
+        np.asarray(out), np.asarray(ref), rtol=1e-4, atol=1e-4
+    )
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tabfm/src/jax/tabfm_v1_0_0.py b/tabfm/src/jax/tabfm_v1_0_0.py
index bad624b..746948c 100644
--- a/tabfm/src/jax/tabfm_v1_0_0.py
+++ b/tabfm/src/jax/tabfm_v1_0_0.py
@@ -251,8 +251,12 @@ def load(
     checkpoint_path: Local directory containing the 'orbax/' checkpoint, or
       None to download from Hugging Face.
     step: Checkpoint step to restore (for local loading).
-    col_attention_impl: Attention impl for column-attention layers ('jax' or
-      'flash').
+    col_attention_impl: Attention impl for column-attention layers ('jax',
+      'flash', or 'cudnn'). 'flash' is the chunked memory-efficient path
+      (runs anywhere); 'cudnn' uses the fused cuDNN flash kernel (GPU only,
+      fp16/bf16, much faster on long sequences; head dims > 128 -- e.g. the
+      v1.0.0 ICL stage's 256 -- require a Hopper-or-later GPU). 'splash'
+      uses the fused Pallas splash-attention kernel (TPU only).
     row_attention_impl: Attention impl for row-attention layers.
     icl_attention_impl: Attention impl for ICL layers.
     dtype: JAX compute dtype.
diff --git a/tabfm/src/pytorch/seqpar.py b/tabfm/src/pytorch/seqpar.py
new file mode 100644
index 0000000..a23eada
--- /dev/null
+++ b/tabfm/src/pytorch/seqpar.py
@@ -0,0 +1,335 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sequence-parallel (row-sharded) inference for the PyTorch TabFM backend.
+
+TabFM reads the whole training fold as one in-context sequence, so a single
+forward's activations grow with the number of rows and exceed one device's
+memory for very large contexts (roughly >450k rows on an 80GB GPU). This
+module shards the *rows* of one ensemble member's sequence across the ranks of
+a ``torch.distributed`` process group, computing bit-equivalent attention (up
+to floating-point summation order) without any approximation:
+
+  * Cell embedding, row interaction (attention over features), feed-forward
+    blocks and the decoder are row-independent and run locally on each shard.
+  * The column set-transformer's inducing attention (``mab1``) attends over all
+    context rows: each rank computes an online-softmax over its shard and the
+    partial results are combined exactly across ranks via log-sum-exp weights.
+  * The ICL blocks' self-attention attends from every row to all context rows:
+    each rank projects K/V for its context shard, all ranks gather them, and a
+    single fused SDPA runs locally per rank. Test rows are never attention
+    keys anywhere in the model, so only context K/V is communicated.
+
+Typical use (one process per GPU, e.g. under ``torchrun``)::
+
+    dist.init_process_group("nccl")
+    model = tabfm_v1_0_0.load(model_type="regression", device=f"cuda:{rank}")
+    reg = TabFMRegressor(model=model, n_estimators=1)
+    reg.fit(X_train, y_train)                      # cheap; no GPU forward
+    preds = seqpar.predict(reg, X_test)            # sharded across the group
+
+``predict`` runs every ensemble member through the sharded forward and applies
+the same ensemble combination as the estimator's own ``predict`` /
+``predict_proba``, so results match the single-device path up to bf16 noise.
+"""
+
+import math
+from typing import Any, Optional
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+_MAB1_KEY_CHUNK = 8192
+_ROW_CHUNK = 32768
+
+
+def _proj_q(attn, xq):
+  """Projects + normalizes queries the way ``MultiheadAttention`` does."""
+  b, tq, _ = xq.shape
+  q = attn.q_proj(xq).view(b, tq, attn.nhead, attn.hd)
+  q = attn.query_ln(q)
+  scale = (
+      1.442695041 / math.sqrt(attn.hd) * F.softplus(attn.per_dim_scale.float())
+  )
+  return (q * scale.to(q.dtype)).transpose(1, 2).contiguous()  # [B,nh,T,hd]
+
+
+def _dist_mab1(mab, ind_vectors, k_src):
+  """``mab1`` (inducing queries over all context rows) with sharded keys.
+
+  Args:
+    mab: The ``MultiheadAttentionBlock`` used as mab1.
+    ind_vectors: ``[num_inds, E]`` inducing vectors (identical on every rank).
+    k_src: ``[B, T_ctx_local, E]`` this rank's context rows. Ranks may hold
+      different numbers of context rows; only softmax statistics (whose shapes
+      do not depend on the key count) cross ranks.
+
+  Returns:
+    ``[B, num_inds, E]`` block output, identical on every rank.
+  """
+  attn = mab.attn
+  b = k_src.shape[0]
+  q0 = ind_vectors.unsqueeze(0).expand(b, -1, -1)
+  q = _proj_q(attn, mab.pre_attn_ln(q0))
+  nh, ni, hd = q.shape[1], q.shape[2], q.shape[3]
+
+  # Online softmax over local key chunks with fp32 accumulators. Keys are
+  # pre-normed/projected chunk-by-chunk so no full-length copy materializes.
+  m = torch.full((b, nh, ni), -float("inf"), device=q.device)
+  den = torch.zeros((b, nh, ni), device=q.device)
+  num = torch.zeros((b, nh, ni, hd), device=q.device)
+  qf = q.float()
+  for s in range(0, k_src.shape[1], _MAB1_KEY_CHUNK):
+    xk = mab.pre_attn_ln(k_src[:, s : s + _MAB1_KEY_CHUNK])
+    tk = xk.shape[1]
+    ks = attn.key_ln(attn.k_proj(xk).view(b, tk, nh, hd)).transpose(1, 2)
+    vs = attn.v_proj(xk).view(b, tk, nh, hd).transpose(1, 2)
+    scores = qf @ ks.float().transpose(-1, -2)  # [B,nh,I,chunk] fp32
+    m_new = torch.maximum(m, scores.amax(-1))
+    alpha = torch.exp(m - m_new)
+    p = torch.exp(scores - m_new[..., None])
+    den = den * alpha + p.sum(-1)
+    num = num * alpha[..., None] + p @ vs.float()
+    m = m_new
+
+  lse = m + torch.log(den)
+  out_local = num / den[..., None]
+
+  # Exact combine: out = sum_r softmax_r(lse_r) * out_r.
+  world = dist.get_world_size()
+  lse_all = [torch.empty_like(lse) for _ in range(world)]
+  out_all = [torch.empty_like(out_local) for _ in range(world)]
+  dist.all_gather(lse_all, lse.contiguous())
+  dist.all_gather(out_all, out_local.contiguous())
+  w = torch.softmax(torch.stack(lse_all), dim=0)[..., None]
+  out = (torch.stack(out_all) * w).sum(0).to(q.dtype)
+
+  out = out.transpose(1, 2).reshape(b, ni, nh * hd)
+  x = q0 + mab.post_attn_ln(attn.out_proj(out))
+  return x + mab._ff(x)  # pylint: disable=protected-access
+
+
+def _row_chunked(fn, src):
+  """Applies a row-independent fn over row slices of ``[B, T, E]``.
+
+  Bounds the fp32 RMSNorm / FFN transients inside stock module calls, which
+  otherwise materialize full-sequence-length copies.
+  """
+  out = None
+  for s in range(0, src.shape[1], _ROW_CHUNK):
+    o = fn(src[:, s : s + _ROW_CHUNK])
+    if out is None:
+      out = torch.empty(
+          src.shape[0],
+          src.shape[1],
+          o.shape[-1],
+          dtype=o.dtype,
+          device=o.device,
+      )
+    out[:, s : s + _ROW_CHUNK] = o
+  return out
+
+
+def _dist_col_embedding(col, emb, c_local):
+  """``ColEmbedding`` with the row axis sharded. ``emb``: [B,T_local,HC,E]."""
+  b, t, hc, e = emb.shape
+  src = emb.permute(0, 2, 1, 3).reshape(b * hc, t, e)
+  del emb
+  for blk in col.tf_col.blocks:
+    hidden = _dist_mab1(blk.mab1, blk.ind_vectors, src[:, :c_local])
+    # mab2 rows attend only to the replicated inducing outputs: local.
+    src = _row_chunked(lambda s: blk.mab2(s, hidden, hidden), src)  # pylint: disable=cell-var-from-loop
+  out = _row_chunked(lambda s: col.ln_w(col.out_w(s)), src)
+  return out.reshape(b, hc, t, e).permute(0, 2, 1, 3)
+
+
+def _dist_icl_block(blk, x, c_local, c_max, key_mask):
+  """One ICL self-attention block with keys gathered from every rank.
+
+  Ranks may hold unequal context shards: each rank zero-pads its projected
+  K/V to ``c_max`` rows before the all-gather and the padded slots are
+  excluded via ``key_mask``.
+  """
+  attn = blk.attn
+  nh, hd = attn.nhead, attn.hd
+  xn = blk.pre_attn_ln(x)
+  q = _proj_q(attn, xn)
+  kn = xn[:, :c_local]
+  b = kn.shape[0]
+  k = attn.key_ln(attn.k_proj(kn).view(b, c_local, nh, hd)).transpose(1, 2)
+  v = attn.v_proj(kn).view(b, c_local, nh, hd).transpose(1, 2)
+  if c_local < c_max:
+    pad = (0, 0, 0, c_max - c_local)  # pad the row (dim -2) axis
+    k, v = F.pad(k, pad), F.pad(v, pad)
+  world = dist.get_world_size()
+  k_all = [torch.empty_like(k) for _ in range(world)]
+  v_all = [torch.empty_like(v) for _ in range(world)]
+  dist.all_gather(k_all, k.contiguous())
+  dist.all_gather(v_all, v.contiguous())
+  key = torch.cat(k_all, dim=2)
+  val = torch.cat(v_all, dim=2)
+  del k_all, v_all
+  o = F.scaled_dot_product_attention(q, key, val, attn_mask=key_mask, scale=1.0)
+  del key, val
+  b, nh, tq, hd = o.shape
+  o = o.transpose(1, 2).reshape(b, tq, nh * hd)
+  x = x + blk.post_attn_ln(attn.out_proj(o))
+  return x + blk._ff(x)  # pylint: disable=protected-access
+
+
+@torch.inference_mode()
+def seqpar_forward(model, x_local, y_local, c_local, cat_mask=None, d=None):
+  """Sharded forward for one ensemble member. Call on every rank.
+
+  Args:
+    model: Loaded PyTorch ``TabFM`` (classifier or regressor variant).
+    x_local: ``[1, T_local, H]`` float array; this rank's context rows first,
+      then its share of the test rows.
+    y_local: ``[1, T_local]`` float array; ``-100.0`` on test positions.
+    c_local: Number of context rows in this rank's shard (may differ by rank).
+    cat_mask: Optional ``[1, H]`` bool array of categorical-feature flags.
+    d: Optional ``[1]`` int array with the active-feature count (for feature
+      padding).
+
+  Returns:
+    ``[T_local_test, L_out]`` numpy array of this rank's test outputs (scaled
+    predictions for regression, per-class logits for classification).
+  """
+  dev = next(model.parameters()).device
+  x = torch.as_tensor(np.asarray(x_local, dtype=np.float32), device=dev)
+  y = torch.as_tensor(np.asarray(y_local, dtype=np.float32), device=dev)
+  ts = torch.tensor([c_local], device=dev, dtype=torch.long)
+  cm = (
+      torch.as_tensor(np.asarray(cat_mask), device=dev, dtype=torch.bool)
+      if cat_mask is not None
+      else None
+  )
+  dt = (
+      torch.as_tensor(np.asarray(d), device=dev, dtype=torch.long)
+      if d is not None
+      else None
+  )
+
+  # Every rank must agree on the padded per-rank key length for the gathers.
+  world = dist.get_world_size()
+  c_locals = [None] * world
+  dist.all_gather_object(c_locals, c_local)
+  c_max = max(c_locals)
+  key_valid = torch.cat([torch.arange(c_max, device=dev) < c for c in c_locals])
+  key_mask = (
+      None
+      if all(c == c_max for c in c_locals)
+      else key_valid[None, None, None, :]
+  )
+
+  x = torch.nan_to_num(x, nan=-100.0).to(model.cls_tokens.dtype)
+  emb = model.cell_embedder(x, y, ts, cat_mask=cm, d=dt)
+  emb = _dist_col_embedding(model.col_embedder, emb, c_local)
+  b, t, _, _ = emb.shape
+  cls = model.cls_tokens.expand(b, t, -1, -1)
+  emb = torch.cat([cls, emb], dim=2)
+  emb = model.row_interactor(emb, d=dt)
+  emb = _dist_col_embedding(model.col_embedder_2, emb, c_local)
+  reps = model.row_interactor_2(emb, d=dt)
+  del emb
+
+  icl = model.icl_predictor
+  tm = torch.arange(t, device=dev)[None, :] < ts[:, None]
+  if icl.is_classifier:
+    y_enc = icl.y_encoder(y)
+  else:
+    y_enc = icl.y_encoder(y[..., None].to(reps.dtype))
+  r = reps + y_enc * tm[..., None]
+  del reps, y_enc
+  for blk in icl.tf_icl.blocks:
+    r = _dist_icl_block(blk, r, c_local, c_max, key_mask)
+  out = icl.decoder(icl.ln(r))
+  return out[0, c_local:, :].float().cpu().numpy()
+
+
+def _shard_bounds(n, world, rank):
+  """Contiguous near-equal split of ``n`` items; returns (start, stop)."""
+  base, rem = divmod(n, world)
+  start = rank * base + min(rank, rem)
+  return start, start + base + (1 if rank < rem else 0)
+
+
+def _member_outputs(estimator, X, rank, world):
+  """Runs every ensemble member through the sharded forward.
+
+  Returns ``[n_members, n_test, L_out]`` outputs, replicated on every rank.
+  """
+  x_enc = estimator.X_encoder_.transform(X)
+  data = estimator.ensemble_generator_.transform(x_enc)
+  xs, ys, cat_masks, ds, _ = (
+      estimator.ensemble_generator_.prepare_ensemble_tensors(data)
+  )
+  n_members = xs.shape[0]
+  n_train = ys.shape[1]
+  n_test = xs.shape[1] - n_train
+
+  c0, c1 = _shard_bounds(n_train, world, rank)
+  t0, t1 = _shard_bounds(n_test, world, rank)
+  outs = []
+  for mi in range(n_members):
+    x_local = np.concatenate(
+        [xs[mi, c0:c1], xs[mi, n_train + t0 : n_train + t1]], axis=0
+    )[None]
+    y_local = np.concatenate([ys[mi, c0:c1], np.full(t1 - t0, -100.0)], axis=0)[
+        None
+    ]
+    out = seqpar_forward(
+        estimator.model,
+        x_local,
+        y_local,
+        c1 - c0,
+        cat_mask=cat_masks[mi : mi + 1] if cat_masks is not None else None,
+        d=ds[mi : mi + 1] if ds is not None else None,
+    )
+    gathered = [None] * world
+    dist.all_gather_object(gathered, out)
+    outs.append(np.concatenate(gathered, axis=0))
+  return np.stack(outs, axis=0)
+
+
+def predict(estimator, X):
+  """Sharded equivalent of ``TabFMRegressor.predict``.
+
+  Must be called collectively on every rank of the process group with the same
+  fitted estimator (fit is deterministic given ``random_state``) and the same
+  ``X``. Every rank returns the full prediction vector.
+  """
+  outputs = _member_outputs(
+      estimator, X, dist.get_rank(), dist.get_world_size()
+  )
+  predictions = outputs.squeeze(-1)  # [E, T]
+  return estimator._combine_predictions(predictions)  # pylint: disable=protected-access
+
+
+def predict_proba(estimator, X):
+  """Sharded equivalent of ``TabFMClassifier.predict_proba`` (all ranks)."""
+  outputs = _member_outputs(
+      estimator, X, dist.get_rank(), dist.get_world_size()
+  )
+  outputs = outputs[..., : estimator.n_classes_]
+  offsets = []
+  for offs in estimator.ensemble_generator_.class_shift_offsets_.values():
+    offsets.extend(offs)
+  logits = np.stack([
+      np.concatenate([out[..., off:], out[..., :off]], axis=-1)
+      for out, off in zip(outputs, offsets)
+  ])
+  return estimator._process_logits(logits)  # pylint: disable=protected-access
diff --git a/tabfm/src/pytorch/seqpar_test.py b/tabfm/src/pytorch/seqpar_test.py
new file mode 100644
index 0000000..535d351
--- /dev/null
+++ b/tabfm/src/pytorch/seqpar_test.py
@@ -0,0 +1,163 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for sequence-parallel PyTorch inference.
+
+These run with the gloo backend on CPU processes, so they exercise the
+sharded forward (including the cross-rank softmax combine and padded K/V
+gathers) in CI without needing GPUs. Outputs are compared against the plain
+single-process model forward.
+"""
+
+import functools
+import multiprocessing
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from tabfm.src.pytorch import model as tabfm_model
+from tabfm.src.pytorch import seqpar
+
+# pylint: disable=invalid-name
+
+_WORLD = 2
+
+
+def _tiny_model(is_classifier):
+  torch.manual_seed(0)
+  m = tabfm_model.TabFM(
+      embed_dim=8,
+      max_classes=3,
+      col_num_blocks=1,
+      col_nhead=2,
+      col_num_inds=4,
+      row_num_blocks=1,
+      row_nhead=2,
+      row_num_cls=2,
+      icl_num_blocks=2,
+      icl_nhead=2,
+      is_classifier=is_classifier,
+  )
+  # Randomize parameters so the comparison is not trivially zeros.
+  with torch.no_grad():
+    for p in m.parameters():
+      p.uniform_(-0.05, 0.05)
+  return m.eval()
+
+
+def _make_data(n_train, n_test, h, is_classifier, seed=0):
+  rng = np.random.default_rng(seed)
+  x = rng.standard_normal((1, n_train + n_test, h)).astype(np.float32)
+  if is_classifier:
+    y_train = rng.integers(0, 2, n_train).astype(np.float32)
+  else:
+    y_train = rng.standard_normal(n_train).astype(np.float32)
+  y = np.concatenate([y_train, np.full(n_test, -100.0)]).astype(np.float32)[
+      None
+  ]
+  return x, y
+
+
+def _reference_forward(model, x, y, n_train, cat_mask=None, d=None):
+  ts = torch.tensor([n_train], dtype=torch.long)
+  y_pad = torch.from_numpy(y)
+  with torch.inference_mode():
+    out = model(
+        torch.from_numpy(x),
+        y_pad,
+        ts,
+        cat_mask=torch.from_numpy(cat_mask) if cat_mask is not None else None,
+        d=torch.from_numpy(d) if d is not None else None,
+    )
+  return out[0, n_train:, :].float().numpy()
+
+
+def _worker(rank, port, is_classifier, n_train, n_test, h, use_cat_and_d, q):
+  """Runs the sharded forward on one CPU rank and reports rank-0's result."""
+  os.environ["MASTER_ADDR"] = "127.0.0.1"
+  os.environ["MASTER_PORT"] = str(port)
+  dist.init_process_group("gloo", rank=rank, world_size=_WORLD)
+  try:
+    model = _tiny_model(is_classifier)
+    x, y = _make_data(n_train, n_test, h, is_classifier)
+    cat_mask = d = None
+    if use_cat_and_d:
+      cat_mask = np.zeros((1, h), dtype=bool)
+      cat_mask[0, 0] = True
+      d = np.array([h - 1], dtype=np.int64)  # last feature column is padding
+
+    c0, c1 = seqpar._shard_bounds(n_train, _WORLD, rank)  # pylint: disable=protected-access
+    t0, t1 = seqpar._shard_bounds(n_test, _WORLD, rank)  # pylint: disable=protected-access
+    x_local = np.concatenate(
+        [x[0, c0:c1], x[0, n_train + t0 : n_train + t1]], axis=0
+    )[None]
+    y_local = np.concatenate([y[0, c0:c1], np.full(t1 - t0, -100.0)])[None]
+    out = seqpar.seqpar_forward(
+        model, x_local, y_local, c1 - c0, cat_mask=cat_mask, d=d
+    )
+    gathered = [None] * _WORLD
+    dist.all_gather_object(gathered, out)
+    if rank == 0:
+      q.put(np.concatenate(gathered, axis=0))
+  finally:
+    dist.destroy_process_group()
+
+
+class SeqparForwardTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ("regressor_even", False, 64, 16, 5, False),
+      ("classifier_even", True, 64, 16, 5, False),
+      ("regressor_ragged", False, 63, 15, 5, False),
+      ("classifier_ragged", True, 61, 17, 5, False),
+      ("regressor_cat_and_d", False, 64, 16, 5, True),
+  )
+  def test_matches_single_process(
+      self, is_classifier, n_train, n_test, h, use_cat_and_d
+  ):
+    ctx = multiprocessing.get_context("spawn")
+    q = ctx.Queue()
+    port = 29500 + hash(self._testMethodName) % 1000
+    procs = [
+        ctx.Process(
+            target=_worker,
+            args=(r, port, is_classifier, n_train, n_test, h, use_cat_and_d, q),
+        )
+        for r in range(_WORLD)
+    ]
+    for p in procs:
+      p.start()
+    sharded = q.get(timeout=120)
+    for p in procs:
+      p.join(timeout=120)
+      self.assertEqual(p.exitcode, 0)
+
+    model = _tiny_model(is_classifier)
+    x, y = _make_data(n_train, n_test, h, is_classifier)
+    cat_mask = d = None
+    if use_cat_and_d:
+      cat_mask = np.zeros((1, h), dtype=bool)
+      cat_mask[0, 0] = True
+      d = np.array([h - 1], dtype=np.int64)
+    ref = _reference_forward(model, x, y, n_train, cat_mask=cat_mask, d=d)
+
+    np.testing.assert_allclose(sharded, ref, rtol=1e-4, atol=1e-4)
+
+
+if __name__ == "__main__":
+  absltest.main()