[perf] FastH3 two-Spark: gather packed FP4 inputs before head projections

Opt-in h3_vsa_fp4.sp_input_gather (default false; requires sp_projection_overlap) gathers the
already-packed FP4 inputs and rotary tables, then computes each rank's assigned Q/K/V heads locally;
the gate projection overlaps the smaller transfer. Supported on the measured GB10/cuDNN shape only;
other shapes keep the existing exchange. Enabled in the Spark-pair example.
T2AV on 2 Sparks: 19.35 -> 18.76 s/clip (merge gate check, same SP reduction order rules).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

docs/getting_started/installation/spark_pair.md

@@ -162,6 +162,14 @@ projection is packed, overlapping traffic with the remaining projection GEMMs. I
 QKV/gate projections with shared input quantization. Other configurations use the ordinary
 head-group exchange. Packing arithmetic and FP8 transport precision are unchanged.
 
+The optional `h3_vsa_fp4.sp_input_gather: true` gathers the existing packed FP4 inputs and
+rotary tables before computing each rank's assigned Q/K/V heads. The gate projection overlaps
+that smaller input transfer. It defaults to `false`, requires `sp_projection_overlap: true`,
+and currently supports the measured GB10/cuDNN shape: 7430 local rows, 5376 input channels,
+56 heads and bias-free NVFP4 QKV projections. Other shapes use the existing exchange path.
+All projections, normalization, rotary operations and FP8 packet rounding remain present;
+no steps or tokens are omitted. Validate generated output when changing this GEMM shape.
+
 The optional `h3_vsa_fp4.sp_pooled_key_mean: true` reuses tile statistics for key smoothing
 on the grouped tile-64 path with FP8 prefix attention and `smooth_k: true`. It defaults to
 `false`. All key rows still contribute, including zero padding in the normalization length,

examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml

@@ -52,6 +52,7 @@ generator:
         sp_pooled_key_mean: true
         sp_direct_keys: true
         sp_projection_overlap: true
+        sp_input_gather: true
       VSA_sparsity: 0.8
       VSA_tile_size: 64
       h3_sequential_load: false

fastvideo/distributed/projection_exchange.py

@@ -99,3 +99,22 @@ class ProjectionExchange:
         return [(PendingInput(self.payload[c], self.remote[c], self.rank,
                               self.completion), PendingScale(self.scale[c], self.received_scale[c], self.completion))
                 for c in range(self.payload.shape[0])]
+
+
+def completed_projection_packets(payload: torch.Tensor, scale: torch.Tensor,
+                                 rank: int) -> list[tuple[PendingInput, PendingScale]]:
+    """Expose already source-ordered QKV packets through the usual consumers.
+
+    Input is [head groups, source ranks, QKV, rows, heads, 128]. An event orders
+    consumers on other CUDA streams after production; handles retain both
+    source slices and scales throughout their consumption.
+    """
+    if (payload.ndim != 6 or payload.shape[1:3] != (2, 3) or payload.shape[-1] != 128 or rank not in (0, 1)
+            or not payload.is_cuda or not payload.is_contiguous() or payload.dtype != torch.float8_e4m3fn
+            or scale.shape != payload.shape[:-1] or scale.device != payload.device or scale.dtype != torch.float32
+            or not scale.is_contiguous()):
+        raise ValueError("Completed projection packets require two contiguous source-rank slices")
+    completion = Completion(payload.device)
+    completion.ready.record()
+    return [(PendingInput(payload[c], payload[c, 1 - rank].view(torch.uint8), rank,
+                          completion), PendingScale(scale[c], scale[c], completion)) for c in range(payload.shape[0])]

fastvideo/fastvideo_args.py

@@ -448,12 +448,12 @@ class FastVideoArgs:
                 self.attention_backend = env_backend
         allowed_fp4 = {
             "enabled", "prefix_precision", "two_level_p", "smooth_k", "sp_head_chunks", "sp_pooled_key_mean",
-            "sp_projection_overlap", "sp_direct_keys"
+            "sp_projection_overlap", "sp_direct_keys", "sp_input_gather"
         }
         if not isinstance(self.h3_vsa_fp4, dict) or self.h3_vsa_fp4.keys() - allowed_fp4:
             raise ValueError(f"h3_vsa_fp4 must be a mapping with keys from {sorted(allowed_fp4)}")
         for key in ("enabled", "two_level_p", "smooth_k", "sp_pooled_key_mean", "sp_projection_overlap",
-                    "sp_direct_keys"):
+                    "sp_direct_keys", "sp_input_gather"):
             if key in self.h3_vsa_fp4 and not isinstance(self.h3_vsa_fp4[key], bool):
                 raise ValueError(f"h3_vsa_fp4.{key} must be a boolean")
         head_chunks = self.h3_vsa_fp4.get("sp_head_chunks", 1)

