From ebdb160af9d95d717bbb2ddaf473975bb1f41858 Mon Sep 17 00:00:00 2001 From: Matthew Noto Date: Mon, 10 Nov 2025 00:25:12 -0800 Subject: [PATCH] checkpoint --- fastvideo/attention/backends/sage_attn3.py | 82 ++- nvfp4_utils.py | 357 ++++++++++ qat_attn.py | 762 +++++++++++++++++++++ test_quantization_precision.py | 116 +++- 4 files changed, 1304 insertions(+), 13 deletions(-) create mode 100644 nvfp4_utils.py create mode 100644 qat_attn.py diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index 495e50228..dc1fb4ea2 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/fastvideo/attention/backends/sage_attn3.py @@ -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) diff --git a/nvfp4_utils.py b/nvfp4_utils.py new file mode 100644 index 000000000..0c0cb6a4b --- /dev/null +++ b/nvfp4_utils.py @@ -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) \ No newline at end of file diff --git a/qat_attn.py b/qat_attn.py new file mode 100644 index 000000000..d7a44f8a5 --- /dev/null +++ b/qat_attn.py @@ -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) \ No newline at end of file diff --git a/test_quantization_precision.py b/test_quantization_precision.py index 023d40a19..c72a01ab3 100644 --- a/test_quantization_precision.py +++ b/test_quantization_precision.py @@ -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()