Add sequence-parallel multi-GPU inference (PyTorch and JAX) and a cuDNN flash attention option

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

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:

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()