[perf]: Add guarded MiniMax-H3 Sol Engine inference optimizations

.buildkite/scripts/lanes/ssim.sh

@@ -25,6 +25,12 @@ if [ "$selected" != all ]; then
   done
 fi
 
+# This lane owns a whole four-GPU tray. Exercise the exact NCCL + Triton
+# MiniMax-H3 packed-SP route before output-quality jobs, while retaining the
+# scheduler-provided rendezvous port in the nested torchrun invocation.
+FASTVIDEO_MINIMAX_H3_PACKED_SP_STRICT_CUDA=1 \
+  pytest ./fastvideo/tests/distributed/test_minimax_h3_packed_sp.py -vs
+
 # MoGe's utils3d dependency builds glcontext from source on ARM64. The current
 # runner image predates the baked-in X11 headers below, so keep this guarded
 # bootstrap until every deployed image digest contains libx11-dev.

.buildkite/scripts/lanes/transformer.sh

@@ -2,4 +2,4 @@
 # Canonical Slurm CI selection for the transformer lane.
 set -euo pipefail
 
-exec pytest ./fastvideo/tests/transformers -vs
+exec pytest ./fastvideo/tests/transformers ./fastvideo/tests/distributed/test_minimax_h3_packed_sp.py -vs

.github/scripts/plan_merge_ci.py

@@ -373,6 +373,13 @@ def classify_paths(paths: list[str]) -> MergePlan:
                 plan.add_lanes("ssim", reason=f"shared SSIM harness/reference: {path}")
             continue
 
+        if path == "fastvideo/tests/distributed/test_minimax_h3_packed_sp.py":
+            # Fastcheck owns the portable world-4 Gloo contract. Also select
+            # the existing four-GPU SSIM lane, whose preflight makes the same
+            # test require the production NCCL + Triton route.
+            plan.add_ssim(("test_minimax_h3_similarity.py", ), reason=f"MiniMax-H3 packed-SP CUDA preflight: {path}")
+            continue
+
         if path.startswith("fastvideo/tests/performance/") or path.startswith(".buildkite/performance-benchmarks/"):
             plan.add_lanes("performance", reason=f"performance coverage: {path}")
             continue

docs/contributing/ci_architecture.md

@@ -164,6 +164,14 @@ selection never deletes or dynamically invents a Buildkite step.
 | Modular train framework | `train_framework` | 1 | `fastvideo/train/` and its tests |
 | Eval metrics | `eval` | 1 | `fastvideo/eval/` and its tests |
 
+The transformer Fastcheck lane also runs the MiniMax-H3 packed sequence-
+parallel world-4 contract with Gloo. That portable test proves rank ordering,
+padding/trim behavior, and Q/K/V-dependent scatter/gather semantics; it is not
+CUDA performance or NCCL/Triton evidence. The existing four-GPU SSIM lane runs
+the same contract first in strict mode, which fails unless all four CUDA
+devices are visible and the production NCCL plus Triton relayout route runs.
+No separate distributed lane is introduced.
+
 Golden-gate and SSIM selections are basenames, not arbitrary pytest arguments.
 The private host checks the comma-separated allowlist before staging, and the
 container checks it again before invoking pytest. Shared quality-harness
@@ -352,7 +360,8 @@ table binds each internal `*_ci` type to that script, its GPU count, wall-clock
 limit, dependency extras, kernel-build policy, secrets, and artifacts. The
 internal suffix is an implementation detail; there is only one active backend.
 
-SSIM uses `fastvideo/tests/ssim/ci_runner.py` inside a single four-GPU lease.
+SSIM first runs the strict MiniMax-H3 packed-SP production preflight, then uses
+`fastvideo/tests/ssim/ci_runner.py` inside the same four-GPU lease.
 It discovers `REQUIRED_GPUS` and `*_MODEL_TO_PARAMS` with AST parsing, then
 packs independent pytest subprocesses across the visible GPUs with fail-fast
 termination. Performance writes reports to a host-mounted artifact directory;

