import math from functools import partial, wraps import os import torch import timeit import statistics import torch_xla import torch_xla.core.xla_model as xm from torch_xla.amp import autocast def wrap_with_mark_step(func): @wraps(func) def wrapper(*args, **kwargs): res = func(*args, **kwargs) xm.mark_step() return res return wrapper @wrap_with_mark_step def sdpa(query, key, value): return torch.nn.functional.scaled_dot_product_attention(query, key, value) @wrap_with_mark_step def standard_attention(query, key, value): # from sdpa's doc implementation (it's pretty straightforward) scale_factor = 1 / math.sqrt(query.size(-1)) attn_weight = query @ key.transpose(-2, -1) * scale_factor attn_weight = torch.softmax(attn_weight, dim=-1) return attn_weight @ value def time_attention(attention_fn, query, key, value): # Warm-up for _ in range(100): attention_fn(query, key, value) # Run timings with repeats repeats = 10 number = 100 times = timeit.repeat(partial(attention_fn, query, key, value), repeat=repeats, number=number) mean_time = statistics.mean(times) / number * 1000 # ms per call std_dev = statistics.stdev(times) / number * 1000 # ms per call print(f" {attention_fn.__name__} average time: {mean_time:.3f} ms per call, dev: {std_dev:.3f} ms per call") if __name__ == '__main__': os.environ["XLA_DISABLE_FUNCTIONALIZATION"] = "1" batch_size = 32 seq_len = 256 num_heads = 16 head_dim = 64 device = xm.xla_device() # Generate input tensors with correct shape: (batch_size, num_heads, seq_len, head_dim) query = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device) key = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device) value = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device) assert torch.allclose(sdpa(query, key, value), standard_attention(query, key, value), rtol=1e-02, atol=1e-02) print("Both attention functions produce about the same output, test time performance now") print("No mixed precision:") time_attention(sdpa, query, key, value) time_attention(standard_attention, query, key, value) print("With mixed precision:") with autocast(xm.xla_device(), enabled=True, dtype=torch.bfloat16): time_attention(sdpa, query, key, value) time_attention(standard_attention, query, key, value) print(f"torch version: {torch.__version__}, torch_xla version: {torch_xla.__version__}")