Fmha fwd remove ck dependency

[JaxChen29]

Motivation

The FMHA V3 forward and backward modules (module_fmha_v3_fwd, module_fmha_v3_varlen_fwd, module_fmha_v3_bwd, module_fmha_v3_varlen_bwd) use only precompiled ASM kernel binaries (.co files) at runtime -- they have no source-level dependency on the Composable Kernel (CK) library. However, the build system still required the full CK submodule at 3rdparty/composable_kernel/ for header includes (fmha_fwd.hpp, fmha_bwd.hpp, mask.hpp, ck_tile/core.hpp), CK example include paths, and CK-specific compile flags.

This unnecessary coupling caused:

• Slower builds: CK headers pull in a deep include chain, increasing compile time for modules that don't use any CK code paths.

• Larger dependency footprint: The full CK submodule (~600MB) must be cloned and available even when only V3 ASM kernels are needed.

• Blocked CK-free packaging: setup.py supports ENABLE_CK=0 to build without CK, but V3 modules still failed to compile without it.

Technical Details

1. Introduced ENABLE_CK preprocessor macro

Replaced the module-specific ONLY_FAV3 guards in public headers with a unified ENABLE_CK macro:

ENABLE_CK=1 (default): use real CK headers -- for full CK modules (fwd, bwd, splitkv, batch_prefill)
ENABLE_CK=0: use the lightweight ck_tile_shim.h -- for V3 ASM-only modules
The default is injected by core.py's build_module(), reading the same ENABLE_CK env var that setup.py uses:

enable_ck = int(os.environ.get("ENABLE_CK", "1"))
if not any("ENABLE_CK" in f for f in flags_extra_cc):
flags_cc.append(f"-DENABLE_CK={enable_ck}")
V3 modules explicitly set -DENABLE_CK=0 in their optCompilerConfig.json.

2. Guarded CK-only code in public headers

aiter_hip_common.h: #if !ENABLE_CK selects ck_tile_shim.h vs ck_tile/core.hpp
mha_fwd.h: #if ENABLE_CK guards fmha_fwd.hpp/mask.hpp includes, CK trait structs (mha_fwd_traits, mha_batch_prefill_traits, mha_fwd_splitkv_traits), and CK function declarations (mha_fwd_splitkv(), mha_batch_prefill())
mha_bwd.h: #if ENABLE_CK guards fmha_bwd.hpp include
py_itfs_common.h: #if ENABLE_CK guards CK type mapping templates (t2ck)
Structs and declarations used by the V3 path (mha_fwd_args, fmha_fwd_v3_args, mha_fwd(), fmha_fwd_v3()) remain always available -- they use ck_tile::index_t and ck_tile::stream_config which the shim provides.

3. Cleaned V3 module build config

For module_fmha_v3_fwd and module_fmha_v3_varlen_fwd:

Cleared extra_include (removed CK_DIR/example/ck_tile/01_fmha)
Cleared flags_extra_hip (removed -DCK_TILE_FMHA_FWD_FAST_EXP2=1 and related CK flags)
Added '-DENABLE_CK=0' to flags_extra_cc
4. Updated C++ test infrastructure

compile.py: Added -DENABLE_CK=0 when ck_exclude=True for both fwd and bwd paths
build_mha.sh: Added -DENABLE_CK=1 to benchmark binary hipcc commands (benchmarks need real CK for reference computation)

Test Plan

CK-free compilation: Hide 3rdparty/composable_kernel/, clean build artifacts, compile all 4 V3 modules -- verify they build successfully without CK headers.
ENABLE_CK default verification: Inspect generated ninja build files to confirm ENABLE_CK=1 for non-V3 modules and ENABLE_CK=0 for V3 modules.
Symbol leak check: Run nm -D on V3 .so files to confirm zero CK symbols leaked.
C++ standalone smoke test: Build and run smoke_test_fwd_v3.sh -a gfx942 and smoke_test_bwd_v3.sh -a gfx942.
Non-regression: Full CK modules (module_mha_fwd, module_mha_bwd) still compile and work with ENABLE_CK=1 (default).

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.