8b51380466
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com> Co-authored-by: Aryan Kumar <aryank@Aryans-Mac-Studio.local> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
339 lines
18 KiB
Python
339 lines
18 KiB
Python
"""Per-component timing of one MiniMax-H3 block on the current GPU. Imported inside the Modal container."""
|
|
import json
|
|
import math
|
|
import time
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
HID, HEADS, HD, FFN = 5376, 56, 128, 14336
|
|
|
|
|
|
def timed(fn, iters=5, warmup=2):
|
|
for _ in range(warmup):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
|
|
start.record()
|
|
for _ in range(iters):
|
|
fn()
|
|
end.record()
|
|
torch.cuda.synchronize()
|
|
return round(start.elapsed_time(end) / iters, 3)
|
|
|
|
|
|
def smooth_tokens(prefix_len, shape, dev, gen, noise=0.5):
|
|
"""Token features with spatial-temporal correlation (proxy for real q/k), packed [prefix | video]."""
|
|
t, h, w = shape
|
|
low = torch.randn(1, HEADS * HD, max(2, t // 6), max(2, h // 6), max(2, w // 6), device=dev, generator=gen)
|
|
vid = F.interpolate(low, size=(t, h, w), mode="trilinear", align_corners=False)
|
|
vid = vid.reshape(HEADS, HD, -1).permute(2, 0, 1)
|
|
vid = vid + noise * torch.randn(vid.shape, device=dev, generator=gen)
|
|
pre = torch.randn(prefix_len, HEADS, HD, device=dev, generator=gen)
|
|
return torch.cat([pre, vid]).to(torch.bfloat16)[None] # [1, L, H, D]
|
|
|
|
|
|
def run(shapes):
|
|
from flashinfer import SfLayout, mm_fp4, nvfp4_quantize
|
|
from fastvideo.attention.backends.video_sparse_attn_h3 import (_build_block_mask, _h3_tile_geometry, _pool_tiles)
|
|
from fastvideo_kernel.block_sparse_attn import block_sparse_attn as bsa64
|
|
from fastvideo.models.dits.minimax_h3_fusions import (fused_qknorm_rope, fused_residual_gate_rmsnorm_modulate,
|
|
fused_rmsnorm_modulate, minimax_h3_swiglu)
|
|
import attn_qat_infer.api as fa
|
|
from attn_qat_infer.api import BLOCK_M, sageattn_blackwell, sageattn_blackwell_sparse
|
|
|
|
dev = torch.device("cuda")
|
|
gen = torch.Generator(device=dev).manual_seed(0)
|
|
report = {"device": torch.cuda.get_device_name(0), "torch": torch.__version__}
|
|
unit = torch.tensor(1.0, device=dev)
|
|
|
|
def fp4_weight(n, k):
|
|
w = torch.randn(n, k, device=dev, dtype=torch.bfloat16) * 0.02
|
|
gsf = (448 * 6) / w.float().abs().max()
|
|
wq, ws = nvfp4_quantize(w, gsf, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
|
|
return w, wq, ws, (1.0 / gsf).float()
|
|
|
|
for name, prefix_segments, video_shape in shapes:
|
|
prefix_segments = tuple(prefix_segments)
|
|
video_shape = tuple(video_shape)
|
|
L = sum(prefix_segments) + math.prod(video_shape)
|
|
r = {"tokens": L}
|
|
torch.cuda.empty_cache()
|
|
|
|
# ---------------- linears ----------------
|
|
lin = {}
|
|
for lname, n, k in (("qkv_one", 7168, HID), ("out", HID, 7168), ("fc_in", 2 * FFN, HID),
|
|
("fc_out", HID, FFN)):
|
|
w, wq, ws, alpha = fp4_weight(n, k)
|
|
x = torch.randn(L, k, device=dev, dtype=torch.bfloat16)
|
|
xq, xs = nvfp4_quantize(x, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
|
|
entry = {"quant_ms": timed(lambda: nvfp4_quantize(x, unit, sfLayout=SfLayout.layout_128x4,
|
|
do_shuffle=False))}
|
|
for backend in ("cudnn", "cutlass", "auto"):
|
|
try:
|
|
entry[f"mm_{backend}_ms"] = timed(
|
|
lambda: mm_fp4(xq, wq.T, xs, ws.T, alpha, torch.bfloat16, None, backend=backend))
|
|
except Exception as exc: # noqa: BLE001
|
|
entry[f"mm_{backend}_ms"] = f"ERR {type(exc).__name__}: {str(exc)[:120]}"
|
|
entry["bf16_ms"] = timed(lambda: F.linear(x, w))
|
|
entry["tflops_fp4_auto"] = (round(2 * L * n * k / entry["mm_auto_ms"] / 1e9, 1)
|
|
if isinstance(entry["mm_auto_ms"], float) else None)
|
|
lin[lname] = entry
|
|
del w, wq, ws, x, xq, xs
|
|
r["linear"] = lin
|
|
|
|
# ---------------- attention ----------------
|
|
geom = _h3_tile_geometry(prefix_segments, video_shape, dev, (4, 4, 4))
|
|
_, vbs, untile, n_prefix, n_video = geom
|
|
n_tiles = vbs.numel()
|
|
Lpad = n_tiles * 64
|
|
att = {"tiles64": n_tiles, "prefix_tiles": n_prefix}
|
|
prefix_len = sum(prefix_segments)
|
|
q = smooth_tokens(prefix_len, video_shape, dev, gen)
|
|
k = (q + 0.5 * torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16)).to(torch.bfloat16)
|
|
v = torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16)
|
|
g = torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16) * 0.1
|
|
|
|
def tile(x):
|
|
buf = torch.zeros((1, Lpad, HEADS, HD), device=dev, dtype=x.dtype)
|
|
buf[:, untile] = x
|
|
return buf
|
|
|
|
qt, kt, vt, gt = tile(q), tile(k), tile(v), tile(g)
|
|
att["tile_scatter4_gather1_ms"] = timed(lambda: (tile(q), tile(k), tile(v), tile(g), qt[:, untile]))
|
|
|
|
def vsa_full(sparsity=0.8, with_gate=True):
|
|
qp = _pool_tiles(qt, vbs, 64)
|
|
kp = _pool_tiles(kt, vbs, 64)
|
|
scores = torch.matmul(qp, kp.transpose(-2, -1)) / HD**0.5
|
|
mask = _build_block_mask(scores, n_prefix, sparsity, True, ((n_prefix, n_prefix + n_video), ), (sparsity, ))
|
|
out, _ = bsa64(qt.transpose(1, 2).contiguous(), kt.transpose(1, 2).contiguous(),
|
|
vt.transpose(1, 2).contiguous(), mask, vbs)
|
|
out = out.transpose(1, 2).contiguous()
|
|
if with_gate:
|
|
vp = _pool_tiles(vt, vbs, 64)
|
|
oc = torch.matmul(torch.softmax(scores, dim=-1), vp).permute(0, 2, 1, 3).to(out.dtype)
|
|
out = (out.view(1, n_tiles, 64, HEADS, HD) + oc.unsqueeze(2) * gt.view(1, n_tiles, 64, HEADS, HD)).view(
|
|
1, Lpad, HEADS, HD)
|
|
return out, mask
|
|
|
|
_, mask = vsa_full()
|
|
att["tile_density"] = round(mask.float().mean().item(), 4)
|
|
att["vsa_triton_total_ms"] = timed(vsa_full)
|
|
qb, kb, vb = (x.transpose(1, 2).contiguous() for x in (qt, kt, vt))
|
|
att["vsa_triton_kernel_only_ms"] = timed(lambda: bsa64(qb, kb, vb, mask, vbs))
|
|
|
|
# dense FP4 on the unpadded packed sequence
|
|
qd, kd, vd = (x.transpose(1, 2).contiguous() for x in (q, k, v))
|
|
att["sage3_dense_fp4_ms"] = timed(lambda: sageattn_blackwell(qd, kd, vd), iters=3)
|
|
# same with delta_s cached + per_block_mean=False (no 9.5 GB memset)
|
|
zero_ds = {}
|
|
|
|
def sage_nods(qx, kx, vx):
|
|
QL, KL = qx.size(2), kx.size(2)
|
|
qx, kx, vx = (F.pad(x, (0, 0, 0, (BLOCK_M - x.size(2) % BLOCK_M) % BLOCK_M)).contiguous() for x in (qx, kx, vx))
|
|
key = (qx.shape[0], qx.shape[1], kx.shape[2])
|
|
if key not in zero_ds:
|
|
zero_ds[key] = torch.zeros((qx.shape[0], qx.shape[1], 1, kx.shape[2]), device=dev, dtype=torch.float32)
|
|
ql, kl, vl = fa.scale_and_quant_fp4(qx), fa.scale_and_quant_fp4_permute(kx), fa.scale_and_quant_fp4_transpose(vx)
|
|
return fa.blockscaled_fp4_attn(ql, kl, vl, zero_ds[key], KL, False, False, True, True, None)[0][:, :, :QL]
|
|
|
|
att["sage3_dense_fp4_no_deltas_ms"] = timed(lambda: sage_nods(qd, kd, vd), iters=3)
|
|
out_a = sageattn_blackwell(qd, kd, vd)
|
|
out_b = sage_nods(qd, kd, vd)
|
|
att["no_deltas_max_abs_diff"] = (out_a.float() - out_b.float()).abs().max().item()
|
|
Lq = qd.shape[2]
|
|
att["dense_fp4_tflops"] = round(4 * Lq * Lq * HD * HEADS / att["sage3_dense_fp4_no_deltas_ms"] / 1e9, 1)
|
|
if L < 30000:
|
|
att["sdpa_bf16_dense_ms"] = timed(lambda: F.scaled_dot_product_attention(qd, kd, vd), iters=3)
|
|
|
|
# sparse FP4 on 128x128 blocks with 64x64 quadrant masks (exact VSA tile-64 semantics)
|
|
nt2 = n_tiles + (n_tiles % 2)
|
|
qs, ks, vs = (F.pad(x.transpose(1, 2), (0, 0, 0, nt2 * 64 - Lpad)).contiguous() for x in (qt, kt, vt))
|
|
q2k_idx, q2k_num, kvv, quad = fa.vsa_tile_mask_to_fp4_blocks(mask, 64, vbs, validate=True)
|
|
att["block128_density"] = round(q2k_num.float().mean().item() / (nt2 // 2), 4)
|
|
att["sparse_fp4_quad_ms"] = timed(lambda: sageattn_blackwell_sparse(qs, ks, vs, q2k_idx, q2k_num, kvv, quad))
|
|
att["mask_to_blocks_ms"] = timed(lambda: fa.vsa_tile_mask_to_fp4_blocks(mask, 64, vbs))
|
|
o_fp4 = sageattn_blackwell_sparse(qs, ks, vs, q2k_idx, q2k_num, kvv, quad)[:, :, :Lpad]
|
|
o_tri, _ = bsa64(qb, kb, vb, mask, vbs)
|
|
rows = untile # valid (non-pad) rows of the tile buffer
|
|
diff = (o_fp4[:, :, rows].float() - o_tri[:, :, rows].float())
|
|
att["quad_vs_triton_rel_l2"] = round((diff.norm() / o_tri[:, :, rows].float().norm()).item(), 4)
|
|
od = sageattn_blackwell(qd, kd, vd)
|
|
from torch.nn.attention import SDPBackend, sdpa_kernel
|
|
if L < 30000:
|
|
ref_dense = F.scaled_dot_product_attention(qd, kd, vd)
|
|
att["dense_fp4_vs_bf16_rel_l2"] = round(((od.float() - ref_dense.float()).norm() /
|
|
ref_dense.float().norm()).item(), 4)
|
|
att["quad_finite"] = bool(torch.isfinite(o_fp4[:, :, rows]).all())
|
|
r["attention"] = att
|
|
del q, k, v, g, qt, kt, vt, gt, qd, kd, vd, qs, ks, vs, qb, kb, vb, out_a, out_b, zero_ds, o_fp4, o_tri, od
|
|
|
|
# ---------------- elementwise ----------------
|
|
ew = {}
|
|
x = torch.randn(1, L, HID, device=dev, dtype=torch.bfloat16)
|
|
br = torch.randn_like(x)
|
|
nw = torch.randn(HID, device=dev, dtype=torch.bfloat16)
|
|
tab = [torch.randn(3, HID, device=dev, dtype=torch.bfloat16) * 0.1 for _ in range(6)]
|
|
idx = torch.randint(0, 3, (L,), device=dev)
|
|
norm = torch.nn.RMSNorm(HID, eps=1e-6, device=dev, dtype=torch.bfloat16)
|
|
|
|
def eager_mod():
|
|
n1 = norm(x) * (1.0 + tab[1].index_select(0, idx)) + tab[0].index_select(0, idx)
|
|
h = x + tab[2].index_select(0, idx) * br
|
|
n2 = norm(h) * (1.0 + tab[4].index_select(0, idx)) + tab[3].index_select(0, idx)
|
|
out = h + tab[5].index_select(0, idx) * br
|
|
return n1, n2, out
|
|
|
|
def fused_mod():
|
|
n1 = fused_rmsnorm_modulate(x, nw, tab[1], tab[0], idx, 1e-6)
|
|
h, n2 = fused_residual_gate_rmsnorm_modulate(x, br, tab[2], nw, tab[4], tab[3], idx, 1e-6)
|
|
out = h + tab[5].index_select(0, idx) * br
|
|
return n1, n2, out
|
|
|
|
with torch.no_grad():
|
|
ew["modulate_eager_ms"] = timed(eager_mod)
|
|
ew["modulate_fused_ms"] = timed(fused_mod)
|
|
ew["modulate_compiled_ms"] = timed(torch.compile(eager_mod))
|
|
packed = torch.randn(1, L, 2 * FFN, device=dev, dtype=torch.bfloat16)
|
|
|
|
def eager_swiglu():
|
|
a, b = packed.chunk(2, dim=-1)
|
|
return a * F.silu(b)
|
|
|
|
ew["swiglu_eager_ms"] = timed(eager_swiglu)
|
|
ew["swiglu_fused_ms"] = timed(lambda: minimax_h3_swiglu(packed))
|
|
del packed
|
|
qq = torch.randn(1, L, HEADS, HD, device=dev, dtype=torch.bfloat16)
|
|
cos = torch.randn(L, 96, device=dev, dtype=torch.bfloat16)
|
|
sin = torch.randn(L, 96, device=dev, dtype=torch.bfloat16)
|
|
qn = torch.nn.RMSNorm(HD, eps=1e-6, device=dev, dtype=torch.bfloat16)
|
|
|
|
def eager_rope():
|
|
outs = []
|
|
for t in (qq, qq):
|
|
t = qn(t)
|
|
rot, pas = t[..., :96], t[..., 96:]
|
|
c, s = cos[None, :, None, :], sin[None, :, None, :]
|
|
a, b = rot.chunk(2, dim=-1)
|
|
outs.append(torch.cat((rot * c + torch.cat((-b, a), -1) * s, pas), -1).contiguous())
|
|
return outs
|
|
|
|
ew["qknorm_rope_eager_ms"] = timed(eager_rope)
|
|
ew["qknorm_rope_fused_ms"] = timed(
|
|
lambda: (fused_qknorm_rope(qq, qn.weight, cos, sin, 1e-6), fused_qknorm_rope(qq, qn.weight, cos, sin, 1e-6)))
|
|
r["elementwise"] = ew
|
|
report[name] = r
|
|
print("BENCH", name, json.dumps(r), flush=True)
|
|
return report
|
|
|
|
|
|
def check_tile64():
|
|
"""Exact semantics check of the quadrant path on small H3-style layouts."""
|
|
from fastvideo.attention.backends.video_sparse_attn_h3 import _build_block_mask, _h3_tile_geometry, _pool_tiles
|
|
import attn_qat_infer.api as fa
|
|
from attn_qat_infer.api import sageattn_blackwell, sageattn_blackwell_sparse
|
|
|
|
dev = torch.device("cuda")
|
|
gen = torch.Generator(device=dev).manual_seed(1)
|
|
out = {}
|
|
cases = [("odd_tiles", (77, 46), (9, 6, 10), 64), ("even_tiles", (64, 40), (8, 8, 8), 64),
|
|
("tile256", (100, 70), (8, 8, 16), 256)]
|
|
for name, prefix, vshape, tt in cases:
|
|
H = 4
|
|
shape = {64: (4, 4, 4), 256: (4, 8, 8)}[tt]
|
|
_, vbs, untile, n_prefix, n_video = _h3_tile_geometry(prefix, vshape, dev, shape)
|
|
n_tiles = vbs.numel()
|
|
Lpad = n_tiles * tt
|
|
L = sum(prefix) + math.prod(vshape)
|
|
q = torch.randn(1, L, H, HD, device=dev, generator=gen, dtype=torch.bfloat16)
|
|
k = torch.randn(1, L, H, HD, device=dev, generator=gen, dtype=torch.bfloat16)
|
|
v = torch.randn(1, L, H, HD, device=dev, generator=gen, dtype=torch.bfloat16)
|
|
|
|
def tile(x):
|
|
buf = torch.zeros((1, Lpad, H, HD), device=dev, dtype=x.dtype)
|
|
buf[:, untile] = x
|
|
return buf.transpose(1, 2).contiguous() # BHSD
|
|
|
|
qt, kt, vt = tile(q), tile(k), tile(v)
|
|
scores = torch.matmul(_pool_tiles(qt.transpose(1, 2), vbs, tt), _pool_tiles(kt.transpose(1, 2), vbs, tt).transpose(-2, -1))
|
|
mask = _build_block_mask(scores, n_prefix, 0.8, True, ((n_prefix, n_prefix + n_video), ), (0.8, ))
|
|
q2k_idx, q2k_num, kvv, quad = fa.vsa_tile_mask_to_fp4_blocks(mask, tt, vbs, validate=True)
|
|
Lk = q2k_idx.shape[2] * 128
|
|
pad = lambda x: F.pad(x, (0, 0, 0, Lk - Lpad)).contiguous()
|
|
o = sageattn_blackwell_sparse(pad(qt), pad(kt), pad(vt), q2k_idx, q2k_num, kvv, quad)[:, :, :Lpad]
|
|
tok_tile = torch.arange(n_tiles, device=dev).repeat_interleave(tt)
|
|
tok_valid = torch.zeros(Lpad, dtype=torch.bool, device=dev)
|
|
tok_valid[untile] = True
|
|
tm = mask[:, :, tok_tile][:, :, :, tok_tile] & tok_valid[None, None, None, :]
|
|
ref = F.scaled_dot_product_attention(qt.float(), kt.float(), vt.float(), attn_mask=tm)
|
|
rows = untile
|
|
err = (o.float() - ref)[:, :, rows]
|
|
# FP4 noise floor: dense FP4 vs dense fp32 on the same valid tokens
|
|
qv, kv_, vv = (x[:, :, rows].contiguous() for x in (qt, kt, vt))
|
|
od = sageattn_blackwell(qv, kv_, vv)
|
|
rd = F.scaled_dot_product_attention(qv.float(), kv_.float(), vv.float())
|
|
out[name] = {
|
|
"tiles": n_tiles, "density": round(mask.float().mean().item(), 3),
|
|
"rel_l2": round((err.norm() / ref[:, :, rows].norm()).item(), 4),
|
|
"dense_fp4_rel_l2_floor": round(((od.float() - rd).norm() / rd.norm()).item(), 4),
|
|
"cosine": round(F.cosine_similarity(o[:, :, rows].float().flatten(), ref[:, :, rows].flatten(), dim=0).item(), 5),
|
|
"finite": bool(torch.isfinite(o[:, :, rows]).all()),
|
|
}
|
|
if tt == 64:
|
|
from fastvideo_kernel.block_sparse_attn import block_sparse_attn as bsa64
|
|
ot, _ = bsa64(qt, kt, vt, mask, vbs)
|
|
out[name]["triton_bf16_rel_l2"] = round(((ot.float() - ref)[:, :, rows].norm() / ref[:, :, rows].norm()).item(), 4)
|
|
print("CHECK", name, json.dumps(out[name]), flush=True)
|
|
return out
|
|
|
|
|
|
def density_study(prefix_segments=(256, 810), video_shape=(72, 24, 42)):
|
|
"""Block densities each kernel granularity would compute, and FP4 sparse timing at 0.8 / 0.9."""
|
|
from fastvideo.attention.backends.video_sparse_attn_h3 import _build_block_mask, _h3_tile_geometry, _pool_tiles
|
|
import attn_qat_infer.api as fa
|
|
|
|
dev = torch.device("cuda")
|
|
gen = torch.Generator(device=dev).manual_seed(0)
|
|
_, vbs, untile, n_prefix, n_video = _h3_tile_geometry(tuple(prefix_segments), tuple(video_shape), dev, (4, 4, 4))
|
|
n_tiles = vbs.numel()
|
|
Lpad = n_tiles * 64
|
|
q = smooth_tokens(sum(prefix_segments), tuple(video_shape), dev, gen)
|
|
k = (q + 0.5 * torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16)).to(torch.bfloat16)
|
|
v = torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16)
|
|
|
|
def tile(x):
|
|
buf = torch.zeros((1, Lpad, HEADS, HD), device=dev, dtype=x.dtype)
|
|
buf[:, untile] = x
|
|
return buf
|
|
|
|
qt, kt, vt = tile(q), tile(k), tile(v)
|
|
scores = torch.matmul(_pool_tiles(qt, vbs, 64), _pool_tiles(kt, vbs, 64).transpose(-2, -1))
|
|
nt2 = n_tiles + n_tiles % 2
|
|
rows = nt2 * 64
|
|
qs, ks, vs = (F.pad(x, (0, 0, 0, 0, 0, rows - Lpad)).contiguous() for x in (qt, kt, vt))
|
|
out = {}
|
|
for sparsity in (0.8, 0.9):
|
|
mask = _build_block_mask(scores, n_prefix, sparsity, True, ((n_prefix, n_prefix + n_video), ), (sparsity, ))
|
|
m = F.pad(mask, (0, nt2 - n_tiles, 0, nt2 - n_tiles), value=False)
|
|
B, H = m.shape[:2]
|
|
d = {"tile64x64": m.float().mean().item()}
|
|
q2 = m.view(B, H, nt2 // 2, 2, nt2) # query pairs
|
|
d["q128_k64"] = q2.any(3).float().mean().item()
|
|
k2 = m.view(B, H, nt2, nt2 // 2, 2)
|
|
d["q64_k128"] = k2.any(4).float().mean().item()
|
|
d["q128_k128"] = m.view(B, H, nt2 // 2, 2, nt2 // 2, 2).any(5).any(3).float().mean().item()
|
|
idx, num, kvv, quad = fa.vsa_tile_mask_to_fp4_blocks(mask, 64, vbs)
|
|
def bshd():
|
|
qh, kh, vh = (x.transpose(1, 2) for x in (qs, ks, vs))
|
|
ds = fa._zero_delta_s(1, HEADS, rows, dev)
|
|
return fa.blockscaled_fp4_attn_sparse(fa.scale_and_quant_fp4(qh), fa.scale_and_quant_fp4_permute(kh),
|
|
fa.scale_and_quant_fp4_transpose(vh), ds, rows, idx, num, kvv, quad,
|
|
False, True, True, None)[0]
|
|
|
|
d["sparse_fp4_quad_ms"] = timed(bshd, iters=3)
|
|
out[str(sparsity)] = {key: round(val, 4) for key, val in d.items()}
|
|
print("DENSITY", sparsity, json.dumps(out[str(sparsity)]), flush=True)
|
|
out["dense_fp4_bshd_ms"] = timed(lambda: fa.sageattn_blackwell(*(x.transpose(1, 2) for x in (qs, ks, vs))), iters=2)
|
|
return out
|