Unable to Reproduce Results from SageAttention3 Paper on RTX 5090

[xc826]

Thank you for your impressive work.

I’m currently trying to reproduce the TOPS results reported in the SageAttention3 paper, and I followed the instructions in the README, but haven’t been able to match the performance described. My test script is modified based on your benchmark file, and I’ve confirmed that both my environment correct. Perhaps there’s an issue with my test script—if you could provide an official one, that would be even better.
I’d really appreciate your help in how to reproducing the TOPS performance numbers reported in the paper.

[jwu10003]

What the SageAttention3 performance do you get?

[xc826]

What the SageAttention3 performance do you get?

Below are my results and the test code.

SageAttention3 Benchmark
dtype: bf16
batch: 4, head: 32, headdim: 128
is_causal: False
1024 flops:277.8569557380028
2048 flops:392.10339627965277
4096 flops:534.928247942938
8192 flops:655.0855152583704
16384 flops:726.705857896177
32768 flops:778.033702084239
is_causal: True
1024 flops:134.22071485552993
2048 flops:237.14272379602738
4096 flops:352.7122170264404
8192 flops:487.75702109387106
16384 flops:589.6134207134803
32768 flops:658.0092382391954

`
import torch
from flash_attn.utils.benchmark import benchmark_forward
from sageattn3 import sageattn3_blackwell
import argparse

parser = argparse.ArgumentParser(description='Benchmark SageAttention3')
parser.add_argument('--dtype', type=str, default='bf16', choices=['fp16', 'bf16'], help='Data type (fp16 or bf16)')
parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
parser.add_argument('--num_heads', type=int, default=32, help='Number of heads')
parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
args = parser.parse_args()

head = args.num_heads
batch = args.batch_size
headdim = args.head_dim
dtype = torch.bfloat16 if args.dtype == 'bf16' else torch.float16

if headdim >= 256:
print(f"Error: SageAttention3 only supports head_dim < 256, got {headdim}")
exit(1)

print(f"SageAttention3 Benchmark")
print(f"dtype: {args.dtype}")
print(f"batch: {batch}, head: {head}, headdim: {headdim}")

is_causal = False
print(f"is_causal: {is_causal}")
for seq_len in {1024, 2048, 4096, 8192, 16384, 32768}:
flops = 4 * head * batch * headdim * seq_len * seq_len // (2 if is_causal else 1)
# SageAttention3 expects (batch_size, head_num, seq_len, head_dim)
q = torch.randn(batch, head, seq_len, headdim, dtype=dtype, device="cuda")
k = torch.randn(batch, head, seq_len, headdim, dtype=dtype, device="cuda")
v = torch.randn(batch, head, seq_len, headdim, dtype=dtype, device="cuda")
for i in range(5):
sageattn3_blackwell(q, k, v, is_causal=is_causal)
torch.cuda.synchronize()
_, time = benchmark_forward(sageattn3_blackwell, q, k, v, is_causal=is_causal, repeats=100, verbose=False, desc='SageAttention3')
print(f'{seq_len} flops:{flops/time.mean*1e-12}')

is_causal = True
print(f"is_causal: {is_causal}")
for seq_len in {1024, 2048, 4096, 8192, 16384, 32768}:
flops = 4 * head * batch * headdim * seq_len * seq_len // (2 if is_causal else 1)
q = torch.randn(batch, head, seq_len, headdim, dtype=dtype, device="cuda")
k = torch.randn(batch, head, seq_len, headdim, dtype=dtype, device="cuda")
v = torch.randn(batch, head, seq_len, headdim, dtype=dtype, device="cuda")
for i in range(5):
sageattn3_blackwell(q, k, v, is_causal=is_causal)
torch.cuda.synchronize()
_, time = benchmark_forward(sageattn3_blackwell, q, k, v, is_causal=is_causal, repeats=100, verbose=False, desc='SageAttention3')
print(f'{seq_len} flops:{flops/time.mean*1e-12}')
`