SLA and SageSLA (SLA based on SageAttention) Code Update

[jt-zhang]

1. We’ve just updated the Triton implementation of SLA. Training is now more stable, faster, and typically achieves better training results.

2. We’ve released the code for SageSLA, a very fast SLA (Sparse-Linear Attention) forward pass based on SageAttention. It uses some code from SpargeAttn. Please refer to the SageSLA/ directory for how to use SageSLA.

Feel free to try it out!

[leeruibin]

Thanks for your excellent work. However, when I try to test the SageSLA on Hopper GPU, I meet the following error.

Warning: Sage2++ NOT enabled
Traceback (most recent call last):
  File /code/tmp.py", line 84, in <module>
    _ = sage_attn(q, k, v)
        ^^^^^^^^^^^^^^^^^^
  File "/home/tiger/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tiger/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/code/SLA/SageSLA/core.py", line 129, in forward
    qattn.qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'spas_sage_attn._qattn' has no attribute 'qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold'. Did you mean: 'qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold'?

Here is my test code.

import time
import torch
from sparse_linear_attention import SparseLinearAttention
from wan.modules.attention import attention as dense_attention # this call the FA2
from SageSLA import SageSparseLinearAttention as SageLinearAttention

sage_attn = SageLinearAttention(
    head_dim=128,
    topk=0.1,                 # = 1 - sparsity
    feature_map="softmax",    # options: elu, relu, softmax
    BLKQ=128,                 # SageSLA only supports BLKQ=128 and BLKK=64
    BLKK=64,
).cuda()

# B, H, L, D = 2, 4, 4096, 128
# q = torch.randn((B, H, L, D), dtype=torch.bfloat16, device='cuda')
# k = torch.randn((B, H, L, D), dtype=torch.bfloat16, device='cuda')
# v = torch.randn((B, H, L, D), dtype=torch.bfloat16, device='cuda')

# o = attn(q, k, v)

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

device = "cuda"

# ----------------------------
# Config
# ----------------------------
B, H, L, D = 2, 4, 4096, 128
dtype = torch.bfloat16
num_warmup = 10
num_iters = 50

# ----------------------------
# Inputs
# ----------------------------
q = torch.randn((B, H, L, D), device=device, dtype=dtype)
k = torch.randn((B, H, L, D), device=device, dtype=dtype)
v = torch.randn((B, H, L, D), device=device, dtype=dtype)

# ----------------------------
# Sparse Attention
# ----------------------------
sparse_attn = SparseLinearAttention(
    head_dim=D,
    topk=0.1,                 # sparsity
    feature_map="softmax",
    BLKQ=64,
    BLKK=64,
).to(device)

# ----------------------------
# Warmup
# ----------------------------
with torch.no_grad():
    for _ in range(num_warmup):
        _ = sparse_attn(q, k, v)
        _ = dense_attention(q, k, v)
        _ = sage_attn(q, k, v)

torch.cuda.synchronize()

# ----------------------------
# Timing helper
# ----------------------------
def benchmark(fn, name):
    torch.cuda.synchronize()
    start = time.time()

    with torch.no_grad():
        for _ in range(num_iters):
            out = fn()

    torch.cuda.synchronize()
    elapsed = time.time() - start
    print(f"{name}: {elapsed / num_iters * 1000:.3f} ms / iter")

# ----------------------------
# Run benchmarks
# ----------------------------
benchmark(lambda: dense_attention(q, k, v), "Dense Attention")
benchmark(lambda: sparse_attn(q, k, v), "SparseLinearAttention")
benchmark(lambda: sage_attn(q, k, v), "SageLinearAttention")

Meanwhile, in my test, in hopper GPU, the FA2 seems to be faster than the SparseLinearAttention. Here are my evaluation results.

SparseLinearAttention: 0.573 ms / iter
Dense Attention: 0.434 ms / iter