Utilities#
Helper functions for visualization, benchmarking, and profiling.
Visualization#
attn_gym.utils.visualize_attention_scores(query, key, score_mod=None, mask_mod=None, device='cuda', name='attention_scores', path=None, batch_idx=0, head_idx=0, scale=None)
#
Generate and save a visualization of attention scores.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
query
|
Tensor
|
Query tensor of shape (batch_size, num_heads, seq_len_q, head_dim). |
required |
key
|
Tensor
|
Key tensor of shape (batch_size, num_heads, seq_len_k, head_dim). |
required |
score_mod
|
Optional[Callable]
|
If this is set this will take precedence over the mask_mod. |
None
|
mask_mod
|
Optional[Callable]
|
The mask_mod function used to create block_mask |
None
|
device
|
str
|
Device to run computations on (default: "cuda"). |
'cuda'
|
name
|
str
|
Base name for the file and title (default: 'attention_scores'). |
'attention_scores'
|
path
|
Path
|
Path to save the visualization. If None, will be saved to the current working directory. |
None
|
batch_idx
|
int
|
Index of the batch to visualize (default: 0). |
0
|
head_idx
|
int
|
Index of the head to visualize (default: 0). |
0
|
scale
|
float
|
Scale factor to apply to the attention scores. If None, will be set to 1 / sqrt(head_dim). |
None
|
Returns:
| Type | Description |
|---|---|
None
|
None |
attn_gym.utils.create_score_mod(query, key, score_mod, mask_mod, device='cuda', _compile=False, scale=None, batch_idx=0, head_idx=0)
#
Benchmarking#
attn_gym.utils.benchmark_cuda_function_in_microseconds(func, *args, **kwargs)
#
Thin wrapper around do_bench_using_profiling
attn_gym.utils.calculate_tflops(batch, heads, seq_q, seq_kv, head_dim, time_ms, sparsity=0.0)
#
Calculate TFLOPs for attention forward pass.
Profiling#
attn_gym.utils.cuda_kernel_profiler(kernel_pattern='flash_attncute')
#
Context manager that profiles CUDA kernels and checks for a pattern.
Usage
with cuda_kernel_profiler("flash_attncute") as result: flex_attention(...) print(result["found"]) # True if flash kernel was called print(result["kernel_names"]) # List of all CUDA kernel names
Misc#
attn_gym.utils.get_flash_block_size(device='cuda')
#
Get block size for Flash backend based on GPU compute capability.
On SM100+ (Blackwell): Q block must be 256, KV block is 128. On SM80/SM90: Both use default 128.