[feat]: shared FastH3 NVFP4 FP4-attention QAD recipe (OmniRef + V2 T2AV) (#1989)
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
co-authored by
Aryan Kumar
Aryan Kumar
parent
181725e3fd
commit
fcec554456
@@ -32,6 +32,7 @@ from __future__ import annotations
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
QUANT_GROUP = 16
|
||||
BLOCK = 128
|
||||
@@ -99,30 +100,73 @@ def key_valid_mask(kv_valid: torch.Tensor | None, n_blocks: int, device: torch.d
|
||||
return (col % HALF)[None, :] < limit
|
||||
|
||||
|
||||
def tile128_mask_to_fp4_blocks(tile_mask: torch.Tensor,
|
||||
tile_valid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""``attn_qat_infer.api.vsa_tile_mask_to_fp4_blocks(mask, 128, tile_valid)`` without the CUDA extensions.
|
||||
def tile_mask_to_fp4_blocks(tile_mask: torch.Tensor,
|
||||
tile_tokens: int,
|
||||
tile_valid: torch.Tensor | None = None) -> tuple[
|
||||
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
"""Pure-PyTorch ``attn_qat_infer.api.vsa_tile_mask_to_fp4_blocks`` for tile 64 or 128.
|
||||
|
||||
``attn_qat_infer`` imports its sm_12x extensions at module load, so data
|
||||
center GPUs cannot call it; this is the same computation for 128-token
|
||||
tiles (one tile per kernel block, so no quadrant masks and no padding
|
||||
block): descending lists and per-64-key-half valid counts. The fidelity
|
||||
check script asserts equality with the real function on sm_12x.
|
||||
The inference API imports sm_12x extensions at module load, so data center
|
||||
GPUs cannot call it. The kernel works on 128x128 blocks. With 64-token VSA
|
||||
tiles two query tiles and two key tiles are paired into one block; a block
|
||||
is listed when any of its four 64x64 quadrants is selected and ``q2k_quad``
|
||||
carries the selected quadrants (bit ``2 * query_half + key_half``). With
|
||||
128-token tiles the mapping is 1:1 and no quadrant mask is needed.
|
||||
|
||||
Returns ``(q2k_idx [B, H, M, N], q2k_num [B, H, M], kv_valid [2N],
|
||||
q2k_quad [B, H, M, N] | None)``; queries cover ``M * 128`` padded rows.
|
||||
"""
|
||||
n_tiles = tile_mask.shape[-1]
|
||||
half_valid = (tile_valid.to(torch.int32)[:, None] -
|
||||
torch.tensor([0, HALF], device=tile_mask.device, dtype=torch.int32)[None, :]).clamp(0, HALF)
|
||||
block_mask = tile_mask & (tile_valid > 0)[None, None, None, :]
|
||||
valid = half_valid.reshape(-1)
|
||||
n_blocks = n_tiles
|
||||
if tile_tokens % HALF or (tile_tokens > HALF and tile_tokens % BLOCK):
|
||||
raise ValueError(f"tile_tokens={tile_tokens} must be {HALF} or a multiple of {BLOCK}")
|
||||
batch, heads, n_tiles, _ = tile_mask.shape
|
||||
device = tile_mask.device
|
||||
halves_per_tile = tile_tokens // HALF
|
||||
half_mask = tile_mask
|
||||
if halves_per_tile > 1:
|
||||
half_mask = half_mask.repeat_interleave(halves_per_tile, dim=2).repeat_interleave(halves_per_tile, dim=3)
|
||||
n_halves = n_tiles * halves_per_tile
|
||||
if tile_valid is None:
|
||||
half_valid = torch.full((n_halves, ), HALF, device=device, dtype=torch.int32)
|
||||
else:
|
||||
offsets = torch.arange(halves_per_tile, device=device, dtype=torch.int32) * HALF
|
||||
half_valid = (tile_valid.to(torch.int32)[:, None] - offsets[None, :]).clamp(0, HALF).reshape(-1)
|
||||
if n_halves % 2:
|
||||
# Pad to whole 128-blocks: the extra key half is empty, the extra query
|
||||
# half (discarded output) attends block 0 only.
|
||||
half_mask = F.pad(half_mask, (0, 1, 0, 1), value=False)
|
||||
half_mask[:, :, -1, 0] = True
|
||||
half_valid = F.pad(half_valid, (0, 1), value=0)
|
||||
n_halves += 1
|
||||
half_mask = half_mask & (half_valid > 0)[None, None, None, :]
|
||||
n_blocks = n_halves // 2
|
||||
quads = half_mask.view(batch, heads, n_blocks, 2, n_blocks, 2)
|
||||
weights = torch.tensor([[1, 2], [4, 8]], device=device, dtype=torch.uint8) # [row_half, col_half]
|
||||
quad = (quads.to(torch.uint8) * weights[None, None, None, :, None, :]).sum(dim=(3, 5), dtype=torch.uint8)
|
||||
block_mask = quad != 0
|
||||
# Descending compaction without a sort: position = running count from the right.
|
||||
rev = block_mask.flip(-1)
|
||||
pos = rev.cumsum(-1, dtype=torch.int32) - 1
|
||||
q2k_num = (pos[..., -1] + 1).contiguous()
|
||||
cols = torch.arange(n_blocks - 1, -1, -1, device=tile_mask.device, dtype=torch.int32).expand_as(pos)
|
||||
cols = torch.arange(n_blocks - 1, -1, -1, device=device, dtype=torch.int32).expand_as(pos)
|
||||
slot = torch.where(rev, pos, torch.full_like(pos, n_blocks)).long()
|
||||
q2k_idx = torch.zeros((*tile_mask.shape[:2], n_blocks, n_blocks + 1), device=tile_mask.device, dtype=torch.int32)
|
||||
q2k_idx = torch.zeros((batch, heads, n_blocks, n_blocks + 1), device=device, dtype=torch.int32)
|
||||
q2k_idx.scatter_(-1, slot, cols)
|
||||
return q2k_idx[..., :n_blocks].contiguous(), q2k_num, valid.contiguous()
|
||||
q2k_idx = q2k_idx[..., :n_blocks].contiguous()
|
||||
q2k_quad = None
|
||||
if halves_per_tile == 1:
|
||||
q2k_quad = torch.zeros((batch, heads, n_blocks, n_blocks + 1), device=device, dtype=torch.uint8)
|
||||
q2k_quad.scatter_(-1, slot, quad.flip(-1))
|
||||
q2k_quad = q2k_quad[..., :n_blocks].contiguous()
|
||||
return q2k_idx, q2k_num, half_valid.contiguous(), q2k_quad
|
||||
|
||||
|
||||
def tile128_mask_to_fp4_blocks(tile_mask: torch.Tensor,
|
||||
tile_valid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Backward-compatible 128-token-tile wrapper."""
|
||||
idx, num, valid, quad = tile_mask_to_fp4_blocks(tile_mask, BLOCK, tile_valid)
|
||||
if quad is not None:
|
||||
raise AssertionError("128-token tiles must not produce a quadrant mask")
|
||||
return idx, num, valid
|
||||
|
||||
|
||||
def invert_block_lists(q2k_idx: torch.Tensor, q2k_num: torch.Tensor,
|
||||
@@ -162,6 +206,7 @@ def fp4_vsa_attention_reference(q: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
kv_valid: torch.Tensor | None = None,
|
||||
q2k_quad: torch.Tensor | None = None,
|
||||
*,
|
||||
two_level_p: bool = False,
|
||||
smooth_k: bool = False,
|
||||
@@ -199,17 +244,31 @@ def fp4_vsa_attention_reference(q: torch.Tensor,
|
||||
row_sum = torch.zeros(BLOCK, device=q.device)
|
||||
acc = torch.zeros(BLOCK, dim, device=q.device)
|
||||
visited: list[tuple[int, torch.Tensor, torch.Tensor]] = []
|
||||
for position, n in enumerate(visit_order(q2k_idx, q2k_num, b, h, m)):
|
||||
count = int(q2k_num[b, h, m])
|
||||
entries = list(range(count))[::-1]
|
||||
for position, entry in enumerate(entries):
|
||||
n = int(q2k_idx[b, h, m, entry])
|
||||
cols = slice(n * BLOCK, (n + 1) * BLOCK)
|
||||
scores = q_hat[b, h, rows] @ k_hat[b, h, cols].T
|
||||
scores = scores.masked_fill(~valid[n][None, :], -math.inf)
|
||||
if q2k_quad is not None:
|
||||
quad = int(q2k_quad[b, h, m, entry])
|
||||
half = torch.arange(BLOCK, device=q.device) // HALF
|
||||
qvalid = ((quad >> (half[:, None] * 2 + half[None, :])) & 1).bool()
|
||||
scores = scores.masked_fill(~qvalid, -math.inf)
|
||||
group_max = scores.unflatten(-1, (-1, QUANT_GROUP)).amax(dim=-1)
|
||||
if position == 0 and first_block_max_floor is not None:
|
||||
group_max = group_max.clamp(min=first_block_max_floor)
|
||||
new_max = torch.maximum(row_max, group_max.amax(dim=-1))
|
||||
m_scaled = new_max * scale_log2 + offset
|
||||
p_tilde = torch.exp2(scores * scale_log2 - m_scaled[:, None])
|
||||
alpha = torch.exp2((row_max - new_max) * scale_log2)
|
||||
# A listed block can have no selected keys for this query
|
||||
# half; both maxima are -inf and the usual exp2 terms are
|
||||
# NaN, so emit zeros and keep the running state unchanged.
|
||||
empty = torch.isinf(new_max) & (new_max < 0)
|
||||
m_scaled = torch.where(empty, torch.zeros_like(new_max), new_max * scale_log2 + offset)
|
||||
p_tilde = torch.where(empty[:, None], torch.zeros_like(scores),
|
||||
torch.exp2(scores * scale_log2 - m_scaled[:, None]))
|
||||
alpha = torch.where(empty, torch.zeros_like(new_max),
|
||||
torch.exp2((row_max - new_max) * scale_log2))
|
||||
alpha = torch.where(torch.isinf(row_max), torch.zeros_like(alpha), alpha)
|
||||
row_sum = row_sum * alpha + p_tilde.sum(dim=-1)
|
||||
if quantize:
|
||||
@@ -226,7 +285,9 @@ def fp4_vsa_attention_reference(q: torch.Tensor,
|
||||
out[b, h, rows] = acc * inv[:, None]
|
||||
if weights is not None:
|
||||
for n, p_hat, block_max in visited:
|
||||
decay = torch.exp2((block_max - row_max) * scale_log2) * inv
|
||||
finite = torch.isfinite(block_max) & torch.isfinite(row_max)
|
||||
decay = torch.where(finite, torch.exp2((block_max - row_max) * scale_log2) * inv,
|
||||
torch.zeros_like(inv))
|
||||
weights[b, h, rows, n * BLOCK:(n + 1) * BLOCK] = p_hat * decay[:, None]
|
||||
if weights is not None:
|
||||
return out, weights
|
||||
|
||||
@@ -141,9 +141,9 @@ def _key_ok(KV_VALID, n_block, cols, HAS_VALID: tl.constexpr, HALF: tl.constexpr
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, BIAS, QGROUP, n_groups, seq_len, n_qblocks,
|
||||
list_width, scale_log2, OFFSET: tl.constexpr, QUANT: tl.constexpr, HAS_VALID: tl.constexpr,
|
||||
HAS_BIAS: tl.constexpr, FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, HIGH_PREC_O: tl.constexpr,
|
||||
def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_QUAD, BIAS, QGROUP, n_groups, seq_len,
|
||||
n_qblocks, list_width, scale_log2, OFFSET: tl.constexpr, QUANT: tl.constexpr, HAS_VALID: tl.constexpr,
|
||||
HAS_QUAD: tl.constexpr, HAS_BIAS: tl.constexpr, FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, HIGH_PREC_O: tl.constexpr,
|
||||
SAVE_MRUN: tl.constexpr,
|
||||
D: tl.constexpr, ROWS: tl.constexpr, BLOCK_N: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
@@ -174,18 +174,28 @@ def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, BIAS, QG
|
||||
if HAS_BIAS: # per-key score correction (Q smoothing: mu_q . k_j, high precision)
|
||||
s = s + tl.load(bias_row + keys)[None, :]
|
||||
ok = _key_ok(KV_VALID, n_block, offs_n, HAS_VALID, BLOCK_N // 2)
|
||||
if HAS_QUAD:
|
||||
quad = tl.load(Q2K_QUAD + meta * list_width + slot).to(tl.int32)
|
||||
qh = pid % halves
|
||||
bit = tl.where(offs_n < ROWS, (quad >> (qh * 2)) & 1, (quad >> (qh * 2 + 1)) & 1)
|
||||
ok = ok & (bit != 0)
|
||||
s = tl.where(ok[None, :], s, float("-inf"))
|
||||
group_max = tl.max(tl.reshape(s, (ROWS, BLOCK_N // 16, 16)), axis=2)
|
||||
if FLOOR:
|
||||
group_max = tl.where(it == 0, tl.maximum(group_max, FLOOR_VALUE), group_max)
|
||||
m_new = tl.maximum(m_i, tl.max(group_max, axis=1))
|
||||
# Tile 64 pairs two query halves per block, so a listed block can have no selected keys for this half:
|
||||
# both maxima are -inf and exp2(-inf - -inf) is NaN. Keep the running state and emit zero probabilities.
|
||||
empty = m_new == float("-inf")
|
||||
m_scaled = m_new * scale_log2 + OFFSET
|
||||
alpha = tl.math.exp2((m_i - m_new) * scale_log2)
|
||||
alpha = tl.where(empty, 0.0, tl.math.exp2((m_i - m_new) * scale_log2))
|
||||
if QUANT:
|
||||
p_tilde, p_hat = _quantized_p(s, m_scaled, it == 0, scale_log2, FLOOR, FLOOR_VALUE, ROWS, BLOCK_N)
|
||||
else:
|
||||
p_tilde = tl.math.exp2(s * scale_log2 - m_scaled[:, None])
|
||||
p_hat = p_tilde
|
||||
p_tilde = tl.where(empty[:, None], 0.0, p_tilde)
|
||||
p_hat = tl.where(empty[:, None], 0.0, p_hat)
|
||||
l_i = l_i * alpha + tl.sum(p_tilde, axis=1)
|
||||
v = tl.load(V + base + keys[:, None] * D + offs_d[None, :])
|
||||
acc = acc * alpha[:, None] + tl.dot(p_hat.to(tl.bfloat16), v)
|
||||
@@ -202,9 +212,9 @@ def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, BIAS, QG
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, seq_len, n_qblocks, list_width,
|
||||
scale_log2, sm_scale, HAS_VALID: tl.constexpr, D: tl.constexpr, ROWS: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr):
|
||||
def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_QUAD, seq_len, n_qblocks,
|
||||
list_width, scale_log2, sm_scale, HAS_VALID: tl.constexpr, HAS_QUAD: tl.constexpr,
|
||||
D: tl.constexpr, ROWS: tl.constexpr, BLOCK_N: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
bh = tl.program_id(1).to(tl.int64)
|
||||
halves: tl.constexpr = BLOCK_N // ROWS
|
||||
@@ -229,7 +239,13 @@ def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, seq_
|
||||
v = tl.load(V + base + keys[:, None] * D + offs_d[None, :])
|
||||
s = tl.dot(q, tl.trans(k))
|
||||
ok = _key_ok(KV_VALID, n_block, cols, HAS_VALID, BLOCK_N // 2)
|
||||
p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse[:, None]), 0.0)
|
||||
if HAS_QUAD:
|
||||
quad = tl.load(Q2K_QUAD + meta * list_width + (it // halves)).to(tl.int32)
|
||||
bit = (quad >> ((pid % halves) * 2 + (it % halves))) & 1
|
||||
ok = ok & (bit != 0)
|
||||
# A row with no valid keys has lse = -inf; clamp it so it contributes zero instead of NaN.
|
||||
lse_safe = tl.where(lse == float("-inf"), 0.0, lse)
|
||||
p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse_safe[:, None]), 0.0)
|
||||
dp = tl.dot(do, tl.trans(v))
|
||||
ds = p * (dp - delta[:, None])
|
||||
dq += tl.dot(ds.to(tl.bfloat16), k)
|
||||
@@ -237,10 +253,11 @@ def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, seq_
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K2Q_NUM, Q2K_NUM, KV_VALID, seq_len,
|
||||
n_qblocks, n_kblocks, list_width, scale_log2, sm_scale, OFFSET: tl.constexpr,
|
||||
QUANT: tl.constexpr, HAS_VALID: tl.constexpr, FLOOR: tl.constexpr,
|
||||
FLOOR_VALUE: tl.constexpr, D: tl.constexpr, ROWS: tl.constexpr, BLOCK_N: tl.constexpr):
|
||||
def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K2Q_NUM, Q2K_NUM, KV_VALID,
|
||||
Q2K_QUAD, seq_len, n_qblocks, n_kblocks, list_width, scale_log2, sm_scale,
|
||||
OFFSET: tl.constexpr, QUANT: tl.constexpr, HAS_VALID: tl.constexpr, HAS_QUAD: tl.constexpr,
|
||||
FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, D: tl.constexpr, ROWS: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
bh = tl.program_id(1).to(tl.int64)
|
||||
halves: tl.constexpr = BLOCK_N // ROWS
|
||||
@@ -268,15 +285,22 @@ def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K
|
||||
lse = tl.load(LSE + bh * seq_len + rows)
|
||||
delta = tl.load(DELTA + bh * seq_len + rows)
|
||||
s = tl.dot(q, tl.trans(k))
|
||||
s = tl.where(ok[None, :], s, float("-inf"))
|
||||
p = tl.math.exp2(s * scale_log2 - lse[:, None])
|
||||
if HAS_QUAD:
|
||||
quad = tl.load(Q2K_QUAD + (bh * n_qblocks + m_block) * list_width + slot).to(tl.int32)
|
||||
bit = (quad >> ((it % halves) * 2 + (pid % halves))) & 1
|
||||
s = tl.where(ok[None, :] & (bit != 0), s, float("-inf"))
|
||||
else:
|
||||
s = tl.where(ok[None, :], s, float("-inf"))
|
||||
lse_safe = tl.where(lse == float("-inf"), 0.0, lse)
|
||||
p = tl.math.exp2(s * scale_log2 - lse_safe[:, None])
|
||||
if QUANT:
|
||||
qmeta = bh * n_qblocks + m_block
|
||||
first_visit = slot == tl.load(Q2K_NUM + qmeta) - 1
|
||||
m_run = tl.load(MRUN + (qmeta * list_width + slot) * BLOCK_N + row0 + offs_r)
|
||||
m_scaled = m_run * scale_log2 + OFFSET
|
||||
m_ok = m_run != float("-inf")
|
||||
m_scaled = tl.where(m_ok, m_run * scale_log2 + OFFSET, 0.0)
|
||||
_, p_hat = _quantized_p(s, m_scaled, first_visit, scale_log2, FLOOR, FLOOR_VALUE, ROWS, ROWS)
|
||||
weight = p_hat * tl.math.exp2(m_scaled - lse)[:, None]
|
||||
weight = tl.where(m_ok[:, None], p_hat * tl.math.exp2(m_scaled - lse_safe)[:, None], 0.0)
|
||||
else:
|
||||
weight = p
|
||||
dv += tl.dot(tl.trans(weight.to(tl.bfloat16)), do)
|
||||
@@ -290,8 +314,8 @@ def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K
|
||||
class _FP4VSAAttnQAT(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, q, k, v, q2k_idx, q2k_num, kv_valid, quantize, quantize_pv, two_level_p, first_block_max_floor,
|
||||
sm_scale, high_prec_o, key_bias=None, q_group=None, need_backward=False): # type: ignore[override]
|
||||
def forward(ctx, q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize, quantize_pv, two_level_p,
|
||||
first_block_max_floor, sm_scale, high_prec_o, key_bias=None, q_group=None, need_backward=False): # type: ignore[override]
|
||||
batch, heads, seq_len, dim = q.shape
|
||||
if seq_len % BLOCK or dim not in (64, 128):
|
||||
raise ValueError(f"need L % {BLOCK} == 0 and D in (64, 128), got L={seq_len}, D={dim}")
|
||||
@@ -306,6 +330,8 @@ class _FP4VSAAttnQAT(torch.autograd.Function):
|
||||
q2k_num = q2k_num.to(torch.int32).contiguous()
|
||||
has_valid = kv_valid is not None
|
||||
kv_valid_arg = kv_valid.to(torch.int32).contiguous() if has_valid else q2k_num
|
||||
has_quad = q2k_quad is not None
|
||||
q2k_quad_arg = q2k_quad.to(torch.uint8).contiguous() if has_quad else q2k_idx
|
||||
width = q2k_idx.shape[-1]
|
||||
out = torch.empty_like(q_hat)
|
||||
out_hp = torch.empty_like(q_hat) if high_prec_o else out
|
||||
@@ -326,15 +352,17 @@ class _FP4VSAAttnQAT(torch.autograd.Function):
|
||||
q_group = (q_group if q_group is not None else torch.zeros(n_blocks, dtype=torch.int32,
|
||||
device=q.device)).to(torch.int32).contiguous()
|
||||
grid = (n_blocks * (BLOCK // _ROWS), batch * heads)
|
||||
_fwd_kernel[grid](q_hat, k_hat, v_hat, out, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg,
|
||||
_fwd_kernel[grid](q_hat, k_hat, v_hat, out, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg, q2k_quad_arg,
|
||||
key_bias if has_bias else lse, q_group if has_bias else q2k_num,
|
||||
key_bias.shape[2] if has_bias else 1, seq_len, n_blocks, width, scale_log2, OFFSET=offset,
|
||||
QUANT=bool(quantize), HAS_VALID=has_valid, HAS_BIAS=has_bias, FLOOR=floor, FLOOR_VALUE=float(first_block_max_floor or 0.0),
|
||||
QUANT=bool(quantize), HAS_VALID=has_valid, HAS_QUAD=has_quad, HAS_BIAS=has_bias, FLOOR=floor, FLOOR_VALUE=float(first_block_max_floor or 0.0),
|
||||
HIGH_PREC_O=bool(high_prec_o), SAVE_MRUN=save_mrun, D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, num_warps=_WARPS,
|
||||
num_stages=2)
|
||||
if need_backward:
|
||||
ctx.save_for_backward(q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg)
|
||||
ctx.options = (bool(quantize), has_valid, floor, float(first_block_max_floor or 0.0), offset, sm_scale)
|
||||
ctx.save_for_backward(q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg,
|
||||
q2k_quad_arg)
|
||||
ctx.options = (bool(quantize), has_valid, has_quad, floor, float(first_block_max_floor or 0.0), offset,
|
||||
sm_scale)
|
||||
ctx.has_bias = has_bias
|
||||
return out
|
||||
|
||||
@@ -342,8 +370,8 @@ class _FP4VSAAttnQAT(torch.autograd.Function):
|
||||
def backward(ctx, grad_out): # type: ignore[override]
|
||||
if ctx.has_bias:
|
||||
raise NotImplementedError("the Q-smoothing score bias is an eval-only emulator option (no backward)")
|
||||
q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid = ctx.saved_tensors
|
||||
quantize, has_valid, floor, floor_value, offset, sm_scale = ctx.options
|
||||
q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid, q2k_quad = ctx.saved_tensors
|
||||
quantize, has_valid, has_quad, floor, floor_value, offset, sm_scale = ctx.options
|
||||
batch, heads, seq_len, dim = q_hat.shape
|
||||
n_blocks = seq_len // BLOCK
|
||||
width = q2k_idx.shape[-1]
|
||||
@@ -355,15 +383,15 @@ class _FP4VSAAttnQAT(torch.autograd.Function):
|
||||
dv = torch.empty_like(v_hat)
|
||||
scale_log2 = sm_scale * LOG2E
|
||||
grid = (n_blocks * (BLOCK // _ROWS), batch * heads)
|
||||
_bwd_dq_kernel[grid](q_hat, k_hat, v_hat, grad_out, lse, delta, dq, q2k_idx, q2k_num, kv_valid, seq_len,
|
||||
n_blocks, width, scale_log2, sm_scale, HAS_VALID=has_valid, D=dim, ROWS=_ROWS,
|
||||
BLOCK_N=BLOCK, num_warps=_WARPS, num_stages=2)
|
||||
_bwd_dq_kernel[grid](q_hat, k_hat, v_hat, grad_out, lse, delta, dq, q2k_idx, q2k_num, kv_valid, q2k_quad,
|
||||
seq_len, n_blocks, width, scale_log2, sm_scale, HAS_VALID=has_valid,
|
||||
HAS_QUAD=has_quad, D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, num_warps=_WARPS, num_stages=2)
|
||||
_bwd_dkdv_kernel[grid](q_hat, k_hat, v_hat, grad_out, lse, delta, m_run, dk, dv, k2q_idx, k2q_slot, k2q_num,
|
||||
q2k_num, kv_valid, seq_len, n_blocks, n_blocks, width, scale_log2, sm_scale,
|
||||
OFFSET=offset, QUANT=quantize, HAS_VALID=has_valid, FLOOR=floor,
|
||||
q2k_num, kv_valid, q2k_quad, seq_len, n_blocks, n_blocks, width, scale_log2, sm_scale,
|
||||
OFFSET=offset, QUANT=quantize, HAS_VALID=has_valid, HAS_QUAD=has_quad, FLOOR=floor,
|
||||
FLOOR_VALUE=floor_value, D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, num_warps=_WARPS,
|
||||
num_stages=2)
|
||||
return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None
|
||||
return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None
|
||||
|
||||
|
||||
def _hadamard(dim: int, device: torch.device) -> torch.Tensor:
|
||||
@@ -426,6 +454,7 @@ def fp4_vsa_attn_qat(q: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
kv_valid: torch.Tensor | None = None,
|
||||
q2k_quad: torch.Tensor | None = None,
|
||||
*,
|
||||
quantize: bool = True,
|
||||
quantize_pv: bool = True,
|
||||
@@ -442,7 +471,8 @@ def fp4_vsa_attn_qat(q: torch.Tensor,
|
||||
|
||||
``q``/``k``/``v`` are ``[B, H, L, D]`` (L a multiple of 128, tile-ordered,
|
||||
pad rows zero); the lists are ``vsa_tile_mask_to_fp4_blocks(mask, 128, vbs)``
|
||||
outputs (``q2k_quad`` must be None: 128-token tiles map 1:1 onto blocks).
|
||||
outputs (``q2k_quad`` is None for 128-token tiles, which map 1:1 onto
|
||||
blocks, and carries the four 64x64 quadrant bits for 64-token tiles).
|
||||
``smooth_k`` subtracts each head's token mean from K first (exact for
|
||||
softmax, differentiable through autograd). ``quantize_pv=False`` quantizes
|
||||
Q/K only and keeps P and V in BF16 (the FA4 NVFP4-QK + BF16-PV kernel).
|
||||
@@ -483,7 +513,7 @@ def fp4_vsa_attn_qat(q: torch.Tensor,
|
||||
# Capture the caller's grad mode here, including no_grad on grad-enabled inputs.
|
||||
need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q, k, v))
|
||||
if not (hadamard or smooth_q or prefix_precision):
|
||||
return _FP4VSAAttnQAT.apply(q, k, v, q2k_idx, q2k_num, kv_valid, quantize, quantize_pv, two_level_p,
|
||||
return _FP4VSAAttnQAT.apply(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize, quantize_pv, two_level_p,
|
||||
first_block_max_floor, scale, high_prec_o, None, None, need_backward)
|
||||
n_blocks = q.shape[2] // BLOCK
|
||||
hp_tiles = min(int(num_prefix_tiles or 0), n_blocks) if prefix_precision else 0
|
||||
@@ -500,14 +530,14 @@ def fp4_vsa_attn_qat(q: torch.Tensor,
|
||||
n_groups = 2 if smooth_q == "segment" else 1
|
||||
groups = _query_groups(n_blocks, num_prefix_tiles, smooth_q == "segment", q.device)
|
||||
q_lo, bias = _smooth_queries(q_lo, k_lo, groups, n_groups)
|
||||
out = _FP4VSAAttnQAT.apply(q_lo, k_lo, v, q2k_idx, num_fp4, kv_valid, quantize, quantize_pv, two_level_p,
|
||||
out = _FP4VSAAttnQAT.apply(q_lo, k_lo, v, q2k_idx, num_fp4, kv_valid, q2k_quad, quantize, quantize_pv, two_level_p,
|
||||
first_block_max_floor, scale, high_prec_o, bias, groups, need_backward)
|
||||
if not hp_tiles:
|
||||
return out
|
||||
num_hp = q2k_num - num_fp4
|
||||
hp_quant = {"bf16": lambda t: t, "fp8": fp8_per_token_fake_quant, "mxfp8": mxfp8_fake_quant}[prefix_precision]
|
||||
q_hp, k_hp = hp_quant(q), hp_quant(k)
|
||||
out_hp = _FP4VSAAttnQAT.apply(q_hp, k_hp, v, q2k_idx, num_hp, kv_valid, False, True, False, None, scale,
|
||||
out_hp = _FP4VSAAttnQAT.apply(q_hp, k_hp, v, q2k_idx, num_hp, kv_valid, q2k_quad, False, True, False, None, scale,
|
||||
high_prec_o, None, None, need_backward)
|
||||
return out + out_hp
|
||||
|
||||
|
||||
Reference in New Issue
Block a user