diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index 9f1a87dd7..3c69e61a8 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -246,10 +246,12 @@ in a local directory. It declares the trained eight-forward ladder and video/audio shifts of 10/3. The recipes use `num_inference_steps: 9` for nine sigma points and eight DiT forwards. -On GB10 with FlashInfer 0.6.18, FastVideo fences activation quantization before -releasing its padded input. Without this completion fence, identical H3 -requests produced different DiT latents and occasionally corrupt video. -The fence applies to `sm_121`; other architectures retain asynchronous execution. +On GB10 (FlashInfer 0.6.18 and 0.7.0), FastVideo completes each NVFP4 +activation quantization on the host before its GEMM. Without this completion +fence, identical H3 requests produced different DiT latents and occasionally +corrupt video. The fence costs no measurable denoise time, because the GPU stays +saturated. It applies to `sm_121`; other architectures retain asynchronous +execution. Run `examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml` from the repository root. It uses 832x480, 124 frames and seed 1234, VSA sparsity 0.8 with diff --git a/fastvideo/layers/quantization/nvfp4_config.py b/fastvideo/layers/quantization/nvfp4_config.py index 17460f46f..e7d6b14d8 100644 --- a/fastvideo/layers/quantization/nvfp4_config.py +++ b/fastvideo/layers/quantization/nvfp4_config.py @@ -203,9 +203,10 @@ def nvfp4_quantize_fenced( sfLayout=SfLayout(sf_layout), do_shuffle=do_shuffle, enable_pdl=False) - # With FlashInfer 0.6.18 on GB10, queued activation quantization - # plus GEMM can diverge. Completing quantization while its padded - # input is alive prevents the observed intermittent corruption. + # With FlashInfer 0.6.18 and 0.7.0 on GB10, queued activation + # quantization plus GEMM can diverge: identical seeded H3 requests + # produce different latents. Completing each quantization on the + # host prevents it, at no measurable denoise cost. torch.cuda.current_stream(x.device).synchronize() return quantized, scales @@ -329,21 +330,24 @@ def _nvfp4_quantize( elif global_sf.device != x.device: global_sf = global_sf.to(device=x.device) if sf_layout == SfLayout.layout_linear.value: - x_for_quant = x - logical_rows = x.shape[0] - else: - # Sequence-parallel can feed either logical rows or row-padded - # rows. Normalize to the kernel tile shape for swizzled layouts - # so both paths share a stable quantization contract. - row_tile = 8 if sf_layout == SfLayout.layout_8x4.value else 128 - logical_rows = x.shape[0] - pad_rows = (-logical_rows) % row_tile - x_for_quant = F.pad(x, (0, 0, 0, pad_rows)) - + return torch.ops.fastvideo_fp4.nvfp4_quantize(x, global_sf, sf_layout, do_shuffle) + # Sequence-parallel can feed either logical rows or row-padded + # rows. Normalize to the kernel tile shape for swizzled layouts + # so both paths share a stable quantization contract. + row_tile = 8 if sf_layout == SfLayout.layout_8x4.value else 128 + logical_rows = x.shape[0] + padded_rows = logical_rows + (-logical_rows) % row_tile + if sf_layout == SfLayout.layout_128x4.value: + # FlashInfer pads 128x4 scale factors to whole row tiles itself, and + # quantizes the logical rows bit for bit like a zero-padded copy, so + # skip that copy (a full activation write per linear) when the + # returned scales already span the padded tile. + quantized, scales = torch.ops.fastvideo_fp4.nvfp4_quantize(x, global_sf, sf_layout, do_shuffle) + if scales.numel() >= padded_rows * (x.shape[-1] // 16): + return quantized, scales + x_for_quant = F.pad(x, (0, 0, 0, padded_rows - logical_rows)) quantized, scales = torch.ops.fastvideo_fp4.nvfp4_quantize(x_for_quant, global_sf, sf_layout, do_shuffle) - if sf_layout != SfLayout.layout_linear.value: - quantized = quantized.narrow(0, 0, logical_rows) - return quantized, scales + return quantized.narrow(0, 0, logical_rows), scales def _mm_fp4( diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py b/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py index fda1aedc8..f70801fb5 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py @@ -48,6 +48,8 @@ def _spy_sync(monkeypatch): return syncs +@pytest.mark.skipif(torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 1), + reason="needs a non-GB10 Blackwell GPU") def test_fenced_quantize_matches_flashinfer_without_a_fence_off_spark(monkeypatch) -> None: real, pdl = _spy_quantize(monkeypatch) syncs = _spy_sync(monkeypatch) @@ -94,3 +96,18 @@ def test_qat_linear_quantizes_through_the_fence_with_unchanged_numerics(monkeypa reference = torch.nn.functional.linear(x.float(), weight.float(), bias.float()) relative = (on_spark.float() - reference).norm() / reference.norm() assert relative < 0.15 + + +@pytest.mark.parametrize("rows", [200, 256]) +def test_swizzled_quantize_skips_the_pad_copy_bit_exactly(monkeypatch, rows: int) -> None: + torch.manual_seed(2) + x = torch.randn(rows, 256, device="cuda", dtype=torch.bfloat16) + global_sf = (448.0 * 6.0) / x.float().abs().amax() + expected = flashinfer.nvfp4_quantize(_pad_rows(x), global_sf, sfLayout=LAYOUT_128X4, do_shuffle=False) + pads = [] + real_pad = torch.nn.functional.pad + monkeypatch.setattr(nv.F, "pad", lambda *args, **kwargs: pads.append(args) or real_pad(*args, **kwargs)) + quantized, scales = nv._nvfp4_quantize(x, global_sf, sfLayout=LAYOUT_128X4) + assert pads == [] + assert quantized.shape[0] == rows + assert torch.equal(quantized, expected[0][:rows]) and torch.equal(scales, expected[1])