97 lines
3.4 KiB
Python
97 lines
3.4 KiB
Python
import torch
|
|
from flex_sta_ref import get_sliding_tile_attention_mask
|
|
from st_attn import sliding_tile_attention
|
|
from torch.nn.attention.flex_attention import flex_attention
|
|
# from flash_attn_interface import flash_attn_func
|
|
from tqdm import tqdm
|
|
|
|
flex_attention = torch.compile(flex_attention, dynamic=False)
|
|
|
|
|
|
def flex_test(Q, K, V, kernel_size):
|
|
mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (36, 48, 48), 39, 'cuda', 0)
|
|
output = flex_attention(Q, K, V, block_mask=mask)
|
|
|
|
return output
|
|
|
|
|
|
def h100_fwd_kernel_test(Q, K, V, kernel_size):
|
|
o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 39, False)
|
|
return o
|
|
|
|
|
|
def generate_tensor(shape, mean, std, dtype, device):
|
|
tensor = torch.randn(shape, dtype=dtype, device=device)
|
|
|
|
magnitude = torch.norm(tensor, dim=-1, keepdim=True)
|
|
scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
|
|
|
|
return scaled_tensor.contiguous()
|
|
|
|
|
|
def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mode='all'):
|
|
results = {
|
|
'TK vs FLEX': {
|
|
'sum_diff': 0,
|
|
'sum_abs': 0,
|
|
'max_diff': 0
|
|
},
|
|
}
|
|
kernel_size_ls = [(6, 1, 6), (6, 6, 1)]
|
|
from tqdm import tqdm
|
|
for kernel_size in tqdm(kernel_size_ls):
|
|
for _ in range(num_iterations):
|
|
torch.manual_seed(0)
|
|
|
|
Q = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
|
|
K = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
|
|
V = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
|
|
tk_o = h100_fwd_kernel_test(Q, K, V, kernel_size)
|
|
pt_o = flex_test(Q, K, V, kernel_size)
|
|
|
|
diff = pt_o - tk_o
|
|
abs_diff = torch.abs(diff)
|
|
results['TK vs FLEX']['sum_diff'] += torch.sum(abs_diff).item()
|
|
results['TK vs FLEX']['max_diff'] = max(results['TK vs FLEX']['max_diff'], torch.max(abs_diff).item())
|
|
|
|
torch.cuda.empty_cache()
|
|
print("kernel_size", kernel_size)
|
|
print("max_diff", torch.max(abs_diff).item())
|
|
print(
|
|
"avg_diff",
|
|
torch.sum(abs_diff).item() / (b * h * n * d *
|
|
(1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4)))
|
|
|
|
total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else
|
|
3 if error_mode == 'backward' else 4) * len(kernel_size_ls)
|
|
for name, data in results.items():
|
|
avg_diff = data['sum_diff'] / total_elements
|
|
max_diff = data['max_diff']
|
|
results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
|
|
|
|
return results
|
|
|
|
|
|
def generate_error_graphs(b, h, d, causal, mean, std, error_mode='all'):
|
|
seq_lengths = [82944]
|
|
|
|
tk_avg_errors, tk_max_errors = [], []
|
|
|
|
for n in tqdm(seq_lengths, desc="Generating error data"):
|
|
results = check_correctness(b, h, n, d, causal, mean, std, error_mode=error_mode)
|
|
|
|
tk_avg_errors.append(results['TK vs FLEX']['avg_diff'])
|
|
tk_max_errors.append(results['TK vs FLEX']['max_diff'])
|
|
|
|
|
|
# Example usage
|
|
b, h, d = 2, 24, 128
|
|
causal = False
|
|
mean = 1e-1
|
|
std = 10
|
|
|
|
for mode in ['output']:
|
|
generate_error_graphs(b, h, d, causal, mean, std, error_mode=mode)
|
|
|
|
print("Error graphs generated and saved for all modes.")
|