GFX1250 changes with updated AITER - #732
Conversation
…d for FW integration
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| } | ||
|
|
||
|
|
||
| void checkMxFP8Support(const TestParams& params, const cudaDeviceProp& prop, bool &use_mxfp8, bool &use_hipkittens_mxfp8) { |
There was a problem hiding this comment.
GTEST_SKIP() inside a helper does not abort the caller.
GTEST_SKIP() expands to return GTEST_MESSAGE_(msg, TestPartResult::kSkip) — it returns from the enclosing function only. Previously all of these skips lived directly in performTest/performDqTest, so hitting one aborted the test body. Now they return from checkMxFP8Support, and performTest continues executing the full GEMM afterwards.
Consequence: on a device where MXFP8 isn't supported (e.g. gfx942 running OperatorTestMXFP8), or on a shape failing the M/N/K alignment guards, the test is reported as skipped but the unsupported GEMM still runs — likely an NVTE_CHECK abort, HIP error, or bogus comparison rather than a clean skip. Same for performDqTest.
Suggested fix: have the helper report rather than skip, and skip at the call site, e.g. return a std::optional<std::string> (reason) or an out-param skip_reason, then in the caller:
if (auto reason = checkMxFP8Support(params, prop, use_mxfp8, use_hipkittens_mxfp8)) {
GTEST_SKIP() << *reason;
}
| GTEST_SKIP() << "MXFP8 is not supported in current config"; | ||
| } | ||
| bool _unused = false; | ||
| checkMxFP8Support(params, prop, _unused, _unused); |
There was a problem hiding this comment.
Passing the same object for two independent non-const out-params aliases them. It happens to be harmless today only because DqGEMMTestSuite is instantiated with use mxfp8 = Values(true) (line 992), so the early if (!use_mxfp8) { use_hipkittens_mxfp8 = false; return; } branch is never taken. If a non-MXFP8 param is ever added to that suite, the aliasing makes the helper silently skip all of the alignment/arch guards that performDqTest used to do unconditionally.
Two separate locals would remove the coupling and read better:
| checkMxFP8Support(params, prop, _unused, _unused); | |
| bool dq_use_mxfp8 = false; | |
| bool dq_use_hipkittens = false; | |
| checkMxFP8Support(params, prop, dq_use_mxfp8, dq_use_hipkittens); |
| // launching any kernel. Builds the same args as ck_attn_fwd and relies on AITER's | ||
| // v3_api_check dry-run (returns 1 when v3 is available, -1 otherwise). | ||
| bool ck_attn_fwd_uses_v3(const CKAttnFwdArgs& args){ | ||
| #if 1 |
There was a problem hiding this comment.
#if 1 / #else / #endif around a permanently-thrown error looks like debugging scaffolding that got committed. Two things worth resolving before merge:
ck_attn_fwd_uses_v3has no callers anywhere in the tree (only the definition here and the declaration inck_fused_attn.hpp:163). If the probe is genuinely unusable now that the split-KV and sink dispatchers ignorev3_api_check, deleting the function and its declaration is cleaner than shipping a function whose only behavior isthrow.- If it's kept as a placeholder, drop the
#if 1/#elsewrapper (the dead branch is preserved in git history) and turn the comment at lines 194-196 into an accurate description — it still says the helper exists so "the probe can never disagree with the launch", which is now the opposite of what the code does.
| fmha_args.splitkv_workspace_ptr = args.splitkv_workspace_ptr; | ||
|
|
||
| #if FA_WITH_SINK | ||
| if(args.h <= kSinkBufMaxHeads && QOLA_NS(mha_fwd_with_sink_supported)(fmha_args)) { |
There was a problem hiding this comment.
When args.h > kSinkBufMaxHeads (256) but mha_fwd_with_sink_supported() returns true, sink_ptr is silently left null and the kernel is launched anyway. Per your own comment at lines 18-19 the sink kernel "requires non-null sink_ptr of shape [nhead]", so this path substitutes undefined behavior for a diagnosable error.
Since the buffer is a fixed fill value, sizing it to the actual head count is straightforward — key the cache on (device, h) (or just allocate max(h, 256) and grow), or at minimum fail loudly:
if(QOLA_NS(mha_fwd_with_sink_supported)(fmha_args)) {
NVTE_CHECK(args.h <= kSinkBufMaxHeads, ...); // or grow the buffer
fmha_args.sink_ptr = get_gfx1250_sink_buf(device_for_stream(stream));
}
| int ck_attn_fwd_num_splits(const CKAttnFwdArgs& args){ | ||
| #if FAV_NATIVE_ON | ||
| #if FA_WITH_NATIVE_SPLITKV | ||
| aiter::mha_fwd_args fmha_args = build_fwd_fmha_args(args); |
There was a problem hiding this comment.
ck_attn_fwd_num_splits and ck_attn_fwd_workspace_size call build_fwd_fmha_args(args) without a stream, so the new sink-attach block runs here too. Both are reached from the sizing-only pass in fused_attn_ck.cpp:567-575 (before planner.is_sizing() returns), which means a pure workspace-size query can now:
- perform a
hipMalloc+ blockinghipMemcpyas a side effect, and - throw
std::runtime_errorfromget_gfx1250_sink_buf/device_for_stream.
With no stream, device_for_stream falls back to hipGetDevice(), so in a multi-GPU process the sizing pass can allocate the buffer on a different device than the one the launch later uses — leaking a redundant per-device buffer.
Since sink_ptr's value doesn't affect the split/workspace computation (only its null-ness would), consider giving build_fwd_fmha_args a flag to skip the allocation on the query paths, or hoisting the sink attach into ck_attn_fwd where the stream is known.
| NVTE_CHECK((prop.major == 9 && prop.minor == 5) || prop.major >= 12, | ||
| "MXFP4 quantization requires gfx950 and newer (detected gfx", | ||
| prop.major, prop.minor, "x)"); |
There was a problem hiding this comment.
prop.major >= 12 is broader than the intent. It admits gfx1200/gfx1201 (RDNA4, major == 12, minor == 0), which don't have the gfx1250 MXFP4 path — the check would pass and the kernel would then fail (or silently misbehave) further down. Matching the Python side, which is explicit about the arch tuple, would be safer:
| NVTE_CHECK((prop.major == 9 && prop.minor == 5) || prop.major >= 12, | |
| "MXFP4 quantization requires gfx950 and newer (detected gfx", | |
| prop.major, prop.minor, "x)"); | |
| NVTE_CHECK((prop.major == 9 && prop.minor == 5) || (prop.major == 12 && prop.minor == 5), | |
| "MXFP4 quantization requires gfx950 or gfx1250 (detected gfx", | |
| prop.major, prop.minor, "x)"); |
Separately, the message text regressed: prop.minor * 10 rendered a concrete arch (gfx950), whereas prop.minor + "x)" now prints gfx95x) / gfx125x), dropping the stepping digit that makes the diagnostic actionable.
| if is_hip_extension(): | ||
| # only GFX12.5 machines support nvfp4 | ||
| return False #TODO add gfx1250 (gpu_arch == 125) when ready | ||
| return gpu_arch >= 120 |
There was a problem hiding this comment.
The condition contradicts the comment directly above it. gpu_arch >= 120 also matches gfx1200/gfx1201 (RDNA4), which per the comment do not support NVFP4 — only gfx1250 does. On those parts the example would now claim NVFP4 support and fail at runtime.
| return gpu_arch >= 120 | |
| return gpu_arch == 125 |
|
|
||
|
|
||
| @functools.lru_cache(maxsize=None) | ||
| @torch.compiler.assume_constant_result |
There was a problem hiding this comment.
Dropping @functools.lru_cache here removes the last use of functools in this module — import functools on line 11 is now unused. unused-import is not in the disable= list in pylintrc, and qa/L0_pytorch_lint runs pylint --recursive=y ... transformer_engine/pytorch, so this will fail lint. The import needs to go along with the decorator.
|
|
||
| def _compute_mxfp4_support() -> Tuple[bool, str]: | ||
| """Return if mxfp4 support is available""" | ||
| if IS_HIP_EXTENSION: | ||
| gpu_arch = get_device_compute_capability() | ||
| if gpu_arch in ((9, 5),): # TODO: enable for gfx1250 when GEMM is available | ||
| return True, "" | ||
| return False, "Device arch gfx95x or newer is required for MXFP4 execution." | ||
| return False, "Only ROCm supports MXFP4" | ||
|
|
There was a problem hiding this comment.
Formatting won't survive the black pre-commit hook (.pre-commit-config.yaml runs black on types: [python]): top-level defs need two surrounding blank lines, and an inline # comment needs two spaces before it (line 207 right above already uses two).
| def _compute_mxfp4_support() -> Tuple[bool, str]: | |
| """Return if mxfp4 support is available""" | |
| if IS_HIP_EXTENSION: | |
| gpu_arch = get_device_compute_capability() | |
| if gpu_arch in ((9, 5),): # TODO: enable for gfx1250 when GEMM is available | |
| return True, "" | |
| return False, "Device arch gfx95x or newer is required for MXFP4 execution." | |
| return False, "Only ROCm supports MXFP4" | |
| def _compute_mxfp4_support() -> Tuple[bool, str]: | |
| """Return if mxfp4 support is available""" | |
| if IS_HIP_EXTENSION: | |
| gpu_arch = get_device_compute_capability() | |
| if gpu_arch in ((9, 5),): # TODO: enable for gfx1250 when GEMM is available | |
| return True, "" | |
| return False, "Device arch gfx95x or newer is required for MXFP4 execution." | |
| return False, "Only ROCm supports MXFP4" | |
Also worth a second look: this message and the one in _compute_mxfp8_support (line 181) both say "gfx95x or newer", but both conditions were just narrowed to reject gfx1250. On a gfx1250 part the user now gets told their arch is too old, which is the opposite of the real reason (GEMM support pending).
| if FusedAttnBackend["CK"] not in fused_backends: | ||
| pytest.skip("CK backend not available for this config") | ||
|
|
There was a problem hiding this comment.
This checks that the CK backend is available, but nothing in _compare verifies that the V3 / sink kernel actually ran — which is the entire point of this file. NVTE_CK_USES_FWD_V3=1 is a request, not a guarantee: if the gfx1250 V3 dispatch declines the config (head count, seqlen, mask type, missing blob in ci/ck_jit_prebuild.txt, …), CK silently falls back to the V2 path and every assertion here still passes against the float32 SDPA reference. The suite would go green while covering nothing new.
That gap is wider now that ck_attn_fwd_uses_v3 throws unconditionally (ck_fused_attn_fwd.cpp:319-325), since there's no longer any probe to query.
Is there something observable to assert on — e.g. parsing the kernel name emitted under CK_FUSED_ATTN_LOG_CONFIG=1 / NVTE_LOG_CK_CONFIG=1, or asserting sink_ptr != 0 in that log? Without some positive confirmation these tests can't fail for the reason they exist.
| FWD V3 notes (fmha_fwd_gfx1250_batched / fmha_fwd_with_sink_asm): | ||
| - Both D64 and D128 require a non-null sink_addr (fixed kernarg layout). | ||
| TE supplies a static [256] fp32 buffer initialized to -1e30f so that | ||
| exp(-1e30f) ≈ 0.0f adds no effective weight for non-fully-masked rows. | ||
| - D64 (ENABLE_SINK=1): kernel reads and uses the sink values as a logit | ||
| floor. Top-left causal; sq ≤ sk (rectangular) safe because even with | ||
| sink≈0 the real attention weights dominate. | ||
| - D128 (ENABLE_SINK=0): kernel ignores the sink values entirely. | ||
| Causal (top-left or bottom-right); sq == sk only — rectangular shapes | ||
| risk NaN on fully-masked KV tiles because the sink floor is disabled. | ||
| - No SWA (window_size_left must be -1). |
There was a problem hiding this comment.
This block disagrees with both the C++ side and the configs below it, which makes it actively misleading for the next reader:
- Line 12-13 says "Both D64 and D128 require a non-null sink_addr", but line 18 says D128 has
ENABLE_SINK=0, and the new comment inck_fused_attn_fwd.cpp:20states the opposite of both: "D128 (ENABLE_SINK=0): dispatch guard rejects sink_ptr!=nullptr; leave null." If the D128 guard rejects a non-null pointer, "requires a non-null sink_addr" can't be right. - Line 19-20 says D128 is "sq == sk only — rectangular shapes risk NaN", yet lines 364-376 add rectangular D128 configs (
sq=128/256, sk=2048;sq=130, sk=2300) with the inline note "bottom-right causal, sq < sk safe". - Line 16-17 says D64 is "Top-left causal", but every D64 config below uses
attn_mask_type="causal_bottom_right", and line 306-307 says the kernel implements bottom-right.
Worth reconciling to a single accurate description — right now three places state three different things about the same kernels.
| @@ -0,0 +1,382 @@ | |||
| # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. | |||
| # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
Copyright audit. This is a brand-new, AMD-authored file with no upstream NVIDIA ancestor, so asserting an NVIDIA copyright dating to 2022 on it isn't right. The convention for wholly-new AMD-only files in this repo is an AMD line only — e.g. tests/pytorch/test_gemm_autotune.py and tests/pytorch/test_sanity_hipified_cast_transpose.py:
| # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # License for AMD contributions = MIT. See LICENSE for more information |
(AMD year 2026 on line 1 is correct.) Every other file touched by this PR passes the audit — AMD end-years are all 2026 and no NVIDIA year ranges were modified.
| for key in [ | ||
| "NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN", | ||
| "NVTE_FUSED_ATTN_CK", "NVTE_FUSED_ATTN_AOTRITON", | ||
| "NVTE_CK_USES_FWD_V3", "NVTE_CK_USES_BWD_V3", | ||
| ]: | ||
| os.environ.pop(key, None) | ||
| for key, val in backend_env.items(): | ||
| os.environ[key] = str(val) | ||
| _attention_backends["backend_selection_requires_update"] = True | ||
|
|
||
| device = "cuda" | ||
| b = config.batch_size | ||
| sq = config.max_seqlen_q | ||
| sk = config.max_seqlen_kv | ||
| hq = config.num_heads | ||
| hk = config.num_gqa_groups | ||
| dqk = config.head_dim_qk |
There was a problem hiding this comment.
Two formatting/scope notes on this helper:
-
black will rewrite this file. The pre-commit hook runs black with
--line-length=100on all Python files. The multi-item-per-line list at lines 83-87 has a magic trailing comma, so black explodes it to one element per line; the aligned assignments at lines 94-99 (b =,sq =) get collapsed to single spaces. Same pattern recurs at lines 248-252. Worth runningbash qa/format.shbefore merge. -
The
is_training=Truepath is dead. Both test functions call_compare(..., is_training=False), so the backward branch here (lines 133-137), the entire grad-reconstruction block in_pytorch_ref(lines 164-167, 176-182, 216-229), and the grad assertions in_compare(lines 284-290) never execute. The module docstring says both D64 and D128 are "FWD-only", so this looks like intentional groundwork — but ~40 lines of untested GQA grad-reduction logic is the kind of thing that silently rots. Either add a training config that exercises it or drop it until BWD V3 lands.
Related: has_swa (line 189) is likewise always False — ModelConfig defaults window_size to (-1, -1) and no config here sets it, and the docstring says SWA is unsupported — so the SWA masking block at lines 200-204 is unreachable too.
| message(STATUS "test_operator hipified sources: ${test_hip_sources}") | ||
| set_target_properties(test_operator PROPERTIES SOURCES "${test_hip_sources}") | ||
| target_include_directories(test_operator BEFORE PRIVATE |
There was a problem hiding this comment.
This include directory was previously added only when gfx1250 was in CMAKE_HIP_ARCHITECTURES; it's now unconditional for every ROCm test build. Two consequences worth confirming:
3rdparty/composable_kernelbecomes a hard build requirement for all ROCm test builds, not just gfx1250 ones. Fine ifgit submodule update --init --recursiveis always assumed, but it's a new failure mode for partial checkouts.BEFOREputs CK'sinclude/ahead of every other include path for the wholetest_operatortarget. CK ships generically-named headers, so this can shadow same-named headers from TE/HIP/gtest in translation units that never needed CK. Scoping it to the one consumer would be safer —test_ck_grouped_mxfp8.cuis the only file that includesck_tile/*, and it already guards itself at runtime (GTEST_SKIPunless gfx1250, line 504-508). A per-sourceset_source_files_properties(... INCLUDE_DIRECTORIES ...)would avoid the global reordering.
| ::testing::ValuesIn(kLayouts), //transa,transb | ||
| ::testing::Values(false), //use mxfp8 | ||
| ::testing::Values(false)), //force hipblaslt | ||
| ::testing::Values(true)), //force hipblaslt |
There was a problem hiding this comment.
Is this flip intentional (same on line 961)? Both suites pass use mxfp8 = false, and NVTE_ROCM_USE_HIPBLASLT_MXFP8 is only consulted under if (is_mxfp8 && kittens_mxfp8_supported()) (rocm_gemm.cu:2280), so the execution path is unchanged.
What does change is every test name in OperatorTest/OperatorTestFP8: GEMMTestName appends std::get<5>(...) ? "HB" : "HK", so all of these rename from …xHK to …xHB. Any saved --gtest_filter or downstream job pinned to the old names silently stops matching. If the goal is just to stop labelling non-MXFP8 tests as "HipKittens", it may be cleaner to make the suffix conditional on use_mxfp8 rather than flipping a parameter that has no effect.
Claude review — GFX1250 changes with updated AITERFirst automated review of this PR (no prior Claude or human review comments found). Scope: full diff against Verdict: changes look directionally right, but there is one issue I'd treat as blocking and several worth resolving before merge. The core direction is sound — removing the blanket gfx1250 Main points (details inline):
Copyright headers: 1 file needs attention — the new |
1 participant
Description
Gfx1250 changes on top of PR#709 and the latest dev
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: