checkpoint
This commit is contained in:
Matthew Noto
@@ -12,6 +12,7 @@ from fastvideo.attention.backends.abstract import (AttentionBackend,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder)
|
||||
from fastvideo.logger import init_logger
|
||||
from test_quantization_precision import get_fake_quant
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -100,6 +101,43 @@ def attn_forward_4bit_fwd_16bit_bwd(q, k, v, is_causal=False, per_block_mean=Fal
|
||||
out = _Matmul3d4bitFWD16bitBWD.apply(p.squeeze(0), v.squeeze(0).transpose(1, 2).contiguous(), scale_row.squeeze(0)).unsqueeze(0)
|
||||
return out.transpose(1, 2).contiguous()
|
||||
|
||||
def preprocess(q, k, v):
|
||||
return get_fake_quant(q), get_fake_quant(k), get_fake_quant(v)
|
||||
|
||||
def qat_attn_qk_torch(fq_q, fq_k):
|
||||
"""
|
||||
Inputs: q/k/v in [B, H, L, D]
|
||||
Returns: out in [B, H, L, D]
|
||||
"""
|
||||
B, H, qL, D = fq_q.shape
|
||||
B, H, kL, D = fq_k.shape
|
||||
assert B == 1
|
||||
q, k = fq_q.squeeze(0), fq_k.squeeze(0)
|
||||
qk = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(q.shape[-1], device=q.device, dtype=q.dtype))
|
||||
qk = qk / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype))
|
||||
p = torch.softmax(qk, dim=-1)
|
||||
return p.unsqueeze(0)
|
||||
|
||||
def qat_attn_backward_16bit(fq_q_BLHD, fq_k_BLHD, fq_v_BLHD, high_prec_o_BLHD, grad_out_BLHD, is_causal=False):
|
||||
assert is_causal == False
|
||||
fq_q, fq_k, fq_v, high_prec_o, do = fq_q_BLHD.transpose(1 ,2).contiguous(), fq_k_BLHD.transpose(1 ,2).contiguous(), fq_v_BLHD.transpose(1 ,2).contiguous(), high_prec_o_BLHD.transpose(1 ,2).contiguous(), grad_out_BLHD.transpose(1 ,2).contiguous()
|
||||
B, H, _, D = fq_q.shape
|
||||
p = qat_attn_qk_torch(fq_q, fq_k).to(do.dtype)
|
||||
fq_p = get_fake_quant(p).unsqueeze(0)
|
||||
dV = fq_p.transpose(-2, -1) @ do # [B,H,L,D]
|
||||
|
||||
dP = do @ fq_v.transpose(-2, -1) # [B,H,L,L]
|
||||
|
||||
do_dot_o = (do * high_prec_o).sum(dim=-1, keepdim=True) # [B,H,L,1]
|
||||
dS = p * (dP - do_dot_o) # [B,H,L,L]
|
||||
|
||||
|
||||
dQ = (dS @ fq_k) / torch.sqrt(torch.tensor(D, device=fq_q.device, dtype=fq_q.dtype)) # [B,H,L,D]
|
||||
dK = (dS.transpose(-2, -1) @ fq_q) / torch.sqrt(torch.tensor(D, device=fq_q.device, dtype=fq_q.dtype)) # [B,H,L,D]
|
||||
|
||||
# grads w.r.t inputs, None for non-tensor inputs
|
||||
return dQ.transpose(1, 2).contiguous(), dK.transpose(1, 2).contiguous(), dV.transpose(1, 2).contiguous(), None, None, None
|
||||
|
||||
|
||||
def sage_attn_qk_torch(q, k, per_block_mean=False):
|
||||
"""
|
||||
@@ -195,6 +233,7 @@ def flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_caus
|
||||
|
||||
USE_TORCH_4BIT_FWD = True
|
||||
USE_TORCH_16BIT_BWD = True
|
||||
USE_QAT_ATTN = True
|
||||
class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=False):
|
||||
@@ -211,8 +250,21 @@ class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function):
|
||||
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
with torch.no_grad():
|
||||
if USE_TORCH_4BIT_FWD:
|
||||
if USE_QAT_ATTN:
|
||||
fq_q, fq_k, fq_v = preprocess(q_BHLD, k_BHLD, v_BHLD)
|
||||
# Compute attention probabilities and output
|
||||
p = qat_attn_qk_torch(fq_q, fq_k) # [B, H, L, L] where B=1
|
||||
fq_p = get_fake_quant(p) # [B, H, L, L] - preserves shape
|
||||
out_BHLD = fq_p @ fq_v
|
||||
high_prec_o = p @ fq_v
|
||||
# Save fake quantized tensors and high precision output for backward
|
||||
ctx.save_for_backward(fq_q, fq_k, fq_v, high_prec_o)
|
||||
out_BLHD = out_BHLD.permute(0, 2, 1, 3).contiguous()
|
||||
elif USE_TORCH_4BIT_FWD:
|
||||
out_BHLD, p = attn_forward_4bit(q_BHLD, k_BHLD, v_BHLD, is_causal)
|
||||
# Convert to BLHD for saving
|
||||
out_BLHD = out_BHLD.permute(0, 2, 1, 3).contiguous()
|
||||
ctx.save_for_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD)
|
||||
else:
|
||||
out_BHLD = sageattn_blackwell(
|
||||
q_BHLD, k_BHLD, v_BHLD,
|
||||
@@ -220,22 +272,34 @@ class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function):
|
||||
is_causal=is_causal,
|
||||
per_block_mean=per_block_mean,
|
||||
)
|
||||
|
||||
# Back to BLHD; keep for FA bwd
|
||||
out_BLHD = out_BHLD.permute(0, 2, 1, 3).contiguous()
|
||||
ctx.save_for_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD)
|
||||
# Convert to BLHD for saving
|
||||
out_BLHD = out_BHLD.permute(0, 2, 1, 3).contiguous()
|
||||
ctx.save_for_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD)
|
||||
|
||||
return out_BLHD
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out_BLHD):
|
||||
q_BLHD, k_BLHD, v_BLHD, out_BLHD = ctx.saved_tensors
|
||||
is_causal = ctx.is_causal
|
||||
|
||||
if USE_TORCH_16BIT_BWD:
|
||||
return attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal)
|
||||
if USE_QAT_ATTN:
|
||||
fq_q, fq_k, fq_v, high_prec_o = ctx.saved_tensors
|
||||
# Convert from BHLD (saved format) to BLHD (expected by qat_attn_backward_16bit)
|
||||
fq_q_BLHD = fq_q.permute(0, 2, 1, 3).contiguous()
|
||||
fq_k_BLHD = fq_k.permute(0, 2, 1, 3).contiguous()
|
||||
fq_v_BLHD = fq_v.permute(0, 2, 1, 3).contiguous()
|
||||
high_prec_o_BLHD = high_prec_o.permute(0, 2, 1, 3).contiguous()
|
||||
# qat_attn_backward_16bit expects BLHD format and returns BLHD format
|
||||
grad_q_BLHD, grad_k_BLHD, grad_v_BLHD, _, _, _ = qat_attn_backward_16bit(
|
||||
fq_q_BLHD, fq_k_BLHD, fq_v_BLHD, high_prec_o_BLHD, grad_out_BLHD, is_causal
|
||||
)
|
||||
return grad_q_BLHD, grad_k_BLHD, grad_v_BLHD, None, None
|
||||
else:
|
||||
return flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal)
|
||||
q_BLHD, k_BLHD, v_BLHD, out_BLHD = ctx.saved_tensors
|
||||
if USE_TORCH_16BIT_BWD:
|
||||
return attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal)
|
||||
else:
|
||||
return flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal)
|
||||
|
||||
|
||||
|
||||
|
||||
+357
@@ -0,0 +1,357 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/numerics_details/mxfp_details/_upcast_from_mxfp.py
|
||||
# and https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/numerics_details/mxfp_details/_downcast_to_mxfp.py
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from triton_kernels.target_info import cuda_capability_geq
|
||||
|
||||
MXFP_BLOCK_SIZE = tl.constexpr(32)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _get_max_quant_val(dtype: tl.constexpr):
|
||||
if dtype == tl.uint8:
|
||||
return 6.0
|
||||
elif dtype == tl.float8e5:
|
||||
return 57344.0
|
||||
elif dtype == tl.float8e4nv:
|
||||
return 448.0
|
||||
else:
|
||||
tl.static_assert(False, f"Invalid {dtype=}")
|
||||
|
||||
@triton.jit
|
||||
def _get_max_power_of_2_quant_val(dtype: tl.constexpr):
|
||||
if dtype == tl.uint8:
|
||||
return 4.0
|
||||
elif dtype == tl.float8e5:
|
||||
return 32768.0
|
||||
elif dtype == tl.float8e4nv:
|
||||
return 256.0
|
||||
|
||||
@triton.jit
|
||||
def _compute_quant_and_scale(src_tensor, valid_src_mask, mx_tensor_dtype: tl.constexpr,
|
||||
DEQUANT_SCALE_ROUNDING_MODE: tl.constexpr = 0):
|
||||
is_fp8: tl.constexpr = mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5
|
||||
BLOCK_SIZE_OUT_DIM: tl.constexpr = src_tensor.shape[0]
|
||||
BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1]
|
||||
BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // MXFP_BLOCK_SIZE
|
||||
|
||||
# Explicit cast to fp32 since most ops are not supported on bfloat16. We avoid needless conversions to and from bf16
|
||||
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
|
||||
abs_tensor = tl.reshape(abs_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
|
||||
max_val = tl.max(abs_tensor, axis=2, keep_dims=True)
|
||||
if DEQUANT_SCALE_ROUNDING_MODE == 0:
|
||||
# DequantScaleRoundingMode.ROUND_UP
|
||||
# compute 2 ** ceil(log2(dequant_scale))
|
||||
# Adding 0x007FFFFF adds exponent by 1 unless mantissa is all zeros
|
||||
# A corner case: exponent is 0xFF that will overflow but that's already
|
||||
# NaN so assume we don't care.
|
||||
dequant_scale = max_val / _get_max_quant_val(mx_tensor_dtype)
|
||||
dequant_scale_exponent = (dequant_scale.to(tl.uint32, bitcast=True) + 0x007FFFFF) & 0x7F800000
|
||||
else:
|
||||
# DequantScaleRoundingMode.ROUND_DOWN
|
||||
# compute 2 ** floor(log2(dequant_scale))
|
||||
assert DEQUANT_SCALE_ROUNDING_MODE == 1
|
||||
dequant_scale = max_val / _get_max_power_of_2_quant_val(mx_tensor_dtype)
|
||||
dequant_scale_exponent = dequant_scale.to(tl.uint32, bitcast=True) & 0x7F800000
|
||||
dequant_scale_rounded = dequant_scale_exponent.to(tl.float32, bitcast=True)
|
||||
quant_scale = tl.where(dequant_scale_rounded == 0, 0, 1.0 / dequant_scale_rounded)
|
||||
|
||||
f32_tensor = tl.reshape(f32_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
|
||||
quant_tensor = f32_tensor * quant_scale
|
||||
|
||||
# Reshape the tensors after scaling
|
||||
quant_tensor = quant_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM])
|
||||
# Set the invalid portions of the tensor to 0. This will ensure that any padding tensors are 0 in the mx format.
|
||||
quant_tensor = tl.where(valid_src_mask, quant_tensor, 0)
|
||||
dequant_scale_exponent = dequant_scale_exponent.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE])
|
||||
|
||||
# First, we simply extract the exponent part of the scales and store the result
|
||||
dequant_scale_exponent = (dequant_scale_exponent >> 23).to(tl.uint8)
|
||||
# Now we must convert the tensors to the mx format.
|
||||
if is_fp8:
|
||||
out_tensor = quant_tensor.to(mx_tensor_dtype)
|
||||
elif cuda_capability_geq(10, 0):
|
||||
# Convert scaled values to two f32 lanes and use PTX cvt to e2m1x2 with two f32 operands.
|
||||
pairs = tl.reshape(quant_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2])
|
||||
lo_f, hi_f = tl.split(pairs)
|
||||
lo_f32 = lo_f.to(tl.float32)
|
||||
hi_f32 = hi_f.to(tl.float32)
|
||||
|
||||
# Inline PTX: cvt.rn.satfinite.e2m1x2.f32 takes two f32 sources and produces one .b8 packed e2m1x2.
|
||||
out_tensor = tl.inline_asm_elementwise(
|
||||
"""
|
||||
{
|
||||
.reg .b8 r;
|
||||
cvt.rn.satfinite.e2m1x2.f32 r, $1, $2;
|
||||
mov.b32 $0, {r, r, r, r};
|
||||
}
|
||||
""",
|
||||
constraints="=r,f,f",
|
||||
args=[hi_f32, lo_f32],
|
||||
dtype=tl.uint8,
|
||||
is_pure=True,
|
||||
pack=1,
|
||||
)
|
||||
else:
|
||||
quant_tensor = quant_tensor.to(tl.uint32, bitcast=True)
|
||||
signs = quant_tensor & 0x80000000
|
||||
exponents = (quant_tensor >> 23) & 0xFF
|
||||
mantissas_orig = (quant_tensor & 0x7FFFFF)
|
||||
|
||||
# For RTNE: 0.25 < x < 0.75 maps to 0.5 (denormal); exactly 0.25 maps to 0.0
|
||||
E8_BIAS = 127
|
||||
E2_BIAS = 1
|
||||
# Move implicit bit 1 at the beginning to mantissa for denormals
|
||||
is_subnormal = exponents < E8_BIAS
|
||||
adjusted_exponents = tl.core.sub(E8_BIAS, exponents + 1, sanitize_overflow=False)
|
||||
mantissas_pre = (0x400000 | (mantissas_orig >> 1))
|
||||
mantissas = tl.where(is_subnormal, mantissas_pre >> adjusted_exponents, mantissas_orig)
|
||||
|
||||
# For normal numbers, we change the bias from 127 to 1, and for subnormals, we keep exponent as 0.
|
||||
exponents = tl.maximum(exponents, E8_BIAS - E2_BIAS) - (E8_BIAS - E2_BIAS)
|
||||
|
||||
# Combine sign, exponent, and mantissa, while saturating
|
||||
# Round to nearest, ties to even (RTNE): use guard/sticky and LSB to decide increment
|
||||
m2bits = mantissas >> 21
|
||||
lsb_keep = (m2bits >> 1) & 0x1
|
||||
guard = m2bits & 0x1
|
||||
IS_SRC_FP32: tl.constexpr = src_tensor.dtype == tl.float32
|
||||
if IS_SRC_FP32:
|
||||
bit0_dropped = (mantissas_orig & 0x1) != 0
|
||||
mask = (1 << tl.minimum(adjusted_exponents, 31)) - 1
|
||||
dropped_post = (mantissas_pre & mask) != 0
|
||||
sticky = is_subnormal & (bit0_dropped | dropped_post)
|
||||
sticky |= ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)
|
||||
else:
|
||||
sticky = ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)
|
||||
round_inc = guard & (sticky | lsb_keep)
|
||||
e2m1_tmp = tl.minimum((((exponents << 2) | m2bits) + round_inc) >> 1, 0x7)
|
||||
e2m1_value = ((signs >> 28) | e2m1_tmp).to(tl.uint8)
|
||||
|
||||
e2m1_value = tl.reshape(e2m1_value, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2])
|
||||
evens, odds = tl.split(e2m1_value)
|
||||
out_tensor = evens | (odds << 4)
|
||||
|
||||
return out_tensor, dequant_scale_exponent
|
||||
|
||||
@triton.jit
|
||||
def _downcast_to_mxfp(
|
||||
mx_tensor_ptr, stride_mxt_outer, stride_mxt_quant: tl.constexpr,
|
||||
mx_scale_ptr, stride_mx_scale_outer, stride_mx_scale_quant,
|
||||
src_ptr, stride_src_outer, stride_src_quant, outer_dim, quant_dim,
|
||||
BLOCK_SIZE_OUT_DIM:tl.constexpr,
|
||||
BLOCK_SIZE_QUANT_DIM: tl.constexpr,
|
||||
DEQUANT_SCALE_ROUNDING_MODE: tl.constexpr,
|
||||
):
|
||||
|
||||
tl.static_assert(stride_mxt_quant == 1, f"Output stride, {stride_mxt_quant=} must be 1.")
|
||||
tl.static_assert(BLOCK_SIZE_QUANT_DIM % MXFP_BLOCK_SIZE == 0, f"{BLOCK_SIZE_QUANT_DIM=} must be a multiple of 32")
|
||||
|
||||
# uint8 signifies two fp4 e2m1 values packed into a single byte
|
||||
mx_tensor_dtype: tl.constexpr = mx_tensor_ptr.dtype.element_ty
|
||||
tl.static_assert(mx_tensor_dtype == tl.uint8 or (mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5),
|
||||
f"Invalid {mx_tensor_dtype=}. Must be uint8 or float8.")
|
||||
|
||||
src_dtype: tl.constexpr = src_ptr.dtype.element_ty
|
||||
tl.static_assert(mx_scale_ptr.dtype.element_ty == tl.uint8, f"{mx_scale_ptr.dtype.element_ty=} must be uint8")
|
||||
tl.static_assert((src_dtype == tl.bfloat16) or (src_dtype == tl.float16) or (src_dtype == tl.float32), f"{src_dtype=} must be bfloat16 or float16 or float32")
|
||||
is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8
|
||||
|
||||
outer_block = tl.program_id(0).to(tl.int64)
|
||||
quant_block = tl.program_id(1).to(tl.int64)
|
||||
|
||||
K_DIVISOR: tl.constexpr = 2 if is_fp4 else 1
|
||||
BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // MXFP_BLOCK_SIZE
|
||||
BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // K_DIVISOR
|
||||
|
||||
start_src_quant = quant_block * BLOCK_SIZE_QUANT_DIM
|
||||
start_mx_scale_quant = quant_block * BLOCK_SIZE_QUANT_MX_SCALE
|
||||
start_mx_quant = quant_block * BLOCK_SIZE_QUANT_MX_TENSOR
|
||||
start_out = outer_block * BLOCK_SIZE_OUT_DIM
|
||||
|
||||
src_ptr += start_src_quant * stride_src_quant + start_out * stride_src_outer
|
||||
mx_scale_ptr += start_mx_scale_quant * stride_mx_scale_quant + start_out * stride_mx_scale_outer
|
||||
mx_tensor_ptr += start_mx_quant * stride_mxt_quant + start_out * stride_mxt_outer
|
||||
|
||||
offs_src_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :].to(tl.int64)
|
||||
offs_mxt_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :].to(tl.int64)
|
||||
offs_scale_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :].to(tl.int64)
|
||||
offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None].to(tl.int64)
|
||||
|
||||
mask_src_quant = start_src_quant + offs_src_quant < quant_dim
|
||||
mask_n = start_out + offs_outer < outer_dim
|
||||
full_mask_src = mask_src_quant & mask_n
|
||||
|
||||
mask_mxt_quant = start_mx_quant + offs_mxt_quant < quant_dim // K_DIVISOR # requires quant_dim % K_DIVISOR == 0
|
||||
full_mask_mxt = mask_mxt_quant & mask_n
|
||||
|
||||
scale_mask_k = start_mx_scale_quant + offs_scale_quant < quant_dim // MXFP_BLOCK_SIZE # requires quant_dim % MXFP_BLOCK_SIZE == 0
|
||||
full_scale_mask = scale_mask_k & mask_n
|
||||
|
||||
src_tensor_offsets = offs_src_quant * stride_src_quant + offs_outer * stride_src_outer
|
||||
mx_scale_offsets = offs_scale_quant * stride_mx_scale_quant + offs_outer * stride_mx_scale_outer
|
||||
mx_tensor_offsets = offs_mxt_quant * stride_mxt_quant + offs_outer * stride_mxt_outer
|
||||
src_tensor = tl.load(src_ptr + src_tensor_offsets, mask=full_mask_src)
|
||||
|
||||
out_tensor, scale_tensor = _compute_quant_and_scale(src_tensor, full_mask_src, mx_tensor_dtype,
|
||||
DEQUANT_SCALE_ROUNDING_MODE)
|
||||
|
||||
tl.store(mx_scale_ptr + mx_scale_offsets, scale_tensor, mask=full_scale_mask)
|
||||
tl.store(mx_tensor_ptr + mx_tensor_offsets, out_tensor, mask=full_mask_mxt)
|
||||
|
||||
@triton.jit
|
||||
def _upcast_from_mxfp(
|
||||
out_desc,
|
||||
mx_tensor_desc,
|
||||
mx_scale_ptr,
|
||||
stride_scale_outer,
|
||||
stride_scale_quant,
|
||||
outer_dim,
|
||||
quant_dim,
|
||||
BLOCK_SIZE_OUT_DIM: tl.constexpr,
|
||||
BLOCK_SIZE_QUANT_DIM: 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=}")
|
||||
# uint8 signifies two fp4 e2m1 values packed into a single byte
|
||||
mx_tensor_dtype: tl.constexpr = mx_tensor_desc.dtype
|
||||
dst_dtype: tl.constexpr = out_desc.dtype
|
||||
tl.static_assert(dst_dtype == tl.float16 or dst_dtype == tl.bfloat16 or dst_dtype == tl.float32)
|
||||
tl.static_assert(
|
||||
mx_tensor_dtype == tl.uint8
|
||||
or ((mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5) or mx_tensor_dtype == dst_dtype),
|
||||
"mx_tensor_ptr must be uint8 or float8 or dst_dtype")
|
||||
tl.static_assert(mx_scale_ptr.dtype.element_ty == tl.uint8, "mx_scale_ptr must be uint8")
|
||||
|
||||
# Determine if we are dealing with fp8 types.
|
||||
is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8
|
||||
is_fp8: tl.constexpr = mx_tensor_dtype == tl.float8e4nv or mx_tensor_dtype == tl.float8e5
|
||||
K_DIVISOR: tl.constexpr = 2 if is_fp4 else 1
|
||||
BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // MXFP_BLOCK_SIZE
|
||||
BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // K_DIVISOR
|
||||
|
||||
# Compute starting indices for the quantized (packed) dimension and the outer dimension.
|
||||
outer_block = tl.program_id(0).to(tl.int64)
|
||||
quant_block = tl.program_id(1).to(tl.int64)
|
||||
|
||||
start_mxt_quant = quant_block * BLOCK_SIZE_QUANT_MX_TENSOR
|
||||
start_out_quant = quant_block * BLOCK_SIZE_QUANT_DIM
|
||||
start_mx_scale_quant = quant_block * BLOCK_SIZE_QUANT_MX_SCALE
|
||||
start_out = outer_block * BLOCK_SIZE_OUT_DIM
|
||||
|
||||
# Load the quantized value tensor.
|
||||
tensor = mx_tensor_desc.load([start_out.to(tl.int32), start_mxt_quant.to(tl.int32)])
|
||||
|
||||
offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None].to(tl.int64)
|
||||
offs_scale = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :].to(tl.int64)
|
||||
mask_outer = start_out + offs_outer < outer_dim
|
||||
mask_scale = start_mx_scale_quant + offs_scale < tl.cdiv(quant_dim, MXFP_BLOCK_SIZE)
|
||||
full_scale_mask = mask_scale & mask_outer
|
||||
scale_offsets = offs_scale * stride_scale_quant + offs_outer * stride_scale_outer
|
||||
scale_ptr_base = mx_scale_ptr + start_out * stride_scale_outer + start_mx_scale_quant * stride_scale_quant
|
||||
scale = tl.load(scale_ptr_base + scale_offsets, mask=full_scale_mask)
|
||||
|
||||
# Upcast the scale to the destination type.
|
||||
if dst_dtype == tl.bfloat16:
|
||||
dst_scale = (scale.to(tl.uint16) << 7).to(dst_dtype, bitcast=True)
|
||||
else:
|
||||
dst_scale = (scale.to(tl.uint32) << 23).to(tl.float32, bitcast=True)
|
||||
if dst_dtype == tl.float16:
|
||||
dst_scale = dst_scale.to(tl.float16)
|
||||
|
||||
# Now upcast the tensor.
|
||||
intermediate_dtype: tl.constexpr = tl.bfloat16 if dst_dtype == tl.float32 else dst_dtype
|
||||
if is_fp8:
|
||||
dst_tensor = tensor.to(intermediate_dtype)
|
||||
if tensor.dtype == tl.float8e5:
|
||||
from_e_bits: tl.constexpr = 5
|
||||
from_m_bits: tl.constexpr = 2
|
||||
to_e_bits: tl.constexpr = 8 if intermediate_dtype == tl.bfloat16 else 5
|
||||
to_m_bits: tl.constexpr = 7 if intermediate_dtype == tl.bfloat16 else 10
|
||||
|
||||
# Preserve infs and nans. FIXME Fp8E5M2_to_Bf16 doesn't preserve them!
|
||||
non_finite_mask_src: tl.constexpr = ((1 << from_e_bits) - 1) << from_m_bits
|
||||
non_finite_mask_dst: tl.constexpr = ((1 << to_e_bits) - 1) << to_m_bits
|
||||
dst_tensor = tl.where(
|
||||
(tensor.to(tl.uint8, bitcast=True) & non_finite_mask_src) == non_finite_mask_src,
|
||||
(dst_tensor.to(tl.uint16, bitcast=True) | non_finite_mask_dst).to(intermediate_dtype, bitcast=True),
|
||||
dst_tensor,
|
||||
)
|
||||
|
||||
elif cuda_capability_geq(10, 0):
|
||||
assert is_fp4
|
||||
packed_u32 = tl.inline_asm_elementwise(
|
||||
asm="""
|
||||
{
|
||||
.reg .b8 in_8;
|
||||
.reg .f16x2 out;
|
||||
cvt.u8.u32 in_8, $1;
|
||||
cvt.rn.f16x2.e2m1x2 out, in_8;
|
||||
mov.b32 $0, out;
|
||||
}
|
||||
""",
|
||||
constraints="=r,r",
|
||||
args=[tensor], # tl.uint8 passed in as a 32-bit reg with value in low 8 bits
|
||||
dtype=tl.uint32,
|
||||
is_pure=True,
|
||||
pack=1,
|
||||
)
|
||||
lo_u16 = (packed_u32 & 0xFFFF).to(tl.uint16)
|
||||
hi_u16 = (packed_u32 >> 16).to(tl.uint16)
|
||||
lo_f16 = lo_u16.to(tl.float16, bitcast=True)
|
||||
hi_f16 = hi_u16.to(tl.float16, bitcast=True)
|
||||
|
||||
if intermediate_dtype == tl.float16:
|
||||
x0, x1 = lo_f16, hi_f16
|
||||
else:
|
||||
x0 = lo_f16.to(intermediate_dtype)
|
||||
x1 = hi_f16.to(intermediate_dtype)
|
||||
|
||||
dst_tensor = tl.interleave(x0, x1)
|
||||
|
||||
else:
|
||||
assert is_fp4
|
||||
dst_bias: tl.constexpr = 127 if intermediate_dtype == tl.bfloat16 else 15
|
||||
dst_0p5: tl.constexpr = 16128 if intermediate_dtype == tl.bfloat16 else 0x3800
|
||||
dst_m_bits: tl.constexpr = 7 if intermediate_dtype == tl.bfloat16 else 10
|
||||
# e2m1
|
||||
em0 = tensor & 0x07
|
||||
em1 = tensor & 0x70
|
||||
x0 = (em0.to(tl.uint16) << (dst_m_bits - 1)) | ((tensor & 0x08).to(tl.uint16) << 12)
|
||||
x1 = (em1.to(tl.uint16) << (dst_m_bits - 5)) | ((tensor & 0x80).to(tl.uint16) << 8)
|
||||
# Three cases:
|
||||
# 1) x is normal and non-zero: Correct bias
|
||||
x0 = tl.where((em0 & 0x06) != 0, x0 + ((dst_bias - 1) << dst_m_bits), x0)
|
||||
x1 = tl.where((em1 & 0x60) != 0, x1 + ((dst_bias - 1) << dst_m_bits), x1)
|
||||
# 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in the dst type
|
||||
x0 = tl.where(em0 == 0x01, dst_0p5 | (x0 & 0x8000), x0)
|
||||
x1 = tl.where(em1 == 0x10, dst_0p5 | (x1 & 0x8000), x1)
|
||||
# 3) x is zero, do nothing
|
||||
dst_tensor = tl.interleave(x0, x1).to(intermediate_dtype, bitcast=True)
|
||||
|
||||
dst_tensor = dst_tensor.to(dst_dtype)
|
||||
|
||||
# Reshape for proper broadcasting: the scale was stored with a 32‐sized “inner” grouping.
|
||||
dst_tensor = dst_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
|
||||
dst_scale = dst_scale.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1])
|
||||
scale = scale.reshape(dst_scale.shape)
|
||||
|
||||
out_tensor = dst_tensor * dst_scale
|
||||
if dst_dtype == tl.float32:
|
||||
max_fin = 3.4028234663852886e+38
|
||||
elif dst_dtype == tl.bfloat16:
|
||||
max_fin = 3.3895313892515355e+38
|
||||
else:
|
||||
tl.static_assert(dst_dtype == tl.float16)
|
||||
max_fin = 65504
|
||||
# TODO: handle infinity same as upcast_from_mxfp_torch together with the
|
||||
# above FIXME
|
||||
out_tensor = tl.clamp(out_tensor, min=-max_fin, max=max_fin)
|
||||
# Correct any NaNs encoded via the scale.
|
||||
out_tensor = tl.where(scale == 0xFF, float("nan"), out_tensor)
|
||||
out_tensor = out_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM])
|
||||
out_desc.store([start_out.to(tl.int32), start_out_quant.to(tl.int32)], out_tensor)
|
||||
+762
@@ -0,0 +1,762 @@
|
||||
"""
|
||||
Fused Attention
|
||||
===============
|
||||
|
||||
This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao (https://tridao.me/publications/flash2/flash2.pdf)
|
||||
|
||||
Credits: OpenAI kernel team
|
||||
|
||||
Extra Credits:
|
||||
|
||||
* Original flash attention paper (https://arxiv.org/abs/2205.14135)
|
||||
* Rabe and Staats (https://arxiv.org/pdf/2112.05682v2.pdf)
|
||||
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import os
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from triton.tools.tensor_descriptor import TensorDescriptor
|
||||
|
||||
DEVICE = triton.runtime.driver.active.get_active_torch_device()
|
||||
|
||||
|
||||
def is_hip():
|
||||
return triton.runtime.driver.active.get_current_target().backend == "hip"
|
||||
|
||||
|
||||
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_blackwell():
|
||||
return is_cuda() and torch.cuda.get_device_capability()[0] == 10
|
||||
|
||||
|
||||
def is_hopper():
|
||||
return is_cuda() and torch.cuda.get_device_capability()[0] == 9
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_fwd_inner(acc, l_i, m_i, q, #
|
||||
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):
|
||||
# range of values handled by this stage
|
||||
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
|
||||
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)
|
||||
# -- compute qk ----
|
||||
k = desc_k.load([offsetk_y, 0]).T
|
||||
qk = tl.dot(q, k)
|
||||
if STAGE == 2:
|
||||
mask = offs_m[:, None] >= (start_n + offs_n[None, :])
|
||||
qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
|
||||
m_ij = tl.maximum(m_i, tl.max(qk, 1))
|
||||
qk -= m_ij[:, None]
|
||||
else:
|
||||
m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
|
||||
qk = qk * qk_scale - m_ij[:, None]
|
||||
p = tl.math.exp2(qk)
|
||||
# -- compute correction factor
|
||||
alpha = tl.math.exp2(m_i - m_ij)
|
||||
l_ij = tl.sum(p, 1)
|
||||
# -- 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]
|
||||
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])
|
||||
else:
|
||||
acc = 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])
|
||||
p = p.to(dtype)
|
||||
# note that this non transposed v for FP8 is only supported on Blackwell
|
||||
acc = tl.dot(p, v, 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
|
||||
return 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 = [BLOCK_M, HEAD_DIM]
|
||||
|
||||
|
||||
if is_hip():
|
||||
NUM_STAGES_OPTIONS = [1]
|
||||
elif supports_host_descriptor():
|
||||
NUM_STAGES_OPTIONS = [2, 3, 4]
|
||||
else:
|
||||
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 = kwargs["N_CTX"]
|
||||
|
||||
# Filter out configs where BLOCK_M > N_CTX
|
||||
return [conf for conf in configs if conf.kwargs.get("BLOCK_M", 0) <= N_CTX]
|
||||
|
||||
|
||||
@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", "HEAD_DIM", "FP8_OUTPUT", "warp_specialize"],
|
||||
prune_configs_by={'early_config_prune': prune_invalid_configs})
|
||||
@triton.jit
|
||||
def _attn_fwd(sm_scale, M, #
|
||||
Z, H, desc_q, desc_k, desc_v, desc_o, N_CTX, #
|
||||
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, #
|
||||
):
|
||||
dtype = tl.float8e5 if FP8_OUTPUT else tl.float16
|
||||
tl.static_assert(BLOCK_N <= HEAD_DIM)
|
||||
start_m = tl.program_id(0)
|
||||
off_hz = tl.program_id(1)
|
||||
off_z = off_hz // H
|
||||
off_h = off_hz % H
|
||||
|
||||
y_dim = Z * H * N_CTX
|
||||
desc_q = _maybe_make_tensor_desc(desc_q, shape=[y_dim, 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], strides=[N_CTX, 1],
|
||||
block_shape=[HEAD_DIM, BLOCK_N])
|
||||
else:
|
||||
desc_v = _maybe_make_tensor_desc(desc_v, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1],
|
||||
block_shape=[BLOCK_N, HEAD_DIM])
|
||||
desc_k = _maybe_make_tensor_desc(desc_k, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1],
|
||||
block_shape=[BLOCK_N, HEAD_DIM])
|
||||
desc_o = _maybe_make_tensor_desc(desc_o, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1],
|
||||
block_shape=[BLOCK_M, HEAD_DIM])
|
||||
|
||||
offset_y = off_z * (N_CTX * H) + off_h * N_CTX
|
||||
qo_offset_y = offset_y + start_m * BLOCK_M
|
||||
# initialize offsets
|
||||
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
offs_n = tl.arange(0, BLOCK_N)
|
||||
# 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)
|
||||
# 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])
|
||||
# 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, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, #
|
||||
desc_k, desc_v, #
|
||||
offset_y, dtype, start_m, qk_scale, #
|
||||
BLOCK_M, HEAD_DIM, BLOCK_N, #
|
||||
4 - STAGE, offs_m, offs_n, N_CTX, #
|
||||
warp_specialize, IS_HOPPER)
|
||||
# stage 2: on-band
|
||||
if STAGE & 2:
|
||||
acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, #
|
||||
desc_k, desc_v, #
|
||||
offset_y, dtype, start_m, qk_scale, #
|
||||
BLOCK_M, HEAD_DIM, BLOCK_N, #
|
||||
2, offs_m, offs_n, N_CTX, #
|
||||
warp_specialize, IS_HOPPER)
|
||||
# epilogue
|
||||
m_i += tl.math.log2(l_i)
|
||||
acc = acc / l_i[:, None]
|
||||
m_ptrs = M + off_hz * N_CTX + offs_m
|
||||
tl.store(m_ptrs, m_i)
|
||||
desc_o.store([qo_offset_y, 0], acc.to(dtype))
|
||||
|
||||
|
||||
@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)
|
||||
# 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)
|
||||
delta = tl.sum(o * do, axis=1)
|
||||
# write-back
|
||||
tl.store(Delta + off_hz * N_CTX + off_m, delta)
|
||||
|
||||
|
||||
# The main inner-loop logic for computing dK and dV.
|
||||
@triton.jit
|
||||
def _attn_bwd_dkdv(dk, dv, #
|
||||
Q, k, v, sm_scale, #
|
||||
DO, #
|
||||
M, D, #
|
||||
# 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):
|
||||
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)
|
||||
qT_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
|
||||
# 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):
|
||||
qT = tl.load(qT_ptrs)
|
||||
# Load m before computing qk to reduce pipeline stall.
|
||||
offs_m = curr_m + tl.arange(0, BLOCK_M1)
|
||||
m = tl.load(M + offs_m)
|
||||
qkT = tl.dot(k, qT)
|
||||
pT = tl.math.exp2(qkT - m[None, :])
|
||||
# Autoregressive masking.
|
||||
if MASK:
|
||||
mask = (offs_m[None, :] >= offs_n[:, None])
|
||||
pT = tl.where(mask, pT, 0.0)
|
||||
do = tl.load(do_ptrs)
|
||||
# Compute dV.
|
||||
ppT = pT
|
||||
ppT = ppT.to(tl.float16)
|
||||
dv += tl.dot(ppT, do)
|
||||
# D (= delta) is pre-divided by ds_scale.
|
||||
Di = tl.load(D + offs_m)
|
||||
# Compute dP and dS.
|
||||
dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
|
||||
dsT = pT * (dpT - Di[None, :])
|
||||
dsT = dsT.to(tl.float16)
|
||||
dk += tl.dot(dsT, tl.trans(qT))
|
||||
# Increment pointers.
|
||||
curr_m += step_m
|
||||
qT_ptrs += step_m * stride_tok
|
||||
do_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,
|
||||
# shared by Q/K/V/DO.
|
||||
stride_tok, stride_d, #
|
||||
H, N_CTX, #
|
||||
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):
|
||||
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)
|
||||
kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
|
||||
vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
|
||||
# D (= delta) is pre-divided by ds_scale.
|
||||
Di = tl.load(D + offs_m)
|
||||
# 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
|
||||
for blk_idx in range(num_steps):
|
||||
kT = tl.load(kT_ptrs)
|
||||
vT = tl.load(vT_ptrs)
|
||||
qk = tl.dot(q, kT)
|
||||
p = tl.math.exp2(qk - m)
|
||||
# Autoregressive masking.
|
||||
if MASK:
|
||||
offs_n = curr_n + tl.arange(0, BLOCK_N2)
|
||||
mask = (offs_m[:, None] >= offs_n[None, :])
|
||||
p = tl.where(mask, p, 0.0)
|
||||
# Compute dP and dS.
|
||||
dp = tl.dot(do, vT).to(tl.float32)
|
||||
ds = p * (dp - Di[:, None])
|
||||
ds = ds.to(tl.float16)
|
||||
# Compute dQ.
|
||||
# NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
|
||||
dq += tl.dot(ds, tl.trans(kT))
|
||||
# Increment pointers.
|
||||
curr_n += step_n
|
||||
kT_ptrs += step_n * stride_tok
|
||||
vT_ptrs += step_n * stride_tok
|
||||
return dq
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_bwd(Q, K, V, sm_scale, #
|
||||
DO, #
|
||||
DQ, DK, DV, #
|
||||
M, D,
|
||||
# shared by Q/K/V/DO.
|
||||
stride_z, stride_h, stride_tok, stride_d, #
|
||||
H, N_CTX, #
|
||||
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):
|
||||
LN2: tl.constexpr = 0.6931471824645996 # = ln(2)
|
||||
|
||||
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
|
||||
|
||||
# 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)
|
||||
|
||||
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)
|
||||
v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
|
||||
|
||||
num_steps = BLOCK_N1 // MASK_BLOCK_M1
|
||||
|
||||
dk, dv = _attn_bwd_dkdv(dk, dv, #
|
||||
Q, k, v, sm_scale, #
|
||||
DO, #
|
||||
M, D, #
|
||||
stride_tok, stride_d, #
|
||||
H, N_CTX, #
|
||||
MASK_BLOCK_M1, BLOCK_N1, HEAD_DIM, #
|
||||
start_n, start_m, num_steps, #
|
||||
MASK=True #
|
||||
)
|
||||
|
||||
start_m += num_steps * MASK_BLOCK_M1
|
||||
num_steps = (N_CTX - start_m) // BLOCK_M1
|
||||
|
||||
# Compute dK and dV for non-masked blocks.
|
||||
dk, dv = _attn_bwd_dkdv( #
|
||||
dk, dv, #
|
||||
Q, k, v, sm_scale, #
|
||||
DO, #
|
||||
M, D, #
|
||||
stride_tok, stride_d, #
|
||||
H, N_CTX, #
|
||||
BLOCK_M1, BLOCK_N1, HEAD_DIM, #
|
||||
start_n, start_m, num_steps, #
|
||||
MASK=False #
|
||||
)
|
||||
|
||||
dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
||||
tl.store(dv_ptrs, dv)
|
||||
|
||||
# Write back dK.
|
||||
dk *= sm_scale
|
||||
dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
|
||||
tl.store(dk_ptrs, dk)
|
||||
|
||||
# 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 = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
|
||||
dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
|
||||
do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
|
||||
|
||||
m = tl.load(M + offs_m)
|
||||
m = m[:, 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.
|
||||
num_steps = BLOCK_M2 // MASK_BLOCK_N2
|
||||
dq = _attn_bwd_dq(dq, q, K, V, #
|
||||
do, m, D, #
|
||||
stride_tok, stride_d, #
|
||||
H, N_CTX, #
|
||||
BLOCK_M2, MASK_BLOCK_N2, HEAD_DIM, #
|
||||
start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps, #
|
||||
MASK=True #
|
||||
)
|
||||
end_n -= num_steps * MASK_BLOCK_N2
|
||||
# stage 2
|
||||
num_steps = end_n // BLOCK_N2
|
||||
dq = _attn_bwd_dq(dq, q, K, V, #
|
||||
do, m, D, #
|
||||
stride_tok, stride_d, #
|
||||
H, N_CTX, #
|
||||
BLOCK_M2, BLOCK_N2, HEAD_DIM, #
|
||||
start_m, end_n - num_steps * BLOCK_N2, num_steps, #
|
||||
MASK=False #
|
||||
)
|
||||
# Write back dQ.
|
||||
dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
|
||||
dq *= LN2
|
||||
tl.store(dq_ptrs, dq)
|
||||
|
||||
|
||||
class _attention(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, q, k, v, causal, sm_scale, warp_specialize=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}
|
||||
o = torch.empty_like(q)
|
||||
stage = 3 if causal else 1
|
||||
extra_kern_args = {}
|
||||
# Tuning for AMD target
|
||||
if is_hip():
|
||||
waves_per_eu = 3 if HEAD_DIM_K <= 64 else 2
|
||||
extra_kern_args = {"waves_per_eu": waves_per_eu, "allow_flush_denorm": True}
|
||||
|
||||
M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
|
||||
# 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.shape[0] * q.shape[1] * q.shape[2]
|
||||
|
||||
dummy_block = [1, 1]
|
||||
desc_q = TensorDescriptor(q, shape=[y_dim, 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], strides=[q.shape[2], 1],
|
||||
block_shape=dummy_block)
|
||||
else:
|
||||
desc_v = TensorDescriptor(v, shape=[y_dim, HEAD_DIM_K], strides=[HEAD_DIM_K, 1],
|
||||
block_shape=dummy_block)
|
||||
desc_k = TensorDescriptor(k, shape=[y_dim, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block)
|
||||
desc_o = TensorDescriptor(o, shape=[y_dim, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block)
|
||||
else:
|
||||
desc_q = q
|
||||
desc_v = v
|
||||
desc_k = k
|
||||
desc_o = o
|
||||
|
||||
def alloc_fn(size: int, align: int, _):
|
||||
return torch.empty(size, dtype=torch.int8, device="cuda")
|
||||
|
||||
triton.set_allocator(alloc_fn)
|
||||
|
||||
def grid(META):
|
||||
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
|
||||
_attn_fwd[grid](
|
||||
sm_scale, M, #
|
||||
q.shape[0], q.shape[1], #
|
||||
desc_q, desc_k, desc_v, desc_o, #
|
||||
N_CTX=q.shape[2], #
|
||||
HEAD_DIM=HEAD_DIM_K, #
|
||||
FP8_OUTPUT=q.dtype == torch.float8_e5m2, #
|
||||
STAGE=stage, #
|
||||
warp_specialize=warp_specialize, #
|
||||
IS_HOPPER=is_hopper(), #
|
||||
**extra_kern_args)
|
||||
|
||||
ctx.save_for_backward(q, k, v, o, M)
|
||||
ctx.sm_scale = sm_scale
|
||||
ctx.HEAD_DIM = HEAD_DIM_K
|
||||
ctx.causal = causal
|
||||
return o
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, do):
|
||||
q, k, v, o, M = ctx.saved_tensors
|
||||
assert do.is_contiguous()
|
||||
assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
|
||||
dq = torch.empty_like(q)
|
||||
dk = torch.empty_like(k)
|
||||
dv = torch.empty_like(v)
|
||||
BATCH, N_HEAD, N_CTX = q.shape[:3]
|
||||
PRE_BLOCK = 128
|
||||
NUM_WARPS, NUM_STAGES = 4, 5
|
||||
BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 128, 128, 32
|
||||
BLK_SLICE_FACTOR = 2
|
||||
RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2)
|
||||
arg_k = k
|
||||
arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
|
||||
PRE_BLOCK = 128
|
||||
assert N_CTX % PRE_BLOCK == 0
|
||||
pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
|
||||
delta = torch.empty_like(M)
|
||||
_attn_bwd_preprocess[pre_grid](
|
||||
o, do, #
|
||||
delta, #
|
||||
BATCH, N_HEAD, N_CTX, #
|
||||
BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM #
|
||||
)
|
||||
grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
|
||||
_attn_bwd[grid](
|
||||
q, arg_k, v, ctx.sm_scale, do, dq, dk, dv, #
|
||||
M, delta, #
|
||||
q.stride(0), q.stride(1), q.stride(2), q.stride(3), #
|
||||
N_HEAD, N_CTX, #
|
||||
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, #
|
||||
num_warps=NUM_WARPS, #
|
||||
num_stages=NUM_STAGES #
|
||||
)
|
||||
|
||||
return dq, dk, dv, None, None, None, None
|
||||
|
||||
|
||||
attention = _attention.apply
|
||||
|
||||
TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')
|
||||
|
||||
|
||||
@pytest.mark.parametrize("Z", [1, 4])
|
||||
@pytest.mark.parametrize("H", [2, 48])
|
||||
@pytest.mark.parametrize("N_CTX", [128, 1024, (2 if is_hip() else 4) * 1024])
|
||||
@pytest.mark.parametrize("HEAD_DIM", [64, 128])
|
||||
@pytest.mark.parametrize("causal", [True]) # FIXME: Non-causal tests do not pass at the moment.
|
||||
@pytest.mark.parametrize("warp_specialize", [False, True] if is_blackwell() else [False])
|
||||
@pytest.mark.parametrize("mode", ["fwd", "bwd"])
|
||||
@pytest.mark.parametrize("provider", ["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []))
|
||||
def test_op(Z, H, N_CTX, HEAD_DIM, causal, warp_specialize, mode, provider, dtype=torch.float16):
|
||||
if mode == "fwd" and "fp16" in provider:
|
||||
pytest.skip("Avoid running the forward computation twice.")
|
||||
if mode == "bwd" and "fp8" in provider:
|
||||
pytest.skip("Backward pass with FP8 is not supported.")
|
||||
torch.manual_seed(20)
|
||||
q = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_())
|
||||
k = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_())
|
||||
v = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_())
|
||||
sm_scale = 0.5
|
||||
# reference implementation
|
||||
ref_dtype = dtype
|
||||
if mode == "fwd" and "fp8" in provider:
|
||||
ref_dtype = torch.float32
|
||||
q = q.to(ref_dtype)
|
||||
k = k.to(ref_dtype)
|
||||
v = v.to(ref_dtype)
|
||||
M = torch.tril(torch.ones((N_CTX, N_CTX), device=DEVICE))
|
||||
p = torch.matmul(q, k.transpose(2, 3)) * sm_scale
|
||||
if causal:
|
||||
p[:, :, M == 0] = float("-inf")
|
||||
p = torch.softmax(p.float(), dim=-1)
|
||||
p = p.to(ref_dtype)
|
||||
# p = torch.exp(p)
|
||||
ref_out = torch.matmul(p, v).half()
|
||||
if mode == "bwd":
|
||||
dout = torch.randn_like(q)
|
||||
ref_out.backward(dout)
|
||||
ref_dv, v.grad = v.grad.clone(), None
|
||||
ref_dk, k.grad = k.grad.clone(), None
|
||||
ref_dq, q.grad = q.grad.clone(), None
|
||||
# triton implementation
|
||||
if mode == "fwd" and "fp8" in provider:
|
||||
q = q.to(torch.float8_e5m2)
|
||||
k = k.to(torch.float8_e5m2)
|
||||
v = v.permute(0, 1, 3, 2).contiguous()
|
||||
v = v.permute(0, 1, 3, 2)
|
||||
v = v.to(torch.float8_e5m2)
|
||||
tri_out = attention(q, k, v, causal, sm_scale, warp_specialize).half()
|
||||
if mode == "fwd":
|
||||
atol = 3 if "fp8" in provider else 1e-2
|
||||
torch.testing.assert_close(tri_out, ref_out, atol=atol, rtol=0)
|
||||
return
|
||||
tri_out.backward(dout)
|
||||
tri_dv, v.grad = v.grad.clone(), None
|
||||
tri_dk, k.grad = k.grad.clone(), None
|
||||
tri_dq, q.grad = q.grad.clone(), None
|
||||
# compare
|
||||
torch.testing.assert_close(tri_out, ref_out, atol=1e-2, rtol=0)
|
||||
rtol = 0.0
|
||||
# Relative tolerance workaround for known hardware limitation of CDNA2 GPU.
|
||||
# For details see https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices
|
||||
if torch.version.hip is not None and triton.runtime.driver.active.get_current_target().arch == "gfx90a":
|
||||
rtol = 1e-2
|
||||
torch.testing.assert_close(tri_dv, ref_dv, atol=1e-2, rtol=rtol)
|
||||
torch.testing.assert_close(tri_dk, ref_dk, atol=1e-2, rtol=rtol)
|
||||
torch.testing.assert_close(tri_dq, ref_dq, atol=1e-2, rtol=rtol)
|
||||
|
||||
|
||||
try:
|
||||
from flash_attn.flash_attn_interface import \
|
||||
flash_attn_qkvpacked_func as flash_attn_func
|
||||
HAS_FLASH = True
|
||||
except BaseException:
|
||||
HAS_FLASH = False
|
||||
|
||||
TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')
|
||||
BATCH, N_HEADS = 4, 32
|
||||
# vary seq length for fixed head and batch=4
|
||||
configs = []
|
||||
for HEAD_DIM in [64, 128]:
|
||||
for mode in ["fwd", "bwd"]:
|
||||
for causal in [True, False]:
|
||||
# Enable warpspec for causal fwd on Hopper
|
||||
enable_ws = mode == "fwd" and (is_blackwell() or (is_hopper() and not causal))
|
||||
for warp_specialize in [False, True] if enable_ws else [False]:
|
||||
configs.append(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["N_CTX"],
|
||||
x_vals=[2**i for i in range(10, 15)],
|
||||
line_arg="provider",
|
||||
line_vals=["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []) +
|
||||
(["flash"] if HAS_FLASH else []),
|
||||
line_names=["Triton [FP16]"] + (["Triton [FP8]"] if TORCH_HAS_FP8 else []) +
|
||||
(["Flash-2"] if HAS_FLASH else []),
|
||||
styles=[("red", "-"), ("blue", "-"), ("green", "-")],
|
||||
ylabel="TFLOPS",
|
||||
plot_name=
|
||||
f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}-causal={causal}-warp_specialize={warp_specialize}",
|
||||
args={
|
||||
"H": N_HEADS,
|
||||
"BATCH": BATCH,
|
||||
"HEAD_DIM": HEAD_DIM,
|
||||
"mode": mode,
|
||||
"causal": causal,
|
||||
"warp_specialize": warp_specialize,
|
||||
},
|
||||
))
|
||||
|
||||
|
||||
@triton.testing.perf_report(configs)
|
||||
def bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, causal, warp_specialize, mode, provider, device=DEVICE):
|
||||
assert mode in ["fwd", "bwd"]
|
||||
dtype = torch.float16
|
||||
if "triton" in provider:
|
||||
q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
|
||||
k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
|
||||
v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
|
||||
if mode == "fwd" and "fp8" in provider:
|
||||
q = q.to(torch.float8_e5m2)
|
||||
k = k.to(torch.float8_e5m2)
|
||||
v = v.permute(0, 1, 3, 2).contiguous()
|
||||
v = v.permute(0, 1, 3, 2)
|
||||
v = v.to(torch.float8_e5m2)
|
||||
sm_scale = 1.3
|
||||
fn = lambda: attention(q, k, v, causal, sm_scale, warp_specialize)
|
||||
if mode == "bwd":
|
||||
o = fn()
|
||||
do = torch.randn_like(o)
|
||||
fn = lambda: o.backward(do, retain_graph=True)
|
||||
ms = triton.testing.do_bench(fn)
|
||||
|
||||
if provider == "flash":
|
||||
qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
|
||||
fn = lambda: flash_attn_func(qkv, causal=causal)
|
||||
if mode == "bwd":
|
||||
o = fn()
|
||||
do = torch.randn_like(o)
|
||||
fn = lambda: o.backward(do, retain_graph=True)
|
||||
ms = triton.testing.do_bench(fn)
|
||||
flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM
|
||||
total_flops = 2 * flops_per_matmul
|
||||
if causal:
|
||||
total_flops *= 0.5
|
||||
if mode == "bwd":
|
||||
total_flops *= 2.5 # 2.0(bwd) + 0.5(recompute)
|
||||
return total_flops * 1e-12 / (ms * 1e-3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# only works on post-Ampere GPUs right now
|
||||
bench_flash_attention.run(save_path=".", print_data=True)
|
||||
@@ -9,7 +9,7 @@ Uses FlashInfer API for 4-bit quantization producing bf16 matrices.
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from flashinfer import SfLayout, mm_fp4, nvfp4_quantize
|
||||
from flashinfer import SfLayout, mm_fp4, nvfp4_quantize, e2m1_and_ufp8sf_scale_to_float
|
||||
import unittest
|
||||
from typing import Tuple, Dict, Any
|
||||
|
||||
@@ -20,6 +20,16 @@ import triton.profiler as proton
|
||||
from triton.tools.tensor_descriptor import TensorDescriptor
|
||||
from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
|
||||
|
||||
def get_fake_quant(x: torch.Tensor):
|
||||
# TODO: dtype needs to be float32?
|
||||
orig_shape = x.shape
|
||||
orig_dtype = x.dtype
|
||||
device = x.device
|
||||
x = x.view(-1, x.shape[-1])
|
||||
x_global_sf = (448 * 6) / x.float().abs().nan_to_num().max()
|
||||
x_fp4, x_scale = nvfp4_quantize(x, x_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
|
||||
x_dequant = e2m1_and_ufp8sf_scale_to_float(x_fp4, x_scale, 1 / x_global_sf) # Note: this function does dequantization on the cpu
|
||||
return x_dequant.view(orig_shape).to(orig_dtype).to(device)
|
||||
|
||||
|
||||
def generate_fake_quant(a_f32: torch.Tensor,
|
||||
@@ -344,10 +354,108 @@ def main():
|
||||
# Run comprehensive test
|
||||
print("\nRunning comprehensive tests...")
|
||||
results = tester.run_comprehensive_test()
|
||||
|
||||
|
||||
|
||||
|
||||
def test_get_fake_quant():
|
||||
"""Test the get_fake_quant function with various inputs."""
|
||||
print("Testing get_fake_quant function...")
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA not available, skipping test")
|
||||
return
|
||||
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
|
||||
# Test 1: Basic functionality - 2D tensor
|
||||
print("\nTest 1: Basic 2D tensor")
|
||||
x = torch.randn(128, 64, device=device, dtype=dtype)
|
||||
x_fq = get_fake_quant(x)
|
||||
assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}"
|
||||
assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}"
|
||||
assert x_fq.device == x.device, f"Device mismatch: {x_fq.device} vs {x.device}"
|
||||
max_diff = (x - x_fq).abs().max().item()
|
||||
print(f" Max difference: {max_diff:.6f}")
|
||||
assert max_diff < 1.0, f"Difference too large: {max_diff}"
|
||||
|
||||
# Test 2: 3D tensor (B, H, D)
|
||||
print("\nTest 2: 3D tensor [B, H, D]")
|
||||
x = torch.randn(2, 8, 128, device=device, dtype=dtype)
|
||||
x_fq = get_fake_quant(x)
|
||||
assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}"
|
||||
assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}"
|
||||
max_diff = (x - x_fq).abs().max().item()
|
||||
print(f" Max difference: {max_diff:.6f}")
|
||||
assert max_diff < 1.0, f"Difference too large: {max_diff}"
|
||||
|
||||
# Test 3: 4D tensor (B, H, L, D)
|
||||
print("\nTest 3: 4D tensor [B, H, L, D]")
|
||||
x = torch.randn(1, 8, 256, 128, device=device, dtype=dtype)
|
||||
x_fq = get_fake_quant(x)
|
||||
assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}"
|
||||
assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}"
|
||||
max_diff = (x - x_fq).abs().max().item()
|
||||
print(f" Max difference: {max_diff:.6f}")
|
||||
assert max_diff < 1.0, f"Difference too large: {max_diff}"
|
||||
|
||||
# Test 4: Small values
|
||||
print("\nTest 4: Small values")
|
||||
x = torch.randn(128, 64, device=device, dtype=dtype) * 0.01
|
||||
x_fq = get_fake_quant(x)
|
||||
assert x_fq.shape == x.shape
|
||||
max_diff = (x - x_fq).abs().max().item()
|
||||
print(f" Max difference: {max_diff:.6f}")
|
||||
|
||||
# Test 5: Large values
|
||||
print("\nTest 5: Large values")
|
||||
x = torch.randn(128, 64, device=device, dtype=dtype) * 10.0
|
||||
x_fq = get_fake_quant(x)
|
||||
assert x_fq.shape == x.shape
|
||||
max_diff = (x - x_fq).abs().max().item()
|
||||
print(f" Max difference: {max_diff:.6f}")
|
||||
|
||||
# Test 6: Zeros
|
||||
print("\nTest 6: Zeros")
|
||||
x = torch.zeros(128, 64, device=device, dtype=dtype)
|
||||
x_fq = get_fake_quant(x)
|
||||
assert x_fq.shape == x.shape
|
||||
max_diff = (x - x_fq).abs().max().item()
|
||||
print(f" Max difference: {max_diff:.6f}")
|
||||
assert max_diff < 1e-5, f"Zeros should remain zeros: {max_diff}"
|
||||
|
||||
# Test 7: Different dtypes
|
||||
print("\nTest 7: Float16 dtype")
|
||||
x = torch.randn(128, 64, device=device, dtype=torch.float16)
|
||||
x_fq = get_fake_quant(x)
|
||||
assert x_fq.shape == x.shape
|
||||
assert x_fq.dtype == x.dtype
|
||||
max_diff = (x - x_fq).abs().max().item()
|
||||
print(f" Max difference: {max_diff:.6f}")
|
||||
|
||||
# Test 8: Idempotency check (quantizing twice should be similar)
|
||||
print("\nTest 8: Idempotency")
|
||||
x = torch.randn(128, 64, device=device, dtype=dtype)
|
||||
x_fq1 = get_fake_quant(x)
|
||||
x_fq2 = get_fake_quant(x_fq1)
|
||||
# Second quantization should be close to first
|
||||
max_diff = (x_fq1 - x_fq2).abs().max().item()
|
||||
print(f" Max difference between two quantizations: {max_diff:.6f}")
|
||||
|
||||
# Test 9: Requires grad
|
||||
print("\nTest 9: Gradient flow")
|
||||
x = torch.randn(128, 64, device=device, dtype=dtype, requires_grad=True)
|
||||
x_fq = get_fake_quant(x)
|
||||
# Note: get_fake_quant uses no_grad internally, so x_fq won't have grad
|
||||
# But we can check it doesn't crash
|
||||
assert x_fq.shape == x.shape
|
||||
print(" Gradient test passed (no crash)")
|
||||
|
||||
print("\n✓ All tests passed!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
import sys
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "test_fake_quant":
|
||||
test_get_fake_quant()
|
||||
else:
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user