376e445f0e
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
58 lines
2.5 KiB
Python
58 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Row chunking for the NVFP4 mm_fp4 wrapper: chunks must tile the rows exactly, on 128-row boundaries."""
|
|
import math
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.layers.quantization import nvfp4_config
|
|
from fastvideo.layers.quantization.nvfp4_config import _mm_fp4_row_chunks
|
|
|
|
|
|
@pytest.mark.parametrize("rows", [35328, 35000, 16384, 20481])
|
|
@pytest.mark.parametrize("parts", [1, 2, 4])
|
|
def test_row_chunks_cover_rows_in_order(rows: int, parts: int) -> None:
|
|
a = torch.arange(rows * 2, dtype=torch.float32).view(rows, 2)
|
|
scale = torch.arange(rows + (-rows) % 128, dtype=torch.float32).view(-1, 1)
|
|
chunks = _mm_fp4_row_chunks(a, scale, parts)
|
|
assert torch.equal(torch.cat([c for c, _ in chunks]), a)
|
|
# every chunk start is a 128-row scale tile boundary
|
|
offset = 0
|
|
for chunk, chunk_scale in chunks:
|
|
assert offset % 128 == 0
|
|
assert chunk_scale[0, 0].item() == offset
|
|
scale_hi = min(math.ceil((offset + chunk.shape[0]) / 128) * 128, scale.shape[0])
|
|
assert chunk_scale.shape[0] == scale_hi - offset
|
|
offset += chunk.shape[0]
|
|
assert offset == rows
|
|
|
|
|
|
@pytest.mark.parametrize(("use_8x4", "a_scale_rows"), [(True, 20000), (False, None)])
|
|
def test_row_split_stays_whole_without_128x4_scale_rows(monkeypatch, use_8x4, a_scale_rows):
|
|
calls = []
|
|
|
|
def fake_chunked(*args):
|
|
calls.append(args[-1])
|
|
return None
|
|
|
|
monkeypatch.setattr(nvfp4_config, "_mm_fp4_chunked", fake_chunked)
|
|
a = torch.zeros(256, 64)
|
|
a_scale = torch.zeros(256, 1) if a_scale_rows else torch.zeros(256)
|
|
assert nvfp4_config._mm_fp4_row_split(a, torch.zeros(32, 64), a_scale, torch.zeros(32, 1), None,
|
|
torch.bfloat16, 16, use_8x4, "cudnn", True) is None
|
|
assert calls == [1]
|
|
|
|
|
|
@pytest.mark.parametrize("policy,expected", [("off", "whole"), ("table", "split")])
|
|
def test_large_cudnn_call_splits_only_when_opted_in(monkeypatch, policy, expected):
|
|
from fastvideo import envs
|
|
|
|
calls = []
|
|
monkeypatch.setattr(nvfp4_config, "_register_ops_once", lambda: None)
|
|
monkeypatch.setattr(torch.ops.fastvideo_fp4, "mm_fp4", lambda *args: calls.append("whole"), raising=False)
|
|
monkeypatch.setattr(nvfp4_config, "_mm_fp4_row_split", lambda *args: calls.append("split"))
|
|
with envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.override(policy):
|
|
nvfp4_config._mm_fp4(torch.empty(16384, 32), torch.empty(16, 64), torch.empty(16384, 1),
|
|
torch.empty(16, 1), None, torch.bfloat16, None, backend="cudnn")
|
|
assert calls == [expected]
|