nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[misc] yapf pass under CI's interpreter (3.12) + pin hook language_version (#1702)

Browse Source
This commit is contained in:
William Lin
2026-08-08 22:24:17 -07:00
committed by GitHub
parent 6eb95693a1
commit 9028953625
582 changed files with 19387 additions and 23526 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo-kernel/python/fastvideo_kernel/__init__.py
+1 -2
View File
@@ -26,8 +26,7 @@ from fastvideo_kernel.turbodiffusion_ops import (
)
from fastvideo_kernel.block_sparse_attn_varlen import (
block_sparse_attn_varlen,
)
block_sparse_attn_varlen, )
from fastvideo_kernel.vsa_utils import (
VSA_TILE_SIZE,
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py
+18 -49
View File
@@ -7,7 +7,6 @@ from typing import Tuple
import torch
# ---------------------------------------------------------------------------
# Backend selection helpers
# ---------------------------------------------------------------------------
@@ -37,10 +36,8 @@ def _force_triton() -> bool:
Honors both the new ``FASTVIDEO_VSA_TRITON`` and the legacy
``FASTVIDEO_KERNEL_VSA_FORCE_TRITON`` for backward compatibility.
"""
return (
os.environ.get("FASTVIDEO_VSA_TRITON", "0") == "1"
or os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1"
)
return (os.environ.get("FASTVIDEO_VSA_TRITON", "0") == "1"
or os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1")
def _force_tk() -> bool:
@@ -62,25 +59,19 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
if block_map.dim() == 3:
block_map = block_map.unsqueeze(0)
if block_map.dim() != 4:
raise ValueError(
f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), "
f"got shape={tuple(block_map.shape)}"
)
raise ValueError(f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), "
f"got shape={tuple(block_map.shape)}")
if block_map.dtype != torch.bool:
block_map = block_map.to(torch.bool)
if not block_map.is_cuda:
raise RuntimeError(
"block_map must be a CUDA tensor (Triton map_to_index required)."
)
raise RuntimeError("block_map must be a CUDA tensor (Triton map_to_index required).")
try:
from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index
except Exception as e: # pragma: no cover - environment issue
raise ImportError(
"Triton map_to_index is required but not available. "
"Ensure Triton is installed and "
"fastvideo_kernel.triton_kernels.index is importable."
) from e
raise ImportError("Triton map_to_index is required but not available. "
"Ensure Triton is installed and "
"fastvideo_kernel.triton_kernels.index is importable.") from e
return triton_map_to_index(block_map)
@@ -123,8 +114,7 @@ def block_sparse_attn_triton(
variable_block_sizes: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import (
triton_block_sparse_attn_forward,
)
triton_block_sparse_attn_forward, )
o, M = triton_block_sparse_attn_forward(
q.contiguous(),
@@ -172,13 +162,10 @@ def block_sparse_attn_backward_triton(
variable_block_sizes: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import (
triton_block_sparse_attn_backward,
)
triton_block_sparse_attn_backward, )
num_kv_blocks = int(variable_block_sizes.numel())
k2q_idx, k2q_num = _invert_indices_for_backward(
q2k_idx, q2k_num, num_kv_blocks
)
k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks)
# q/k/v are saved from the user-facing inputs and may be non-contiguous;
# o/M are kernel outputs so are already contiguous.
dq, dk, dv = triton_block_sparse_attn_backward(
@@ -223,16 +210,11 @@ def _setup_context_triton(ctx, inputs, output):
def _backward_triton(ctx, grad_o, grad_M):
q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors
dq, dk, dv = block_sparse_attn_backward_triton(
grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes
)
dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes)
return dq, dk, dv, None, None, None
block_sparse_attn_triton.register_autograd(
_backward_triton, setup_context=_setup_context_triton
)
block_sparse_attn_triton.register_autograd(_backward_triton, setup_context=_setup_context_triton)
# ---------------------------------------------------------------------------
# SM90 backend custom ops (index-native)
@@ -306,9 +288,7 @@ def block_sparse_attn_backward_sm90(
raise ImportError("fastvideo_kernel_ops.block_sparse_bwd is not available")
num_kv_blocks = int(variable_block_sizes.numel())
k2q_idx, k2q_num = _invert_indices_for_backward(
q2k_idx, q2k_num, num_kv_blocks
)
k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks)
# q/k/v are saved from user-facing inputs; o/lse are kernel outputs.
dq, dk, dv = block_sparse_bwd(
@@ -353,16 +333,11 @@ def _setup_context_sm90(ctx, inputs, output):
def _backward_sm90(ctx, grad_o, grad_lse):
q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors
dq, dk, dv = block_sparse_attn_backward_sm90(
grad_o, q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes
)
dq, dk, dv = block_sparse_attn_backward_sm90(grad_o, q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes)
return dq, dk, dv, None, None, None
block_sparse_attn_sm90.register_autograd(
_backward_sm90, setup_context=_setup_context_sm90
)
block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90)
# ---------------------------------------------------------------------------
# Public API
@@ -385,11 +360,7 @@ def block_sparse_attn_from_indices(
variable_block_sizes = _as_int32_contig(variable_block_sizes, "variable_block_sizes")
block_sparse_fwd, block_sparse_bwd = _get_sm90_ops()
sm90_available = (
_is_sm90()
and block_sparse_fwd is not None
and block_sparse_bwd is not None
)
sm90_available = (_is_sm90() and block_sparse_fwd is not None and block_sparse_bwd is not None)
# Backend resolution:
# - FASTVIDEO_VSA_TRITON forces Triton everywhere.
@@ -419,6 +390,4 @@ def block_sparse_attn(
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Bool-mask compat wrapper; prefer block_sparse_attn_from_indices."""
q2k_idx, q2k_num = _map_to_index(block_map)
return block_sparse_attn_from_indices(
q, k, v, q2k_idx, q2k_num, variable_block_sizes
)
return block_sparse_attn_from_indices(q, k, v, q2k_idx, q2k_num, variable_block_sizes)
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py
+5 -14
View File
@@ -30,7 +30,6 @@ from .block_sparse_attn import block_sparse_attn_triton, _force_triton
# FA4 CuTe build (``flash_attn.cute``) and make it a hard dependency of the
# default Triton path.
_KV_BLOCK_PHYS = 128 # FA4 CuTe BSA forward uses 128-token KV blocks.
_KV_BLOCK_TRITON = 64 # Existing Triton path uses 64-token KV blocks.
@@ -66,7 +65,7 @@ def _expand_mask_and_sizes_256_to_128(
child0 = torch.clamp(sizes_i32, min=0, max=_KV_BLOCK_PHYS)
child1 = torch.clamp(sizes_i32 - _KV_BLOCK_PHYS, min=0, max=_KV_BLOCK_PHYS)
expanded_sizes = torch.empty(
(sizes_i32.numel() * 2,),
(sizes_i32.numel() * 2, ),
dtype=torch.int32,
device=sizes_i32.device,
)
@@ -108,9 +107,7 @@ def _triton_via_route_a(
) -> Tuple[torch.Tensor, torch.Tensor]:
from .triton_kernels.index import map_to_index as triton_map_to_index
mask_64, sizes_64 = _expand_mask_and_sizes_256_to_64(
logical_mask_256, logical_kv_sizes_256
)
mask_64, sizes_64 = _expand_mask_and_sizes_256_to_64(logical_mask_256, logical_kv_sizes_256)
q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool))
return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
@@ -127,13 +124,9 @@ def block_sparse_attn_256(
logical_block_map_256 = logical_block_map_256.unsqueeze(0)
if _resolve_backend() == "triton":
return _triton_via_route_a(
q, k, v, logical_block_map_256, logical_variable_block_sizes_256
)
return _triton_via_route_a(q, k, v, logical_block_map_256, logical_variable_block_sizes_256)
mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(
logical_block_map_256, logical_variable_block_sizes_256
)
mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256)
from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
return block_sparse_attn_cute_fwd(q, k, v, mask_128, sizes_128)
@@ -163,8 +156,6 @@ def block_sparse_attn_256_bshd(
)
return out_bhsd.transpose(1, 2).contiguous(), aux
mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(
logical_block_map_256, logical_variable_block_sizes_256
)
mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256)
from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
return block_sparse_attn_cute_fwd_bshd(q, k, v, mask_128, sizes_128)
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py
+11 -20
View File
@@ -22,13 +22,11 @@ from typing import Tuple
import torch
_FA4_IMPORT_HINT = (
"VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
"provides `flash_attn.cute` with block-sparsity support (plus "
"`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional "
"dependency; the default VSA-256 path is Triton. Install the FA4 CuTe "
"build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath."
)
_FA4_IMPORT_HINT = ("VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
"provides `flash_attn.cute` with block-sparsity support (plus "
"`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional "
"dependency; the default VSA-256 path is Triton. Install the FA4 CuTe "
"build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath.")
def _load_fa4_cute():
@@ -54,17 +52,14 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
if block_map.dim() == 3:
block_map = block_map.unsqueeze(0)
if block_map.dim() != 4:
raise ValueError(
f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), "
f"got shape={tuple(block_map.shape)}"
)
raise ValueError(f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), "
f"got shape={tuple(block_map.shape)}")
if block_map.dtype != torch.bool:
block_map = block_map.to(torch.bool)
if not block_map.is_cuda:
raise RuntimeError("block_map must be a CUDA tensor.")
from fastvideo_kernel.triton_kernels.index import (
map_to_index as triton_map_to_index,
)
map_to_index as triton_map_to_index, )
return triton_map_to_index(block_map)
@@ -84,10 +79,8 @@ def _aggregate_q_block_map(
) -> torch.Tensor:
factor = q_sparse_block_size // q_block_size
if factor <= 0 or q_sparse_block_size % q_block_size != 0:
raise ValueError(
f"q_sparse_block_size must be a positive multiple of "
f"q_block_size ({q_block_size}), got {q_sparse_block_size}"
)
raise ValueError(f"q_sparse_block_size must be a positive multiple of "
f"q_block_size ({q_block_size}), got {q_sparse_block_size}")
bsz, nhead, q_blocks, kv_blocks = block_map.shape
q_blocks_sparse = (q_blocks + factor - 1) // factor
pad_q = q_blocks_sparse * factor - q_blocks
@@ -164,9 +157,7 @@ def _cute_forward(
q_block_size=q_block_size,
)
kv_full = (variable_block_sizes == kv_block_size).view(1, 1, 1, -1)
kv_partial = (
(variable_block_sizes > 0) & (variable_block_sizes < kv_block_size)
).view(1, 1, 1, -1)
kv_partial = ((variable_block_sizes > 0) & (variable_block_sizes < kv_block_size)).view(1, 1, 1, -1)
full_map = sparse_map & kv_full
mask_map = sparse_map & kv_partial
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_varlen.py
+48 -23
View File
@@ -41,9 +41,7 @@ def _scatter_to_padded(
for actual in sizes:
actual = min(actual, src_end - src_pos)
if actual > 0:
dst[:, :, dst_pos:dst_pos + actual, :] = (
src[src_pos:src_pos + actual].transpose(0, 1).unsqueeze(0)
)
dst[:, :, dst_pos:dst_pos + actual, :] = (src[src_pos:src_pos + actual].transpose(0, 1).unsqueeze(0))
src_pos += actual
dst_pos += block_size
@@ -68,9 +66,7 @@ def _gather_from_padded(
for actual in sizes:
actual = min(actual, dst_end - dst_pos)
if actual > 0:
dst[dst_pos:dst_pos + actual] = (
src[0, :, src_pos:src_pos + actual, :].transpose(0, 1)
)
dst[dst_pos:dst_pos + actual] = (src[0, :, src_pos:src_pos + actual, :].transpose(0, 1))
dst_pos += actual
src_pos += block_size
@@ -131,9 +127,7 @@ def block_sparse_attn_varlen(
if q_variable_block_sizes_list is not None:
q_vbs_resolved.append(q_variable_block_sizes_list[i])
else:
q_vbs_resolved.append(
torch.full((n_q_blocks,), block_size, dtype=torch.int32)
)
q_vbs_resolved.append(torch.full((n_q_blocks, ), block_size, dtype=torch.int32))
total_padded_q = sum(padded_q_lens)
total_padded_kv = sum(padded_kv_lens)
@@ -146,16 +140,31 @@ def block_sparse_attn_varlen(
kv_offset = 0
for i in range(num_seqs):
_scatter_to_padded(
q, q_vbs_resolved[i], block_size,
q_packed, q_offset, cu_q[i], cu_q[i + 1],
q,
q_vbs_resolved[i],
block_size,
q_packed,
q_offset,
cu_q[i],
cu_q[i + 1],
)
_scatter_to_padded(
k, variable_block_sizes_list[i], block_size,
k_packed, kv_offset, cu_kv[i], cu_kv[i + 1],
k,
variable_block_sizes_list[i],
block_size,
k_packed,
kv_offset,
cu_kv[i],
cu_kv[i + 1],
)
_scatter_to_padded(
v, variable_block_sizes_list[i], block_size,
v_packed, kv_offset, cu_kv[i], cu_kv[i + 1],
v,
variable_block_sizes_list[i],
block_size,
v_packed,
kv_offset,
cu_kv[i],
cu_kv[i + 1],
)
q_offset += padded_q_lens[i]
kv_offset += padded_kv_lens[i]
@@ -164,12 +173,19 @@ def block_sparse_attn_varlen(
max_kv_per_q = max(t.shape[-1] for t in q2k_idx_list)
global_q2k_idx = torch.zeros(
1, num_heads, total_q_blocks, max_kv_per_q,
dtype=torch.int32, device=device,
1,
num_heads,
total_q_blocks,
max_kv_per_q,
dtype=torch.int32,
device=device,
)
global_q2k_num = torch.zeros(
1, num_heads, total_q_blocks,
dtype=torch.int32, device=device,
1,
num_heads,
total_q_blocks,
dtype=torch.int32,
device=device,
)
global_vbs_parts = []
@@ -191,16 +207,25 @@ def block_sparse_attn_varlen(
global_vbs = torch.cat(global_vbs_parts, dim=0).to(torch.int32).contiguous()
out_packed, _ = block_sparse_attn_from_indices(
q_packed, k_packed, v_packed,
global_q2k_idx, global_q2k_num, global_vbs,
q_packed,
k_packed,
v_packed,
global_q2k_idx,
global_q2k_num,
global_vbs,
)
out = torch.zeros(cu_q[-1], num_heads, head_dim, device=device, dtype=dtype)
q_offset = 0
for i in range(num_seqs):
_gather_from_padded(
out_packed, q_vbs_resolved[i], block_size,
out, q_offset, cu_q[i], cu_q[i + 1],
out_packed,
q_vbs_resolved[i],
block_size,
out,
q_offset,
cu_q[i],
cu_q[i + 1],
)
q_offset += padded_q_lens[i]
fastvideo-kernel/python/fastvideo_kernel/ops.py
+30 -61
View File
@@ -27,9 +27,7 @@ def sliding_tile_attention(
) -> torch.Tensor:
# Check if the specific op is available
if sta_fwd is None:
return sliding_tile_attention_triton(
q, k, v, window_size, text_length, has_text, seq_shape
)
return sliding_tile_attention_triton(q, k, v, window_size, text_length, has_text, seq_shape)
seq_length = q.shape[2]
shape_map = {"30x48x80": 1, "36x48x48": 2, "18x48x80": 3}
@@ -53,11 +51,7 @@ def sliding_tile_attention(
k_h = k[:, head_idx:head_idx + 1].contiguous()
v_h = v[:, head_idx:head_idx + 1].contiguous()
o_h = torch.empty_like(q_h)
sta_fwd(
q_h, k_h,
v_h, o_h,
t, h, w, text_length, False, has_text, flag
)
sta_fwd(q_h, k_h, v_h, o_h, t, h, w, text_length, False, has_text, flag)
output[:, head_idx:head_idx + 1] = o_h
if has_text:
@@ -96,35 +90,27 @@ def video_sparse_attn(
if k.shape[0] != batch or v.shape[0] != batch or k.shape[1] != heads or v.shape[1] != heads:
raise ValueError("Expected q/k/v to have the same batch and head dimensions.")
if v.shape[2] != kv_seq_len:
raise ValueError(
f"Expected k and v to have the same sequence length, got "
f"k.shape[2]={kv_seq_len}, v.shape[2]={v.shape[2]}"
)
raise ValueError(f"Expected k and v to have the same sequence length, got "
f"k.shape[2]={kv_seq_len}, v.shape[2]={v.shape[2]}")
if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0:
raise ValueError(
f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, "
f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}"
)
raise ValueError(f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, "
f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}")
q_num_blocks = q_seq_len // block_elements
kv_num_blocks = kv_seq_len // block_elements
if variable_block_sizes.numel() != kv_num_blocks:
raise ValueError(
f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, "
f"got {variable_block_sizes.numel()}"
)
raise ValueError(f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, "
f"got {variable_block_sizes.numel()}")
if q_variable_block_sizes.numel() != q_num_blocks:
raise ValueError(
f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
f"got {q_variable_block_sizes.numel()}"
)
raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
f"got {q_variable_block_sizes.numel()}")
# Compression branch (fused Triton: bf16 read → fp32 accumulate → div → bf16 write)
q_c = fused_block_mean(q, q_variable_block_sizes, block_elements)
k_c = fused_block_mean(k, variable_block_sizes, block_elements)
v_c = fused_block_mean(v, variable_block_sizes, block_elements)
scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5)
scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5)
attn = torch.softmax(scores, dim=-1)
out_c = torch.matmul(attn, v_c)
out_c = out_c.view(batch, heads, q_num_blocks, 1, dim)
@@ -163,61 +149,46 @@ def video_sparse_attn_bshd(
block_size = (block_size, block_size, block_size)
block_elements = block_size[0] * block_size[1] * block_size[2]
if block_elements != 256:
raise ValueError(
"video_sparse_attn_bshd is only defined for block_elements=256 "
f"(got {block_elements}); use video_sparse_attn for the 64-block path."
)
raise ValueError("video_sparse_attn_bshd is only defined for block_elements=256 "
f"(got {block_elements}); use video_sparse_attn for the 64-block path.")
batch, q_seq_len, heads, dim = q.shape
kv_seq_len = k.shape[1]
if k.shape[0] != batch or v.shape[0] != batch or k.shape[2] != heads or v.shape[2] != heads:
raise ValueError("Expected q/k/v to have the same batch and head dimensions.")
if v.shape[1] != kv_seq_len:
raise ValueError(
f"Expected k and v to have the same sequence length, got "
f"k.shape[1]={kv_seq_len}, v.shape[1]={v.shape[1]}"
)
raise ValueError(f"Expected k and v to have the same sequence length, got "
f"k.shape[1]={kv_seq_len}, v.shape[1]={v.shape[1]}")
if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0:
raise ValueError(
f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, "
f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}"
)
raise ValueError(f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, "
f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}")
q_num_blocks = q_seq_len // block_elements
kv_num_blocks = kv_seq_len // block_elements
if variable_block_sizes.numel() != kv_num_blocks:
raise ValueError(
f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, "
f"got {variable_block_sizes.numel()}"
)
raise ValueError(f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, "
f"got {variable_block_sizes.numel()}")
if q_variable_block_sizes.numel() != q_num_blocks:
raise ValueError(
f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
f"got {q_variable_block_sizes.numel()}"
)
raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
f"got {q_variable_block_sizes.numel()}")
# Compression branch (BSHD-native: mean over the 256-token axis).
token_idx = torch.arange(block_elements, device=q.device, dtype=torch.int32)
q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1, 1)).view(
1, q_num_blocks, block_elements, 1, 1
)
kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1, 1)).view(
1, kv_num_blocks, block_elements, 1, 1
)
q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1,
1)).view(1, q_num_blocks, block_elements, 1, 1)
kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1,
1)).view(1, kv_num_blocks, block_elements, 1, 1)
q_c = q.view(batch, q_num_blocks, block_elements, heads, dim)
k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim)
v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim)
q_c = ((q_c.float() * q_token_valid).sum(dim=2)
/ q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
k_c = ((k_c.float() * kv_token_valid).sum(dim=2)
/ variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
v_c = ((v_c.float() * kv_token_valid).sum(dim=2)
/ variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
q_c = ((q_c.float() * q_token_valid).sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
k_c = ((k_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
v_c = ((v_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
q_ch = q_c.permute(0, 2, 1, 3).contiguous()
k_ch = k_c.permute(0, 2, 1, 3).contiguous()
v_ch = v_c.permute(0, 2, 1, 3).contiguous()
scores = torch.matmul(q_ch, k_ch.transpose(-2, -1)) / (dim ** 0.5)
scores = torch.matmul(q_ch, k_ch.transpose(-2, -1)) / (dim**0.5)
attn = torch.softmax(scores, dim=-1)
out_c_ch = torch.matmul(attn, v_ch)
out_c_blk = out_c_ch.permute(0, 2, 1, 3).contiguous()
@@ -229,9 +200,7 @@ def video_sparse_attn_bshd(
out = out_s
out_view = out.view(batch, q_num_blocks, block_elements, heads, dim)
if compress_attn_weight is not None:
gate_view = compress_attn_weight.view(
batch, q_num_blocks, block_elements, heads, dim
)
gate_view = compress_attn_weight.view(batch, q_num_blocks, block_elements, heads, dim)
out_view.add_(out_c_blk.unsqueeze(2) * gate_view)
else:
out_view.add_(out_c_blk.unsqueeze(2))
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py
+520 -345
View File
File diff suppressed because it is too large Load Diff
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py
+9 -16
View File
@@ -104,8 +104,7 @@ def _attn_fwd_sparse(
block_shape=(HEAD_DIM, BLOCK_N),
order=(0, 1))
v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1,
0)
v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
V_base = tl.make_block_ptr(base=V + v_off,
shape=(N_CTX_KV, HEAD_DIM),
strides=(stride_vk, stride_vn),
@@ -178,10 +177,8 @@ def _attn_bwd_preprocess(
off_hz = tl.program_id(1)
off_n = tl.arange(0, HEAD_DIM)
# load
o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
off_n[None, :])
do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
off_n[None, :]).to(tl.float32)
o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
delta = tl.sum(o * do, axis=1)
# write-back
tl.store(Delta + off_hz * N_CTX + off_m, delta)
@@ -235,8 +232,7 @@ def _attn_bwd_dkdv(
block_size = tl.load(variable_block_sizes + kv_blk)
for blk_idx in range(q_blocks * 2):
block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 +
blk_idx % 2) * step_m
block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 + blk_idx % 2) * step_m
qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
# Load m before computing qk to reduce pipeline stall.
offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
@@ -497,7 +493,7 @@ def _attn_bwd_dkdv_kernel(
stride_tok,
stride_d, #
# batch/head strides (may differ between Q and KV)
stride_qz,
stride_qz,
stride_qh,
stride_kz,
stride_kh,
@@ -607,7 +603,7 @@ def _attn_bwd_dq_kernel(
stride_tok,
stride_d, #
# batch/head strides (may differ between Q and KV)
stride_qz,
stride_qz,
stride_qh,
stride_kz,
stride_kh,
@@ -689,8 +685,7 @@ def _attn_bwd_dq_kernel(
# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
variable_block_sizes):
def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes):
B, H, Tq, D = q.shape
Tkv = k.shape[2]
sm_scale = 1.0 / math.sqrt(D)
@@ -701,8 +696,7 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
-1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
assert variable_block_sizes.numel() == Tkv // 64, (
f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, "
f"got {variable_block_sizes.numel()}"
)
f"got {variable_block_sizes.numel()}")
o = torch.empty_like(q)
M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device)
@@ -743,8 +737,7 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
return o, M
def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
k2q_index, k2q_num, variable_block_sizes):
def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes):
assert do.is_contiguous()
B, H, Tq, D = q.shape
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/fused_compress_topk.py
+24 -10
View File
@@ -122,9 +122,13 @@ def _fused_block_mean_bwd(
grid = (num_blocks, B * H)
_fused_block_mean_bwd_kernel[grid](
go_flat, gx_flat, variable_block_sizes,
go_flat.stride(0), go_flat.stride(1),
gx_flat.stride(0), gx_flat.stride(1),
go_flat,
gx_flat,
variable_block_sizes,
go_flat.stride(0),
go_flat.stride(1),
gx_flat.stride(0),
gx_flat.stride(1),
num_blocks,
BLOCK_ELEMENTS=block_elements,
HEAD_DIM=D,
@@ -152,9 +156,13 @@ def _fused_block_mean_fwd(
grid = (num_blocks, B * H)
_fused_block_mean_kernel[grid](
x_flat, out_flat, variable_block_sizes,
x_flat.stride(0), x_flat.stride(1),
out_flat.stride(0), out_flat.stride(1),
x_flat,
out_flat,
variable_block_sizes,
x_flat.stride(0),
x_flat.stride(1),
out_flat.stride(0),
out_flat.stride(1),
num_blocks,
BLOCK_ELEMENTS=block_elements,
HEAD_DIM=D,
@@ -311,7 +319,8 @@ def fused_topk_mask(
logger.debug(
"fused_topk_mask: kv_blocks=%d exceeds Triton limit %d, "
"falling back to PyTorch topk (slower)",
kv_blocks, MAX_KV_BLOCK_SIZE,
kv_blocks,
MAX_KV_BLOCK_SIZE,
)
return _pytorch_topk_mask_fallback(scores, topk)
@@ -323,9 +332,14 @@ def fused_topk_mask(
grid = (q_blocks, B * H)
_fused_topk_mask_kernel[grid](
scores_flat, mask_flat,
scores_flat.stride(0), scores_flat.stride(1), scores_flat.stride(2),
mask_flat.stride(0), mask_flat.stride(1), mask_flat.stride(2),
scores_flat,
mask_flat,
scores_flat.stride(0),
scores_flat.stride(1),
scores_flat.stride(2),
mask_flat.stride(0),
mask_flat.stride(1),
mask_flat.stride(2),
kv_blocks=kv_blocks,
topk=topk,
KV_BLOCK_SIZE=KV_BLOCK_SIZE,
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py
+43 -60
View File
@@ -5,6 +5,7 @@ import triton
import triton.language as tl
import torch
@triton.jit
def topk_index_to_map_kernel(
map_ptr,
@@ -57,14 +58,10 @@ def map_to_index_kernel(
tl.store(index_ptr_base + num * index_kv_stride, i)
num += 1
tl.store(
index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
q * index_num_q_stride, num)
tl.store(index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + q * index_num_q_stride, num)
def topk_index_to_map(index: torch.Tensor,
num_kv_blocks: int,
transpose_map: bool = False):
def topk_index_to_map(index: torch.Tensor, num_kv_blocks: int, transpose_map: bool = False):
"""
Convert topk indices to a map.
@@ -83,13 +80,9 @@ def topk_index_to_map(index: torch.Tensor,
bs, h, num_q_blocks, topk = index.shape
if transpose_map is False:
block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
dtype=torch.bool,
device=index.device)
block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), dtype=torch.bool, device=index.device)
else:
block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
dtype=torch.bool,
device=index.device)
block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), dtype=torch.bool, device=index.device)
block_map = block_map.transpose(2, 3)
grid = (bs, h, num_q_blocks)
@@ -126,13 +119,8 @@ def map_to_index(block_map: torch.Tensor):
"""
bs, h, num_q_blocks, num_kv_blocks = block_map.shape
index = torch.full((block_map.shape),
-1,
dtype=torch.int32,
device=block_map.device)
index_num = torch.empty((bs, h, num_q_blocks),
dtype=torch.int32,
device=block_map.device)
index = torch.full((block_map.shape), -1, dtype=torch.int32, device=block_map.device)
index_num = torch.empty((bs, h, num_q_blocks), dtype=torch.int32, device=block_map.device)
grid = (bs, h, num_q_blocks)
map_to_index_kernel[grid](
@@ -162,10 +150,20 @@ def _invert_indices_kernel(
q2k_num_ptr,
k2q_idx_ptr,
k2q_num_ptr,
q2k_idx_b, q2k_idx_h, q2k_idx_q, q2k_idx_k,
q2k_num_b, q2k_num_h, q2k_num_q,
k2q_idx_b, k2q_idx_h, k2q_idx_k, k2q_idx_q,
k2q_num_b, k2q_num_h, k2q_num_k,
q2k_idx_b,
q2k_idx_h,
q2k_idx_q,
q2k_idx_k,
q2k_num_b,
q2k_num_h,
q2k_num_q,
k2q_idx_b,
k2q_idx_h,
k2q_idx_k,
k2q_idx_q,
k2q_num_b,
k2q_num_h,
k2q_num_k,
MAX_KV_PER_Q: tl.constexpr,
):
# One program per (b, h, q): reserve a slot in k2q via atomicAdd, write q.
@@ -173,36 +171,17 @@ def _invert_indices_kernel(
pid_h = tl.program_id(1)
pid_q = tl.program_id(2)
n = tl.load(
q2k_num_ptr
+ pid_b * q2k_num_b
+ pid_h * q2k_num_h
+ pid_q * q2k_num_q
)
n = tl.load(q2k_num_ptr + pid_b * q2k_num_b + pid_h * q2k_num_h + pid_q * q2k_num_q)
q2k_row = (
q2k_idx_ptr
+ pid_b * q2k_idx_b
+ pid_h * q2k_idx_h
+ pid_q * q2k_idx_q
)
q2k_row = (q2k_idx_ptr + pid_b * q2k_idx_b + pid_h * q2k_idx_h + pid_q * q2k_idx_q)
for i in tl.range(0, MAX_KV_PER_Q):
if i < n:
kv = tl.load(q2k_row + i * q2k_idx_k)
count_ptr = (
k2q_num_ptr
+ pid_b * k2q_num_b
+ pid_h * k2q_num_h
+ kv * k2q_num_k
)
count_ptr = (k2q_num_ptr + pid_b * k2q_num_b + pid_h * k2q_num_h + kv * k2q_num_k)
pos = tl.atomic_add(count_ptr, 1)
tl.store(
k2q_idx_ptr
+ pid_b * k2q_idx_b
+ pid_h * k2q_idx_h
+ kv * k2q_idx_k
+ pos * k2q_idx_q,
k2q_idx_ptr + pid_b * k2q_idx_b + pid_h * k2q_idx_h + kv * k2q_idx_k + pos * k2q_idx_q,
pid_q,
)
@@ -214,22 +193,16 @@ def invert_indices(
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Transpose a Q->KV index list into a K->Q one via atomic compaction (GPU)."""
if q2k_idx.dim() != 4:
raise ValueError(
f"q2k_idx must be [B, H, Nq, Mk], got shape={tuple(q2k_idx.shape)}"
)
raise ValueError(f"q2k_idx must be [B, H, Nq, Mk], got shape={tuple(q2k_idx.shape)}")
if q2k_num.dim() != 3:
raise ValueError(
f"q2k_num must be [B, H, Nq], got shape={tuple(q2k_num.shape)}"
)
raise ValueError(f"q2k_num must be [B, H, Nq], got shape={tuple(q2k_num.shape)}")
if not q2k_idx.is_cuda or not q2k_num.is_cuda:
raise RuntimeError("invert_indices requires CUDA tensors.")
B, H, Nq, Mk = q2k_idx.shape
if q2k_num.shape != (B, H, Nq):
raise ValueError(
f"q2k_num shape {tuple(q2k_num.shape)} does not match q2k_idx "
f"[B, H, Nq] = {(B, H, Nq)}"
)
raise ValueError(f"q2k_num shape {tuple(q2k_num.shape)} does not match q2k_idx "
f"[B, H, Nq] = {(B, H, Nq)}")
q2k_idx = q2k_idx.contiguous()
q2k_num = q2k_num.contiguous()
@@ -257,10 +230,20 @@ def invert_indices(
q2k_num,
k2q_idx,
k2q_num,
q2k_idx.stride(0), q2k_idx.stride(1), q2k_idx.stride(2), q2k_idx.stride(3),
q2k_num.stride(0), q2k_num.stride(1), q2k_num.stride(2),
k2q_idx.stride(0), k2q_idx.stride(1), k2q_idx.stride(2), k2q_idx.stride(3),
k2q_num.stride(0), k2q_num.stride(1), k2q_num.stride(2),
q2k_idx.stride(0),
q2k_idx.stride(1),
q2k_idx.stride(2),
q2k_idx.stride(3),
q2k_num.stride(0),
q2k_num.stride(1),
q2k_num.stride(2),
k2q_idx.stride(0),
k2q_idx.stride(1),
k2q_idx.stride(2),
k2q_idx.stride(3),
k2q_num.stride(0),
k2q_num.stride(1),
k2q_num.stride(2),
MAX_KV_PER_Q=Mk,
)
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py
+9 -7
View File
@@ -8,6 +8,7 @@ from triton.language.target_info import cuda_capability_geq
MXFP_BLOCK_SIZE = tl.constexpr(16)
@triton.jit
def _compute_quant_and_scale(
src_tensor,
@@ -22,9 +23,7 @@ def _compute_quant_and_scale(
is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8
tl.static_assert(
is_fp4
or mx_tensor_dtype == tl.float8e4nv
or mx_tensor_dtype == tl.float8e5,
is_fp4 or mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5,
"mx_tensor_dtype must be uint8, float8e4nv, or float8e5",
)
@@ -32,7 +31,7 @@ def _compute_quant_and_scale(
f32_tensor = src_tensor.to(tl.float32)
abs_tensor = tl.abs(f32_tensor)
abs_tensor = tl.where(valid_src_mask, abs_tensor, -1.0) # Don't consider padding tensors in scale computation
if two_level_quant_P:
# row max from SageAttn3 paper
global_max_val = tl.max(f32_tensor, axis=1, keep_dims=True) # (BLOCK_SIZE_OUT_DIM, 1)
@@ -51,8 +50,9 @@ def _compute_quant_and_scale(
elif not two_level_quant_P and not use_global_sf:
s_dec = 1.0
s_enc = 1.0
max_val = tl.max(abs_tensor, axis=2, keep_dims=True) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) # per block maxima
max_val = tl.max(abs_tensor, axis=2,
keep_dims=True) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1) # per block maxima
s_dec_b = max_val / 6 # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
s_dec_b_e4m3 = (s_dec_b * s_enc).to(tl.float8e4nv) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
s_enc_b = 1 / (s_dec_b_e4m3.to(tl.float32) * s_dec) # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
@@ -132,6 +132,7 @@ def _compute_quant_and_scale(
return out_tensor, dequant_scale, s_dec
@triton.jit
def _compute_dequant(
mx_tensor,
@@ -141,7 +142,8 @@ def _compute_dequant(
BLOCK_SIZE_QUANT_DIM: tl.constexpr,
dst_dtype: tl.constexpr,
):
tl.static_assert(BLOCK_SIZE_QUANT_DIM % MXFP_BLOCK_SIZE == 0, f"Block size along quantization block must be a multiple of {MXFP_BLOCK_SIZE=}")
tl.static_assert(BLOCK_SIZE_QUANT_DIM % MXFP_BLOCK_SIZE == 0,
f"Block size along quantization block must be a multiple of {MXFP_BLOCK_SIZE=}")
# uint8 signifies two fp4 e2m1 values packed into a single byte
mx_tensor_dtype: tl.constexpr = mx_tensor.dtype
tl.static_assert(dst_dtype == tl.float16 or dst_dtype == tl.bfloat16 or dst_dtype == tl.float32)
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py
+76 -32
View File
@@ -3,33 +3,44 @@ import triton.language as tl
from .nvfp4_utils import _compute_quant_and_scale, _compute_dequant
@triton.jit
def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr,
BLOCK_SIZE_QUANT_DIM: tl.constexpr,
dst_dtype: tl.constexpr,
mx_tensor_dtype: tl.constexpr = tl.uint8,
use_global_sf: tl.constexpr = True,
two_level_quant_P: tl.constexpr = False):
def fake_quantize(src_tensor,
valid_src_mask,
BLOCK_SIZE_OUT_DIM: tl.constexpr,
BLOCK_SIZE_QUANT_DIM: tl.constexpr,
dst_dtype: tl.constexpr,
mx_tensor_dtype: tl.constexpr = tl.uint8,
use_global_sf: tl.constexpr = True,
two_level_quant_P: tl.constexpr = False):
high_prec_src_tensor = src_tensor
src_tensor, src_scale, src_s_dec = _compute_quant_and_scale(src_tensor=src_tensor,
valid_src_mask=valid_src_mask,
mx_tensor_dtype=mx_tensor_dtype,
use_global_sf=use_global_sf,
two_level_quant_P=two_level_quant_P)
src_tensor = _compute_dequant(mx_tensor=src_tensor,
scale=src_scale,
s_dec=src_s_dec,
BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM,
BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM,
src_tensor = _compute_dequant(mx_tensor=src_tensor,
scale=src_scale,
s_dec=src_s_dec,
BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM,
BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM,
dst_dtype=dst_dtype)
return src_tensor, high_prec_src_tensor.to(src_tensor.dtype)
@triton.jit
def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q,
stride_tok_q, stride_d_q,
fake_stride_z_q, fake_stride_h_q,
fake_stride_tok_q, fake_stride_d_q,
H, N_CTX_Q,
def fake_quantize_q(Q,
fake_Q,
stride_z_q,
stride_h_q,
stride_tok_q,
stride_d_q,
fake_stride_z_q,
fake_stride_h_q,
fake_stride_tok_q,
fake_stride_d_q,
H,
N_CTX_Q,
BLOCK_M: tl.constexpr,
HEAD_DIM: tl.constexpr,
use_global_sf: tl.constexpr = True):
@@ -46,18 +57,33 @@ def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q,
q_valid = offs_m < N_CTX_Q
q = tl.load(Q + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None], other=0.0)
q, _ = fake_quantize(src_tensor=q, valid_src_mask=q_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=q.dtype, use_global_sf=use_global_sf)
q, _ = fake_quantize(src_tensor=q,
valid_src_mask=q_valid[:, None],
BLOCK_SIZE_OUT_DIM=BLOCK_M,
BLOCK_SIZE_QUANT_DIM=HEAD_DIM,
dst_dtype=q.dtype,
use_global_sf=use_global_sf)
tl.store(fake_Q + offs_m[:, None] * fake_stride_tok_q + offs_k[None, :] * fake_stride_d_q, q, mask=q_valid[:, None])
@triton.jit
def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv,
stride_tok_kv, stride_d_kv,
fake_stride_z_kv, fake_stride_h_kv,
fake_stride_tok_kv, fake_stride_d_kv,
H, N_CTX_KV,
BLOCK_N: tl.constexpr,
HEAD_DIM: tl.constexpr,
use_global_sf: tl.constexpr = True):
def fake_quantize_kv(K,
V,
fake_K,
fake_V,
stride_z_kv,
stride_h_kv,
stride_tok_kv,
stride_d_kv,
fake_stride_z_kv,
fake_stride_h_kv,
fake_stride_tok_kv,
fake_stride_d_kv,
H,
N_CTX_KV,
BLOCK_N: tl.constexpr,
HEAD_DIM: tl.constexpr,
use_global_sf: tl.constexpr = True):
bhid = tl.program_id(1)
adj_kv = (stride_h_kv * (bhid % H) + stride_z_kv * (bhid // H))
fake_adj_kv = (fake_stride_h_kv * (bhid % H) + fake_stride_z_kv * (bhid // H))
@@ -70,11 +96,29 @@ def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv,
start_n = pid * BLOCK_N
offs_n = start_n + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, HEAD_DIM)
kv_valid = offs_n < N_CTX_KV
k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0)
v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0)
k, _ = fake_quantize(src_tensor=k_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=k_block.dtype, use_global_sf=use_global_sf)
v, _ = fake_quantize(src_tensor=v_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=v_block.dtype, use_global_sf=use_global_sf)
tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, k, mask=kv_valid[:, None])
tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, v, mask=kv_valid[:, None])
k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv,
mask=kv_valid[:, None],
other=0.0)
v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv,
mask=kv_valid[:, None],
other=0.0)
k, _ = fake_quantize(src_tensor=k_block,
valid_src_mask=kv_valid[:, None],
BLOCK_SIZE_OUT_DIM=BLOCK_N,
BLOCK_SIZE_QUANT_DIM=HEAD_DIM,
dst_dtype=k_block.dtype,
use_global_sf=use_global_sf)
v, _ = fake_quantize(src_tensor=v_block,
valid_src_mask=kv_valid[:, None],
BLOCK_SIZE_OUT_DIM=BLOCK_N,
BLOCK_SIZE_QUANT_DIM=HEAD_DIM,
dst_dtype=v_block.dtype,
use_global_sf=use_global_sf)
tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv,
k,
mask=kv_valid[:, None])
tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv,
v,
mask=kv_valid[:, None])
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py
+97 -48
View File
@@ -19,10 +19,14 @@ import triton.language as tl
@triton.jit
def _attn_fwd(
Q, K, V,
Q,
K,
V,
qk_scale: tl.constexpr,
topk: tl.constexpr,
LUT, LSE, OS,
LUT,
LSE,
OS,
L: tl.constexpr,
M_BLOCKS: tl.constexpr,
D: tl.constexpr,
@@ -45,7 +49,7 @@ def _attn_fwd(
OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
LUT_ptr = LUT + lut_offset
LSE_ptrs = LSE + lse_offset + offs_m
m_i = tl.full([BLOCK_M], -float('inf'), dtype=tl.float32)
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
@@ -54,7 +58,7 @@ def _attn_fwd(
for block_idx in tl.range(topk):
idx_n = tl.load(LUT_ptr + block_idx)
n_mask = offs_n < L - idx_n * BLOCK_N
k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # = 1 / ln(2)
if L - idx_n * BLOCK_N < BLOCK_N:
@@ -76,14 +80,16 @@ def _attn_fwd(
o_s = o_s / l_i[:, None]
tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
m_i += tl.math.log2(l_i)
tl.store(LSE_ptrs, m_i, mask=offs_m < L)
@triton.jit
def _attn_bwd_preprocess(
OS, DOS, DELTAS,
OS,
DOS,
DELTAS,
L,
D: tl.constexpr,
BLOCK_M: tl.constexpr,
@@ -100,15 +106,21 @@ def _attn_bwd_preprocess(
o_s = tl.load(OS + offs_m[:, None] * D + offs_d[None, :], mask=offs_m[:, None] < L)
do_s = tl.load(DOS + offs_m[:, None] * D + offs_d[None, :], mask=offs_m[:, None] < L)
delta_s = tl.sum(o_s * do_s, axis=1).to(DELTAS.type.element_ty)
tl.store(DELTAS + offs_m, delta_s, mask=offs_m < L)
@triton.jit
def _attn_bwd_dq(
Q, K, V, LSE, DELTAS,
DOS, DQ, LUT,
Q,
K,
V,
LSE,
DELTAS,
DOS,
DQ,
LUT,
qk_scale: tl.constexpr,
topk: tl.constexpr,
L: tl.constexpr,
@@ -141,28 +153,36 @@ def _attn_bwd_dq(
do_s = tl.load(DOS_ptrs, mask=offs_m[:, None] < L)
delta_s = tl.load(DELTAS_ptrs, mask=offs_m < L)
lse = tl.load(LSE_ptrs, mask=offs_m < L, other=float("inf"))
dq = tl.zeros([BLOCK_M, D], dtype=tl.float32)
for block_idx in tl.range(topk, num_stages=2):
idx_n = tl.load(LUT_ptr + block_idx)
n_mask = offs_n < L - idx_n * BLOCK_N
k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
qk = tl.dot(q, k.T) * (qk_scale * 1.4426950408889634)
p = tl.math.exp2(qk - lse[:, None])
p = tl.where(n_mask[None, :], p, 0.0)
dp = tl.dot(do_s, v.T).to(tl.float32)
ds = p * (dp - delta_s[:, None])
dq += tl.dot(ds.to(k.dtype), k)
tl.store(DQ_ptrs, dq * qk_scale, mask=offs_m[:, None] < L)
@triton.jit
def _attn_bwd_dkdv(
Q, K, V, DOS, DK, DV,
qk_scale, KBID, LSE, DELTAS,
Q,
K,
V,
DOS,
DK,
DV,
qk_scale,
KBID,
LSE,
DELTAS,
L: tl.constexpr,
M_BLOCKS: tl.constexpr,
N_BLOCKS: tl.constexpr,
@@ -196,7 +216,7 @@ def _attn_bwd_dkdv(
k = tl.load(K_ptrs, mask=offs_n[:, None] < L)
v = tl.load(V_ptrs, mask=offs_n[:, None] < L)
dk = tl.zeros([BLOCK_N, D], dtype=tl.float32)
dv = tl.zeros([BLOCK_N, D], dtype=tl.float32)
for idx_m in tl.range(0, L, BLOCK_M2):
@@ -215,7 +235,7 @@ def _attn_bwd_dkdv(
dpT = tl.dot(v, tl.trans(do))
dsT = pT * (dpT - delta[None, :])
dk += tl.dot(dsT.to(q.dtype), q)
Q_ptrs += BLOCK_M2 * D
DOS_ptrs += BLOCK_M2 * D
LSE_ptrs += BLOCK_M2
@@ -225,11 +245,11 @@ def _attn_bwd_dkdv(
tl.store(DK_ptrs, dk * qk_scale, mask=offs_n[:, None] < L)
tl.store(DV_ptrs, dv, mask=offs_n[:, None] < L)
class _attention(torch.autograd.Function):
"""Sparse attention forward/backward with autograd support."""
@staticmethod
def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None):
assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
@@ -248,15 +268,22 @@ class _attention(torch.autograd.Function):
lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32)
grid = (M_BLOCKS, B * H)
_attn_fwd[grid](
q, k, v, qk_scale, topk,
lut, lse, o_s,
L, M_BLOCKS,
D, BLOCK_M, BLOCK_N,
num_warps=4 if q.shape[-1] == 64 else 8,
num_stages=3
)
_attn_fwd[grid](q,
k,
v,
qk_scale,
topk,
lut,
lse,
o_s,
L,
M_BLOCKS,
D,
BLOCK_M,
BLOCK_N,
num_warps=4 if q.shape[-1] == 64 else 8,
num_stages=3)
ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s)
ctx.qk_scale = qk_scale
ctx.topk = topk
@@ -282,30 +309,52 @@ class _attention(torch.autograd.Function):
grid = (M_BLOCKS, B * H)
_attn_bwd_preprocess[grid](
o_s, do_s, delta_s,
L, D, BLOCK_M,
o_s,
do_s,
delta_s,
L,
D,
BLOCK_M,
)
grid = (M_BLOCKS, B * H)
_attn_bwd_dq[grid](
q, k, v, lse, delta_s,
do_s, dq, lut,
ctx.qk_scale, ctx.topk,
L, M_BLOCKS,
D, BLOCK_M, BLOCK_N,
num_warps=4 if q.shape[-1] == 64 else 8,
num_stages=4 if q.shape[-1] == 64 else 5
)
_attn_bwd_dq[grid](q,
k,
v,
lse,
delta_s,
do_s,
dq,
lut,
ctx.qk_scale,
ctx.topk,
L,
M_BLOCKS,
D,
BLOCK_M,
BLOCK_N,
num_warps=4 if q.shape[-1] == 64 else 8,
num_stages=4 if q.shape[-1] == 64 else 5)
grid = (N_BLOCKS, B * H)
_attn_bwd_dkdv[grid](
q, k, v, do_s, dk, dv,
ctx.qk_scale, k_block_id, lse, delta_s,
L, M_BLOCKS, N_BLOCKS,
D, BLOCK_M, BLOCK_N,
BLOCK_SLICE_FACTOR=BLOCK_M // 64,
num_warps=4 if q.shape[-1] == 64 else 8,
num_stages=4 if q.shape[-1] == 64 else 5
)
_attn_bwd_dkdv[grid](q,
k,
v,
do_s,
dk,
dv,
ctx.qk_scale,
k_block_id,
lse,
delta_s,
L,
M_BLOCKS,
N_BLOCKS,
D,
BLOCK_M,
BLOCK_N,
BLOCK_SLICE_FACTOR=BLOCK_M // 64,
num_warps=4 if q.shape[-1] == 64 else 8,
num_stages=4 if q.shape[-1] == 64 else 5)
return dq, dk, dv, None, None, None, None, None, None
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py
+111 -71
View File
@@ -3,6 +3,7 @@ import torch
import triton
import triton.language as tl
def is_cuda():
return triton.runtime.driver.active.get_current_target().backend == "cuda"
@@ -14,7 +15,8 @@ def is_hip():
def is_cdna3_cdna4():
target = triton.runtime.driver.active.get_current_target()
return (target.arch == 'gfx1201' or target.arch == 'gfx1101' or target.arch == 'gfx1100' or target.arch == 'gfx1030')
return (target.arch == 'gfx1201' or target.arch == 'gfx1101' or target.arch == 'gfx1100'
or target.arch == 'gfx1030')
def get_common_autotune_config():
@@ -56,10 +58,17 @@ def clamp_int(value, min_val, max_val):
@triton.jit
def _attn_fwd_loop(
q, k, v, kv_mask, m, l, acc, sm_scale,
q,
k,
v,
kv_mask,
m,
l,
acc,
sm_scale,
MASK_KV: tl.constexpr,
):
scores = tl.dot(q, k.T) #[BLOCK_Q, BLOCK_KV]
scores = tl.dot(q, k.T) #[BLOCK_Q, BLOCK_KV]
scores = scores * sm_scale
if MASK_KV:
scores = tl.where(kv_mask[None, :], scores, -float('inf'))
@@ -86,13 +95,25 @@ def _attn_fwd_loop(
)
@triton.jit
def triton_sta_kernel(
Q, K, V, output,
batch_size: int, num_heads: int, seq_len: int, head_dim: int,
Q,
K,
V,
output,
batch_size: int,
num_heads: int,
seq_len: int,
head_dim: int,
img_seq_len: int,
text_length: int,
canvas_t: int, canvas_h: int, canvas_w: int,
kernel_t: int, kernel_h: int, kernel_w: int,
tile_t: int, tile_h: int, tile_w: int,
canvas_t: int,
canvas_h: int,
canvas_w: int,
kernel_t: int,
kernel_h: int,
kernel_w: int,
tile_t: int,
tile_h: int,
tile_w: int,
scale: float,
has_text: tl.constexpr,
text_q: tl.constexpr,
@@ -111,8 +132,8 @@ def triton_sta_kernel(
q_tile_flat = tl.program_id(2) // q_block_per_tile
q_block_idx = tl.program_id(2) % q_block_per_tile
m = tl.full((BLOCK_Q,), -float('inf'), dtype=tl.float32)
l = tl.zeros((BLOCK_Q,), dtype=tl.float32)
m = tl.full((BLOCK_Q, ), -float('inf'), dtype=tl.float32)
l = tl.zeros((BLOCK_Q, ), dtype=tl.float32)
acc = tl.zeros((BLOCK_Q, BLOCK_DIM), dtype=tl.float32)
q_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim
@@ -125,11 +146,9 @@ def triton_sta_kernel(
q_idx = q_base_idx + q_offset_in_tile
q_mask = (q_block_idx * BLOCK_Q + tl.arange(0, BLOCK_Q)) < total_tile_size
q = tl.load(
Q + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=q_mask[:, None],
other=0.0
) # [BLOCK_Q, BLOCK_DIM]
q = tl.load(Q + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=q_mask[:, None],
other=0.0) # [BLOCK_Q, BLOCK_DIM]
# Scale sm_scale by log_2(e) and use 2^x instead of exp
sm_scale = scale * 1.4426950408889634
@@ -175,7 +194,8 @@ def triton_sta_kernel(
for kv_tile_t in tl.range(kv_tile_start_t, kv_tile_end_t):
for kv_tile_h in tl.range(kv_tile_start_h, kv_tile_end_h):
for kv_tile_w in tl.range(kv_tile_start_w, kv_tile_end_w):
kv_base_idx = (kv_tile_t * num_tiles_h * num_tiles_w + kv_tile_h * num_tiles_w + kv_tile_w) * total_tile_size
kv_base_idx = (kv_tile_t * num_tiles_h * num_tiles_w + kv_tile_h * num_tiles_w +
kv_tile_w) * total_tile_size
for kv_block_idx in tl.range(0, total_tile_size, BLOCK_KV):
kv_offset_in_block = tl.arange(0, BLOCK_KV)
@@ -184,20 +204,15 @@ def triton_sta_kernel(
kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim
k = tl.load(
K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0
) # [BLOCK_KV, BLOCK_DIM]
v = tl.load(
V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0
) # [BLOCK_KV, BLOCK_DIM]
k = tl.load(K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0) # [BLOCK_KV, BLOCK_DIM]
v = tl.load(V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0) # [BLOCK_KV, BLOCK_DIM]
m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, False)
# for kv_text
if has_text:
kv_base_idx = img_seq_len
@@ -208,35 +223,32 @@ def triton_sta_kernel(
kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim
k = tl.load(
K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0
) # [BLOCK_KV, BLOCK_DIM]
v = tl.load(
V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0
) # [BLOCK_KV, BLOCK_DIM]
k = tl.load(K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0) # [BLOCK_KV, BLOCK_DIM]
v = tl.load(V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
mask=kv_mask[:, None],
other=0.0) # [BLOCK_KV, BLOCK_DIM]
m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, True)
output_acc = acc / l[:, None]
tl.store(
output + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
output_acc,
mask=q_mask[:, None]
) # [BLOCK_Q, BLOCK_DIM]
tl.store(output + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :],
output_acc,
mask=q_mask[:, None]) # [BLOCK_Q, BLOCK_DIM]
def sliding_tile_attention_triton(
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
window_size, text_length: int,
has_text=True, dit_seq_shape='30x48x80') -> torch.Tensor:
def sliding_tile_attention_triton(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
window_size,
text_length: int,
has_text=True,
dit_seq_shape='30x48x80') -> torch.Tensor:
seq_length = q.shape[2]
if has_text:
assert q.shape[2] >= 115200 and q.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '30x48x80' for HunyuanVideo"
assert q.shape[2] >= 115200 and q.shape[
2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '30x48x80' for HunyuanVideo"
target_size = math.ceil(seq_length / 384) * 384
pad_size = target_size - seq_length
if pad_size > 0:
@@ -244,22 +256,24 @@ def sliding_tile_attention_triton(
k = torch.cat([k, k[:, :, -pad_size:]], dim=2)
v = torch.cat([v, v[:, :, -pad_size:]], dim=2)
else:
if dit_seq_shape == '36x48x48': # Stepvideo
if dit_seq_shape == '36x48x48': # Stepvideo
assert q.shape[2] == 82944
elif dit_seq_shape == '18x48x80': # Wan
elif dit_seq_shape == '18x48x80': # Wan
assert q.shape[2] == 69120
else:
raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
raise ValueError(
f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan"
)
assert q.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
batch_size, num_heads, seq_len, head_dim = q.shape
if dit_seq_shape == '30x48x80': # Hunyuan
if dit_seq_shape == '30x48x80': # Hunyuan
canvas_t, canvas_h, canvas_w = 30, 48, 80
tile_t, tile_h, tile_w = 6, 8, 8
elif dit_seq_shape == '36x48x48': # Stepvideo
elif dit_seq_shape == '36x48x48': # Stepvideo
canvas_t, canvas_h, canvas_w = 36, 48, 48
tile_t, tile_h, tile_w = 6, 8, 8
elif dit_seq_shape == '18x48x80': # Wan
elif dit_seq_shape == '18x48x80': # Wan
canvas_t, canvas_h, canvas_w = 18, 48, 80
tile_t, tile_h, tile_w = 6, 8, 8
@@ -283,22 +297,36 @@ def sliding_tile_attention_triton(
# This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
for head_index, (kernel_t, kernel_h, kernel_w) in enumerate(window_size):
for batch in range(batch_size):
q_head, k_head, v_head, o_head = (q[batch:batch + 1, head_index:head_index + 1],
k[batch:batch + 1, head_index:head_index + 1],
v[batch:batch + 1, head_index:head_index + 1],
output[batch:batch + 1, head_index:head_index + 1])
q_head, k_head, v_head, o_head = (q[batch:batch + 1,
head_index:head_index + 1], k[batch:batch + 1,
head_index:head_index + 1],
v[batch:batch + 1,
head_index:head_index + 1], output[batch:batch + 1,
head_index:head_index + 1])
# triton_sta_kernel[(1, 1, num_tiles * triton.cdiv(total_tile_size, BLOCK_Q))](
grid = lambda META: (1, 1, num_tiles * triton.cdiv(total_tile_size, META['BLOCK_Q']))
triton_sta_kernel[grid](
q_head, k_head, v_head, o_head,
1, 1, seq_len, head_dim,
q_head,
k_head,
v_head,
o_head,
1,
1,
seq_len,
head_dim,
img_seq_len,
text_length,
canvas_t, canvas_h, canvas_w,
kernel_t, kernel_h, kernel_w,
tile_t, tile_h, tile_w,
scale=1.0 / (head_dim ** 0.5),
canvas_t,
canvas_h,
canvas_w,
kernel_t,
kernel_h,
kernel_w,
tile_t,
tile_h,
tile_w,
scale=1.0 / (head_dim**0.5),
has_text=has_text,
text_q=False,
# BLOCK_Q=BLOCK_Q,
@@ -312,15 +340,27 @@ def sliding_tile_attention_triton(
# triton_sta_kernel[(batch_size, num_heads, triton.cdiv(total_tile_size, BLOCK_Q))](
grid = lambda META: (batch_size, num_heads, triton.cdiv(total_tile_size, META['BLOCK_Q']))
triton_sta_kernel[grid](
q, k, v, output,
batch_size, num_heads, seq_len, head_dim,
q,
k,
v,
output,
batch_size,
num_heads,
seq_len,
head_dim,
img_seq_len,
text_length,
canvas_t, canvas_h, canvas_w,
3, 3, 3,
canvas_t,
canvas_h,
canvas_w,
3,
3,
3,
#kernel_t, kernel_h, kernel_w,
tile_t, tile_h, tile_w,
scale=1.0 / (head_dim ** 0.5),
tile_t,
tile_h,
tile_w,
scale=1.0 / (head_dim**0.5),
has_text=has_text,
text_q=True,
# BLOCK_Q=BLOCK_Q,
@@ -332,4 +372,4 @@ def sliding_tile_attention_triton(
if pad_size > 0:
output = output[:, :, :seq_length]
return output
return output
fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py
+98 -96
View File
@@ -20,6 +20,7 @@ except ImportError:
rms_norm_cuda = None
layer_norm_cuda = None
def int8_quant(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Quantize a floating-point tensor to int8 using a custom CUDA kernel.
@@ -67,6 +68,7 @@ def int8_linear(
gemm_cuda(x_q, x_s, w_q, w_s, y)
return y.reshape(*shape[:-1], n)
def flatten_if_batched(*tensors):
"""
Flattens all input tensors from (B, N, D_i) to (B * N, D_i) if they are batched (3D).
@@ -92,9 +94,8 @@ def flatten_if_batched(*tensors):
batched = True
batch_size = first.shape[0]
assert all(t.shape[0] == batch_size for t in tensors), "All input tensors must have the same batch size"
assert all(
t.shape[1] == first.shape[1] for t in tensors
), "All input tensors must have the same sequence length"
assert all(t.shape[1] == first.shape[1]
for t in tensors), "All input tensors must have the same sequence length"
flat_tensors = [t.reshape(-1, t.shape[-1]) for t in tensors]
else:
batched = False
@@ -167,21 +168,21 @@ def rmsnorm(x, w, eps):
# allocate output
M, N = x.shape
y = torch.empty_like(x, dtype=x.dtype)
rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
# heuristics for number of warps
num_warps = 8
# Avoid illegal memory access
N2 = triton.next_power_of_2(N)
if N <= 512:
BLOCK_M = 32
else:
BLOCK_M = 1
# Call the triton kernel
_rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( #
_rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( #
x,
y,
w,
@@ -201,6 +202,7 @@ def rmsnorm(x, w, eps):
return y, rstd
@triton.jit
def _layer_norm_param_fwd_fused(
X, # pointer to the input
@@ -245,7 +247,7 @@ def _layer_norm_param_fwd_fused(
w = tl.load(W + cols)
b = tl.load(B + cols)
x_hat = x_hat * w + b
# Write output
@@ -260,20 +262,20 @@ def layernorm_param(x, w, b, eps):
# allocate output
M, N = x.shape
y = torch.empty_like(x, dtype=torch.float32)
mean = torch.empty((M,), dtype=torch.float32, device=x.device)
rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
mean = torch.empty((M, ), dtype=torch.float32, device=x.device)
rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
# heuristics for number of warps
num_warps = 8
N2 = triton.next_power_of_2(N)
if N <= 512:
BLOCK_M = 32
else:
BLOCK_M = 1
# enqueue kernel
_layer_norm_param_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( #
_layer_norm_param_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( #
x,
y,
w,
@@ -355,20 +357,20 @@ def layernorm_noparam(x, eps):
# allocate output
M, N = x.shape
y = torch.empty_like(x, dtype=torch.float32)
mean = torch.empty((M,), dtype=torch.float32, device=x.device)
rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
mean = torch.empty((M, ), dtype=torch.float32, device=x.device)
rstd = torch.empty((M, ), dtype=torch.float32, device=x.device)
# heuristics for number of warps
num_warps = 8
N2 = triton.next_power_of_2(N)
if N <= 512:
BLOCK_M = 32
else:
BLOCK_M = 1
# enqueue kernel
_layer_norm_noparam_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( #
_layer_norm_noparam_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( #
x,
y,
mean,
@@ -388,6 +390,7 @@ def layernorm_noparam(x, eps):
return y, mean, rstd
def layernorm(x, w, b, eps, elementwise_affine=True):
if elementwise_affine:
assert w is not None and b is not None
@@ -396,9 +399,11 @@ def layernorm(x, w, b, eps, elementwise_affine=True):
assert w is None and b is None
return layernorm_noparam(x, eps)
def cdiv(a: int, b: int):
return (a + b - 1) // b
def dequantize_weight(w_q, w_s, dtype):
# w_q: (N, K) int8
# w_s: (NB, KB) float32
@@ -412,12 +417,14 @@ def dequantize_weight(w_q, w_s, dtype):
w_s_exp = w_s_exp[:N, :K]
return w_q.to(dtype) * w_s_exp.to(dtype)
class Int8LinearFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, x, w_q, w_s, bias=None):
ctx.save_for_backward(x, w_q, w_s, bias)
ctx.bias_requires_grad = bias.requires_grad if bias is not None else False
out = int8_linear(x, w_q, w_s)
if bias is not None:
out = out + bias
@@ -426,24 +433,26 @@ class Int8LinearFunction(torch.autograd.Function):
@staticmethod
def backward(ctx, grad_output):
x, w_q, w_s, bias = ctx.saved_tensors
grad_input = None
grad_weight = None
grad_weight = None
grad_scale = None
grad_bias = None
if ctx.needs_input_grad[0]:
if ctx.needs_input_grad[0]:
# grad_input = grad_output @ W
w_float = dequantize_weight(w_q, w_s, grad_output.dtype)
grad_input = torch.matmul(grad_output, w_float)
if ctx.bias_requires_grad:
dim_to_sum = list(range(grad_output.dim() - 1))
grad_bias = grad_output.sum(dim=dim_to_sum)
return grad_input, grad_weight, grad_scale, grad_bias
class RMSNormFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, x, w, eps):
y, rstd = rmsnorm(x, w, eps)
@@ -459,14 +468,14 @@ class RMSNormFunction(torch.autograd.Function):
x, w, rstd = ctx.saved_tensors
eps = ctx.eps
N = x.shape[-1]
# dL/dy * w
dx = grad_output * w
# Expand rstd
rstd = rstd.unsqueeze(-1) # (B, 1) or (M, 1)
rstd = rstd.unsqueeze(-1) # (B, 1) or (M, 1)
if x.dim() == 3:
# Flatten if necessary or handle dimensions.
# Flatten if necessary or handle dimensions.
# forward flattens x to (M, N) internally but saved x might be 3D?
# rmsnorm takes x, flattens it. But does it modify x in place? No.
# But ctx.save_for_backward saves the original x (3D if input was 3D).
@@ -479,36 +488,36 @@ class RMSNormFunction(torch.autograd.Function):
x_flat = x
grad_output_flat = grad_output
dx = dx
# Standard RMSNorm backward
# dy = grad_output_flat
# x_hat = x * rstd
# w * dy
# c1 = mean(dy * w * x) * rstd^2
# dx = (dy * w - c1 * x) * rstd
# More precise:
# y = x * rstd * w
# dy = grad_output
# dw = sum(dy * x * rstd, dim=0)
grad_w = None
if ctx.needs_input_grad[1]: # w
# grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0)
# More accurate gradient for weight when considering rstd was computed from x
# Actually, for Affine part (w), it is just sum(dL/dy * x_hat).
# y = x_hat * w
# dL/dw = sum(dL/dy * x_hat)
x_hat = x_flat * rstd
grad_w = (grad_output_flat * x_hat).sum(dim=0)
if ctx.needs_input_grad[1]: # w
# grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0)
# More accurate gradient for weight when considering rstd was computed from x
# Actually, for Affine part (w), it is just sum(dL/dy * x_hat).
# y = x_hat * w
# dL/dw = sum(dL/dy * x_hat)
x_hat = x_flat * rstd
grad_w = (grad_output_flat * x_hat).sum(dim=0)
grad_input = None
if ctx.needs_input_grad[0]: # x
if ctx.needs_input_grad[0]: # x
# x_hat = x * rstd
# dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat)
# dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat)
# but for RMSNorm:
# dx = rstd * (w * dy - (x * rstd^2) * mean(w * dy * x)) -> check formula
# Using PyTorch autograd for reference logic:
# y = x / sqrt(mean(x^2) + eps) * w
# Let sigma = sqrt(...)
@@ -518,12 +527,12 @@ class RMSNormFunction(torch.autograd.Function):
# dL/dx = (dL/dy * w)/sigma - sum(dL/dy * w * x) * x / (N * sigma^3)
# = (1/sigma) * [ (dL/dy * w) - x * sum(dL/dy * w * x) / (N * sigma^2) ]
# = rstd * [ (dL/dy * w) - x * rstd^2 * mean(dL/dy * w * x) ] # Wait, sum/N is mean
# Let's compute:
dy_w = grad_output_flat * w
# term2 = (dy_w * x_flat).sum(dim=-1, keepdim=True) / N # mean(dy*w*x)
# grad_input = rstd * (dy_w - x_flat * (rstd ** 2) * term2)
# Re-derivation:
# x_hat = x * rstd
# y = x_hat * w
@@ -531,20 +540,22 @@ class RMSNormFunction(torch.autograd.Function):
# dy/dx = w * dx_hat/dx
# dx_hat/dx = rstd * (I - x * x^T * rstd^2 / N) ?? No.
# dx_hat_i / dx_j = rstd * (delta_ij - x_i * x_j * rstd^2 / N)
# dL/dx_hat = dL/dy * w
# dL/dx = rstd * (dL/dx_hat - x * rstd^2 * mean(dL/dx_hat * x))
dx_hat = grad_output_flat * w
dx_hat_x_mean = (dx_hat * x_flat).mean(dim=-1, keepdim=True)
grad_input = rstd * (dx_hat - x_flat * (rstd ** 2) * dx_hat_x_mean)
grad_input = rstd * (dx_hat - x_flat * (rstd**2) * dx_hat_x_mean)
if x.dim() == 3:
grad_input = grad_input.reshape(x.shape)
return grad_input, grad_w, None
class LayerNormFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, x, w, b, eps, elementwise_affine):
if elementwise_affine:
@@ -553,7 +564,7 @@ class LayerNormFunction(torch.autograd.Function):
else:
y, mean, rstd = layernorm_noparam(x, eps)
ctx.save_for_backward(x, None, None, mean, rstd)
ctx.eps = eps
ctx.elementwise_affine = elementwise_affine
return y
@@ -562,60 +573,62 @@ class LayerNormFunction(torch.autograd.Function):
def backward(ctx, grad_output):
x, w, b, mean, rstd = ctx.saved_tensors
N = x.shape[-1]
if x.dim() == 3:
x_flat = x.reshape(-1, N)
grad_output_flat = grad_output.reshape(-1, N)
else:
x_flat = x
grad_output_flat = grad_output
if w is None:
w_eff = 1.0
w_eff = 1.0
else:
w_eff = w
w_eff = w
# dx calculation
# x_hat = (x - mean) * rstd
# y = x_hat * w + b
# dL/dx_hat = dL/dy * w
dy = grad_output_flat
dx_hat = dy * w_eff
# dL/dvar = sum(dL/dx_hat * (x-mean) * (-0.5) * (var+eps)^(-1.5))
# = -0.5 * rstd^3 * sum(dx_hat * (x-mean))
# dL/dmean = sum(dL/dx_hat * (-rstd)) + dL/dvar * (-2/N) * sum(x-mean)
# term sum(x-mean) is 0. So second part vanishes.
# dL/dmean = -rstd * sum(dx_hat)
# dL/dx = dL/dx_hat * rstd + dL/dvar * 2(x-mean)/N + dL/dmean * 1/N
# = dx_hat * rstd + (-0.5 * rstd^3 * sum(dx_hat * (x-mean))) * 2(x-mean)/N + (-rstd * sum(dx_hat))/N
# = rstd * [ dx_hat - mean(dx_hat) - (x-mean)*rstd^2 * mean(dx_hat * (x-mean)) ]
x_centered = x_flat - mean.unsqueeze(-1)
dx_hat_mean = dx_hat.mean(dim=-1, keepdim=True)
dx_hat_x_centered_mean = (dx_hat * x_centered).mean(dim=-1, keepdim=True)
grad_input = rstd.unsqueeze(-1) * (dx_hat - dx_hat_mean - x_centered * (rstd.unsqueeze(-1)**2) * dx_hat_x_centered_mean)
grad_input = rstd.unsqueeze(-1) * (dx_hat - dx_hat_mean - x_centered *
(rstd.unsqueeze(-1)**2) * dx_hat_x_centered_mean)
if x.dim() == 3:
grad_input = grad_input.reshape(x.shape)
grad_w = None
grad_b = None
if ctx.elementwise_affine:
if ctx.needs_input_grad[1]:
x_hat = x_centered * rstd.unsqueeze(-1)
grad_w = (dy * x_hat).sum(dim=0)
x_hat = x_centered * rstd.unsqueeze(-1)
grad_w = (dy * x_hat).sum(dim=0)
if ctx.needs_input_grad[2]:
grad_b = dy.sum(dim=0)
grad_b = dy.sum(dim=0)
return grad_input, grad_w, grad_b, None, None
class Int8Linear(nn.Module):
def __init__(self, in_features, out_features, bias=True, dtype=torch.bfloat16):
super().__init__()
self.in_features = in_features
@@ -623,27 +636,24 @@ class Int8Linear(nn.Module):
row_blocks = cdiv(out_features, b=128)
col_blocks = cdiv(in_features, b=128)
self.register_buffer("int8_weight", torch.empty((out_features, in_features), dtype=torch.int8))
self.register_buffer("scale", torch.empty((row_blocks, col_blocks), dtype=torch.float32))
if bias:
self.register_buffer("bias", torch.empty(out_features, dtype=dtype))
else:
self.bias = None
def forward(self, x):
return Int8LinearFunction.apply(x, self.int8_weight, self.scale, self.bias)
@classmethod
def from_linear(cls, original_linear: nn.Linear, quantize: bool = True):
int8_layer = cls(
original_linear.in_features,
original_linear.out_features,
bias=original_linear.bias is not None,
dtype=original_linear.weight.dtype
)
int8_layer = cls(original_linear.in_features,
original_linear.out_features,
bias=original_linear.bias is not None,
dtype=original_linear.weight.dtype)
if quantize:
w_data = original_linear.weight.data.cuda()
int8_w, scale = int8_quant(w_data)
@@ -652,10 +662,12 @@ class Int8Linear(nn.Module):
int8_layer.scale.copy_(scale)
if original_linear.bias is not None:
int8_layer.bias.data.copy_(original_linear.bias.data.cuda())
return int8_layer
class FastRMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-5):
super().__init__()
self.dim = dim
@@ -667,23 +679,15 @@ class FastRMSNorm(nn.Module):
@classmethod
def from_rmsnorm(cls, original_rmsnorm):
rmsnorm_layer = cls(
dim=original_rmsnorm.dim,
eps=original_rmsnorm.eps
)
rmsnorm_layer = cls(dim=original_rmsnorm.dim, eps=original_rmsnorm.eps)
if original_rmsnorm.weight.device != torch.device('meta'):
rmsnorm_layer.weight.data.copy_(original_rmsnorm.weight.float().data)
return rmsnorm_layer
class FastLayerNorm(nn.Module):
def __init__(
self,
dim: int,
eps: float = 1e-5,
elementwise_affine: bool = False,
bias: bool = True
) :
def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = False, bias: bool = True):
super().__init__()
self.dim = dim # type: ignore[arg-type]
self.eps = eps
@@ -700,15 +704,13 @@ class FastLayerNorm(nn.Module):
def forward(self, x):
return LayerNormFunction.apply(x.float(), self.weight, self.bias, self.eps, self.elementwise_affine).to(x.dtype)
@classmethod
def from_layernorm(cls, original_layernorm):
layernorm_layer = cls(
dim=original_layernorm.normalized_shape[0],
eps=original_layernorm.eps,
elementwise_affine=False if original_layernorm.weight is None else True,
bias=original_layernorm.bias is not None
)
layernorm_layer = cls(dim=original_layernorm.normalized_shape[0],
eps=original_layernorm.eps,
elementwise_affine=False if original_layernorm.weight is None else True,
bias=original_layernorm.bias is not None)
if original_layernorm.weight is not None and original_layernorm.weight.device != torch.device('meta'):
layernorm_layer.weight.data.copy_(original_layernorm.weight.data)
if original_layernorm.bias is not None and original_layernorm.bias.device != torch.device('meta'):
fastvideo-kernel/python/fastvideo_kernel/vmoba.py
+137 -281
View File
@@ -10,9 +10,7 @@ try:
except ImportError:
def _unsupported(*args, **kwargs):
raise ImportError(
"flash-attn is not installed. Please install it, e.g., `uv pip install flash-attn`."
)
raise ImportError("flash-attn is not installed. Please install it, e.g., `uv pip install flash-attn`.")
_flash_attn_varlen_forward = _unsupported
_flash_attn_varlen_backward = _unsupported
@@ -39,25 +37,17 @@ def calc_chunks(cu_seqlen, moba_chunk_size):
)
cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0)
num_chunk = cu_num_chunk[-1]
chunk_sizes = torch.full((num_chunk + 1, ),
moba_chunk_size,
dtype=torch.int32,
device=cu_seqlen.device)
chunk_sizes = torch.full((num_chunk + 1, ), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device)
chunk_sizes[0] = 0
batch_last_chunk_size = batch_sizes - (batch_num_chunk -
1) * moba_chunk_size
batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size
chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size
cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32)
chunk_to_batch = torch.zeros((num_chunk, ),
dtype=torch.int32,
device=cu_seqlen.device)
chunk_to_batch = torch.zeros((num_chunk, ), dtype=torch.int32, device=cu_seqlen.device)
chunk_to_batch[cu_num_chunk[1:-1]] = 1
chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32)
# Do not filter out any chunk
filtered_chunk_indices = torch.arange(num_chunk,
device=cu_seqlen.device,
dtype=torch.int32)
filtered_chunk_indices = torch.arange(num_chunk, device=cu_seqlen.device, dtype=torch.int32)
num_filtered_chunk = num_chunk
return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch
@@ -66,9 +56,7 @@ def calc_chunks(cu_seqlen, moba_chunk_size):
# --- Threshold Selection Helper Functions ---
def _select_threshold_query_head(gate: torch.Tensor,
valid_gate_mask: torch.Tensor,
gate_self_chunk_mask: torch.Tensor,
def _select_threshold_query_head(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor,
simsum_threshold: float) -> torch.Tensor:
"""
Selects chunks for each <query, head> pair based on threshold.
@@ -78,16 +66,13 @@ def _select_threshold_query_head(gate: torch.Tensor,
eps = 1e-6
# LSE‐style normalization per <head, query> (across chunks)
gate_masked = torch.where(valid_gate_mask, gate,
-torch.inf) # Use -inf for max
gate_min_val = torch.where(valid_gate_mask, gate,
torch.inf) # Use +inf for min
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
row_min = gate_min_val.amin(dim=0) # (H, S)
row_max = gate_masked.amax(dim=0) # (H, S)
denom = row_max - row_min
denom = torch.where(denom <= eps, torch.ones_like(denom),
denom) # avoid divide‑by‑zero
denom = torch.where(denom <= eps, torch.ones_like(denom), denom) # avoid divide‑by‑zero
gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0)
gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
@@ -97,18 +82,15 @@ def _select_threshold_query_head(gate: torch.Tensor,
# 2) compute how much more normalized weight we need beyond self
total_norm_sum = gate_norm.sum(dim=0) # (H, S)
remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps
) # (H, S)
remain_ratio = torch.clamp(remain_ratio,
min=0.0) # if already ≥ thresh, no extra needed
remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps) # (H, S)
remain_ratio = torch.clamp(remain_ratio, min=0.0) # if already ≥ thresh, no extra needed
# 3) zero out the self‐chunk in a copy, so we only sort “others”
others_norm = gate_norm.clone()
others_norm[gate_self_chunk_mask] = 0.0
# 4) sort the other chunks by descending norm, per <head,seq>
sorted_norm, sorted_idx = torch.sort(others_norm, descending=True,
dim=0) # (C, H, S)
sorted_norm, sorted_idx = torch.sort(others_norm, descending=True, dim=0) # (C, H, S)
# 5) cumulative‑sum the sorted norms per <head,seq>
cumsum_others = sorted_norm.cumsum(dim=0) # (C, H, S)
@@ -118,9 +100,7 @@ def _select_threshold_query_head(gate: torch.Tensor,
cond = ratio >= remain_ratio.unsqueeze(0) # (C, H, S) boolean mask
any_cond = cond.any(dim=0) # (H, S)
# Find the index of the first True value along dim 0. If none, use C-1.
cutoff = torch.where(any_cond,
cond.float().argmax(dim=0),
torch.full_like(any_cond, fill_value=C - 1)) # (H, S)
cutoff = torch.where(any_cond, cond.float().argmax(dim=0), torch.full_like(any_cond, fill_value=C - 1)) # (H, S)
# 7) build a mask in sorted order up to that cutoff
idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1) # (C, 1, 1)
@@ -136,8 +116,7 @@ def _select_threshold_query_head(gate: torch.Tensor,
return final_gate_mask
def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
gate_self_chunk_mask: torch.Tensor,
def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor,
simsum_threshold: float) -> torch.Tensor:
"""
Selects <query, head> pairs for each block based on threshold.
@@ -148,16 +127,13 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
eps = 1e-6
# LSE‐style normalization per block (across heads and queries)
gate_masked = torch.where(valid_gate_mask, gate,
-torch.inf) # Use -inf for max
gate_min_val = torch.where(valid_gate_mask, gate,
torch.inf) # Use +inf for min
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1)
block_min = gate_min_val.amin(dim=(1, 2), keepdim=True) # (C, 1, 1)
block_denom = block_max - block_min
block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom),
block_denom) # (C, 1, 1)
block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom), block_denom) # (C, 1, 1)
gate_norm = (gate - block_min) / block_denom # (C, H, S)
gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
@@ -169,8 +145,7 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
# 2) compute how much more normalized weight each block needs beyond its self-chunk contributions
total_norm_sum_per_block = gate_norm.sum(dim=(1, 2)) # (C,)
remain_ratio = simsum_threshold - self_norm_sum_per_block / (
total_norm_sum_per_block + eps) # (C,)
remain_ratio = simsum_threshold - self_norm_sum_per_block / (total_norm_sum_per_block + eps) # (C,)
remain_ratio = torch.clamp(remain_ratio, min=0.0) # (C,)
# 3) zero out the self‐chunk entries in a copy, so we only sort “others”
@@ -179,31 +154,25 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
# 4) sort the other <head, seq> pairs by descending norm, per block
others_flat = others_norm.contiguous().view(C, HS) # (C, H*S)
sorted_others_flat, sorted_indices_flat = torch.sort(
others_flat, dim=1, descending=True) # (C, H*S)
sorted_others_flat, sorted_indices_flat = torch.sort(others_flat, dim=1, descending=True) # (C, H*S)
# 5) cumulative‑sum the sorted norms per block
cumsum_others_flat = sorted_others_flat.cumsum(dim=1) # (C, H*S)
# 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio
ratio_flat = cumsum_others_flat / (
total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S)
ratio_flat = cumsum_others_flat / (total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S)
cond_flat = ratio_flat >= remain_ratio.unsqueeze(1) # (C, H*S) boolean mask
any_cond = cond_flat.any(dim=1) # (C,)
# Find the index of the first True value along dim 1. If none, use HS-1.
cutoff_flat = torch.where(any_cond,
cond_flat.float().argmax(dim=1),
torch.full_like(any_cond,
fill_value=HS - 1)) # (C,)
cond_flat.float().argmax(dim=1), torch.full_like(any_cond, fill_value=HS - 1)) # (C,)
# 7) build a mask in sorted order up to that cutoff per block
idx_range_flat = torch.arange(HS,
device=gate.device).unsqueeze(0) # (1, H*S)
idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0) # (1, H*S)
sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1) # (C, H*S)
# 8) scatter it back to original <head, seq> order per block
others_mask_flat = torch.zeros_like(others_flat,
dtype=torch.bool) # (C, H*S)
others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C, H*S)
others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat)
others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
@@ -213,8 +182,7 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
return final_gate_mask
def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
gate_self_chunk_mask: torch.Tensor,
def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor,
simsum_threshold: float) -> torch.Tensor:
"""
Selects <chunk, query, head> triplets globally based on threshold.
@@ -225,17 +193,14 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
eps = 1e-6
# LSE‐style normalization globally across all valid entries
gate_masked = torch.where(valid_gate_mask, gate,
-torch.inf) # Use -inf for max
gate_min_val = torch.where(valid_gate_mask, gate,
torch.inf) # Use +inf for min
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
overall_max = gate_masked.max() # scalar
overall_min = gate_min_val.min() # scalar
overall_denom = overall_max - overall_min
overall_denom = torch.where(
overall_denom <= eps,
torch.tensor(1.0, device=gate.device, dtype=gate.dtype), overall_denom)
overall_denom = torch.where(overall_denom <= eps, torch.tensor(1.0, device=gate.device, dtype=gate.dtype),
overall_denom)
gate_norm = (gate - overall_min) / overall_denom # (C, H, S)
gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
@@ -247,8 +212,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
# 2) compute how much more normalized weight is needed globally beyond self-chunk contributions
total_norm_sum_overall = gate_norm.sum() # scalar
remain_ratio = simsum_threshold - self_norm_sum_overall / (
total_norm_sum_overall + eps) # scalar
remain_ratio = simsum_threshold - self_norm_sum_overall / (total_norm_sum_overall + eps) # scalar
remain_ratio = torch.clamp(remain_ratio, min=0.0) # scalar
# 3) zero out the self‐chunk entries in a copy, so we only sort “others”
@@ -264,18 +228,14 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
valid_others_indices = torch.where(valid_others_mask_flat)[0]
valid_others_values = others_flat[valid_others_indices]
sorted_others_values, sort_perm = torch.sort(
valid_others_values, descending=True) # (N_valid_others,)
sorted_original_indices = valid_others_indices[
sort_perm] # Original indices in C*H*S space, sorted by value
sorted_others_values, sort_perm = torch.sort(valid_others_values, descending=True) # (N_valid_others,)
sorted_original_indices = valid_others_indices[sort_perm] # Original indices in C*H*S space, sorted by value
# 5) cumulative‑sum the sorted valid 'other' norms globally
cumsum_others_values = sorted_others_values.cumsum(
dim=0) # (N_valid_others,)
cumsum_others_values = sorted_others_values.cumsum(dim=0) # (N_valid_others,)
# 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally
ratio_values = cumsum_others_values / (total_norm_sum_overall + eps
) # (N_valid_others,)
ratio_values = cumsum_others_values / (total_norm_sum_overall + eps) # (N_valid_others,)
cond_values = ratio_values >= remain_ratio # (N_valid_others,) boolean mask
any_cond = cond_values.any() # scalar
@@ -283,19 +243,15 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
cutoff_idx_in_sorted = torch.where(
any_cond,
cond_values.float().argmax(dim=0),
torch.tensor(len(sorted_others_values) - 1,
device=gate.device,
dtype=torch.long))
torch.tensor(len(sorted_others_values) - 1, device=gate.device, dtype=torch.long))
# 7) build a mask selecting the top-k others based on the cutoff
# Select the original indices corresponding to the top entries in the sorted list
selected_other_indices = sorted_original_indices[:cutoff_idx_in_sorted + 1]
# 8) create the mask in the original flat shape
others_mask_flat = torch.zeros_like(others_flat,
dtype=torch.bool) # (C*H*S,)
if selected_other_indices.numel(
) > 0: # Check if any 'other' indices were selected
others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,)
if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected
others_mask_flat[selected_other_indices] = True
others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
@@ -305,9 +261,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
return final_gate_mask
def _select_threshold_head_global(gate: torch.Tensor,
valid_gate_mask: torch.Tensor,
gate_self_chunk_mask: torch.Tensor,
def _select_threshold_head_global(gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor,
simsum_threshold: float) -> torch.Tensor:
"""
Selects <chunk, query> globally for each head based on threshold.
@@ -334,8 +288,7 @@ def _select_threshold_head_global(gate: torch.Tensor,
total_norm_sum = gate_norm.sum(dim=(0, 2)) # (H,)
# 4) how much more normalized weight needed per head
remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps
) # (H,)
remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps) # (H,)
remain_ratio = torch.clamp(remain_ratio, min=0.0)
# 5) zero out self‐chunk entries to focus on "others"
@@ -349,10 +302,8 @@ def _select_threshold_head_global(gate: torch.Tensor,
.permute(1, 0, 2).reshape(H, CS) # (H, C*S)
# 7) vectorized selection of “others” per head
masked_flat = torch.where(valid_flat, others_flat,
torch.zeros_like(others_flat))
sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1,
descending=True) # (H, C*S)
masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat))
sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1, descending=True) # (H, C*S)
cumsum_vals = sorted_vals.cumsum(dim=1) # (H, C*S)
ratio_vals = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps) # (H, C*S)
@@ -360,8 +311,7 @@ def _select_threshold_head_global(gate: torch.Tensor,
has_cutoff = cond.any(dim=1) # (H,)
default = torch.full((H, ), CS - 1, device=gate.device, dtype=torch.long)
cutoff = torch.where(has_cutoff,
cond.float().argmax(dim=1), default) # (H,)
cutoff = torch.where(has_cutoff, cond.float().argmax(dim=1), default) # (H,)
idx_range = torch.arange(CS, device=gate.device).unsqueeze(0) # (1, C*S)
sorted_mask = idx_range <= cutoff.unsqueeze(1) # (H, C*S)
@@ -430,25 +380,19 @@ class MixedAttention(torch.autograd.Function):
self_attn_lse_sh = self_attn_lse_hs.t().contiguous()
moba_attn_lse = moba_attn_lse_hs.t().contiguous()
output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]),
device=q.device,
dtype=torch.float32)
output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
output_2d = output.view(-1, q.shape[2])
max_lse_1d = self_attn_lse_sh.view(-1)
max_lse_1d = max_lse_1d.index_reduce(0, moba_q_sh_indices,
moba_attn_lse.view(-1), "amax")
self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(
self_attn_lse_sh)
moba_attn_lse = (moba_attn_lse.view(-1).sub(
max_lse_1d.index_select(
0, moba_q_sh_indices)).reshape_as(moba_attn_lse))
max_lse_1d = max_lse_1d.index_reduce(0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax")
self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh)
moba_attn_lse = (moba_attn_lse.view(-1).sub(max_lse_1d.index_select(
0, moba_q_sh_indices)).reshape_as(moba_attn_lse))
mixed_attn_se_sh = self_attn_lse_sh.exp()
moba_attn_se = moba_attn_lse.exp()
mixed_attn_se_sh.view(-1).index_add_(0, moba_q_sh_indices,
moba_attn_se.view(-1))
mixed_attn_se_sh.view(-1).index_add_(0, moba_q_sh_indices, moba_attn_se.view(-1))
mixed_attn_lse_sh = mixed_attn_se_sh.log()
# Combine self-attention output
@@ -457,15 +401,13 @@ class MixedAttention(torch.autograd.Function):
output_2d += self_attn_out_sh.reshape_as(output_2d)
# Combine MOBA attention output
mixed_attn_lse = (mixed_attn_lse_sh.view(-1).index_select(
0, moba_q_sh_indices).view_as(moba_attn_lse))
mixed_attn_lse = (mixed_attn_lse_sh.view(-1).index_select(0, moba_q_sh_indices).view_as(moba_attn_lse))
factor = (moba_attn_lse - mixed_attn_lse).exp() # [S, H]
moba_attn_out = moba_attn_out * factor.unsqueeze(-1)
raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1])
output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out)
output = output.to(q.dtype)
mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(
mixed_attn_se_sh)
mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh)
ctx.save_for_backward(
output,
mixed_attn_lse_sh,
@@ -508,37 +450,33 @@ class MixedAttention(torch.autograd.Function):
dq = torch.empty_like(q)
dk = torch.empty_like(k)
dv = torch.empty_like(v)
_ = _flash_attn_varlen_backward(
dout=d_output,
q=q,
k=k,
v=v,
out=output,
softmax_lse=mixed_attn_vlse_sh.t().contiguous(),
dq=dq,
dk=dk,
dv=dv,
cu_seqlens_q=self_attn_cu_seqlen,
cu_seqlens_k=self_attn_cu_seqlen,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
softmax_scale=softmax_scale,
causal=False,
dropout_p=0.0,
softcap=0.0,
alibi_slopes=None,
deterministic=True,
window_size_left=-1,
window_size_right=-1)
_ = _flash_attn_varlen_backward(dout=d_output,
q=q,
k=k,
v=v,
out=output,
softmax_lse=mixed_attn_vlse_sh.t().contiguous(),
dq=dq,
dk=dk,
dv=dv,
cu_seqlens_q=self_attn_cu_seqlen,
cu_seqlens_k=self_attn_cu_seqlen,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
softmax_scale=softmax_scale,
causal=False,
dropout_p=0.0,
softcap=0.0,
alibi_slopes=None,
deterministic=True,
window_size_left=-1,
window_size_right=-1)
headdim = q.shape[-1]
d_moba_output = (d_output.view(-1, headdim).index_select(
0, moba_q_sh_indices).unsqueeze(1))
moba_output = (output.view(-1, headdim).index_select(
0, moba_q_sh_indices).unsqueeze(1))
d_moba_output = (d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1))
moba_output = (output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1))
mixed_attn_vlse = (mixed_attn_vlse_sh.view(-1).index_select(
0, moba_q_sh_indices).view(1, -1))
mixed_attn_vlse = (mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1))
dmq = torch.empty_like(moba_q)
dmkv = torch.empty_like(moba_kv)
@@ -595,8 +533,7 @@ def moba_attn_varlen(
seqlen, num_head, head_dim = q.shape
# Compute chunk boundaries.
cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(
cu_seqlens, moba_chunk_size)
cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(cu_seqlens, moba_chunk_size)
self_attn_cu_seqlen = cu_chunk
@@ -609,20 +546,15 @@ def moba_attn_varlen(
chunk_lengths = chunk_ends - chunk_starts # [num_filtered_chunk]
max_chunk_len = int(chunk_lengths.max().item())
range_tensor = torch.arange(max_chunk_len,
device=kv.device,
dtype=chunk_starts.dtype).unsqueeze(0)
range_tensor = torch.arange(max_chunk_len, device=kv.device, dtype=chunk_starts.dtype).unsqueeze(0)
indices = chunk_starts.unsqueeze(1) + range_tensor
indices = torch.clamp(indices, max=kv.shape[0] - 1)
valid_mask = range_tensor < chunk_lengths.unsqueeze(1)
gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len,
*kv.shape[1:])
gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(
-1).type_as(gathered)
gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:])
gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered)
# Compute key_gate_weight over valid tokens.
key_values = gathered[:, :, 0].float(
) # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
key_values = gathered[:, :, 0].float() # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1)
key_sum = (key_values * valid_mask_exp).sum(dim=1)
divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1)
@@ -631,19 +563,15 @@ def moba_attn_varlen(
# Compute gate logits between key_gate_weight and queries.
q_float = q.float()
# gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float) # [num_filtered_chunk, num_head, seqlen]
gate = torch.bmm(key_gate_weight.permute(1, 0, 2),
q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2)
gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2)
# Amplify the diagonal (self chunk) contributions.
gate_seq_idx = torch.arange(seqlen, device=q.device,
dtype=torch.int32).unsqueeze(0).expand(
num_filtered_chunk, seqlen)
dtype=torch.int32).unsqueeze(0).expand(num_filtered_chunk, seqlen)
chunk_start = cu_chunk[filtered_chunk_indices] # [num_filtered_chunk]
chunk_end = cu_chunk[filtered_chunk_indices + 1] # [num_filtered_chunk]
gate_self_chunk_mask = (
(gate_seq_idx >= chunk_start.unsqueeze(1)) &
(gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand(
-1, num_head, -1)
gate_self_chunk_mask = ((gate_seq_idx >= chunk_start.unsqueeze(1)) &
(gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand(-1, num_head, -1)
amplification_factor = 1e9 # Example factor; adjust as needed.
origin_gate = gate.clone()
gate = gate.clone()
@@ -664,11 +592,7 @@ def moba_attn_varlen(
if threshold_type == 'query_head':
# === per‐<head,seq> top-k across chunks (original behavior) ===
# gate: (C, H, S)
_, gate_topk_idx = torch.topk(gate,
k=moba_topk,
dim=0,
largest=True,
sorted=False)
_, gate_topk_idx = torch.topk(gate, k=moba_topk, dim=0, largest=True, sorted=False)
gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool)
gate_idx_mask.scatter_(0, gate_topk_idx, True)
gate_mask = valid_gate_mask & gate_idx_mask
@@ -679,85 +603,52 @@ def moba_attn_varlen(
flat_mask = valid_gate_mask.flatten()
flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf"))
# pick topk global entries
vals, idx = torch.topk(flat_gate_masked,
k=moba_topk * H * S,
largest=True,
sorted=False)
vals, idx = torch.topk(flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False)
others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool)
others_mask_flat[idx] = True
gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(
gate.shape)
gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape)
elif threshold_type == 'head_global':
# per-head top-k across all chunks and sequence positions
C, H, S = gate.shape
CS = C * S
flat_gate = gate.permute(1, 0, 2).reshape(H, CS)
flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS)
flat_gate_masked = torch.where(
flat_valid, flat_gate, torch.full_like(flat_gate,
-float('inf')))
flat_gate_masked = torch.where(flat_valid, flat_gate, torch.full_like(flat_gate, -float('inf')))
# pick top-k indices per head
_, topk_idx = torch.topk(flat_gate_masked,
k=moba_topk * S,
dim=1,
largest=True,
sorted=False)
_, topk_idx = torch.topk(flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False)
gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool)
gate_idx_flat.scatter_(1, topk_idx, True)
gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2)
else:
raise ValueError(
f"Invalid threshold_type for topk: {threshold_type}. "
"Choose 'query_head', 'block', or 'overall'.")
raise ValueError(f"Invalid threshold_type for topk: {threshold_type}. "
"Choose 'query_head', 'block', or 'overall'.")
elif select_mode == 'threshold':
# Delegate to the specific thresholding function
valid_gate_mask = gate != -float("inf") # (num_chunk, num_head, seqlen)
if threshold_type == 'query_head':
gate_mask = _select_threshold_query_head(gate, valid_gate_mask,
gate_self_chunk_mask,
simsum_threshold)
gate_mask = _select_threshold_query_head(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
elif threshold_type == 'block':
gate_mask = _select_threshold_block(gate, valid_gate_mask,
gate_self_chunk_mask,
simsum_threshold)
gate_mask = _select_threshold_block(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
elif threshold_type == 'overall':
gate_mask = _select_threshold_overall(gate, valid_gate_mask,
gate_self_chunk_mask,
simsum_threshold)
gate_mask = _select_threshold_overall(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
elif threshold_type == 'head_global':
gate_mask = _select_threshold_head_global(gate, valid_gate_mask,
gate_self_chunk_mask,
simsum_threshold)
gate_mask = _select_threshold_head_global(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
else:
raise ValueError(
f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'."
)
raise ValueError(f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'.")
else:
raise ValueError(
f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'."
)
raise ValueError(f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'.")
# eliminate self_chunk in MoBA branch
gate_mask = gate_mask & ~gate_self_chunk_mask
# if gate_mask is all false, perform flash_attn instead
if gate_mask.sum() == 0:
return flash_attn_varlen_func(q,
k,
v,
cu_seqlens,
cu_seqlens,
max_seqlen,
max_seqlen,
causal=False)
return flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False)
# Determine which query positions are selected.
# nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index].
moba_q_indices = gate_mask.reshape(
gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)]
moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (
moba_q_indices // seqlen)
moba_q = rearrange(q, "s h d -> (h s) d").index_select(
0, moba_q_indices).unsqueeze(1)
moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)]
moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen)
moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1)
# Build cumulative sequence lengths for the selected queries.
moba_seqlen_q = gate_mask.sum(dim=-1).flatten()
@@ -775,23 +666,19 @@ def moba_attn_varlen(
# Rearrange gathered KV for the MOBA branch.
experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d")
valid_expert_lengths = chunk_lengths.unsqueeze(1).expand(
num_filtered_chunk, num_head).reshape(-1).to(torch.int32)
valid_expert_lengths = chunk_lengths.unsqueeze(1).expand(num_filtered_chunk, num_head).reshape(-1).to(torch.int32)
if q_zero_mask.sum() > 0:
experts_tensor = experts_tensor[valid_expert_mask]
valid_expert_lengths = valid_expert_lengths[valid_expert_mask]
seq_range = torch.arange(experts_tensor.shape[1],
device=experts_tensor.device).unsqueeze(0)
seq_range = torch.arange(experts_tensor.shape[1], device=experts_tensor.device).unsqueeze(0)
mask = seq_range < valid_expert_lengths.unsqueeze(1)
moba_kv = experts_tensor[mask] # Shape: ((nc h cl_valid) two d)
moba_kv = moba_kv.unsqueeze(2) # Shape: ((nc h cl_valid) two 1 d)
moba_cu_seqlen_kv = torch.cat(
[
torch.zeros(1, device=experts_tensor.device, dtype=torch.int32),
valid_expert_lengths.cumsum(dim=0)
],
[torch.zeros(1, device=experts_tensor.device, dtype=torch.int32),
valid_expert_lengths.cumsum(dim=0)],
dim=0,
).to(torch.int32)
@@ -831,18 +718,14 @@ def process_moba_input(
torch.Tensor: Processed input tensor.
"""
if isinstance(chunk_size, float) or isinstance(chunk_size, int):
moba_chunk_size = int(chunk_size * patch_resolution[1] *
patch_resolution[2])
moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2])
else:
assert isinstance(
chunk_size, (tuple, list)
), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}"
assert isinstance(chunk_size,
(tuple, list)), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}"
if len(chunk_size) == 2:
assert patch_resolution[1] % chunk_size[
0] == 0 and patch_resolution[2] % chunk_size[
1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}"
nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[
2] // chunk_size[1]
assert patch_resolution[1] % chunk_size[0] == 0 and patch_resolution[2] % chunk_size[
1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}"
nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[2] // chunk_size[1]
x = rearrange(x,
"b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d",
t=patch_resolution[0],
@@ -852,26 +735,22 @@ def process_moba_input(
cw=chunk_size[1])
moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1]
elif len(chunk_size) == 3:
assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[
1] % chunk_size[1] == 0 and patch_resolution[2] % chunk_size[
assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[1] % chunk_size[
1] == 0 and patch_resolution[2] % chunk_size[
2] == 0, f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}"
nct, nch, ncw = patch_resolution[0] // chunk_size[
0], patch_resolution[1] // chunk_size[1], patch_resolution[
2] // chunk_size[2]
x = rearrange(
x,
"b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d",
nct=nct,
nch=nch,
ncw=ncw,
ct=chunk_size[0],
ch=chunk_size[1],
cw=chunk_size[2])
nct, nch, ncw = patch_resolution[0] // chunk_size[0], patch_resolution[1] // chunk_size[
1], patch_resolution[2] // chunk_size[2]
x = rearrange(x,
"b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d",
nct=nct,
nch=nch,
ncw=ncw,
ct=chunk_size[0],
ch=chunk_size[1],
cw=chunk_size[2])
moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2]
else:
raise ValueError(
f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}"
)
raise ValueError(f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}")
return x, moba_chunk_size
@@ -892,15 +771,14 @@ def process_moba_output(
ch=chunk_size[0],
cw=chunk_size[1])
elif len(chunk_size) == 3:
x = rearrange(
x,
"b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d",
nct=patch_resolution[0] // chunk_size[0],
nch=patch_resolution[1] // chunk_size[1],
ncw=patch_resolution[2] // chunk_size[2],
ct=chunk_size[0],
ch=chunk_size[1],
cw=chunk_size[2])
x = rearrange(x,
"b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d",
nct=patch_resolution[0] // chunk_size[0],
nch=patch_resolution[1] // chunk_size[1],
ncw=patch_resolution[2] // chunk_size[2],
ct=chunk_size[0],
ch=chunk_size[1],
cw=chunk_size[2])
return x
@@ -912,18 +790,11 @@ def generate_data(batch_size, seqlen, num_head, head_dim, dtype):
torch.cuda.manual_seed(0)
device = torch.cuda.current_device()
q = torch.randn((batch_size, seqlen, num_head, head_dim),
requires_grad=True).to(dtype=dtype, device='cuda')
k = torch.randn((batch_size, seqlen, num_head, head_dim),
requires_grad=True).to(dtype=dtype, device='cuda')
v = torch.randn((batch_size, seqlen, num_head, head_dim),
requires_grad=True).to(dtype=dtype, device='cuda')
q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}")
cu_seqlens = torch.arange(0,
q.shape[0] * q.shape[1] + 1,
q.shape[1],
dtype=torch.int32,
device='cuda')
cu_seqlens = torch.arange(0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device='cuda')
max_seqlen = q.shape[1]
q = rearrange(q, "b s ... -> (b s) ...")
k = rearrange(k, "b s ... -> (b s) ...")
@@ -944,8 +815,7 @@ def test_attn_varlen_moba_speed(batch,
threshold_type='query_head'):
"""Speed test comparing flash_attn vs moba_attention"""
# Get data
q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head,
head_dim, dtype)
q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype)
print(
f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}"
)
@@ -957,27 +827,13 @@ def test_attn_varlen_moba_speed(batch,
# Warmup
for _ in range(warmup_iters):
o = flash_attn_varlen_func(q,
k,
v,
cu_seqlen,
cu_seqlen,
max_seqlen,
max_seqlen,
causal=False)
o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
torch.autograd.backward(o, vo_grad)
torch.cuda.synchronize()
start_flash = time.perf_counter()
for _ in range(perf_test_iters):
o = flash_attn_varlen_func(q,
k,
v,
cu_seqlen,
cu_seqlen,
max_seqlen,
max_seqlen,
causal=False)
o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
torch.autograd.backward(o, vo_grad)
torch.cuda.synchronize()
fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py
+4 -10
View File
@@ -78,7 +78,7 @@ def construct_variable_block_sizes(
n_t, n_h, n_w = num_tiles
def _sizes(dim_len: int, tile: int, n: int) -> torch.LongTensor:
sizes = torch.full((n,), tile, dtype=torch.int, device=device)
sizes = torch.full((n, ), tile, dtype=torch.int, device=device)
remainder = dim_len - (n - 1) * tile
sizes[-1] = remainder if remainder > 0 else tile
return sizes
@@ -87,11 +87,7 @@ def construct_variable_block_sizes(
h_sizes = _sizes(h, ts_h, n_h)
w_sizes = _sizes(w, ts_w, n_w)
return (
t_sizes[:, None, None]
* h_sizes[None, :, None]
* w_sizes[None, None, :]
).reshape(-1)
return (t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :]).reshape(-1)
def get_non_pad_index(
@@ -134,10 +130,8 @@ def build_vsa_metadata(
ts_t, ts_h, ts_w = tile_size
max_block_size = math.prod(tile_size)
if max_block_size not in _SUPPORTED_VSA_BLOCK_VOLUMES:
raise ValueError(
f"Unsupported VSA tile volume {max_block_size} for tile_size={tile_size}; "
f"supported volumes are {_SUPPORTED_VSA_BLOCK_VOLUMES}."
)
raise ValueError(f"Unsupported VSA tile volume {max_block_size} for tile_size={tile_size}; "
f"supported volumes are {_SUPPORTED_VSA_BLOCK_VOLUMES}.")
num_tiles = (
math.ceil(T / ts_t),