From d0b91810b2c8dff0107817f61a067b87aef2a8d4 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Tue, 6 Oct 2026 23:08:46 -0700 Subject: [PATCH] [perf]: skip the NVFP4 activation pad copy for 128x4 scale layouts _nvfp4_quantize copied every activation into a zero-padded buffer (F.pad copies even when no rows are added) before FlashInfer quantization. For the 128x4 swizzled layout FlashInfer already pads the scale factors to whole row tiles and quantizes the logical rows exactly as it would the padded copy, so quantize the input directly and keep the padded path as a fallback when the returned scales do not span the padded tile. On DGX Spark the pad accounted for about 2.7k device copies plus the matching fills per FastH3 Trim denoise (about 6 s of 62 s). Quantized bytes, scale factors and mm_fp4 outputs are bit-identical for the H3 projection shapes, and a FastH3 Trim 832x480x124 clip is bit-identical with denoise at 72 s instead of 78 s. The GB10 quantization fence stays: without it, identical seeded requests still diverge with or without the pad copy, on FlashInfer 0.6.18 and 0.7.0. Its comment and the Spark guide now describe it without the padded input. The off-Spark fence test skips on GB10, where the fence is expected. Co-Authored-By: Claude Opus 5.5 --- .../installation/spark_performance.md | 10 +++-- fastvideo/layers/quantization/nvfp4_config.py | 38 ++++++++++--------- .../quantization/test_nvfp4_spark_fence.py | 17 +++++++++ 3 files changed, 44 insertions(+), 21 deletions(-) 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])