[perf] Fuse VSA64 tile layout for SM100 inference (#1900)

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.
2026-10-07 15:31:58 -07:00
parent e5f31c1b4c
commit ef63de1f60
9 changed files with 491 additions and 16 deletions
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused VSA64 tile scatter and BSHD->BHSD layout conversion for SM100 inference."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _tile_to_bhsd_kernel(
+    x_ptr,
+    source_index_ptr,
+    output_ptr,
+    sequence,
+    padded_sequence,
+    heads,
+    head_dim: tl.constexpr,
+    stride_batch,
+    stride_sequence,
+    stride_head,
+    stride_dim,
+    block: tl.constexpr,
+):
+    pid0 = tl.program_id(0).to(tl.int64)
+    batch_head = tl.program_id(1).to(tl.int64)
+    linear = pid0 * block + tl.arange(0, block).to(tl.int64)
+    batch = batch_head // heads
+    head = batch_head % heads
+    padded_pos = linear // head_dim
+    dim = linear % head_dim
+    valid = linear < padded_sequence * head_dim
+    source_pos = tl.load(source_index_ptr + padded_pos, mask=valid, other=-1).to(tl.int64)
+    x_offset = (batch * stride_batch + source_pos * stride_sequence + head * stride_head + dim * stride_dim)
+    value = tl.load(x_ptr + x_offset, mask=valid & (source_pos >= 0) & (source_pos < sequence), other=0)
+    out_offset = batch_head * padded_sequence * head_dim + linear
+    tl.store(output_ptr + out_offset, value, mask=valid)
+
+
+def tile_to_bhsd(
+    x: torch.Tensor,
+    source_index: torch.Tensor,
+    out: torch.Tensor,
+    sequence: int,
+    padded_sequence: int,
+    heads: int,
+    head_dim: int,
+) -> torch.Tensor:
+    """Scatter-tile ``x`` [B, S, H, D] into ``out`` [B, H, padded_S, D] in BHSD order."""
+    block = 1024
+    grid = (triton.cdiv(padded_sequence * head_dim, block), x.shape[0] * heads)
+    _tile_to_bhsd_kernel[grid](
+        x,
+        source_index,
+        out,
+        sequence,
+        padded_sequence,
+        heads,
+        head_dim,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        x.stride(3),
+        block,
+        num_warps=4,
+    )
+    return out