docs/contributing/testing.md

@@ -16,6 +16,7 @@ slash-command mappings, and workflow ownership live in
 | Inference tests | `fastvideo/tests/inference` | Validate specialized inference paths such as LoRA inference and V-MoBA. |
 | Performance tests | `fastvideo/tests/performance` | Gate latency, throughput, peak memory, and stage timings. See [Performance Benchmarks](performance_benchmarks.md). |
 | Eval tests | `fastvideo/tests/eval` | Check eval metrics against pinned reference scores and assets. |
+| Distributed contracts | `fastvideo/tests/distributed` | Exercise selected real multi-rank communication contracts; most remain manual unless an existing lane names them explicitly. |
 | DreamVerse app tests | `apps/dreamverse` | Validate the DreamVerse backend, frontend, and mock-backed browser flows. |
 
 ## Running Tests Locally
@@ -151,6 +152,14 @@ by the changed model family. Shared SSIM harness changes still select the
 complete lane. Independently, `main` runs the full SSIM matrix every Sunday at
 05:00 UTC so infrequently touched model families retain periodic coverage.
 
+Before scheduling SSIM cases, the lane uses its existing four-GPU allocation
+to run the MiniMax-H3 packed-SP world-4 contract in strict CUDA mode. Strict
+mode rejects a Gloo fallback and therefore covers the production NCCL and
+Triton relayout route. The one-GPU transformer Fastcheck lane runs the same
+test with its portable Gloo fallback; that result covers collective ordering
+and Q/K/V semantics only, not CUDA behavior or performance. Both nested
+torchrun invocations preserve the Slurm runner's assigned `MASTER_PORT`.
+
 For a focused developer run, invoke pytest directly and optionally select one
 model from a parameterized test through `FASTVIDEO_SSIM_MODEL_ID`:
 

docs/inference/optimizations.md

@@ -113,6 +113,51 @@ the Preview checkpoint's sparse VSA blocks. Packed-varlen changes floating-point
 reduction order relative to fixed-length FA4, so treat it as a speed/quality
 evaluation option rather than an exact-parity mode.
 
+### MiniMax-H3 AdaLN trajectory and packed sequence parallelism (opt-in)
+
+Two additional released Sol-Engine-style DiT optimizations are independently
+available for controlled MiniMax-H3 inference experiments:
+
+```bash
+export FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE=1
+export FASTVIDEO_MINIMAX_H3_PACKED_SP=1
+```
+
+Both flags default to off. They do not enable VSA, cross-step caches,
+quantization, or VAE optimizations.
+
+- `FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE` evaluates each full-rank block AdaLN
+  projection once for every step in the fixed denoising schedule, installs
+  non-persistent lookup tables, and releases the projection modules. The table
+  values match the original projections, but the loaded transformer is then
+  tied to that schedule for the rest of its lifetime. The runtime cursor is a
+  non-persistent buffer and follows transformer device moves, including full
+  CPU offload. The replacement tables and cursor are deliberately omitted
+  from `state_dict`, while the original projection parameters no longer
+  exist: a post-precompute state dict is therefore not reloadable. Save or
+  convert checkpoints before enabling precompute, and reload the stock model
+  to change schedules. It requires a stock
+  full-rank checkpoint with replicated, materialized weights:
+  `dit_layerwise_offload=False` and `use_fsdp_inference=False`. Setup briefly
+  holds the original projections and all replacement tables together. In the
+  reported 50-step workload the tables used 0.88 GB and replaced 24.23 GB of
+  projection weights; other schedules and checkpoints have different costs.
+- `FASTVIDEO_MINIMAX_H3_PACKED_SP` fuses the Q/K/V communication layout around
+  two direct NCCL all-to-all collectives. It engages only for inference with
+  sequence-parallel world size greater than one, batch size one, dense
+  FlashAttention, no replicated tokens or deferred RoPE, and no VSA. A
+  grad-enabled forward automatically retains the autograd-aware generic
+  Ulysses route. The packed direct collective is separate from the optional
+  `FASTVIDEO_ULYSSES_A2A=auto` transport.
+
+At SP=1 the packed flag is deliberately inert, so it provides no acceleration
+on a one-GPU GB10 and does not force FlashAttention over the configured GB10
+backend. The contribution's performance evidence is scoped to dense BF16 FA4
+on 4× GB200, SP=4, 1344×768×124, with one post-warmup sample; it is not evidence
+for a sparse/VSA route or a one-GPU system. These exact table/relayout features
+also do not make regional `torch.compile` eager-parity-safe: the MiniMax-H3
+accuracy caveat below still applies to the complete serving profile.
+
 ### FP4 Flash Attention 4 (Blackwell only)
 
 **`FLASH_ATTN`** with **`--nvfp4_fa4`**

