Using dummy/padded attention slices with `flex_flash_attn_func` for `torch.compile` static shapes

[keunhong]

Hi, thanks for releasing MagiAttention -- it’s a very nice!

I have a question about how to make sure the attention slices have a stable shape to avoid torch.compile recompilations when the attention pattern is dynamic per batch:

• Different examples have different numbers of "attention slices" (flex ranges).
• These are stored in q_ranges, k_ranges, and attn_type_map tensors as documented.
• The total number of tokens (shapes of qkv) are fixed, but the number of ranges may vary per batch.

WIth torch.compile, changing the shape of q_ranges etc. will generally cause a new graph to be compiled. To avoid excessive recompilations, I'd like to keep the shapes of these tensors fixed and only vary their values.

This naturally leads to the idea of "dummy" or "padded" ranges to signal to FFA that those ranges are unused.

Roughly I'd like to do something like

R_max = 16  # Max number of flex slices per batch.

def make_padded_ranges(raw_q_ranges, raw_k_ranges, raw_types, total_tokens):
    num_real = raw_q_ranges.shape[0]

    q_ranges = torch.zeros(
        (R_max, 2),
        dtype=raw_q_ranges.dtype,
        device=raw_q_ranges.device,
    )
    k_ranges = torch.zeros(
        (R_max, 2),
        dtype=raw_k_ranges.dtype,
        device=raw_k_ranges.device,
    )
    attn_type_map = torch.zeros(
        (R_max,),
        dtype=raw_types.dtype,
        device=raw_types.device,
    )

    # Fill the first `num_real` rows with “real” slices.
    q_ranges[:num_real] = raw_q_ranges
    k_ranges[:num_real] = raw_k_ranges
    attn_type_map[:num_real] = raw_types

    # The remaining rows are “dummy” slices.
    # Question: what is the safe way to define these?
    return q_ranges, k_ranges, attn_type_map

then in a compiled function call

@torch.compile
def attn_step(q, k, v, q_ranges, k_ranges, attn_type_map):
    return flex_flash_attn_func(
        q,
        k,
        v,
        q_ranges=q_ranges,
        k_ranges=k_ranges,
        attn_type_map=attn_type_map,
        # other args kept static,
    )

From torch.compile’s perspective, the argument shapes are now static, so it won’t recompile just because a batch has more or fewer ranges.

Question:
Is there any notion of “dummy” / no-op ranges for flex_flash_attn_func? For example zero-length ranges (start = end) or padded ranges that are guaranteed to be ignored.

[keunhong]

It does seem like dummy ranges with start == end works as expected, though this behavior isn't documented.

import torch
from magi_attention.api import flex_flash_attn_func


def max_rel_diff(a: torch.Tensor, b: torch.Tensor) -> float:
    eps = 1e-6
    num = (a - b).abs()
    denom = torch.maximum(a.abs(), b.abs()).clamp_min(eps)
    return float((num / denom).max().item())


@torch.compile()
def attn_fn(q, k, v, total_seqlen, q_ranges, k_ranges, attn_type_map):
    out, _ = flex_flash_attn_func(
        q=q,
        k=k,
        v=v,
        max_seqlen_q=total_seqlen,
        max_seqlen_k=total_seqlen,
        q_ranges=q_ranges,
        k_ranges=k_ranges,
        attn_type_map=attn_type_map,
        softmax_scale=None,
        softcap=0,
    )
    return out


device = "cuda"
dtype = torch.bfloat16

total_seqlen = 2048
num_heads_q = 8
num_heads_kv = 2
head_dim = 64
max_num_slices = 20  # padded number of ranges

torch.manual_seed(0)
q = torch.randn(total_seqlen, num_heads_q, head_dim, dtype=dtype, device=device, requires_grad=True)
k = torch.randn(total_seqlen, num_heads_kv, head_dim, dtype=dtype, device=device, requires_grad=True)
v = torch.randn(total_seqlen, num_heads_kv, head_dim, dtype=dtype, device=device, requires_grad=True)
do = torch.randn_like(q)

# Real mask: 2 slices, same as docs.
half = total_seqlen // 2
q_ranges = torch.tensor(
    [[0, half], [half, total_seqlen]],
    dtype=torch.int32,
    device=device,
)
k_ranges = torch.tensor(
    [[0, half], [0, total_seqlen]],
    dtype=torch.int32,
    device=device,
)
# 0 = FULL, 1 = CAUSAL.
attn_type_map = torch.tensor([0, 1], dtype=torch.int32, device=device)

# Baseline call.
out0 = attn_fn(
    q=q,
    k=k,
    v=v,
    total_seqlen=total_seqlen,
    q_ranges=q_ranges,
    k_ranges=k_ranges,
    attn_type_map=attn_type_map,
)
out0.backward(do)
dq0, dk0, dv0 = q.grad.clone(), k.grad.clone(), v.grad.clone()

q.grad = k.grad = v.grad = None

# Build padded ranges of fixed shape.
num_real = q_ranges.shape[0]
padded_q = torch.zeros((max_num_slices, 2), dtype=torch.int32, device=device)
padded_k = torch.zeros((max_num_slices, 2), dtype=torch.int32, device=device)
padded_t = torch.zeros((max_num_slices,), dtype=torch.int32, device=device)
padded_q[:num_real] = q_ranges
padded_k[:num_real] = k_ranges
padded_t[:num_real] = attn_type_map

# Padded call.
out1 = attn_fn(
    q=q,
    k=k,
    v=v,
    total_seqlen=total_seqlen,
    q_ranges=padded_q,
    k_ranges=padded_k,
    attn_type_map=padded_t,
)
out1.backward(do)
dq1, dk1, dv1 = q.grad, k.grad, v.grad

print("out  max rel diff:", max_rel_diff(out0, out1))
print("dq   max rel diff:", max_rel_diff(dq0, dq1))
print("dk   max rel diff:", max_rel_diff(dk0, dk1))
print("dv   max rel diff:", max_rel_diff(dv0, dv1))

out  max rel diff: 0.0
dq   max rel diff: 0.007232666015625
dk   max rel diff: 0.0067138671875
dv   max rel diff: 0.007293701171875

[hanwen-sun]

Hi, thanks for using magiattention.
You are correct that the shape of qkrange in MagiAttention can indeed change at every step, which triggers frequent recompilations in torch.compile.
No there are two approaches to address this issue:

• padding: As you suggested, You can pad qkrange and attn_map to a pre-defined maximum size at each step.

  • Mechanism: MagiAttention is designed to automatically ignore any qkrange where start == end for either the Q-range or K-range. You can visualize this as a rectangle in the attention map: if either the length or the width is zero, the effective area is zero, and the kernel skips it. You can refer to our blog for more details: https://sandai-org.github.io/MagiAttention/blog/.
  • Performance: The impact on kernel performance is negligible. While there is a very minor overhead in the tile_scheduler to iterate over and skip these empty ranges, it introduces no actual computation overhead.
  • Trade-off: The main downside is that you need to know the maximum possible size in advance, and there is some extra overhead associated with the padding operations on the Python side.

• Enable Dynamic Shapes: You can enable dynamic shapes by setting torch.compile(..., dynamic=True)

  • Mechanism: This instructs TorchDynamo to relax its checks on input parameter shapes, preventing recompilation when shape change.
  • Trade-off: Please use this with caution, as dynamic=True can sometimes limit the compiler's ability to perform specific optimizations, potentially resulting in lower performance compared to static compilation.

We are actively investigating more efficient methods to handle dynamic shapes with lower overhead. We will inform you immediately if we implement a better solution.