diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 8a9a5e29e..eeb02fb27 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/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. diff --git a/fastvideo/tests/attention/test_nvfp4_fa4_pad.py b/fastvideo/tests/attention/test_nvfp4_fa4_pad.py new file mode 100644 index 000000000..9fb68f660 --- /dev/null +++ b/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)