Skip to content

Update FMHA calls with new MX qscale related args#2008

Draft
ex-rzr wants to merge 1 commit intomainfrom
agorenko/fmha-update-api-new-mx-params
Draft

Update FMHA calls with new MX qscale related args#2008
ex-rzr wants to merge 1 commit intomainfrom
agorenko/fmha-update-api-new-mx-params

Conversation

@ex-rzr
Copy link

@ex-rzr ex-rzr commented Feb 9, 2026

Motivation

ROCm/rocm-libraries#4368 adds these new args for CK_TILE FMHA:

  • seqstart_v_scale_ptr
  • stride_q_descale
  • stride_k_descale
  • stride_v_descale

Technical Details

These update the relevant code to match the new API. None of the args and quant_scale_enum::mx are actually used by Aiter.

Caution

CK has moved to the monorepo while Aiter still uses the old separate repo as a submodule.
This PR does not include a new hash for 3rdparty/composable_kernel so it must not be merged before updating Aiter for monorepo!

Test Plan

Run the following tests to check that
aiter/op_tests/test_mha.py
aiter/op_tests/test_mha_varlen.py

The tests have been run manually because of monorepo.

Test Result

The existing tests work as usual.

Submission Checklist

CK_TILE FMHA new args:
 * seqstart_v_scale_ptr
 * stride_q_descale
 * stride_k_descale
 * stride_v_descale
See ROCm/rocm-libraries#4368
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant