[TRITON] Sagev2 patch

[juuso-oskari]

This PR concerns the sage attention kernels (the vanilla and the mxfp4):

• Use int64 offsets to prevent B x S x H x D from overflowing int32 values
• Move the sage attention quantization kernels to their own files.
• Revert to un-fused way of quantization for perf boost for mxfp4 sage attention.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:sglang	SGLang integration tests
ci:atom	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
ci:multi-gpu	Multi-GPU op tests (8 GPU)
ci:vllm	vLLM benchmark
ci:all	All of the above

Add labels via the sidebar or gh pr edit 2240 --add-label <label>

[Copilot]

Pull request overview

This PR updates SageAttention (vanilla + MXFP4) Triton paths to avoid int32 overflow in pointer arithmetic by promoting offsets to int64, and refactors MXFP4 quantization to use an unfused quantization path for better performance.

Changes:

• Cast key program ids / offsets to tl.int64 in Sage attention kernels to prevent overflow in pointer calculations.
• Move Sage quantization logic into dedicated quant wrapper / kernel modules and switch MXFP4 wrapper to the unfused quantization path.
• Update MXFP4 benchmark to use CUDA-graph benchmarking and add a -test flag to optionally run accuracy checks.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
op_tests/op_benchmarks/triton/bench_fav3_sage_mxfp4.py	Switches benchmark timing method and gates correctness tests behind a new CLI flag.
aiter/ops/triton/quant/sage_attention_quant_wrappers.py	Introduces Python-level quantization wrappers (rotation/smoothing/downcast) used by Sage attention.
aiter/ops/triton/attention/fav3_sage_attention_mxfp4_wrapper.py	Rewires MXFP4 forward wrapper to use the new unfused quantization wrapper.
aiter/ops/triton/attention/fav3_sage.py	Redirects Sage quant import to the new quant wrapper module.
aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py	Adds Triton kernels for Sage quantization (including int64 pid handling).
aiter/ops/triton/_triton_kernels/attention/fav3_sage_attention_mxfp4.py	Promotes program ids to int64 and removes in-file quantization helpers.
aiter/ops/triton/_triton_kernels/attention/fav3_sage_attention.py	Promotes program ids to int64 and removes in-file quantization helpers.
3rdparty/composable_kernel	Updates the CK submodule revision.

Comments suppressed due to low confidence (1)

aiter/ops/triton/attention/fav3_sage_attention_mxfp4_wrapper.py:1

• The wrapper’s public flags (hadamard_rotation, q_smooth, R, BLOCK_R) are no longer passed into quantization, so toggling these options will not affect behavior as the API suggests. Consider either (a) plumbing these args through to a quant path that honors them (e.g., call fused_sage_quant_mxfp4 when requested, or extend the unfused path to accept/implement them), or (b) explicitly disallow/raise when these flags are enabled to avoid silent misconfiguration.

# SPDX-License-Identifier: MIT

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[azaidy]

LGTM!