fastvideo/attention/layer.py

@@ -8,12 +8,13 @@
 
 from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
 from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather,
-                                                    sequence_model_parallel_all_to_all_4D)
+                                                    sequence_model_parallel_all_to_all_4D,
+                                                    sequence_model_parallel_direct_all_to_all)
 from fastvideo.distributed.parallel_state import (get_sp_parallel_rank, get_sp_world_size)
 from fastvideo.forward_context import ForwardContext, get_forward_context
+from fastvideo.layers.rotary_embedding import _apply_rotary_emb
 from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.utils import get_compute_dtype
-from fastvideo.layers.rotary_embedding import _apply_rotary_emb
 
 
 def _attention_compile_disabled() -> bool:
@@ -70,6 +71,7 @@ def __init__(self,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
                  prefix: str = "",
+                 packed_qkv_relayout: bool = False,
                  **extra_impl_args) -> None:
         super().__init__()
         if softmax_scale is None:
@@ -103,6 +105,9 @@ def __init__(self,
         # inference loader may enable this one instance after validating the
         # transformer's resolved backend; no process-global default changes.
         self._compile_forward_enabled = not _attention_compile_disabled()
+        self.packed_qkv_relayout = packed_qkv_relayout
+        if packed_qkv_relayout and self.backend != AttentionBackendEnum.FLASH_ATTN:
+            raise ValueError("MiniMax-H3 packed QKV relayout currently supports only dense FLASH_ATTN")
 
     def _set_compile_forward_enabled(self, enabled: bool) -> None:
         self._compile_forward_enabled = enabled
@@ -144,6 +149,43 @@ def forward(
         forward_context: ForwardContext = get_forward_context()
         ctx_attn_metadata = forward_context.attn_metadata
 
+        # The direct packed collective is deliberately inference-only. A
+        # training process may inherit the opt-in environment variable, but a
+        # grad-enabled forward must retain the established autograd-aware
+        # Ulysses path.
+        if self.packed_qkv_relayout and world_size > 1 and not torch.is_grad_enabled():
+            if batch_size != 1:
+                raise ValueError("MiniMax-H3 packed QKV relayout currently requires batch size 1")
+            if num_heads % world_size:
+                raise ValueError(
+                    f"MiniMax-H3 packed QKV relayout requires {num_heads} heads to be divisible by SP={world_size}")
+            if any(t is not None for t in (replicated_q, replicated_k, replicated_v, freqs_cis)):
+                raise ValueError("MiniMax-H3 packed QKV relayout does not support replicated tokens or deferred RoPE")
+            from fastvideo.models.dits.minimax_h3_fusions.relayout import (merge_heads, pack_qkv_destination_major)
+
+            rows_local = q.shape[1]
+            packed = pack_qkv_destination_major(q[0], k[0], v[0], world_size)
+            packed = sequence_model_parallel_direct_all_to_all(packed)
+            heads_local = num_heads // world_size
+            packed = packed.reshape(world_size * rows_local, heads_local, 3 * self.head_size)
+            q_full, k_full, v_full = packed.split(self.head_size, dim=-1)
+            if original_seq_len is None:
+                original_seq_len = q_full.shape[0]
+            if original_seq_len < 1 or original_seq_len > q_full.shape[0]:
+                raise ValueError(f"original_seq_len must be in [1, {q_full.shape[0]}], got {original_seq_len}")
+            pad_seq_len = q_full.shape[0] - original_seq_len
+            output = self.attn_impl.forward(
+                q_full[:original_seq_len].unsqueeze(0),
+                k_full[:original_seq_len].unsqueeze(0),
+                v_full[:original_seq_len].unsqueeze(0),
+                ctx_attn_metadata,
+            )
+            output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
+            output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
+            output = sequence_model_parallel_direct_all_to_all(output.squeeze(0).contiguous())
+            output = merge_heads(output.reshape(world_size, rows_local, heads_local, self.head_size))
+            return output.unsqueeze(0), None
+
         # Stack QKV
         qkv = torch.cat([q, k, v], dim=0)  # [3*batch, seq_len, num_heads, head_dim]
 

fastvideo/distributed/communication_op.py

@@ -14,6 +14,36 @@
 _sp_warmup_done = False
 
 
+def _validate_direct_all_to_all_group(group: object, input_: torch.Tensor, expected_backend: str) -> None:
+    """Validate the live process group immediately before a direct collective."""
+    process_group = getattr(group, "device_group", None)
+    if process_group is None or not torch.distributed.is_initialized():
+        raise RuntimeError("direct all-to-all requires a live distributed process group")
+    try:
+        actual_world = torch.distributed.get_world_size(process_group)
+        torch.distributed.get_rank(process_group)
+        backend = str(torch.distributed.get_backend(process_group)).lower()
+    except (RuntimeError, ValueError) as error:
+        raise RuntimeError("direct all-to-all process group is not live") from error
+    configured_world = int(getattr(group, "world_size", 0))
+    if actual_world != configured_world:
+        raise RuntimeError(
+            f"direct all-to-all group world size mismatch: coordinator={configured_world}, process_group={actual_world}"
+        )
+    if backend != expected_backend:
+        raise RuntimeError(
+            f"direct all-to-all requires the {expected_backend} backend for {input_.device.type} tensors, got {backend}"
+        )
+    configured_device = getattr(group, "device", None)
+    if configured_device is None:
+        raise RuntimeError("direct all-to-all coordinator does not declare its collective device")
+    expected_device = torch.device(configured_device)
+    if (expected_device.type != input_.device.type
+            or (expected_device.index is not None and expected_device.index != input_.device.index)):
+        raise RuntimeError(
+            f"direct all-to-all tensor device {input_.device} does not match coordinator device {expected_device}")
+
+
 def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
     """All-reduce the input tensor across model parallel group."""
     return get_tp_group().all_reduce(input_)
@@ -32,6 +62,40 @@ def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor,
     return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim)
 
 
+def sequence_model_parallel_direct_all_to_all(input_: torch.Tensor) -> torch.Tensor:
+    """Synchronous equal-split all-to-all used by the packed H3 inference path.
+
+    This primitive intentionally has no autograd formula. The owning attention
+    route falls back to ``all_to_all_4D`` whenever gradients are enabled; this
+    guard keeps other callers from discovering the restriction at backward.
+    """
+    group = get_sp_group()
+    # Packed SP is deliberately inert at SP=1 and must not require distributed
+    # initialization merely to preserve the identity path.
+    if group.world_size == 1:
+        return input_
+    if torch.is_grad_enabled() and input_.requires_grad:
+        raise RuntimeError("sequence_model_parallel_direct_all_to_all is inference-only and has no autograd formula; "
+                           "use sequence_model_parallel_all_to_all_4D for grad-enabled execution")
+    if input_.ndim < 1 or input_.shape[0] % group.world_size:
+        raise ValueError(
+            "direct all-to-all requires the leading dimension to be evenly divisible by the SP world size; "
+            f"got shape {tuple(input_.shape)} and SP={group.world_size}")
+    if not input_.is_contiguous():
+        raise ValueError("direct all-to-all requires a contiguous input tensor")
+    # Dynamo reaches the CUDA custom op below, whose opaque runtime body
+    # repeats these checks. Eager and CPU/Gloo execution validate here.
+    if not torch.compiler.is_compiling():
+        _validate_direct_all_to_all_group(group, input_, "nccl" if input_.is_cuda else "gloo")
+    # CPU/Gloo is useful for the real multi-rank contract test. The production
+    # packed H3 route is CUDA/Triton and takes the compiler-visible custom op.
+    if not input_.is_cuda:
+        output = torch.empty_like(input_)
+        torch.distributed.all_to_all_single(output, input_, group=group.device_group)
+        return output
+    return torch.ops.fastvideo.direct_all_to_all_single(input_, group.unique_name)
+
+
 def sequence_model_parallel_all_gather(input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
     """All-gather the input tensor across model parallel group."""
     return get_sp_group().all_gather(input_, dim)

fastvideo/distributed/parallel_state.py

@@ -114,6 +114,59 @@ def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
     return torch.empty_like(tensor)
 
 
+@torch.library.custom_op(
+    "fastvideo::direct_all_to_all_single",
+    mutates_args=(),
+    device_types="cuda",
+)
+def direct_all_to_all_single(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
+    """Issue a synchronous all-to-all without the functional-collective wrapper."""
+    if group_name not in _groups:
+        raise RuntimeError(f"Group {group_name} is not registered.")
+    group = _groups[group_name]()
+    if group is None:
+        raise RuntimeError(f"Group {group_name} is destroyed.")
+    process_group = getattr(group, "device_group", None)
+    if process_group is None or not torch.distributed.is_initialized():
+        raise RuntimeError("direct all-to-all requires a live NCCL process group")
+    try:
+        actual_world = torch.distributed.get_world_size(process_group)
+        torch.distributed.get_rank(process_group)
+        backend = str(torch.distributed.get_backend(process_group)).lower()
+    except (RuntimeError, ValueError) as error:
+        raise RuntimeError("direct all-to-all process group is not live") from error
+    configured_world = int(getattr(group, "world_size", 0))
+    if actual_world != configured_world:
+        raise RuntimeError(
+            f"direct all-to-all group world size mismatch: coordinator={configured_world}, process_group={actual_world}"
+        )
+    if backend != "nccl":
+        raise RuntimeError(f"direct all-to-all CUDA tensors require the NCCL backend, got {backend}")
+    configured_device = getattr(group, "device", None)
+    if configured_device is None:
+        raise RuntimeError("direct all-to-all coordinator does not declare its CUDA device")
+    expected_device = torch.device(configured_device)
+    if (expected_device.type != "cuda"
+            or (expected_device.index is not None and expected_device.index != tensor.device.index)):
+        raise RuntimeError(
+            f"direct all-to-all tensor device {tensor.device} does not match coordinator device {expected_device}")
+    if tensor.ndim < 1 or tensor.shape[0] % actual_world:
+        raise ValueError(
+            "direct all-to-all requires the leading dimension to be evenly divisible by the live group world size; "
+            f"got shape {tuple(tensor.shape)} and world={actual_world}")
+    if not tensor.is_contiguous():
+        raise ValueError("direct all-to-all requires a contiguous input tensor")
+    output = torch.empty_like(tensor)
+    torch.distributed.all_to_all_single(output, tensor, group=group.device_group)
+    return output
+
+
+@torch.library.register_fake("fastvideo::direct_all_to_all_single")
+def direct_all_to_all_single_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
+    del group_name
+    return torch.empty_like(tensor)
+
+
 class GroupCoordinator:
     """
     PyTorch ProcessGroup wrapper for a group of processes.