SageAttention Compatibility Wrapper (sageattention)

sageattention is a compatibility wrapper that preserves selected SageAttention-style Python entry points, running on top of MATE attention operators on MUSA. Matching a function name or signature does not imply that the upstream CUDA backend, quantization algorithm, or every optional argument has identical behavior.

Overview

This wrapper is designed for projects that already target SageAttention-style Python APIs, allowing you to run on MUSA through MATE with minimal integration effort.

The current compatibility scope covers the public APIs listed below. Read the compatibility notes before porting code that depends on a specific upstream backend or optional parameter.

Package and import

  • Package name: sageattention

  • Import path: sageattention

  • Runtime backend: MATE dense quantized attention operators on MUSA

MUSA wrapper releases use the PEP 440 local version suffix +musa, for example 0.2.4+musa. Use python -m pip show sageattention to distinguish this wrapper from the native package.

Requirements

Before using this wrapper, make sure the following are available:

  • TorchMUSA and the MUSA runtime environment are available.

  • The target workload is configured to execute on MUSA devices.

Build

Build a wheel from the wrappers/SageAttention directory:

python -m build --wheel

The generated wheel will be placed under:

dist/

Installation

For delivered packages, install from the external MUSA wheel source:

python -m pip install sageattention \
  --index-url https://dl.mthreads.com/repo/api/pypi/pypi/simple

This installs the matching mate dependency automatically.

For local wrapper development, install from source:

python -m pip install --no-build-isolation --no-deps -e ../.. -v
python -m pip install --no-build-isolation --no-deps -e .

Install a built local wheel:

python -m pip install --no-deps dist/sageattention-*.whl

Import

Import the package directly:

import sageattention

Import individual APIs:

from sageattention import (
    sageattn,
    sageattn_varlen,
    sageattn_qk_int8_pv_fp16_cuda,
    sageattn_qk_int8_pv_fp16_triton,
    sageattn_qk_int8_pv_fp8_cuda,
    sageattn_qk_int8_pv_fp8_cuda_sm90,
)

Public APIs

The wrapper currently exposes:

  • sageattn: primary SageAttention-compatible dense attention entry

  • sageattn_varlen: SageAttention-compatible varlen attention entry

  • sageattn_qk_int8_pv_fp8_cuda and sageattn_qk_int8_pv_fp8_cuda_sm90: compatibility entrypoints for the supported dense quantized attention path

  • sageattn_qk_int8_pv_fp16_cuda and sageattn_qk_int8_pv_fp16_triton: API compatibility entrypoints routed to the MATE FP8-PV kernel; the Triton-style entrypoint uses per-block QK quantization

0.2.7 Compatibility Additions

MATE 0.2.7 adds the following SageAttention-style entry points. They preserve the upstream-facing names and call shapes where practical, but they do not claim full upstream behavioral parity.

Entry point

MATE implementation

Compatibility boundary

sageattn_varlen

MATE FlashAttention 3 varlen path

Uses packed q, k, and v with cumulative sequence lengths. It is not the dense quantized SageAttention path; smooth_k is accepted for compatibility and ignored.

sageattn_qk_int8_pv_fp8_cuda

MATE dense SageAttention kernel

pv_accum_dtype and smooth_v are accepted for the upstream-style signature, but the current MATE path uses fixed PV accumulation and does not apply smooth_v.

sageattn_qk_int8_pv_fp16_cuda

MATE FP8-PV dense kernel

The name is FP16-PV compatible, but MATE does not provide a native FP16-PV kernel. The wrapper routes through FP8-PV while preserving the normal output dtype.

sageattn_qk_int8_pv_fp16_triton

MATE native dense kernel

The name is Triton-style; MATE does not dispatch to Triton. The default recipe is per-block (128, 128, -1, 1), quantization_backend is accepted but does not select a backend, and a non-None attn_mask raises NotImplementedError.

The existing sageattn and sageattn_qk_int8_pv_fp8_cuda_sm90 entries remain the main dense quantized path. The sm90 suffix is retained for source compatibility and does not indicate that a CUDA SM90 kernel is executed on MUSA.

Quick Start

Minimal dense attention example:

import torch
from sageattention import sageattn

device = "musa"
dtype = torch.bfloat16

q = torch.randn((1, 8, 128, 128), device=device, dtype=dtype)
k = torch.randn((1, 8, 128, 128), device=device, dtype=dtype)
v = torch.randn((1, 8, 128, 128), device=device, dtype=dtype)

out = sageattn(
    q,
    k,
    v,
    tensor_layout="HND",
    is_causal=False,
    qk_quant_dtype="int8",
)

FP8 output example:

out_fp8, out_scale = sageattn(
    q,
    k,
    v,
    tensor_layout="HND",
    is_causal=False,
    qk_quant_dtype="int8",
    fp8_output=True,
)

out_dequant = out_fp8.to(torch.float32) * out_scale

When return_lse=True, the FP8 output form returns (out_fp8, out_scale, lse). Without FP8 output, the return forms are out or (out, lse).

Tests

Wrapper-level tests are available in:

tests/test_sageattn_interface.py

Run them from the wrappers/SageAttention directory:

pytest tests/test_sageattn_interface.py

Dense Path Contract

  • The dense path requires q, k, and v on the same MUSA device, with torch.float16 or torch.bfloat16 dtype and matching dtypes.

  • Supported public tensor_layout values are "HND" and "NHD".

  • Supported head dimensions are positive values up to 128.

  • qk_quant_dtype supports int8 and fp8.

  • The default quantization recipe is (128, 16, -1, 1); passing quant_recipe overrides qk_quant_gran.

  • Only qk_quant_gran="per_thread" is supported as a shortcut. Express other supported granularities with an explicit quant_recipe.

  • fp8_output=True returns an FP8 tensor plus a torch.float32 out_scale tensor; out_scale has the same public tensor layout as the output and a final scale dimension of 1.

The dense wrapper quantizes Q and K using the selected qk_quant_dtype, while V is quantized through the MATE FP8-PV path. return_lse=True returns (out, lse) for normal output and (out_fp8, out_scale, lse) for FP8 output.

Varlen Path Contract

sageattn_varlen accepts packed tensors and cumulative sequence lengths:

q: [total_q, nheads_q, headdim]
k: [total_k, nheads_k, headdim]
v: [total_k, nheads_k, headdim_v]
cu_seqlens_q: [batch_size + 1]
cu_seqlens_k: [batch_size + 1]

max_seqlen_q and max_seqlen_k are required by the MATE FlashAttention 3 varlen call. The function returns the packed attention output with shape [total_q, nheads_q, headdim_v]. This entry point does not expose the dense wrapper’s quantization controls or return_lse output.

Unsupported or MATE-Specific Behavior

  • MATE does not currently provide a native FP16-PV SageAttention kernel; the FP16-PV compatibility names are routed through the supported MATE FP8-PV implementation. The output retains the expected public output dtype.

  • attn_mask is not supported by the Triton-style compatibility entrypoint.

  • Unsupported in this wrapper package: KV-cache wrapper entrypoints, arbitrary attention masks, and low-level pre-quantized public APIs