Skip to content

Fused q/k/v views on the hand-written kernels (v0.3.0)

Choose a tag to compare

@IxMxAMAR IxMxAMAR released this 10 Oct 23:09
· 5 commits to main since this release

Fused q/k/v views on the hand-written kernels, and Linux support when built from source.

python -m pip install --no-deps sageattention-2.2.0+amd.gfx12.3-cp312-cp312-win_amd64.whl

Requirements (unchanged): Windows 10/11, Python 3.12, PyTorch 2.13.0+rocm10.0.0, an RX 9070 or RX 9070 XT. Linux is supported from source: the loader carries @boxwrench's Linux branch (tested on a Radeon AI PRO R9700, Ubuntu 24.04, torch 2.9.1+rocm7.2.1), but this wheel is Windows-only and there is no prebuilt Linux wheel yet.

What changed

  • MiniMax-H3's q/k/v are now served. H3 fuses q, k and v into one projection, so they arrive as views into a single tensor and k's row stride is 3*H*D instead of its own width. Earlier releases refused that layout, which meant the call quietly ran PR #368's path even with SAGEATTN_SK1_BF16=1. The acceptance rule is now the raw launch's own requirement, stated once in sk1_loader.strided_view_reason and used by both the pre-launch envelope and the loader's own check, so the two cannot disagree; a refusal carries its reason instead of a catch-all.
  • Against PR #368 at six production sizes (N = 27 280 to 99 792), the full call is 6–13 % faster (0.87–0.94× its time). In a real H3 w4a8 int4 render at 864x480 and 124 frames, the sampling phase is 5 % faster at 4 steps and 15 % faster at 8 (0.950× and 0.847× its time), same checkpoint and seed.
  • The version strings agree again: setup.py, sageattention/__init__.py and sageattention/_version.py all say 2.2.0+amd.gfx12.3.

Worth knowing

  • Those margins are ratios, not milliseconds. This GPU's clock cannot be pinned, so every timing block is ratio-only.
  • The sweep stops at N = 99 792 because that is where the launch's 32-bit addressing runs out, not because the margin does.
  • head_dim 64 still uses PR #368's kernel. It has its own entry on the list.
  • The native extension is byte-identical to v0.2.0's, deliberately: no C++ or HIP source changed in this release, so nothing was rebuilt.

There's a note from me at the end of the README about where this is going.

See the README for the numbers and docs/FINDINGS.md (F160–F161) for how they were measured.

sha256 f23517243970fc57d014661e8d568e21decf327106bf7de01ac36c65e04a7a96