[perf] FA4 FP4 attention: zero only the pad rows before NVFP4 quantization

fastvideo/attention/backends/flash_attn.py

@@ -104,7 +104,11 @@ def _nvfp4_quantize_for_fa4_impl(tensor_4d: torch.Tensor) -> tuple[torch.Tensor,
     tile_m = 128
     seqlen_padded = (seqlen + tile_m - 1) // tile_m * tile_m
     if seqlen_padded != seqlen:
-        tensor_4d = F.pad(tensor_4d, (0, 0, 0, 0, 0, seqlen_padded - seqlen))
+        # F.pad fills the whole padded tensor before copying; only the tail rows need zeros.
+        padded = tensor_4d.new_empty((batch, seqlen_padded, nheads, headdim))
+        padded[:, :seqlen] = tensor_4d
+        padded[:, seqlen:].zero_()
+        tensor_4d = padded
 
     # Quantize with nheads squashed into K dimension so M=batch*seqlen (divisible by 128)
     # and K=nheads*headdim. This ensures 128-row SF tiles align with seqlen boundaries.

fastvideo/tests/attention/test_nvfp4_fa4_pad.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4 quantization needs CUDA")
+
+
+@pytest.mark.parametrize("seqlen", [77, 300])
+def test_fa4_quantize_pads_like_zero_padding(seqlen: int) -> None:
+    flash_attn = pytest.importorskip("fastvideo.attention.backends.flash_attn")
+    try:
+        flash_attn._resolve_fa4_quant_ops()
+    except Exception as exc:  # noqa: BLE001
+        pytest.skip(f"NVFP4 quantize ops unavailable: {exc}")
+    torch.manual_seed(0)
+    x = torch.randn(1, seqlen, 4, 128, device="cuda", dtype=torch.bfloat16)
+    padded = F.pad(x, (0, 0, 0, 0, 0, (-seqlen) % 128))
+
+    fp4, sf = flash_attn._nvfp4_quantize_for_fa4_impl(x)
+    ref_fp4, ref_sf = flash_attn._nvfp4_quantize_for_fa4_impl(padded)
+
+    assert torch.equal(fp4.view(torch.uint8), ref_fp4.view(torch.uint8))
+    assert torch.equal(sf, ref_sf)