Wenbo Ji 嵇文博
ef63de1f60
For Wan VSA inference with 64-token tiles, each attention layer scattered Q, K, V and the gate into a tile buffer and then made four transposed copies. On SM100, one Triton kernel (tile_to_bhsd) now does the scatter and the transposes in one pass. Training, 256-token tiles, other GPUs and builds without the kernel use the previous path. FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT=1 turns the fused path off. On GB200 the layout and attention step goes from 15.7 ms to 8.6 ms. FastWan 1.3B frames are bit-exact with the previous path at sequence parallel 1 and 2.
46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Bitwise parity for the fused VSA64 tile-to-BHSD Triton kernel.
|
|
|
|
Production only routes SM100 here, but the kernel is plain Triton, so parity is
|
|
checked on any CUDA device.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
pytest.importorskip("triton")
|
|
|
|
from fastvideo_kernel.triton_kernels.vsa_tile_layout import tile_to_bhsd # noqa: E402
|
|
|
|
|
|
@pytest.mark.parametrize("dit_shape", [(5, 5, 6), (16, 28, 52)])
|
|
def test_tile_to_bhsd_matches_scatter_and_four_transposes(dit_shape):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA device required")
|
|
|
|
from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionImpl,
|
|
VideoSparseAttentionMetadataBuilder)
|
|
|
|
impl = object.__new__(VideoSparseAttentionImpl)
|
|
metadata = VideoSparseAttentionMetadataBuilder().build(
|
|
current_timestep=0,
|
|
raw_latent_shape=dit_shape,
|
|
patch_size=(1, 1, 1),
|
|
VSA_sparsity=0.8,
|
|
device=torch.device("cuda"),
|
|
cache_tile_buf=True,
|
|
)
|
|
sequence = metadata.total_seq_length
|
|
padded_sequence = metadata.variable_block_sizes.numel() * 64
|
|
source = torch.randn((4, sequence + 7, 2, 128), device="cuda", dtype=torch.bfloat16)
|
|
qkvg = source[:, :sequence]
|
|
expected = tuple(piece.transpose(1, 2).contiguous() for piece in impl.tile(qkvg, metadata).chunk(4))
|
|
|
|
source_index = torch.full((padded_sequence, ), -1, device="cuda", dtype=torch.int32)
|
|
source_index[metadata.non_pad_index] = metadata.tile_partition_indices.to(torch.int32)
|
|
out = torch.empty((4, 2, padded_sequence, 128), device="cuda", dtype=torch.bfloat16)
|
|
actual = tile_to_bhsd(qkvg, source_index, out, sequence, padded_sequence, 2, 128).chunk(4)
|
|
for name, old, new in zip(("q", "k", "v", "gate"), expected, actual):
|
|
assert old.shape == new.shape, name
|
|
assert torch.equal(old, new), name
|