diff --git a/docs/getting_started/installation/spark_pair.md b/docs/getting_started/installation/spark_pair.md index 024684aa1..f3c897c47 100644 --- a/docs/getting_started/installation/spark_pair.md +++ b/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, diff --git a/examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml index 8a4890534..ab1616e19 100644 --- a/examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml +++ b/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 diff --git a/fastvideo/distributed/projection_exchange.py b/fastvideo/distributed/projection_exchange.py index 59beba0d2..6a7d29832 100644 --- a/fastvideo/distributed/projection_exchange.py +++ b/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])] diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 7dc7daa14..007ddf84c 100644 --- a/fastvideo/fastvideo_args.py +++ b/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) diff --git a/fastvideo/models/dits/minimax_h3_fusions/sp_head_pipeline.py b/fastvideo/models/dits/minimax_h3_fusions/sp_head_pipeline.py index 0e3ffdd88..2484e6c15 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/sp_head_pipeline.py +++ b/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 = [] diff --git a/fastvideo/models/dits/minimax_h3_fusions/sp_input_gather.py b/fastvideo/models/dits/minimax_h3_fusions/sp_input_gather.py new file mode 100644 index 000000000..0bbd62566 --- /dev/null +++ b/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 diff --git a/fastvideo/tests/transformers/test_minimax_h3_sp_input_gather.py b/fastvideo/tests/transformers/test_minimax_h3_sp_input_gather.py new file mode 100644 index 000000000..5e26f79d6 --- /dev/null +++ b/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)