[MiniMax-H3] Reuse Ulysses shard lengths across an inference forward

[LeslieWylie]

What does this PR do?

MiniMax-H3 keeps one packed self-attention sequence across its transformer blocks, but Ulysses Anything currently collects local sequence lengths separately for the Q, K, V, and output exchanges in every attention call. This draft adds an explicit inference opt-in, attention_kwargs={"reuse_ulysses_metadata": True}, to construct and pass one shard-length plan across the transformer blocks in a single forward.

Every participating rank gathers the full packed length on every forward and checks agreement before deriving the existing tensor_split lengths. The plan lives in a per-call copy of ParallelConfig; it does not mutate the model configuration or cache state between requests. The token refiner keeps its existing collection path. All ranks must use the same opt-in setting and receive the same full packed layout. Grad-enabled calls with the option are rejected. Default H3, other models, Ring Attention, and normal Ulysses behavior remain unchanged.

The lower-level seq_lengths plumbing overlaps with open PR #14834. That proposal reuses lengths within one attention invocation and includes backward support. This draft's separate increment is the H3-specific, per-forward plan shared across transformer layers. It does not add backward support, and the shared helper changes should be reconciled if that PR lands first.

Validation

The focused CPU/Gloo check ran the real tiny H3 model on two processes with successive packed lengths 18 -> 20 -> 18. Opt-in and default outputs were bitwise equal, profiler all-gather counts decreased, and the original configuration object remained unchanged. A direct uneven attention check covered 9 -> 10 -> 9, where rank zero retains five local rows while the other rank changes length. A mismatched global layout raised on both ranks without retaining a plan. The invalid inference/parallelism case also passed.

DIFFUSERS_TEST_DEVICE=cpu PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 OMP_NUM_THREADS=1 \
  python -m pytest tests/models/transformers/test_models_transformer_minimax_h3.py::TestMiniMaxH3SequencePlan -q

Affected Ruff 0.9.10 checks, copy/dummy consistency, AI-instruction checks, forward-docstring checks, and whitespace checks passed. The only change after the focused execution was documentation clarifying the shared-rank opt-in requirement and token-refiner scope.

CUDA/NCCL, torch.compile, full-size checkpoint generation, isolated latency, and output-quality validation are pending. This draft makes no measured speedup claim. Earlier combined-stack experiments are not presented as a benchmark of this patch.

Final self-review notes

READY for Draft review. The repository self-review guidance and an independent read-only review found no blocking correctness issue or dead code. The opt-in is reachable through the documented call and actual model checks. Under the standard H3 split plan and matching rank settings, every forward executes the same initial collective; no rank-local cache hit can skip it. The two review suggestions about synchronized opt-in and transformer-block scope were applied in documentation.

The collective-order guarantee does not extend to inconsistent rank settings or arbitrary external custom processors/split plans. GPU behavior and performance remain unverified. No maintainer approval or human execution of validation is claimed. AI assistance was used for implementation, validation, and this description; the requester authorized publication as a Draft.

Before submitting

• AI assistance was used and disclosed.
• Read the Coding with AI agents guide and applicable repository instructions.
• Applied the repository self-review skill and included its final notes.
• Read the contributor guideline and philosophy guidance.
• Discussed/approved by maintainers in an issue or forum.
• Updated documentation.
• Added the focused distributed regression check needed for rank-local length changes and per-forward state isolation.
• Author or member of the original MiniMax-H3 model team.