1554 lines
57 KiB
Python
1554 lines
57 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from https://github.com/triton-lang/triton/blob/main/python/tutorials/06-fused-attention.py
|
|
|
|
import os
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
from triton.tools.tensor_descriptor import TensorDescriptor
|
|
|
|
from .quant_utils import fake_quantize_q, fake_quantize_kv, fake_quantize
|
|
# from attn_qat_infer.api import triton_group_mean
|
|
|
|
|
|
def is_cuda():
|
|
return triton.runtime.driver.active.get_current_target().backend == "cuda"
|
|
|
|
|
|
def supports_host_descriptor():
|
|
return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
|
|
|
|
|
|
def is_sm100(device=None):
|
|
return is_cuda() and torch.cuda.get_device_capability(device) == (10, 0)
|
|
|
|
|
|
def is_blackwell():
|
|
return is_cuda() and torch.cuda.get_device_capability()[0] == 10
|
|
|
|
|
|
def is_consumer_blackwell():
|
|
return is_cuda() and torch.cuda.get_device_capability()[0] == 12
|
|
|
|
|
|
def is_hopper():
|
|
return is_cuda() and torch.cuda.get_device_capability()[0] == 9
|
|
|
|
|
|
def _sm100_optimization_enabled():
|
|
return os.environ.get("FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED", "1") != "0"
|
|
|
|
|
|
def _sm100_exact_m_enabled():
|
|
return os.environ.get("FASTVIDEO_ATTN_QAT_FWD_EXACT_M", "0") != "0"
|
|
|
|
|
|
def _consumer_blackwell_join_qat_pv_enabled():
|
|
"""Return whether SM120 uses the joined quantized/STE P@V path."""
|
|
return os.environ.get("FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV", "1") != "0"
|
|
|
|
|
|
def _use_sm100_optimized_qat(
|
|
device,
|
|
head_dim: int,
|
|
causal: bool,
|
|
is_qat: bool,
|
|
fake_quant_p: bool,
|
|
two_level_quant_p: bool,
|
|
use_global_sf_p: bool,
|
|
) -> bool:
|
|
"""Return whether this call matches the validated SM100 fast path."""
|
|
return (_sm100_optimization_enabled() and is_sm100(device) and head_dim == 128 and not causal and is_qat
|
|
and fake_quant_p and not two_level_quant_p and not use_global_sf_p)
|
|
|
|
|
|
def _select_sm100_forward_config(n_ctx_q: int, n_ctx_kv: int, mode: str):
|
|
n_ctx = max(n_ctx_q, n_ctx_kv)
|
|
if mode == "reference":
|
|
return 32, 32, 4, 4 if n_ctx >= 16_384 else 5
|
|
if n_ctx <= 2_048:
|
|
return 32, 32, 4, 5
|
|
if mode == "balanced":
|
|
return 64, 32, 4, 4
|
|
return 128, 128, 8, 3
|
|
|
|
|
|
@triton.jit
|
|
def _mul_alpha(acc, alpha, BM: tl.constexpr, BN: tl.constexpr):
|
|
acc0, acc1 = acc.reshape([BM, 2, BN // 2]).permute(0, 2, 1).split()
|
|
acc0 = acc0 * alpha[:, None]
|
|
acc1 = acc1 * alpha[:, None]
|
|
acc = tl.join(acc0, acc1).permute(0, 2, 1).reshape([BM, BN])
|
|
return acc
|
|
|
|
|
|
@triton.jit
|
|
def _attn_fwd_inner(acc,
|
|
high_prec_acc,
|
|
l_i,
|
|
m_i,
|
|
q,
|
|
q_valid,
|
|
desc_k,
|
|
desc_v,
|
|
offset_y,
|
|
dtype: tl.constexpr,
|
|
start_m,
|
|
qk_scale,
|
|
BLOCK_M: tl.constexpr,
|
|
HEAD_DIM: tl.constexpr,
|
|
BLOCK_N: tl.constexpr,
|
|
STAGE: tl.constexpr,
|
|
offs_m: tl.constexpr,
|
|
offs_n: tl.constexpr,
|
|
N_CTX: tl.constexpr,
|
|
warp_specialize: tl.constexpr,
|
|
IS_HOPPER: tl.constexpr,
|
|
IS_QAT: tl.constexpr,
|
|
fake_quant_P: tl.constexpr = True,
|
|
two_level_quant_P: tl.constexpr = False,
|
|
use_global_sf_P: tl.constexpr = True,
|
|
JOIN_QAT_PV: tl.constexpr = False):
|
|
# range of values handled by this stage (kv blocks)
|
|
if STAGE == 1:
|
|
lo, hi = 0, start_m * BLOCK_M
|
|
elif STAGE == 2:
|
|
lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
|
|
lo = tl.multiple_of(lo, BLOCK_M)
|
|
# causal = False
|
|
else:
|
|
lo, hi = 0, N_CTX
|
|
offsetk_y = offset_y + lo # offset from the start of the current batch-head combination
|
|
if dtype == tl.float8e5:
|
|
offsetv_y = offset_y * HEAD_DIM + lo
|
|
else:
|
|
offsetv_y = offset_y + lo
|
|
# loop over k, v and update accumulator
|
|
for start_n in tl.range(lo, hi, BLOCK_N, warp_specialize=warp_specialize):
|
|
start_n = tl.multiple_of(start_n, BLOCK_N)
|
|
kv_valid = (start_n + offs_n) < N_CTX
|
|
# -- compute qk ----
|
|
k = desc_k.load([offsetk_y, 0])
|
|
k = tl.where(kv_valid[:, None], k, 0.0)
|
|
qk = tl.dot(q, tl.trans(k))
|
|
if STAGE == 2:
|
|
mask = offs_m[:, None] >= (start_n + offs_n[None, :])
|
|
qk = qk * qk_scale + tl.where(mask & kv_valid[None, :], 0, -1.0e6)
|
|
m_ij = tl.maximum(m_i, tl.max(qk, 1))
|
|
qk -= m_ij[:, None]
|
|
else:
|
|
qk = tl.where(kv_valid[None, :], qk, -1.0e6)
|
|
m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
|
|
qk = qk * qk_scale - m_ij[:, None]
|
|
p = tl.math.exp2(qk)
|
|
if IS_QAT:
|
|
p, high_prec_p = fake_quantize(src_tensor=p,
|
|
valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0,
|
|
BLOCK_SIZE_OUT_DIM=BLOCK_M,
|
|
BLOCK_SIZE_QUANT_DIM=BLOCK_N,
|
|
dst_dtype=dtype,
|
|
two_level_quant_P=two_level_quant_P,
|
|
use_global_sf=use_global_sf_P)
|
|
l_ij = tl.sum(high_prec_p, 1)
|
|
else:
|
|
l_ij = tl.sum(p, 1)
|
|
# -- compute correction factor
|
|
alpha = tl.math.exp2(m_i - m_ij)
|
|
# -- update output accumulator --
|
|
if not IS_HOPPER and warp_specialize and BLOCK_M == 128 and HEAD_DIM == 128:
|
|
BM: tl.constexpr = acc.shape[0]
|
|
BN: tl.constexpr = acc.shape[1]
|
|
acc = _mul_alpha(acc, alpha, BM, BN)
|
|
if IS_QAT:
|
|
high_prec_acc = _mul_alpha(high_prec_acc, alpha, BM, BN)
|
|
else:
|
|
acc = acc * alpha[:, None]
|
|
if IS_QAT:
|
|
high_prec_acc = high_prec_acc * alpha[:, None]
|
|
# prepare p and v for the dot
|
|
if dtype == tl.float8e5:
|
|
v = desc_v.load([0, offsetv_y]).T
|
|
else:
|
|
v = desc_v.load([offsetv_y, 0])
|
|
v = tl.where(kv_valid[:, None], v, 0.0)
|
|
p = p.to(dtype)
|
|
# Keep the quantized and STE paths in one tensor-core operation. They
|
|
# share V, so joining along M avoids issuing two small, independent
|
|
# dot operations for every KV tile.
|
|
if IS_QAT and JOIN_QAT_PV:
|
|
joined_p = tl.join(p, high_prec_p).permute(2, 0, 1).reshape([2 * BLOCK_M, BLOCK_N])
|
|
joined_acc = tl.join(acc, high_prec_acc).permute(2, 0, 1).reshape([2 * BLOCK_M, HEAD_DIM])
|
|
joined_acc = tl.dot(joined_p, v.to(dtype), joined_acc)
|
|
acc, high_prec_acc = joined_acc.reshape([2, BLOCK_M, HEAD_DIM]).permute(1, 2, 0).split()
|
|
else:
|
|
# note that this non transposed v for FP8 is only supported on Blackwell
|
|
acc = tl.dot(p, v.to(dtype), acc)
|
|
if IS_QAT:
|
|
high_prec_acc = tl.dot(high_prec_p, v, high_prec_acc)
|
|
# update m_i and l_i
|
|
# place this at the end of the loop to reduce register pressure
|
|
l_i = l_i * alpha + l_ij
|
|
m_i = m_ij
|
|
offsetk_y += BLOCK_N
|
|
offsetv_y += BLOCK_N
|
|
if IS_QAT:
|
|
return acc, high_prec_acc, l_i, m_i
|
|
else:
|
|
return acc, acc, l_i, m_i
|
|
|
|
|
|
def _host_descriptor_pre_hook(nargs):
|
|
BLOCK_M = nargs["BLOCK_M"]
|
|
BLOCK_N = nargs["BLOCK_N"]
|
|
HEAD_DIM = nargs["HEAD_DIM"]
|
|
if not isinstance(nargs["desc_q"], TensorDescriptor):
|
|
return
|
|
nargs["desc_q"].block_shape = [BLOCK_M, HEAD_DIM]
|
|
if nargs["FP8_OUTPUT"]:
|
|
nargs["desc_v"].block_shape = [HEAD_DIM, BLOCK_N]
|
|
else:
|
|
nargs["desc_v"].block_shape = [BLOCK_N, HEAD_DIM]
|
|
nargs["desc_k"].block_shape = [BLOCK_N, HEAD_DIM]
|
|
nargs["desc_o"].block_shape = [1, BLOCK_M, HEAD_DIM]
|
|
if "desc_high_prec_o" in nargs and isinstance(nargs["desc_high_prec_o"], TensorDescriptor):
|
|
nargs["desc_high_prec_o"].block_shape = [1, BLOCK_M, HEAD_DIM]
|
|
|
|
|
|
NUM_STAGES_OPTIONS = [2, 3, 4]
|
|
|
|
configs = [
|
|
triton.Config({
|
|
'BLOCK_M': BM,
|
|
'BLOCK_N': BN
|
|
}, num_stages=s, num_warps=w, pre_hook=_host_descriptor_pre_hook) for BM in [64, 128] for BN in [32, 64, 128]
|
|
for s in NUM_STAGES_OPTIONS for w in [4, 8]
|
|
]
|
|
if "PYTEST_VERSION" in os.environ:
|
|
# Use a single config in testing for reproducibility
|
|
configs = [
|
|
triton.Config(dict(BLOCK_M=128, BLOCK_N=64), num_stages=2, num_warps=4, pre_hook=_host_descriptor_pre_hook),
|
|
]
|
|
|
|
|
|
def keep(conf):
|
|
BLOCK_M = conf.kwargs["BLOCK_M"]
|
|
BLOCK_N = conf.kwargs["BLOCK_N"]
|
|
return not (is_cuda() and torch.cuda.get_device_capability()[0] == 9 and BLOCK_M * BLOCK_N < 128 * 128
|
|
and conf.num_warps == 8)
|
|
|
|
|
|
def prune_invalid_configs(configs, named_args, **kwargs):
|
|
N_CTX_Q = kwargs["N_CTX_Q"]
|
|
N_CTX_KV = kwargs["N_CTX_KV"]
|
|
return [
|
|
conf for conf in configs
|
|
if conf.kwargs.get("BLOCK_M", 0) <= N_CTX_Q and conf.kwargs.get("BLOCK_N", 0) <= N_CTX_KV
|
|
and conf.kwargs.get("BLOCK_N", 0) % conf.kwargs.get("BLOCK_M", 0) == 0
|
|
]
|
|
|
|
|
|
@triton.jit
|
|
def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape):
|
|
if isinstance(desc_or_ptr, tl.tensor_descriptor):
|
|
return desc_or_ptr
|
|
else:
|
|
return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape)
|
|
|
|
|
|
# @triton.autotune(configs=list(filter(keep, configs)), key=["N_CTX_Q", "N_CTX_KV", "HEAD_DIM", "FP8_OUTPUT", "warp_specialize"],
|
|
# prune_configs_by={'early_config_prune': prune_invalid_configs}, cache_results=True)
|
|
@triton.jit
|
|
def _attn_fwd(
|
|
sm_scale,
|
|
M,
|
|
Z,
|
|
H,
|
|
desc_q,
|
|
desc_k,
|
|
desc_v,
|
|
desc_o,
|
|
desc_high_prec_o,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
HEAD_DIM: tl.constexpr,
|
|
BLOCK_M: tl.constexpr,
|
|
BLOCK_N: tl.constexpr,
|
|
FP8_OUTPUT: tl.constexpr,
|
|
STAGE: tl.constexpr,
|
|
warp_specialize: tl.constexpr,
|
|
IS_HOPPER: tl.constexpr,
|
|
IS_QAT: tl.constexpr,
|
|
fake_quant_P: tl.constexpr = True,
|
|
two_level_quant_P: tl.constexpr = False,
|
|
use_global_sf_P: tl.constexpr = True,
|
|
JOIN_QAT_PV: tl.constexpr = False,
|
|
):
|
|
dtype = tl.float8e5 if FP8_OUTPUT else tl.bfloat16
|
|
tl.static_assert(BLOCK_N <= HEAD_DIM)
|
|
start_m = tl.program_id(0)
|
|
off_hz = tl.program_id(1)
|
|
off_z = off_hz // H # since it's Z then H in the shape
|
|
off_h = off_hz % H
|
|
|
|
y_dim_q = Z * H * N_CTX_Q
|
|
y_dim_kv = Z * H * N_CTX_KV
|
|
desc_q = _maybe_make_tensor_desc(desc_q,
|
|
shape=[y_dim_q, HEAD_DIM],
|
|
strides=[HEAD_DIM, 1],
|
|
block_shape=[BLOCK_M, HEAD_DIM])
|
|
if FP8_OUTPUT:
|
|
desc_v = _maybe_make_tensor_desc(desc_v,
|
|
shape=[HEAD_DIM, y_dim_kv],
|
|
strides=[N_CTX_KV, 1],
|
|
block_shape=[HEAD_DIM, BLOCK_N])
|
|
else:
|
|
desc_v = _maybe_make_tensor_desc(desc_v,
|
|
shape=[y_dim_kv, HEAD_DIM],
|
|
strides=[HEAD_DIM, 1],
|
|
block_shape=[BLOCK_N, HEAD_DIM])
|
|
desc_k = _maybe_make_tensor_desc(desc_k,
|
|
shape=[y_dim_kv, HEAD_DIM],
|
|
strides=[HEAD_DIM, 1],
|
|
block_shape=[BLOCK_N, HEAD_DIM])
|
|
desc_o = _maybe_make_tensor_desc(desc_o,
|
|
shape=[Z * H, N_CTX_Q, HEAD_DIM],
|
|
strides=[N_CTX_Q * HEAD_DIM, HEAD_DIM, 1],
|
|
block_shape=[1, BLOCK_M, HEAD_DIM])
|
|
if IS_QAT:
|
|
desc_high_prec_o = _maybe_make_tensor_desc(desc_high_prec_o,
|
|
shape=[Z * H, N_CTX_Q, HEAD_DIM],
|
|
strides=[N_CTX_Q * HEAD_DIM, HEAD_DIM, 1],
|
|
block_shape=[1, BLOCK_M, HEAD_DIM])
|
|
|
|
offset_y_q = off_z * (N_CTX_Q * H) + off_h * N_CTX_Q # offset for query tensor
|
|
offset_y_kv = off_z * (N_CTX_KV * H) + off_h * N_CTX_KV # offset for key/value tensors
|
|
qo_offset_y = offset_y_q + start_m * BLOCK_M
|
|
# initialize offsets
|
|
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
offs_n = tl.arange(0, BLOCK_N) # for kv blocks
|
|
q_valid = offs_m < N_CTX_Q
|
|
# initialize pointer to m and l
|
|
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
|
|
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
|
|
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
|
if IS_QAT:
|
|
high_prec_acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
|
else:
|
|
# dummy value
|
|
high_prec_acc = acc
|
|
# load scales
|
|
qk_scale = sm_scale
|
|
qk_scale *= 1.44269504 # 1/log(2)
|
|
# load q: it will stay in SRAM throughout
|
|
q = desc_q.load([qo_offset_y, 0]) # load from start of q block and start of the entire head dimension
|
|
q = tl.where(q_valid[:, None], q, 0.0)
|
|
# stage 1: off-band
|
|
# For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE
|
|
# For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE
|
|
if STAGE & 1:
|
|
acc, high_prec_acc, l_i, m_i = _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, desc_k, desc_v,
|
|
offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM,
|
|
BLOCK_N, 4 - STAGE, offs_m, offs_n, N_CTX_KV, warp_specialize,
|
|
IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P,
|
|
use_global_sf_P, JOIN_QAT_PV)
|
|
# stage 2: on-band
|
|
if STAGE & 2:
|
|
acc, high_prec_acc, l_i, m_i = _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, desc_k, desc_v,
|
|
offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM,
|
|
BLOCK_N, 2, offs_m, offs_n, N_CTX_KV, warp_specialize, IS_HOPPER,
|
|
IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P,
|
|
JOIN_QAT_PV)
|
|
# epilogue
|
|
m_i += tl.math.log2(l_i)
|
|
acc = acc / l_i[:, None]
|
|
if IS_QAT:
|
|
high_prec_acc = high_prec_acc / l_i[:, None]
|
|
m_ptrs = M + off_hz * N_CTX_Q + offs_m
|
|
tl.store(m_ptrs, m_i, mask=q_valid)
|
|
desc_o.store([off_hz, start_m * BLOCK_M, 0], acc[None, :, :])
|
|
if IS_QAT:
|
|
desc_high_prec_o.store([off_hz, start_m * BLOCK_M, 0], high_prec_acc[None, :, :])
|
|
|
|
|
|
@triton.jit
|
|
def _attn_fwd_exact_m(
|
|
desc_q,
|
|
desc_k,
|
|
M,
|
|
sm_scale,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
HEAD_DIM: tl.constexpr,
|
|
BLOCK_M: tl.constexpr,
|
|
BLOCK_N: tl.constexpr,
|
|
):
|
|
"""Reproduce the legacy 32x32 forward softmax statistic exactly."""
|
|
start_m = tl.program_id(0) * BLOCK_M
|
|
off_hz = tl.program_id(1)
|
|
offs_m = start_m + tl.arange(0, BLOCK_M)
|
|
offs_n = tl.arange(0, BLOCK_N)
|
|
q_valid = offs_m < N_CTX_Q
|
|
|
|
q_base = off_hz * N_CTX_Q
|
|
kv_base = off_hz * N_CTX_KV
|
|
q = desc_q.load([q_base + start_m, 0])
|
|
q = tl.where(q_valid[:, None], q, 0.0)
|
|
|
|
m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
|
|
l_i = tl.full([BLOCK_M], 1.0, tl.float32)
|
|
qk_scale = sm_scale * 1.44269504
|
|
|
|
for start_n in tl.range(0, N_CTX_KV, BLOCK_N):
|
|
start_n = tl.multiple_of(start_n, BLOCK_N)
|
|
kv_valid = start_n + offs_n < N_CTX_KV
|
|
k = desc_k.load([kv_base + start_n, 0])
|
|
k = tl.where(kv_valid[:, None], k, 0.0)
|
|
qk = tl.dot(q, tl.trans(k))
|
|
qk = tl.where(kv_valid[None, :], qk, -1.0e6)
|
|
m_ij = tl.maximum(m_i, tl.max(qk, axis=1) * qk_scale)
|
|
p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
|
|
l_ij = tl.sum(p.to(tl.bfloat16), axis=1)
|
|
alpha = tl.math.exp2(m_i - m_ij)
|
|
l_i = l_i * alpha + l_ij
|
|
m_i = m_ij
|
|
|
|
m_i += tl.math.log2(l_i)
|
|
tl.store(M + off_hz * N_CTX_Q + offs_m, m_i, mask=q_valid)
|
|
|
|
|
|
@triton.jit
|
|
def _attn_bwd_preprocess(O, DO, Delta, Z, H, N_CTX, BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr):
|
|
off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
off_hz = tl.program_id(1)
|
|
off_n = tl.arange(0, HEAD_DIM)
|
|
valid = off_m < N_CTX
|
|
# load
|
|
o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None])
|
|
do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :],
|
|
mask=valid[:, None]).to(tl.float32)
|
|
delta = tl.sum(o * do, axis=1)
|
|
# write-back
|
|
tl.store(Delta + off_hz * N_CTX + off_m, delta, mask=valid)
|
|
|
|
|
|
# The main inner-loop logic for computing dK and dV.
|
|
@triton.jit
|
|
def _attn_bwd_dkdv(
|
|
dk,
|
|
dv,
|
|
Q,
|
|
k,
|
|
v,
|
|
qk_scale,
|
|
DO,
|
|
M,
|
|
D,
|
|
Q_MEAN,
|
|
# shared by Q/K/V/DO.
|
|
stride_tok,
|
|
stride_d,
|
|
H,
|
|
N_CTX,
|
|
BLOCK_M1: tl.constexpr,
|
|
BLOCK_N1: tl.constexpr,
|
|
HEAD_DIM: tl.constexpr,
|
|
# Filled in by the wrapper.
|
|
start_n,
|
|
start_m,
|
|
num_steps,
|
|
MASK: tl.constexpr,
|
|
IS_QAT: tl.constexpr,
|
|
two_level_quant_P: tl.constexpr = False,
|
|
fake_quant_P: tl.constexpr = True,
|
|
SMOOTH_Q: tl.constexpr = False,
|
|
use_global_sf_P: tl.constexpr = True,
|
|
warp_specialize: tl.constexpr = False):
|
|
offs_m = start_m + tl.arange(0, BLOCK_M1)
|
|
offs_n = start_n + tl.arange(0, BLOCK_N1)
|
|
offs_k = tl.arange(0, HEAD_DIM)
|
|
q_ptrs = Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
if SMOOTH_Q:
|
|
q_m_ptrs = Q_MEAN + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
# BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
|
|
tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
|
|
curr_m = start_m
|
|
step_m = BLOCK_M1
|
|
for blk_idx in range(num_steps):
|
|
offs_m = curr_m + tl.arange(0, BLOCK_M1)
|
|
q_valid = offs_m < N_CTX
|
|
q = tl.load(q_ptrs, mask=q_valid[:, None])
|
|
if SMOOTH_Q:
|
|
q_m = tl.load(q_m_ptrs, mask=q_valid[:, None])
|
|
# Load m before computing qk to reduce pipeline stall.
|
|
m = tl.load(M + offs_m, mask=q_valid)
|
|
qk = tl.dot(q, tl.trans(k))
|
|
qk = qk * qk_scale
|
|
# Autoregressive masking - apply BEFORE exp2 to match forward pass behavior
|
|
if MASK:
|
|
mask = (offs_m[:, None] >= offs_n[None, :])
|
|
qk = tl.where(mask, qk, -1.0e6)
|
|
p = tl.math.exp2(qk - m[:, None])
|
|
do = tl.load(do_ptrs, mask=q_valid[:, None])
|
|
# Compute dV.
|
|
p_quant = p
|
|
if IS_QAT and fake_quant_P:
|
|
p_quant, _ = fake_quantize(src_tensor=p,
|
|
valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0,
|
|
BLOCK_SIZE_OUT_DIM=BLOCK_M1,
|
|
BLOCK_SIZE_QUANT_DIM=BLOCK_N1,
|
|
dst_dtype=p.dtype,
|
|
two_level_quant_P=two_level_quant_P,
|
|
use_global_sf=use_global_sf_P)
|
|
dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do)
|
|
# D (= delta) is pre-divided by ds_scale.
|
|
Di = tl.load(D + offs_m, mask=q_valid)
|
|
# Compute dP and dS.
|
|
dp = tl.dot(do, tl.trans(v))
|
|
ds = p * (dp - Di[:, None])
|
|
ds = ds.to(tl.bfloat16)
|
|
dk += tl.dot(tl.trans(ds), q)
|
|
if SMOOTH_Q:
|
|
dk += tl.sum(ds, axis=1, keep_dims=True) * q_m
|
|
# Increment pointers.
|
|
curr_m += step_m
|
|
q_ptrs += step_m * stride_tok
|
|
do_ptrs += step_m * stride_tok
|
|
if SMOOTH_Q:
|
|
q_m_ptrs += step_m * stride_tok
|
|
return dk, dv
|
|
|
|
|
|
# The main inner-loop logic for computing dQ
|
|
@triton.jit
|
|
def _attn_bwd_dq(
|
|
dq,
|
|
q,
|
|
K,
|
|
V,
|
|
do,
|
|
m,
|
|
D,
|
|
qk_scale,
|
|
# shared by Q/K/V/DO.
|
|
stride_tok,
|
|
stride_d,
|
|
H,
|
|
N_CTX,
|
|
K_MEAN,
|
|
BLOCK_M2: tl.constexpr,
|
|
BLOCK_N2: tl.constexpr,
|
|
HEAD_DIM: tl.constexpr,
|
|
# Filled in by the wrapper.
|
|
start_m,
|
|
start_n,
|
|
num_steps,
|
|
MASK: tl.constexpr,
|
|
SMOOTH_K: tl.constexpr,
|
|
warp_specialize: tl.constexpr = False):
|
|
offs_m = start_m + tl.arange(0, BLOCK_M2)
|
|
offs_n = start_n + tl.arange(0, BLOCK_N2)
|
|
offs_k = tl.arange(0, HEAD_DIM)
|
|
k_ptrs = K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
v_ptrs = V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
q_valid = offs_m < N_CTX
|
|
# D (= delta) is pre-divided by ds_scale.
|
|
Di = tl.load(D + offs_m, mask=q_valid)
|
|
# BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
|
|
tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
|
|
curr_n = start_n
|
|
step_n = BLOCK_N2
|
|
|
|
if SMOOTH_K:
|
|
k_m = tl.load(K_MEAN + offs_k)
|
|
|
|
for blk_idx in range(num_steps):
|
|
# bounds checking for kv block (dynamic)
|
|
offs_n = curr_n + tl.arange(0, BLOCK_N2)
|
|
kv_valid = offs_n < N_CTX
|
|
k = tl.load(k_ptrs, mask=kv_valid[:, None])
|
|
v = tl.load(v_ptrs, mask=kv_valid[:, None])
|
|
|
|
qk = tl.dot(q, tl.trans(k))
|
|
qk = qk * qk_scale
|
|
# Autoregressive masking - apply BEFORE exp2 to match forward pass behavior
|
|
if MASK:
|
|
mask = (offs_m[:, None] >= offs_n[None, :])
|
|
qk = tl.where(mask, qk, -1.0e6)
|
|
p = tl.math.exp2(qk - m)
|
|
# Compute dP and dS.
|
|
dp = tl.dot(do, tl.trans(v))
|
|
ds = p * (dp - Di[:, None])
|
|
# Compute dQ.
|
|
# NOTE: We need to de-scale dq in the end, because k was pre-scaled.
|
|
dq += tl.dot(ds.to(tl.bfloat16), k)
|
|
|
|
if SMOOTH_K:
|
|
dq += tl.sum(ds, axis=1, keep_dims=True) * k_m[None, :]
|
|
# Increment pointers.
|
|
curr_n += step_n
|
|
k_ptrs += step_n * stride_tok
|
|
v_ptrs += step_n * stride_tok
|
|
return dq
|
|
|
|
|
|
@triton.jit
|
|
def _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_kv, stride_h_q, stride_h_kv):
|
|
"""Helper function to compute pointer offsets for cross-attention backward kernels."""
|
|
off_chz = (bhid * N_CTX_Q).to(tl.int64)
|
|
adj_q = (stride_h_q * (bhid % H) + stride_z_q * (bhid // H))
|
|
adj_kv = (stride_h_kv * (bhid % H) + stride_z_kv * (bhid // H))
|
|
return off_chz, adj_q, adj_kv
|
|
|
|
|
|
@triton.jit
|
|
def _attn_bwd_dq_cross(Q,
|
|
K,
|
|
V,
|
|
sm_scale,
|
|
DO,
|
|
DQ,
|
|
M,
|
|
D,
|
|
stride_z_q,
|
|
stride_z_kv,
|
|
stride_h_q,
|
|
stride_h_kv,
|
|
stride_tok_q,
|
|
stride_tok_kv,
|
|
stride_d_q,
|
|
stride_d_kv,
|
|
H,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
K_MEAN,
|
|
BLOCK_M2: tl.constexpr,
|
|
BLOCK_N2: tl.constexpr,
|
|
HEAD_DIM: tl.constexpr,
|
|
SMOOTH_K: tl.constexpr,
|
|
warp_specialize: tl.constexpr = False):
|
|
# Apply scale AFTER dot product for better precision
|
|
RCP_LN2: tl.constexpr = 1.4426950408889634 # = 1.0 / ln(2)
|
|
qk_scale = sm_scale * RCP_LN2
|
|
|
|
bhid = tl.program_id(2)
|
|
off_chz, adj_q, adj_kv = _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_kv, stride_h_q,
|
|
stride_h_kv)
|
|
|
|
Q += adj_q
|
|
K += adj_kv
|
|
V += adj_kv
|
|
DO += adj_q
|
|
DQ += adj_q
|
|
M += off_chz
|
|
D += off_chz
|
|
|
|
pid = tl.program_id(0)
|
|
start_m = pid * BLOCK_M2
|
|
offs_m = start_m + tl.arange(0, BLOCK_M2)
|
|
offs_k = tl.arange(0, HEAD_DIM)
|
|
|
|
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])
|
|
dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
|
|
do = tl.load(DO + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None])
|
|
m = tl.load(M + offs_m, mask=q_valid)[:, None]
|
|
|
|
Di = tl.load(D + offs_m, mask=q_valid)
|
|
num_steps = (N_CTX_KV + BLOCK_N2 - 1) // BLOCK_N2
|
|
if SMOOTH_K:
|
|
k_m = tl.load(K_MEAN + offs_k)
|
|
for step in range(num_steps):
|
|
start_n = step * BLOCK_N2
|
|
offs_n = start_n + tl.arange(0, BLOCK_N2)
|
|
kv_valid = offs_n < N_CTX_KV
|
|
|
|
k = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None])
|
|
v = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None])
|
|
|
|
qk = tl.dot(q, tl.trans(k))
|
|
qk = qk * qk_scale
|
|
p = tl.math.exp2(qk - m)
|
|
|
|
dp = tl.dot(do, tl.trans(v))
|
|
ds = p * (dp - Di[:, None])
|
|
ds = ds.to(tl.bfloat16)
|
|
dq += tl.dot(ds, k)
|
|
|
|
if SMOOTH_K:
|
|
dq += tl.sum(ds, axis=1, keep_dims=True) * k_m[None, :]
|
|
|
|
# NOTE: dq is scaled by sm_scale since K is not pre-scaled
|
|
dq *= sm_scale
|
|
dq_ptrs = DQ + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q
|
|
tl.store(dq_ptrs, dq, mask=q_valid[:, None])
|
|
|
|
|
|
@triton.jit
|
|
def _attn_bwd_dkdv_cross(Q,
|
|
K,
|
|
V,
|
|
sm_scale,
|
|
DO,
|
|
DK,
|
|
DV,
|
|
M,
|
|
D,
|
|
Q_MEAN,
|
|
stride_z_q,
|
|
stride_z_kv,
|
|
stride_h_q,
|
|
stride_h_kv,
|
|
stride_tok_q,
|
|
stride_tok_kv,
|
|
stride_d_q,
|
|
stride_d_kv,
|
|
H,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
BLOCK_M1: tl.constexpr,
|
|
BLOCK_N1: tl.constexpr,
|
|
HEAD_DIM: tl.constexpr,
|
|
IS_QAT: tl.constexpr,
|
|
two_level_quant_P: tl.constexpr = False,
|
|
fake_quant_P: tl.constexpr = True,
|
|
SMOOTH_Q: tl.constexpr = False,
|
|
use_global_sf_P: tl.constexpr = True,
|
|
warp_specialize: tl.constexpr = False):
|
|
# Apply scale AFTER dot product for better precision
|
|
RCP_LN2: tl.constexpr = 1.4426950408889634 # = 1.0 / ln(2)
|
|
qk_scale = sm_scale * RCP_LN2
|
|
|
|
bhid = tl.program_id(2)
|
|
off_chz, adj_q, adj_kv = _compute_cross_attn_pointer_offsets(bhid, H, N_CTX_Q, stride_z_q, stride_z_kv, stride_h_q,
|
|
stride_h_kv)
|
|
|
|
Q += adj_q
|
|
K += adj_kv
|
|
V += adj_kv
|
|
DO += adj_q
|
|
DK += adj_kv
|
|
DV += adj_kv
|
|
M += off_chz
|
|
D += off_chz
|
|
|
|
if SMOOTH_Q:
|
|
Q_MEAN += adj_q
|
|
|
|
pid = tl.program_id(0)
|
|
start_n = pid * BLOCK_N1
|
|
offs_n = start_n + tl.arange(0, BLOCK_N1)
|
|
offs_k = tl.arange(0, HEAD_DIM)
|
|
|
|
dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
|
|
dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
|
|
|
|
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])
|
|
v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None])
|
|
|
|
num_q_steps = (N_CTX_Q + BLOCK_M1 - 1) // BLOCK_M1
|
|
for step in range(num_q_steps):
|
|
start_m = step * BLOCK_M1
|
|
offs_m = start_m + tl.arange(0, BLOCK_M1)
|
|
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])
|
|
do = tl.load(DO + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None])
|
|
m = tl.load(M + offs_m, mask=q_valid)
|
|
|
|
qk = tl.dot(q, tl.trans(k_block))
|
|
# Apply scale AFTER dot product (matches forward pass, better precision)
|
|
qk = qk * qk_scale
|
|
p = tl.math.exp2(qk - m[:, None])
|
|
p_quant = p
|
|
if IS_QAT and fake_quant_P:
|
|
p_quant, _ = fake_quantize(src_tensor=p,
|
|
valid_src_mask=tl.full(shape=p.shape, value=1.0, dtype=p.dtype) == 1.0,
|
|
BLOCK_SIZE_OUT_DIM=BLOCK_M1,
|
|
BLOCK_SIZE_QUANT_DIM=BLOCK_N1,
|
|
dst_dtype=p.dtype,
|
|
two_level_quant_P=two_level_quant_P,
|
|
use_global_sf=use_global_sf_P)
|
|
dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do)
|
|
|
|
dp = tl.dot(do, tl.trans(v_block))
|
|
Di = tl.load(D + offs_m, mask=q_valid)
|
|
ds = p * (dp - Di[:, None])
|
|
ds = ds.to(tl.bfloat16)
|
|
dk += tl.dot(tl.trans(ds), q)
|
|
|
|
if SMOOTH_Q:
|
|
q_m = tl.load(Q_MEAN + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None])
|
|
dk += tl.sum(ds, axis=1, keep_dims=True) * q_m
|
|
|
|
dv_ptrs = DV + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv
|
|
tl.store(dv_ptrs, dv, mask=kv_valid[:, None])
|
|
|
|
dk *= sm_scale
|
|
dk_ptrs = DK + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv
|
|
tl.store(dk_ptrs, dk, mask=kv_valid[:, None])
|
|
|
|
|
|
@triton.jit
|
|
def _attn_bwd(
|
|
Q,
|
|
K,
|
|
V,
|
|
sm_scale,
|
|
DO,
|
|
DQ,
|
|
DK,
|
|
DV,
|
|
M,
|
|
D,
|
|
Q_MEAN,
|
|
# shared by Q/K/V/DO.
|
|
stride_z,
|
|
stride_h,
|
|
stride_tok,
|
|
stride_d,
|
|
H,
|
|
N_CTX,
|
|
K_MEAN,
|
|
BLOCK_M1: tl.constexpr,
|
|
BLOCK_N1: tl.constexpr,
|
|
BLOCK_M2: tl.constexpr,
|
|
BLOCK_N2: tl.constexpr,
|
|
BLK_SLICE_FACTOR: tl.constexpr,
|
|
HEAD_DIM: tl.constexpr,
|
|
CAUSAL: tl.constexpr,
|
|
IS_QAT: tl.constexpr,
|
|
SMOOTH_K: tl.constexpr,
|
|
two_level_quant_P: tl.constexpr = False,
|
|
fake_quant_P: tl.constexpr = True,
|
|
SMOOTH_Q: tl.constexpr = False,
|
|
use_global_sf_P: tl.constexpr = True,
|
|
warp_specialize: tl.constexpr = False):
|
|
|
|
bhid = tl.program_id(2)
|
|
off_chz = (bhid * N_CTX).to(tl.int64)
|
|
adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
|
|
pid = tl.program_id(0)
|
|
|
|
# offset pointers for batch/head
|
|
Q += adj
|
|
K += adj
|
|
V += adj
|
|
DO += adj
|
|
DQ += adj
|
|
DK += adj
|
|
DV += adj
|
|
M += off_chz
|
|
D += off_chz
|
|
|
|
if SMOOTH_Q:
|
|
Q_MEAN += adj
|
|
|
|
# load scales
|
|
offs_k = tl.arange(0, HEAD_DIM)
|
|
|
|
start_n = pid * BLOCK_N1
|
|
start_m = start_n
|
|
|
|
MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 // BLK_SLICE_FACTOR
|
|
offs_n = start_n + tl.arange(0, BLOCK_N1)
|
|
kv_valid = offs_n < N_CTX
|
|
|
|
dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
|
|
dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
|
|
|
|
# load K and V: they stay in SRAM throughout the inner loop.
|
|
k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None])
|
|
v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None])
|
|
|
|
# qk_scale is expressed in base-2 domain because we use exp2/log2.
|
|
# Two equivalent ways to form qk_scale:
|
|
# - **Post-scale (default)**: qk = dot(q, k^T) * (sm_scale / ln2)
|
|
# - **Pre-scale (Flash-aligned)**: k is pre-multiplied by (sm_scale / ln2) and qk_scale = 1
|
|
RCP_LN2: tl.constexpr = 1.4426950408889634 # = 1.0 / ln(2)
|
|
qk_scale = sm_scale * RCP_LN2
|
|
|
|
# For causal attention, process diagonal block with masking, then rest without
|
|
# For non-causal attention, process all blocks without masking
|
|
if CAUSAL:
|
|
num_steps = BLOCK_N1 // MASK_BLOCK_M1
|
|
|
|
dk, dv = _attn_bwd_dkdv(dk,
|
|
dv,
|
|
Q,
|
|
k,
|
|
v,
|
|
qk_scale,
|
|
DO,
|
|
M,
|
|
D,
|
|
Q_MEAN,
|
|
stride_tok,
|
|
stride_d,
|
|
H,
|
|
N_CTX,
|
|
MASK_BLOCK_M1,
|
|
BLOCK_N1,
|
|
HEAD_DIM,
|
|
start_n,
|
|
start_m,
|
|
num_steps,
|
|
MASK=True,
|
|
IS_QAT=IS_QAT,
|
|
two_level_quant_P=two_level_quant_P,
|
|
fake_quant_P=fake_quant_P,
|
|
SMOOTH_Q=SMOOTH_Q,
|
|
use_global_sf_P=use_global_sf_P,
|
|
warp_specialize=warp_specialize)
|
|
|
|
start_m += num_steps * MASK_BLOCK_M1
|
|
num_steps = (N_CTX - start_m + BLOCK_M1 - 1) // BLOCK_M1
|
|
else:
|
|
# For non-causal, start from 0 and process all Q blocks
|
|
start_m = 0
|
|
num_steps = (N_CTX + BLOCK_M1 - 1) // BLOCK_M1
|
|
|
|
# Compute dK and dV for non-masked blocks.
|
|
dk, dv = _attn_bwd_dkdv(dk,
|
|
dv,
|
|
Q,
|
|
k,
|
|
v,
|
|
qk_scale,
|
|
DO,
|
|
M,
|
|
D,
|
|
Q_MEAN,
|
|
stride_tok,
|
|
stride_d,
|
|
H,
|
|
N_CTX,
|
|
BLOCK_M1,
|
|
BLOCK_N1,
|
|
HEAD_DIM,
|
|
start_n,
|
|
start_m,
|
|
num_steps,
|
|
MASK=False,
|
|
IS_QAT=IS_QAT,
|
|
two_level_quant_P=two_level_quant_P,
|
|
fake_quant_P=fake_quant_P,
|
|
SMOOTH_Q=SMOOTH_Q,
|
|
use_global_sf_P=use_global_sf_P,
|
|
warp_specialize=warp_specialize)
|
|
|
|
dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
tl.store(dv_ptrs, dv, mask=kv_valid[:, None])
|
|
|
|
# Write back dK.
|
|
dk *= sm_scale
|
|
dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
tl.store(dk_ptrs, dk, mask=kv_valid[:, None])
|
|
|
|
# THIS BLOCK DOES DQ:
|
|
start_m = pid * BLOCK_M2
|
|
end_n = start_m + BLOCK_M2
|
|
|
|
MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 // BLK_SLICE_FACTOR
|
|
offs_m = start_m + tl.arange(0, BLOCK_M2)
|
|
q_valid = offs_m < N_CTX
|
|
|
|
q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=q_valid[:, None])
|
|
|
|
dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
|
|
do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=q_valid[:, None])
|
|
|
|
m = tl.load(M + offs_m, mask=q_valid)[:, None]
|
|
|
|
# Compute dQ for masked (diagonal) blocks.
|
|
# NOTE: This code scans each row of QK^T backward (from right to left,
|
|
# but inside each call to _attn_bwd_dq, from left to right), but that's
|
|
# not due to anything important. I just wanted to reuse the loop
|
|
# structure for dK & dV above as much as possible.
|
|
if CAUSAL:
|
|
num_steps = BLOCK_M2 // MASK_BLOCK_N2
|
|
dq = _attn_bwd_dq(
|
|
dq,
|
|
q,
|
|
K,
|
|
V,
|
|
do,
|
|
m,
|
|
D,
|
|
qk_scale,
|
|
stride_tok,
|
|
stride_d,
|
|
H,
|
|
N_CTX,
|
|
K_MEAN,
|
|
BLOCK_M2,
|
|
MASK_BLOCK_N2,
|
|
HEAD_DIM,
|
|
start_m,
|
|
end_n - num_steps * MASK_BLOCK_N2,
|
|
num_steps,
|
|
MASK=True,
|
|
SMOOTH_K=SMOOTH_K,
|
|
warp_specialize=warp_specialize,
|
|
)
|
|
end_n -= num_steps * MASK_BLOCK_N2
|
|
# stage 2: process KV blocks from 0 to end_n (before diagonal), ceiling division to include remainder
|
|
num_steps = (end_n + BLOCK_N2 - 1) // BLOCK_N2
|
|
start_n = 0
|
|
else:
|
|
# For non-causal, process all KV blocks from 0 to N_CTX (ceiling division to include remainder)
|
|
num_steps = (N_CTX + BLOCK_N2 - 1) // BLOCK_N2
|
|
start_n = 0
|
|
# stage 2
|
|
dq = _attn_bwd_dq(
|
|
dq,
|
|
q,
|
|
K,
|
|
V,
|
|
do,
|
|
m,
|
|
D,
|
|
qk_scale,
|
|
stride_tok,
|
|
stride_d,
|
|
H,
|
|
N_CTX,
|
|
K_MEAN,
|
|
BLOCK_M2,
|
|
BLOCK_N2,
|
|
HEAD_DIM,
|
|
start_m,
|
|
start_n,
|
|
num_steps,
|
|
MASK=False,
|
|
SMOOTH_K=SMOOTH_K,
|
|
warp_specialize=warp_specialize,
|
|
)
|
|
# Write back dQ.
|
|
# NOTE: dq is scaled by sm_scale since K is not pre-scaled (unlike original which used LN2)
|
|
dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
|
|
dq *= sm_scale
|
|
tl.store(dq_ptrs, dq, mask=q_valid[:, None])
|
|
|
|
|
|
class _attention(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(
|
|
ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
causal,
|
|
sm_scale,
|
|
use_qat_qkv_backward=True,
|
|
smooth_k=True,
|
|
warp_specialize=True,
|
|
IS_QAT=True,
|
|
two_level_quant_P=True,
|
|
fake_quant_P=True,
|
|
use_high_prec_o=False,
|
|
smooth_q=False,
|
|
use_global_sf_P=True,
|
|
use_global_sf_QKV=True,
|
|
):
|
|
# shape constraints
|
|
HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
|
|
# when v is in float8_e5m2 it is transposed.
|
|
HEAD_DIM_V = v.shape[-1]
|
|
assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
|
|
assert HEAD_DIM_K in {16, 32, 64, 128, 256}
|
|
|
|
# Triton 3.7's NVWS pass aborts for this kernel on Blackwell. Keep the
|
|
# architecture guard next to the kernel so direct callers and the
|
|
# FastVideo backend follow the same supported path on sm_100/sm_120.
|
|
consumer_blackwell = is_consumer_blackwell()
|
|
blackwell = is_blackwell() or consumer_blackwell
|
|
warp_specialize = warp_specialize and not blackwell
|
|
|
|
# Support different sequence lengths for q and k/v (needed for cross attention)
|
|
N_CTX_Q = q.shape[2] # Query sequence length
|
|
N_CTX_KV = k.shape[2] # Key/Value sequence length (may differ from query)
|
|
assert k.shape[2] == v.shape[2], "k and v must have the same sequence length"
|
|
sm100_optimized = (q.dtype == torch.bfloat16 and k.dtype == q.dtype and v.dtype == q.dtype
|
|
and k.device == q.device and v.device == q.device and _use_sm100_optimized_qat(
|
|
q.device,
|
|
HEAD_DIM_K,
|
|
causal,
|
|
IS_QAT,
|
|
fake_quant_P,
|
|
two_level_quant_P,
|
|
use_global_sf_P,
|
|
))
|
|
|
|
# smoothing k from SageAttn
|
|
ctx.k_mean = None
|
|
if smooth_k:
|
|
k_mean = k.mean(dim=(0, 1, 2), keepdim=True)
|
|
k = k - k_mean
|
|
ctx.k_mean = k_mean.view(-1)
|
|
|
|
q_orig, k_orig, v_orig = None, None, None
|
|
if not use_qat_qkv_backward:
|
|
q_orig, k_orig, v_orig = q, k, v
|
|
ctx.q_orig, ctx.k_orig, ctx.v_orig = q_orig, k_orig, v_orig
|
|
|
|
o = torch.empty_like(q)
|
|
if IS_QAT:
|
|
high_prec_o = torch.empty_like(q)
|
|
else:
|
|
# Initialize to a dummy value
|
|
high_prec_o = o
|
|
stage = 3 if causal else 1
|
|
extra_kern_args = {}
|
|
|
|
M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) # (Z, H, N_CTX_Q)
|
|
# Use device_descriptor for Hopper + warpspec.
|
|
if supports_host_descriptor() and not (is_hopper() and warp_specialize):
|
|
# Note that on Hopper we cannot perform a FP8 dot with a non-transposed second tensor
|
|
y_dim_q = q.shape[0] * q.shape[1] * q.shape[2]
|
|
y_dim_kv = k.shape[0] * k.shape[1] * k.shape[2]
|
|
|
|
dummy_block = [1, 1]
|
|
# (Z, H, N_CTX_Q, HEAD_DIM_K) -> (Z*H*N_CTX_Q, HEAD_DIM_K)
|
|
desc_q = TensorDescriptor(q, shape=[y_dim_q, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block)
|
|
if q.dtype == torch.float8_e5m2:
|
|
desc_v = TensorDescriptor(v,
|
|
shape=[HEAD_DIM_K, y_dim_kv],
|
|
strides=[k.shape[2], 1],
|
|
block_shape=dummy_block)
|
|
else:
|
|
desc_v = TensorDescriptor(v,
|
|
shape=[y_dim_kv, HEAD_DIM_K],
|
|
strides=[HEAD_DIM_K, 1],
|
|
block_shape=dummy_block)
|
|
desc_k = TensorDescriptor(k, shape=[y_dim_kv, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block)
|
|
# Use 3D descriptor for output to handle N_CTX_Q boundaries correctly
|
|
dummy_block3d = [1, 1, 1]
|
|
ZH = q.shape[0] * q.shape[1]
|
|
desc_o = TensorDescriptor(o,
|
|
shape=[ZH, N_CTX_Q, HEAD_DIM_K],
|
|
strides=[N_CTX_Q * HEAD_DIM_K, HEAD_DIM_K, 1],
|
|
block_shape=dummy_block3d)
|
|
if IS_QAT:
|
|
desc_high_prec_o = TensorDescriptor(high_prec_o,
|
|
shape=[ZH, N_CTX_Q, HEAD_DIM_K],
|
|
strides=[N_CTX_Q * HEAD_DIM_K, HEAD_DIM_K, 1],
|
|
block_shape=dummy_block3d)
|
|
else:
|
|
desc_high_prec_o = desc_o # Use regular output descriptor when not in QAT mode
|
|
else:
|
|
desc_q = q
|
|
desc_v = v
|
|
desc_k = k
|
|
desc_o = o
|
|
if IS_QAT:
|
|
desc_high_prec_o = high_prec_o
|
|
else:
|
|
desc_high_prec_o = o # Use regular output when not in QAT mode
|
|
|
|
def alloc_fn(size: int, align: int, _):
|
|
return torch.empty(size, dtype=torch.int8, device="cuda")
|
|
|
|
triton.set_allocator(alloc_fn)
|
|
|
|
def grid(META):
|
|
# (ceil(N_CTX / BLOCK_M), Z * H, 1)
|
|
return (triton.cdiv(q.shape[2], META["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
|
|
|
|
ctx.grid = grid
|
|
if is_blackwell() and warp_specialize:
|
|
if HEAD_DIM_K == 128 and q.dtype == torch.float16:
|
|
extra_kern_args["maxnreg"] = 168
|
|
else:
|
|
extra_kern_args["maxnreg"] = 80
|
|
|
|
qkv_block_m, qkv_block_n = 32, 32
|
|
fwd_block_m, fwd_block_n = 32, 32
|
|
fwd_num_warps, fwd_num_stages = 4, 2
|
|
fwd_mode = "legacy"
|
|
if sm100_optimized:
|
|
fwd_mode = os.environ.get("FASTVIDEO_ATTN_QAT_FWD_MODE", "fast").lower()
|
|
if fwd_mode not in {"fast", "balanced", "reference"}:
|
|
raise ValueError(f"FASTVIDEO_ATTN_QAT_FWD_MODE={fwd_mode!r} "
|
|
"(want fast|balanced|reference)")
|
|
fwd_block_m, fwd_block_n, fwd_num_warps, fwd_num_stages = _select_sm100_forward_config(
|
|
N_CTX_Q, N_CTX_KV, fwd_mode)
|
|
if IS_QAT:
|
|
fake_q = torch.empty_like(q)
|
|
fake_k = torch.empty_like(k)
|
|
fake_v = torch.empty_like(v)
|
|
|
|
# override desc_q, desc_k, desc_v with fake_q, fake_k, fake_v
|
|
if supports_host_descriptor() and not (is_hopper() and warp_specialize):
|
|
# (Z, H, N_CTX_Q, HEAD_DIM_K) -> (Z*H*N_CTX_Q, HEAD_DIM_K)
|
|
desc_q = TensorDescriptor(fake_q,
|
|
shape=[y_dim_q, HEAD_DIM_K],
|
|
strides=[HEAD_DIM_K, 1],
|
|
block_shape=dummy_block)
|
|
desc_k = TensorDescriptor(fake_k,
|
|
shape=[y_dim_kv, HEAD_DIM_K],
|
|
strides=[HEAD_DIM_K, 1],
|
|
block_shape=dummy_block)
|
|
desc_v = TensorDescriptor(fake_v,
|
|
shape=[y_dim_kv, HEAD_DIM_K],
|
|
strides=[HEAD_DIM_K, 1],
|
|
block_shape=dummy_block)
|
|
else:
|
|
desc_q = fake_q
|
|
desc_k = fake_k
|
|
desc_v = fake_v
|
|
|
|
H = q.shape[1]
|
|
grid_1 = (triton.cdiv(q.shape[2], qkv_block_m), q.shape[0] * q.shape[1], 1)
|
|
grid_2 = (triton.cdiv(k.shape[2], qkv_block_n), q.shape[0] * q.shape[1], 1)
|
|
|
|
fake_quantize_q[grid_1](
|
|
q,
|
|
fake_q,
|
|
q.stride(0),
|
|
q.stride(1),
|
|
q.stride(2),
|
|
q.stride(3),
|
|
fake_q.stride(0),
|
|
fake_q.stride(1),
|
|
fake_q.stride(2),
|
|
fake_q.stride(3),
|
|
H,
|
|
N_CTX_Q,
|
|
BLOCK_M=qkv_block_m,
|
|
HEAD_DIM=HEAD_DIM_K,
|
|
use_global_sf=use_global_sf_QKV,
|
|
)
|
|
fake_quantize_kv[grid_2](
|
|
k,
|
|
v,
|
|
fake_k,
|
|
fake_v,
|
|
k.stride(0),
|
|
k.stride(1),
|
|
k.stride(2),
|
|
k.stride(3),
|
|
fake_k.stride(0),
|
|
fake_k.stride(1),
|
|
fake_k.stride(2),
|
|
fake_k.stride(3),
|
|
H,
|
|
N_CTX_KV,
|
|
BLOCK_N=qkv_block_n,
|
|
HEAD_DIM=HEAD_DIM_K,
|
|
use_global_sf=use_global_sf_QKV,
|
|
)
|
|
|
|
# Apply pre-hook to set block shapes on tensor descriptors
|
|
_host_descriptor_pre_hook({
|
|
"BLOCK_M": fwd_block_m,
|
|
"BLOCK_N": fwd_block_n,
|
|
"HEAD_DIM": HEAD_DIM_K,
|
|
"desc_q": desc_q,
|
|
"desc_k": desc_k,
|
|
"desc_v": desc_v,
|
|
"desc_o": desc_o,
|
|
"desc_high_prec_o": desc_high_prec_o,
|
|
"FP8_OUTPUT": q.dtype == torch.float8_e5m2,
|
|
})
|
|
|
|
_attn_fwd[grid](sm_scale,
|
|
M,
|
|
q.shape[0],
|
|
q.shape[1],
|
|
desc_q,
|
|
desc_k,
|
|
desc_v,
|
|
desc_o,
|
|
desc_high_prec_o,
|
|
N_CTX_Q=N_CTX_Q,
|
|
N_CTX_KV=N_CTX_KV,
|
|
HEAD_DIM=HEAD_DIM_K,
|
|
BLOCK_M=fwd_block_m,
|
|
BLOCK_N=fwd_block_n,
|
|
FP8_OUTPUT=q.dtype == torch.float8_e5m2,
|
|
STAGE=stage,
|
|
warp_specialize=warp_specialize,
|
|
IS_HOPPER=is_hopper(),
|
|
IS_QAT=IS_QAT,
|
|
fake_quant_P=fake_quant_P,
|
|
two_level_quant_P=two_level_quant_P,
|
|
use_global_sf_P=use_global_sf_P,
|
|
JOIN_QAT_PV=(consumer_blackwell and _consumer_blackwell_join_qat_pv_enabled()),
|
|
num_warps=fwd_num_warps,
|
|
num_stages=fwd_num_stages,
|
|
**extra_kern_args)
|
|
|
|
exact_m = _sm100_exact_m_enabled()
|
|
if (sm100_optimized and fwd_mode != "reference" and exact_m and (fwd_block_m != 32 or fwd_block_n != 32)):
|
|
# The large forward tile changes a legal reduction order. Restore
|
|
# the legacy statistic so dV remains bitwise-compatible while the
|
|
# two output paths retain the faster large-tile PV computation.
|
|
assert isinstance(desc_q, TensorDescriptor)
|
|
assert isinstance(desc_k, TensorDescriptor)
|
|
desc_q.block_shape = [32, HEAD_DIM_K]
|
|
desc_k.block_shape = [32, HEAD_DIM_K]
|
|
stats_grid = (triton.cdiv(N_CTX_Q, 32), q.shape[0] * q.shape[1])
|
|
_attn_fwd_exact_m[stats_grid](
|
|
desc_q,
|
|
desc_k,
|
|
M,
|
|
sm_scale,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
HEAD_DIM=HEAD_DIM_K,
|
|
BLOCK_M=32,
|
|
BLOCK_N=32,
|
|
num_warps=8,
|
|
num_stages=4,
|
|
)
|
|
o_for_bwd = high_prec_o if IS_QAT and use_high_prec_o else o
|
|
|
|
if IS_QAT:
|
|
q = fake_q
|
|
k = fake_k
|
|
v = fake_v
|
|
|
|
ctx.save_for_backward(q, k, v, o_for_bwd, M)
|
|
ctx.sm_scale = sm_scale
|
|
ctx.HEAD_DIM = HEAD_DIM_K
|
|
ctx.causal = causal
|
|
ctx.IS_QAT = IS_QAT
|
|
ctx.use_qat_qkv_backward = use_qat_qkv_backward
|
|
ctx.smooth_k = smooth_k
|
|
ctx.two_level_quant_P = two_level_quant_P
|
|
ctx.fake_quant_P = fake_quant_P
|
|
ctx.smooth_q = smooth_q
|
|
ctx.use_global_sf_P = use_global_sf_P
|
|
ctx.warp_specialize = warp_specialize
|
|
ctx.sm100_optimized = sm100_optimized
|
|
return o
|
|
|
|
@staticmethod
|
|
def backward(ctx, do):
|
|
q, k, v, o_for_bwd, M = ctx.saved_tensors
|
|
do = do.contiguous()
|
|
assert do.is_contiguous()
|
|
dq = torch.empty_like(q)
|
|
dk = torch.empty_like(k)
|
|
dv = torch.empty_like(v)
|
|
BATCH, N_HEAD, N_CTX_Q = q.shape[:3]
|
|
N_CTX_KV = k.shape[2]
|
|
assert k.shape[2] == v.shape[2], "k and v must have the same sequence length"
|
|
PRE_BLOCK = 128
|
|
# Long video sequences are occupancy-bound on consumer Blackwell: a
|
|
# third software-pipeline stage consumes shared memory without hiding
|
|
# additional latency. Shorter sequences retain the deeper pipeline.
|
|
NUM_STAGES = 2 if is_consumer_blackwell() and max(N_CTX_Q, N_CTX_KV) >= 8192 else 3
|
|
NUM_WARPS = 4
|
|
BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 32, 32, 32
|
|
if not ctx.use_qat_qkv_backward:
|
|
q = ctx.q_orig
|
|
k = ctx.k_orig
|
|
v = ctx.v_orig
|
|
BLK_SLICE_FACTOR = 1 if ctx.IS_QAT else 2 # must be 1 for QAT
|
|
# NOTE: K is NOT pre-scaled here - scaling is applied AFTER qk dot product in kernels
|
|
# This improves precision by avoiding rounding errors in K before the dot product
|
|
arg_k = k
|
|
pre_grid = ((N_CTX_Q + PRE_BLOCK - 1) // PRE_BLOCK, BATCH * N_HEAD)
|
|
delta = torch.empty_like(M)
|
|
_attn_bwd_preprocess[pre_grid](o_for_bwd,
|
|
do,
|
|
delta,
|
|
BATCH,
|
|
N_HEAD,
|
|
N_CTX_Q,
|
|
BLOCK_M=PRE_BLOCK,
|
|
HEAD_DIM=ctx.HEAD_DIM)
|
|
|
|
q_m = None
|
|
if ctx.smooth_q:
|
|
# _, q_m = triton_group_mean(q)
|
|
q_m = q_m.repeat_interleave(q.shape[2] // q_m.shape[2], dim=2) # B,H,L,D
|
|
|
|
sm100_optimized_backward = (getattr(ctx, "sm100_optimized", False) and ctx.use_qat_qkv_backward
|
|
and not ctx.smooth_k and not ctx.smooth_q and N_CTX_KV % 16 == 0)
|
|
if sm100_optimized_backward:
|
|
# Keeping dQ and dK/dV in separate programs allows 64x64 tiles
|
|
# without carrying all three fp32 accumulators at once. On SM100
|
|
# this is substantially faster than the legacy 32x32 combined
|
|
# self-attention program with the same math and BF16 parity bounds.
|
|
block_m, block_n = 64, 64
|
|
grid_dq = ((N_CTX_Q + block_m - 1) // block_m, 1, BATCH * N_HEAD)
|
|
_attn_bwd_dq_cross[grid_dq](
|
|
q,
|
|
arg_k,
|
|
v,
|
|
ctx.sm_scale,
|
|
do,
|
|
dq,
|
|
M,
|
|
delta,
|
|
q.stride(0),
|
|
k.stride(0),
|
|
q.stride(1),
|
|
k.stride(1),
|
|
q.stride(2),
|
|
k.stride(2),
|
|
q.stride(3),
|
|
k.stride(3),
|
|
N_HEAD,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
ctx.k_mean,
|
|
BLOCK_M2=block_m,
|
|
BLOCK_N2=block_n,
|
|
HEAD_DIM=ctx.HEAD_DIM,
|
|
SMOOTH_K=False,
|
|
warp_specialize=False,
|
|
num_warps=8,
|
|
num_stages=2,
|
|
)
|
|
grid_dkdv = ((N_CTX_KV + block_n - 1) // block_n, 1, BATCH * N_HEAD)
|
|
_attn_bwd_dkdv_cross[grid_dkdv](
|
|
q,
|
|
arg_k,
|
|
v,
|
|
ctx.sm_scale,
|
|
do,
|
|
dk,
|
|
dv,
|
|
M,
|
|
delta,
|
|
q_m,
|
|
q.stride(0),
|
|
k.stride(0),
|
|
q.stride(1),
|
|
k.stride(1),
|
|
q.stride(2),
|
|
k.stride(2),
|
|
q.stride(3),
|
|
k.stride(3),
|
|
N_HEAD,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
BLOCK_M1=block_m,
|
|
BLOCK_N1=block_n,
|
|
HEAD_DIM=ctx.HEAD_DIM,
|
|
IS_QAT=True,
|
|
two_level_quant_P=False,
|
|
fake_quant_P=True,
|
|
SMOOTH_Q=False,
|
|
use_global_sf_P=False,
|
|
warp_specialize=False,
|
|
num_warps=8,
|
|
num_stages=3,
|
|
)
|
|
elif N_CTX_Q == N_CTX_KV:
|
|
# Use existing kernel for self-attention (same sequence lengths)
|
|
grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD)
|
|
_attn_bwd[grid](q,
|
|
arg_k,
|
|
v,
|
|
ctx.sm_scale,
|
|
do,
|
|
dq,
|
|
dk,
|
|
dv,
|
|
M,
|
|
delta,
|
|
q_m,
|
|
q.stride(0),
|
|
q.stride(1),
|
|
q.stride(2),
|
|
q.stride(3),
|
|
N_HEAD,
|
|
N_CTX_KV,
|
|
ctx.k_mean,
|
|
BLOCK_M1=BLOCK_M1,
|
|
BLOCK_N1=BLOCK_N1,
|
|
BLOCK_M2=BLOCK_M2,
|
|
BLOCK_N2=BLOCK_N2,
|
|
BLK_SLICE_FACTOR=BLK_SLICE_FACTOR,
|
|
HEAD_DIM=ctx.HEAD_DIM,
|
|
CAUSAL=ctx.causal,
|
|
IS_QAT=ctx.IS_QAT,
|
|
SMOOTH_K=ctx.smooth_k,
|
|
two_level_quant_P=ctx.two_level_quant_P,
|
|
fake_quant_P=ctx.fake_quant_P,
|
|
SMOOTH_Q=ctx.smooth_q,
|
|
use_global_sf_P=ctx.use_global_sf_P,
|
|
warp_specialize=ctx.warp_specialize,
|
|
num_warps=NUM_WARPS,
|
|
num_stages=NUM_STAGES)
|
|
else:
|
|
# Use separate kernels for cross-attention (different sequence lengths)
|
|
grid_dq = ((N_CTX_Q + BLOCK_M2 - 1) // BLOCK_M2, 1, BATCH * N_HEAD)
|
|
_attn_bwd_dq_cross[grid_dq](
|
|
q,
|
|
arg_k,
|
|
v,
|
|
ctx.sm_scale,
|
|
do,
|
|
dq,
|
|
M,
|
|
delta,
|
|
q.stride(0),
|
|
k.stride(0),
|
|
q.stride(1),
|
|
k.stride(1),
|
|
q.stride(2),
|
|
k.stride(2),
|
|
q.stride(3),
|
|
k.stride(3),
|
|
N_HEAD,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
ctx.k_mean,
|
|
BLOCK_M2=BLOCK_M2,
|
|
BLOCK_N2=BLOCK_N2,
|
|
HEAD_DIM=ctx.HEAD_DIM,
|
|
SMOOTH_K=ctx.smooth_k,
|
|
warp_specialize=ctx.warp_specialize,
|
|
num_warps=NUM_WARPS,
|
|
num_stages=NUM_STAGES,
|
|
)
|
|
grid_dkdv = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD)
|
|
_attn_bwd_dkdv_cross[grid_dkdv](
|
|
q,
|
|
arg_k,
|
|
v,
|
|
ctx.sm_scale,
|
|
do,
|
|
dk,
|
|
dv,
|
|
M,
|
|
delta,
|
|
q_m,
|
|
q.stride(0),
|
|
k.stride(0),
|
|
q.stride(1),
|
|
k.stride(1),
|
|
q.stride(2),
|
|
k.stride(2),
|
|
q.stride(3),
|
|
k.stride(3),
|
|
N_HEAD,
|
|
N_CTX_Q,
|
|
N_CTX_KV,
|
|
BLOCK_M1=BLOCK_M1,
|
|
BLOCK_N1=BLOCK_N1,
|
|
HEAD_DIM=ctx.HEAD_DIM,
|
|
IS_QAT=ctx.IS_QAT,
|
|
two_level_quant_P=ctx.two_level_quant_P,
|
|
fake_quant_P=ctx.fake_quant_P,
|
|
SMOOTH_Q=ctx.smooth_q,
|
|
use_global_sf_P=ctx.use_global_sf_P,
|
|
warp_specialize=ctx.warp_specialize,
|
|
num_warps=NUM_WARPS,
|
|
num_stages=NUM_STAGES,
|
|
)
|
|
|
|
return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None
|
|
|
|
|
|
attention = _attention.apply
|