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

[weihaokong]

Summary

Three additions for large-context inference:

1. tabfm/src/pytorch/seqpar.py — sequence-parallel (row-sharded)
   multi-GPU inference for the PyTorch backend under torch.distributed.
2. tabfm/src/jax/seqpar.py — the equivalent for the JAX backend: a
   single process shards the rows across all local devices with
   jax.shard_map.
3. AttentionImplementation.CUDNN — fused cuDNN flash attention option
   for the JAX backend.

Sequence-parallel inference (both backends)

TabFM reads the whole training fold as one in-context sequence, so
activations grow with the number of rows and exceed a single device's memory
for very large contexts (roughly >450k rows on an 80GB GPU). These modules
shard the rows of one ensemble member's sequence across devices with
mathematically exact attention — no approximation; outputs match the
single-device path up to floating-point summation order.

• Cell embedding, row interaction (attention over features), feed-forward
  blocks and the decoder are row-independent and run locally on each shard,
  reusing the stock modules (so all fp32-upcast behavior is inherited
  unchanged).
• The column set-transformer's inducing attention runs an fp32 online
  softmax over each device's context shard and the partial results are
  combined exactly across devices via log-sum-exp weights; only softmax
  statistics cross devices, so unequal shards work natively. The JAX
  version additionally scans the keys in chunks so no full-length score
  matrix materializes.
• ICL self-attention gathers each device's projected context K/V (test rows
  are never attention keys anywhere in the model) and runs one fused SDPA
  (PyTorch) or the chunked memory-efficient attention (JAX) per device,
  with padded key slots masked.
• seqpar.predict / seqpar.predict_proba mirror TabFMRegressor.predict
  and TabFMClassifier.predict_proba (class-shift undo, ensemble
  combination) for any n_estimators, with cat_mask / feature-padding
  support. A runnable PyTorch example is included under examples/.

Measured (PyTorch, 4x H100): a 1M-row context runs at ~35GB/GPU, ~104s per
ensemble member — a single 80GB device cannot run it at all. At
single-device-feasible sizes the sharded path is faster (TabArena
GiveMeSomeCredit fold 0, 135k-row context, n_estimators=1: 17.8s -> 5.2s
with AUC unchanged at bf16 noise level). Fold 0 of all 13 TabArena
regression tasks with the 32-member ensemble preset and sharded predict
reproduces the published TabFM-Ensemble reference RMSEs (results/) with a
median deviation of 0.20%.

cuDNN flash attention (JAX backend)

AttentionImplementation.CUDNN, selectable via
load(..., col_attention_impl='cudnn', icl_attention_impl='cudnn'). The
chunked FLASH path stays the default; CUDNN dispatches to
jax.nn.dot_product_attention(implementation='cudnn'). TabFM's boolean
masks are always key-prefix (padding) masks, so they are mapped onto
cuDNN's key_value_seq_lengths instead of the dense [B, N, T, T_src]
mask cuDNN would otherwise require. Head dims > 128 (e.g. the v1.0.0 ICL
stage's 256) require a Hopper-or-later GPU; cuDNN raises
NotImplementedError otherwise.

Measured on 1x H100 (same fold-0 task, 135k-row context): single-member
predict_proba 634s -> 7.4s, with predictions matching the PyTorch backend
(probability correlation 0.999; the chunked path's bf16 accumulation was
the numerical outlier).

Tests

• PyTorch: seqpar_test.py runs 2 CPU processes with the gloo backend and
  compares the sharded forward against the single-process model forward at
  1e-4, covering regression/classification, even and ragged shards, and
  cat_mask/d feature padding.
• JAX: seqpar_test.py runs on two simulated CPU host devices (same
  pattern as classifier_and_regressor_multidevice_test.py) and compares
  sharded predict/predict_proba against the estimators' single-device
  paths at 1e-4 in float32, including multi-member ensembles.
  seqpar_multiprocess_test.py rehearses the multi-host execution model
  (two spawned processes joined via jax.distributed.initialize, e.g. for
  multi-host TPU slices) against the same reference.
• Full repo suite passes with all commits (100 tests).