fix: fp8 ref matches kernel

flash_attn/flash_attn_triton_amd/fwd_prefill.py

@@ -183,10 +183,10 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
         if IS_FP8:
             p_scale = 1
             p_scaled = (p / p_scale)
-            acc += tl.dot(p.to(v.type.element_ty), (v*v_scale).to(v.type.element_ty)) # if you want to use p_scaled: tl.dot(p_scaled.to(v.type.element_ty), v.to(v.type.element_ty)) * v_scale * p_scale
+            acc += tl.dot(p_scaled.to(v.type.element_ty), v.to(v.type.element_ty)).to(tl.float32) * v_scale * p_scale # if you want to use p_scaled: tl.dot(p_scaled.to(v.type.element_ty), v.to(v.type.element_ty)) * v_scale * p_scale
         else:
-            acc += tl.dot(p.to(v.type.element_ty), v).to(tl.float32) # NOTE: acc += tl.dot(p.to(tl.float16), v.to(tl.float16)) PASSES
-
+            # NOTE: if you make the below operation tl.float16 + set FLASH_ATTENTION_TRITON_AMD_REMOVE_QUANT_SCALE=1. It passes. --> acc += tl.dot(p.to(tl.float16), v.to(tl.float16)) PASSES
+            acc += tl.dot(p.to(v.type.element_ty), v).to(tl.float32)
 
         k_ptrs += BLOCK_N * stride_kn
         v_ptrs += BLOCK_N * stride_vk

flash_attn/flash_attn_triton_amd/fwd_ref.py

@@ -1,10 +1,23 @@
 import torch
 import math
-from .utils import DEBUG
+from .utils import create_scale_tensors, check_is_fp8, DEBUG
 
 DEBUG_CORE = False
 
-def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox_seed, philox_offset, use_exp2):
+def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, layout, dropout_p, philox_seed, philox_offset, use_exp2):
+    is_fp8 = check_is_fp8(q)
+    if is_fp8:
+        # if qkv are fp8, then find scaling factor for quantization
+        q_scale, k_scale, v_scale = create_scale_tensors(q, k, v, SCALE_PER_HEAD=True, layout=layout) # TODO: if SCALE_PER_HEAD: within the kernel itself just compute qkv_scale = tl.max(q or k or v)
+        q_scale_stride_z = q_scale.stride(0)
+        kv_scale_stride_z = k_scale.stride(0)
+
+        # scale qkv tensors if FP8
+        q = q / q_scale
+        k = k / k_scale
+        v = v / v_scale
+    else:
+        q_scale = k_scale = v_scale = 1
     if DEBUG_CORE:
         print()
         print("attention_forward_core_ref_impl")
@@ -17,14 +30,22 @@ def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox
         print("philox_seed:", philox_seed)
         print("philox_offset:", philox_offset)
         print("use_exp2:", use_exp2)
+        print('layout:', layout)
+        print('is_fp8:', is_fp8)
+        print('q_scale:', q_scale)
+        print('k_scale:', k_scale)
+        print('v_scale:', v_scale)
 
     # cast to float32
     q = q.to(torch.float32)
     k = k.to(torch.float32)
     v = v.to(torch.float32)
 
     # Compute attention scores
-    attention_scores = torch.matmul(q, k.transpose(-2, -1))
+    if is_fp8:
+        attention_scores = torch.matmul(q.to(torch.float32), k.transpose(-2, -1).to(torch.float32)) * q_scale * v_scale
+    else:
+        attention_scores = torch.matmul(q.to(torch.float32), k.transpose(-2, -1).to(torch.float32))
     if DEBUG_CORE:
         print("attention_scores:", attention_scores, attention_scores.shape)
 
@@ -129,7 +150,10 @@ def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, dropout_p, philox
         print("softmax_lse:", softmax_lse, softmax_lse.shape)
 
     # Compute output
-    o = torch.matmul(p, v)
+    if is_fp8:
+        o = torch.matmul(p, v.to(torch.float32)) * v_scale
+    else:
+        o = torch.matmul(p, v)
     if DEBUG_CORE:
         print("o:", o, o.shape)
 
@@ -176,7 +200,7 @@ def attention_vanilla_forward_pytorch_ref_impl(q, k, v, sm_scale, causal, layout
 
     # Call the core attention function
     o, softmax_lse, sd_mask = attention_forward_core_ref_impl(
-        q, k, v, sm_scale, causal, dropout_p, philox_seed, philox_offset, use_exp2
+        q, k, v, sm_scale, causal, layout, dropout_p, philox_seed, philox_offset, use_exp2
     )
 
     if group_size != 1:
@@ -278,7 +302,7 @@ def attention_varlen_forward_pytorch_ref_impl(
             v_i = v_i.reshape(nheads_k, seqlen_k, head_dim)
 
         # Call the core attention function for this sequence
-        o_i, softmax_lse_i, sd_mask_i = attention_forward_core_ref_impl(q_i, k_i, v_i, sm_scale, causal, dropout_p, philox_seed, philox_offset, use_exp2)
+        o_i, softmax_lse_i, sd_mask_i = attention_forward_core_ref_impl(q_i, k_i, v_i, sm_scale, causal, layout, dropout_p, philox_seed, philox_offset, use_exp2)
 
         # Reshape outputs back to original dimensions
         if group_size != 1:

flash_attn/flash_attn_triton_amd/test.py

@@ -10,7 +10,7 @@
 from .fwd_decode import dequantize_kv_fp16, quantize_kv_int4
 
 # defailt fp16 tolerance is ATOL, RTOL = 1e-5, 1e-3. See table https://pytorch.org/docs/stable/testing.html
-ATOL, RTOL = 1e-2, 0 # old standard. maybe to lose. 
+ATOL, RTOL = 1e-2, 1e-2 # old standard. maybe to lose. 
 # ATOL, RTOL = 1e-3, 1e-3  # catchs fa mismatch issues
 # ATOL, RTOL = 1e-4, 1e-3 # to strict. there will be small diffs
 # ATOL, RTOL = 1e-5, 1e-3 # # default fp16. there will be small diffs