[PERFORMANCE] Top‑k windowed mask generation dominates end‑to‑end latency for large sequences

[LoserCheems]

Performance issue description

When enabling top‑k windowing (key_len > window_size), the mask generation step (topk + scatter) often dominates end‑to‑end latency, outweighing the attention kernel speedup. This is most pronounced for larger Lq/Lk and for bias layouts with more “head-like” planes (e.g., (B,H) vs (B,1)). With small windows, attention gets slightly faster, but the Python-side mask materialization increases latency significantly, making E2E slower.

Current performance metrics

• Device: CUDA
• Dtype: float16
• B=2, H=8, KVH=4, D=64, Lq=4096, Lk=4096
• Numbers are per-iteration latency (ms). total = mask + attn.

bias_kind=B,H, mask_kind=None

case	mask	attn	total
window=None	0.006	0.826	0.832
window=Lk(4096)	0.002	0.860	0.862
window=256	3.625	0.720	4.345
window=512	4.031	0.749	4.780
window=1024	4.899	0.849	5.748
window=2048	6.861	1.002	7.863

bias_kind=B,H, mask_kind=2D

case	mask	attn	total
window=None	0.003	0.331	0.334
window=Lk(4096)	0.004	0.329	0.333
window=256	6.136	0.657	6.793
window=512	6.576	0.675	7.251
window=1024	7.477	0.660	8.137
window=2048	9.475	0.656	10.131

bias_kind=B,H, mask_kind=4D-like-bias

case	mask	attn	total
window=None	0.004	0.652	0.656
window=Lk(4096)	0.001	0.670	0.671
window=256	7.202	0.670	7.872
window=512	7.656	0.654	8.310
window=1024	8.533	0.656	9.189
window=2048	10.518	0.679	11.197

bias_kind=B,KVH, mask_kind=None

case	mask	attn	total
window=None	0.004	0.508	0.512
window=Lk(4096)	0.002	0.512	0.514
window=256	1.743	0.426	2.169
window=512	1.965	0.455	2.420
window=1024	2.407	0.517	2.924
window=2048	3.402	0.625	4.027

bias_kind=B,KVH, mask_kind=4D-like-bias

case	mask	attn	total
window=None	0.002	0.425	0.427
window=Lk(4096)	0.002	0.419	0.421
window=256	3.570	0.436	4.006
window=512	3.781	0.423	4.204
window=1024	4.236	0.414	4.650
window=2048	5.217	0.418	5.635

bias_kind=1,H, mask_kind=None

case	mask	attn	total
window=None	0.022	0.516	0.538
window=Lk(4096)	0.001	0.505	0.506
window=256	1.725	0.441	2.166
window=512	1.963	0.462	2.425
window=1024	2.388	0.496	2.884
window=2048	3.406	0.594	4.000

bias_kind=B,1, mask_kind=None

case	mask	attn	total
window=None	0.030	0.504	0.534
window=Lk(4096)	0.003	0.510	0.513
window=256	0.410	0.307	0.717
window=512	0.463	0.326	0.789
window=1024	0.605	0.384	0.989
window=2048	0.848	0.488	1.336

Observations

• The attention kernel time drops modestly for small windows, but mask creation quickly dominates as window_size grows (and as head-like planes increase).
• B,1 is the cheapest case for mask creation; B,H with a 4D mask is the most expensive.
• Baseline (no windowing) is consistently faster end‑to‑end in these runs.

Expected performance

• Expect enabling windowing to reduce end‑to‑end latency when ws << Lk (e.g., ws ≤ Lk/16), or at least not degrade it.
• Ideally, mask generation should be on par with or cheaper than the kernel savings. A target could be E2E <= baseline for ws ≤ Lk/16 across bias/mask variants.

Environment information

PyTorch: 2.9.0a0+50eac811a6.nv25.09
CUDA: 13.0
GPU: NVIDIA GeForce RTX 4090

Benchmark code

Profiling information

No response

System information

No response

Additional context

No response

[LoserCheems]

However, according to the description in the original paper, without introducing an external 4D padding mask, our mask and bias should be batch_size, num_kv_heads, 1, key_len.

When key_len>window_size, there will not be and additional speed loss. However, when key_len = window_size, since the kernel needs to judge the mask and add the bias additionally, it will naturally be slower. This is as expected.

from flash_dmattn.flash_dmattn_interface import flash_dmattn_func
from flash_dmattn.utils.mask import create_mask
import torch
import time
import statistics as stats


def _format_mib(n_bytes: int) -> str:
    return f"{n_bytes / (1024 * 1024):.2f} MiB"


