Added AITER V3 API check mechanism

[Micky774]

Description

Added mechanism to check for V3 API usage for future flexibility.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[github-actions]

Claude Walkthrough

Intent. Expose a lightweight way to ask AITER whether its v3 (asm) fused-attention path would run for a given config, without launching a kernel. This gives TE a hook to make future backend-selection / dispatch decisions (e.g. fall back to CK v2 when v3 is not available) based on AITER own capability check rather than TE-side heuristics.

Key changes.

• Bumps 3rdparty/QoLA from 47e0f550 -> 0d6665b6 to pick up AITER v3_api_check dry-run mode.
• Adds two new public entry points in transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp:146-148: ck_attn_fwd_uses_v3 and ck_attn_bwd_uses_v3.
• Refactors ck_attn_fwd / ck_attn_bwd to split out the mha_fwd_args / mha_bwd_args population into shared static builders (build_fwd_fmha_args, build_bwd_fmha_args), so the probe and the real launch can never disagree on what config is being asked about.

Walkthrough.

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp — the body of ck_attn_fwd that built aiter::mha_fwd_args from CKAttnFwdArgs is lifted verbatim into a new static aiter::mha_fwd_args build_fwd_fmha_args(...). The real launch path (ck_attn_fwd) now calls the builder and keeps everything after arg-construction unchanged (env-var handling, NVTE_CK_RUNTIME_MAX_SEQLEN override, logging, dispatch). The new ck_attn_fwd_uses_v3 calls the same builder, sets fmha_args.v3_api_check = true, passes a null-stream / no-log stream_config (nothing is launched in check mode), and returns QOLA_NS(mha_fwd)(...) == 1 — AITER convention is that v3_api_check returns 1 when v3 is available and -1 otherwise.

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_bwd.cpp — same shape. build_bwd_fmha_args becomes the single source of truth for populating aiter::mha_bwd_args. ck_attn_bwd_uses_v3 mirrors the forward probe. One wrinkle: the bwd path needs is_mqa_gqa, has_dbias, and bias_shape after the AITER call for the post-dispatch dk/dv reduction and dbias handling, so ck_attn_bwd recomputes these locally after invoking the builder (they are also computed inside the builder for arg population). The comment at the top of the caller flags this intentional duplication.

ck_fused_attn.hpp — declares the two new probe functions next to the existing ck_attn_fwd / ck_attn_bwd, with a comment describing the AITER v3_api_check semantics.

Testing. No tests added. The probe is a pure dry-run against AITER, so behavior is only exercised once a caller (a future PR, per the "for future flexibility" framing in the description) wires it into dispatch.

Notes for reviewers.

• is_mqa_gqa / has_dbias / bias_shape are now computed in two places in the bwd path — a comment marks this, but if the derivation ever changes it must be updated in both. Worth deciding whether to have the builder return these via out-params instead.
• The probe deliberately does not apply the NVTE_CK_RUNTIME_MAX_SEQLEN override; the comment on the builder notes this does not affect v3 kernel selection. If that assumption ever changes, the probe and the real launch will diverge silently.
• QoLA pin bump must land together with this PR — the new API surface (v3_api_check field, dry-run return convention) is provided by the bumped commit.
• No caller uses the new probes yet, so this PR is API-only; behavior of existing ck_attn_fwd / ck_attn_bwd should be unchanged modulo the refactor.

---

Generated by Claude. To request a code review, comment /claude review.

[github-actions]

Minor doc-clarity nit for future callers: the probe's result depends on the caller's uses_fwd_v3 / uses_bwd_v3 flags, because build_{fwd,bwd}_fmha_args copies them into fmha_args.use_asm_v3 before flipping v3_api_check. So this reports "would v3 dispatch given this config as-configured", not "is v3 available for this config regardless of the opt-in flag".

If future callers use the probe to decide whether to enable v3 (rather than confirm an already-decided choice), they'll need to pre-set args.uses_{fwd,bwd}_v3 = true before calling — worth mentioning in this comment block so that's not a surprise. Not blocking.

[github-actions]

Reviewed the diff (HEAD^1..HEAD^2, 3 source files + QoLA submodule bump).

Verdict: Looks good. The refactor extracts build_{fwd,bwd}_fmha_args as static helpers and layers the new ck_attn_{fwd,bwd}_uses_v3 probes on top by flipping only v3_api_check, so the probe path and the launch path can't diverge on mha_{fwd,bwd}_args field population. ck_attn_bwd recomputes is_mqa_gqa / has_dbias / bias_shape for the post-dispatch MQA-GQA and dbias reductions — redundant but derived purely from args, so no drift risk, and the inline comment already flags it.

Left one non-blocking inline nit on the header docstring around how the probe interacts with the caller's uses_{fwd,bwd}_v3 flag — worth spelling out for future callers since no callers land in this PR.

Copyright headers: OK — all three edited files already carry Copyright (c) 2024-2026, Advanced Micro Devices, Inc. and no new files were introduced.

[wenchenvincent]

Can we add a test?

[Micky774]

Can we add a test?

What kind of test would you want added? Since this is pretty internal to the CK FA API we'd need to either test directly in C w/ known kernel configs, or create a python binding to expose to pytest. I personally think it should be okay without a test (we had this mechanism before) and that it'll actually be indirectly tested by other downstream mechanisms that rely on it (e.g. your PR).

[wenchenvincent]

@Micky774 Could you check if the CI failures are related?

[Micky774]

@Micky774 Could you check if the CI failures are related?

Failures seem unrelated. They're a combo of group GEMM failures, timeouts and resource errors.