fastvideo/models/dits/minimax_h3_fusions/sp_head_pipeline.py

@@ -113,7 +113,15 @@ def attention_pipelined(attn, hidden_states, rotary_emb, meta, sp_group, chunks,
                    and hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
                    and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
                    and base._shared_qkv_gate_scale(attn, hidden_states) is not None)
-    if early_sends:
+    input_gather = False
+    if early_sends and options.get('sp_input_gather', False):
+        from .sp_input_gather import supports_input_gather
+        input_gather = supports_input_gather(attn, hidden_states)
+    if input_gather:
+        from .sp_input_gather import project_gathered_inputs
+        with base.STAGES.span('pipe_project_pack'):
+            packets, gate = project_gathered_inputs(attn, hidden_states, rotary_emb, sp_group, chunks, pre_quantized)
+    elif early_sends:
         from .sp_projection import project_and_exchange
         with base.STAGES.span('pipe_project_pack'):
             packets, gate_pre = project_and_exchange(attn, hidden_states, rotary_emb, sp_group.device_group,
@@ -142,12 +150,13 @@ def attention_pipelined(attn, hidden_states, rotary_emb, meta, sp_group, chunks,
                 chunks,
             )
         packets = [_start_exchange(payload[chunk], scale[chunk], sp_group.device_group) for chunk in range(chunks)]
-    with base.STAGES.span('gate_proj'):
-        if gate_pre is None:
-            gate, _ = attn.to_gate_compress(hidden_states)
-        else:
-            linear = attn.to_gate_compress
-            gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
+    if not input_gather:
+        with base.STAGES.span('gate_proj'):
+            if gate_pre is None:
+                gate, _ = attn.to_gate_compress(hidden_states)
+            else:
+                linear = attn.to_gate_compress
+                gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
     vbs = meta.variable_block_sizes
     api = base._api()
     pending = []

fastvideo/models/dits/minimax_h3_fusions/sp_input_gather.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Gather exact packed inputs before head-sharded two-rank H3 projections.
+
+The packed inputs already carry the shared unit activation scale. Only weight
+views are sharded; each row/head computes the same Q/K/V projection, norm/RoPE,
+and transport FP8 rounding as the sequence-sharded producer. Scale reassembly
+handles the unaligned join between independently padded 128-row scale tiles.
+"""
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import all_gather_tensor_async
+from fastvideo.distributed.projection_exchange import completed_projection_packets
+from fastvideo.layers.quantization.nvfp4_config import _mm_fp4, _mm_fp4_backend
+
+from .sp_projection import allocate
+from .sp_qkv_ordered import _normalize
+
+
+@triton.jit
+def _join_scale_rows(Source, Out, ROWS: tl.constexpr, PADDED: tl.constexpr, GROUPS: tl.constexpr,
+                     TOTAL_PADDED: tl.constexpr, BLOCK: tl.constexpr):
+    index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    row, group = index // GROUPS, index % GROUPS
+    rank, local_row = row // ROWS, row % ROWS
+    local_tile = (local_row // 128) * (GROUPS // 4) + group // 4
+    local_index = ((local_tile * 32 + local_row % 32) * 4 + (local_row // 32) % 4) * 4 + group % 4
+    value = tl.load(Source + rank * PADDED * GROUPS + local_index, row < 2 * ROWS, other=0)
+    tile = (row // 128) * (GROUPS // 4) + group // 4
+    destination = ((tile * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4
+    tl.store(Out + destination, value, index < TOTAL_PADDED * GROUPS)
+
+
+def join_scale_rows(source: torch.Tensor, rows: int) -> torch.Tensor:
+    assert source.ndim == 3 and source.shape[0] == 2 and source.is_contiguous()
+    assert source.dtype == torch.uint8 and source.shape[2] % 4 == 0
+    padded, groups = source.shape[1:]
+    assert padded == triton.cdiv(rows, 128) * 128
+    total_padded = triton.cdiv(rows * 2, 128) * 128
+    output = torch.empty((total_padded, groups), device=source.device, dtype=source.dtype)
+    _join_scale_rows[(triton.cdiv(total_padded * groups, 512), )](
+        source, output, rows, padded, groups, total_padded, 512, num_warps=4)
+    return output
+
+
+@triton.jit
+def _pack_gathered(X, Weight, Cos, Sin, Payload, Scale, LOCAL_ROWS: tl.constexpr, HEADS: tl.constexpr,
+                   CHUNKS: tl.constexpr, ROTARY: tl.constexpr, EPS: tl.constexpr, PROJECTION: tl.constexpr,
+                   BLOCK_ROWS: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
+    row, head = row_head // HEADS, row_head % HEADS
+    column = tl.arange(0, 128)
+    valid = row[:, None] < LOCAL_ROWS * 2
+    if PROJECTION == 2:
+        values = tl.load(X + row_head[:, None] * 128 + column[None, :], valid, other=0.0).to(tl.float32)
+    else:
+        values = _normalize(X, Weight, Cos, Sin, row, head, HEADS, LOCAL_ROWS * 2, 128, ROTARY, EPS, 128)
+    maximum = tl.max(tl.abs(values), 1)
+    scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    chunk_heads: tl.constexpr = HEADS // CHUNKS
+    chunk = head // chunk_heads
+    source = row // LOCAL_ROWS
+    destination = ((
+        (chunk * 2 + source) * 3 + PROJECTION) * LOCAL_ROWS + row % LOCAL_ROWS) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * 128 + column[None, :], quantized, valid)
+    tl.store(Scale + destination, scale, row < LOCAL_ROWS * 2)
+
+
+def pack_gathered(x, payload, scale, projection, weight=None, cosine=None, sine=None, epsilon=0.0) -> None:
+    chunks, world, projections, rows, heads, dim = payload.shape
+    assert world == 2 and projections == 3 and dim == 128
+    assert x.shape == (2 * rows, chunks * heads * dim) and x.is_contiguous()
+    assert payload.is_contiguous() and scale.is_contiguous()
+    if projection == 2:
+        weight, cosine, sine = x, x, x
+        rotary = 0
+    else:
+        assert weight.shape == (128, ) and cosine.shape == sine.shape
+        assert cosine.shape[0] == rows * 2 and cosine.is_contiguous() and sine.is_contiguous()
+        rotary = cosine.shape[1]
+    block_rows = 16 if projection == 2 else 8
+    _pack_gathered[(triton.cdiv(2 * rows * chunks * heads, block_rows), )](
+        x, weight, cosine, sine, payload, scale, rows, chunks * heads, chunks,
+        rotary, epsilon, projection, block_rows, num_warps=4 if projection == 2 else 8)
+
+
+def supports_input_gather(attn, hidden: torch.Tensor) -> bool:
+    """Keep other shapes/backends on the established projection exchange."""
+    if (not hidden.is_cuda or hidden.dtype != torch.bfloat16
+            or hidden.shape != (1, 7430, 5376) or attn.num_attention_heads != 56
+            or torch.cuda.get_device_capability(hidden.device) != (12, 1)
+            or _mm_fp4_backend() != "cudnn"):
+        return False
+    for linear in (attn.to_q, attn.to_k, attn.to_v):
+        weight = getattr(linear, "_nvfp4_weight", None)
+        scale = getattr(linear, "_nvfp4_weight_scale", None)
+        if (weight is None or scale is None or weight.shape != (7168, 2688)
+                or scale.shape != (7168, 336) or linear.bias is not None
+                or weight.dtype != torch.uint8 or scale.dtype != torch.uint8
+                or weight.device != hidden.device or scale.device != hidden.device
+                or not weight.is_contiguous() or not scale.is_contiguous()):
+            return False
+    return True
+
+
+def project_gathered_inputs(attn, hidden, rotary, sp_group, chunks, pre_quantized=None):
+    """Return completed head packets and the gate computed while inputs travel.
+
+    The caller checks eager, two-rank, head-128 inference and the common unit
+    activation-scale contract through _shared_qkv_gate_scale. No activations
+    or packets survive this call except through the returned consumer handles.
+    """
+    linears = (attn.to_q, attn.to_k, attn.to_v)
+    pre = pre_quantized if pre_quantized is not None else linears[0].quant_method.quantize_input(hidden)
+    a, sa, input_scale = pre
+    a, sa = a.reshape(-1, a.shape[-1]), sa.reshape(-1, sa.shape[-1])
+    cosine, sine = (item.to(hidden.dtype) for item in rotary)
+    pending = [all_gather_tensor_async(value, sp_group.device_group, 2) for value in (a, sa, cosine, sine)]
+    linear = attn.to_gate_compress
+    gate = linear.quant_method.apply(linear, hidden, linear.bias, pre_quantized=pre)
+    gathered_a, gathered_sa, gathered_cosine, gathered_sine = (item.wait() for item in pending)
+    full_a = gathered_a.flatten(0, 1)
+    full_sa = join_scale_rows(gathered_sa, a.shape[0])
+    full_cosine, full_sine = gathered_cosine.flatten(0, 1), gathered_sine.flatten(0, 1)
+    rows, heads, rank = hidden.shape[1], attn.num_attention_heads, sp_group.rank_in_group
+    payload, scale = allocate(rows, heads, 2, chunks, hidden.device)
+    for projection, linear in enumerate(linears):
+        half = linear._nvfp4_weight.shape[0] // 2
+        weight = linear._nvfp4_weight.narrow(0, rank * half, half)
+        weight_scale = linear._nvfp4_weight_scale.narrow(0, rank * half, half)
+        alpha = (linear._nvfp4_alpha / input_scale if hasattr(linear, "_nvfp4_alpha")
+                 else 1.0 / (input_scale * linear._weight_global_sf))
+        value = _mm_fp4(full_a, weight.T, full_sa, weight_scale.T, alpha, torch.bfloat16, None, backend="cudnn")
+        if projection == 2:
+            pack_gathered(value, payload, scale, projection)
+        else:
+            norm = attn.norm_q if projection == 0 else attn.norm_k
+            pack_gathered(value, payload, scale, projection, norm.weight, full_cosine, full_sine, norm.eps)
+    return completed_projection_packets(payload, scale, rank), gate

fastvideo/tests/transformers/test_minimax_h3_sp_input_gather.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact FP4 input gathering and source-ordered head packets on Blackwell."""
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and Triton")
+
+
+def _swizzle(logical):
+    rows, groups = logical.shape
+    return logical.reshape(rows // 128, 4, 32, groups // 4, 4).permute(0, 3, 2, 1, 4).contiguous().reshape(rows, groups)
+
+
+@pytest.mark.parametrize("rows", [1, 127, 128, 129, 7430])
+@pytest.mark.parametrize("groups", [4, 336])
+def test_join_scale_rows_preserves_unaligned_sources_and_zeros_padding(rows, groups):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_input_gather import join_scale_rows
+    torch.manual_seed(217)
+    padded, total = (rows + 127) // 128 * 128, (2 * rows + 127) // 128 * 128
+    logical = torch.randint(0, 256, (2, padded, groups), dtype=torch.uint8, device="cuda")
+    source = torch.stack([_swizzle(part) for part in logical])
+    expected = torch.zeros((total, groups), dtype=torch.uint8, device="cuda")
+    expected[:2 * rows] = logical[:, :rows].reshape(2 * rows, groups)
+    assert torch.equal(join_scale_rows(source, rows), _swizzle(expected))
+
+
+@pytest.mark.parametrize("rows", [1, 127, 7430])
+@pytest.mark.parametrize("projection", [0, 1, 2])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_gathered_head_packets_match_original_exchange(rows, projection, rank):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_input_gather import pack_gathered
+    from fastvideo.models.dits.minimax_h3_fusions.sp_projection import allocate, pack_one, pack_value
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("ordered QK packing requires Blackwell")
+    torch.manual_seed(29176)
+    values = torch.randn((2, rows, 56, 128), device="cuda", dtype=torch.bfloat16)
+    values[:, ::127] = 0
+    weight = torch.randn((128,), device="cuda", dtype=torch.bfloat16)
+    weight[::17] = 0
+    cosine, sine = [torch.randn((2, rows, 96), device="cuda", dtype=torch.bfloat16) for _ in range(2)]
+    gathered = values[:, :, rank * 28:(rank + 1) * 28].contiguous().reshape(2 * rows, 3584)
+    payload, scale = allocate(rows, 56, 2, 7, values.device)
+    pack_gathered(gathered, payload, scale, projection, weight, cosine.flatten(0, 1), sine.flatten(0, 1), 1e-5)
+    for source in range(2):
+        expected, expected_scale = allocate(rows, 56, 2, 7, values.device)
+        if projection == 2:
+            pack_value(values[source], expected, expected_scale)
+        else:
+            pack_one(values[source], weight, cosine[source], sine[source], expected, expected_scale, 1e-5, projection)
+        assert torch.equal(payload[:, source, projection].contiguous().view(torch.uint8),
+                           expected[:, rank, projection].contiguous().view(torch.uint8))
+        assert torch.equal(scale[:, source, projection], expected_scale[:, rank, projection])
+
+
+def test_gathered_input_head_gemm_matches_sequence_projection():
+    from flashinfer import SfLayout, nvfp4_quantize
+    from fastvideo.layers.quantization.nvfp4_config import _mm_fp4
+    from fastvideo.models.dits.minimax_h3_fusions.sp_input_gather import join_scale_rows
+    if torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip("production route is restricted to GB10")
+    torch.manual_seed(62071)
+    rows, inner, columns = 7430, 5376, 7168
+    unit = torch.ones((), device="cuda", dtype=torch.float32)
+    x = torch.randn((2, rows, inner), device="cuda", dtype=torch.bfloat16)
+    x[1] *= 8
+    x[:, ::127] = 0
+    w = torch.randn((columns, inner), device="cuda", dtype=torch.bfloat16) / inner**.5
+    global_scale = torch.full((), 13721.5, device="cuda", dtype=torch.float32)
+    b, sb = nvfp4_quantize(w, global_scale, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
+    inputs = [nvfp4_quantize(part, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False) for part in x]
+    full_a = torch.cat([p[0] for p in inputs])
+    full_sa = join_scale_rows(torch.stack([p[1] for p in inputs]), rows)
+    expected = [_mm_fp4(a, b.T, sa, sb.T, 1.0 / global_scale, torch.bfloat16, None, backend="cudnn")
+                for a, sa in inputs]
+    for rank in (0, 1):
+        start, stop = rank * 3584, (rank + 1) * 3584
+        output = _mm_fp4(full_a, b[start:stop].T, full_sa, sb[start:stop].T,
+                         1.0 / global_scale, torch.bfloat16, None, backend="cudnn")
+        target = torch.cat([part[:, start:stop] for part in expected])
+        assert torch.equal(output.view(torch.int16), target.view(torch.int16))
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_completed_packets_order_other_stream_and_retain_source(rank):
+    from fastvideo.distributed.projection_exchange import completed_projection_packets
+    producer, consumer = torch.cuda.Stream(), torch.cuda.Stream()
+    with torch.cuda.stream(producer):
+        payload = torch.empty((7, 2, 3, 127, 4, 128), device="cuda", dtype=torch.float8_e4m3fn)
+        scale = torch.empty(payload.shape[:-1], device="cuda", dtype=torch.float32)
+        torch.cuda._sleep(5_000_000)
+        payload.view(torch.uint8).fill_(37)
+        scale.fill_(2.5)
+        packets = completed_projection_packets(payload, scale, rank)
+    with torch.cuda.stream(consumer):
+        received = [(p.wait().materialize(), s.wait().clone()) for p, s in packets]
+    consumer.synchronize()
+    del payload, scale
+    for payload, scale in received:
+        assert torch.all(payload.view(torch.uint8) == 37)
+        assert torch.all(scale == 2.5)
+
+
+def test_input_gather_guard_rejects_other_shapes_and_backend(monkeypatch):
+    from fastvideo.models.dits.minimax_h3_fusions import sp_input_gather
+    if torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip("production route is restricted to GB10")
+    hidden = torch.empty((1, 7430, 5376), device="cuda", dtype=torch.bfloat16)
+    linear = SimpleNamespace(_nvfp4_weight=torch.empty((7168, 2688), dtype=torch.uint8, device="cuda"),
+                             _nvfp4_weight_scale=torch.empty((7168, 336), dtype=torch.uint8, device="cuda"),
+                             bias=None)
+    attn = SimpleNamespace(num_attention_heads=56, to_q=linear, to_k=linear, to_v=linear)
+    monkeypatch.setattr(sp_input_gather, "_mm_fp4_backend", lambda: "cudnn")
+    assert sp_input_gather.supports_input_gather(attn, hidden)
+    assert not sp_input_gather.supports_input_gather(attn, hidden[:, :-1])
+    monkeypatch.setattr(sp_input_gather, "_mm_fp4_backend", lambda: "cutlass")
+    assert not sp_input_gather.supports_input_gather(attn, hidden)
+    monkeypatch.setattr(sp_input_gather, "_mm_fp4_backend", lambda: "cudnn")
+    linear.bias = torch.zeros((7168,), device="cuda")
+    assert not sp_input_gather.supports_input_gather(attn, hidden)