Skip to content

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

Distributed#

Copy-engine transfers into symmetric memory, for groups inside one NVLink domain (one node or an NVL72 rack). Payload copies do not occupy SMs, so they overlap attention kernels well; they still share memory and NVLink bandwidth.

attn_gym.distributed.CopyEngineAllGather #

All-gather into a persistent symmetric buffer, written directly in the consumer's layout.

Each rank contributes num_parts tensors of part_shape (for example K and V). The output has shape (num_parts, *part_shape) with part_shape[dim] multiplied by the world size, and rank r's shard of every part lands at offset r * part_shape[dim] along dim: there is no gathered-then-reordered intermediate. Every rank pushes its shard into every peer's buffer (one contiguous copy per row in front of dim), so the side stream does P * num_parts * prod(part_shape[:dim]) copies.

start(parts) issues the pushes; wait(parts) returns the gathered buffer. When any part requires grad, wait is differentiable: its backward reduce-scatters the gradient over NCCL. Every start first tells peers this rank's buffer is free, then waits for each peer's free signal before writing into it. The signal is queued behind everything already on the current stream (including any backward that reads the previous gather), so one gather may be outstanding at a time, and every call issues the same signals, so a captured start replays correctly whether or not the gather ran eagerly first.

reduce_scatter(grad) #

Sum grad (shaped like the output) over ranks; return this rank's shard per part.

attn_gym.distributed.CopyEngineChannel #

Double-buffered pushes to dst and receives from src (ranks within group).

Each rank owns two receive slots of shape (num_parts, *part_shape) that src writes alternately, one per step. A receiver frees a slot only after its compute (and any onward push from that slot) is queued, so a slot is never overwritten while it is being read.

Per step i: send(parts, i) pushes into the slot dst reads at i + 1, recv(i) waits for and returns this rank's slot for i, and release(i) hands it back. Call join() before work that must see every pending push finished (for example the end of a CUDA graph capture).

True when every rank of group shares one NVLink domain, so copy-engine transfers work.

Collective: every rank of group must call it.

Identify the NVLink domain of the current GPU.

On GPUs attached to an NVLink fabric (GB200/GB300 NVL72 racks) this is the fabric cluster UUID and clique reported by NVML (needs nvidia-ml-py). Otherwise it falls back to the host name, which assumes the GPUs of one host are NVLink peers (true on HGX/DGX-style systems, not on PCIe-only boxes).

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.