def benchmark(fn, *args, repeats: int = 5, warmup: int = 1, label: str = "fn"):
    torch.cuda.synchronize()
    for _ in range(warmup):
        _ = fn(*args)
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    times = []
    for _ in range(repeats):
        t0 = time.perf_counter()
        _ = fn(*args)
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        times.append(t1 - t0)
    mean = sum(times) / len(times) * 1000
    stdev = (stats.pstdev(times) if len(times) > 1 else 0.0) * 1000

    mem_str = ""
    torch.cuda.synchronize()
    cur_alloc = torch.cuda.memory_allocated()
    cur_reserved = torch.cuda.memory_reserved()
    peak_alloc = torch.cuda.max_memory_allocated()
    peak_reserved = torch.cuda.max_memory_reserved()
    mem_str = (
        f" | CUDA mem: alloc={_format_mib(cur_alloc)}, "
        f"reserved={_format_mib(cur_reserved)}, "
        f"peak_alloc={_format_mib(peak_alloc)}, "
        f"peak_reserved={_format_mib(peak_reserved)}"
    )

    print(f"{label}: {mean:.6f}ms ± {stdev:.6f}ms over {repeats} runs{mem_str}")
    return mean, stdev


device = 'cuda'

batch, num_heads, num_kv_heads, query_len, key_len, head_dim = 1, 32, 8, 65536, 65536, 128
window_size = 4096
dtype = torch.bfloat16
min_dtype = torch.finfo(dtype).min

query = torch.randn(batch, query_len, num_heads, head_dim, device=device, dtype=dtype)
key = torch.randn(batch, key_len, num_kv_heads, head_dim, device=device, dtype=dtype)
value = torch.randn(batch, key_len, num_kv_heads, head_dim, device=device, dtype=dtype)
attn_bias = torch.ones(batch, num_kv_heads, 1, key_len, device=device, dtype=dtype)
attn_mask = create_mask(attn_bias, None, batch, query_len, key_len, window_size, min_dtype)

create_mask_compiled = torch.compile(create_mask, backend="inductor", mode="max-autotune", fullgraph=False)

print("\nBenchmarking mask and attn implementations:\n")
print(f"Configuration: batch={batch}, num_heads={num_heads}, num_kv_heads={num_kv_heads}, query_len={query_len}, key_len={key_len}, head_dim={head_dim}, window_size={window_size}, dtype={dtype}\n")

benchmark(create_mask, attn_bias, None, batch, query_len, key_len, window_size, min_dtype, warmup=3, repeats=1_000, label="create_mask")
benchmark(create_mask_compiled, attn_bias, None, batch, query_len, key_len, window_size, min_dtype, warmup=3, repeats=1_000, label="create_mask_compiled")

benchmark(flash_dmattn_func, query, key, value, attn_mask, attn_bias, None, True, None, None, False, warmup=3, repeats=1_000, label="DMA")
benchmark(flash_dmattn_func, query, key, value, None, None, None, True, None, None, False, warmup=3, repeats=1_000, label="MHA")

Benchmarking mask and attn implementations:

Configuration: batch=1, num_heads=32, num_kv_heads=8, query_len=65536, key_len=65536, head_dim=128, window_size=4096, dtype=torch.bfloat16

create_mask: 0.131156ms ± 0.046332ms over 1000 runs | CUDA mem: alloc=770.00 MiB, reserved=772.00 MiB, peak_alloc=770.34 MiB, peak_reserved=772.00 MiB
create_mask_compiled: 0.092347ms ± 0.017928ms over 1000 runs | CUDA mem: alloc=769.50 MiB, reserved=846.00 MiB, peak_alloc=769.50 MiB, peak_reserved=846.00 MiB
DMA: 123.959192ms ± 0.520455ms over 1000 runs | CUDA mem: alloc=1281.50 MiB, reserved=1378.00 MiB, peak_alloc=1289.50 MiB, peak_reserved=1378.00 MiB
MHA: 217.424891ms ± 0.890010ms over 1000 runs | CUDA mem: alloc=1281.50 MiB, reserved=1378.00 MiB, peak_alloc=1289.50 MiB, peak_reserved=1378.00 MiB

Benchmarking mask and attn implementations:

Configuration: batch=1, num_heads=32, num_kv_heads=8, query_len=4096, key_len=4096, head_dim=128, window_size=4096, dtype=torch.bfloat16

create_mask: 0.067283ms ± 0.025501ms over 1000 runs | CUDA mem: alloc=48.12 MiB, reserved=54.00 MiB, peak_alloc=48.47 MiB, peak_reserved=54.00 MiB
create_mask_compiled: 0.079385ms ± 0.020779ms over 1000 runs | CUDA mem: alloc=48.09 MiB, reserved=130.00 MiB, peak_alloc=48.09 MiB, peak_reserved=130.00 MiB
DMA: 1.195452ms ± 0.085983ms over 1000 runs | CUDA mem: alloc=80.09 MiB, reserved=162.00 MiB, peak_alloc=80.59 MiB, peak_reserved=162.00 MiB
MHA: 1.054148ms ± 0.133294ms over 1000 runs | CUDA mem: alloc=80.09 MiB, reserved=162.00 MiB, peak_alloc=80.59 MiB, peak_reserved=162.00 MiB

cc: @ftgreat @zacliu2023