diff --git a/projects/composablekernel/Dockerfile b/projects/composablekernel/Dockerfile index 91aeee87a1e8..3389c1a75d35 100644 --- a/projects/composablekernel/Dockerfile +++ b/projects/composablekernel/Dockerfile @@ -100,7 +100,15 @@ RUN set -x && \ wget https://github.com/Yelp/dumb-init/releases/download/v1.2.0/dumb-init_1.2.0_amd64.deb && \ dpkg -i dumb-init_*.deb && rm dumb-init_*.deb && \ # Install packages for processing the performance results - pip3 install --break-system-packages --upgrade pytest pymysql pandas==2.2.3 sqlalchemy==2.0.3 setuptools-rust setuptools sshtunnel==0.4.0 && \ +# ml_dtypes is pinned like its neighbours, and for a sharper reason than most: +# it supplies both float8_e4m3fn/float8_e5m2 (OCP) and float8_e4m3fnuz/ +# float8_e5m2fnuz (FNUZ), and the dispatcher tests pick between those pairs per +# arch. An unpinned upgrade would silently move fp8 encoding semantics -- and +# the reference values every parity/correctness test compares against -- from +# under CI on an image rebuild. 0.6.0 is what this image already resolved on +# Ubuntu 24.04 / Python 3.12, so this records current behaviour rather than +# bumping it. + pip3 install --break-system-packages --upgrade pytest pymysql pandas==2.2.3 sqlalchemy==2.0.3 setuptools-rust setuptools sshtunnel==0.4.0 ml_dtypes==0.6.0 && \ # Add render group groupadd -f render && \ # Install the new rocm-cmake version diff --git a/projects/composablekernel/Jenkinsfile b/projects/composablekernel/Jenkinsfile index 98928d326f64..b52fe8142e38 100644 --- a/projects/composablekernel/Jenkinsfile +++ b/projects/composablekernel/Jenkinsfile @@ -91,7 +91,7 @@ def showCompilerInfo(boolean deferBinaryInfo = false, String dockerImage = "") { //launch develop branch daily jobs CRON_SETTINGS = BRANCH_NAME == "develop" ? '''0 23 * * * % RUN_FULL_QA=true;RUN_CK_TILE_FMHA_TESTS=true;RUN_PERFORMANCE_TESTS=true;FORCE_CI=true - 0 22 * * * % RUN_FULL_QA=true;DISABLE_DL_KERNELS=true;RUN_TILE_ENGINE_BASIC_TESTS=true;RUN_TILE_ENGINE_GEMM_TESTS=true;RUN_PERFORMANCE_TESTS=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true + 0 22 * * * % RUN_FULL_QA=true;DISABLE_DL_KERNELS=true;RUN_DISPATCHER_PERF_TESTS=true;RUN_DISPATCHER_CORRECTNESS_TESTS=true;RUN_PERFORMANCE_TESTS=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true 0 21 * * * % RUN_GROUPED_CONV_LARGE_CASES_TESTS=true;hipTensor_test=true;BUILD_GFX101=false;BUILD_GFX908=false;BUILD_GFX942=true;BUILD_GFX950=true;RUN_PERFORMANCE_TESTS=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true;BUILD_PACKAGES=true 0 19 * * * % BUILD_DOCKER=true;COMPILER_VERSION=develop;BUILD_COMPILER=/llvm-project/build/bin/clang++;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true 0 17 * * * % BUILD_DOCKER=true;COMPILER_VERSION=therock;USE_SCCACHE=false;NINJA_BUILD_TRACE=true;RUN_ALL_UNIT_TESTS=true;FORCE_CI=true @@ -182,13 +182,13 @@ pipeline { defaultValue: false, description: "Run the ck_tile FMHA tests (default: OFF)") booleanParam( - name: "RUN_TILE_ENGINE_BASIC_TESTS", - defaultValue: true, - description: "Run the tile_engine_basic tests (default: ON)") + name: "RUN_DISPATCHER_CORRECTNESS_TESTS", + defaultValue: false, + description: "Run Correctness Tier for Dispatcher") booleanParam( - name: "RUN_TILE_ENGINE_GEMM_TESTS", + name: "RUN_DISPATCHER_PERF_TESTS", defaultValue: false, - description: "Run the tile_engine_gemm tests (default: OFF)") + description: "Run Performance Tier for Dispatcher") booleanParam( name: "BUILD_INSTANCES_ONLY", defaultValue: false, @@ -383,92 +383,52 @@ pipeline { beforeAgent true expression { env.SHOULD_RUN_CI.toBoolean() } } - parallel - { - stage("Run Pytorch Tests on gfx942") - { - when { - beforeAgent true - expression { params.RUN_PYTORCH_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - ck.run_downstream_tests(image: "${env.CK_PYTORCH_IMAGE}", timeoutHours: 2, execute_cmds: ck.getPytorchTestsCmds()) - cleanWs() + agent none + steps { + script { + loadCk() + parallel([ + "Run Pytorch Tests on gfx942": { + if (params.RUN_PYTORCH_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx942")) { + ck.run_downstream_tests(image: "${env.CK_PYTORCH_IMAGE}", timeoutHours: 2, execute_cmds: ck.getPytorchTestsCmds()) + cleanWs() + } } - } - } - } - stage("Run AITER Tests on gfx942") - { - when { - beforeAgent true - expression { params.RUN_AITER_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) - cleanWs() + }, + "Run AITER Tests on gfx942": { + if (params.RUN_AITER_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx942")) { + ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) + cleanWs() + } } - } - } - } - stage("Run AITER Tests on gfx950") - { - when { - beforeAgent true - expression { params.RUN_AITER_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx950")) { - ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) - cleanWs() + }, + "Run AITER Tests on gfx950": { + if (params.RUN_AITER_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx950")) { + ck.run_downstream_tests(image: "${env.CK_AITER_IMAGE}", timeoutHours: 5, execute_cmds: ck.getAiterTestsCmds()) + cleanWs() + } } - } - } - } - stage("Run FA Tests on gfx942") - { - when { - beforeAgent true - expression { params.RUN_FA_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) - cleanWs() + }, + "Run FA Tests on gfx942": { + if (params.RUN_FA_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx942")) { + ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) + cleanWs() + } } - } - } - } - stage("Run FA Tests on gfx950") - { - when { - beforeAgent true - expression { params.RUN_FA_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx950")) { - ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) - cleanWs() + }, + "Run FA Tests on gfx950": { + if (params.RUN_FA_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx950")) { + ck.run_downstream_tests(image: "${env.CK_FA_IMAGE}", timeoutHours: 5, execute_cmds: ck.getFaTestsCmds()) + cleanWs() + } } } - } + ]) } } } @@ -562,364 +522,154 @@ pipeline { beforeAgent true expression { env.SHOULD_RUN_CI.toBoolean() } } - parallel - { - stage("Run CK_TILE_FMHA Tests on gfx90a") - { - when { - beforeAgent true - expression { params.RUN_CK_TILE_FMHA_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx90a")) { - deleteDir() - ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx90a")) - cleanWs() - } - } - } - } - stage("Run CK_TILE_FMHA Tests on gfx942") - { - when { - beforeAgent true - expression { params.RUN_CK_TILE_FMHA_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - deleteDir() - ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx942")) - cleanWs() + agent none + steps { + script { + loadCk() + parallel([ + "Run CK_TILE_FMHA Tests on gfx90a": { + if (params.RUN_CK_TILE_FMHA_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx90a")) { + deleteDir() + ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx90a")) + cleanWs() + } } - } - } - } - stage("Run CK_TILE_FMHA Tests on gfx950") - { - when { - beforeAgent true - expression { params.RUN_CK_TILE_FMHA_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx950")) { - deleteDir() - ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx950")) - cleanWs() + }, + "Run CK_TILE_FMHA Tests on gfx942": { + if (params.RUN_CK_TILE_FMHA_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx942")) { + deleteDir() + ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx942")) + cleanWs() + } } - } - } - } - stage("Run CK_TILE_FMHA Tests on gfx1201") - { - when { - beforeAgent true - expression { params.RUN_CK_TILE_FMHA_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx1201")) { - deleteDir() - ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx1201")) - cleanWs() + }, + "Run CK_TILE_FMHA Tests on gfx950": { + if (params.RUN_CK_TILE_FMHA_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx950")) { + deleteDir() + ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx950")) + cleanWs() + } } - } - } - } - } - } - stage("Run TILE_ENGINE_BASIC Tests") - { - when { - beforeAgent true - expression { env.SHOULD_RUN_CI.toBoolean() } - } - parallel - { - stage("Run TILE_ENGINE_BASIC Tests on gfx942") - { - when { - beforeAgent true - expression { params.RUN_TILE_ENGINE_BASIC_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - deleteDir() - ck.runTileEngineBasicTests(params.BUILD_COMPILER) - cleanWs() + }, + "Run CK_TILE_FMHA Tests on gfx1201": { + if (params.RUN_CK_TILE_FMHA_TESTS.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx1201")) { + deleteDir() + ck.buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: ck.build_and_run_fmha("gfx1201")) + cleanWs() + } } } - } + ]) } } } - stage("Run TILE_ENGINE_GEMM Tests") + stage("Run DISPATCHER Tests") { when { beforeAgent true - expression { env.SHOULD_RUN_CI.toBoolean() } + expression { env.SHOULD_RUN_CI.toBoolean() && (params.RUN_DISPATCHER_CORRECTNESS_TESTS.toBoolean() || params.RUN_DISPATCHER_PERF_TESTS.toBoolean()) } } - parallel - { - stage("Run TILE_ENGINE_GEMM Tests on gfx942") - { - when { - beforeAgent true - expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - deleteDir() - ck.runTileEngineGemmTests("gfx942", params.BUILD_COMPILER) - cleanWs() - } - } - } - } - stage("Run TILE_ENGINE_GEMM Tests on gfx950") - { - when { - beforeAgent true - expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx950")) { - deleteDir() - ck.runTileEngineGemmTests("gfx950", params.BUILD_COMPILER) - cleanWs() - } - } - } - } - stage("Run TILE_ENGINE_GEMM Tests on gfx1201") - { - when { - beforeAgent true - expression { params.RUN_TILE_ENGINE_GEMM_TESTS.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx1201")) { - deleteDir() - ck.runTileEngineGemmTests("gfx1201", params.BUILD_COMPILER) - cleanWs() - } - } - } + agent none + steps { + script { + loadCk() + ck.runDispatcherTests( + this.&rocmnode, + params.RUN_DISPATCHER_CORRECTNESS_TESTS.toBoolean(), + params.RUN_DISPATCHER_PERF_TESTS.toBoolean(), + params.BUILD_COMPILER) } } } - stage("Build CK and run Tests") + stage("Build CK and run Tests") { when { beforeAgent true expression { env.SHOULD_RUN_CI.toBoolean() } } - parallel - { - stage("Build CK and run Tests on gfx942") - { - when { - beforeAgent true - expression { (params.BUILD_GFX942.toBoolean() || params.RUN_FULL_QA.toBoolean()) && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - deleteDir() - ck.runBuildCKAndTests("gfx942") - cleanWs() + agent none + steps { + script { + loadCk() + parallel([ + "Build CK and run Tests on gfx942": { + if ((params.BUILD_GFX942.toBoolean() || params.RUN_FULL_QA.toBoolean()) && !params.BUILD_INSTANCES_ONLY.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx942")) { + deleteDir() + ck.runBuildCKAndTests("gfx942") + cleanWs() + } } - } - } - } - stage("Build CK and run Tests on gfx950") - { - when { - beforeAgent true - expression { params.BUILD_GFX950.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx950")) { - deleteDir() - ck.runBuildCKAndTests("gfx950") - cleanWs() + }, + "Build CK and run Tests on gfx950": { + if (params.BUILD_GFX950.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx950")) { + deleteDir() + ck.runBuildCKAndTests("gfx950") + cleanWs() + } } - } - } - } - /* - stage("Build CK and run Tests on gfx908") - { - when { - beforeAgent true - expression { params.BUILD_GFX908.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent{ label rocmnode("gfx908") } - steps{ - deleteDir() - script { loadCk(); ck.runBuildCKAndTests("gfx908") } - cleanWs() - } - } - */ - stage("Build CK and run Tests on gfx90a") - { - when { - beforeAgent true - expression { params.BUILD_GFX90A.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx90a")) { - deleteDir() - ck.runBuildCKAndTests("gfx90a") - cleanWs() + }, + "Build CK and run Tests on gfx90a": { + if (params.BUILD_GFX90A.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx90a")) { + deleteDir() + ck.runBuildCKAndTests("gfx90a") + cleanWs() + } } - } - } - } - stage("Build CK instances for all supported targets") - { - when { - beforeAgent true - expression { params.BUILD_INSTANCES_ONLY.toBoolean() && !params.RUN_FULL_QA.toBoolean() } - } - agent none - environment{ - setup_args = "NO_CK_BUILD" - execute_args = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -DCMAKE_CXX_COMPILER="${params.BUILD_COMPILER}" \ - -DCMAKE_HIP_COMPILER="${params.BUILD_COMPILER}" \ - -DGPU_ARCHS="gfx908;gfx90a;gfx942;gfx950;gfx10-3-generic;gfx11-generic;gfx12-generic" \ - -D CMAKE_BUILD_TYPE=Release .. && ninja -j64 """ - } - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx942")) { - deleteDir() - ck.runBuildInstancesOnly(params.BUILD_COMPILER) - cleanWs() + }, + "Build CK instances for all supported targets": { + if (params.BUILD_INSTANCES_ONLY.toBoolean() && !params.RUN_FULL_QA.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx942")) { + deleteDir() + ck.runBuildInstancesOnly(params.BUILD_COMPILER) + cleanWs() + } } - } - } - } - /* - stage("Build CK and run Tests on gfx1010") - { - when { - beforeAgent true - expression { params.BUILD_GFX101.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent{ label rocmnode("gfx1010") } - steps{ - deleteDir() - script { loadCk(); ck.runBuildCKAndTests("gfx10-1-generic") } - cleanWs() - } - } - */ - stage("Build CK and run Tests on gfx1030") - { - when { - beforeAgent true - expression { params.BUILD_GFX103.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx1030")) { - deleteDir() - ck.runBuildCKAndTests("gfx10-3-generic") - cleanWs() + }, + "Build CK and run Tests on gfx1030": { + if (params.BUILD_GFX103.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx1030")) { + deleteDir() + ck.runBuildCKAndTests("gfx10-3-generic") + cleanWs() + } } - } - } - } - stage("Build CK and run Tests on gfx11") - { - when { - beforeAgent true - expression { params.BUILD_GFX11.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode('miopen && (gfx1101 || gfx1100)') { - deleteDir() - ck.runBuildCKAndTests("gfx11-generic") - cleanWs() + }, + "Build CK and run Tests on gfx11": { + if (params.BUILD_GFX11.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean()) { + ck.runOnHealthyNode('miopen && (gfx1101 || gfx1100)') { + deleteDir() + ck.runBuildCKAndTests("gfx11-generic") + cleanWs() + } } - } - } - } - stage("Build CK and run Tests on gfx1201") - { - when { - beforeAgent true - expression { params.BUILD_GFX12.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx1201")) { - deleteDir() - ck.runBuildCKAndTests("gfx12-generic") - cleanWs() + }, + "Build CK and run Tests on gfx1201": { + if (params.BUILD_GFX12.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx1201")) { + deleteDir() + ck.runBuildCKAndTests("gfx12-generic") + cleanWs() + } } - } - } - } - stage("Build CK for gfx1250") - { - when { - beforeAgent true - expression { params.BUILD_GFX1250.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean() } - } - agent none - steps{ - script { - loadCk() - ck.runOnHealthyNode(rocmnode("gfx90a")) { - deleteDir() - ck.runBuildCKAndTests("gfx1250") - cleanWs() + }, + "Build CK for gfx1250": { + if (params.BUILD_GFX1250.toBoolean() && !params.RUN_FULL_QA.toBoolean() && !params.BUILD_INSTANCES_ONLY.toBoolean()) { + ck.runOnHealthyNode(rocmnode("gfx90a")) { + deleteDir() + ck.runBuildCKAndTests("gfx1250") + cleanWs() + } } } - } + ]) } } post { @@ -927,7 +677,6 @@ pipeline { node(rocmnode("nogpu")) { script { loadCk() - // Simulate capture ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx11.json") ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx12.json") ck.generateAndArchiveBuildTraceVisualization("ck_build_trace_gfx90a.json") @@ -941,7 +690,6 @@ pipeline { script { node(rocmnode("nogpu")) { loadCk() - // Report the parent stage build ck and run tests status ck.setGithubStatus("${env.STAGE_NAME}", 'success', "Stage ${env.STAGE_NAME} passed") echo "Reporting success status for build ck and run tests" } diff --git a/projects/composablekernel/dispatcher/python/batched_contraction_utils.py b/projects/composablekernel/dispatcher/python/batched_contraction_utils.py index 2c0244d3d45a..be6158a19223 100644 --- a/projects/composablekernel/dispatcher/python/batched_contraction_utils.py +++ b/projects/composablekernel/dispatcher/python/batched_contraction_utils.py @@ -693,11 +693,40 @@ def expand_sweep(config: dict, dtype: str = "fp16", layout: str = "rcr") -> List return out -def default_fp16_config(gfx_arch: Optional[str] = None) -> BatchedContractionKernelConfig: +# Warp tile per dtype, chosen from the XDL allow-list in is_valid(). fp16 and +# bf16 share the same MFMA shape; fp32 has no 32x32x16 instruction and needs +# 32x32x8 instead. Picking a shape outside the allow-list is not an error you +# would notice -- is_valid() just returns False and the config is silently +# dropped from the sweep (test_batched_contraction_bridge.py asserts exactly +# that), so a "default" config with the wrong tile would look like a dtype that +# simply has no instances. +_DEFAULT_WARP_TILE = { + "fp16": (32, 32, 16), + "bf16": (32, 32, 16), + "fp32": (32, 32, 8), +} + + +def default_config( + dtype: str = "fp16", gfx_arch: Optional[str] = None +) -> BatchedContractionKernelConfig: + """A known-valid single config for ``dtype``, for smoke tests and defaults.""" + try: + wt_m, wt_n, wt_k = _DEFAULT_WARP_TILE[dtype] + except KeyError: + raise ValueError( + f"no default warp tile for dtype {dtype!r}; " + f"known: {', '.join(sorted(_DEFAULT_WARP_TILE))}" + ) from None gfx_arch = _validate_arch(gfx_arch) if gfx_arch else _get_arch() return BatchedContractionKernelConfig( - dtype="fp16", layout="rcr", pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + dtype=dtype, layout="rcr", pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", tile_m=128, tile_n=128, tile_k=64, warp_m=2, warp_n=2, warp_k=1, - warp_tile_m=32, warp_tile_n=32, warp_tile_k=16, + warp_tile_m=wt_m, warp_tile_n=wt_n, warp_tile_k=wt_k, num_dim_g=1, num_dim_m=1, num_dim_n=1, num_dim_k=1, gfx_arch=gfx_arch, ) + + +def default_fp16_config(gfx_arch: Optional[str] = None) -> BatchedContractionKernelConfig: + """Thin alias kept for existing callers; prefer :func:`default_config`.""" + return default_config("fp16", gfx_arch=gfx_arch) diff --git a/projects/composablekernel/dispatcher/python/dispatcher_common.py b/projects/composablekernel/dispatcher/python/dispatcher_common.py index 5254df6d4684..b0304e0d804c 100644 --- a/projects/composablekernel/dispatcher/python/dispatcher_common.py +++ b/projects/composablekernel/dispatcher/python/dispatcher_common.py @@ -101,6 +101,120 @@ def detect_gpu_arch(fallback: str = "gfx942") -> str: return fallback +# ============================================================================ +# fp8 / bf8 encoding format per architecture +# ============================================================================ +# +# CK has two incompatible 8-bit float encodings and picks between them per arch: +# +# OCP (gfx950, gfx12): e4m3fn / e5m2 -- exponent bias 7 / 15 +# FNUZ (everything else, notably gfx942): e4m3fnuz / e5m2fnuz -- bias 8 / 16 +# +# include/ck_tile/core/config.hpp:361-371 resolves this at compile time, but only +# the *device* pass sees __gfx950__; the host pass of the very same header falls +# back to FNUZ. So a host-side reference encoder that guesses from the header ends +# up disagreeing with the kernel it is validating. Every caller must therefore +# decide from the target arch string, and pass -DCK_TILE_USE_OCP_FP8 explicitly so +# both compiler passes agree. These helpers are the single source of truth for +# both halves of that contract. + + +def normalize_arch(arch: Optional[str]) -> str: + """Lowercase `arch` and strip the target-feature suffix, if any. + + rocminfo, hipcc and the HSA runtime all hand back full target triples -- + ``gfx950:sramecc+:xnack-`` -- while configs, CI parameters and these helpers + are written in terms of the bare name. Every arch predicate below matches on + a prefix, so a triple happens to survive unnormalized, but anything carrying + a leading component (an ``amdgcn-amd-amdhsa--gfx950`` offload target, say) + silently falls through to the FNUZ default. That miss drops + ``-DCK_USE_OCP_FP8``, and host and device then disagree on the fp8 encoding + with no diagnostic at all -- the kernel builds and returns wrong numbers. + + Normalizing in one place, up front, is what keeps that from depending on + which spelling of the arch a given caller happened to be handed. + """ + a = (arch or "").lower().strip() + # Feature suffix: gfx950:sramecc+:xnack- -> gfx950 + a = a.split(":", 1)[0] + # Offload-target prefix: amdgcn-amd-amdhsa--gfx950 -> gfx950 + idx = a.rfind("gfx") + return a[idx:] if idx > 0 else a + + +def fp8_uses_ocp(arch: Optional[str]) -> bool: + """True iff `arch` uses the OCP fp8/bf8 encoding rather than FNUZ. + + Mirrors the __gfx950__ / __gfx12__ test in include/ck_tile/core/config.hpp. + Use this to select the host-side codec (ml_dtypes.float8_e4m3fn vs + float8_e4m3fnuz) so it matches the bytes the kernel actually produces. + + Accepts bare names and full target triples alike; see `normalize_arch`. + """ + a = normalize_arch(arch) + return a.startswith("gfx950") or a.startswith("gfx12") + + +def ocp_arch_defines(arch: Optional[str]) -> List[str]: + """hipcc defines that pin the fp8/bf8 encoding for `arch`. + + Returned for OCP archs only; FNUZ is the default for both compiler passes and + needs no define. Passing these makes the host pass agree with the device pass + instead of silently falling back to FNUZ. + """ + if not fp8_uses_ocp(arch): + return [] + return ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"] + + +def arch_feature_defines(arch: Optional[str]) -> List[str]: + """`ocp_arch_defines` plus the per-arch feature-enablement defines. + + The top-level CMakeLists.txt (:456-512) sets these for a normal build; they + are absent in the standalone hipcc JIT path, so any bridge that compiles a + kernel out-of-tree has to re-supply them. This mirrors that block for the + arches the dispatcher targets: + + gfx950 -> CK_USE_NATIVE_MX_SUPPORT, CK_GFX950_SUPPORT + gfx1250 -> CK_USE_GFX1250, CK_USE_NATIVE_MX_SUPPORT, CK_GFX1250_SUPPORT + gfx11/gfx12 -> CK_TILE_USE_WMMA=1 (gfx12 also CK_GFX12_SUPPORT) + + CK_TILE_USE_WMMA is the one that must be passed even when it is 0: CMake + always defines it (:480), and the JIT path leaving it undefined only happens + to work because the preprocessor reads an undefined identifier as 0, which is + the right answer on gfx942/gfx950 and the wrong one on every WMMA part. + + CK_USE_GFX950 is deliberately *not* emitted -- CMake sets it, but it is read + only by the legacy ck/library conv instances, never by ck_tile, so it is + dead weight on this path. + + On the warp tile: these defines do select the branch of + tile_gemm_shape.hpp get_k_warp_tile() (CK_TILE_USE_WMMA outermost, then + CK_USE_GFX1250 / CK_GFX950_SUPPORT), but the emitted kernels never *call* + that function -- codegen writes a literal warp_tile_k, mirroring it in Python + via codegen_common.fp8_warp_tile_k_for_arch. The hazard is therefore a + mismatch between that literal and the warp-gemm the target arch actually has: + there is no 16x16x128 fp8 warp-gemm on gfx942, and asking for one compiles + cleanly and returns all zeros. Keep warp_tile_k arch-aware alongside these + defines, and derive it from the one helper rather than a local copy. + """ + a = normalize_arch(arch) + defines = ocp_arch_defines(arch) + if a.startswith("gfx950"): + defines = defines + ["-DCK_USE_NATIVE_MX_SUPPORT", "-DCK_GFX950_SUPPORT"] + elif a.startswith("gfx11") or a.startswith("gfx12"): + defines = defines + ["-DCK_TILE_USE_WMMA=1"] + if a.startswith("gfx12"): + defines = defines + ["-DCK_GFX12_SUPPORT"] + if a.startswith("gfx1250"): + defines = defines + [ + "-DCK_USE_GFX1250", + "-DCK_USE_NATIVE_MX_SUPPORT", + "-DCK_GFX1250_SUPPORT", + ] + return defines + + # ============================================================================ # Architecture Filter Data # ============================================================================ diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 0475a58e6a29..767bebc3428c 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -66,12 +66,24 @@ try: - # Reuse the single canonical amd-smi bridge instead of re-implementing it. + # Reuse the single canonical amd-smi bridge instead of re-implementing it, + # and the single source of truth for the fp8/bf8 encoding format per arch. from dispatcher_common import _detect_gpu_arch_via_amd_smi + from dispatcher_common import fp8_uses_ocp as _fp8_uses_ocp + from dispatcher_common import ocp_arch_defines as _ocp_arch_defines except Exception: # noqa: BLE001 - standalone use without dispatcher_common on path def _detect_gpu_arch_via_amd_smi() -> Optional[str]: return None + def _fp8_uses_ocp(arch: Optional[str]) -> bool: + a = (arch or "").lower() + return a.startswith("gfx950") or a.startswith("gfx12") + + def _ocp_arch_defines(arch: Optional[str]) -> List[str]: + if not _fp8_uses_ocp(arch): + return [] + return ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"] + @functools.lru_cache(maxsize=1) def _get_arch() -> str: @@ -179,11 +191,13 @@ def _cap(flag: bool) -> str: } -def numpy_dtype_for(dtype: str): +def numpy_dtype_for(dtype: str, use_ocp: Optional[bool] = None): """Return the numpy dtype object used for host operands of ``dtype``. fp16 -> np.float16; bf16/fp8/bf8 require the ``ml_dtypes`` package (imported - lazily) and use FNUZ fp8 encodings for gfx942 parity. + lazily). The fp8/bf8 encoding follows the target arch: OCP (e4m3fn/e5m2) on + gfx950/gfx12, FNUZ (e4m3fnuz/e5m2fnuz) everywhere else. ``use_ocp=None`` + resolves from the local GPU; pass it explicitly when building for another arch. """ token = _DTYPE_ALIASES.get(str(dtype).lower()) if token is None: @@ -198,10 +212,11 @@ def numpy_dtype_for(dtype: str): ) from exc if token == "bf16": return np.dtype(ml_dtypes.bfloat16) + ocp = _resolve_use_ocp(use_ocp) if token == "fp8": - return np.dtype(ml_dtypes.float8_e4m3fnuz) + return np.dtype(ml_dtypes.float8_e4m3fn if ocp else ml_dtypes.float8_e4m3fnuz) if token == "bf8": - return np.dtype(ml_dtypes.float8_e5m2fnuz) + return np.dtype(ml_dtypes.float8_e5m2 if ocp else ml_dtypes.float8_e5m2fnuz) raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}") # pragma: no cover @@ -225,14 +240,16 @@ def output_dtype_for(dtype: str) -> str: return CommonTypeMappings.get_output_dtype(token) -def output_numpy_dtype_for(dtype: str): +def output_numpy_dtype_for(dtype: str, use_ocp: Optional[bool] = None): """Numpy dtype of a kernel's OUTPUT buffer for input ``dtype``. Composition of :func:`output_dtype_for` + :func:`numpy_dtype_for`. For fp8/bf8 this resolves to ``np.float16`` (2 bytes) because the kernel's - ``CDataType`` is fp16; for fp16/bf16 it equals the input dtype. + ``CDataType`` is fp16; for fp16/bf16 it equals the input dtype. ``use_ocp`` + is therefore inert today, but is threaded through so an fp8-output kernel + would not silently pick the wrong encoding. """ - return numpy_dtype_for(output_dtype_for(dtype)) + return numpy_dtype_for(output_dtype_for(dtype), use_ocp=use_ocp) # ============================================================================ @@ -872,7 +889,7 @@ def _bf16_u16_to_fp32(u16: np.ndarray) -> np.ndarray: # --------------------------------------------------------------------------- -# fp8 (E4M3) / bf8 (E5M2) -- FNUZ ("NANOO") encoding used by gfx942/MI300. +# fp8 (E4M3) / bf8 (E5M2) -- both the FNUZ and the OCP encoding. # # numpy has no native 8-bit float, and the C ABI only cares about the 1-byte # memory layout (sizeof(fp8_t) == sizeof(bf8_t) == 1). We carry the value as a @@ -881,32 +898,79 @@ def _bf16_u16_to_fp32(u16: np.ndarray) -> np.ndarray: # NumPy reference multiplies bit-for-bit what the GPU multiplies. The ENCODE only # needs to land on the nearest representable byte. # -# FNUZ format (gfx942): bias = 2^(exp_bits-1); the all-1s exponent is a normal -# number (no Inf), the sole NaN is the sign=1/exp=0/mant=0 byte (0x80), and there -# is no negative zero. gfx950/MI350 uses the OCP fp8 format instead; this codec -# targets the gfx942 default and the OCP path needs separate handling. +# The two formats differ in exponent bias and in what the all-1s exponent means: +# +# FNUZ ("NANOO", gfx942 and every non-gfx950/gfx12 arch) +# bias = 2^(exp_bits-1) (8 / 16); the all-1s exponent is an ordinary normal +# number (no Inf); the sole NaN is 0x80 (sign=1, exp=0, mant=0), so there is +# no negative zero. Max finite: 240 (e4m3fnuz) / 57344 (e5m2fnuz). +# +# OCP (gfx950, gfx12) +# bias = 2^(exp_bits-1) - 1 (7 / 15); 0x80 is -0.0. e4m3fn has no Inf and +# reserves only exp=1111,mant=111 for NaN (max finite 448); e5m2 is IEEE-like, +# with Inf at exp=11111,mant=0 and NaN otherwise (max finite 57344). +# +# Which one applies is a property of the *target arch*, not of the header: the +# host pass of config.hpp always resolves to FNUZ (see dispatcher_common), so the +# selection must be made here, from the arch, and mirrored into the compile flags +# via ocp_arch_defines(). Getting this wrong does not crash -- it silently shifts +# every reference value by a factor of two. # --------------------------------------------------------------------------- +@functools.lru_cache(maxsize=1) +def _default_use_ocp() -> bool: + """Whether the host fp8/bf8 codec should emit OCP bytes for the local GPU. + + Falls back to FNUZ when no arch can be detected -- that is CK's default for + every arch outside gfx950/gfx12, and matches what a header compiled without + an explicit -DCK_TILE_USE_OCP_FP8 would pick on the host pass. + """ + try: + return _fp8_uses_ocp(_get_arch()) + except Exception: # noqa: BLE001 - no GPU / unknown arch + return False + + +def _resolve_use_ocp(use_ocp: Optional[bool]) -> bool: + return _default_use_ocp() if use_ocp is None else bool(use_ocp) + + @functools.lru_cache(maxsize=None) -def _fnuz_decode_table(exp_bits: int, mant_bits: int) -> np.ndarray: - """Build the 256-entry byte -> fp32 value table for an 8-bit FNUZ float. +def _fp8_decode_table(exp_bits: int, mant_bits: int, use_ocp: bool) -> np.ndarray: + """Build the 256-entry byte -> fp32 value table for an 8-bit float format. - The table is a pure function of (exp_bits, mant_bits), so it is cached; the - returned array is marked read-only because callers share the one instance. + The table is a pure function of (exp_bits, mant_bits, use_ocp), so it is + cached; the returned array is marked read-only because callers share the one + instance. """ - bias = (1 << (exp_bits - 1)) + bias = (1 << (exp_bits - 1)) - 1 if use_ocp else (1 << (exp_bits - 1)) mant_max = 1 << mant_bits sign_shift = exp_bits + mant_bits exp_mask = (1 << exp_bits) - 1 + # OCP e4m3fn has no Inf and steals only the all-1s mantissa for NaN; OCP e5m2 + # is IEEE-like. FNUZ has neither -- its all-1s exponent is just a big normal. + ocp_has_inf = use_ocp and mant_bits == 2 table = np.zeros(256, dtype=np.float32) for b in range(256): sign = (b >> sign_shift) & 1 exp = (b >> mant_bits) & exp_mask mant = b & (mant_max - 1) if exp == 0 and mant == 0: - # +0 (0x00); the negative-zero slot (0x80) is the lone NaN. - table[b] = np.float32(np.nan) if sign else np.float32(0.0) + if use_ocp: + table[b] = np.float32(-0.0 if sign else 0.0) + else: + # +0 (0x00); the negative-zero slot (0x80) is the lone NaN. + table[b] = np.float32(np.nan) if sign else np.float32(0.0) + continue + if use_ocp and exp == exp_mask: + if ocp_has_inf: + val = np.inf if mant == 0 else np.nan + elif mant == mant_max - 1: + val = np.nan + else: + val = (1.0 + mant / mant_max) * (2.0 ** (exp - bias)) + table[b] = np.float32(-val if (sign and val == val) else val) continue if exp == 0: val = (mant / mant_max) * (2.0 ** (1 - bias)) # subnormal @@ -917,16 +981,35 @@ def _fnuz_decode_table(exp_bits: int, mant_bits: int) -> np.ndarray: return table -def _fnuz_encode(x: np.ndarray, exp_bits: int, mant_bits: int) -> np.ndarray: - """Encode fp32 -> nearest 8-bit FNUZ float, returned as a uint8 bit pattern.""" - table = _fnuz_decode_table(exp_bits, mant_bits) +def _fnuz_decode_table(exp_bits: int, mant_bits: int) -> np.ndarray: + """FNUZ-only view of :func:`_fp8_decode_table` (kept for existing callers).""" + return _fp8_decode_table(exp_bits, mant_bits, False) + + +def _fp8_encode( + x: np.ndarray, exp_bits: int, mant_bits: int, use_ocp: Optional[bool] = None +) -> np.ndarray: + """Encode fp32 -> nearest 8-bit float, returned as a uint8 bit pattern. + + Rounding is round-to-nearest, ties-to-even -- the hardware convert's rule and + ml_dtypes'. Saturation on overflow is deliberate and is where this encoder + *does* diverge from ml_dtypes: CK's convert clamps to the max finite value + (1e30 -> 448.0 for e4m3), while ``np.astype(ml_dtypes.float8_e4m3fn)`` + produces NaN. Matching the kernel is the point, so the divergence stays; + ``test_gemm_utils.TestFp8EncodeMatchesMlDtypes`` pins both halves of that. + """ + use_ocp = _resolve_use_ocp(use_ocp) + table = _fp8_decode_table(exp_bits, mant_bits, use_ocp) sign_byte = np.uint8(1 << (exp_bits + mant_bits)) # 0x80 - # Positive half (bytes 0..127) holds every non-negative magnitude, sorted. + # Positive half (bytes 0..127) holds every non-negative magnitude. Drop the + # non-finite slots (OCP NaN/Inf) so out-of-range inputs saturate to the max + # finite value instead of rounding to Inf, matching CK's saturating convert. # Compare in float64: for very large inputs the gap between the two top # magnitudes is below fp32 resolution, which would tie and mis-saturate. pos_mag = table[: int(sign_byte)].astype(np.float64) - order = np.argsort(pos_mag) + finite = np.flatnonzero(np.isfinite(pos_mag)) + order = finite[np.argsort(pos_mag[finite])] sorted_mag = pos_mag[order] sorted_byte = order.astype(np.uint8) @@ -937,35 +1020,52 @@ def _fnuz_encode(x: np.ndarray, exp_bits: int, mant_bits: int) -> np.ndarray: raw = np.searchsorted(sorted_mag, ax) hi = np.clip(raw, 0, sorted_mag.size - 1) lo = np.clip(raw - 1, 0, sorted_mag.size - 1) - pick_lo = np.abs(sorted_mag[lo] - ax) <= np.abs(sorted_mag[hi] - ax) + d_lo = np.abs(sorted_mag[lo] - ax) + d_hi = np.abs(sorted_mag[hi] - ax) + # Ties-to-even, not ties-toward-zero. Adjacent entries in `sorted_mag` are + # adjacent byte values (the positive half of both encodings is monotonic in + # the byte, subnormals through normals), so the mantissa LSB is just bit 0 of + # the byte. Breaking ties toward `lo` instead -- as this did -- biases the + # host reference low on exactly the midpoints and disagrees with both the + # hardware convert and ml_dtypes on 63 of 256 e4m3 values. + # + # The degenerate ends are unaffected: raw==0 and raw==size both give lo==hi, + # so the tie branch picks the same index either way and saturation holds. + lo_is_even = (sorted_byte[lo] & np.uint8(1)) == 0 + pick_lo = np.where(d_lo == d_hi, lo_is_even, d_lo < d_hi) chosen = np.where(pick_lo, lo, hi) out = sorted_byte[chosen] - # Apply sign, but never the 0x80 (-0 == NaN) slot: zeros stay +0. + # Apply sign. Under FNUZ 0x80 is the NaN slot rather than -0, so zeros must + # stay positive there; under OCP -0.0 is representable but +0 is equivalent + # for every consumer here, so keep the same rule for both. is_zero = sorted_mag[chosen] == 0 out = np.where((xf < 0) & ~is_zero, out | sign_byte, out) - out = np.where(np.isnan(xf), sign_byte, out) # NaN inputs -> NaN byte + # NaN byte: 0x80 under FNUZ, all-1s exponent and mantissa under OCP. + nan_byte = np.uint8(((1 << exp_bits) - 1) << mant_bits | ((1 << mant_bits) - 1)) \ + if use_ocp else sign_byte + out = np.where(np.isnan(xf), nan_byte, out) return out.astype(np.uint8).reshape(np.shape(x)) -def _fp32_to_fp8_u8(x: np.ndarray) -> np.ndarray: - """Encode fp32 -> fp8 E4M3 (FNUZ) bit pattern in a uint8 array.""" - return _fnuz_encode(x, exp_bits=4, mant_bits=3) +def _fp32_to_fp8_u8(x: np.ndarray, use_ocp: Optional[bool] = None) -> np.ndarray: + """Encode fp32 -> fp8 E4M3 bit pattern in a uint8 array (arch-selected format).""" + return _fp8_encode(x, exp_bits=4, mant_bits=3, use_ocp=use_ocp) -def _fp8_u8_to_fp32(u8: np.ndarray) -> np.ndarray: - """Decode an fp8 E4M3 (FNUZ) bit pattern back to fp32.""" - return _fnuz_decode_table(4, 3)[u8.astype(np.intp)] +def _fp8_u8_to_fp32(u8: np.ndarray, use_ocp: Optional[bool] = None) -> np.ndarray: + """Decode an fp8 E4M3 bit pattern back to fp32 (arch-selected format).""" + return _fp8_decode_table(4, 3, _resolve_use_ocp(use_ocp))[u8.astype(np.intp)] -def _fp32_to_bf8_u8(x: np.ndarray) -> np.ndarray: - """Encode fp32 -> bf8 E5M2 (FNUZ) bit pattern in a uint8 array.""" - return _fnuz_encode(x, exp_bits=5, mant_bits=2) +def _fp32_to_bf8_u8(x: np.ndarray, use_ocp: Optional[bool] = None) -> np.ndarray: + """Encode fp32 -> bf8 E5M2 bit pattern in a uint8 array (arch-selected format).""" + return _fp8_encode(x, exp_bits=5, mant_bits=2, use_ocp=use_ocp) -def _bf8_u8_to_fp32(u8: np.ndarray) -> np.ndarray: - """Decode a bf8 E5M2 (FNUZ) bit pattern back to fp32.""" - return _fnuz_decode_table(5, 2)[u8.astype(np.intp)] +def _bf8_u8_to_fp32(u8: np.ndarray, use_ocp: Optional[bool] = None) -> np.ndarray: + """Decode a bf8 E5M2 bit pattern back to fp32 (arch-selected format).""" + return _fp8_decode_table(5, 2, _resolve_use_ocp(use_ocp))[u8.astype(np.intp)] def _fp32_to_fp8_ocp_u8(x): @@ -1031,12 +1131,16 @@ class GpuGemmRunner: numpy arrays straight to the .so. """ - def __init__(self, lib_path: Path): + def __init__(self, lib_path: Path, arch: Optional[str] = None): self.lib = GemmDispatcherLib(lib_path) if not self.lib.initialize(): raise RuntimeError(f"Failed to initialize dispatcher .so: {lib_path}") names = self.lib.kernel_names self._kernel_name = names[0] if names else "unknown" + # fp8/bf8 encoding must match the arch the .so was compiled for. Running + # a .so requires the matching GPU, so the local arch is the right default; + # `arch` is the explicit pin for callers that already know it. + self._use_ocp = _fp8_uses_ocp(arch) if arch else None @property def kernel_name(self) -> str: @@ -1074,13 +1178,11 @@ def run( A_h = _fp32_to_bf16_u16(A_lay) B_h = _fp32_to_bf16_u16(B_lay) elif dtype == "fp8": - _enc = _fp32_to_fp8_ocp_u8 if _use_ocp_fp8() else _fp32_to_fp8_u8 - A_h = _enc(A_lay) - B_h = _enc(B_lay) + A_h = _fp32_to_fp8_u8(A_lay, use_ocp=self._use_ocp) + B_h = _fp32_to_fp8_u8(B_lay, use_ocp=self._use_ocp) elif dtype == "bf8": - _enc = _fp32_to_bf8_ocp_u8 if _use_ocp_fp8() else _fp32_to_bf8_u8 - A_h = _enc(A_lay) - B_h = _enc(B_lay) + A_h = _fp32_to_bf8_u8(A_lay, use_ocp=self._use_ocp) + B_h = _fp32_to_bf8_u8(B_lay, use_ocp=self._use_ocp) elif dtype == "int8": A_h = np.ascontiguousarray(A_lay, dtype=np.int8) B_h = np.ascontiguousarray(B_lay, dtype=np.int8) @@ -1140,7 +1242,13 @@ class GpuGroupedGemmRunner: overrun the host buffer (heap corruption). See :func:`output_numpy_dtype_for`. """ - def __init__(self, lib_path: Path, dtype: str = "fp16", layout: str = "rcr"): + def __init__( + self, + lib_path: Path, + dtype: str = "fp16", + layout: str = "rcr", + arch: Optional[str] = None, + ): self.lib = GemmDispatcherLib(lib_path) if not self.lib.initialize(): raise RuntimeError( @@ -1149,10 +1257,13 @@ def __init__(self, lib_path: Path, dtype: str = "fp16", layout: str = "rcr"): names = self.lib.kernel_names self._kernel_name = names[0] if names else "unknown" self._dtype = dtype + # fp8/bf8 encoding must match the arch the .so was compiled for; see + # GpuGemmRunner.__init__. None -> resolve from the local GPU. + self._use_ocp = _fp8_uses_ocp(arch) if arch else None # A/B (input) codec vs C (output) codec: they differ for fp8/bf8 # (output is fp16), so keep them distinct to size the C buffer correctly. - self._np_dtype = numpy_dtype_for(dtype) - self._c_np_dtype = output_numpy_dtype_for(dtype) + self._np_dtype = numpy_dtype_for(dtype, use_ocp=self._use_ocp) + self._c_np_dtype = output_numpy_dtype_for(dtype, use_ocp=self._use_ocp) if len(layout) != 3 or any(ch not in ("r", "c") for ch in layout): raise ValueError(f"layout must be a 3-char r/c string, got {layout!r}") self._layout = layout @@ -1891,6 +2002,12 @@ def _build_compile_jobs( "-D__HIP_PLATFORM_AMD__", f"--offload-arch={config.gfx_arch}", f'-DGFX_ARCH="{config.gfx_arch}"', + # Pin the fp8/bf8 encoding so BOTH compiler passes agree. Without this the + # device pass of config.hpp sees __gfx950__ and picks OCP while the host + # pass falls back to FNUZ -- and the numpy reference, which follows the + # arch (see numpy_dtype_for), then disagrees with the kernel by a factor + # of two. FNUZ archs need no define; the list is empty for them. + *_ocp_arch_defines(config.gfx_arch), # Match Tile Engine's AMDGPU codegen flags exactly (see variant_flags / # _tile_engine_codegen_flags). Without them the kernel is compiled with # different inlining/register allocation, which changes occupancy; @@ -2434,6 +2551,12 @@ def _as_list(v, default): tm // m_div, tn // n_div, wtm, wtn ): continue + # Stream-K supports only the cshuffle epilogue, and the codegen skips + # anything else (unified_gemm_codegen.py, GemmVariant.STREAM_K branch). + # Without this gate expand_sweep hands back configs whose header is + # never emitted, which surfaces downstream as a spurious build failure. + if variant == "stream_k" and epi != "cshuffle": + continue for (m_na, m_nb, m_nd, m_aop, m_bop, m_cdeop) in mabd_combos: for ew_op, md_nd in md_combos: diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py index 5248d34b5ab8..8d6b440c07bc 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_abquant_utils.py @@ -54,6 +54,11 @@ sys.path.insert(0, _codegen_dir) from codegen_common import make_abquant_kernel_name # noqa: E402 +_python_dir = str(Path(__file__).parent) +if _python_dir not in sys.path: + sys.path.insert(0, _python_dir) +from dispatcher_common import arch_feature_defines # noqa: E402 + _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" @@ -546,9 +551,16 @@ def _compile_abquant_kernel( ) -> bool: ck_include = _get_ck_include_dir() + # Arch-specific defines: gfx950 uses OCP fp8 (not FNUZ) and native MX support. + # These mirror the CMakeLists.txt definitions that are normally injected by CMake + # but are absent in the standalone hipcc build path. Without them the host pass of + # config.hpp silently falls back to FNUZ while the device pass picks OCP. + arch_defines = arch_feature_defines(gfx_arch) + cmd = [hipcc] + _HIPCC_BASE_FLAGS + [ f"--offload-arch={gfx_arch}", f"-DGFX_ARCH=\"{gfx_arch}\"", + *arch_defines, "-include", str(hpp_path), str(_CTYPES_LIB_SRC), "-o", str(so_path), @@ -593,6 +605,18 @@ def setup_multiple_abquant_dispatchers( if not configs: return [] + _ABQUANT_SUPPORTED_LAYOUTS = ("rcr",) + bad = [c for c in configs if c.layout not in _ABQUANT_SUPPORTED_LAYOUTS] + if bad: + raise ValueError( + f"grouped_gemm_abquant bridge only supports layouts " + f"{_ABQUANT_SUPPORTED_LAYOUTS}; " + f"got unsupported layouts: " + f"{sorted({c.layout for c in bad})}. " + f"Non-rcr layout support requires changes to the ctypes stride " + f"derivation in grouped_gemm_abquant_ctypes_lib.cpp (plan Step 7)." + ) + arch = gfx_arch or _detect_gpu_arch() base_dir = output_dir or Path(tempfile.mkdtemp(prefix="abquant_dispatcher_")) base_dir.mkdir(parents=True, exist_ok=True) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py index 8ae843cc0c43..aebc66ab821a 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_aquant_utils.py @@ -54,6 +54,11 @@ sys.path.insert(0, _codegen_dir) from codegen_common import make_aquant_kernel_name # noqa: E402 +_python_dir = str(Path(__file__).parent) +if _python_dir not in sys.path: + sys.path.insert(0, _python_dir) +from dispatcher_common import arch_feature_defines # noqa: E402 + _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" @@ -498,9 +503,15 @@ def _compile_aquant_kernel( """Compile a generated .hpp into a .so via hipcc. Returns True on success.""" ck_include = _get_ck_include_dir() + # Arch-specific defines: gfx950 uses OCP fp8 (not FNUZ) and native MX support. + # These mirror the CMakeLists.txt definitions that are normally injected by CMake + # but are absent in the standalone hipcc build path. + arch_defines = arch_feature_defines(gfx_arch) + cmd = [hipcc] + _HIPCC_BASE_FLAGS + [ f"--offload-arch={gfx_arch}", f"-DGFX_ARCH=\"{gfx_arch}\"", + *arch_defines, "-include", str(hpp_path), str(_CTYPES_LIB_SRC), "-o", str(so_path), @@ -553,6 +564,18 @@ def setup_multiple_aquant_dispatchers( if not configs: return [] + _AQUANT_SUPPORTED_LAYOUTS = ("rcr",) + bad = [c for c in configs if c.layout not in _AQUANT_SUPPORTED_LAYOUTS] + if bad: + raise ValueError( + f"grouped_gemm_aquant bridge only supports layouts " + f"{_AQUANT_SUPPORTED_LAYOUTS}; " + f"got unsupported layouts: " + f"{sorted({c.layout for c in bad})}. " + f"Non-rcr layout support requires changes to the ctypes stride " + f"derivation in grouped_gemm_aquant_ctypes_lib.cpp (plan Step 7)." + ) + arch = gfx_arch or _detect_gpu_arch() base_dir = output_dir or Path(tempfile.mkdtemp(prefix="aquant_dispatcher_")) base_dir.mkdir(parents=True, exist_ok=True) @@ -737,6 +760,18 @@ def default_bf8i4_config( ) +def _preshuffleaq_warp_tile_k(gfx_arch: str) -> int: + """warp_tile_k for preshuffleAQ (FlatMM): 128 on gfx950, 64 on gfx942. + + get_k_warp_tile() in tile_gemm_shape.hpp branches on CK_GFX950_SUPPORT: + FlatMM selects 128 on gfx950 (the 16x16x128 f8f6f4 MFMA) and 64 on gfx942 + (the 16x16x64 fp8 MFMA). A mismatched value compiles but silently produces + zeros rather than a build error, so this helper and arch_feature_defines + must stay in sync. + """ + return 128 if gfx_arch.startswith("gfx950") else 64 + + def default_fp8_preshuffleaq_config( quant_group_k: int = 128, quant_group_m: int = 1, @@ -745,7 +780,7 @@ def default_fp8_preshuffleaq_config( """fp8 preshuffle-AQ config (GemmConfigPreshuffleQuantDecode, CompV3 pipeline). APreshuffleQuant=true, BPreshuffleQuant=true — both scale tensors are preshuffled. - Tile: 16x64x256, warp_tile_k=128 (gfx950 FlatMM, 8-bit float). + Tile: 16x64x256. warp_tile_k is arch-aware: 128 on gfx950 (FlatMM), 64 on gfx942. """ return AQuantKernelConfig( variant_key="fp8", @@ -755,7 +790,7 @@ def default_fp8_preshuffleaq_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_preshuffleaq_warp_tile_k(gfx_arch), quant_group_m=quant_group_m, quant_group_n=1, quant_group_k=quant_group_k, @@ -778,7 +813,7 @@ def default_bf8_preshuffleaq_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_preshuffleaq_warp_tile_k(gfx_arch), quant_group_m=quant_group_m, quant_group_n=1, quant_group_k=quant_group_k, diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py index 6f68a5b37573..0692d5fd6bcc 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_bquant_utils.py @@ -52,6 +52,11 @@ sys.path.insert(0, _codegen_dir) from codegen_common import make_bquant_kernel_name # noqa: E402 +_python_dir = str(Path(__file__).parent) +if _python_dir not in sys.path: + sys.path.insert(0, _python_dir) +from dispatcher_common import arch_feature_defines # noqa: E402 + _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" @@ -524,11 +529,7 @@ def _compile_bquant_kernel( # Arch-specific defines: gfx950 uses OCP fp8 (not FNUZ) and native MX support. # These mirror the CMakeLists.txt definitions that are normally injected by CMake # but are absent in the standalone hipcc build path. - arch_defines = [] - if "gfx12" in gfx_arch or "gfx950" in gfx_arch: - arch_defines += ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"] - if "gfx950" in gfx_arch: - arch_defines += ["-DCK_USE_NATIVE_MX_SUPPORT", "-DCK_GFX950_SUPPORT"] + arch_defines = arch_feature_defines(gfx_arch) compile_cmd = [hipcc, "-c", "-fPIC", "-O3", "-std=c++17", "-DCK_TILE_SINGLE_KERNEL_INCLUDE", "-w", @@ -620,6 +621,21 @@ def setup_multiple_bquant_dispatchers( if not configs: return [] + # The bquant ctypes lib asserts packed rcr strides and will raise at runtime + # for any other layout. Fail loudly here, before any build work, so the + # caller knows which config is wrong rather than getting a ctypes error later. + _BQUANT_SUPPORTED_LAYOUTS = ("rcr",) + bad = [c for c in configs if c.layout not in _BQUANT_SUPPORTED_LAYOUTS] + if bad: + raise ValueError( + f"grouped_gemm_bquant bridge only supports layouts " + f"{_BQUANT_SUPPORTED_LAYOUTS}; " + f"got unsupported layouts: " + f"{sorted({c.layout for c in bad})}. " + f"Non-rcr layout support requires changes to the ctypes stride " + f"derivation in grouped_gemm_bquant_ctypes_lib.cpp (plan Step 7)." + ) + arch = gfx_arch or _detect_gpu_arch() base_dir = output_dir or Path(tempfile.mkdtemp(prefix="bquant_dispatcher_")) base_dir.mkdir(parents=True, exist_ok=True) @@ -791,6 +807,32 @@ def expand_bquant_sweep( # ============================================================================= +# NOTE: the gfx942 branch of both helpers below is correct but not yet exercised +# by CI -- the bquant lane is gated to gfx950 (ck.groovy) and +# test_bquant_gpu_correctness.SUPPORTED_ARCHS matches it. They are written +# arch-aware now so that enabling gfx942 is a one-line change in those two +# places; until that follow-up lands, treat the gfx942 values as unvalidated. + + +def _fp8_warp_tile_k(gfx_arch: str) -> int: + """warp_tile_k for fp8/bf8 compv3 (non-FlatMM): 128 on gfx950, 32 on gfx942. + + get_k_warp_tile() in tile_gemm_shape.hpp returns 128 + on gfx950 (CK_GFX950_SUPPORT set) and 32 on gfx942 (is_8bit_float, standard + MFMA, mfma_f32_16x16x32_fp8_fp8). A mismatched value compiles but silently + returns zeros rather than a build error. + """ + return 128 if gfx_arch.startswith("gfx950") else 32 + + +def _preshuffleb_warp_tile_k(gfx_arch: str) -> int: + """warp_tile_k for preshuffleB FlatMM: 128 on gfx950, 64 on gfx942. + + Mirrors abquant's _preshuffleb_warp_tile_k; same CK_GFX950_SUPPORT branch. + """ + return 128 if gfx_arch.startswith("gfx950") else 64 + + def default_fp8_config( quant_group_k: int = 128, quant_group_n: int = 1, @@ -798,8 +840,9 @@ def default_fp8_config( ) -> BQuantKernelConfig: """Return the default fp8 BQuant config (tile = 16x64x256, warp = 1x4x1). - WarpTileK=128: on gfx950 get_k_warp_tile() returns 128 - (is_8bit_float=true, M_Warp_Tile!=32 → K_warp=128). Using 16 causes zero output. + warp_tile_k is arch-aware: 128 on gfx950 (get_k_warp_tile returns 128 for + CK_GFX950_SUPPORT + is_8bit_float + M_Warp_Tile=16), 32 on gfx942 (standard + mfma_f32_16x16x32_fp8_fp8). A mismatched value silently returns zeros. """ return BQuantKernelConfig( variant_key="fp8", @@ -809,7 +852,7 @@ def default_fp8_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_fp8_warp_tile_k(gfx_arch), quant_group_m=1, quant_group_n=quant_group_n, quant_group_k=quant_group_k, @@ -824,7 +867,7 @@ def default_bf8_config( ) -> BQuantKernelConfig: """Return the default bf8 BQuant config (tile = 16x64x256, warp = 1x4x1). - WarpTileK=128: on gfx950 get_k_warp_tile() returns 128. + warp_tile_k is arch-aware: 128 on gfx950, 32 on gfx942. See default_fp8_config. """ return BQuantKernelConfig( variant_key="bf8", @@ -834,7 +877,7 @@ def default_bf8_config( scheduler="intrawave", tile_m=16, tile_n=64, tile_k=256, warp_m=1, warp_n=4, warp_k=1, - warp_tile_m=16, warp_tile_n=16, warp_tile_k=128, + warp_tile_m=16, warp_tile_n=16, warp_tile_k=_fp8_warp_tile_k(gfx_arch), quant_group_m=1, quant_group_n=quant_group_n, quant_group_k=quant_group_k, @@ -931,8 +974,12 @@ def default_fp8_preshuffleb_config( quant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> BQuantKernelConfig: - """fp8 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill).""" - return _preshuffleb_config("fp8", warp_tile_k=128, quant_group_k=quant_group_k, + """fp8 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill). + + warp_tile_k is arch-aware via _preshuffleb_warp_tile_k: 128 on gfx950, 64 on gfx942. + """ + return _preshuffleb_config("fp8", warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), + quant_group_k=quant_group_k, quant_group_n=quant_group_n, gfx_arch=gfx_arch) @@ -941,8 +988,12 @@ def default_bf8_preshuffleb_config( quant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> BQuantKernelConfig: - """bf8 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill).""" - return _preshuffleb_config("bf8", warp_tile_k=128, quant_group_k=quant_group_k, + """bf8 preshuffle_b prefill config (GemmConfigPreshuffleB_BQuant_Prefill). + + warp_tile_k is arch-aware via _preshuffleb_warp_tile_k: 128 on gfx950, 64 on gfx942. + """ + return _preshuffleb_config("bf8", warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), + quant_group_k=quant_group_k, quant_group_n=quant_group_n, gfx_arch=gfx_arch) @@ -1015,8 +1066,12 @@ def default_fp8_preshufflequant_config( quant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> BQuantKernelConfig: - """fp8 preshuffle_bquant prefill config (GemmConfigPreshuffleBQuantPrefill).""" - return _preshufflequant_config("fp8", warp_tile_k=128, quant_group_k=quant_group_k, + """fp8 preshuffle_bquant prefill config (GemmConfigPreshuffleBQuantPrefill). + + warp_tile_k is arch-aware via _preshuffleb_warp_tile_k: 128 on gfx950, 64 on gfx942. + """ + return _preshufflequant_config("fp8", warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), + quant_group_k=quant_group_k, quant_group_n=quant_group_n, gfx_arch=gfx_arch) @@ -1025,8 +1080,12 @@ def default_bf8_preshufflequant_config( quant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> BQuantKernelConfig: - """bf8 preshuffle_bquant prefill config (GemmConfigPreshuffleBQuantPrefill).""" - return _preshufflequant_config("bf8", warp_tile_k=128, quant_group_k=quant_group_k, + """bf8 preshuffle_bquant prefill config (GemmConfigPreshuffleBQuantPrefill). + + warp_tile_k is arch-aware via _preshuffleb_warp_tile_k: 128 on gfx950, 64 on gfx942. + """ + return _preshufflequant_config("bf8", warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), + quant_group_k=quant_group_k, quant_group_n=quant_group_n, gfx_arch=gfx_arch) @@ -1091,8 +1150,12 @@ def default_fp8_preshuffleb_bquant_config( quant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> BQuantKernelConfig: - """fp8 preshuffle_b+preshuffle_bquant config (GemmConfigPreshuffleB_PreshuffleBQuant_Prefill).""" - return _preshuffleb_bquant_config("fp8", warp_tile_k=128, quant_group_k=quant_group_k, + """fp8 preshuffle_b+preshuffle_bquant config (GemmConfigPreshuffleB_PreshuffleBQuant_Prefill). + + warp_tile_k is arch-aware via _preshuffleb_warp_tile_k: 128 on gfx950, 64 on gfx942. + """ + return _preshuffleb_bquant_config("fp8", warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), + quant_group_k=quant_group_k, quant_group_n=quant_group_n, gfx_arch=gfx_arch) @@ -1101,8 +1164,12 @@ def default_bf8_preshuffleb_bquant_config( quant_group_n: int = 1, gfx_arch: str = _DEFAULT_GFX_ARCH, ) -> BQuantKernelConfig: - """bf8 preshuffle_b+preshuffle_bquant config.""" - return _preshuffleb_bquant_config("bf8", warp_tile_k=128, quant_group_k=quant_group_k, + """bf8 preshuffle_b+preshuffle_bquant config. + + warp_tile_k is arch-aware via _preshuffleb_warp_tile_k: 128 on gfx950, 64 on gfx942. + """ + return _preshuffleb_bquant_config("bf8", warp_tile_k=_preshuffleb_warp_tile_k(gfx_arch), + quant_group_k=quant_group_k, quant_group_n=quant_group_n, gfx_arch=gfx_arch) diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py index 7bb7d1338770..34a086480b06 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_rowcolquant_utils.py @@ -54,6 +54,11 @@ make_rowcolquant_kernel_name, ) +_python_dir = str(Path(__file__).parent) +if _python_dir not in sys.path: + sys.path.insert(0, _python_dir) +from dispatcher_common import arch_feature_defines # noqa: E402 + _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" @@ -500,11 +505,7 @@ def _compile_rowcolquant_kernel( obj_path = so_path.with_suffix(".o") - arch_defines = [] - if "gfx12" in gfx_arch or "gfx950" in gfx_arch: - arch_defines += ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"] - if "gfx950" in gfx_arch: - arch_defines += ["-DCK_USE_NATIVE_MX_SUPPORT", "-DCK_GFX950_SUPPORT"] + arch_defines = arch_feature_defines(gfx_arch) compile_cmd = [hipcc, "-c", "-fPIC", "-O3", "-std=c++17", "-DCK_TILE_SINGLE_KERNEL_INCLUDE", "-w", diff --git a/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py b/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py index e240b0f9604a..ef3f02e33b42 100644 --- a/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py +++ b/projects/composablekernel/dispatcher/python/grouped_gemm_tensorquant_utils.py @@ -54,6 +54,11 @@ make_tensorquant_kernel_name, ) +_python_dir = str(Path(__file__).parent) +if _python_dir not in sys.path: + sys.path.insert(0, _python_dir) +from dispatcher_common import arch_feature_defines # noqa: E402 + _DEFAULT_HIPCC = "hipcc" _DEFAULT_GFX_ARCH = "gfx950" @@ -488,11 +493,7 @@ def _compile_tensorquant_kernel( obj_path = so_path.with_suffix(".o") - arch_defines = [] - if "gfx12" in gfx_arch or "gfx950" in gfx_arch: - arch_defines += ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"] - if "gfx950" in gfx_arch: - arch_defines += ["-DCK_USE_NATIVE_MX_SUPPORT", "-DCK_GFX950_SUPPORT"] + arch_defines = arch_feature_defines(gfx_arch) compile_cmd = [hipcc, "-c", "-fPIC", "-O3", "-std=c++17", "-DCK_TILE_SINGLE_KERNEL_INCLUDE", "-w", diff --git a/projects/composablekernel/dispatcher/tests/CMakeLists.txt b/projects/composablekernel/dispatcher/tests/CMakeLists.txt index 899c2081821e..bbf6019e0734 100644 --- a/projects/composablekernel/dispatcher/tests/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/tests/CMakeLists.txt @@ -13,12 +13,20 @@ find_package(Python3 COMPONENTS Interpreter REQUIRED) # Bridge on-device GPU correctness test (SKIPs cleanly without a GPU/built lib). add_test( NAME dispatcher_test_preshuffle_gpu_correctness - COMMAND ${Python3_EXECUTABLE} -m unittest discover -s ${CMAKE_CURRENT_SOURCE_DIR} -p test_preshuffle_gpu_correctness.py -v + # Run the script directly: this module is script-style (a bare function plus + # a main() guard), not unittest.TestCase, so `unittest discover` collects 0 + # tests -- and unittest exits NONZERO on "NO TESTS RAN", turning an + # unrunnable box into a hard ctest failure. main() skips (exit 77) without a + # GPU/built lib and returns nonzero only on a real correctness failure. + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_preshuffle_gpu_correctness.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ) set_tests_properties(dispatcher_test_preshuffle_gpu_correctness PROPERTIES LABELS "dispatcher;python;bridge;gpu" TIMEOUT 600 + # Without this a CPU-only runner reports a green PASS for a test that never + # reached the GPU; 77 makes ctest record it as skipped instead. + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) @@ -27,7 +35,7 @@ add_test( NAME dispatcher_test_batched_gemm_gpu_correctness # Run the script directly: this module is script-style (a parameterized # function + main() guard), not unittest.TestCase, so `unittest discover` - # would collect 0 tests and pass vacuously. main() skips (exit 0) without a + # would collect 0 tests and pass vacuously. main() skips (exit 77) without a # GPU/built lib and returns nonzero only on a real correctness failure. COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_batched_gemm_gpu_correctness.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} @@ -35,6 +43,9 @@ add_test( set_tests_properties(dispatcher_test_batched_gemm_gpu_correctness PROPERTIES LABELS "dispatcher;python;bridge;gpu" TIMEOUT 600 + # Without this a CPU-only runner reports a green PASS for a test that never + # reached the GPU; 77 makes ctest record it as skipped instead. + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) @@ -43,7 +54,7 @@ add_test( NAME dispatcher_test_batched_contraction_gpu_correctness # Run the script directly: this module is script-style (a parameterized # function + main() guard), not unittest.TestCase, so `unittest discover` - # would collect 0 tests and pass vacuously. main() skips (exit 0) without a + # would collect 0 tests and pass vacuously. main() skips (exit 77) without a # GPU/built lib and returns nonzero only on a real correctness failure. COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_batched_contraction_gpu_correctness.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} @@ -51,18 +62,25 @@ add_test( set_tests_properties(dispatcher_test_batched_contraction_gpu_correctness PROPERTIES LABELS "dispatcher;python;bridge;gpu" TIMEOUT 600 + # Without this a CPU-only runner reports a green PASS for a test that never + # reached the GPU; 77 makes ctest record it as skipped instead. + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) # Bridge on-device GPU correctness test (SKIPs cleanly without a GPU/built lib). add_test( NAME dispatcher_test_tensor_quant_gpu_correctness - COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_tensor_quant_gpu_correctness.py + # Run the script directly rather than under pytest: pytest exits 0 when every + # test skips, so a CPU-only or non-fp8 runner reported a green PASS for work + # that never reached the GPU. main() exits 77 on those boxes instead. + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_tensor_quant_gpu_correctness.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ) set_tests_properties(dispatcher_test_tensor_quant_gpu_correctness PROPERTIES LABELS "dispatcher;python;bridge;gpu" TIMEOUT 600 + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) @@ -71,7 +89,7 @@ add_test( NAME dispatcher_test_contraction_multi_abd_gpu_correctness # Run the script directly: this module is script-style (a parameterized # function + main() guard), not unittest.TestCase, so `unittest discover` - # would collect 0 tests and pass vacuously. main() skips (exit 0) without a + # would collect 0 tests and pass vacuously. main() skips (exit 77) without a # GPU/hipcc and returns nonzero only on a real correctness failure. COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_contraction_multi_abd_gpu_correctness.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} @@ -81,6 +99,9 @@ set_tests_properties(dispatcher_test_contraction_multi_abd_gpu_correctness PROPE # Builds a kernel from source via hipcc before running; needs more than the # 600s the sibling bridges use. TIMEOUT 1800 + # Without this a CPU-only runner reports a green PASS for a test that never + # reached the GPU; 77 makes ctest record it as skipped instead. + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) @@ -114,6 +135,19 @@ set_tests_properties(dispatcher_test_arch_filter_constraints PROPERTIES ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) +# dispatcher_common unit tests: path helpers, validators, and -- since this is +# the only place they are covered -- the fp8-encoding arch predicates +# (fp8_uses_ocp / ocp_arch_defines / arch_feature_defines). Those decide whether +# -DCK_USE_OCP_FP8 reaches hipcc; getting them wrong makes host and device +# disagree on the fp8 encoding with no build error, so they need a test that +# actually runs. This file was previously not registered at all. +add_test( + NAME dispatcher_test_dispatcher_common + COMMAND ${Python3_EXECUTABLE} -m unittest discover -s ${CMAKE_CURRENT_SOURCE_DIR} -p test_dispatcher_common.py -v + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) + +set_tests_properties(dispatcher_test_dispatcher_common PROPERTIES # Architecture-aware LDS staging budget regression test add_test( NAME dispatcher_test_lds_capacity_arch_aware @@ -147,24 +181,102 @@ set_tests_properties(dispatcher_test_gemm_utils PROPERTIES # numpy fp32 reference (scale-relative tolerance). They SKIP cleanly (no fail) # when no GPU / hipcc / static lib is present, so they cost nothing on a # CPU-only runner and only run on a GPU runner. mx arch-gates to gfx950. +# +# These three are unittest.TestCase-based and skip themselves in setUp, so they +# are driven through run_unittest_77.py rather than `python -m unittest +# discover` directly. discover exits 0 when every test skipped, which made ctest +# report Passed for a run that asserted nothing -- a CPU-only runner showed the +# same green tick as a real GPU run. The driver keeps the discovery semantics and +# only maps an all-skipped run onto 77, which SKIP_RETURN_CODE turns into +# "Skipped". multi_d is script-style-77 in its own right and is registered below. # ============================================================================= foreach(_gpu_test test_grouped_gemm_gpu_correctness test_multi_abd_gpu_correctness - test_multi_d_gpu_correctness test_mx_gemm_gpu_correctness) add_test( NAME dispatcher_${_gpu_test} - COMMAND ${Python3_EXECUTABLE} -m unittest discover -s ${CMAKE_CURRENT_SOURCE_DIR} -p ${_gpu_test}.py -v + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/run_unittest_77.py ${_gpu_test}.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ) set_tests_properties(dispatcher_${_gpu_test} PROPERTIES LABELS "dispatcher;python;bridge;gpu" TIMEOUT 600 + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) endforeach() +# Multi-D GEMM: E = MultiDAdd(A @ B, D0, D1) over all four A/B layouts. +# Invoked as a script, not via `unittest discover`: it returns 77 on a missing +# GPU / static lib / hipcc, which SKIP_RETURN_CODE turns into ctest "Skipped". +# The discover path above exits 0 on an all-skipped run and would report Passed. +# +# TIMEOUT is 1800, not the 600 its neighbours use: the layout sweep is four +# separate hipcc builds (setup_multiple_gemm_dispatchers runs once per layout), +# so the wall-clock budget is roughly 4x theirs. +add_test( + NAME dispatcher_test_multi_d_gpu_correctness + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_multi_d_gpu_correctness.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/.. +) + +set_tests_properties(dispatcher_test_multi_d_gpu_correctness PROPERTIES + LABELS "dispatcher;python;bridge;gpu" + TIMEOUT 1800 + SKIP_RETURN_CODE 77 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + +# Config search-space sweep: enumerate configs for a variant, build and run a +# sampled subset, verify each against a numpy reference. Previously registered +# nowhere -- it was reachable only from the Jenkins dispatcher lane, so a local +# `ctest` never exercised it and neither did any pre-merge gate. +# +# --budget 16, one per (dtype x layout) stratum: the sampler splits the budget +# evenly across the 16 strata the standard variant enumerates, so anything below +# 16 leaves whole strata unsampled on any given run. At --budget 4 the run was +# not a small slice of the space but a slice of a quarter of it, and no single +# run ever compiled an fp8 or bf8 config. Each sampled config is its own hipcc +# build, hence the 1800s timeout. +# +# The runner returns 77 when no GPU is present and --arch is not given, so +# SKIP_RETURN_CODE keeps a CPU-only runner honest instead of green. That skip +# path and the sampler's stratum coverage are both pinned by +# test_search_space_sampling.py, which needs no GPU. +add_test( + NAME dispatcher_test_gemm_search_space + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_gemm_search_space.py --budget 16 + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) + +set_tests_properties(dispatcher_test_gemm_search_space PROPERTIES + LABELS "dispatcher;python;gemm;gpu" + TIMEOUT 1800 + SKIP_RETURN_CODE 77 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + +# Unit tests for the search-space runner's sampler and skip path. Pure CPU, no +# hipcc and no built kernels, so unlike the sweep above these run on every +# machine and gate every build. +# +# Registered separately because neither behaviour is observable from the sweep's +# own result: a sampler that draws from only a quarter of the strata still +# passes every config it does run, and a broken skip path shows up as a green +# GPU lane, not a red one. Both regressed silently once already. +add_test( + NAME dispatcher_test_search_space_sampling + COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_search_space_sampling.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) + +set_tests_properties(dispatcher_test_search_space_sampling PROPERTIES + LABELS "dispatcher;python;gemm;validation" + TIMEOUT 60 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + # ============================================================================= # RowColQuant / TensorQuant bridge tests. # @@ -174,8 +286,13 @@ endforeach() # silently collect zero tests and report success. # # The *_utils suites are pure Python (naming, config, codegen text) and run -# anywhere. The *_gpu_correctness suites arch-gate to gfx942/gfx950 and skip -# cleanly without a GPU or hipcc, so they cost nothing on a CPU-only runner. +# anywhere, and stay on `-m pytest`. +# +# The *_gpu_correctness suites are dual-mode: importable by pytest, but also +# runnable as scripts with a `main()` that returns 77 to skip. ctest invokes +# them as scripts, because `-m pytest` reports exit 0 -- i.e. PASSED -- when +# every test in the file skips, which would hide a missing GPU behind a green +# result. Run as scripts with SKIP_RETURN_CODE 77 they report Skipped instead. # ============================================================================= foreach(_quant_test test_grouped_gemm_rowcolquant_utils @@ -197,12 +314,13 @@ foreach(_quant_gpu_test test_tensorquant_gpu_correctness) add_test( NAME dispatcher_${_quant_gpu_test} - COMMAND ${Python3_EXECUTABLE} -m pytest ${CMAKE_CURRENT_SOURCE_DIR}/${_quant_gpu_test}.py -v + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/${_quant_gpu_test}.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ) set_tests_properties(dispatcher_${_quant_gpu_test} PROPERTIES LABELS "dispatcher;python;bridge;quant;gpu" TIMEOUT 600 + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) endforeach() @@ -331,11 +449,88 @@ add_test( set_tests_properties(dispatcher_test_streamk_registry PROPERTIES LABELS "dispatcher;python;streamk;gpu" + # 96 cases (4 dtypes x 4 layouts x 3 strategies x 2 shapes), each a separate + # hipcc build of the driver -- compile time dominates, so the budget scales + # with the matrix, not with GPU speed. Measured end-to-end at ~55 min on a + # local box; 900 was not survivable by any arch. + TIMEOUT 4200 + SKIP_RETURN_CODE 77 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + +# GPU numeric-correctness tests for the quantized / Stream-K bridges. +# +# These are script-style (argparse + `return 77` to skip), not unittest, so each +# needs its own add_test + SKIP_RETURN_CODE pair -- they cannot join the +# `unittest discover` foreach above, which has no skip code and a fixed timeout. +# +# Deliberately NOT given a --gfx argument: each test hipcc-compiles kernels with +# --offload-arch and then runs them, so the arch must match the PHYSICAL GPU. +# rocminfo auto-detection is the right source for that; the CMake-configured +# GPU_TARGETS may legitimately differ from the card in the box. Every one of +# these exits 77 (Skipped, not Passed) when no GPU is present or the detected +# arch is not the one the operator supports. + +# Stream-K: atomic / linear / tree reduction through the ctypes bridge (gfx942). +# 12 hipcc builds plus a 4096-cubed reference, hence the larger timeout. +add_test( + NAME dispatcher_test_streamk_gpu_correctness + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_streamk_gpu_correctness.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/.. +) + +set_tests_properties(dispatcher_test_streamk_gpu_correctness PROPERTIES + LABELS "dispatcher;python;bridge;gpu" TIMEOUT 900 SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) +# AQuant grouped GEMM: fp8 / bf8 per-(m,k)-group A-side scales (gfx950). +add_test( + NAME dispatcher_test_aquant_gpu_correctness + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_aquant_gpu_correctness.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/.. +) + +set_tests_properties(dispatcher_test_aquant_gpu_correctness PROPERTIES + LABELS "dispatcher;python;bridge;gpu" + TIMEOUT 600 + SKIP_RETURN_CODE 77 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + +# BQuant grouped GEMM: C4 (fp8/bf8) + H3 (mx_*) B-side scales (gfx950). +# The test already existed but ran only from the Jenkins gfx950 lane; this makes +# it reachable from ctest too. +add_test( + NAME dispatcher_test_bquant_gpu_correctness + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_bquant_gpu_correctness.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/.. +) + +set_tests_properties(dispatcher_test_bquant_gpu_correctness PROPERTIES + LABELS "dispatcher;python;bridge;gpu" + TIMEOUT 600 + SKIP_RETURN_CODE 77 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + +# ABQuant grouped GEMM: fp8 / bf8 scales on BOTH operands, compv3 + the +# gfx950-native eightwaves pipeline (gfx950). +add_test( + NAME dispatcher_test_abquant_gpu_correctness + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_abquant_gpu_correctness.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/.. +) + +set_tests_properties(dispatcher_test_abquant_gpu_correctness PROPERTIES + LABELS "dispatcher;python;bridge;gpu" + TIMEOUT 600 + SKIP_RETURN_CODE 77 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + # Stress Test Script add_test( NAME dispatcher_stress_test @@ -604,43 +799,24 @@ message(STATUS " Run verbose: ctest -R dispatcher_test_autocorrect_verbose") # dispatcher_test_rowcolquant_gpu_correctness -- distinct test, distinct file. add_test( NAME dispatcher_test_gemm_rowcolquant_gpu_correctness - COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_gemm_rowcolquant_gpu_correctness.py + # Run the script directly rather than under pytest: pytest exits 0 when every + # test skips, so a CPU-only or non-fp8 runner reported a green PASS for work + # that never reached the GPU. main() exits 77 on those boxes instead. + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_gemm_rowcolquant_gpu_correctness.py WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ) set_tests_properties(dispatcher_test_gemm_rowcolquant_gpu_correctness PROPERTIES LABELS "dispatcher;python;bridge;gpu" TIMEOUT 600 - ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" -) - - -# Bridge on-device GPU correctness test (SKIPs cleanly without a GPU/built lib). -add_test( - NAME dispatcher_test_aquant_gpu_correctness - COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_aquant_gpu_correctness.py - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} -) -set_tests_properties(dispatcher_test_aquant_gpu_correctness PROPERTIES - LABELS "dispatcher;python;bridge;gpu" - TIMEOUT 600 - ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" -) - - -# Bridge on-device GPU correctness test (SKIPs cleanly without a GPU/built lib). -add_test( - NAME dispatcher_test_abquant_gpu_correctness - COMMAND ${Python3_EXECUTABLE} -m pytest -q ${CMAKE_CURRENT_SOURCE_DIR}/test_abquant_gpu_correctness.py - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} -) -set_tests_properties(dispatcher_test_abquant_gpu_correctness PROPERTIES - LABELS "dispatcher;python;bridge;gpu" - TIMEOUT 600 + SKIP_RETURN_CODE 77 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) # ============================================================================= -# Block-scale quant bridge tests (per-op unit tests + self-tests + bquant GPU) +# Block-scale quant bridge tests (per-op unit tests + self-tests) +# +# The *_gpu_correctness counterparts are registered above, next to the other +# script-style tests that need SKIP_RETURN_CODE 77. # ============================================================================= # Per-op bridge name-contract / codegen / scope unit tests (CPU-only, unittest). @@ -680,16 +856,3 @@ set_tests_properties(dispatcher_test_grouped_gemm_bquant_utils PROPERTIES TIMEOUT 120 ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) - -# bquant on-device GPU correctness test (SKIPs cleanly without a GPU/built lib), -# registered consistently with the sibling *_gpu_correctness bridge tests. -add_test( - NAME dispatcher_test_bquant_gpu_correctness - COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_bquant_gpu_correctness.py - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} -) -set_tests_properties(dispatcher_test_bquant_gpu_correctness PROPERTIES - LABELS "dispatcher;python;bridge;gpu" - TIMEOUT 600 - ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" -) diff --git a/projects/composablekernel/dispatcher/tests/run_unittest_77.py b/projects/composablekernel/dispatcher/tests/run_unittest_77.py new file mode 100644 index 000000000000..41e0a7db7be7 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/run_unittest_77.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +""" +ctest driver for unittest.TestCase suites that skip themselves. + +``python -m unittest discover`` exits 0 when every test skipped, so ctest +records Passed for a run that executed nothing. On a CPU-only runner -- or a +gfx942 box for a gfx950-only suite -- that is a green tick backed by zero +assertions, which is exactly the signal SKIP_RETURN_CODE exists to remove. + +This runs the same discovery and maps the outcome onto the convention the +script-style tests here already use: + + 0 at least one test actually ran, and everything passed + 77 every collected test skipped -> ctest "Skipped" + 1 a test failed or errored, or discovery collected nothing at all + +Discovery collecting nothing is a failure rather than a skip on purpose: a +renamed file or a typo'd pattern otherwise reports Skipped forever and nobody +notices the suite left the build. + +Used instead of converting the suites to the script-style-77 form so the +skip-vs-run decision stays where it belongs -- in each test's own setUp. +""" + +import sys +import unittest +from pathlib import Path + +_SKIP = 77 + + +def main(argv) -> int: + if len(argv) != 2: + print(f"usage: {Path(argv[0]).name} ", file=sys.stderr) + return 2 + + pattern = argv[1] + start_dir = str(Path(__file__).resolve().parent) + + suite = unittest.defaultTestLoader.discover(start_dir, pattern=pattern) + result = unittest.TextTestRunner(verbosity=2).run(suite) + + if not result.wasSuccessful(): + return 1 + + if result.testsRun == 0: + print( + f"ERROR: pattern {pattern!r} collected no tests under {start_dir}. " + "The file was probably renamed or removed -- reporting failure " + "rather than Skipped so this does not go unnoticed.", + file=sys.stderr, + ) + return 1 + + ran = result.testsRun - len(result.skipped) + if ran == 0: + reasons = sorted({reason for _, reason in result.skipped}) + print( + f"SKIP: all {result.testsRun} test(s) in {pattern} skipped " + f"({'; '.join(reasons)})" + ) + return _SKIP + + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/projects/composablekernel/dispatcher/tests/test_abquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_abquant_gpu_correctness.py index a309e87ad79b..c85bddb07da2 100644 --- a/projects/composablekernel/dispatcher/tests/test_abquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_abquant_gpu_correctness.py @@ -4,252 +4,411 @@ # SPDX-License-Identifier: MIT """ -On-device GPU correctness test for the ABQuant (A- and B-scaled block-scale) -GEMM dispatcher bridge (PR #9981). +GPU correctness tests for ABQuant GEMM dispatcher — scales on *both* operands. -Mirrors test_bquant_gpu_correctness.py: build the op's block-scale dispatcher -.so, run it on the GPU via ABQuantGpuGemmRunner, and compare against a NumPy -fp32 reference: +Completes the quant trio alongside test_aquant_gpu_correctness.py (A-side only) +and test_bquant_gpu_correctness.py (B-side only). ABQuant scales A per (m, k) +group and B per (k, n) group in the same kernel: C = dequant(A, AQ) @ dequant(B, BQ) - -where - A_deq[m, k] = A[m, k] * AQ[m, k // aquant_group_k] (per-row, per-K-block) - B_deq[k, n] = B[k, n] * BQ[k // bquant_group_k, n // bquant_group_n] (per-K-block, per-N-block) - -The reference rounds A/B through the SAME fp8/bf8 quantization the kernel sees, -so the tolerance only absorbs GEMM accumulation error. - -Block-scale traps handled explicitly: - 1. fp8 flavour: FNUZ on gfx942, OCP on gfx950/gfx12 (wrong flavour NaNs / - decorrelates the reference). Selected from the detected arch. - 2. warp_tile_k=32 on gfx942 (128 silently all-zeros): default_fp8_config / - default_bf8_config derive it from the arch; the test asserts it matches. - 3. AQ is row-major [M, QK_A]; BQ is col-major [QK_B, QN_B] (passed Fortran- - contiguous); B is a C-contiguous [K, N] buffer consumed directly. - -TOOLCHAIN GATE (build-only skip, never a false pass): - ABQuant's codegen injects the `-mllvm -amdgpu-coerce-illegal-types=1` flag, - which older LLVM accepts but clang >= 22 (ROCm 7.2's roc-7.2 toolchain) - REJECTS ("Unknown command line argument"). On such a toolchain the .so cannot - be built, so this test SKIPS with that exact reason rather than reporting a - spurious failure. The numeric logic below was validated bit-accurate on gfx942 - (fp8/bf8 max_rel ~4e-4) after removing only that flag, so the skip is purely a - toolchain gate, not a correctness gap. The op's CMake/codegen registration for - this flag should likewise be made arch/flag-gated. + A[m, k] *= AQ[m // aM, k // aK] + B[k, n] *= BQ[k // bK, n // bN] + +Requires a gfx942 (MI300X) or gfx950 (MI350X / MI355X) GPU and hipcc in PATH. +Skips cleanly (exit 77) when no GPU is visible or the detected arch is neither, +so it is safe to invoke unconditionally from a CI lane. + +The fp8/bf8 host codec follows the target arch — OCP on gfx950, FNUZ on gfx942 — +via the same dispatcher_common.fp8_uses_ocp predicate that drives the JIT +defines, so the reference cannot end up encoding differently from the kernel. + +Tests: + fp8/compv3, bf8/compv3 — GPU output is non-zero, non-constant, finite, and + within 5% max-relative-error vs. a fp32 CPU reference; time_ms is + positive. M=N=256 spans a 2x2 grid of 128x128 output tiles, so a + constant tile index cannot pass. + fp8/eightwaves — the 8-wave pipeline (192x256 tiles, TransposeC=True, + bquant_group_n=128), which exercises a different MFMA path and + B-scale granularity than compv3. gfx950 only for now; see + EIGHTWAVES_SUPPORTED_ARCHS. Run: - python3 -m pytest test_abquant_gpu_correctness.py -v - python3 test_abquant_gpu_correctness.py # standalone + python3 test_abquant_gpu_correctness.py + python3 test_abquant_gpu_correctness.py -v # verbose hipcc output + python3 test_abquant_gpu_correctness.py --gfx gfx950 """ -import math -import subprocess +import argparse +import logging import sys import tempfile from pathlib import Path import numpy as np -import pytest - -_HERE = Path(__file__).resolve() -sys.path.insert(0, str(_HERE.parent)) # for conftest helpers under standalone run - -# Shared GPU probes, arch-aware fp8/bf8 codecs and build->run->verify harness. -from conftest import ( # noqa: E402 - encode_fp8 as _encode_arch, - qdq_fp8 as _qdq_arch, - gpu_available as _have_gpu, - ml_dtypes_available as _have_ml_dtypes, - run_and_verify, -) - - -def _find_python_dir() -> Path: - for c in (_HERE.parent.parent / "python", _HERE.parent / "python"): - if (c / "gemm_abquant_utils.py").is_file(): - return c - for parent in _HERE.parents: - cand = parent / "dispatcher" / "python" - if (cand / "gemm_abquant_utils.py").is_file(): - return cand - raise RuntimeError("could not locate dispatcher/python/gemm_abquant_utils.py") +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) -sys.path.insert(0, str(_find_python_dir())) +from dispatcher_common import fp8_uses_ocp -from gemm_abquant_utils import ( # noqa: E402 +from grouped_gemm_abquant_utils import ( ABQuantGemmProblem, ABQuantGpuGemmRunner, setup_multiple_abquant_dispatchers, - default_fp8_config, - default_bf8_config, - _warp_tile_k_for, - _detect_gpu_arch, + default_fp8_compv3_config, + default_bf8_compv3_config, + default_fp8_eightwaves_config, + _eightwaves_warp_tile_k, + _preshuffleb_warp_tile_k, ) - -# Thin local aliases so the case body reads the same as before (the codecs are -# the single canonical copies from conftest.py). -def _encode(arr, dtype: str, arch: str) -> np.ndarray: - return _encode_arch(arr, dtype, arch) - - -def _qdq(arr, dtype: str, arch: str) -> np.ndarray: - return _qdq_arch(arr, dtype, arch) - - -def _coerce_flag_supported() -> bool: - """True if the local clang accepts -amdgpu-coerce-illegal-types (LLVM opt). - - clang >= 22 (ROCm 7.2 roc-7.2) rejects it, which blocks the abquant .so - build; used to SKIP (not fail) on such toolchains. +log = logging.getLogger(__name__) + +TOLERANCE = 0.05 # 5% max relative error — fp8/bf8 precision floor + +PASS = "PASS" +FAIL = "FAIL" +# Per-case skip, for a pipeline not yet enabled on the detected arch. Distinct +# from SKIP_EXIT, which skips the whole file: here some cases did run, so the +# process still exits 0/1 on their result. +SKIP = "SKIP" + +# ctest reports this as "skipped" rather than passed. Without it, a CPU-only or +# non-gfx950 runner would report every ABQuant test as a hard FAIL for a reason +# that has nothing to do with the code under test. +SKIP_EXIT = 77 + +# ABQuant's compv3 pipeline uses standard fp8 MFMA (mfma_f32_16x16x32_fp8_fp8), +# which gfx942 has; gemm_abquant/CMakeLists.txt lists gfx942 in DESIRED_TARGETS +# too. The former gfx950-only restriction was a conservative default, not a +# hardware limit -- what it was really working around was the fp8 encoding +# mismatch, now fixed by threading the arch through the host codec. +SUPPORTED_ARCHS = ("gfx942", "gfx950") + +# eightwaves is held back one release. Its warp_tile_k is arch-aware +# (_eightwaves_warp_tile_k: 128 on gfx950, 32 elsewhere), so it *should* work, +# but 8-wave scheduling with 192x256 tiles has never been run on gfx942 and the +# failure mode if the tile is wrong is silent zeros rather than a build error. +# Widen this once one CI run has shown compv3 green on gfx942. +EIGHTWAVES_SUPPORTED_ARCHS = ("gfx950",) + + +# --------------------------------------------------------------------------- +# Dtype helpers +# --------------------------------------------------------------------------- + +def _fp8_ml_dtype(dtype: str, gfx_arch: str): + """The ml_dtypes fp8/bf8 type matching what the kernel for `gfx_arch` produces. + + gfx950 kernels are compiled OCP (e4m3fn / e5m2); gfx942 is native FNUZ + (e4m3fnuz / e5m2fnuz), which differs by one in the exponent bias. Encoding + with the wrong one shifts every value by a factor of two -- large enough to + be real error, small enough for a 5% relative gate to sometimes swallow. The + predicate is shared with the JIT flags (dispatcher_common.ocp_arch_defines) + so the reference and the kernel cannot disagree. """ - try: - clang = subprocess.run( - ["hipcc", "-print-prog-name=clang++"], capture_output=True, text=True, timeout=30 - ).stdout.strip() or "clang++" - probe = subprocess.run( - [clang, "-x", "c++", "-c", "-o", "/dev/null", - "-mllvm", "-amdgpu-coerce-illegal-types=1", "-"], - input="int main(){return 0;}", capture_output=True, text=True, timeout=60, - ) - return "Unknown command line argument" not in (probe.stderr or "") - except Exception: - # If we cannot probe, don't block; let the build itself report. - return True + import ml_dtypes + if fp8_uses_ocp(gfx_arch): + return ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return ml_dtypes.float8_e4m3fnuz if dtype == "fp8" else ml_dtypes.float8_e5m2fnuz -# ============================================================================= -# NumPy reference: C = dequant(A, AQ) @ dequant(B, BQ) -# ============================================================================= +def _encode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str) -> np.ndarray: + """Encode float32 → fp8/bf8 bytes (uint8 view of the bit pattern).""" + return arr.astype(_fp8_ml_dtype(dtype, gfx_arch)).view(np.uint8) -def reference_abquant_gemm(A_dec, B_dec, AQ, BQ, problem: ABQuantGemmProblem) -> np.ndarray: - M, N, K = problem.M, problem.N, problem.K - agK = problem.aquant_group_k - bgK = problem.bquant_group_k - bgN = problem.bquant_group_n - - A_deq = A_dec.astype(np.float32).copy() - for qi in range(problem.QK_A): - k0, k1 = qi * agK, min(qi * agK + agK, K) - A_deq[:, k0:k1] *= AQ.astype(np.float32)[:, qi][:, None] - B_deq = B_dec.astype(np.float32).copy() - for qi in range(problem.QK_B): - for qj in range(problem.QN_B): - k0, k1 = qi * bgK, min(qi * bgK + bgK, K) - n0, n1 = qj * bgN, min(qj * bgN + bgN, N) - B_deq[k0:k1, n0:n1] *= float(BQ[qi, qj]) +def _decode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str) -> np.ndarray: + """Decode fp8/bf8 bytes (uint8 view) back to float32.""" + return arr.view(_fp8_ml_dtype(dtype, gfx_arch)).astype(np.float32) - return (A_deq @ B_deq).astype(np.float32) +# --------------------------------------------------------------------------- +# CPU reference +# --------------------------------------------------------------------------- -# ============================================================================= -# build + run + verify -# ============================================================================= +def _reference_gemm(A_f32: np.ndarray, B_f32: np.ndarray, + AQ: np.ndarray, BQ: np.ndarray, + problem: ABQuantGemmProblem) -> np.ndarray: + """C = dequant(A, AQ) @ dequant(B, BQ) in fp32. + Mirrors reference_abquant_gemm in examples/gemm/python/15_grouped_gemm_abquant.py. + Unlike AQuant, both operands carry scales, and their group axes differ: + AQ is indexed (m, k) and BQ is indexed (k, n). + """ + M, N, K = problem.M, problem.N, problem.K + aM, aK = problem.aquant_group_m, problem.aquant_group_k + bK, bN = problem.bquant_group_k, problem.bquant_group_n + + A_dq = A_f32.astype(np.float32).copy() + for qi in range(problem.QM_A): + for qk in range(problem.QK_A): + m0, m1 = qi * aM, min((qi + 1) * aM, M) + k0, k1 = qk * aK, min((qk + 1) * aK, K) + A_dq[m0:m1, k0:k1] *= float(AQ[qi, qk]) + + B_dq = B_f32.astype(np.float32).copy() + for qk in range(problem.QK_B): + for qn in range(problem.QN_B): + k0, k1 = qk * bK, min((qk + 1) * bK, K) + n0, n1 = qn * bN, min((qn + 1) * bN, N) + B_dq[k0:k1, n0:n1] *= float(BQ[qk, qn]) + + return (A_dq @ B_dq).astype(np.float32) + + +def _max_rel_err(C_gpu: np.ndarray, C_ref: np.ndarray) -> float: + C_gpu_f = C_gpu.astype(np.float32) + C_ref_f = C_ref.astype(np.float32) + num = np.abs(C_gpu_f - C_ref_f) + # Use 1% of the global max magnitude as the denominator floor to avoid + # inflating the relative error when individual elements are near zero + # (a common occurrence with random inputs that partially cancel in GEMM). + ref_max = float(np.abs(C_ref_f).max()) + den = np.abs(C_ref_f) + max(ref_max * 1e-2, 1e-6) + return float(np.max(num / den)) + + +# --------------------------------------------------------------------------- +# Core test helper +# --------------------------------------------------------------------------- + +def _detect_arch() -> "str | None": + """Detected gfx arch, or None when no GPU is visible. + + The empty fallback keeps absence detectable; detect_gpu_arch would otherwise + invent a default and turn "no GPU" into a confusing build failure. + """ + try: + from dispatcher_common import detect_gpu_arch -def _run_case(dtype: str, M: int, N: int, K: int, agK: int, bgK: int, bgN: int, out_dir: Path): - arch = _detect_gpu_arch() - make_cfg = default_fp8_config if dtype == "fp8" else default_bf8_config - config = make_cfg(bquant_group_n=bgN, gfx_arch=arch) + return detect_gpu_arch(fallback="") or None + except Exception: + return None - expected_wtk = _warp_tile_k_for(dtype, arch) - assert config.warp_tile_k == expected_wtk, ( - f"warp_tile_k arch trap: got {config.warp_tile_k}, expected {expected_wtk} for {arch}" - ) - rng = np.random.default_rng(1234) - QK_A = math.ceil(K / agK) - QK_B = math.ceil(K / bgK) - QN_B = math.ceil(N / bgN) - A_f = rng.uniform(-1.0, 1.0, (M, K)).astype(np.float32) - B_f = rng.uniform(-1.0, 1.0, (K, N)).astype(np.float32) - AQ = rng.uniform(0.5, 1.5, (M, QK_A)).astype(np.float32) - BQ = rng.uniform(0.5, 1.5, (QK_B, QN_B)).astype(np.float32) +def _make_inputs(problem: ABQuantGemmProblem, dtype: str, gfx_arch: str, + seed: int = 42): + """Generate fp8-encoded A/B plus float32 AQ/BQ scales, and the decoded references. - A_raw = _encode(A_f, dtype, arch) - B_raw = _encode(B_f, dtype, arch) # C-contiguous [K, N], consumed directly - A_dec = _qdq(A_f, dtype, arch) - B_dec = _qdq(B_f, dtype, arch) + A and B are encoded to fp8 and decoded back so the CPU reference consumes the + same rounded values the kernel does; otherwise the comparison would be + dominated by quantisation error rather than kernel correctness. + AQ is row-major [QM_A, QK_A]; BQ is column-major [QK_B, QN_B] — the runner + applies the Fortran-order view, so the shape here stays logical. + """ + rng = np.random.default_rng(seed) + M, N, K = problem.M, problem.N, problem.K + A_f32 = rng.uniform(-2.0, 2.0, (M, K)).astype(np.float32) + B_f32 = rng.uniform(-2.0, 2.0, (K, N)).astype(np.float32) + AQ = rng.uniform(0.5, 2.0, (problem.QM_A, problem.QK_A)).astype(np.float32) + BQ = rng.uniform(0.5, 2.0, (problem.QK_B, problem.QN_B)).astype(np.float32) + A_raw = _encode_fp8(A_f32, dtype, gfx_arch) + B_raw = _encode_fp8(B_f32, dtype, gfx_arch) + return (A_raw, _decode_fp8(A_raw, dtype, gfx_arch), + B_raw, _decode_fp8(B_raw, dtype, gfx_arch), AQ, BQ) + + +def _expected_warp_tile_k(pipeline: str, gfx_arch: str) -> int: + """The warp_tile_k get_k_warp_tile() will pick for `pipeline` on `gfx_arch`. + + compv3 is a standard compute pipeline (not FlatMM), so 32 -- the standard fp8 + MFMA -- is correct on both arches and is deliberately arch-independent. + eightwaves and preshuffleb are FlatMM and do vary. + """ + if pipeline == "compv3": + return 32 + if pipeline == "preshuffleb": + return _preshuffleb_warp_tile_k(gfx_arch) + return _eightwaves_warp_tile_k(gfx_arch) + + +def _run_one(label: str, config, M: int, N: int, K: int, dtype: str, + out_dir: Path, gfx_arch: str, + seed: int = 42) -> "tuple[str, str]": + """Build, run, and verify one kernel. Returns (PASS|FAIL, detail_message).""" + # Arch trap, checked before the build: a warp_tile_k that disagrees with the + # branch get_k_warp_tile() takes for this arch compiles cleanly and returns + # zeros. The degenerate-C guards below do catch that, but only after a full + # hipcc build, and they report it as "kernel did not compute" rather than + # naming the config field that drifted. + expected_wtk = _expected_warp_tile_k(config.pipeline, gfx_arch) + if config.warp_tile_k != expected_wtk: + return FAIL, (f"{label}: warp_tile_k arch trap: got {config.warp_tile_k}, " + f"expected {expected_wtk} for {config.pipeline} on {gfx_arch}") + + # The problem must repeat the config's quant grouping: the kernel bakes the + # group sizes into the generated header, and the host strides are derived + # from them, so a mismatch silently reads the wrong scale. problem = ABQuantGemmProblem( - M=M, N=N, K=K, aquant_group_k=agK, bquant_group_k=bgK, bquant_group_n=bgN + M=M, N=N, K=K, + aquant_group_m=config.aquant_group_m, + aquant_group_n=config.aquant_group_n, + aquant_group_k=config.aquant_group_k, + bquant_group_m=config.bquant_group_m, + bquant_group_n=config.bquant_group_n, + bquant_group_k=config.bquant_group_k, ) - def _build(): - so_paths = setup_multiple_abquant_dispatchers( - configs=[config], output_dir=out_dir, gfx_arch=arch - ) - return so_paths[0] if so_paths else None - - def _on_build_fail(): - # Distinguish the known toolchain gate from a genuine failure. - if not _coerce_flag_supported(): - pytest.skip( - "abquant .so build blocked: clang>=22 rejects " - "-amdgpu-coerce-illegal-types=1 (ROCm 7.2 toolchain); test logic " - "verified bit-accurate after removing only that flag" - ) - pytest.fail("abquant kernel build failed for an unexpected reason") - - def _run(so_path): - runner = ABQuantGpuGemmRunner(so_path) - result = runner.run( - A=A_raw, B=B_raw, AQ=AQ, BQ=np.asfortranarray(BQ), problem=problem - ) - return result.C, result.time_ms - - res = run_and_verify( - build_so=_build, - run_kernel=_run, - reference=lambda: reference_abquant_gemm(A_dec, B_dec, AQ, BQ, problem), - tol=0.05, # fp8/bf8 block-scale ~1e-2 .. 5e-2 - label=f"abquant {dtype} {M}x{N}x{K}", - on_build_fail=_on_build_fail, + so_paths = setup_multiple_abquant_dispatchers( + configs=[config], + output_dir=out_dir, + gfx_arch=gfx_arch, ) - return res.max_rel, res.time_ms - + if not so_paths or so_paths[0] is None: + return FAIL, f"{label}: kernel build failed" -_SKIP_NO_GPU = pytest.mark.skipif(not _have_gpu(), reason="no ROCm GPU detected") -_SKIP_NO_MLD = pytest.mark.skipif(not _have_ml_dtypes(), reason="ml_dtypes not installed") + A_raw, A_dec, B_raw, B_dec, AQ, BQ = _make_inputs( + problem, dtype, gfx_arch, seed=seed + ) + runner = ABQuantGpuGemmRunner(so_paths[0]) + + result = runner.run(A=A_raw, B=B_raw, AQ=AQ, BQ=BQ, problem=problem, + c_dtype=np.float16) + C_gpu = result.C.astype(np.float32) + + # A degenerate C is the signature of a kernel that launched but never did the + # work -- e.g. a warp_tile_k tuned for gfx950 built for gfx942, where + # get_k_warp_tile() takes the other branch of CK_GFX950_SUPPORT and the + # result is silently all zeros rather than a compile error. Random A/B/AQ/BQ + # make any correct output wildly varying, so a zero standard deviation is + # conclusive; check it before the tolerance gate, whose relative-error + # denominator would otherwise report a confusing near-1.0 on a zero matrix. + if np.all(C_gpu == 0): + return FAIL, f"{label}: GPU output is all-zero" + if float(C_gpu.std()) <= 0.0: + return FAIL, (f"{label}: GPU output is constant " + f"(value={float(C_gpu.flat[0]):.6g}); kernel did not compute") + nan_mask = ~np.isfinite(C_gpu) + if np.any(nan_mask): + nan_frac = nan_mask.mean() + log.debug("%s: NaN/Inf fraction=%.3f, first 8 elements=%s", + label, nan_frac, C_gpu.flat[:8].tolist()) + return FAIL, f"{label}: GPU output contains NaN/Inf (frac={nan_frac:.3f})" + + C_ref = _reference_gemm(A_dec, B_dec, AQ, BQ, problem) + mre = _max_rel_err(C_gpu, C_ref) + if mre > TOLERANCE: + return FAIL, (f"{label}: max_rel_err={mre:.4f} > tol={TOLERANCE:.4f} " + f"(shape M={M} N={N} K={K})") + + # Timing sanity: a zero time_ms means the HIP event pair never bracketed a launch. + result_timed = runner.run(A=A_raw, B=B_raw, AQ=AQ, BQ=BQ, problem=problem, + c_dtype=np.float16) + if result_timed.time_ms <= 0.0: + return FAIL, f"{label}: time_ms={result_timed.time_ms:.4f} is not positive" + + return PASS, (f"{label}: max_rel_err={mre:.4f}, " + f"time_ms={result_timed.time_ms:.3f}") + + +# --------------------------------------------------------------------------- +# Individual test cases +# --------------------------------------------------------------------------- +# +# kPadM/kPadN/kPadK are false for every ABQuant prefill config, so M, N and K +# must be exact multiples of TileM, TileN and TileK. Shapes below are the +# smallest multi-tile sizes that satisfy that for each pipeline. + +def case_fp8_compv3(out_dir: Path, gfx_arch: str) -> "tuple[str, str]": + # compv3 tiles are 128x128x128; M=N=256 spans a 2x2 tile grid. + # bquant_group_n=1 gives one B scale per column, the finest granularity. + M, N, K, gK, bN = 256, 256, 512, 128, 1 + cfg = default_fp8_compv3_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=gfx_arch) + return _run_one("fp8/compv3", cfg, M, N, K, "fp8", out_dir, + gfx_arch=gfx_arch, seed=42) + + +def case_bf8_compv3(out_dir: Path, gfx_arch: str) -> "tuple[str, str]": + M, N, K, gK, bN = 256, 256, 512, 128, 1 + cfg = default_bf8_compv3_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=gfx_arch) + return _run_one("bf8/compv3", cfg, M, N, K, "bf8", out_dir, + gfx_arch=gfx_arch, seed=43) + + +def case_fp8_eightwaves(out_dir: Path, gfx_arch: str) -> "tuple[str, str]": + if gfx_arch not in EIGHTWAVES_SUPPORTED_ARCHS: + return SKIP, (f"fp8/eightwaves: not yet enabled on {gfx_arch} " + f"(see EIGHTWAVES_SUPPORTED_ARCHS)") + # eightwaves tiles are 192x256x128, so M must be a multiple of 192 and N of + # 256. bquant_group_n=128 is the only granularity the C++ tests validate for + # this pipeline. + M, N, K, gK, bN = 384, 512, 512, 128, 128 + cfg = default_fp8_eightwaves_config(quant_group_k=gK, bquant_group_n=bN, + gfx_arch=gfx_arch) + return _run_one("fp8/eightwaves", cfg, M, N, K, "fp8", out_dir, + gfx_arch=gfx_arch, seed=44) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +TESTS = [ + ("fp8/compv3", case_fp8_compv3), + ("bf8/compv3", case_bf8_compv3), + ("fp8/eightwaves", case_fp8_eightwaves), +] + + +def main(): + parser = argparse.ArgumentParser(description="ABQuant GPU correctness tests") + # No hardcoded default: it must stay possible to tell "user asked for + # gfx950" from "we are on an unrelated box", so the skip below can fire. + parser.add_argument("--gfx", default=None, + help="GPU arch override (default: auto-detect; gfx950 only)") + parser.add_argument("-v", "--verbose", action="store_true") + parser.add_argument("--output-dir", type=Path, default=None) + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s: %(message)s", + ) -@_SKIP_NO_GPU -@_SKIP_NO_MLD -@pytest.mark.parametrize("dtype", ["fp8", "bf8"]) -def test_abquant_gpu_matches_reference(dtype, tmp_path): - max_rel, _ = _run_case(dtype, M=256, N=256, K=512, agK=128, bgK=128, bgN=1, out_dir=tmp_path) - assert max_rel <= 0.05 + gfx = args.gfx or _detect_arch() + if not gfx: + print("SKIP: no supported GPU detected (rocminfo); ABQuant GPU tests skipped") + return SKIP_EXIT + if gfx not in SUPPORTED_ARCHS: + print(f"SKIP: ABQuant is {'/'.join(SUPPORTED_ARCHS)}-only; detected {gfx}") + return SKIP_EXIT + try: + import ml_dtypes # noqa: F401 + except ImportError: + # Without ml_dtypes there is no trustworthy fp8 codec here, and a + # stand-in would compare the kernel against the wrong values. + print("SKIP: ml_dtypes not installed; ABQuant fp8/bf8 encoding unavailable") + return SKIP_EXIT + + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="abquant_gpu_test_")) + log.info("Kernel output dir: %s", out_dir) + + results = [] + for name, fn in TESTS: + log.info("--- Running %s ---", name) + try: + status, detail = fn(out_dir, gfx) + except Exception as exc: + status, detail = FAIL, f"{name}: exception: {exc}" + results.append((name, status, detail)) + log.info("[%s] %s", status, detail) + print("\n=== Summary ===") + for name, status, detail in results: + print(f" [{status:4s}] {detail}") + passed = sum(1 for _, s, _ in results if s == PASS) + skipped = sum(1 for _, s, _ in results if s == SKIP) + failed = len(results) - passed - skipped + print(f"\n{passed}/{len(results) - skipped} passed" + + (f", {skipped} skipped" if skipped else "")) -@_SKIP_NO_GPU -@_SKIP_NO_MLD -def test_abquant_gpu_not_all_zeros(tmp_path): - _run_case("fp8", M=256, N=256, K=512, agK=128, bgK=128, bgN=1, out_dir=tmp_path) + return 0 if failed == 0 else 1 if __name__ == "__main__": - if not _have_gpu(): - print("SKIP: no GPU"); raise SystemExit(0) - if not _have_ml_dtypes(): - print("SKIP: ml_dtypes not installed"); raise SystemExit(0) - d = Path(tempfile.mkdtemp(prefix="abquant_gputest_")) - ok = True - for dt in ("fp8", "bf8"): - try: - mr, t = _run_case(dt, 256, 256, 512, 128, 128, 1, d) - print(f"PASS abquant {dt}: max_rel={mr:.4f} time_ms={t:.3f}") - except Exception as e: - # A pytest.skip raised outside pytest surfaces as an exception; treat - # the known toolchain gate as a non-failure so standalone runs on a - # clang>=22 box report BLOCKED, not FAIL. - msg = str(e) - if "amdgpu-coerce-illegal-types" in msg or "build blocked" in msg: - print(f"BLOCKED abquant {dt}: {msg}") - else: - ok = False - print(f"FAIL abquant {dt}: {msg}") - raise SystemExit(0 if ok else 1) + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py index 9cf953b9631e..05557955e6cf 100644 --- a/projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_aquant_gpu_correctness.py @@ -4,190 +4,328 @@ # SPDX-License-Identifier: MIT """ -On-device GPU correctness test for the AQuant (A-scaled block-scale) GEMM -dispatcher bridge (PR #9980). +GPU correctness tests for AQuant GEMM dispatcher — A-side per-group scales. -Mirrors the structure of test_bquant_gpu_correctness.py: build the op's -block-scale dispatcher .so, run it on the GPU via AQuantGpuGemmRunner, and -compare against a NumPy fp32 reference +Companion to test_bquant_gpu_correctness.py. Where BQuant scales the B operand +per (k, n) group, AQuant scales A per (m, k) group: - C = dequant(A, AQ) @ B + C = dequant(A, AQ) @ B, A[m, k] *= AQ[m // gM, k // gK] -where AQ applies a per-(row, K-group) scale to A: +Requires a gfx942 (MI300X) or gfx950 (MI350X / MI355X) GPU and hipcc in PATH. +Skips cleanly (exit 77) when no GPU is visible or the detected arch is neither. - A_deq[m, k] = A[m, k] * AQ[m, k // quant_group_k] +The non-preshuffleaq tests (fp8, bf8, fp8/tiled) use standard fp8 MFMA +(warp_tile_k=32) which gfx942 has. The preshuffleaq tests are gated behind +PRESHUFFLEAQ_SUPPORTED_ARCHS = ("gfx950",): they use FlatMM (warp_tile_k=64 +on gfx942), and that path has not yet been CI-verified on gfx942. -Two block-scale correctness traps are handled explicitly (both were needed to -get a bit-accurate match on gfx942): +The host fp8 codec follows the target arch — OCP on gfx950, FNUZ on gfx942. - 1. fp8 encoding flavour. gfx942 uses FNUZ fp8/bf8 (e4m3fnuz / e5m2fnuz); - gfx950/gfx12 use OCP (e4m3 / e5m2). Encoding A/B with the wrong flavour - makes the reference NaN or uncorrelated. Selected from the detected arch. - 2. warp_tile_k. fp8/bf8 needs warp_tile_k=32 on gfx942 (128 silently outputs - all-zeros). default_fp8_config/default_bf8_config derive this from the - arch; the test asserts it matches to guard against regressions. - -Also note the rcr layout uses ColMajor B, so the [K, N] logical B is stored -transposed (a C-contiguous [N, K] buffer == [K, N] col-major) before it is -handed to the kernel. +Tests: + fp8, bf8 — GPU output is non-zero, non-constant, finite, and within 5% + max-relative-error vs. a fp32 CPU reference; time_ms is positive. + fp8/tiled — same checks on a shape spanning a 4x4 grid of output tiles + (tile_m=16, tile_n=64), which the 1x1-tile cases cannot catch. Run: - python3 -m pytest test_aquant_gpu_correctness.py -v - python3 test_aquant_gpu_correctness.py # standalone + python3 test_aquant_gpu_correctness.py + python3 test_aquant_gpu_correctness.py -v # verbose hipcc output + python3 test_aquant_gpu_correctness.py --gfx gfx942 """ +import argparse +import logging import math import sys import tempfile from pathlib import Path import numpy as np -import pytest - -_HERE = Path(__file__).resolve() -sys.path.insert(0, str(_HERE.parent)) # for conftest helpers under standalone run - -# Shared GPU probes, arch-aware fp8/bf8 codecs and build->run->verify harness. -from conftest import ( # noqa: E402 - encode_fp8 as _encode_arch, - qdq_fp8 as _qdq_arch, - gpu_available as _have_gpu, - ml_dtypes_available as _have_ml_dtypes, - run_and_verify, -) - - -def _find_python_dir() -> Path: - for c in (_HERE.parent.parent / "python", _HERE.parent / "python"): - if (c / "gemm_aquant_utils.py").is_file(): - return c - for parent in _HERE.parents: - cand = parent / "dispatcher" / "python" - if (cand / "gemm_aquant_utils.py").is_file(): - return cand - raise RuntimeError("could not locate dispatcher/python/gemm_aquant_utils.py") +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) -sys.path.insert(0, str(_find_python_dir())) +from dispatcher_common import fp8_uses_ocp -from gemm_aquant_utils import ( # noqa: E402 +from grouped_gemm_aquant_utils import ( AQuantGemmProblem, AQuantGpuGemmRunner, setup_multiple_aquant_dispatchers, default_fp8_config, default_bf8_config, - _detect_gpu_arch, - _warp_tile_k_for, ) +log = logging.getLogger(__name__) -# Thin local aliases so the case body reads the same as before (the codecs are -# the single canonical copies from conftest.py). -def _encode(arr, dtype: str, arch: str) -> np.ndarray: - return _encode_arch(arr, dtype, arch) +TOLERANCE = 0.05 # 5% max relative error — fp8/bf8 precision floor +PASS = "PASS" +FAIL = "FAIL" +SKIP = "SKIP" -def _qdq(arr, dtype: str, arch: str) -> np.ndarray: - """Round through fp8/bf8 and back to fp32 (reference-side counterpart).""" - return _qdq_arch(arr, dtype, arch) +# ctest reports this as "skipped" rather than passed. Without it, a CPU-only or +# non-supported runner would report every AQuant test as a hard FAIL for a reason +# that has nothing to do with the code under test. +SKIP_EXIT = 77 +# fp8, bf8, fp8/tiled use standard MFMA (warp_tile_k=32), present on gfx942. +# The former gfx950-only restriction was the OCP-hardcoded host codec, now fixed. +SUPPORTED_ARCHS = ("gfx942", "gfx950") -# ============================================================================= -# NumPy reference: C = dequant(A, AQ) @ B -# ============================================================================= +# preshuffleaq uses FlatMM (warp_tile_k=64 on gfx942 per _preshuffleaq_warp_tile_k). +# Held back one CI run: FlatMM has not been empirically exercised on gfx942. +PRESHUFFLEAQ_SUPPORTED_ARCHS = ("gfx950",) -def reference_aquant_gemm(A_dec, B_dec, AQ, problem: AQuantGemmProblem) -> np.ndarray: - """A_dec [M,K], B_dec [K,N] are already fp8-rounded fp32 (kernel-identical).""" - M, N, K = problem.M, problem.N, problem.K - gK = problem.quant_group_k - A_deq = A_dec.astype(np.float32).copy() - for qi in range(problem.QK_A): - k0 = qi * gK - k1 = min(k0 + gK, K) - A_deq[:, k0:k1] *= AQ.astype(np.float32)[:, qi][:, None] - return (A_deq @ B_dec.astype(np.float32)).astype(np.float32) +# --------------------------------------------------------------------------- +# Dtype helpers +# --------------------------------------------------------------------------- +def _fp8_ml_dtype(dtype: str, gfx_arch: str): + """The ml_dtypes fp8/bf8 type matching what the kernel for `gfx_arch` produces.""" + import ml_dtypes -# ============================================================================= -# build + run + verify -# ============================================================================= + if fp8_uses_ocp(gfx_arch): + return ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return ml_dtypes.float8_e4m3fnuz if dtype == "fp8" else ml_dtypes.float8_e5m2fnuz -def _run_case(dtype: str, M: int, N: int, K: int, gK: int, out_dir: Path): - arch = _detect_gpu_arch() - make_cfg = default_fp8_config if dtype == "fp8" else default_bf8_config - config = make_cfg(quant_group_k=gK, quant_group_n=1, layout="rcr", gfx_arch=arch) +def _encode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str) -> np.ndarray: + """Encode float32 → fp8/bf8 bytes (uint8 view). Format follows target arch.""" + return arr.astype(_fp8_ml_dtype(dtype, gfx_arch)).view(np.uint8) - expected_wtk = _warp_tile_k_for(arch, preshuffle_aquant=False) - assert config.warp_tile_k == expected_wtk, ( - f"warp_tile_k arch trap: got {config.warp_tile_k}, expected {expected_wtk} for {arch}" - ) - rng = np.random.default_rng(1234) - QK_A = math.ceil(K / gK) - A_f = rng.uniform(-1.0, 1.0, (M, K)).astype(np.float32) - B_f = rng.uniform(-1.0, 1.0, (K, N)).astype(np.float32) # logical [K, N] - AQ_f = rng.uniform(0.5, 1.5, (M, QK_A)).astype(np.float32) - - A_raw = _encode(A_f, dtype, arch) # logical [M, K] fp8 bytes - B_raw = _encode(B_f, dtype, arch) # logical [K, N]; runner materializes rcr col-major - A_dec = _qdq(A_f, dtype, arch) - B_dec = _qdq(B_f, dtype, arch) - - problem = AQuantGemmProblem(M=M, N=N, K=K, quant_group_k=gK) - - def _build(): - so_paths = setup_multiple_aquant_dispatchers( - configs=[config], output_dir=out_dir, gfx_arch=arch - ) - return so_paths[0] if so_paths else None - - def _run(so_path): - runner = AQuantGpuGemmRunner(so_path, layout="rcr") - result = runner.run(A=A_raw, AQ=AQ_f, B=B_raw, problem=problem) - return result.C, result.time_ms - - res = run_and_verify( - build_so=_build, - run_kernel=_run, - reference=lambda: reference_aquant_gemm(A_dec, B_dec, AQ_f, problem), - tol=0.05, # fp8/bf8 block-scale ~1e-2 .. 5e-2 - label=f"aquant {dtype} {M}x{N}x{K} gK={gK}", - ) - return res.max_rel, res.time_ms +def _decode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str) -> np.ndarray: + """Decode fp8/bf8 bytes (uint8 view) → float32.""" + return arr.view(_fp8_ml_dtype(dtype, gfx_arch)).astype(np.float32) -_SKIP_NO_GPU = pytest.mark.skipif(not _have_gpu(), reason="no ROCm GPU detected") -_SKIP_NO_MLD = pytest.mark.skipif(not _have_ml_dtypes(), reason="ml_dtypes not installed") +# --------------------------------------------------------------------------- +# CPU reference +# --------------------------------------------------------------------------- +def _reference_gemm(A_f32: np.ndarray, B_f32: np.ndarray, + AQ: np.ndarray, problem: AQuantGemmProblem) -> np.ndarray: + """C = dequant(A, AQ) @ B in fp32. -@_SKIP_NO_GPU -@_SKIP_NO_MLD -@pytest.mark.parametrize("dtype", ["fp8", "bf8"]) -def test_aquant_gpu_matches_reference(dtype, tmp_path): - max_rel, _ = _run_case(dtype, M=256, N=256, K=512, gK=128, out_dir=tmp_path) - assert max_rel <= 0.05 + A is scaled per (m, k) group — the transpose of BQuant's (k, n) grouping. + Mirrors reference_aquant_gemm in examples/gemm/python/14_grouped_gemm_aquant.py. + """ + gM, gK = problem.quant_group_m, problem.quant_group_k + A_dq = A_f32.astype(np.float32).copy() + for qi in range(problem.QM_A): + for qk in range(problem.QK_A): + m0, m1 = qi * gM, min((qi + 1) * gM, problem.M) + k0, k1 = qk * gK, min((qk + 1) * gK, problem.K) + A_dq[m0:m1, k0:k1] *= float(AQ[qi, qk]) + return (A_dq @ B_f32.astype(np.float32)).astype(np.float32) -@_SKIP_NO_GPU -@_SKIP_NO_MLD -def test_aquant_gpu_not_all_zeros(tmp_path): - _run_case("fp8", M=256, N=256, K=512, gK=128, out_dir=tmp_path) +def _max_rel_err(C_gpu: np.ndarray, C_ref: np.ndarray) -> float: + C_gpu_f = C_gpu.astype(np.float32) + C_ref_f = C_ref.astype(np.float32) + num = np.abs(C_gpu_f - C_ref_f) + # Use 1% of the global max magnitude as the denominator floor to avoid + # inflating the relative error when individual elements are near zero + # (a common occurrence with random inputs that partially cancel in GEMM). + ref_max = float(np.abs(C_ref_f).max()) + den = np.abs(C_ref_f) + max(ref_max * 1e-2, 1e-6) + return float(np.max(num / den)) -if __name__ == "__main__": - if not _have_gpu(): - print("SKIP: no GPU"); raise SystemExit(0) - if not _have_ml_dtypes(): - print("SKIP: ml_dtypes not installed"); raise SystemExit(0) - d = Path(tempfile.mkdtemp(prefix="aquant_gputest_")) - ok = True - for dt in ("fp8", "bf8"): +# --------------------------------------------------------------------------- +# Core test helper +# --------------------------------------------------------------------------- + +def _detect_arch() -> "str | None": + """Detected gfx arch, or None when no GPU is visible. + + The empty fallback keeps absence detectable; detect_gpu_arch would otherwise + invent a default and turn "no GPU" into a confusing build failure. + """ + try: + from dispatcher_common import detect_gpu_arch + + return detect_gpu_arch(fallback="") or None + except Exception: + return None + + +def _make_inputs(M, N, K, gM, gK, dtype, gfx_arch: str, seed=42): + """Generate fp8-encoded A/B plus float32 AQ scales, and the decoded references. + + A and B are encoded to fp8 and decoded back so the CPU reference consumes the + same rounded values the kernel does; otherwise the comparison would be + dominated by quantisation error rather than kernel correctness. + """ + rng = np.random.default_rng(seed) + QM_A = math.ceil(M / gM) + QK_A = math.ceil(K / gK) + A_f32 = rng.uniform(-2.0, 2.0, (M, K)).astype(np.float32) + B_f32 = rng.uniform(-2.0, 2.0, (K, N)).astype(np.float32) + AQ = rng.uniform(0.5, 2.0, (QM_A, QK_A)).astype(np.float32) + A_raw = _encode_fp8(A_f32, dtype, gfx_arch) + B_raw = _encode_fp8(B_f32, dtype, gfx_arch) + return (A_raw, _decode_fp8(A_raw, dtype, gfx_arch), + B_raw, _decode_fp8(B_raw, dtype, gfx_arch), AQ) + + +def _run_one(label: str, config, M: int, N: int, K: int, + A_raw: np.ndarray, A_f32: np.ndarray, + B_raw: np.ndarray, B_f32: np.ndarray, + AQ: np.ndarray, + out_dir: Path, + gfx_arch: str) -> "tuple[str, str]": + """Build, run, and verify one kernel. Returns (PASS|FAIL, detail_message).""" + problem = AQuantGemmProblem( + M=M, N=N, K=K, + quant_group_m=config.quant_group_m, + quant_group_n=config.quant_group_n, + quant_group_k=config.quant_group_k, + ) + + so_paths = setup_multiple_aquant_dispatchers( + configs=[config], + output_dir=out_dir, + gfx_arch=gfx_arch, + ) + if not so_paths or so_paths[0] is None: + return FAIL, f"{label}: kernel build failed" + + runner = AQuantGpuGemmRunner(so_paths[0]) + + result = runner.run(A=A_raw, B=B_raw, AQ=AQ, problem=problem, c_dtype=np.float16) + C_gpu = result.C.astype(np.float32) + + # A degenerate C is the signature of a kernel that launched but never did the + # work — e.g. a wrong warp_tile_k (CK_GFX950_SUPPORT branch mismatch) returning + # all zeros on gfx942. Check before the tolerance gate; see abquant for rationale. + if np.all(C_gpu == 0): + return FAIL, f"{label}: GPU output is all-zero" + if float(C_gpu.std()) <= 0.0: + return FAIL, (f"{label}: GPU output is constant " + f"(value={float(C_gpu.flat[0]):.6g}); kernel did not compute") + nan_mask = ~np.isfinite(C_gpu) + if np.any(nan_mask): + nan_frac = nan_mask.mean() + log.debug("%s: NaN/Inf fraction=%.3f, first 8 elements=%s", + label, nan_frac, C_gpu.flat[:8].tolist()) + return FAIL, f"{label}: GPU output contains NaN/Inf (frac={nan_frac:.3f})" + + C_ref = _reference_gemm(A_f32, B_f32, AQ, problem) + mre = _max_rel_err(C_gpu, C_ref) + if mre > TOLERANCE: + return FAIL, (f"{label}: max_rel_err={mre:.4f} > tol={TOLERANCE:.4f} " + f"(shape M={M} N={N} K={K})") + + # Timing sanity: a zero time_ms means the HIP event pair never bracketed a launch. + result_timed = runner.run(A=A_raw, B=B_raw, AQ=AQ, problem=problem, + c_dtype=np.float16) + if result_timed.time_ms <= 0.0: + return FAIL, f"{label}: time_ms={result_timed.time_ms:.4f} is not positive" + + return PASS, (f"{label}: max_rel_err={mre:.4f}, " + f"time_ms={result_timed.time_ms:.3f}") + + +# --------------------------------------------------------------------------- +# Individual test cases +# --------------------------------------------------------------------------- + +def case_fp8(out_dir: Path, gfx_arch: str) -> "tuple[str, str]": + # K=768 = 3*TileK(256): num_loop=3 gives the no-hot-loop/TailNumber::Odd path, + # which exercises the AQ scale prefetch more thoroughly than an even tail. + M, N, K, gM, gK = 16, 64, 768, 1, 128 + cfg = default_fp8_config(quant_group_k=gK, quant_group_m=gM, gfx_arch=gfx_arch) + A_raw, A_dec, B_raw, B_dec, AQ = _make_inputs(M, N, K, gM, gK, "fp8", gfx_arch) + return _run_one("fp8", cfg, M, N, K, A_raw, A_dec, B_raw, B_dec, AQ, + out_dir, gfx_arch=gfx_arch) + + +def case_bf8(out_dir: Path, gfx_arch: str) -> "tuple[str, str]": + M, N, K, gM, gK = 16, 64, 768, 1, 128 + cfg = default_bf8_config(quant_group_k=gK, quant_group_m=gM, gfx_arch=gfx_arch) + A_raw, A_dec, B_raw, B_dec, AQ = _make_inputs(M, N, K, gM, gK, "bf8", gfx_arch, + seed=43) + return _run_one("bf8", cfg, M, N, K, A_raw, A_dec, B_raw, B_dec, AQ, + out_dir, gfx_arch=gfx_arch) + + +def case_fp8_tiled(out_dir: Path, gfx_arch: str) -> "tuple[str, str]": + # M=64, N=256 spans 4x4 output tiles (tile_m=16, tile_n=64). The single-tile + # cases above cannot distinguish a correct tile index from a constant one. + M, N, K, gM, gK = 64, 256, 768, 1, 128 + cfg = default_fp8_config(quant_group_k=gK, quant_group_m=gM, gfx_arch=gfx_arch) + A_raw, A_dec, B_raw, B_dec, AQ = _make_inputs(M, N, K, gM, gK, "fp8", gfx_arch, + seed=44) + return _run_one("fp8/tiled", cfg, M, N, K, A_raw, A_dec, B_raw, B_dec, AQ, + out_dir, gfx_arch=gfx_arch) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +TESTS = [ + ("fp8", case_fp8), + ("bf8", case_bf8), + ("fp8/tiled", case_fp8_tiled), +] + + +def main(): + parser = argparse.ArgumentParser(description="AQuant GPU correctness tests") + # No hardcoded default: it must stay possible to tell "user asked for + # gfx950" from "we are on an unrelated box", so the skip below can fire. + parser.add_argument("--gfx", default=None, + help="GPU arch override (default: auto-detect)") + parser.add_argument("-v", "--verbose", action="store_true") + parser.add_argument("--output-dir", type=Path, default=None) + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s: %(message)s", + ) + + gfx = args.gfx or _detect_arch() + if not gfx: + print("SKIP: no supported GPU detected (rocminfo); AQuant GPU tests skipped") + return SKIP_EXIT + if gfx not in SUPPORTED_ARCHS: + print(f"SKIP: AQuant is {'/'.join(SUPPORTED_ARCHS)}-only; detected {gfx}") + return SKIP_EXIT + try: + import ml_dtypes # noqa: F401 + except ImportError: + # Without ml_dtypes there is no trustworthy fp8 codec here, and a + # stand-in would compare the kernel against the wrong values. + print("SKIP: ml_dtypes not installed; AQuant fp8/bf8 encoding unavailable") + return SKIP_EXIT + + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="aquant_gpu_test_")) + log.info("Kernel output dir: %s", out_dir) + + results = [] + for name, fn in TESTS: + log.info("--- Running %s ---", name) try: - mr, t = _run_case(dt, 256, 256, 512, 128, d) - print(f"PASS aquant {dt}: max_rel={mr:.4f} time_ms={t:.3f}") - except Exception as e: - ok = False - print(f"FAIL aquant {dt}: {e}") - raise SystemExit(0 if ok else 1) + status, detail = fn(out_dir, gfx) + except Exception as exc: + status, detail = FAIL, f"{name}: exception: {exc}" + results.append((name, status, detail)) + log.info("[%s] %s", status, detail) + + print("\n=== Summary ===") + for name, status, detail in results: + print(f" [{status:4s}] {detail}") + passed = sum(1 for _, s, _ in results if s == PASS) + skipped = sum(1 for _, s, _ in results if s == SKIP) + failed = len(results) - passed - skipped + print(f"\n{passed}/{len(results) - skipped} passed" + + (f", {skipped} skipped" if skipped else "")) + + return 0 if failed == 0 else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_batched_contraction_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_batched_contraction_gpu_correctness.py index 8788f3b48f66..5a93d5c306e2 100644 --- a/projects/composablekernel/dispatcher/tests/test_batched_contraction_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_batched_contraction_gpu_correctness.py @@ -6,11 +6,13 @@ """ GPU correctness test for the batched-contraction dispatcher bridge. -Builds the batched-contraction dispatcher .so (fp16 / rcr / num_d_tensors=0 — -the PassThrough v1 signature Old-TE's batched_contraction argparse validates), -runs a small contraction on-device via GpuBatchedContractionRunner, and compares -the GPU output to an fp32 numpy einsum reference within an fp16-appropriate -tolerance. Skips cleanly (exit 0) when no GPU / hipcc is available. +Builds the batched-contraction dispatcher .so once per dtype (fp16 / bf16 / +fp32, all rcr / num_d_tensors=0 — the PassThrough v1 signature Old-TE's +batched_contraction argparse validates), runs a small contraction on-device via +GpuBatchedContractionRunner, and compares the GPU output to an fp32 numpy einsum +reference within a per-dtype tolerance (see TOLERANCE). Every dtype is attempted +even if an earlier one fails, so one broken dtype does not mask the others. +Skips cleanly (exit 77) when no GPU / hipcc is available. The kernel computes: E[g,m,n] = sum_k A[g,m,k] * B[g,n,k] @@ -37,7 +39,7 @@ from batched_contraction_utils import ( # noqa: E402 BatchedContractionProblem, GpuBatchedContractionRunner, - default_fp16_config, + default_config, setup_multiple_batched_contraction_dispatchers, _get_arch, _validate_arch, @@ -45,12 +47,26 @@ log = logging.getLogger(__name__) -# fp16 inputs, fp32 accumulate; small K keeps the worst-case error under 1e-2. -TOLERANCE = 1e-2 +# Per-dtype gates. All three accumulate in fp32, so the gate is set by the input +# precision: fp16 (10-bit mantissa) at 1e-2, bf16 (7-bit) proportionally looser, +# fp32 near exact -- it is the same arithmetic as the numpy reference, differing +# only in summation order. Small K keeps the worst case inside these. +TOLERANCE = { + "fp16": 1e-2, + "bf16": 3e-2, + "fp32": 1e-5, +} + +DTYPES = ("fp16", "bf16", "fp32") PASS = "PASS" FAIL = "FAIL" +# ctest reports this as "skipped" rather than passed; see SKIP_RETURN_CODE in +# dispatcher/tests/CMakeLists.txt. Returning 0 here would make a CPU-only runner +# report a green PASS for a test that never touched the GPU. +SKIP_EXIT = 77 + def _has_gpu() -> bool: try: @@ -77,18 +93,24 @@ def _max_rel_err(E_gpu: np.ndarray, E_ref: np.ndarray) -> float: return float(np.max(np.abs(g - r)) / ref_scale) -def test_contraction_fp16(gfx_arch: str) -> tuple[str, str]: - cfg = default_fp16_config(gfx_arch=gfx_arch) +def check_contraction(dtype: str, gfx_arch: str) -> tuple[str, str]: + tol = TOLERANCE[dtype] + cfg = default_config(dtype, gfx_arch=gfx_arch) + # A dtype whose warp tile is off the XDL allow-list is dropped by is_valid() + # rather than rejected, so the build would just produce nothing. Catch it + # here, where the cause is still obvious. + if not cfg.is_valid(): + return FAIL, f"contraction/{dtype}: default config is not valid for this dtype" so_paths = setup_multiple_batched_contraction_dispatchers([cfg], gfx_arch=gfx_arch) if not so_paths or so_paths[0] is None: - return FAIL, "contraction/fp16: kernel build failed" + return FAIL, f"contraction/{dtype}: kernel build failed" runner = GpuBatchedContractionRunner( - so_paths[0], dtype="fp16", num_d_tensors=0, elementwise="PassThrough" + so_paths[0], dtype=dtype, num_d_tensors=0, elementwise="PassThrough" ) - # num_dim_g/m/n/k = 1 each (default_fp16_config): one axis per group. + # num_dim_g/m/n/k = 1 each (default_config): one axis per group. # K=128 gives 2 tile-K iterations (tile_k=64). G, M, N, K = 2, 128, 128, 128 prob = BatchedContractionProblem( @@ -103,22 +125,22 @@ def test_contraction_fp16(gfx_arch: str) -> tuple[str, str]: E_gpu = result.E if E_gpu.shape != (G, M, N): - return FAIL, f"contraction/fp16: output shape {E_gpu.shape} != {(G, M, N)}" + return FAIL, f"contraction/{dtype}: output shape {E_gpu.shape} != {(G, M, N)}" if np.all(E_gpu == 0): - return FAIL, "contraction/fp16: GPU output is all-zero" + return FAIL, f"contraction/{dtype}: GPU output is all-zero" if not np.all(np.isfinite(E_gpu.astype(np.float32))): - return FAIL, "contraction/fp16: GPU output contains NaN/Inf" + return FAIL, f"contraction/{dtype}: GPU output contains NaN/Inf" # runner.reference computes E[g,m,n] = sum_k A[g,m,k]*B[g,n,k] in fp32. E_ref = runner.reference(A, B, prob) mre = _max_rel_err(E_gpu, E_ref) - if mre > TOLERANCE: - return FAIL, (f"contraction/fp16: max_rel_err={mre:.4e} > tol={TOLERANCE:.1e} " + if mre > tol: + return FAIL, (f"contraction/{dtype}: max_rel_err={mre:.4e} > tol={tol:.1e} " f"(G={G} M={M} N={N} K={K})") if result.time_ms <= 0.0: - return FAIL, f"contraction/fp16: time_ms={result.time_ms:.4f} not positive" + return FAIL, f"contraction/{dtype}: time_ms={result.time_ms:.4f} not positive" - return PASS, (f"contraction/fp16: max_rel_err={mre:.4e}, " + return PASS, (f"contraction/{dtype}: max_rel_err={mre:.4e}, " f"time_ms={result.time_ms:.3f}, G={G} MNK={M}/{N}/{K}") @@ -137,20 +159,24 @@ def main() -> int: if not _has_gpu(): print("SKIP: no supported GPU detected (rocminfo); contraction GPU test skipped") - return 0 + return SKIP_EXIT gfx = _validate_arch(args.gfx) if args.gfx else _get_arch() log.info("Running batched-contraction GPU correctness on %s", gfx) - try: - status, detail = test_contraction_fp16(gfx) - except Exception as exc: # noqa: BLE001 - status, detail = FAIL, f"contraction/fp16: exception: {exc}" + results = [] + for dtype in DTYPES: + try: + results.append(check_contraction(dtype, gfx)) + except Exception as exc: # noqa: BLE001 + results.append((FAIL, f"contraction/{dtype}: exception: {exc}")) print("\n=== Summary ===") - print(f" [{status:4s}] {detail}") - print(f"\n{1 if status == PASS else 0}/1 passed") - return 0 if status == PASS else 1 + for status, detail in results: + print(f" [{status:4s}] {detail}") + passed = sum(1 for status, _ in results if status == PASS) + print(f"\n{passed}/{len(results)} passed") + return 0 if passed == len(results) else 1 if __name__ == "__main__": diff --git a/projects/composablekernel/dispatcher/tests/test_batched_gemm_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_batched_gemm_gpu_correctness.py index 240263950813..42845f921c41 100644 --- a/projects/composablekernel/dispatcher/tests/test_batched_gemm_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_batched_gemm_gpu_correctness.py @@ -9,7 +9,7 @@ Builds the batched GEMM dispatcher .so (fp16 / rcr — the only signature Old-TE's batched_gemm_instance_builder validates), runs a small batched problem on-device via GpuBatchedGemmRunner, and compares the GPU output to a per-batch fp32 numpy -reference within an fp16-appropriate tolerance. Skips cleanly (exit 0) when no +reference within an fp16-appropriate tolerance. Skips cleanly (exit 77) when no GPU / hipcc is available so it is safe in a CPU-only CI lane. The kernel computes, for each batch b: @@ -52,6 +52,11 @@ FAIL = "FAIL" SKIP = "SKIP" +# ctest reports this as "skipped" rather than passed; see SKIP_RETURN_CODE in +# dispatcher/tests/CMakeLists.txt. Returning 0 here would make a CPU-only runner +# report a green PASS for a test that never touched the GPU. +SKIP_EXIT = 77 + def _has_gpu() -> bool: """True iff a supported GPU is visible to rocminfo (so the .so can run).""" @@ -159,7 +164,7 @@ def main() -> int: if not _has_gpu(): print("SKIP: no supported GPU detected (rocminfo); batched GPU test skipped") - return 0 + return SKIP_EXIT gfx = args.gfx or _resolve_arch(None) log.info("Running batched GEMM GPU correctness on %s", gfx) diff --git a/projects/composablekernel/dispatcher/tests/test_bquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_bquant_gpu_correctness.py index 5b1819ad9fec..8a514fbd5b80 100644 --- a/projects/composablekernel/dispatcher/tests/test_bquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_bquant_gpu_correctness.py @@ -4,25 +4,39 @@ # SPDX-License-Identifier: MIT """ -GPU correctness tests for BQuant GEMM dispatcher -- C4 and H3 coverage. +GPU correctness tests for BQuant GEMM dispatcher — C4 and H3 coverage. -Requires a gfx950 GPU (MI350X / MI355X) and hipcc in PATH. +Requires a gfx950 (MI350X / MI355X) GPU and hipcc in PATH. Skips cleanly +(exit 77) when no GPU is visible, the detected arch is not gfx950, or +ml_dtypes is missing. + +C4 (fp8/bf8 compv3) uses standard fp8 MFMA, which the gfx942 ISA also has, and +the host fp8 codec now follows the target arch via dispatcher_common.fp8_uses_ocp +rather than hardcoding OCP. gfx942 is nevertheless not in SUPPORTED_ARCHS: CI +gates the whole bquant lane to gfx950 (ck.groovy), so a gfx942 run here would +exercise a path no lane covers. Enabling it is a follow-up, not an oversight. +H3 (mx_bf16*) cannot be enabled on gfx942 at all: the scale-MFMA builtins it +needs (__builtin_amdgcn_mfma_scale_f32_*_f8f6f4) are absent from that ISA and +fail at compile time rather than silently. Tests: - C4 -- fp8, bf8: GPU output is non-zero and within 5% max-relative-error - vs. a fp32 CPU reference. - H3 -- mx_bf16bf16, mx_bf16bf8, mx_bf16fp4: same non-zero / rel-error checks, + C4 — fp8, bf8: GPU output is non-zero, non-constant, and within 5% + max-relative-error vs. a fp32 CPU reference. + C — fp8i4, bf8i4: pk_int4 B with an fp8/bf8-encoded BQ scale. + H3 — mx_bf16bf16, mx_bf16bf8, mx_bf16fp4: same non-zero / rel-error checks, plus verify QuantType::BQuantGrouped + e8m0 pipeline compiles and runs. - M2 -- timing: time_ms is non-zero when timing is requested. + gfx950 only (see MX_SUPPORTED_ARCHS). + M2 — timing: time_ms is non-zero when timing is requested. + +Every case sweeps N over _N_SWEEP to cover the per-N-tile de-permute. Run: python3 test_bquant_gpu_correctness.py python3 test_bquant_gpu_correctness.py -v # verbose hipcc output - python3 test_bquant_gpu_correctness.py --gfx gfx950 + python3 test_bquant_gpu_correctness.py --gfx gfx942 """ import argparse -import pytest import logging import math import sys @@ -31,29 +45,14 @@ import numpy as np -sys.path.insert(0, str(Path(__file__).resolve().parent)) # for conftest helpers sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) -# Canonical numeric codecs (single copies shared with the sibling quant GPU -# tests): e8m0, OCP-fp4 LUT/unpack, bf16<->f32, and the global-max-floored -# max-rel-error helper all live in conftest.py. -from conftest import ( - uses_ocp_fp8 as _uses_ocp_fp8, - ml_fp8_dtype as _ml_fp8_dtype, - encode_e8m0 as _encode_e8m0, - decode_e8m0 as _decode_e8m0, - decode_fp4 as _decode_fp4, - bf16_raw_to_f32 as _bf16_raw_to_f32, - max_rel_err as _conftest_max_rel_err, - FP4_E2M1_LUT as _FP4_E2M1_LUT, -) +from dispatcher_common import fp8_uses_ocp -# Non-grouped gemm_bquant bridge (38_block_scale_gemm), the subject of PR #9982. -from gemm_bquant_utils import ( +from grouped_gemm_bquant_utils import ( BQuantGemmProblem, BQuantGpuGemmRunner, setup_multiple_bquant_dispatchers, - _detect_gpu_arch, default_fp8_config, default_bf8_config, default_fp8i4_config, @@ -65,41 +64,96 @@ log = logging.getLogger(__name__) -TOLERANCE = 0.05 # 5% max relative error -- fp8/bf8 precision floor +TOLERANCE = 0.05 # 5% max relative error — fp8/bf8 precision floor # --------------------------------------------------------------------------- -# Dtype helpers. _uses_ocp_fp8 / _ml_fp8_dtype and the e8m0 / fp4 / bf16 -# codecs come from conftest.py (single canonical copies). The fp8/bf8 -# encode/decode wrappers below keep this test's ml_dtypes-missing fallback and -# its direct (no float32-cast) ml_dtypes path, so the standalone script still -# runs on a box without ml_dtypes. +# Dtype helpers # --------------------------------------------------------------------------- -def _encode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str = "gfx950") -> np.ndarray: - """Encode float32 -> fp8 bytes (uint8 view), ARCH-AWARE (FNUZ on gfx942).""" - try: - ml_t = _ml_fp8_dtype(dtype, gfx_arch) - return arr.astype(ml_t).view(np.uint8) - except ImportError: - return (np.clip(arr, -2.0, 2.0) * 64).astype(np.int8).view(np.uint8) +def _fp8_ml_dtype(dtype: str, gfx_arch: str): + """The ml_dtypes fp8/bf8 type matching what the kernel for `gfx_arch` produces.""" + import ml_dtypes + if fp8_uses_ocp(gfx_arch): + return ml_dtypes.float8_e4m3fn if dtype == "fp8" else ml_dtypes.float8_e5m2 + return ml_dtypes.float8_e4m3fnuz if dtype == "fp8" else ml_dtypes.float8_e5m2fnuz + + +def _encode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str) -> np.ndarray: + """Encode float32 → fp8/bf8 bytes (uint8 view). Format follows target arch.""" + return arr.astype(_fp8_ml_dtype(dtype, gfx_arch)).view(np.uint8) -def _decode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str = "gfx950") -> np.ndarray: - try: - ml_t = _ml_fp8_dtype(dtype, gfx_arch) - return arr.view(ml_t).astype(np.float32) - except ImportError: - return arr.view(np.int8).astype(np.float32) / 64.0 +def _decode_fp8(arr: np.ndarray, dtype: str, gfx_arch: str) -> np.ndarray: + """Decode fp8/bf8 bytes (uint8 view) → float32.""" + return arr.view(_fp8_ml_dtype(dtype, gfx_arch)).astype(np.float32) -def _encode_bf8(arr: np.ndarray, gfx_arch: str = "gfx950") -> np.ndarray: + +def _encode_bf8(arr: np.ndarray, gfx_arch: str) -> np.ndarray: return _encode_fp8(arr, "bf8", gfx_arch) -def _decode_bf8(arr: np.ndarray, gfx_arch: str = "gfx950") -> np.ndarray: + +def _decode_bf8(arr: np.ndarray, gfx_arch: str) -> np.ndarray: return _decode_fp8(arr, "bf8", gfx_arch) +def _encode_e8m0(arr: np.ndarray) -> np.ndarray: + """Encode float32 scale values → e8m0 uint8 (MX block scale format). + + e8m0 stores a power-of-two exponent: byte b represents 2^(b - 127). + Scales must be positive; zero maps to 0 (subnormal/zero in e8m0). + """ + arr = np.asarray(arr, dtype=np.float32) + # Clamp to representable range: 2^-127 … 2^127 + arr = np.clip(arr, 0.0, np.float32(2.0 ** 127)) + nonzero = arr > 0.0 + out = np.zeros(arr.shape, dtype=np.uint8) + # biased exponent = floor(log2(s)) + 127, clamped to [0, 254] + exp = np.floor(np.log2(arr[nonzero])).astype(np.int32) + 127 + out[nonzero] = np.clip(exp, 0, 254).astype(np.uint8) + return out + + +def _decode_e8m0(arr: np.ndarray) -> np.ndarray: + """Decode e8m0 uint8 → float32 scale values (2^(b - 127)).""" + arr = np.asarray(arr, dtype=np.uint8) + return np.exp2(arr.astype(np.float32) - 127.0) + + +# OCP FP4 E2M1 lookup table (from pk_fp4.hpp e2m1_to_fp32_table). +# Index i (0-15) gives the float32 value for the 4-bit code i. +_FP4_E2M1_LUT: np.ndarray = np.array([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +], dtype=np.float32) + + +def _decode_fp4(packed: np.ndarray, K: int, N: int) -> np.ndarray: + """Unpack K*N OCP FP4 E2M1 values from K*N/2 packed bytes. + + pk_fp4_t packing (from pk_fp4.hpp _pack/_unpack): + byte = (element1 << 4) | (element0 & 0xF) + Low nibble = element at flat index 2i; high nibble = element at flat index 2i+1. + The flat layout is row-major [K, N], so each byte contains two consecutive N-elements. + """ + flat = packed.flatten() + lo = (flat & 0x0F).astype(np.uint8) + hi = ((flat >> 4) & 0x0F).astype(np.uint8) + out = np.empty(K * N, dtype=np.float32) + out[0::2] = _FP4_E2M1_LUT[lo] + out[1::2] = _FP4_E2M1_LUT[hi] + return out.reshape(K, N) + + +def _bf16_raw_to_f32(arr: np.ndarray) -> np.ndarray: + """Reinterpret a uint16 array of bf16 bit patterns as float32.""" + u16 = arr.flatten().astype(np.uint16) + words = np.zeros(len(u16) * 2, dtype=np.uint16) + words[1::2] = u16 # bf16 occupies upper 2 bytes of float32 (little-endian) + return words.view(np.float32).reshape(arr.shape) + + # --------------------------------------------------------------------------- # CPU reference # --------------------------------------------------------------------------- @@ -119,11 +173,15 @@ def _reference_gemm(A_f32: np.ndarray, B_f32: np.ndarray, def _max_rel_err(C_gpu: np.ndarray, C_ref: np.ndarray) -> float: - # Canonical helper in conftest.py: max|gpu-ref| / (|ref| + max(1%*max|ref|, - # 1e-6)). The 1%-of-global-max floor avoids inflating the relative error - # when individual elements are near zero (random GEMM inputs that partially - # cancel). - return _conftest_max_rel_err(C_gpu, C_ref, floor_frac=1e-2) + C_gpu_f = C_gpu.astype(np.float32) + C_ref_f = C_ref.astype(np.float32) + num = np.abs(C_gpu_f - C_ref_f) + # Use 1% of the global max magnitude as the denominator floor to avoid + # inflating the relative error when individual elements are near zero + # (a common occurrence with random inputs that partially cancel in GEMM). + ref_max = float(np.abs(C_ref_f).max()) + den = np.abs(C_ref_f) + max(ref_max * 1e-2, 1e-6) + return float(np.max(num / den)) # --------------------------------------------------------------------------- @@ -134,6 +192,37 @@ def _max_rel_err(C_gpu: np.ndarray, C_ref: np.ndarray) -> float: FAIL = "FAIL" SKIP = "SKIP" +# ctest reports this as "skipped" rather than passed. Without it, a CPU-only or +# non-supported runner would report every BQuant test as a hard FAIL for a reason +# that has nothing to do with the code under test. +SKIP_EXIT = 77 + +# Kept in lockstep with the CI gate (`arch == "gfx950"` in ck.groovy): the two +# must not disagree, or the test claims coverage no lane produces. The blockers +# that made this gfx950-only -- the OCP-hardcoded host codec and the warp_tile_k +# mismatches in the utils -- are fixed, so widening this to include "gfx942" is +# a matter of flipping the lane and this tuple together, in that follow-up. +SUPPORTED_ARCHS = ("gfx950",) + +# H3 (mx_*) is gfx950-only at the ISA level: __builtin_amdgcn_mfma_scale_f32_* +# builtins do not exist on gfx942 and cause a compile error there, not a silent +# zero. These cases are not merely deferred -- they cannot be enabled on gfx942. +MX_SUPPORTED_ARCHS = ("gfx950",) + + +def _detect_arch() -> str | None: + """Detected gfx arch, or None when no GPU is visible. + + The empty fallback keeps absence detectable; detect_gpu_arch would otherwise + invent a default and turn "no GPU" into a confusing build failure. + """ + try: + from dispatcher_common import detect_gpu_arch + + return detect_gpu_arch(fallback="") or None + except Exception: + return None + def _run_one(label: str, config, M: int, N: int, K: int, A_raw: np.ndarray, A_f32: np.ndarray, @@ -181,9 +270,14 @@ def _run_one(label: str, config, M: int, N: int, K: int, # Convert raw output to float32 for validation (needed for bf16 output via uint16 buffer) C_gpu = c_decode_fn(C_gpu_raw) if c_decode_fn is not None else C_gpu_raw.astype(np.float32) - # Non-zero check (C4 / H3 smoke) + # Non-trivial check (C4 / H3 smoke). A wrong warp_tile_k (CK_GFX950_SUPPORT + # branch mismatch) produces all-zero or all-constant output on gfx942; catch + # it here before the tolerance gate, whose denominator near zero is unhelpful. if np.all(C_gpu == 0): return FAIL, f"{label}: GPU output is all-zero" + if float(C_gpu.std()) <= 0.0: + return FAIL, (f"{label}: GPU output is constant " + f"(value={float(C_gpu.flat[0]):.6g}); kernel did not compute") nan_mask = ~np.isfinite(C_gpu.astype(np.float32)) if np.any(nan_mask): nan_frac = nan_mask.mean() @@ -212,7 +306,7 @@ def _run_one(label: str, config, M: int, N: int, K: int, # Individual test cases # --------------------------------------------------------------------------- -def _make_fp8_inputs(M, N, K, gK, gN, dtype="fp8", seed=42, gfx_arch="gfx950"): +def _make_fp8_inputs(M, N, K, gK, gN, dtype="fp8", gfx_arch="gfx950", seed=42): rng = np.random.default_rng(seed) QK_B = math.ceil(K / gK) QN_B = math.ceil(N / gN) @@ -227,7 +321,7 @@ def _make_fp8_inputs(M, N, K, gK, gN, dtype="fp8", seed=42, gfx_arch="gfx950"): def _to_bf16_raw(x: np.ndarray) -> np.ndarray: - """Encode float32 array -> uint16 array of bfloat16 bit patterns.""" + """Encode float32 array → uint16 array of bfloat16 bit patterns.""" packed = np.frombuffer(x.astype(np.float32).tobytes(), dtype=np.uint16) # Little-endian: bf16 occupies the upper 2 bytes of each float32, # which are at odd indices (1, 3, 5, ...) in the uint16 view. @@ -257,10 +351,13 @@ def _make_bf16_inputs(M, N, K, gK, gN, seed=42): # N-tile sweep: N=128 is a single N-tile; N=256/512 span 2 and 4 TileN blocks and # exercise the round-6 per-N-tile PermuteN de-permute (the round-5 global riffle # scrambled columns at N>=256). gN must divide N; TileN is 64 (decode) / 128 (MX). +# +# A single pinned N under TileN would leave the de-permute entirely unexercised, +# so this is coverage, not a knob: shrink it only with a replacement for that path. _N_SWEEP = (128, 256, 512) -def _case_c4_fp8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: +def case_c4_fp8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: # K=768 = 3*TileK(256): use num_loop=3 (TailNumber::Odd) for better coverage. # num_loop=2 works but exercises only the no-hot-loop/Even tail path; 3 gives # the no-hot-loop/Odd tail path and exercises the BQ scale prefetch more robustly. @@ -277,8 +374,8 @@ def _case_c4_fp8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: return PASS, f"C4/fp8: PASS for N in {_N_SWEEP}" -def _case_c4_bf8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: - # K=768 = 3*TileK(256): use num_loop=3 for the same reason as test_c4_fp8. +def case_c4_bf8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: + # K=768 = 3*TileK(256): use num_loop=3 for the same reason as case_c4_fp8. M, K, gK, gN = 16, 768, 128, 1 cfg = default_bf8_config(quant_group_k=gK, quant_group_n=gN, gfx_arch=gfx_arch) for N in _N_SWEEP: @@ -292,7 +389,10 @@ def _case_c4_bf8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: return PASS, f"C4/bf8: PASS for N in {_N_SWEEP}" -def _case_h3_mx_bf16bf16(out_dir: Path, gfx_arch: str) -> tuple[str, str]: +def case_h3_mx_bf16bf16(out_dir: Path, gfx_arch: str) -> tuple[str, str]: + if gfx_arch not in MX_SUPPORTED_ARCHS: + return SKIP, (f"H3/mx_bf16bf16: not supported on {gfx_arch} " + f"(requires {'/'.join(MX_SUPPORTED_ARCHS)})") # K=256 = 2*TileK(128): MicroscaleCompV3 needs num_loop>=2 to avoid OOB second prefetch M, K, gK, gN = 128, 256, 32, 1 cfg = default_mx_bf16bf16_config(quant_group_k=gK, quant_group_n=gN, gfx_arch=gfx_arch) @@ -308,7 +408,10 @@ def _case_h3_mx_bf16bf16(out_dir: Path, gfx_arch: str) -> tuple[str, str]: return PASS, f"H3/mx_bf16bf16: PASS for N in {_N_SWEEP}" -def _case_h3_mx_bf16bf8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: +def case_h3_mx_bf16bf8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: + if gfx_arch not in MX_SUPPORTED_ARCHS: + return SKIP, (f"H3/mx_bf16bf8: not supported on {gfx_arch} " + f"(requires {'/'.join(MX_SUPPORTED_ARCHS)})") # K=384 = 3*TileK(128): use num_loop=3 (TailNumber::Odd) for broader pipeline coverage. M, K, gK, gN = 128, 384, 128, 1 cfg = default_mx_bf16bf8_config(quant_group_k=gK, quant_group_n=gN, gfx_arch=gfx_arch) @@ -327,8 +430,11 @@ def _case_h3_mx_bf16bf8(out_dir: Path, gfx_arch: str) -> tuple[str, str]: return PASS, f"H3/mx_bf16bf8: PASS for N in {_N_SWEEP}" -def _case_h3_mx_bf16fp4(out_dir: Path, gfx_arch: str) -> tuple[str, str]: - # pk_fp4: 2 fp4 values per byte -> B buffer is K*N/2 bytes. +def case_h3_mx_bf16fp4(out_dir: Path, gfx_arch: str) -> tuple[str, str]: + if gfx_arch not in MX_SUPPORTED_ARCHS: + return SKIP, (f"H3/mx_bf16fp4: not supported on {gfx_arch} " + f"(requires {'/'.join(MX_SUPPORTED_ARCHS)})") + # pk_fp4: 2 fp4 values per byte → B buffer is K*N/2 bytes. # K=256 = 2*TileK(128): MicroscaleCompV3 needs num_loop>=2 to avoid OOB second prefetch. M, K, gK, gN = 128, 256, 32, 1 cfg = default_mx_bf16fp4_config(quant_group_k=gK, quant_group_n=gN, gfx_arch=gfx_arch) @@ -353,7 +459,7 @@ def _case_h3_mx_bf16fp4(out_dir: Path, gfx_arch: str) -> tuple[str, str]: return PASS, f"H3/mx_bf16fp4: PASS for N in {_N_SWEEP}" -def _case_c_i4(out_dir: Path, gfx_arch: str) -> tuple[str, str]: +def case_c_i4(out_dir: Path, gfx_arch: str) -> tuple[str, str]: """Round-6: fp8i4 / bf8i4 must be exact once BQ is encoded to the kernel's QDataType (fp8/bf8, 1 byte). The round-5 float32 BQ produced NaN. Swept over N to also exercise the per-N-tile de-permute. B is pk_int4 (2 per byte).""" @@ -409,67 +515,46 @@ def _case_c_i4(out_dir: Path, gfx_arch: str) -> tuple[str, str]: # --------------------------------------------------------------------------- TESTS = [ - ("C4/fp8", _case_c4_fp8), - ("C4/bf8", _case_c4_bf8), - ("C/i4", _case_c_i4), - ("H3/mx_bf16bf16", _case_h3_mx_bf16bf16), - ("H3/mx_bf16bf8", _case_h3_mx_bf16bf8), - ("H3/mx_bf16fp4", _case_h3_mx_bf16fp4), + ("C4/fp8", case_c4_fp8), + ("C4/bf8", case_c4_bf8), + ("C/i4", case_c_i4), + ("H3/mx_bf16bf16", case_h3_mx_bf16bf16), + ("H3/mx_bf16bf8", case_h3_mx_bf16bf8), + ("H3/mx_bf16fp4", case_h3_mx_bf16fp4), ] -@pytest.mark.parametrize("case_name,case_fn", TESTS, ids=[t[0] for t in TESTS]) -def test_bquant_gpu_c4_h3(case_name, case_fn, gpu_arch, tmp_path): - """C4/H3 bquant GPU correctness -- one pytest case per TESTS entry. - - Runs on a real GPU (the gpu_arch fixture skips cleanly on CPU-only boxes). - MX (H3) variants need gfx950 e8m0 hardware; on any other arch they skip. - """ - try: - status, detail = case_fn(tmp_path, gpu_arch) - except Exception as exc: # noqa: BLE001 - if "requires gfx950" in str(exc): - pytest.skip(f"{case_name}: MX is gfx950-only (arch={gpu_arch})") - raise - if status == SKIP: - pytest.skip(detail) - assert status == PASS, detail - - -def _gpu_and_hipcc_available() -> bool: - """True only if both hipcc and a ROCm GPU are present. - - Lets the standalone test SKIP cleanly (exit 0) on CPU-only CI runners - instead of failing when the kernel build/run cannot proceed. Delegates to - the canonical probe shared with the sibling quant GPU tests. - """ - from conftest import gpu_available - - return gpu_available() - - def main(): parser = argparse.ArgumentParser(description="BQuant GPU correctness tests (C4 + H3)") - parser.add_argument( - "--gfx", - default=None, - help="GPU arch (default: auto-detect the running device via rocm_agent_enumerator)", - ) + # No hardcoded default: it must stay possible to tell "user asked for + # gfx950" from "we are on an unrelated box", so the skip below can fire. + parser.add_argument("--gfx", default=None, + help="GPU arch override (default: auto-detect)") parser.add_argument("-v", "--verbose", action="store_true") parser.add_argument("--output-dir", type=Path, default=None) args = parser.parse_args() - if not _gpu_and_hipcc_available(): - print("SKIP: no GPU/hipcc detected; skipping bquant GPU correctness") - return 0 - logging.basicConfig( level=logging.DEBUG if args.verbose else logging.INFO, format="%(levelname)s: %(message)s", ) - gfx_arch = args.gfx or _detect_gpu_arch() - log.info("Target GPU arch: %s", gfx_arch) + gfx = args.gfx or _detect_arch() + if not gfx: + print("SKIP: no supported GPU detected (rocminfo); BQuant GPU tests skipped") + return SKIP_EXIT + if gfx not in SUPPORTED_ARCHS: + print(f"SKIP: BQuant is {'/'.join(SUPPORTED_ARCHS)}-only; detected {gfx}") + return SKIP_EXIT + try: + import ml_dtypes # noqa: F401 + except ImportError: + # Gate here rather than letting _encode_fp8 raise: an ImportError inside a + # test function is swallowed by the loop's `except Exception` below and + # reported as FAIL, blaming the kernel for a missing host dependency. + print("SKIP: ml_dtypes not installed; BQuant fp8/bf8 encoding unavailable") + return SKIP_EXIT + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="bquant_gpu_test_")) log.info("Kernel output dir: %s", out_dir) @@ -477,25 +562,20 @@ def main(): for name, fn in TESTS: log.info("--- Running %s ---", name) try: - status, detail = fn(out_dir, gfx_arch) + status, detail = fn(out_dir, gfx) except Exception as exc: - # MX (H3) variants need gfx950 e8m0 microscaling hardware; on any - # other arch (e.g. gfx1250/MI400) that is an expected skip, not a - # failure -- the fp8/bf8/i4 block-scale paths still run there. - if "requires gfx950" in str(exc): - status, detail = SKIP, f"{name}: skipped on {gfx_arch} (MX is gfx950-only)" - else: - status, detail = FAIL, f"{name}: exception: {exc}" + status, detail = FAIL, f"{name}: exception: {exc}" results.append((name, status, detail)) log.info("[%s] %s", status, detail) print("\n=== Summary ===") - passed = sum(1 for _, s, _ in results if s == PASS) - skipped = sum(1 for _, s, _ in results if s == SKIP) - failed = sum(1 for _, s, _ in results if s == FAIL) for name, status, detail in results: print(f" [{status:4s}] {detail}") - print(f"\n{passed}/{len(results)} passed, {skipped} skipped, {failed} failed") + passed = sum(1 for _, s, _ in results if s == PASS) + skipped = sum(1 for _, s, _ in results if s == SKIP) + failed = len(results) - passed - skipped + print(f"\n{passed}/{len(results) - skipped} passed" + + (f", {skipped} skipped" if skipped else "")) return 0 if failed == 0 else 1 diff --git a/projects/composablekernel/dispatcher/tests/test_contraction_multi_abd_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_contraction_multi_abd_gpu_correctness.py index 5053c4bfff59..3a4f99c52a34 100644 --- a/projects/composablekernel/dispatcher/tests/test_contraction_multi_abd_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_contraction_multi_abd_gpu_correctness.py @@ -61,6 +61,10 @@ PASS = "PASS" FAIL = "FAIL" +# ctest SKIP_RETURN_CODE: main() returns this when the box cannot run the test +# at all, so the lane reports Skipped rather than a vacuous Passed. +SKIP_EXIT = 77 + # Must match the kernel's default tile (256x256x64) exactly: the default config # pads nothing, so M and N have to be whole multiples of the tile or # IsSupportedArguments rejects the launch. @@ -189,7 +193,7 @@ def main() -> int: if not _has_gpu(): print("SKIP: no supported GPU or hipcc detected; " "contraction_multi_abd GPU test skipped") - return 0 + return SKIP_EXIT gfx = _validate_arch(args.gfx) if args.gfx else _detect_gpu_arch() log.info("Running contraction_multi_abd GPU correctness on %s", gfx) diff --git a/projects/composablekernel/dispatcher/tests/test_dispatcher_common.py b/projects/composablekernel/dispatcher/tests/test_dispatcher_common.py index 2c0fc8307cdb..c7407f00b8e1 100644 --- a/projects/composablekernel/dispatcher/tests/test_dispatcher_common.py +++ b/projects/composablekernel/dispatcher/tests/test_dispatcher_common.py @@ -39,6 +39,10 @@ print_error, print_info, cleanup_generated_kernels, + normalize_arch, + fp8_uses_ocp, + ocp_arch_defines, + arch_feature_defines, ) @@ -239,5 +243,143 @@ def test_cleanup_nonexistent_dir_no_error(self): cleanup_generated_kernels(Path("/tmp/nonexistent_ck_test_dir_12345")) +class TestNormalizeArch(unittest.TestCase): + """Tests for normalize_arch -- the front door for every arch predicate.""" + + def test_bare_names_pass_through(self): + for arch in ("gfx90a", "gfx942", "gfx950", "gfx1200", "gfx1250"): + with self.subTest(arch=arch): + self.assertEqual(normalize_arch(arch), arch) + + def test_feature_suffix_is_stripped(self): + # What rocminfo and the HSA runtime actually hand back. + self.assertEqual(normalize_arch("gfx950:sramecc+:xnack-"), "gfx950") + self.assertEqual(normalize_arch("gfx942:sramecc+:xnack-"), "gfx942") + self.assertEqual(normalize_arch("gfx90a:xnack+"), "gfx90a") + + def test_offload_target_prefix_is_stripped(self): + # What --offload-arch round-trips through in some toolchain paths. This + # is the form the old startswith() test missed outright. + self.assertEqual(normalize_arch("amdgcn-amd-amdhsa--gfx950"), "gfx950") + self.assertEqual( + normalize_arch("amdgcn-amd-amdhsa--gfx950:sramecc+:xnack-"), "gfx950" + ) + + def test_case_and_whitespace_are_normalized(self): + self.assertEqual(normalize_arch(" GFX950 "), "gfx950") + + def test_none_and_empty_are_empty(self): + self.assertEqual(normalize_arch(None), "") + self.assertEqual(normalize_arch(""), "") + + +class TestFp8ArchPredicates(unittest.TestCase): + """fp8_uses_ocp and the define lists derived from it. + + A miss here is silent: the -DCK_USE_OCP_FP8 pair simply does not get passed, + the host pass of config.hpp falls back to FNUZ, and the kernel builds and + runs while disagreeing with its own reference on what the bytes mean. + """ + + _OCP = ("gfx950", "gfx1200", "gfx1201", "gfx1250") + _FNUZ = ("gfx90a", "gfx908", "gfx942", "gfx1100", "gfx1030") + + def test_ocp_arches(self): + for arch in self._OCP: + with self.subTest(arch=arch): + self.assertTrue(fp8_uses_ocp(arch)) + + def test_fnuz_arches(self): + for arch in self._FNUZ: + with self.subTest(arch=arch): + self.assertFalse(fp8_uses_ocp(arch)) + + def test_full_target_triples_agree_with_bare_names(self): + for arch in self._OCP + self._FNUZ: + triple = f"{arch}:sramecc+:xnack-" + with self.subTest(arch=triple): + self.assertEqual(fp8_uses_ocp(triple), fp8_uses_ocp(arch)) + + def test_offload_target_form_agrees_with_bare_name(self): + for arch in self._OCP + self._FNUZ: + target = f"amdgcn-amd-amdhsa--{arch}" + with self.subTest(arch=target): + self.assertEqual(fp8_uses_ocp(target), fp8_uses_ocp(arch)) + + def test_unknown_and_none_default_to_fnuz(self): + # FNUZ is the compiler default for both passes, so it is also the safe + # answer for an arch we do not recognize. + for arch in (None, "", "not-a-gpu"): + with self.subTest(arch=arch): + self.assertFalse(fp8_uses_ocp(arch)) + + def test_ocp_defines_are_emitted_as_a_pair(self): + # Passing one without the other splits ck/ from ck_tile/. + self.assertEqual( + ocp_arch_defines("gfx950"), + ["-DCK_USE_OCP_FP8", "-DCK_TILE_USE_OCP_FP8"], + ) + self.assertEqual(ocp_arch_defines("gfx942"), []) + + +class TestArchFeatureDefines(unittest.TestCase): + """arch_feature_defines mirrors the top-level CMakeLists.txt arch block.""" + + def test_gfx942_gets_no_feature_defines(self): + self.assertEqual(arch_feature_defines("gfx942"), []) + + def test_gfx950_gets_ocp_plus_mx(self): + self.assertEqual( + arch_feature_defines("gfx950"), + [ + "-DCK_USE_OCP_FP8", + "-DCK_TILE_USE_OCP_FP8", + "-DCK_USE_NATIVE_MX_SUPPORT", + "-DCK_GFX950_SUPPORT", + ], + ) + + def test_gfx12_enables_wmma(self): + # CK_TILE_USE_WMMA must be passed explicitly: it selects the outermost + # branch of get_k_warp_tile(), and leaving it undefined reads as 0 -- + # the MFMA branch, which no gfx12 part has. + defines = arch_feature_defines("gfx1200") + self.assertIn("-DCK_TILE_USE_WMMA=1", defines) + self.assertIn("-DCK_GFX12_SUPPORT", defines) + self.assertIn("-DCK_USE_OCP_FP8", defines) + self.assertNotIn("-DCK_GFX950_SUPPORT", defines) + + def test_gfx1250_gets_its_own_defines_on_top_of_gfx12(self): + defines = arch_feature_defines("gfx1250") + for expected in ( + "-DCK_TILE_USE_WMMA=1", + "-DCK_GFX12_SUPPORT", + "-DCK_USE_GFX1250", + "-DCK_USE_NATIVE_MX_SUPPORT", + "-DCK_GFX1250_SUPPORT", + ): + with self.subTest(define=expected): + self.assertIn(expected, defines) + + def test_gfx950_and_gfx12_defines_are_mutually_exclusive(self): + # get_k_warp_tile() branches on these; both at once is not a state the + # header is written for. + for arch in ("gfx950", "gfx1250"): + with self.subTest(arch=arch): + defines = arch_feature_defines(arch) + self.assertNotEqual( + "-DCK_GFX950_SUPPORT" in defines, + "-DCK_TILE_USE_WMMA=1" in defines, + ) + + def test_target_triple_gets_the_same_defines_as_the_bare_name(self): + for arch in ("gfx942", "gfx950", "gfx1250"): + with self.subTest(arch=arch): + self.assertEqual( + arch_feature_defines(f"{arch}:sramecc+:xnack-"), + arch_feature_defines(arch), + ) + + if __name__ == "__main__": unittest.main() diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_parity.py b/projects/composablekernel/dispatcher/tests/test_gemm_parity.py index b9d1cd1cd9ab..8ed4325acbcb 100644 --- a/projects/composablekernel/dispatcher/tests/test_gemm_parity.py +++ b/projects/composablekernel/dispatcher/tests/test_gemm_parity.py @@ -89,9 +89,15 @@ # Global-relative-error gates. fp16 measured ~3-4e-4 and bf16 ~8e-3 on gfx942. # fp8/bf8 are far coarser (3- and 2-bit mantissa) so their gates are looser; int8 -# is an exact integer accumulation so it must match bit-for-bit. The fp8/bf8 -# gates are first-cut headroom values and may want tightening once measured on a -# GPU. +# is an exact integer accumulation so it must match bit-for-bit. +# +# The fp8/bf8 numbers are still first-cut headroom, and they are deliberately +# unchanged by the OCP/FNUZ fix. That fix removed a real source of error on +# gfx950 -- the reference used to be FNUZ while the kernel ran OCP, a factor-of- +# two shift these gates were wide enough to swallow -- but it changed nothing on +# gfx942, which was FNUZ-correct all along and still needs this much room purely +# for 3-/2-bit quantization. Tightening is a measurement, not a deduction: run +# the sweep on both archs and set each gate from the observed maximum. _TOL = { "fp16": 2e-3, "bf16": 1.5e-2, @@ -102,6 +108,16 @@ _LAYOUT_WORD = {"r": "row", "c": "col"} +# Same convention as the sibling *_gpu_correctness.py tests and as +# SKIP_RETURN_CODE in dispatcher/tests/CMakeLists.txt. This file is not +# ctest-registered, so the exit code is the only signal a CI lane gets: returning +# 0 from _main() on a CPU-only or hipcc-less runner would report a green PASS for +# the 60-case matrix below -- the lane's only int8 coverage -- without a single +# kernel having been built. The caller in groovy/vars/ck.groovy wraps this in +# run_ok, which translates 77 into a logged skip and lets any other non-zero +# exit fail the lane. +SKIP_EXIT = 77 + def _emulate_input(x: np.ndarray, dtype: str) -> np.ndarray: """Round an fp32 operand to the kernel's storage dtype so the CPU reference @@ -259,7 +275,7 @@ def _main() -> int: reason = _gpu_environment_reason() if reason: print(f"SKIP: {reason}") - return 0 + return SKIP_EXIT arch = detect_gpu_arch() print("=" * 78) diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_rowcolquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_gemm_rowcolquant_gpu_correctness.py index 5d86bbafffc6..dbd79155c90c 100644 --- a/projects/composablekernel/dispatcher/tests/test_gemm_rowcolquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_gemm_rowcolquant_gpu_correctness.py @@ -34,11 +34,20 @@ Run: python3 -m pytest test_gemm_rowcolquant_gpu_correctness.py -v python3 test_gemm_rowcolquant_gpu_correctness.py # standalone + python3 test_gemm_rowcolquant_gpu_correctness.py --gfx gfx942 + +The standalone path is what ctest and the Jenkins lane drive; it exits 77 (see +SKIP_EXIT) when the box cannot run the test at all, so an unrunnable runner is +recorded as Skipped rather than as a vacuous Passed or a spurious Failed. """ +import argparse +import logging +import subprocess import sys import tempfile from pathlib import Path +from typing import Optional import numpy as np import pytest @@ -80,11 +89,43 @@ def _find_python_dir() -> Path: ) +log = logging.getLogger(__name__) + +PASS = "PASS" +FAIL = "FAIL" +SKIP_EXIT = 77 # ctest SKIP_RETURN_CODE; see main() + +# RowColQuant is a block-scale fp8/bf8 op: it needs native fp8 MFMA. A GPU +# outside this set clears conftest's gpu_available() probe and then fails at +# hipcc, which reads as a correctness FAIL rather than "this box cannot run it". +# Matches SUPPORTED_ARCHS in the grouped tensorquant/rowcolquant GPU tests. +_SUPPORTED_ARCHS = ("gfx942", "gfx950") + + def _have_ml_dtypes() -> bool: # rowcolquant's fp8 encoders come from ml_dtypes via its own utils. return u.fp8_encoding_available() +def _has_hipcc() -> bool: + try: + subprocess.run(["hipcc", "--version"], capture_output=True, timeout=10) + return True + except Exception: + return False + + +def _safe_detect_arch() -> Optional[str]: + """_detect_gpu_arch() raises on absence/unsupported; the gates want None.""" + try: + return _detect_gpu_arch() + except Exception: + return None + + +_GFX_ARCH = _safe_detect_arch() + + # ============================================================================= # NumPy reference: C[m,n] = AQ[m] * BQ[n] * sum_k A[m,k] B[k,n] # ============================================================================= @@ -102,8 +143,11 @@ def reference_rowcolquant_gemm(A_dec, B_dec, AQ, BQ) -> np.ndarray: # ============================================================================= -def _run_case(dtype: str, M: int, N: int, K: int, out_dir: Path): - arch = _detect_gpu_arch() +def _run_case(dtype: str, M: int, N: int, K: int, out_dir: Path, + arch: Optional[str] = None): + # main() resolves the arch once (honouring --gfx) and passes it down; the + # pytest entry points leave it None and take the detected one. + arch = arch or _detect_gpu_arch() make_cfg = default_fp8_config if dtype == "fp8" else default_bf8_config config = make_cfg(gfx_arch=arch) @@ -149,10 +193,19 @@ def _run(so_path): _SKIP_NO_GPU = pytest.mark.skipif(not _have_gpu(), reason="no ROCm GPU detected") _SKIP_NO_MLD = pytest.mark.skipif(not _have_ml_dtypes(), reason="ml_dtypes not installed") +# gpu_available() only says "a ROCm GPU is present". Without this second gate a +# gfx90a box runs the fp8 build and reports a red FAIL for hardware that was +# never in scope. +_SKIP_BAD_ARCH = pytest.mark.skipif( + _GFX_ARCH not in _SUPPORTED_ARCHS, + reason=f"RowColQuant needs native fp8 ({'/'.join(_SUPPORTED_ARCHS)}); " + f"detected {_GFX_ARCH}", +) @_SKIP_NO_GPU @_SKIP_NO_MLD +@_SKIP_BAD_ARCH @pytest.mark.parametrize("dtype", ["fp8", "bf8"]) def test_rowcolquant_gpu_matches_reference(dtype, tmp_path): max_rel, _ = _run_case(dtype, M=256, N=256, K=512, out_dir=tmp_path) @@ -161,22 +214,91 @@ def test_rowcolquant_gpu_matches_reference(dtype, tmp_path): @_SKIP_NO_GPU @_SKIP_NO_MLD +@_SKIP_BAD_ARCH def test_rowcolquant_gpu_not_all_zeros(tmp_path): _run_case("fp8", M=256, N=256, K=512, out_dir=tmp_path) -if __name__ == "__main__": - if not _have_gpu(): - print("SKIP: no GPU"); raise SystemExit(0) +# ============================================================================= +# Standalone entry point (the one ctest and Jenkins drive) +# ============================================================================= + + +def _case(dtype: str): + """Build a never-raising case fn returning (status, detail). + + The TESTS-loop convention is that a case reports its own verdict; main()'s + ``except Exception`` is a backstop for genuine bugs, not the normal path. + """ + def run(out_dir: Path, arch: str): + max_rel, time_ms = _run_case(dtype, 256, 256, 512, out_dir, arch) + if max_rel > 0.05: + return FAIL, f"rowcolquant {dtype}: max_rel={max_rel:.4f} > tol=0.05" + return PASS, (f"rowcolquant {dtype}: max_rel={max_rel:.4f} " + f"time_ms={time_ms:.3f}") + return run + + +TESTS = [(f"rowcolquant_{dt}", _case(dt)) for dt in ("fp8", "bf8")] + + +def main() -> int: + parser = argparse.ArgumentParser( + description="RowColQuant (non-grouped) GPU correctness tests") + # No hardcoded default: an explicit --gfx is an override the caller is + # trusted on, but falling back to a fixed arch would make a CPU-only or + # gfx90a box compile for gfx950 and crash instead of skipping. + parser.add_argument("--gfx", default=None, + help=f"GPU arch override (default: auto-detect; " + f"{'/'.join(_SUPPORTED_ARCHS)} only)") + parser.add_argument("-v", "--verbose", action="store_true") + parser.add_argument("--output-dir", type=Path, default=None) + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s: %(message)s", + ) + + # The pytest path gates via the skipif marks above; this standalone path is + # what CI drives, so it needs the same gates spelled out. SKIP_EXIT keeps a + # clean skip distinguishable from a failure -- ctest maps it via + # SKIP_RETURN_CODE and the Jenkins lane via its run_ok helper. + gfx = args.gfx or _GFX_ARCH + if not gfx: + print("SKIP: no supported GPU detected (rocm_agent_enumerator)") + return SKIP_EXIT + if gfx not in _SUPPORTED_ARCHS: + print(f"SKIP: RowColQuant needs native fp8 " + f"({'/'.join(_SUPPORTED_ARCHS)}); detected {gfx}") + return SKIP_EXIT + if not _has_hipcc(): + print("SKIP: hipcc not found in PATH; cannot JIT-compile kernels") + return SKIP_EXIT if not _have_ml_dtypes(): - print("SKIP: ml_dtypes not installed"); raise SystemExit(0) - d = Path(tempfile.mkdtemp(prefix="rowcolquant_gputest_")) - ok = True - for dt in ("fp8", "bf8"): + print("SKIP: ml_dtypes not installed; fp8/bf8 encoding unavailable") + return SKIP_EXIT + + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="rowcolquant_gputest_")) + log.info("Kernel output dir: %s", out_dir) + + results = [] + for name, fn in TESTS: + log.info("--- Running %s ---", name) try: - mr, t = _run_case(dt, 256, 256, 512, d) - print(f"PASS rowcolquant {dt}: max_rel={mr:.4f} time_ms={t:.3f}") - except Exception as e: - ok = False - print(f"FAIL rowcolquant {dt}: {e}") - raise SystemExit(0 if ok else 1) + status, detail = fn(out_dir, gfx) + except Exception as exc: # noqa: BLE001 + status, detail = FAIL, f"{name}: exception: {exc}" + results.append((status, detail)) + log.info("[%s] %s", status, detail) + + print("\n=== Summary ===") + passed = sum(1 for s, _ in results if s == PASS) + for status, detail in results: + print(f" [{status:4s}] {detail}") + print(f"\n{passed}/{len(results)} passed") + return 0 if passed == len(results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_search_space.py b/projects/composablekernel/dispatcher/tests/test_gemm_search_space.py new file mode 100644 index 000000000000..41693bf279f0 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_gemm_search_space.py @@ -0,0 +1,492 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +""" +Dispatcher GEMM search space runner (all bridged GEMM variants). + +Enumerates the GemmKernelConfig search space by calling expand_sweep with the +variant's tile_engine default_ci_config.json (multi_d and multi_abd have their +own, carrying the fused-epilogue block) across all (dtype, layout) strata, +samples a budget-limited subset with a daily rotating seed, compiles each config +via the bridge, runs on GPU, and reports correctness + TFLOPS. + +One runner covers every variant expand_sweep understands -- standard, grouped, +multi_d, multi_abd and stream_k -- because they share the same enumerate/sample/ +build/report machinery and differ only in which Gpu*Runner to drive and how the +numpy reference is computed. Select with --variant. + +This is the dispatcher equivalent of tile_engine's per-op benchmark scripts +driven with TILE_ENGINE_SAMPLING_TIER=daily, with one deliberate difference: +those scripts never verify numerics and always exit 0. This one validates every +kernel against a numpy reference and exits 1 on any mismatch or build failure. + +Exit codes: 0 = all pass, 1 = failure, 77 = skipped (no GPU) per the ctest +SKIP_RETURN_CODE convention. + +Usage: + python3 test_gemm_search_space.py --arch gfx942 --budget 500 + python3 test_gemm_search_space.py --variant grouped --groups 4 + python3 test_gemm_search_space.py --variant multi_abd # fp16/rcrr only + python3 test_gemm_search_space.py --variant multi_d # sweeps MultiDAdd/Multiply + python3 test_gemm_search_space.py --variant multi_d --elementwise-op PassThrough + python3 test_gemm_search_space.py --budget 0 # full space, no cap + python3 test_gemm_search_space.py --budget 500 --seed 42 --json results.json +""" + +import argparse +import hashlib +import json +import random +import sys +import time +from dataclasses import replace +from datetime import date +from pathlib import Path + +_DISPATCHER = Path(__file__).parent.parent +sys.path.insert(0, str(_DISPATCHER / "python")) + +import numpy as np + +from ctypes_utils import detect_gpu_arch +from gemm_utils import ( + GemmProblem, + GpuGemmRunner, + GpuGroupedGemmRunner, + GpuMultiABDRunner, + GpuMultiDGemmRunner, + GroupedGemmProblem, + MultiDGemmProblem, + _cde_reference, + expand_sweep, + numpy_dtype_for, + setup_multiple_gemm_dispatchers, +) + +_TE_GEMM = _DISPATCHER.parent / "tile_engine/ops/gemm" +_CI_CONFIG = _TE_GEMM / "configs/default_ci_config.json" + +# multi_d and multi_abd have their own CI configs carrying the +# ``multi_d_config`` / ``multi_abd_config`` blocks that describe the fused +# epilogue (which element-wise op, how many D tensors). The generic config has +# neither, so pointing every variant at it left expand_sweep with nothing to +# enumerate and the fusion untested. gemm_multi_d_full_benchmark.py already +# reads these same files. +_VARIANT_CI_CONFIG = { + "multi_d": _TE_GEMM / "gemm_multi_d/configs/default_ci_config.json", + "multi_abd": _TE_GEMM / "gemm_multi_abd/configs/default_ci_config.json", +} + + +def _ci_config_for(variant: str) -> Path: + return _VARIANT_CI_CONFIG.get(variant, _CI_CONFIG) +_DTYPES = ["fp16", "bf16", "fp8", "bf8"] +_LAYOUTS = ["rcr", "rrr", "crr", "ccr"] +_RTOL = 2e-2 +_SKIP = 77 + +# The arches whose kernels this sweep can actually build and run, matching +# SUPPORTED_ARCHS in the aquant/abquant GPU tests. The no-GPU gate below is not +# sufficient on its own: a box with an *unsupported* GPU (gfx90a, say) passes it, +# then fails every sampled config at the hipcc step and reports FAIL. That is the +# same false signal the exit-77 convention exists to remove, pointing the other +# way -- so gate on the resolved arch, forced or detected, exactly as bquant does. +_SUPPORTED_ARCHS = ("gfx942", "gfx950") + +# Per-variant defaults for --dtypes / --layouts. multi_abd and multi_d are +# fp16-only end-to-end (codegen, ctypes lib and runner all assume fp16) and take +# the 4-char (A,B,C/E,D) layout code rather than the 3-char one. Only the A and +# B chars vary -- the TE epilogue writes C/E and reads D row-major. +# +# Widening multi_d from one layout to four costs no extra runtime: _sample() +# allocates floor(budget / n_strata) configs per stratum, so a fixed --budget is +# redistributed rather than multiplied. multi_abd stays single-layout until its +# bridge is verified the same way multi_d now is. +_MULTI_LAYOUTS = ["rcrr", "rrrr", "ccrr", "crrr"] +_VARIANT_DEFAULTS = { + "standard": (_DTYPES, _LAYOUTS), + "grouped": (_DTYPES, _LAYOUTS), + "stream_k": (_DTYPES, _LAYOUTS), + "multi_d": (["fp16"], _MULTI_LAYOUTS), + "multi_abd": (["fp16"], ["rcrr"]), +} +_VARIANTS = tuple(_VARIANT_DEFAULTS) + + +def _daily_seed() -> int: + return int(hashlib.md5(date.today().isoformat().encode()).hexdigest(), 16) % (2**31) + + +def _sample(strata: dict, budget: int, seed: int) -> list: + """Stratified budget-limited sample: equal share per non-empty stratum. + + Each stratum gets floor(budget/n), with ``budget % n`` strata taking one + extra; a stratum smaller than its share releases the difference to a final + top-up draw, so the full budget is used whenever the space allows. + + The strata are enumerated in sorted order and then *shuffled* before the + remainder is handed out, rather than being visited in sorted order. Sorted + order alone made the remainder deterministic in a way the seed could not + move: when ``budget < n`` every stratum's share is 0, the whole budget is + remainder, and it all lands on the lexicographically-first strata. At the + real 16 strata (bf16|bf8|fp16|fp8 x ccr|crr|rcr|rrr), --budget 4 therefore + picked bf16/{ccr,crr,rcr,rrr} for every seed, so the ctest registration was + permanently bf16-only and the daily seed bought nothing. Seeding the shuffle + from the same rng keeps the result reproducible from the seed alone. + """ + nonempty = {k: v for k, v in strata.items() if v} + if not nonempty or budget <= 0: + return [] + rng = random.Random(seed) + per, remainder = divmod(budget, len(nonempty)) + strata_order = sorted(nonempty.items()) + rng.shuffle(strata_order) + selected, leftover = [], [] + for i, (_, configs) in enumerate(strata_order): + # Clamp to the stratum size: budget < n makes `per` 0, and a stratum + # smaller than its share would otherwise ask random.sample for more + # items than exist. + alloc = min(per + (1 if i < remainder else 0), len(configs)) + # Sample indices, not objects: it makes "which ones were not picked" + # exact without relying on object identity or equality. + picked = set(rng.sample(range(len(configs)), alloc)) if alloc > 0 else set() + selected.extend(configs[j] for j in picked) + leftover.extend(configs[j] for j in range(len(configs)) if j not in picked) + if len(selected) < budget and leftover: + selected.extend(rng.sample(leftover, min(budget - len(selected), len(leftover)))) + rng.shuffle(selected) + return selected[:budget] + + +def _max_rel(out: np.ndarray, ref: np.ndarray) -> float: + return float(np.max(np.abs(out - ref))) / (float(np.max(np.abs(ref))) + 1e-12) + + +def _quantize(x: np.ndarray, dtype: str) -> np.ndarray: + """Round-trip fp32 through the kernel's host operand dtype. + + The grouped/multi_d runners cast operands with ``numpy_dtype_for`` (fp16 or + an ml_dtypes bf16/fp8-FNUZ/bf8-FNUZ type), so the reference must see the same + quantized values the kernel does. + """ + return x.astype(numpy_dtype_for(dtype)).astype(np.float32) + + +def _layout_code(cfg, n: int) -> str: + """``rcr``/``rcrr`` style code from a config's row/col layout words.""" + words = [cfg.layout_a, cfg.layout_b, cfg.layout_c, getattr(cfg, "layout_d", "row")] + return "".join(w[0] for w in words[:n]) + + +class _Operands: + """Shared host operands for one sweep, generated once and reused per kernel.""" + + def __init__(self, size: int, seed: int, groups: int, num_d: int): + rng = np.random.default_rng(seed) + self.size = size + self.seed = seed + self.A = (rng.standard_normal((size, size)) * 0.1).astype(np.float32) + self.B = (rng.standard_normal((size, size)) * 0.1).astype(np.float32) + self.problem = GemmProblem(M=size, N=size, K=size) + + # Grouped: deliberately NON-uniform M so a kernel that mishandles the + # per-group offsets fails instead of silently passing. All Ms stay + # multiples of 64 (the CI config's only tile_m) because pad_m is false. + step = max(64, size // 4) + ms = [max(64, size - i * step) for i in range(groups)] + self.groups = [(m, size, size) for m in ms] + self.grouped_problem = GroupedGemmProblem(groups=self.groups) + self.A_list = [self.A[:m] for m in ms] + self.B_list = [self.B for _ in ms] + + # Multi-D: D tensors are MxN, stored fp16 by the runner. + self.Ds = [ + (rng.standard_normal((size, size)) * 0.1).astype(np.float16) + for _ in range(max(num_d, 1)) + ] + + +def _make_runner(variant: str, cfg, so: Path): + if variant == "grouped": + return GpuGroupedGemmRunner( + lib_path=so, dtype=cfg.dtype_a, layout=_layout_code(cfg, 3) + ) + if variant == "multi_d": + return GpuMultiDGemmRunner(lib_path=so) + if variant == "multi_abd": + return GpuMultiABDRunner( + lib_path=so, + layout4=_layout_code(cfg, 4), + a_elementwise_op=cfg.a_elementwise_op, + b_elementwise_op=cfg.b_elementwise_op, + cde_elementwise_op=cfg.cde_elementwise_op, + ) + # standard and stream_k share the single-problem C ABI. + return GpuGemmRunner(lib_path=so) + + +def _invoke(variant: str, runner, ops: _Operands): + """One timed launch. Returns an object with .success/.time_ms.""" + if variant == "grouped": + return runner.run(ops.A_list, ops.B_list, ops.grouped_problem) + if variant == "multi_d": + nd = runner.num_d_tensors + return runner.run( + ops.A, ops.B, ops.Ds[:nd], + MultiDGemmProblem(M=ops.size, N=ops.size, K=ops.size, num_d=nd), + ) + if variant == "multi_abd": + # Verification is a separate final call (see _verify): keeping it off the + # timed path avoids recomputing a full fp32 reference on every repeat. + return runner.run(ops.problem, seed=ops.seed, verify=False) + return runner.run(ops.A, ops.B, ops.problem) + + +def _verify(variant: str, cfg, runner, result, ops: _Operands) -> float: + """max_rel of the GPU result against a numpy reference.""" + if variant == "multi_abd": + # The runner generates its own operands, so it also owns the reference + # (mirroring ck_tile::reference_gemm_multiple_abd). One extra launch. + verified = runner.run(ops.problem, seed=ops.seed, verify=True, verify_tol=_RTOL) + if verified.max_rel is None: + raise RuntimeError(f"multi_abd verification did not run (status={verified.status})") + return verified.max_rel + + # Only GemmResult carries max_rel; the grouped/multi_d result types do not. + if getattr(result, "max_rel", None) is not None: + return result.max_rel + + dtype = cfg.dtype_a + if variant == "grouped": + worst = 0.0 + Bq = _quantize(ops.B, dtype) + for out, A_g in zip(result.outputs, ops.A_list): + ref = _quantize(A_g, dtype) @ Bq + worst = max(worst, _max_rel(np.asarray(out).astype(np.float32), ref)) + return worst + + if variant == "multi_d": + # multi_d's fused epilogue op lives in `elementwise_op`; the similarly + # named `cde_elementwise_op` is the multi_abd field and stays at its + # PassThrough default here, so reading it would reference the wrong op. + # Read it per-config rather than from --elementwise-op: the sweep now + # enumerates MultiDAdd and MultiDMultiply together, so a single pinned + # op would mis-reference half the configs. + # + # nd comes from the kernel, not the config, and the D pool is sized to + # the largest num_d in the sample -- num_d is itself swept over {1, 2}. + nd = runner.num_d_tensors + acc = _quantize(ops.A, dtype) @ _quantize(ops.B, dtype) + ref = _cde_reference(cfg.elementwise_op, acc, ops.Ds[:nd]) + return _max_rel(result.output.astype(np.float32), ref) + + # standard / stream_k + # + # Quantize operands to the *input* dtype and the result to the *output* + # dtype -- they differ (expand_sweep maps fp8/bf8 inputs to an fp16 C), and + # conflating them is what a naive "round everything to `dtype`" reference + # gets wrong: for fp8/bf8 it would round the operands to fp16, leaving the + # reference strictly more precise than the kernel and failing every fp8/bf8 + # config by construction. + acc = _quantize(ops.A, dtype) @ _quantize(ops.B, dtype) + ref = _quantize(acc, cfg.dtype_c) + return _max_rel(result.output, ref) + + +def run(args) -> int: + # fallback="" so a failed detection is falsy and the skip below can fire. + # detect_gpu_arch defaults to fallback="gfx942" (ctypes_utils.py), which + # never returns empty -- with the default, a CPU-only runner walked straight + # into hipcc instead of reporting skipped, which is the exact false signal + # the exit-77 convention exists to remove. Matches the aquant/abquant/bquant + # GPU tests, which all spell this out the same way. + arch = args.arch or detect_gpu_arch(fallback="") or None + if not arch: + print("SKIP: no GPU detected and --arch not given") + return _SKIP + + if arch not in _SUPPORTED_ARCHS: + print(f"SKIP: search-space sweep is {'/'.join(_SUPPORTED_ARCHS)}-only; " + f"got {arch}") + return _SKIP + + variant = args.variant + def_dtypes, def_layouts = _VARIANT_DEFAULTS[variant] + dtypes = [d.strip() for d in args.dtypes.split(",")] if args.dtypes else def_dtypes + layouts = [l.strip() for l in args.layouts.split(",")] if args.layouts else def_layouts + budget = args.budget + seed = args.seed if args.seed is not None else _daily_seed() + size = args.size + + print(f"variant={variant} arch={arch} dtypes={dtypes} layouts={layouts}") + print(f"budget={budget if budget > 0 else 'unlimited'} seed={seed} size={size}") + + # --- Enumerate: same JSON config as tile_engine, across all strata --- + print("\nEnumerating search space...") + t0 = time.time() + strata = {} + ci_config = _ci_config_for(variant) + for dtype in dtypes: + for layout in layouts: + configs = expand_sweep(str(ci_config), arch=arch, dtype=dtype, + layout=layout, variant=variant) + if variant == "multi_d" and args.elementwise_op is not None: + # Only when explicitly asked for. The config dataclass derives + # .name (and the codegen's elementwise_ops list) from this field, + # so replacing it is enough to select a different kernel -- which + # makes `--elementwise-op PassThrough` a usable A/B experiment. + # It is not a usable default: PassThrough takes Ds as an unnamed + # parameter pack device-side and _cde_reference returns acc + # host-side, so both sides compute A@B and a kernel that never + # loads D still passes. + configs = [replace(c, elementwise_op=args.elementwise_op) + for c in configs] + strata[f"{dtype}/{layout}"] = configs + total = sum(len(v) for v in strata.values()) + print(f" {total} valid configs across {len(strata)} strata ({time.time()-t0:.1f}s)") + for key, cfgs in sorted(strata.items()): + print(f" {key}: {len(cfgs)}") + + # --- Sample --- + if budget > 0 and budget < total: + configs = _sample(strata, budget=budget, seed=seed) + print(f"\nSampled {len(configs)}/{total} configs (budget={budget})") + else: + configs = [c for v in strata.values() for c in v] + print(f"\nRunning all {len(configs)} configs (no cap)") + + if not configs: + print("ERROR: no configs to run") + return 1 + + # --- Build --- + print(f"\nBuilding {len(configs)} .so files (parallel JIT)...") + t0 = time.time() + so_paths = setup_multiple_gemm_dispatchers(configs, verbose=False) + build_ok = sum(1 for p in so_paths if p is not None) + print(f" Built {build_ok}/{len(configs)} in {time.time()-t0:.1f}s " + f"({len(configs)-build_ok} failed)") + + # --- Run --- + print(f"\nRunning {build_ok} kernels on GPU (M=N=K={size}, " + f"warmup={args.warmup}, repeat={args.repeat})...") + # num_d is per-kernel, but every config in a sweep shares the same operand + # shapes, so generate the largest D set once and slice per kernel. + max_num_d = max((getattr(c, "num_d_tensors", 0) for c in configs), default=0) + ops = _Operands(size=size, seed=seed, groups=args.groups, num_d=max_num_d) + flops = ops.grouped_problem.flops if variant == "grouped" else ops.problem.flops + + results = [] + n_pass = n_fail = n_build_fail = 0 + + for cfg, so in zip(configs, so_paths): + if so is None: + n_build_fail += 1 + results.append({"name": cfg.name, "status": "build_fail"}) + continue + try: + runner = _make_runner(variant, cfg, so) + for _ in range(args.warmup): + _invoke(variant, runner, ops) + times = [] + result = None + for _ in range(max(1, args.repeat)): + result = _invoke(variant, runner, ops) + if result.success: + times.append(result.time_ms) + except Exception as exc: + n_fail += 1 + results.append({"name": cfg.name, "status": "run_error", "error": str(exc)}) + continue + if result is None or not result.success: + n_fail += 1 + results.append({"name": cfg.name, "status": "run_fail", + "error": f"status={getattr(result, 'status', 'unknown')}"}) + continue + # Use avg TFLOPS from timed repeat runs; fall back to result.tflops (single run). + if times: + avg_ms = sum(times) / len(times) + tflops = (flops / (avg_ms * 1e-3)) / 1e12 + else: + tflops = result.tflops + try: + mr = _verify(variant, cfg, runner, result, ops) + except Exception as exc: + n_fail += 1 + results.append({"name": cfg.name, "status": "verify_error", "error": str(exc)}) + continue + ok = mr <= _RTOL + n_pass += ok + n_fail += not ok + results.append({"name": cfg.name, "status": "pass" if ok else "fail", + "tflops": round(tflops, 3), "max_rel": round(mr, 6)}) + + # --- Report --- + print(f"\n{'='*60}") + print(f"Results: {n_pass} pass / {n_fail} fail / {n_build_fail} build-fail " + f"/ {len(configs)} total") + + if args.json: + out = {"variant": variant, "arch": arch, "size": size, "seed": seed, + "budget": budget, "total_configs": total, "n_pass": n_pass, + "n_fail": n_fail, "n_build_fail": n_build_fail, "kernels": results} + Path(args.json).write_text(json.dumps(out, indent=2)) + print(f"Results written to {args.json}") + + if n_fail > 0 or n_build_fail > 0: + print("\nFailed kernels:") + for r in results: + if r["status"] != "pass": + print(f" {r['name']}: {r['status']} " + f"{r.get('error', r.get('max_rel', ''))}") + + # A build failure is a failure. Excluding it would mean a total codegen or + # compile breakage -- every kernel failing to build, nothing verified -- + # still reports a green lane. + return 0 if n_fail == 0 and n_build_fail == 0 else 1 + + +def main() -> int: + p = argparse.ArgumentParser( + description="Dispatcher GEMM search space runner (all bridged variants)", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument("--variant", default="standard", choices=_VARIANTS, + help="GEMM variant to sweep (default: standard)") + p.add_argument("--arch", default=None, + help="GPU arch (default: auto-detect via rocminfo)") + p.add_argument("--dtypes", default=None, + help="Comma-separated dtypes (default: per-variant; " + f"{','.join(_DTYPES)} for standard/grouped/stream_k, " + "fp16 for multi_d/multi_abd)") + p.add_argument("--layouts", default=None, + help="Comma-separated layouts (default: per-variant; " + f"{','.join(_LAYOUTS)} for standard/grouped/stream_k, " + "rcrr for multi_d/multi_abd)") + p.add_argument("--groups", type=int, default=4, + help="Sub-problem count for --variant grouped (default: 4)") + p.add_argument("--elementwise-op", default=None, + choices=("PassThrough", "MultiDAdd", "MultiDMultiply"), + help="Pin the epilogue op for --variant multi_d, overriding " + "the sweep (default: sweep MultiDAdd/MultiDMultiply from " + "the config). PassThrough discards the D tensors on both " + "the device and reference sides, so it verifies plain " + "A@B only. Ignored by every other variant.") + p.add_argument("--budget", type=int, default=500, + help="Max configs to run; 0 = no cap (default: 500)") + p.add_argument("--seed", type=int, default=None, + help="RNG seed; default = daily rotating seed") + p.add_argument("--warmup", type=int, default=5, + help="Warmup iterations per kernel (default: 5)") + p.add_argument("--repeat", type=int, default=5, + help="Timed iterations per kernel (default: 5)") + p.add_argument("--size", type=int, default=1024, + help="M=N=K problem size (default: 1024)") + p.add_argument("--json", default=None, + help="Write results JSON to this path") + return run(p.parse_args()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_utils.py b/projects/composablekernel/dispatcher/tests/test_gemm_utils.py index 3c70ff77c586..47856269069e 100644 --- a/projects/composablekernel/dispatcher/tests/test_gemm_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_gemm_utils.py @@ -37,6 +37,9 @@ _fp32_to_bf8_u8, _bf8_u8_to_fp32, _fnuz_decode_table, + _fp8_decode_table, + _fp8_encode, + numpy_dtype_for, _output_dtype, _dtype_from_kernel_name, _layout_from_kernel_name, @@ -44,6 +47,19 @@ ) +def _require_ml_dtypes(): + """Return the ``ml_dtypes`` module, skipping the test when it is absent. + + It is the independent oracle for the fp8/bf8 tables, so a test that needs it + must skip rather than silently assert nothing. + """ + try: + import ml_dtypes # noqa: WPS433 (lazy: optional dep) + except ImportError: # pragma: no cover - env-dependent + raise unittest.SkipTest("ml_dtypes is required for fp8/bf8 table checks") + return ml_dtypes + + class TestBf16Encoding(unittest.TestCase): """bf16 = top 16 bits of fp32 with round-to-nearest-even.""" @@ -129,10 +145,14 @@ def test_decode_is_consistent_with_encode(self): self.assertTrue(np.all(np.isfinite(d))) def test_saturates_no_inf(self): - # FNUZ has no infinity: huge magnitudes clamp to the finite max. + # FNUZ has no infinity: huge magnitudes clamp to the finite max. Pin the + # format explicitly -- the wrappers otherwise follow the local GPU, and on + # a gfx950 runner these are the OCP maxima (448 / -57344) instead. big = np.array([1e30, -1e30], dtype=np.float32) - self.assertEqual(float(_fp8_u8_to_fp32(_fp32_to_fp8_u8(big))[0]), 240.0) - self.assertEqual(float(_bf8_u8_to_fp32(_fp32_to_bf8_u8(big))[1]), -57344.0) + enc8 = _fp32_to_fp8_u8(big, use_ocp=False) + enc5 = _fp32_to_bf8_u8(big, use_ocp=False) + self.assertEqual(float(_fp8_u8_to_fp32(enc8, use_ocp=False)[0]), 240.0) + self.assertEqual(float(_bf8_u8_to_fp32(enc5, use_ocp=False)[1]), -57344.0) def test_dtype_and_size(self): for enc in (_fp32_to_fp8_u8, _fp32_to_bf8_u8): @@ -141,6 +161,173 @@ def test_dtype_and_size(self): self.assertEqual(u8.itemsize, 1) # must match sizeof(fp8_t/bf8_t) +class TestFp8Bf8OcpEncoding(unittest.TestCase): + """The OCP (gfx950/gfx12) half of the fp8/bf8 codec. + + The decode table is the load-bearing half -- it defines what the numpy + reference multiplies -- so it is checked byte-for-byte against ml_dtypes + rather than spot-checked. Getting the format wrong does not crash; it + silently rescales every reference value by a factor of two. + """ + + def test_decode_matches_ml_dtypes_for_every_byte(self): + ml_dtypes = _require_ml_dtypes() + cases = ( + (4, 3, True, ml_dtypes.float8_e4m3fn), + (4, 3, False, ml_dtypes.float8_e4m3fnuz), + (5, 2, True, ml_dtypes.float8_e5m2), + (5, 2, False, ml_dtypes.float8_e5m2fnuz), + ) + all_bytes = np.arange(256, dtype=np.uint8) + for exp_bits, mant_bits, use_ocp, ml_type in cases: + with self.subTest(dtype=ml_type.__name__): + ours = _fp8_decode_table(exp_bits, mant_bits, use_ocp) + theirs = all_bytes.view(ml_type).astype(np.float32) + same = (ours == theirs) | (np.isnan(ours) & np.isnan(theirs)) + self.assertTrue( + np.all(same), + f"bytes {np.flatnonzero(~same).tolist()} disagree with ml_dtypes", + ) + + def test_ocp_format_ranges(self): + # OCP bias is one less than FNUZ, so e4m3fn reaches 448 (not 240); + # e5m2 keeps 57344 but gains a real infinity. + t43 = _fp8_decode_table(4, 3, True) + t52 = _fp8_decode_table(5, 2, True) + self.assertEqual(float(np.nanmax(t43[np.isfinite(t43)])), 448.0) + self.assertEqual(float(np.nanmin(t43[np.isfinite(t43)])), -448.0) + self.assertEqual(float(np.nanmax(t52[np.isfinite(t52)])), 57344.0) + self.assertEqual(np.count_nonzero(np.isinf(t43)), 0) # e4m3fn has no Inf + self.assertEqual(np.count_nonzero(np.isinf(t52)), 2) # e5m2 does + + def test_ocp_zero_slots(self): + # Unlike FNUZ, 0x80 is a genuine -0.0 rather than the lone NaN. + for tab in (_fp8_decode_table(4, 3, True), _fp8_decode_table(5, 2, True)): + self.assertEqual(float(tab[0x00]), 0.0) + self.assertFalse(np.isnan(tab[0x80])) + self.assertEqual(float(tab[0x80]), 0.0) + + def test_every_representable_value_roundtrips_exactly(self): + for exp_bits, mant_bits in ((4, 3), (5, 2)): + for use_ocp in (False, True): + with self.subTest(bits=(exp_bits, mant_bits), ocp=use_ocp): + table = _fp8_decode_table(exp_bits, mant_bits, use_ocp) + vals = table[np.isfinite(table)] + enc = _fp8_encode(vals, exp_bits, mant_bits, use_ocp=use_ocp) + np.testing.assert_array_equal(table[enc.astype(np.intp)], vals) + + def test_ocp_saturates_to_max_finite_not_inf(self): + # e5m2 has an Inf byte; a naive nearest-neighbour search would round + # huge inputs to it. CK's convert saturates, so the codec must too. + big = np.array([1e30, -1e30], dtype=np.float32) + table = _fp8_decode_table(5, 2, True) + got = table[_fp8_encode(big, 5, 2, use_ocp=True).astype(np.intp)] + self.assertEqual(float(got[0]), 57344.0) + self.assertEqual(float(got[1]), -57344.0) + t43 = _fp8_decode_table(4, 3, True) + got8 = t43[_fp8_encode(big, 4, 3, use_ocp=True).astype(np.intp)] + self.assertEqual(float(got8[0]), 448.0) + + def test_nan_maps_to_a_nan_byte_in_both_formats(self): + nan = np.array([np.nan], dtype=np.float32) + for exp_bits, mant_bits in ((4, 3), (5, 2)): + for use_ocp in (False, True): + with self.subTest(bits=(exp_bits, mant_bits), ocp=use_ocp): + table = _fp8_decode_table(exp_bits, mant_bits, use_ocp) + enc = _fp8_encode(nan, exp_bits, mant_bits, use_ocp=use_ocp) + self.assertTrue(np.isnan(table[int(enc[0])])) + + def test_fnuz_aliases_forward_to_the_shared_table(self): + # _fnuz_decode_table is kept as a thin FNUZ view; callers rely on the + # lru_cache identity, so it must return the cached instance itself. + self.assertIs(_fnuz_decode_table(4, 3), _fp8_decode_table(4, 3, False)) + self.assertIsNot(_fnuz_decode_table(4, 3), _fp8_decode_table(4, 3, True)) + + def test_numpy_dtype_follows_the_requested_format(self): + ml_dtypes = _require_ml_dtypes() + self.assertEqual(numpy_dtype_for("fp8", use_ocp=True), + np.dtype(ml_dtypes.float8_e4m3fn)) + self.assertEqual(numpy_dtype_for("fp8", use_ocp=False), + np.dtype(ml_dtypes.float8_e4m3fnuz)) + self.assertEqual(numpy_dtype_for("bf8", use_ocp=True), + np.dtype(ml_dtypes.float8_e5m2)) + self.assertEqual(numpy_dtype_for("bf8", use_ocp=False), + np.dtype(ml_dtypes.float8_e5m2fnuz)) + + +class TestFp8EncodeMatchesMlDtypes(unittest.TestCase): + """_fp8_encode vs ml_dtypes: identical in range, divergent on overflow. + + _fp8_encode exists because ml_dtypes cannot express CK's saturating convert, + not because the rounding differs -- so everything except overflow has to + agree exactly, and the one place it does not has to be pinned as deliberate. + """ + + _CASES = ( + (4, 3, True, "float8_e4m3fn"), + (5, 2, True, "float8_e5m2"), + (4, 3, False, "float8_e4m3fnuz"), + (5, 2, False, "float8_e5m2fnuz"), + ) + + @staticmethod + def _finite_magnitudes(table): + return np.unique(table[np.isfinite(table)].astype(np.float64)) + + def test_midpoints_round_half_to_even(self): + # The regression this guards: ties were broken toward the smaller + # magnitude, which disagreed with the hardware convert (and ml_dtypes) + # on roughly half of all representable midpoints -- 126 of 252 for + # e4m3. Every one of those is a silently low reference value. + ml_dtypes = _require_ml_dtypes() + for exp_bits, mant_bits, use_ocp, ml_name in self._CASES: + with self.subTest(bits=(exp_bits, mant_bits), ocp=use_ocp): + table = _fp8_decode_table(exp_bits, mant_bits, use_ocp) + mags = self._finite_magnitudes(table) + mids = ((mags[:-1] + mags[1:]) / 2.0).astype(np.float32) + ours = table[_fp8_encode(mids, exp_bits, mant_bits, + use_ocp=use_ocp).astype(np.intp)] + theirs = mids.astype(getattr(ml_dtypes, ml_name)) + np.testing.assert_array_equal(ours.astype(np.float64), + theirs.astype(np.float64)) + + def test_agrees_across_the_whole_finite_range(self): + ml_dtypes = _require_ml_dtypes() + for exp_bits, mant_bits, use_ocp, ml_name in self._CASES: + with self.subTest(bits=(exp_bits, mant_bits), ocp=use_ocp): + table = _fp8_decode_table(exp_bits, mant_bits, use_ocp) + mags = self._finite_magnitudes(table) + max_finite = float(mags[-1]) + vals = np.unique(np.concatenate([ + mags, + -mags, + np.linspace(-max_finite, max_finite, 4001), + ])) + vals = vals[np.abs(vals) <= max_finite].astype(np.float32) + ours = table[_fp8_encode(vals, exp_bits, mant_bits, + use_ocp=use_ocp).astype(np.intp)] + theirs = vals.astype(getattr(ml_dtypes, ml_name)) + np.testing.assert_array_equal(ours.astype(np.float64), + theirs.astype(np.float64)) + + def test_overflow_saturates_where_ml_dtypes_does_not(self): + # Intended divergence, not a gap: CK's convert clamps, so the host + # reference must clamp too. Deferring to ml_dtypes here would poison the + # reference with NaN/Inf on any input the kernel merely saturates. + ml_dtypes = _require_ml_dtypes() + big = np.array([1e30, -1e30], dtype=np.float32) + for exp_bits, mant_bits, use_ocp, ml_name in self._CASES: + with self.subTest(bits=(exp_bits, mant_bits), ocp=use_ocp): + table = _fp8_decode_table(exp_bits, mant_bits, use_ocp) + max_finite = float(self._finite_magnitudes(table)[-1]) + ours = table[_fp8_encode(big, exp_bits, mant_bits, + use_ocp=use_ocp).astype(np.intp)] + self.assertEqual(float(ours[0]), max_finite) + self.assertEqual(float(ours[1]), -max_finite) + theirs = big.astype(getattr(ml_dtypes, ml_name)).astype(np.float64) + self.assertFalse(np.all(np.isfinite(theirs))) + + class TestOutputDtype(unittest.TestCase): """Output (C) element dtype must mirror the codegen's get_output_dtype.""" @@ -321,9 +508,15 @@ def test_ocp_fp8_roundtrip_exact_values(self): def test_ocp_differs_from_fnuz(self): # OCP and FNUZ are distinct encodings; the byte patterns must not match # for a generic value (guards against silently reusing the FNUZ codec). + # + # use_ocp=False is required, not incidental: _fp32_to_fp8_u8 defaults to + # autodetecting the local GPU, so on an OCP arch (gfx950/gfx12) both + # sides would encode OCP and the assertion below would fail for a reason + # unrelated to the codecs being distinct. x = np.array([[0.1, 0.3, 1.5]], dtype=np.float32) self.assertFalse( - np.array_equal(_fp32_to_fp8_ocp_u8(x), _fp32_to_fp8_u8(x)) + np.array_equal(_fp32_to_fp8_ocp_u8(x), + _fp32_to_fp8_u8(x, use_ocp=False)) ) def test_ocp_bf8_shape(self): diff --git a/projects/composablekernel/dispatcher/tests/test_multi_d_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_multi_d_gpu_correctness.py index 7632d5321571..cf168e1d2b03 100644 --- a/projects/composablekernel/dispatcher/tests/test_multi_d_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_multi_d_gpu_correctness.py @@ -13,18 +13,23 @@ Real numeric check -- random A/B/D, element-wise validation vs. fp32 reference, plus a non-zero / finite guard so a mis-launched (all-zero) kernel FAILS. -The multi_d TE op is fp16-only, so only fp16 is exercised. +The multi_d TE op is fp16-only, so only fp16 is exercised, but all four A/B +layout combinations are (rcrr / rrrr / ccrr / crrr). -Runs green on gfx942 (MI300X). SKIPs cleanly with no GPU / hipcc / static lib. +Runs green on gfx942 (MI300X). Exits 77 (= ctest "Skipped") with no GPU / +static lib / hipcc, rather than reporting a vacuous PASS. Run: - python3 -m pytest tests/test_multi_d_gpu_correctness.py -v python3 tests/test_multi_d_gpu_correctness.py + python3 tests/test_multi_d_gpu_correctness.py -v # verbose hipcc output + python3 tests/test_multi_d_gpu_correctness.py --gfx gfx942 """ +import argparse +import logging import shutil import sys -import unittest +from functools import partial from pathlib import Path import numpy as np @@ -40,13 +45,30 @@ setup_multiple_gemm_dispatchers, ) +log = logging.getLogger(__name__) + _TOL = 1e-2 # fp16 GEMM + fp16 D-fuse precision band +PASS = "PASS" +FAIL = "FAIL" +SKIP = "SKIP" + +# ctest reports this as "skipped" rather than passed. Without it, a CPU-only +# runner would report every multi_d case as PASSED -- `-m pytest` exits 0 when +# every test in a file skips -- hiding a missing GPU behind a green result. +SKIP_EXIT = 77 + +# 4-char (A, B, C/E, D) layout codes, matching _VARIANT_DEFAULTS["multi_d"] in +# test_gemm_search_space.py. C and E are row-major throughout: the TE multi_d +# epilogue writes row-major, so only the A and B chars actually vary. +_LAYOUTS = ("rcrr", "rrrr", "ccrr", "crrr") +_LAYOUT_WORD = {"r": "row", "c": "col"} + def _detect_arch(): # subprocess.run with a timeout (rocminfo can hang on misconfigured ROCm - # installs and would otherwise stall test discovery). Validate the parsed - # token is a real gfx string before returning it. Mirrors + # installs and would otherwise stall the run). Validate the parsed token is + # a real gfx string before returning it. Mirrors # dispatcher/python/ctypes_utils.py:detect_gpu_arch. import subprocess try: @@ -80,75 +102,133 @@ def _max_rel_err(got: np.ndarray, ref: np.ndarray) -> float: return float(np.max(np.abs(g - r) / den)) -class TestMultiDGemmGpu(unittest.TestCase): - ARCH = _detect_arch() - - def setUp(self): - if self.ARCH is None: - self.skipTest("no GPU / rocminfo not available") - if not _static_lib_present(): - self.skipTest( - "dispatcher static lib (libck_tile_dispatcher.a) not built; " - "multi_d is registry-routed and needs it" - ) - if shutil.which("hipcc") is None and not Path("/opt/rocm/bin/hipcc").exists(): - self.skipTest("hipcc not found") - - def test_fp16_multid_add(self): - num_d = 2 - cfg = GemmKernelConfig( - dtype_a="fp16", dtype_b="fp16", dtype_c="fp16", dtype_acc="fp32", - layout_a="row", layout_b="col", layout_c="row", - tile_m=128, tile_n=128, tile_k=32, - wave_m=2, wave_n=2, wave_k=1, - warp_tile_m=32, warp_tile_n=32, warp_tile_k=16, - pipeline="compv4", scheduler="intrawave", epilogue="cshuffle", - pad_m=True, pad_n=True, pad_k=True, persistent=False, - variant="multi_d", elementwise_op="MultiDAdd", - num_d_tensors=num_d, d_layout="row", - gfx_arch=self.ARCH, - ) - so_paths = setup_multiple_gemm_dispatchers([cfg], verbose=False) - so = so_paths[0] - if so is None: - self.fail("multi_d fp16 kernel failed to build") - - M, N, K = 512, 512, 512 - problem = MultiDGemmProblem(M=M, N=N, K=K, num_d=num_d) - - rng = np.random.default_rng(11) - A = rng.uniform(-1.0, 1.0, (M, K)).astype(np.float32) - B = rng.uniform(-1.0, 1.0, (K, N)).astype(np.float32) - Ds = [rng.uniform(-1.0, 1.0, (M, N)).astype(np.float32) - for _ in range(num_d)] - - runner = GpuMultiDGemmRunner(so) - self.assertEqual(runner.num_d_tensors, num_d, - "kernel D-count mismatch") - result = runner.run(A.astype(np.float16), B.astype(np.float16), - [d.astype(np.float16) for d in Ds], problem) - self.assertEqual(result.status, 0, - f"multi_d run status={result.status}") - self.assertGreater(result.time_ms, 0.0, "multi_d time_ms not positive") - - # Reference at fp16 input precision, accumulated + fused in fp32: - # E = A@B + D0 + D1 (MultiDAdd). - A_q = A.astype(np.float16).astype(np.float32) - B_q = B.astype(np.float16).astype(np.float32) - E_ref = A_q @ B_q - for d in Ds: - E_ref = E_ref + d.astype(np.float16).astype(np.float32) - - E_got = np.asarray(result.output).astype(np.float32) - self.assertFalse(np.all(E_got == 0.0), "multi_d output all-zero") - self.assertTrue(np.all(np.isfinite(E_got)), "multi_d output NaN/Inf") - mre = _max_rel_err(E_got, E_ref) - self.assertLessEqual( - mre, _TOL, - f"multi_d fp16 max_rel={mre:.4f} > tol={_TOL} (M={M} N={N} K={K})", - ) - print(f"[multi_d/fp16] max_rel={mre:.4e}, time_ms={result.time_ms:.3f}") +def _check_layout(layout: str, gfx_arch: str): + """E = A@B + D0 + D1 for one A/B layout combination. + + The runner reads the layout back out of the kernel name and transposes the + host operands itself, so A and B are handed over as logical (M, K) and + (K, N) in every case and the reference below is layout-independent. + + Returns (status, detail) and never raises: each layout is a separate hipcc + build, and a failure in one must name its layout without hiding the other + three. + """ + num_d = 2 + la, lb, lc, ld = (_LAYOUT_WORD[c] for c in layout) + cfg = GemmKernelConfig( + dtype_a="fp16", dtype_b="fp16", dtype_c="fp16", dtype_acc="fp32", + layout_a=la, layout_b=lb, layout_c=lc, + tile_m=128, tile_n=128, tile_k=32, + wave_m=2, wave_n=2, wave_k=1, + warp_tile_m=32, warp_tile_n=32, warp_tile_k=16, + pipeline="compv4", scheduler="intrawave", epilogue="cshuffle", + pad_m=True, pad_n=True, pad_k=True, persistent=False, + variant="multi_d", elementwise_op="MultiDAdd", + num_d_tensors=num_d, d_layout=ld, + gfx_arch=gfx_arch, + ) + so_paths = setup_multiple_gemm_dispatchers( + [cfg], verbose=log.isEnabledFor(logging.DEBUG) + ) + so = so_paths[0] + if so is None: + return FAIL, f"multi_d/fp16/{layout}: kernel failed to build" + + M, N, K = 512, 512, 512 + problem = MultiDGemmProblem(M=M, N=N, K=K, num_d=num_d) + + rng = np.random.default_rng(11) + A = rng.uniform(-1.0, 1.0, (M, K)).astype(np.float32) + B = rng.uniform(-1.0, 1.0, (K, N)).astype(np.float32) + Ds = [rng.uniform(-1.0, 1.0, (M, N)).astype(np.float32) + for _ in range(num_d)] + + runner = GpuMultiDGemmRunner(so) + if runner.num_d_tensors != num_d: + return FAIL, (f"multi_d/fp16/{layout}: kernel D-count mismatch " + f"({runner.num_d_tensors} != {num_d})") + result = runner.run(A.astype(np.float16), B.astype(np.float16), + [d.astype(np.float16) for d in Ds], problem) + if result.status != 0: + return FAIL, f"multi_d/fp16/{layout}: run status={result.status}" + if result.time_ms <= 0.0: + return FAIL, (f"multi_d/fp16/{layout}: time_ms={result.time_ms:.4f} " + f"is not positive") + + # Reference at fp16 input precision, accumulated + fused in fp32: + # E = A@B + D0 + D1 (MultiDAdd). + A_q = A.astype(np.float16).astype(np.float32) + B_q = B.astype(np.float16).astype(np.float32) + E_ref = A_q @ B_q + for d in Ds: + E_ref = E_ref + d.astype(np.float16).astype(np.float32) + + E_got = np.asarray(result.output).astype(np.float32) + if np.all(E_got == 0.0): + return FAIL, f"multi_d/fp16/{layout}: output all-zero" + if not np.all(np.isfinite(E_got)): + return FAIL, f"multi_d/fp16/{layout}: output NaN/Inf" + mre = _max_rel_err(E_got, E_ref) + if mre > _TOL: + return FAIL, (f"multi_d/fp16/{layout}: max_rel={mre:.4f} > tol={_TOL} " + f"(M={M} N={N} K={K})") + return PASS, (f"multi_d/fp16/{layout}: max_rel={mre:.4e}, " + f"time_ms={result.time_ms:.3f}") + + +# One entry per layout rather than one test that loops: each is its own hipcc +# build, and a per-layout PASS/FAIL row is what makes a single bad layout +# legible in the summary. +TESTS = [(layout, partial(_check_layout, layout)) for layout in _LAYOUTS] + + +def main(): + parser = argparse.ArgumentParser(description="Multi-D GEMM GPU correctness tests") + # No hardcoded default: it must stay possible to tell "user asked for + # gfx942" from "we are on a GPU-less box", so the skip below can fire. + parser.add_argument("--gfx", default=None, + help="GPU arch override (default: auto-detect)") + parser.add_argument("-v", "--verbose", action="store_true") + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s: %(message)s", + ) + + gfx = args.gfx or _detect_arch() + if not gfx: + print("SKIP: no GPU detected (rocminfo); multi_d GPU tests skipped") + return SKIP_EXIT + if not _static_lib_present(): + print("SKIP: dispatcher static lib (libck_tile_dispatcher.a) not built; " + "multi_d is registry-routed and needs it") + return SKIP_EXIT + if shutil.which("hipcc") is None and not Path("/opt/rocm/bin/hipcc").exists(): + print("SKIP: hipcc not found; cannot build multi_d kernels") + return SKIP_EXIT + + results = [] + for name, fn in TESTS: + log.info("--- Running %s ---", name) + try: + status, detail = fn(gfx) + except Exception as exc: + status, detail = FAIL, f"{name}: exception: {exc}" + results.append((name, status, detail)) + log.info("[%s] %s", status, detail) + + print("\n=== Summary ===") + for name, status, detail in results: + print(f" [{status:4s}] {detail}") + passed = sum(1 for _, s, _ in results if s == PASS) + skipped = sum(1 for _, s, _ in results if s == SKIP) + failed = len(results) - passed - skipped + print(f"\n{passed}/{len(results) - skipped} passed" + + (f", {skipped} skipped" if skipped else "")) + + return 0 if failed == 0 else 1 if __name__ == "__main__": - unittest.main(verbosity=2) + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_preshuffle_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_preshuffle_gpu_correctness.py index c763872ae370..c3376ea47cfd 100644 --- a/projects/composablekernel/dispatcher/tests/test_preshuffle_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_preshuffle_gpu_correctness.py @@ -9,7 +9,7 @@ Builds the preshuffle GEMM dispatcher .so (fp16 / rcr — the bridge's supported signature; permute_n is pinned False by BRIDGE_PERMUTE_N), runs a small GEMM on-device via GpuGemmRunner, and compares the GPU output to an fp32 numpy -reference within an fp16-appropriate tolerance. Skips cleanly (exit 0) when no +reference within an fp16-appropriate tolerance. Skips cleanly (exit 77) when no GPU / hipcc is available. The preshuffle kernel pre-permutes the B (weight) operand into a packed layout @@ -55,6 +55,10 @@ PASS = "PASS" FAIL = "FAIL" +# ctest SKIP_RETURN_CODE: main() returns this when the box cannot run the test +# at all, so the lane reports Skipped rather than a vacuous Passed or a Failed. +SKIP_EXIT = 77 + def _has_gpu() -> bool: try: @@ -185,13 +189,19 @@ def main() -> int: if not _has_gpu(): print("SKIP: no supported GPU detected (rocminfo); preshuffle GPU test skipped") - return 0 + return SKIP_EXIT gfx = args.gfx or _resolve_arch(None) log.info("Running preshuffle GEMM GPU correctness on %s", gfx) try: status, detail = _run_preshuffle_fp16(gfx) + except FileNotFoundError as exc: + # Same gate the pytest entry point applies: the on-device check needs the + # compiled dispatcher artifacts, and a unit-only stage runs without them. + # Absent artifacts is "cannot run here", not "the bridge is wrong". + print(f"SKIP: dispatcher not built; preshuffle GPU test skipped ({exc})") + return SKIP_EXIT except Exception as exc: # noqa: BLE001 status, detail = FAIL, f"preshuffle/fp16: exception: {exc}" diff --git a/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py index 07a6c49df8f9..f1c4c1d9eebe 100644 --- a/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_rowcolquant_gpu_correctness.py @@ -32,6 +32,8 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) +from dispatcher_common import fp8_uses_ocp as _fp8_uses_ocp + from grouped_gemm_rowcolquant_utils import ( RowColQuantGemmProblem, RowColQuantGpuGemmRunner, @@ -43,6 +45,7 @@ log = logging.getLogger(__name__) TOLERANCE = 0.05 # 5% max relative error — fp8/bf8 precision floor +SKIP_EXIT = 77 # ctest SKIP_RETURN_CODE; see main() # --------------------------------------------------------------------------- @@ -99,17 +102,6 @@ def _require_ml_dtypes(): ) -def _fp8_uses_ocp(arch: str) -> bool: - """Mirror the -DCK_USE_OCP_FP8 compile logic in grouped_gemm_rowcolquant_utils.py. - - gfx950 / gfx12 build the kernel with OCP fp8 (e4m3fn / e5m2); every other - supported arch (notably gfx942) uses the native FNUZ format - (e4m3fnuz / e5m2fnuz). The host-side encoding MUST match the format the kernel - was compiled for, otherwise the reinterpreted bytes decode to NaN/Inf on device. - """ - return "gfx950" in arch or "gfx12" in arch - - def _fp8_ml_dtype(dtype: str): """Return the ml_dtypes fp8 type matching the compiled kernel's arch.""" _require_ml_dtypes() @@ -285,6 +277,12 @@ def test_rowcolquant_timing_positive(tmp_path): # Standalone runner (mirrors test_bquant_gpu_correctness.py style) # --------------------------------------------------------------------------- +# Mirrors the pytest cases above one-for-one. ctest runs this file as a script +# (not via `-m pytest`) so that a missing GPU exits 77 = Skipped, so anything +# absent from this list is not covered by ctest at all. +# +# test_*_timing_positive has no entry of its own because _run_one already fails +# the case when time_ms <= 0. TESTS = [ ("C4/fp8", lambda od, gfx: _run_one( "C4/fp8", default_fp8_config(gfx_arch=gfx), 128, 128, 192, @@ -292,13 +290,22 @@ def test_rowcolquant_timing_positive(tmp_path): ("C4/bf8", lambda od, gfx: _run_one( "C4/bf8", default_bf8_config(gfx_arch=gfx), 128, 128, 192, *_make_inputs(128, 128, 192, "bf8"), Path(od), gfx_arch=gfx)), + # Non-square M/N/K to stress stride math. + ("rect/fp8", lambda od, gfx: _run_one( + "rect/fp8", default_fp8_config(gfx_arch=gfx), 64, 256, 128, + *_make_inputs(64, 256, 128, "fp8"), Path(od), gfx_arch=gfx)), ] def main(): parser = argparse.ArgumentParser( description="RowColQuant GPU correctness tests") - parser.add_argument("--gfx", default="gfx950") + # No hardcoded default: an explicit --gfx is an override the caller is + # trusted on, but falling back to a fixed arch would make a CPU-only or + # gfx90a box compile for gfx950 and crash instead of skipping. + parser.add_argument("--gfx", default=None, + help=f"GPU arch override (default: auto-detect; " + f"{'/'.join(_SUPPORTED_ARCHES)} only)") parser.add_argument("-v", "--verbose", action="store_true") parser.add_argument("--output-dir", type=Path, default=None) args = parser.parse_args() @@ -308,6 +315,27 @@ def main(): format="%(levelname)s: %(message)s", ) + # The pytest path gates on the requires_gpu marker; this standalone path is + # what CI drives, so it needs the same gates spelled out. SKIP_EXIT keeps a + # clean skip distinguishable from a failure -- ctest maps it via + # SKIP_RETURN_CODE and the Jenkins lane via its run_ok helper. + gfx = args.gfx or _GFX_ARCH + if not gfx: + print("SKIP: no supported GPU detected (rocm_agent_enumerator)") + return SKIP_EXIT + if gfx not in _SUPPORTED_ARCHES: + print(f"SKIP: RowColQuant needs native fp8 " + f"({'/'.join(_SUPPORTED_ARCHES)}); detected {gfx}") + return SKIP_EXIT + if not _has_hipcc(): + print("SKIP: hipcc not found in PATH; cannot JIT-compile kernels") + return SKIP_EXIT + if _ml_dtypes is None: + # _require_ml_dtypes raises pytest.skip's Skipped, which derives from + # BaseException and so slips past the per-test except Exception below. + print("SKIP: ml_dtypes not installed; fp8/bf8 encoding unavailable") + return SKIP_EXIT + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="rowcolquant_gpu_test_")) log.info("Kernel output dir: %s", out_dir) @@ -315,7 +343,7 @@ def main(): for name, fn in TESTS: log.info("--- Running %s ---", name) try: - status, detail = fn(out_dir, args.gfx) + status, detail = fn(out_dir, gfx) except Exception as exc: status, detail = "FAIL", f"{name}: exception: {exc}" results.append((name, status, detail)) diff --git a/projects/composablekernel/dispatcher/tests/test_search_space_sampling.py b/projects/composablekernel/dispatcher/tests/test_search_space_sampling.py new file mode 100644 index 000000000000..46eca129138d --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_search_space_sampling.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +""" +Unit tests for test_gemm_search_space.py's sampler and skip path. + +Both behaviours regressed silently once before and neither is observable from +the runner's own output, so they are pinned here rather than left to the GPU +lane: these run on any box, with no GPU, hipcc or built kernels. +""" + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_TESTS = Path(__file__).parent +sys.path.insert(0, str(_TESTS.parent / "python")) + + +def _load_runner(): + """Import the runner by path. + + It is a script whose name begins with ``test_``, so a plain ``import`` would + have pytest collect it as a second test module. + """ + spec = importlib.util.spec_from_file_location( + "_gemm_search_space_runner", _TESTS / "test_gemm_search_space.py" + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +runner = _load_runner() + +# The real 16 strata the standard variant enumerates. +_DTYPES = ["bf16", "bf8", "fp16", "fp8"] +_LAYOUTS = ["ccr", "crr", "rcr", "rrr"] + + +def _strata(per_stratum=10): + return { + f"{d}/{l}": [f"{d}/{l}/cfg{i}" for i in range(per_stratum)] + for d in _DTYPES + for l in _LAYOUTS + } + + +def _dtypes_of(selected): + return {s.split("/")[0] for s in selected} + + +def test_budget_below_stratum_count_varies_which_strata_are_picked(): + """budget < n_strata must let the seed move which *strata* are visited. + + The remainder used to be handed out in sorted stratum order, so with + ``budget=4`` over 16 strata every share was 0, the whole budget was + remainder, and it always landed on bf16/{ccr,crr,rcr,rrr}. The ctest + registration was therefore permanently bf16-only and the daily rotating + seed bought nothing. + + This asserts on the stratum keys, not on the selected configs: the old code + did vary the configs it drew from within those four fixed bf16 strata, so a + whole-sample comparison passes against the bug and pins nothing. + """ + strata = _strata() + visited = { + frozenset(s.rsplit("/", 1)[0] for s in runner._sample(strata, budget=4, seed=s)) + for s in range(20) + } + assert len(visited) > 1, "every seed drew from the identical four strata" + + +def test_budget_below_stratum_count_reaches_every_dtype_across_seeds(): + """Across seeds the union of picks must reach every dtype, not just bf16.""" + seen = set() + for seed in range(50): + seen |= _dtypes_of(runner._sample(_strata(), budget=4, seed=seed)) + assert seen == set(_DTYPES) + + +def test_sample_is_reproducible_from_seed(): + strata = _strata() + assert runner._sample(strata, budget=7, seed=42) == runner._sample( + strata, budget=7, seed=42 + ) + + +@pytest.mark.parametrize("budget", [1, 4, 15, 16, 17, 64, 160]) +def test_sample_never_exceeds_budget(budget): + selected = runner._sample(_strata(), budget=budget, seed=7) + assert len(selected) == budget + assert len(set(selected)) == len(selected), "sampled the same config twice" + + +def test_sample_handles_strata_smaller_than_their_share(): + """A short stratum releases its slack instead of over-drawing. + + ``random.sample`` raises if asked for more items than exist, so the clamp + and the top-up draw have to agree. + """ + strata = {"bf16/rcr": ["a", "b"], "fp16/rcr": [f"c{i}" for i in range(50)]} + selected = runner._sample(strata, budget=20, seed=3) + assert len(selected) == 20 + assert len(set(selected)) == 20 + + +def test_sample_caps_at_total_when_budget_exceeds_space(): + strata = {"bf16/rcr": ["a", "b"], "fp16/rcr": ["c"]} + assert sorted(runner._sample(strata, budget=99, seed=1)) == ["a", "b", "c"] + + +def test_sample_empty_inputs(): + assert runner._sample({}, budget=10, seed=1) == [] + assert runner._sample({"bf16/rcr": []}, budget=10, seed=1) == [] + assert runner._sample(_strata(), budget=0, seed=1) == [] + + +def test_no_gpu_exits_77(monkeypatch, capsys): + """A CPU-only runner must report skipped, not walk into hipcc. + + ``detect_gpu_arch`` defaults to ``fallback="gfx942"`` and so never returns + falsy; calling it bare made this skip path dead code and a GPU-less box + compiled kernels for an arch it does not have. + + The stub returns the fallback it is handed rather than a fixed "", which is + what the real function does when rocminfo finds nothing. That is what makes + this test discriminating: the runner only skips if it asks for ``fallback=""``. + Stubbing a bare "" would pass against the old bare ``detect_gpu_arch()`` call + too, and pin nothing. + """ + monkeypatch.setattr(runner, "detect_gpu_arch", lambda fallback="gfx942": fallback) + monkeypatch.setattr(sys, "argv", ["test_gemm_search_space.py", "--budget", "1"]) + assert runner.main() == runner._SKIP + assert "SKIP" in capsys.readouterr().out + + +def test_unsupported_arch_exits_77(capsys): + """An unsupported GPU skips too, rather than failing every hipcc step.""" + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(sys, "argv", ["x", "--arch", "gfx90a", "--budget", "1"]) + try: + assert runner.main() == runner._SKIP + finally: + monkeypatch.undo() + assert "SKIP" in capsys.readouterr().out + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/projects/composablekernel/dispatcher/tests/test_streamk_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_streamk_gpu_correctness.py new file mode 100644 index 000000000000..1ee73a1d0c40 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_streamk_gpu_correctness.py @@ -0,0 +1,398 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +GPU correctness tests for the Stream-K GEMM dispatcher — all three reduction +strategies. + +The shared sweep (``test_gemm_search_space.py --variant stream_k``) only ever +builds the ``atomic`` strategy: ``default_ci_config.json`` carries no +``streamk_config`` block, so ``expand_sweep`` falls back to ``["atomic"]``. +``linear`` and ``tree`` — which reduce through a device workspace instead of +global atomics — have never been verified through the ctypes bridge. This test +closes that gap. + +Requires a gfx942 GPU (MI300X) and hipcc in PATH. Skips cleanly (exit 77) when +no GPU is visible or the detected arch is not gfx942, so it is safe to invoke +unconditionally from a CI lane. + +Shapes: + S1 M=4096 N=4096 K=2048 — all three strategies. Output tiles vastly exceed + the occupancy limit, so the partitioner leaves few (if any) Stream-K + tiles and ``num_wgs_per_tile`` stays <= 2 under atomic. + S2 M=128 N=128 K=8192 — linear/tree ONLY. This is the Stream-K-heavy + regime: one output tile split across many workgroups. + ``estimate_num_wgs_per_tile`` is gated on the Atomic strategy + (streamk_gemm_tile_partitioner_impl.hpp:279) and returns 1 otherwise, so + linear/tree keep a tight tolerance here where atomic would need ~5e-1 + rtol for bf16 — a vacuous assertion. The atomic side of this regime is + already covered in C++ by test_streamk_registry.py. + +Tolerance is CK's own split-K-aware bound, ported from +``dispatcher/examples/gemm/cpp/streamk_driver_common.hpp`` + +``include/ck_tile/host/check_err.hpp`` — asserted element-wise, not as a +scalar max-relative-error. + +Run: + python3 test_streamk_gpu_correctness.py + python3 test_streamk_gpu_correctness.py -v # verbose hipcc + python3 test_streamk_gpu_correctness.py --gfx gfx942 + python3 test_streamk_gpu_correctness.py --dtypes fp16,bf16 # trim matrix +""" + +import argparse +import logging +import math +import os +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) + +from ctypes_utils import detect_gpu_arch +from gemm_utils import ( + GemmKernelConfig, + GemmProblem, + GpuGemmRunner, + numpy_dtype_for, + output_dtype_for, + setup_multiple_gemm_dispatchers, +) + +log = logging.getLogger("streamk_gpu") + +PASS, FAIL = "PASS", "FAIL" +SKIP_EXIT = 77 + +# gfx90a stays out: it has no fp8/bf8 MFMA, so half the matrix below is +# unrunnable there. gfx950 used to be excluded for a different reason -- the +# gemm_utils host codec was FNUZ-only while the arch encodes fp8 as OCP -- which +# had nothing to do with Stream-K. That codec now follows the arch (see +# dispatcher_common.fp8_uses_ocp), so gfx950 is back in. +SUPPORTED_ARCHS = ("gfx942", "gfx950") + +STRATEGIES = ("atomic", "linear", "tree") +ALL_DTYPES = ("fp16", "bf16", "fp8", "bf8") + +# pad_m/n/k are all False, so every dimension must be a multiple of the tile. +TILE_M = TILE_N = TILE_K = 64 + +S1 = (4096, 4096, 2048) # dp-dominant: all strategies +S2 = (128, 128, 8192) # sk-dominant: linear/tree only + +# Upper bound on the workgroups reducing into a single output tile. Not knowable +# from Python (it needs hipOccupancyMaxActiveBlocksPerMultiprocessor), which is +# exactly why the shapes above pin it to a provable bound. +S1_ATOMIC_WGS_PER_TILE = 2 +NON_ATOMIC_WGS_PER_TILE = 1 + + +# --------------------------------------------------------------------------- +# Tolerance: a Python port of CK's split-K-aware bound +# --------------------------------------------------------------------------- + +# Explicit-mantissa bit counts, matching ck_tile::numeric::mant(): +# fp32 numeric.hpp:92 | fp16 half.hpp:233 | bf16 bfloat16.hpp:453 +# fp8/e4m3 float8.hpp:220 | bf8/e5m2 float8.hpp:238 +_MANT_BITS = {"fp32": 23, "fp16": 10, "bf16": 7, "fp8": 3, "bf8": 2} + + +def _relative_threshold(compute: str, out: str, acc: str, n_accum: int) -> float: + """Port of ck_tile::get_relative_threshold(n_accum).""" + compute_error = 2.0 ** -_MANT_BITS[compute] * 0.5 + output_error = 2.0 ** -_MANT_BITS[out] * 0.5 + midway_error = max(compute_error, output_error) + acc_error = 2.0 ** -_MANT_BITS[acc] * 0.5 * n_accum + return max(acc_error, midway_error) + + +def _absolute_threshold( + compute: str, out: str, acc: str, maxv: float, n_accum: int +) -> float: + """Port of ck_tile::get_absolute_threshold(maxv, n_accum).""" + if maxv == 0.0: + return 0.0 + expo = math.floor(math.log2(abs(maxv))) + compute_error = 2.0 ** (expo - _MANT_BITS[compute]) * 0.5 + # Full ULP, not half: the output error also absorbs the hardware-vs-software + # conversion difference. This asymmetry is deliberate in check_err.hpp. + output_error = 2.0 ** (expo - _MANT_BITS[out]) * 1.0 + midway_error = max(compute_error, output_error) + acc_error = 2.0 ** (expo - _MANT_BITS[acc]) * 0.5 * n_accum + return max(acc_error, midway_error) + + +def _streamk_tolerance(dtype: str, out_dtype: str, K: int, num_wgs_per_tile: int, + maxv: float): + """Port of streamk_tolerance() from streamk_driver_common.hpp:85-101. + + Stream-K accumulates a tile's K-range across ``num_wgs_per_tile`` partials, + so the bound is the looser of two regimes: the per-split MAC error (over + ``K / num_wgs`` accumulations at accumulator precision) and the cross-split + reduction error (over ``num_wgs`` accumulations at OUTPUT precision, which is + what atomics and the workspace both use). + """ + k_per_split = -(-K // num_wgs_per_tile) # integer ceil + rtol_base = _relative_threshold(dtype, out_dtype, "fp32", k_per_split) + atol_base = _absolute_threshold( + dtype, out_dtype, "fp32", maxv / num_wgs_per_tile, k_per_split + ) + rtol_split_k = _relative_threshold(out_dtype, out_dtype, out_dtype, + num_wgs_per_tile) + atol_split_k = _absolute_threshold(out_dtype, out_dtype, out_dtype, maxv, + num_wgs_per_tile) + return max(rtol_base, rtol_split_k), max(atol_base, atol_split_k) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _quantize(x: np.ndarray, dtype: str) -> np.ndarray: + """Round through ``dtype`` and back to fp32. + + The kernel consumes rounded operands, so the reference must too — otherwise + the comparison measures the quantization step, not the kernel. + """ + return x.astype(numpy_dtype_for(dtype)).astype(np.float32) + + +def _max_rel(got: np.ndarray, ref: np.ndarray) -> float: + """Global relative error, for logging only. The assertion is element-wise.""" + denom = float(np.max(np.abs(ref))) + if denom == 0.0: + return float(np.max(np.abs(got))) + return float(np.max(np.abs(got.astype(np.float32) - ref))) / denom + + +def _detect_arch(): + try: + arch = detect_gpu_arch(fallback="") + except Exception: + return None + return arch or None + + +def _make_config(dtype: str, strategy: str, gfx_arch: str) -> GemmKernelConfig: + """One Stream-K kernel config, spelled out rather than defaulted. + + Mirrors what the CI sweep builds green today for variant=stream_k on gfx942. + ``epilogue`` must stay "cshuffle": the codegen emits nothing for any other + value, which surfaces later as a confusing "no .hpp generated" build error. + """ + return GemmKernelConfig( + dtype_a=dtype, + dtype_b=dtype, + dtype_c=output_dtype_for(dtype), + dtype_acc="fp32", + layout_a="row", + layout_b="col", + layout_c="row", + tile_m=TILE_M, + tile_n=TILE_N, + tile_k=TILE_K, + wave_m=2, + wave_n=2, + wave_k=1, + warp_tile_m=32, + warp_tile_n=32, + warp_tile_k=16, + pipeline="compv3", + scheduler="intrawave", + epilogue="cshuffle", + pad_m=False, + pad_n=False, + pad_k=False, + persistent=False, + variant="stream_k", + reduction_strategy=strategy, + gfx_arch=gfx_arch, + ) + + +class _Shape: + """One shape: fp32 operands plus a per-dtype quantized reference. + + A and B are generated once and shared by every dtype/strategy; the fp32 + matmul is memoized per dtype so the 4096-cubed reference is computed four + times, not twelve. + """ + + def __init__(self, label: str, M: int, N: int, K: int, seed: int): + self.label = label + self.M, self.N, self.K = M, N, K + rng = np.random.default_rng(seed) + self.A = (rng.standard_normal((M, K)) * 0.1).astype(np.float32) + self.B = (rng.standard_normal((K, N)) * 0.1).astype(np.float32) + self.problem = GemmProblem(M=M, N=N, K=K) + self._refs = {} + + def reference(self, dtype: str) -> np.ndarray: + ref = self._refs.get(dtype) + if ref is None: + acc = _quantize(self.A, dtype) @ _quantize(self.B, dtype) + ref = _quantize(acc, output_dtype_for(dtype)) + self._refs[dtype] = ref + return ref + + +def _verify(label, so, shape, dtype, strategy): + """Build a runner for ``so``, launch once, and check the output.""" + runner = GpuGemmRunner(lib_path=so) + result = runner.run(shape.A, shape.B, shape.problem) + if not result.success: + return FAIL, f"{label}: launch failed (status={result.status})" + + got = np.asarray(result.output, dtype=np.float32) + if not np.all(np.isfinite(got)): + return FAIL, f"{label}: output contains NaN/Inf" + # An all-zero C is the signature of a Stream-K kernel that never landed its + # partials — the most likely failure mode for an unexercised strategy. + if not np.any(got): + return FAIL, f"{label}: output is all zeros" + + ref = shape.reference(dtype) + maxv = float(np.max(np.abs(ref))) + num_wgs = ( + S1_ATOMIC_WGS_PER_TILE if strategy == "atomic" else NON_ATOMIC_WGS_PER_TILE + ) + rtol, atol = _streamk_tolerance( + dtype, output_dtype_for(dtype), shape.K, num_wgs, maxv + ) + + bad = np.abs(got - ref) > atol + rtol * np.abs(ref) + rel = _max_rel(got, ref) + if bad.any(): + n_bad = int(bad.sum()) + return FAIL, ( + f"{label}: {n_bad}/{got.size} elements outside " + f"rtol={rtol:.3e} atol={atol:.3e} (max_rel={rel:.3e})" + ) + return PASS, ( + f"{label}: OK max_rel={rel:.3e} rtol={rtol:.3e} atol={atol:.3e} " + f"{result.time_ms:.3f} ms" + ) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def _build_matrix(dtypes): + """(label, dtype, strategy, shape_key) for every case we intend to run. + + S2 is deliberately linear/tree-only — see the module docstring. + """ + cases = [] + for dtype in dtypes: + for strategy in STRATEGIES: + cases.append((f"S1/{dtype}/{strategy}", dtype, strategy, "S1")) + for dtype in dtypes: + for strategy in STRATEGIES: + if strategy == "atomic": + continue + cases.append((f"S2/{dtype}/{strategy}", dtype, strategy, "S2")) + return cases + + +def main(): + parser = argparse.ArgumentParser( + description="Stream-K GPU correctness tests (atomic / linear / tree)" + ) + # No hardcoded default: it must stay possible to tell "user asked for gfx942" + # from "we are on an unrelated box", so the skip below can fire. + parser.add_argument("--gfx", default=None, + help="GPU arch override (default: auto-detect; gfx942 only)") + parser.add_argument("--dtypes", default=",".join(ALL_DTYPES), + help=f"comma-separated subset of {','.join(ALL_DTYPES)}") + parser.add_argument("-v", "--verbose", action="store_true") + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s: %(message)s", + ) + + gfx = args.gfx or _detect_arch() + if not gfx: + print("SKIP: no supported GPU detected (rocminfo); Stream-K GPU tests skipped") + return SKIP_EXIT + if gfx not in SUPPORTED_ARCHS: + print(f"SKIP: Stream-K GPU tests are {'/'.join(SUPPORTED_ARCHS)}-only; " + f"detected {gfx}") + return SKIP_EXIT + + dtypes = [d.strip() for d in args.dtypes.split(",") if d.strip()] + unknown = [d for d in dtypes if d not in ALL_DTYPES] + if unknown: + print(f"ERROR: unknown dtype(s) {unknown}; choose from {list(ALL_DTYPES)}") + return 2 + + # bf16/fp8/bf8 host codecs come from ml_dtypes. fp16 needs nothing, so drop + # the others rather than skipping the whole suite. + try: + import ml_dtypes # noqa: F401 + except ImportError: + dropped = [d for d in dtypes if d != "fp16"] + if dropped: + log.warning("ml_dtypes not installed; dropping %s", ",".join(dropped)) + dtypes = [d for d in dtypes if d == "fp16"] + if not dtypes: + print("SKIP: no runnable dtypes (install ml_dtypes for bf16/fp8/bf8)") + return SKIP_EXIT + + # The default 50 warmup / 100 repeat iterations are pure waste for a + # correctness run; each launch here re-zeros and recomputes the full C. + os.environ.setdefault("CK_TILE_BENCH_WARMUP", "1") + os.environ.setdefault("CK_TILE_BENCH_REPEAT", "2") + + cases = _build_matrix(dtypes) + # One kernel per (dtype, strategy); S1 and S2 share it since M/N/K are + # runtime arguments, not codegen constants. + kernel_keys = sorted({(d, s) for _, d, s, _ in cases}) + configs = [_make_config(d, s, gfx) for d, s in kernel_keys] + + log.info("Building %d Stream-K kernels for %d cases ...", + len(configs), len(cases)) + sos = setup_multiple_gemm_dispatchers(configs, verbose=args.verbose) + by_key = dict(zip(kernel_keys, sos)) + + shapes = { + "S1": _Shape("S1", *S1, seed=42), + "S2": _Shape("S2", *S2, seed=43), + } + + results = [] + for label, dtype, strategy, shape_key in cases: + log.info("--- Running %s ---", label) + so = by_key.get((dtype, strategy)) + if so is None: + results.append((label, FAIL, f"{label}: codegen/compile failed")) + log.info("[%s] %s", FAIL, results[-1][2]) + continue + try: + status, detail = _verify( + label, so, shapes[shape_key], dtype, strategy + ) + except Exception as exc: + status, detail = FAIL, f"{label}: exception: {exc}" + results.append((label, status, detail)) + log.info("[%s] %s", status, detail) + + print("\n=== Summary ===") + passed = sum(1 for _, s, _ in results if s == PASS) + for name, status, detail in results: + print(f" [{status:4s}] {detail}") + print(f"\n{passed}/{len(results)} passed") + + return 0 if passed == len(results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_streamk_registry.py b/projects/composablekernel/dispatcher/tests/test_streamk_registry.py index 5e6c92b9d428..e93b173140f2 100644 --- a/projects/composablekernel/dispatcher/tests/test_streamk_registry.py +++ b/projects/composablekernel/dispatcher/tests/test_streamk_registry.py @@ -45,6 +45,15 @@ SKIP = 77 # ctest SKIP_RETURN_CODE +# Arches the Stream-K dispatcher is validated on. The gate matters because the +# arch can arrive from --arch (CMakeLists.txt passes the arch the build was +# configured with) and so bypass detect_arch() entirely. Without it, a box whose +# arch merely *has* a GPU -- gfx90a, say -- builds and runs the whole matrix and +# reports failures: fp16/bf16/fp8 pass there, but every bf8 case aborts with +# HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION. That is an unsupported-arch result, +# not a regression, and it belongs in the Skipped column. +SUPPORTED_ARCHS = ("gfx942", "gfx950") + # One tile config, all three reduction strategies. TILE = "128x128x64_2x2x1_32x32x16" TILE_CONFIG_JSON = json.dumps( @@ -135,6 +144,12 @@ def main(): if not arch: print("SKIP: no GPU / could not detect gfx arch") return SKIP + # Strip the feature suffixes of a full target triple ("gfx950:sramecc+:xnack-") + # before matching, so a configured-target string is gated the same as a bare arch. + if arch.split(":", 1)[0] not in SUPPORTED_ARCHS: + print(f"SKIP: {arch} is not a supported Stream-K arch " + f"(supported: {', '.join(SUPPORTED_ARCHS)})") + return SKIP print(f"Stream-K registry test on {arch} @ {args.m}x{args.n}x{args.k}") inc = ["-I", str(CK_DIR / "include"), "-I", str(DISPATCHER_DIR / "include")] diff --git a/projects/composablekernel/dispatcher/tests/test_tensor_quant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_tensor_quant_gpu_correctness.py index 564328777a12..4fc3e920caa9 100644 --- a/projects/composablekernel/dispatcher/tests/test_tensor_quant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_tensor_quant_gpu_correctness.py @@ -28,11 +28,20 @@ Run: python3 -m pytest test_tensor_quant_gpu_correctness.py -v python3 test_tensor_quant_gpu_correctness.py # standalone + python3 test_tensor_quant_gpu_correctness.py --gfx gfx942 + +The standalone path is what ctest and the Jenkins lane drive; it exits 77 (see +SKIP_EXIT) when the box cannot run the test at all, so an unrunnable runner is +recorded as Skipped rather than as a vacuous Passed or a spurious Failed. """ +import argparse +import logging +import subprocess import sys import tempfile from pathlib import Path +from typing import Optional import numpy as np import pytest @@ -73,6 +82,37 @@ def _find_python_dir() -> Path: _detect_gpu_arch, ) +log = logging.getLogger(__name__) + +PASS = "PASS" +FAIL = "FAIL" +SKIP_EXIT = 77 # ctest SKIP_RETURN_CODE; see main() + +# TensorQuant is a block-scale fp8/bf8 op: it needs native fp8 MFMA. A GPU +# outside this set clears conftest's gpu_available() probe and then fails at +# hipcc, which reads as a correctness FAIL rather than "this box cannot run it". +# Matches SUPPORTED_ARCHS in the grouped tensorquant/rowcolquant GPU tests. +_SUPPORTED_ARCHS = ("gfx942", "gfx950") + + +def _has_hipcc() -> bool: + try: + subprocess.run(["hipcc", "--version"], capture_output=True, timeout=10) + return True + except Exception: + return False + + +def _safe_detect_arch() -> Optional[str]: + """_detect_gpu_arch() raises on absence/unsupported; the gates want None.""" + try: + return _detect_gpu_arch() + except Exception: + return None + + +_GFX_ARCH = _safe_detect_arch() + # Thin local aliases so the case body reads the same as before (the codecs are # the single canonical copies from conftest.py). @@ -99,8 +139,11 @@ def reference_tensor_quant_gemm(A_dec, B_dec, AQ: float, BQ: float) -> np.ndarra # ============================================================================= -def _run_case(dtype: str, M: int, N: int, K: int, out_dir: Path): - arch = _detect_gpu_arch() +def _run_case(dtype: str, M: int, N: int, K: int, out_dir: Path, + arch: Optional[str] = None): + # main() resolves the arch once (honouring --gfx) and passes it down; the + # pytest entry points leave it None and take the detected one. + arch = arch or _detect_gpu_arch() make_cfg = default_fp8_config if dtype == "fp8" else default_bf8_config config = make_cfg(arch) @@ -147,10 +190,19 @@ def _run(so_path): _SKIP_NO_GPU = pytest.mark.skipif(not _have_gpu(), reason="no ROCm GPU detected") _SKIP_NO_MLD = pytest.mark.skipif(not _have_ml_dtypes(), reason="ml_dtypes not installed") +# gpu_available() only says "a ROCm GPU is present". Without this second gate a +# gfx90a box runs the fp8 build and reports a red FAIL for hardware that was +# never in scope. +_SKIP_BAD_ARCH = pytest.mark.skipif( + _GFX_ARCH not in _SUPPORTED_ARCHS, + reason=f"TensorQuant needs native fp8 ({'/'.join(_SUPPORTED_ARCHS)}); " + f"detected {_GFX_ARCH}", +) @_SKIP_NO_GPU @_SKIP_NO_MLD +@_SKIP_BAD_ARCH @pytest.mark.parametrize("dtype", ["fp8", "bf8"]) def test_tensor_quant_gpu_matches_reference(dtype, tmp_path): max_rel, _ = _run_case(dtype, M=256, N=256, K=512, out_dir=tmp_path) @@ -159,22 +211,91 @@ def test_tensor_quant_gpu_matches_reference(dtype, tmp_path): @_SKIP_NO_GPU @_SKIP_NO_MLD +@_SKIP_BAD_ARCH def test_tensor_quant_gpu_not_all_zeros(tmp_path): _run_case("fp8", M=256, N=256, K=512, out_dir=tmp_path) -if __name__ == "__main__": - if not _have_gpu(): - print("SKIP: no GPU"); raise SystemExit(0) +# ============================================================================= +# Standalone entry point (the one ctest and Jenkins drive) +# ============================================================================= + + +def _case(dtype: str): + """Build a never-raising case fn returning (status, detail). + + The TESTS-loop convention is that a case reports its own verdict; main()'s + ``except Exception`` is a backstop for genuine bugs, not the normal path. + """ + def run(out_dir: Path, arch: str): + max_rel, time_ms = _run_case(dtype, 256, 256, 512, out_dir, arch) + if max_rel > 0.05: + return FAIL, f"tensor_quant {dtype}: max_rel={max_rel:.4f} > tol=0.05" + return PASS, (f"tensor_quant {dtype}: max_rel={max_rel:.4f} " + f"time_ms={time_ms:.3f}") + return run + + +TESTS = [(f"tensor_quant_{dt}", _case(dt)) for dt in ("fp8", "bf8")] + + +def main() -> int: + parser = argparse.ArgumentParser( + description="TensorQuant (non-grouped) GPU correctness tests") + # No hardcoded default: an explicit --gfx is an override the caller is + # trusted on, but falling back to a fixed arch would make a CPU-only or + # gfx90a box compile for gfx950 and crash instead of skipping. + parser.add_argument("--gfx", default=None, + help=f"GPU arch override (default: auto-detect; " + f"{'/'.join(_SUPPORTED_ARCHS)} only)") + parser.add_argument("-v", "--verbose", action="store_true") + parser.add_argument("--output-dir", type=Path, default=None) + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s: %(message)s", + ) + + # The pytest path gates via the skipif marks above; this standalone path is + # what CI drives, so it needs the same gates spelled out. SKIP_EXIT keeps a + # clean skip distinguishable from a failure -- ctest maps it via + # SKIP_RETURN_CODE and the Jenkins lane via its run_ok helper. + gfx = args.gfx or _GFX_ARCH + if not gfx: + print("SKIP: no supported GPU detected (rocm_agent_enumerator)") + return SKIP_EXIT + if gfx not in _SUPPORTED_ARCHS: + print(f"SKIP: TensorQuant needs native fp8 " + f"({'/'.join(_SUPPORTED_ARCHS)}); detected {gfx}") + return SKIP_EXIT + if not _has_hipcc(): + print("SKIP: hipcc not found in PATH; cannot JIT-compile kernels") + return SKIP_EXIT if not _have_ml_dtypes(): - print("SKIP: ml_dtypes not installed"); raise SystemExit(0) - d = Path(tempfile.mkdtemp(prefix="tensor_quant_gputest_")) - ok = True - for dt in ("fp8", "bf8"): + print("SKIP: ml_dtypes not installed; fp8/bf8 encoding unavailable") + return SKIP_EXIT + + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="tensor_quant_gputest_")) + log.info("Kernel output dir: %s", out_dir) + + results = [] + for name, fn in TESTS: + log.info("--- Running %s ---", name) try: - mr, t = _run_case(dt, 256, 256, 512, d) - print(f"PASS tensor_quant {dt}: max_rel={mr:.4f} time_ms={t:.3f}") - except Exception as e: - ok = False - print(f"FAIL tensor_quant {dt}: {e}") - raise SystemExit(0 if ok else 1) + status, detail = fn(out_dir, gfx) + except Exception as exc: # noqa: BLE001 + status, detail = FAIL, f"{name}: exception: {exc}" + results.append((status, detail)) + log.info("[%s] %s", status, detail) + + print("\n=== Summary ===") + passed = sum(1 for s, _ in results if s == PASS) + for status, detail in results: + print(f" [{status:4s}] {detail}") + print(f"\n{passed}/{len(results)} passed") + return 0 if passed == len(results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py index 0fbaa9d8221d..426da979811f 100644 --- a/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py +++ b/projects/composablekernel/dispatcher/tests/test_tensorquant_gpu_correctness.py @@ -32,6 +32,8 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) +from dispatcher_common import fp8_uses_ocp as _fp8_uses_ocp + from grouped_gemm_tensorquant_utils import ( TensorQuantGemmProblem, TensorQuantGpuGemmRunner, @@ -43,6 +45,7 @@ log = logging.getLogger(__name__) TOLERANCE = 0.05 # 5% max relative error — fp8/bf8 precision floor +SKIP_EXIT = 77 # ctest SKIP_RETURN_CODE; see main() # --------------------------------------------------------------------------- @@ -99,17 +102,6 @@ def _require_ml_dtypes(): ) -def _fp8_uses_ocp(arch: str) -> bool: - """Mirror the -DCK_USE_OCP_FP8 compile logic in grouped_gemm_tensorquant_utils.py. - - gfx950 / gfx12 build the kernel with OCP fp8 (e4m3fn / e5m2); every other - supported arch (notably gfx942) uses the native FNUZ format - (e4m3fnuz / e5m2fnuz). The host-side encoding MUST match the format the kernel - was compiled for, otherwise the reinterpreted bytes decode to NaN/Inf on device. - """ - return "gfx950" in arch or "gfx12" in arch - - def _fp8_ml_dtype(dtype: str): """Return the ml_dtypes fp8 type matching the compiled kernel's arch.""" _require_ml_dtypes() @@ -283,6 +275,12 @@ def test_tensorquant_timing_positive(tmp_path): # Standalone runner # --------------------------------------------------------------------------- +# Mirrors the pytest cases above one-for-one. ctest runs this file as a script +# (not via `-m pytest`) so that a missing GPU exits 77 = Skipped, so anything +# absent from this list is not covered by ctest at all. +# +# test_*_timing_positive has no entry of its own because _run_one already fails +# the case when time_ms <= 0. TESTS = [ ("C4/fp8", lambda od, gfx: _run_one( "C4/fp8", default_fp8_config(gfx_arch=gfx), 128, 128, 192, @@ -290,13 +288,22 @@ def test_tensorquant_timing_positive(tmp_path): ("C4/bf8", lambda od, gfx: _run_one( "C4/bf8", default_bf8_config(gfx_arch=gfx), 128, 128, 192, *_make_inputs(128, 128, 192, "bf8"), Path(od), gfx_arch=gfx)), + # Non-square M/N/K to stress stride math. + ("rect/fp8", lambda od, gfx: _run_one( + "rect/fp8", default_fp8_config(gfx_arch=gfx), 64, 256, 128, + *_make_inputs(64, 256, 128, "fp8"), Path(od), gfx_arch=gfx)), ] def main(): parser = argparse.ArgumentParser( description="TensorQuant GPU correctness tests") - parser.add_argument("--gfx", default="gfx950") + # No hardcoded default: an explicit --gfx is an override the caller is + # trusted on, but falling back to a fixed arch would make a CPU-only or + # gfx90a box compile for gfx950 and crash instead of skipping. + parser.add_argument("--gfx", default=None, + help=f"GPU arch override (default: auto-detect; " + f"{'/'.join(_SUPPORTED_ARCHES)} only)") parser.add_argument("-v", "--verbose", action="store_true") parser.add_argument("--output-dir", type=Path, default=None) args = parser.parse_args() @@ -306,6 +313,27 @@ def main(): format="%(levelname)s: %(message)s", ) + # The pytest path gates on the requires_gpu marker; this standalone path is + # what CI drives, so it needs the same gates spelled out. SKIP_EXIT keeps a + # clean skip distinguishable from a failure -- ctest maps it via + # SKIP_RETURN_CODE and the Jenkins lane via its run_ok helper. + gfx = args.gfx or _GFX_ARCH + if not gfx: + print("SKIP: no supported GPU detected (rocm_agent_enumerator)") + return SKIP_EXIT + if gfx not in _SUPPORTED_ARCHES: + print(f"SKIP: TensorQuant needs native fp8 " + f"({'/'.join(_SUPPORTED_ARCHES)}); detected {gfx}") + return SKIP_EXIT + if not _has_hipcc(): + print("SKIP: hipcc not found in PATH; cannot JIT-compile kernels") + return SKIP_EXIT + if _ml_dtypes is None: + # _require_ml_dtypes raises pytest.skip's Skipped, which derives from + # BaseException and so slips past the per-test except Exception below. + print("SKIP: ml_dtypes not installed; fp8/bf8 encoding unavailable") + return SKIP_EXIT + out_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="tensorquant_gpu_test_")) log.info("Kernel output dir: %s", out_dir) @@ -313,7 +341,7 @@ def main(): for name, fn in TESTS: log.info("--- Running %s ---", name) try: - status, detail = fn(out_dir, args.gfx) + status, detail = fn(out_dir, gfx) except Exception as exc: status, detail = "FAIL", f"{name}: exception: {exc}" results.append((name, status, detail)) diff --git a/projects/composablekernel/groovy/vars/ck.groovy b/projects/composablekernel/groovy/vars/ck.groovy index b50b515a8608..998445a8c607 100644 --- a/projects/composablekernel/groovy/vars/ck.groovy +++ b/projects/composablekernel/groovy/vars/ck.groovy @@ -1404,113 +1404,505 @@ def runComprehensiveConvDatasetTests() { ) } -def runTileEngineBasicTests(String compiler) { - buildAndTest( - setup_args: "NO_CK_BUILD", - build_type: 'Release', - execute_cmd: """ - cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${compiler}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx942" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_UNIVERSAL_CONFIG_FILE="default_ci_config.json" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_MULTI_D_CONFIG_FILE="default_ci_config.json" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D GEMM_PRESHUFFLE_CONFIG_FILE="default_ci_config.json" .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_preshuffle_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_multi_d_results.json""" - ) -} - -def runTileEngineGemmTests(String arch, String compiler) { - def execute_cmd - if (arch == "gfx942") { - execute_cmd = """ +@NonCPS +String _dispatcherPerfCmd(String compiler, String gpuTarget, String samplingTier, + String gemmDatatype, String gemmLayout, String multiLayout, + String quantDatatype, String quantLayout, + String problemSizes, String problemConfigs, String perOpInstances) { + // All operators whose CMakeLists list both gfx942 and gfx950 in DESIRED_TARGETS + // run on both arches. Two sets remain gfx950-only: + // + // gemm_aquant / gemm_bquant / gemm_abquant: SUPPORTED_ARCHS = ("gfx950",) in + // their Python layers. The C++ side would compile on gfx942 (standard fp8 + // MFMA), but the Python guards have not been lifted yet. Widening is its own + // change; the gate is conservative and likely to lift soon. + // + // mx_gemm: CMakeLists filters with `target MATCHES "^gfx950"`, so + // benchmark_mx_gemm_all is never created for gfx942 and ninja would fail on + // an unknown target. Also has no fp16 path (MX_GEMM_DATATYPE defaults to + // fp4;fp8), hence quantDatatype below. + // + // gemm_preshuffle and contraction_multi_abd are deferred: preshuffle has no + // dispatcher bridge on develop; contraction_multi_abd has no Python benchmark + // driver. Both get lanes in the next PR once those prerequisites land. + def cmd = """ cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ -D BUILD_CK_TILE_ENGINE="ON" \ -D CMAKE_CXX_COMPILER="${compiler}" \ -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx942" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16;bf8;bf16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_STREAMK_DATATYPE="fp8;fp16" \ - -D GEMM_STREAMK_LAYOUT="rcr" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D GROUPED_GEMM_DATATYPE="fp8;fp16" \ - -D GROUPED_GEMM_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_MULTI_ABD_DATATYPE="fp16" \ - -D GEMM_MULTI_ABD_LAYOUT="rcrr" \ - -D BATCHED_CONTRACTION_DATATYPE="fp16" \ - -D BATCHED_CONTRACTION_LAYOUT="rcr" \ - -D GEMM_ROWCOLQUANT_DATATYPE="fp8;bf8" \ - -D GEMM_ROWCOLQUANT_LAYOUT="rcr" \ - -D GEMM_TENSOR_QUANT_DATATYPE="fp8;bf8" \ - -D GEMM_TENSOR_QUANT_LAYOUT="rcr" \ - -D GROUPED_GEMM_ROWCOLQUANT_DATATYPE="fp8;bf8" \ - -D GROUPED_GEMM_ROWCOLQUANT_LAYOUT="rcr" \ - -D GROUPED_GEMM_TENSORQUANT_DATATYPE="fp8;bf8" \ - -D GROUPED_GEMM_TENSORQUANT_LAYOUT="rcr" \ + -D GPU_TARGETS="${gpuTarget}" \ + -D GEMM_UNIVERSAL_DATATYPE="${gemmDatatype}" \ + -D GEMM_UNIVERSAL_LAYOUT="${gemmLayout}" \ + -D GEMM_UNIVERSAL_MAX_INSTANCES=${perOpInstances} \ -D BATCHED_GEMM_DATATYPE="fp16" \ -D BATCHED_GEMM_LAYOUT="rcr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all benchmark_gemm_streamk_all benchmark_grouped_gemm_all benchmark_gemm_multi_abd_all benchmark_batched_contraction_all benchmark_gemm_rowcolquant_all benchmark_gemm_tensor_quant_all benchmark_grouped_gemm_rowcolquant_all benchmark_grouped_gemm_tensorquant_all benchmark_batched_gemm_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_preshuffle_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_multi_d_results.json && \ - python3 ../tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py . --problem-sizes "1024,1024,1024" --group-counts 8 --warmup 5 --repeat 5 --verbose --json grouped_gemm_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_abd/gemm_multi_abd_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_multi_abd_results.json && \ - python3 ../tile_engine/ops/gemm/batched_contraction/batched_contraction_benchmark.py . --problem-configs "g=2;m=1024;n=1024;k=1024" --warmup 5 --repeat 5 --verbose --json batched_contraction_results.json && \ - python3 ../tile_engine/ops/gemm/block_scale_gemm/gemm_rowcolquant/gemm_rowcolquant_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_rowcolquant_results.json && \ - python3 ../tile_engine/ops/gemm/block_scale_gemm/gemm_tensor_quant/gemm_tensor_quant_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_tensor_quant_results.json && \ - python3 ../tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_rowcolquant/grouped_gemm_rowcolquant_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json grouped_gemm_rowcolquant_results.json && \ - python3 ../tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_tensorquant/grouped_gemm_tensorquant_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json grouped_gemm_tensorquant_results.json && \ - python3 ../tile_engine/ops/gemm/batched_gemm/batched_gemm_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json batched_gemm_results.json """ - } else if (arch == "gfx950") { - execute_cmd = """ - cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${compiler}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx950" \ - -D GEMM_UNIVERSAL_DATATYPE="fp8;fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D GEMM_MULTI_D_DATATYPE="fp16" \ - -D GEMM_MULTI_D_LAYOUT="rcrr;rrrr;crrr;ccrr" \ - -D GEMM_PRESHUFFLE_DATATYPE="fp16;fp8;bf16;bf8" \ - -D GEMM_PRESHUFFLE_LAYOUT="rcr" \ - -D MX_GEMM_DATATYPE="fp4;fp8" \ - -D MX_GEMM_LAYOUT="rcr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all benchmark_gemm_preshuffle_all benchmark_gemm_multi_d_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_preshuffle/gemm_preshuffle_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_preshuffle_results.json && \ - python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_multi_d_results.json && \ - python3 ../tile_engine/ops/gemm/mx_gemm/mx_gemm_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json mx_gemm_results.json """ - } else if (arch == "gfx1201") { - execute_cmd = """ - cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ - -D BUILD_CK_TILE_ENGINE="ON" \ - -D CMAKE_CXX_COMPILER="${compiler}" \ - -D CMAKE_BUILD_TYPE=Release \ - -D GPU_TARGETS="gfx1201" \ - -D GEMM_UNIVERSAL_DATATYPE="fp16" \ - -D GEMM_UNIVERSAL_LAYOUT="rcr;rrr;crr;ccr" \ - -D TILE_ENGINE_SAMPLING_TIER=daily .. && \ - ninja -j${nthreads()} benchmark_gemm_universal_all && \ - python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes "1024,1024,1024" --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json""" - } - buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: execute_cmd) + -D BATCHED_GEMM_MAX_INSTANCES=${perOpInstances} \ + -D BATCHED_CONTRACTION_DATATYPE="fp16" \ + -D BATCHED_CONTRACTION_LAYOUT="rcr" \ + -D BATCHED_CONTRACTION_MAX_INSTANCES=${perOpInstances} \ + -D GEMM_STREAMK_DATATYPE="${gemmDatatype}" \ + -D GEMM_STREAMK_LAYOUT="${gemmLayout}" \ + -D GEMM_STREAMK_MAX_INSTANCES=${perOpInstances} \ + -D GROUPED_GEMM_DATATYPE="${gemmDatatype}" \ + -D GROUPED_GEMM_LAYOUT="${gemmLayout}" \ + -D GROUPED_GEMM_MAX_INSTANCES=${perOpInstances} \ + -D GEMM_MULTI_D_DATATYPE="${gemmDatatype}" \ + -D GEMM_MULTI_D_LAYOUT="${multiLayout}" \ + -D GEMM_MULTI_D_MAX_INSTANCES=${perOpInstances} \ + -D GEMM_MULTI_ABD_DATATYPE="${gemmDatatype}" \ + -D GEMM_MULTI_ABD_LAYOUT="${multiLayout}" \ + -D GEMM_MULTI_ABD_MAX_INSTANCES=${perOpInstances} \ + -D GROUPED_GEMM_ROWCOLQUANT_DATATYPE="${quantDatatype}" \ + -D GROUPED_GEMM_ROWCOLQUANT_LAYOUT="${quantLayout}" \ + -D GROUPED_GEMM_ROWCOLQUANT_MAX_INSTANCES=${perOpInstances} \ + -D GROUPED_GEMM_TENSORQUANT_DATATYPE="${quantDatatype}" \ + -D GROUPED_GEMM_TENSORQUANT_LAYOUT="${quantLayout}" \ + -D GROUPED_GEMM_TENSORQUANT_MAX_INSTANCES=${perOpInstances}""" + if (gpuTarget == "gfx950") { + cmd += """ \ + -D GEMM_AQUANT_DATATYPE="${quantDatatype}" \ + -D GEMM_AQUANT_LAYOUT="${quantLayout}" \ + -D GEMM_AQUANT_MAX_INSTANCES=${perOpInstances} \ + -D GEMM_BQUANT_DATATYPE="${quantDatatype}" \ + -D GEMM_BQUANT_LAYOUT="${quantLayout}" \ + -D GEMM_BQUANT_MAX_INSTANCES=${perOpInstances} \ + -D GEMM_ABQUANT_DATATYPE="${quantDatatype}" \ + -D GEMM_ABQUANT_LAYOUT="${quantLayout}" \ + -D GEMM_ABQUANT_MAX_INSTANCES=${perOpInstances} \ + -D MX_GEMM_DATATYPE="${quantDatatype}" \ + -D MX_GEMM_LAYOUT="${quantLayout}" \ + -D MX_GEMM_MAX_INSTANCES=${perOpInstances}""" + } + cmd += """ \ + -D TILE_ENGINE_SAMPLING_TIER=${samplingTier} .. && \ + ninja -j\$(nproc) benchmark_gemm_universal_all benchmark_batched_gemm_all benchmark_batched_contraction_all benchmark_gemm_streamk_all benchmark_grouped_gemm_all benchmark_gemm_multi_d_all benchmark_gemm_multi_abd_all benchmark_grouped_gemm_rowcolquant_all benchmark_grouped_gemm_tensorquant_all""" + if (gpuTarget == "gfx950") { + cmd += " benchmark_gemm_aquant_all benchmark_gemm_bquant_all benchmark_gemm_abquant_all benchmark_mx_gemm_all" + } + cmd += """ && \ + python3 ../tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json gemm_universal_results.json && \ + python3 ../tile_engine/ops/gemm/batched_gemm/batched_gemm_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json batched_gemm_results.json && \ + python3 ../tile_engine/ops/gemm/batched_contraction/batched_contraction_benchmark.py . --problem-configs ${problemConfigs} --warmup 5 --repeat 5 --verbose --json batched_contraction_results.json && \ + python3 ../tile_engine/ops/gemm_streamk/gemm_streamk_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json gemm_streamk_results.json && \ + python3 ../tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json grouped_gemm_results.json && \ + python3 ../tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json gemm_multi_d_results.json && \ + python3 ../tile_engine/ops/gemm/gemm_multi_abd/gemm_multi_abd_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json gemm_multi_abd_results.json && \ + python3 ../tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_rowcolquant/grouped_gemm_rowcolquant_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json grouped_gemm_rowcolquant_results.json && \ + python3 ../tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_tensorquant/grouped_gemm_tensorquant_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json grouped_gemm_tensorquant_results.json""" + if (gpuTarget == "gfx950") { + cmd += """ && \ + python3 ../tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json gemm_aquant_results.json && \ + python3 ../tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json gemm_bquant_results.json && \ + python3 ../tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json gemm_abquant_results.json && \ + python3 ../tile_engine/ops/gemm/mx_gemm/mx_gemm_benchmark.py . --problem-sizes ${problemSizes} --warmup 5 --repeat 5 --verbose --json mx_gemm_results.json""" + } + return cmd +} + +// Benchmark ahead-of-time tile_engine instances at smoke coverage: one dtype, +// one layout, one problem size, and a fixed per-op AOT instance cap. Sized to +// catch codegen/build/launch breakage on demand without booking a GPU node for +// the length of a full sweep. +// +// Covers every operator that has a dispatcher bridge, so the AOT instance path +// gets the same operator coverage runDispatcherCorrectnessTests gives the JIT +// path. The two are gated separately and neither substitutes for the other: +// correctness JIT-builds each config with hipcc and checks it against a host +// reference, while this lane exercises tile_engine's codegen, instance builder +// and benchmark harness. +def runDispatcherPerfTests(String compiler, String gpuTarget = "gfx942") { + def samplingTier = "500" + def gemmDatatype = "fp16" + def gemmLayout = "rcr" + def problemSizes = '"1024,1024,1024"' + def problemConfigs = '"g=2;m=1024;n=1024;k=1024"' + // gemm_multi_d and gemm_multi_abd take 4-character layouts (A, B, D, E), + // not the 3-character rcr the other ops use. Passing rcr here is not a + // narrower selection, it is a parse error in the instance builder. + def multiLayout = "rcrr" + // Pin instances per op rather than letting TILE_ENGINE_SAMPLING_TIER divide + // itself up. That tier is a *total* budget split evenly across whichever ops + // have a non-empty _DATATYPE (tile_engine/ops/gemm/CMakeLists.txt:16-23), + // so every operator added below would silently shrink the coverage of the + // ones already here -- gfx942 alone would have gone 125 -> 55 instances per + // op. An explicit _MAX_INSTANCES wins over the allocated value (same + // file, :50-59), which keeps lane cost fixed and independent of how many ops + // the lane happens to build. The tier is still passed: each op also forwards + // it to the sampler as --tier, where it picks the sampling strategy. + // + // One dtype and one layout per op is also load-bearing, not just smoke + // sizing: the per-op budget is divided again across (dtype x layout) combos + // with integer truncation, so a multi-combo op can floor a combo to zero + // instances and build nothing. + def perOpInstances = "64" + // Operator distribution across arches -- see _dispatcherPerfCmd for the + // per-op rationale. Summary: + // + // Both arches: gemm_universal, batched_gemm, batched_contraction, + // gemm_streamk, grouped_gemm, gemm_multi_d, gemm_multi_abd, + // grouped_gemm_rowcolquant, grouped_gemm_tensorquant. + // All have gfx950 in their DESIRED_TARGETS and no Python-level arch gate. + // + // gfx950 only: gemm_aquant, gemm_bquant, gemm_abquant, mx_gemm. + // gemm_aquant/bquant/abquant: SUPPORTED_ARCHS = ("gfx950",) in the Python + // layer. gfx942 has fp8 (FNUZ), but those guards haven't been lifted yet. + // mx_gemm: CMakeLists filters with `target MATCHES "^gfx950"`, so the + // benchmark_mx_gemm_all target is never created for gfx942; also has no + // fp16 path, hence quantDatatype below. + // + // Deferred (next PR): gemm_preshuffle (no bridge on develop), and + // contraction_multi_abd (no Python benchmark driver yet). + def quantDatatype = "fp8" + def quantLayout = "rcr" + // Stream-K below is tile_engine/ops/gemm_streamk/ (top level), not + // tile_engine/ops/gemm/gemm_streamk/. Two paths exist and only the first is + // real: it holds the CMakeLists, the instance builder and the benchmark. The + // nested one contains a lone configs/default_config.json that nothing in the + // tree reads. Edit the top-level tree. + def execute_cmd = _dispatcherPerfCmd(compiler, gpuTarget, samplingTier, gemmDatatype, gemmLayout, multiLayout, quantDatatype, quantLayout, problemSizes, problemConfigs, perOpInstances) + try { + buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: execute_cmd) + } finally { + // finally, not a trailing call: a red lane is exactly when the per-kernel + // JSON is worth having, and buildAndTest throws on failure. Same path + // reasoning as runDispatcherCorrectnessTests: execute_cmd runs from + // projects/composablekernel/build while archiveArtifacts resolves against + // the workspace root, so the build-relative path is spelled out in full. + // + // *_results.json, NOT dispatcher_*: this lane drives tile_engine's own + // benchmark drivers, and every file it produces is a tile_engine JSON + // (gemm_universal_results.json, gemm_aquant_results.json, ...). Nothing + // here is written with a dispatcher_ prefix -- only the correctness lane + // writes those -- so the narrower glob matched zero files and + // allowEmptyArchive turned that into a silent success. + // + // allowEmptyArchive stays true only because this is a finally block: on a + // cmake/ninja failure no JSON exists yet, and a throwing archive step + // would mask the real exception. The signal that the run actually + // produced results comes from the drivers, which now exit non-zero when + // no kernels were discovered or the result set is empty, and record + // launches_attempted/succeeded/failed in each JSON. + archiveArtifacts artifacts: "projects/composablekernel/build/*_results.json", + allowEmptyArchive: true + } +} + +// Which GEMM variants of test_gemm_search_space.py each arch can run. +// +// stream_k used to be gfx942-only here, attributed to the "TODO: add gfx950" in +// tile_engine/ops/gemm_streamk/CMakeLists.txt. That was the wrong cause: this +// path never touches those AOT targets -- test_gemm_search_space drives +// gemm_utils, which JIT-builds each config with hipcc and --offload-arch, and +// neither gemm_utils nor the sweep tables carry any stream_k arch guard. The +// real blocker was the fp8/bf8 encoding mismatch on gfx950, which is fixed (see +// dispatcherSweepDtypesFor below). The CMake TODO is still worth resolving, but +// as AOT hygiene, not as a gate on this lane. +// +// gfx1201 is absent on purpose, and the Jenkinsfile has no gfx1201 dispatcher +// stage for the same reason: gemm_utils._SUPPORTED_ARCHES is +// (gfx90a, gfx942, gfx950), so test_gemm_search_space.py raises +// ValueError: Unsupported GPU architecture 'gfx1201' +// before it ever reaches the GPU. Widening that guard does not help -- with it +// widened, expand_sweep returns 0 configs for gfx1201 in all 16 dtype x layout +// strata of default_ci_config.json, because RDNA4 is warp_size 32 and every +// warp tile that config names is 64 wide, and the runner then exits 1 on "no +// configs to run". Covering RDNA4 needs real gfx1201 tile coverage first. +def dispatcherGemmVariantsFor(String arch) { + switch (arch) { + case "gfx942": + case "gfx950": return ["grouped", "multi_d", "multi_abd", "stream_k"] + default: return [] + } +} + +// Which dtypes the GEMM sweeps may ask for on each arch. +// +// gfx950 used to be fp16/bf16 only. Not because the hardware lacks fp8, but +// because the dispatcher's fp8 path was gfx942-shaped on both ends: the JIT +// handed hipcc only --offload-arch, so gfx950 kernels fell into the FNUZ #else +// branch of ck_tile/core/numeric/float8.hpp on the host pass while the device +// pass and the fp8 MFMA used OCP -- and the numpy reference in gemm_utils.py was +// unconditionally FNUZ. Three encodings, one test; the strata passed only because +// the parity gates are loose. Both halves are fixed now: +// +// * Device: gemm_utils._build_compile_jobs appends ocp_arch_defines(gfx_arch) +// (-DCK_USE_OCP_FP8 -DCK_TILE_USE_OCP_FP8), matching what CMake does for a +// gfx950 build and what the quant bridges already did for their own JIT. +// * Host: the gemm_utils fp8/bf8 codec is arch-aware -- _fp8_decode_table +// carries both formats and is verified byte-for-byte against ml_dtypes, and +// numpy_dtype_for picks e4m3fn/e5m2 vs e4m3fnuz/e5m2fnuz from the arch. +// +// dispatcher_common.fp8_uses_ocp is the single predicate behind both, so they +// cannot drift apart again. +def dispatcherSweepDtypesFor(String arch) { + switch (arch) { + case "gfx942": + case "gfx950": return "fp16,bf16,fp8,bf8" + default: return "fp16,bf16" + } +} + +// One test_gemm_search_space.py invocation for a non-standard GEMM variant. +// +// layouts are deliberately omitted: the runner picks per-variant defaults +// (multi_d/multi_abd are fp16/rcrr-only, so naming the standard 4x4 matrix here +// would enumerate configs that variant cannot build). --elementwise-op is +// likewise omitted, so multi_d sweeps the MultiDAdd/MultiDMultiply x num_d +// {1,2} matrix from gemm_multi_d/configs/default_ci_config.json. Passing +// PassThrough here would make the lane vacuous: it discards the D tensors on +// both the device and reference sides, so the fusion under test is never run. +// +// dtypes are omitted for the same reason -- except for "grouped", whose +// _VARIANT_DEFAULTS entry in test_gemm_search_space.py is the full +// fp16,bf16,fp8,bf8 list and would otherwise put fp8/bf8 back on gfx950 through +// the side door. multi_d/multi_abd default to fp16 alone and must keep doing so, +// which is why this is not applied to every variant. stream_k shares the full +// list but is gfx942-only, so it is unaffected either way. +// +// The budget stays far under the standard sweep's 64 because each variant +// JIT-compiles its own .so set and these run in addition to, not instead of, +// that sweep -- 16 is enough to prove each variant still builds and matches its +// reference. +def dispatcherVariantCmd(String arch, String variant) { + def budget = 16 + def dtypeArg = (variant == "grouped") ? " --dtypes ${dispatcherSweepDtypesFor(arch)}" : "" + return """python3 ../dispatcher/tests/test_gemm_search_space.py \ + --variant ${variant} \ + --arch ${arch}${dtypeArg} \ + --budget ${budget} \ + --warmup 5 \ + --repeat 5 \ + --size 1024 \ + --json dispatcher_${variant}_results.json""" +} + +// Verify JIT-compiled dispatcher kernels against a host reference. +// +// The lane runs the whole operator set for the arch -- every dispatcher operator +// with a bridge on develop -- 13 of them. A developer touching the bridge needs +// pre-merge signal on all of them, and a lane covering four cannot give that. +// Cost is held down by the search budget instead: --budget 64 for the standard +// sweep, --budget 16 for each variant sweep. +// +// The layout list stays at the full 4 on purpose, as does the dtype list on +// gfx942: the sweep samples per stratum, so even 64 spreads ~4 configs across +// all 16 dtype x layout combinations rather than testing one combination +// deeply. gfx950 runs 2 dtypes x 4 layouts -- see dispatcherSweepDtypesFor. +// +// --------------------------------------------------------------------------- +// Coverage against the 13 dispatcher operators with a bridge on develop. +// +// Operators covered by this lane (all grouped_gemm / dispatcher-bridge variants): +// gemm_universal, grouped_gemm, multi_d_gemm (4 layouts), multi_abd_gemm, +// batched_gemm, batched_contraction (fp16+bf16+fp32), stream_k, +// grouped_gemm_rowcolquant, grouped_gemm_tensorquant, +// grouped_gemm_aquant, grouped_gemm_abquant (both arches), +// grouped_gemm_bquant C4 (fp8/bf8) + C (fp8i4/bf8i4) + H3 (mx_*), +// mx_gemm. +// The last two lines are gfx950 only -- mx_gemm and bquant H3 for an ISA +// limitation (no scale-MFMA builtins on gfx942), bquant C4/C pending one +// validating gfx942 run. See the gate comment further down. +// +// Operators NOT covered by this lane (no dispatcher bridge on develop yet): +// preshuffled_gemm, gemm_aquant, gemm_abquant, gemm_bquant, +// gemm_rowcolquant, gemm_tensorquant (block_scale_gemm tile_engine ops), +// batched_contraction_multi_abd. +// +// All 13 bridged ops are invoked. Twelve are fully covered; one +// (grouped_gemm_bquant non-rcr layouts) runs but over a narrowed surface. +// +// One non-obvious detail: gemm_universal's int8 coverage comes from +// test_gemm_parity.py's _INT_DTYPES, not from the sweep. test_gemm_search_space +// has no int8 stratum, so dropping the parity call would drop int8 entirely. +// That is why it is invoked here and why its skip path now exits 77 rather +// than 0. +// +// Known gaps in the bridged ops, deliberately not closed here: +// +// grouped_gemm_{a,ab,b}quant All three config builders hardcode +// non-rcr layouts layout="rcr", and the ctypes libs assert +// packed rcr strides, so a non-rcr kernel +// would build and then be rejected at every +// call. Fixing requires ctypes stride +// derivation from the compile-time layout +// types (see plan Step 7). +// --------------------------------------------------------------------------- +@NonCPS +String _dispatcherCorrectnessBaseCmd(String arch, String compiler, String dtypes, int budget) { + return """ + run_ok() { rc=0; "\$@" || rc=\$?; if [ \$rc -eq 77 ]; then echo "SKIP(77): \$*"; return 0; fi; return \$rc; } && \ + cmake -G Ninja -D CMAKE_PREFIX_PATH=/opt/rocm \ + -D CMAKE_CXX_COMPILER="${compiler}" \ + -D CMAKE_BUILD_TYPE=Release \ + -D GPU_TARGETS="${arch}" \ + -D CK_TILE_DISPATCHER=ON \ + -D BUILD_DISPATCHER_BINDINGS=ON \ + -D DISPATCHER_RULE_SET=tests .. && \ + ninja -j\$(nproc) ck_tile_dispatcher dispatcher_gemm_lib && \ + python3 ../dispatcher/tests/test_gemm_search_space.py \ + --arch ${arch} \ + --dtypes ${dtypes} \ + --layouts rcr,rrr,crr,ccr \ + --budget ${budget} \ + --warmup 5 \ + --repeat 5 \ + --size 1024 \ + --json dispatcher_gemm_results.json && \ + run_ok python3 ../dispatcher/tests/test_gemm_parity.py && \ + run_ok python3 ../dispatcher/tests/test_batched_gemm_gpu_correctness.py --gfx ${arch} && \ + run_ok python3 ../dispatcher/tests/test_batched_contraction_gpu_correctness.py --gfx ${arch} && \ + run_ok python3 ../dispatcher/tests/test_grouped_gemm_gpu_correctness.py && \ + run_ok python3 ../dispatcher/tests/test_multi_d_gpu_correctness.py --gfx ${arch} && \ + run_ok python3 ../dispatcher/tests/test_multi_abd_gpu_correctness.py && \ + run_ok python3 ../dispatcher/tests/test_rowcolquant_gpu_correctness.py --gfx ${arch} && \ + run_ok python3 ../dispatcher/tests/test_tensorquant_gpu_correctness.py --gfx ${arch}""" +} + +def runDispatcherCorrectnessTests(String arch, String compiler) { + def budget = 64 + def execute_cmd = _dispatcherCorrectnessBaseCmd(arch, compiler, dispatcherSweepDtypesFor(arch), budget) + // The grouped/multi_d/multi_abd tests are the bridge-level companions to the + // --variant sweeps below: the sweep exercises the search space, they exercise + // the ctypes bridge. They were registered in ctest but never invoked from + // here. All three are wrapped with run_ok: multi_d is script-style and exits + // 77 on a clean skip (a bare 77 would break the && chain), and grouped_gemm / + // multi_abd are unittest-based but still need run_ok to guard against exit 1 + // from environment failures (bad import, setUpClass error) that would kill the + // rest of the chain without running the quant operators. + // + // rowcolquant and tensorquant sit here rather than in the gfx950 block + // because, unlike the other quant ops, both support gfx942 as well -- + // _SUPPORTED_ARCHES is ("gfx942", "gfx950") and their _fp8_uses_ocp mirrors + // the FNUZ/OCP split instead of assuming OCP. Their bridges landed in #10010. + dispatcherGemmVariantsFor(arch).each { variant -> + execute_cmd += " && \\\n " + dispatcherVariantCmd(arch, variant) + } + // Stream-K has its own registry-level driver test; its defaults are the + // full 4x4 matrix (48 driver compiles / 96 GPU runs), so CI pins one + // dtype/layout. Note it takes --arch, not --gfx. + // + // test_streamk_gpu_correctness.py is the bridge-level companion: the + // shared sweep above only ever builds the atomic reduction strategy + // (default_ci_config.json has no streamk_config, so expand_sweep falls + // back to ["atomic"]), leaving linear/tree unverified through ctypes. + // It takes --gfx and self-gates via SUPPORTED_ARCHS, now (gfx942, gfx950). + // + // Both are keyed off the variant list rather than a hardcoded arch == "gfx942" + // branch, so the sweep and its companions can never disagree about which + // archs run Stream-K. test_streamk_registry.py has no arch allow-list of its + // own, which is precisely why it needs the gate here. + if (dispatcherGemmVariantsFor(arch).contains("stream_k")) { + execute_cmd += """ && \ + run_ok python3 ../dispatcher/tests/test_streamk_registry.py --arch ${arch} --datatypes fp16 --layouts rcr && \ + run_ok python3 ../dispatcher/tests/test_streamk_gpu_correctness.py --gfx ${arch}""" + } + // bquant stays gfx950-only, but no longer for the original reason: its config + // builders used to hardcode the gfx950 warp_tile_k (128), which on gfx942 took + // the other branch of CK_GFX950_SUPPORT and produced a kernel that built, ran, + // and returned zeros. grouped_gemm_bquant_utils is arch-aware now + // (_fp8_warp_tile_k / _preshuffleb_warp_tile_k), so the gate is held only + // pending one validating run on gfx942 -- not by a known defect. Flipping it + // means widening SUPPORTED_ARCHS in test_bquant_gpu_correctness.py in the same + // change; the two must not drift apart. + // This is why the lane fans out to a gfx950 node as well as gfx942 -- see the + // stage comment in the Jenkinsfile. + // + // mx_gemm is gfx950 for a harder reason: the scale-MFMA builtins + // (__builtin_amdgcn_mfma_scale_f32_*_f8f6f4) do not exist in the gfx942 ISA, + // so it fails at compile time rather than silently. It also has no + // expand_sweep -- mx_gemm_utils exposes only default_fp8_config() / + // default_fp4_config() -- so it cannot join the --variant sweep above and + // instead runs the two-config smoke test merged in #10132. That test is + // unittest-based and self-gates to gfx950, so it takes no --gfx. It is still + // wrapped with run_ok -- see the NOTE below for why every test in this chain is. + if (arch == "gfx950") { + execute_cmd += """ && \ + run_ok python3 ../dispatcher/tests/test_bquant_gpu_correctness.py --gfx ${arch} && \ + run_ok python3 ../dispatcher/tests/test_mx_gemm_gpu_correctness.py""" + } + // aquant runs on both archs. Its non-preshuffleaq builders use standard fp8 + // MFMA (warp_tile_k=32), present on gfx942. _preshuffleaq_warp_tile_k() is + // now arch-aware (128 gfx950, 64 gfx942), and the test's preshuffleaq cases + // are gated behind PRESHUFFLEAQ_SUPPORTED_ARCHS pending one CI run on gfx942. + execute_cmd += """ && \ + run_ok python3 ../dispatcher/tests/test_aquant_gpu_correctness.py --gfx ${arch}""" + // abquant runs on both archs. Its Python layer was already arch-aware -- + // _eightwaves_warp_tile_k / _preshuffleb_warp_tile_k carry correct gfx942 + // branches -- and its compv3 pipeline uses standard fp8 MFMA that gfx942 + // has; only the test's self-gate and its OCP-hardcoded host codec kept it + // off gfx942. Both are fixed, so this goes outside the arch block. The test + // still skips its eightwaves case on gfx942 (EIGHTWAVES_SUPPORTED_ARCHS) + // pending one green run, and reports that as a per-case SKIP, not a 77. + execute_cmd += """ && \ + run_ok python3 ../dispatcher/tests/test_abquant_gpu_correctness.py --gfx ${arch}""" + // NOTE: run_ok is applied to every script-style test AND every unittest-based + // test above. Unittest-based tests exit 0 on an internal skipTest, but a bad + // import or setUpClass crash still exits 1, which would break the && chain and + // silently drop every subsequent operator. run_ok costs nothing and makes the + // chain resilient to environment failures regardless of test style. + try { + buildAndTest(setup_args: "NO_CK_BUILD", build_type: 'Release', execute_cmd: execute_cmd) + } finally { + // finally, not a trailing call: a red lane is exactly when the per-kernel + // JSON is worth having, and buildAndTest throws on failure. + // + // Path, not a bare glob: cmake_build runs execute_cmd from + // projects/composablekernel/build, while archiveArtifacts resolves + // against the workspace root. A bare "dispatcher_*_results.json" matches + // nothing there, and allowEmptyArchive hides that it matched nothing. + // Glob, not one name: each variant writes its own dispatcher__*.json. + archiveArtifacts artifacts: "projects/composablekernel/build/dispatcher_*_results.json", + allowEmptyArchive: true + } +} + +// Orchestrate correctness and perf dispatcher lanes in parallel. +// Called from the Jenkinsfile as a scripted step to keep the declarative +// pipeline method small enough to avoid the JVM 64KB method-size limit. +// +// rocmnode is passed in rather than called directly: it is defined in the +// Jenkinsfile (`def rocmnode(name)`), and a vars/ script does not inherit +// Jenkinsfile-local methods, so calling it here throws MissingMethodException +// before runOnHealthyNode is ever reached. Same convention as +// the caller +// passes `this.&rocmnode`. +def runDispatcherTests(def rocmnode, boolean runCorrectness, boolean runPerf, String compiler) { + def branches = [:] + if (runCorrectness) { + branches["DISPATCHER_CORRECTNESS gfx942"] = { + runOnHealthyNode(rocmnode("gfx942")) { + deleteDir() + runDispatcherCorrectnessTests("gfx942", compiler) + cleanWs() + } + } + branches["DISPATCHER_CORRECTNESS gfx950"] = { + runOnHealthyNode(rocmnode("gfx950")) { + deleteDir() + runDispatcherCorrectnessTests("gfx950", compiler) + cleanWs() + } + } + } + if (runPerf) { + branches["DISPATCHER_PERF gfx942"] = { + runOnHealthyNode(rocmnode("gfx942")) { + deleteDir() + runDispatcherPerfTests(compiler, "gfx942") + cleanWs() + } + } + branches["DISPATCHER_PERF gfx950"] = { + runOnHealthyNode(rocmnode("gfx950")) { + deleteDir() + runDispatcherPerfTests(compiler, "gfx950") + cleanWs() + } + } + } + if (branches) { + parallel branches + } } def runBuildCKAndTests(String arch) { @@ -1564,3 +1956,4 @@ def runBuildInstancesOnly(String compiler) { -D CMAKE_BUILD_TYPE=Release .. && ninja -j${nthreads()}""" ) } + diff --git a/projects/composablekernel/tile_engine/ops/common/benchmark_utils.py b/projects/composablekernel/tile_engine/ops/common/benchmark_utils.py index f94bc4a969ac..b8ec4cffe952 100644 --- a/projects/composablekernel/tile_engine/ops/common/benchmark_utils.py +++ b/projects/composablekernel/tile_engine/ops/common/benchmark_utils.py @@ -143,7 +143,11 @@ def export_best_kernels(best_kernels: Dict, filename: str, verbose: bool = False def export_json( - results: List[Dict], filename: str, best_kernels: Dict = None, verbose: bool = False + results: List[Dict], + filename: str, + best_kernels: Dict = None, + verbose: bool = False, + launch_counts: Dict = None, ): """Export all results and best kernels to JSON with comprehensive metadata""" from datetime import datetime @@ -222,14 +226,24 @@ def export_json( r.get("tflops", 0) for r in relevant_results ) / len(relevant_results) + benchmark_metadata = { + "timestamp": datetime.now().isoformat(), + "total_kernels_tested": len(results), + "unique_kernels": len(set(r.get("name", "unknown") for r in results)), + "successful_runs": len(successful_results), + "failed_runs": len(results) - len(successful_results), + } + + # Launch accounting, when the caller tracks it + if launch_counts is not None: + attempted = launch_counts.get("attempted", 0) + failed = launch_counts.get("failed", 0) + benchmark_metadata["launches_attempted"] = attempted + benchmark_metadata["launches_succeeded"] = attempted - failed + benchmark_metadata["launches_failed"] = failed + output_data = { - "benchmark_metadata": { - "timestamp": datetime.now().isoformat(), - "total_kernels_tested": len(results), - "unique_kernels": len(set(r.get("name", "unknown") for r in results)), - "successful_runs": len(successful_results), - "failed_runs": len(results) - len(successful_results), - }, + "benchmark_metadata": benchmark_metadata, "performance_summary": { "tflops_stats": { "best": max(tflops_values, default=0), diff --git a/projects/composablekernel/tile_engine/ops/gemm/batched_contraction/batched_contraction_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/batched_contraction/batched_contraction_benchmark.py index 599183016bff..f8f9c4817094 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/batched_contraction/batched_contraction_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/batched_contraction/batched_contraction_benchmark.py @@ -16,6 +16,8 @@ def __init__(self, build_dir: str, verbose: bool = False): self.build_dir = Path(build_dir) self.verbose = verbose self.results = [] + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_batched_contraction_* executables in the build directory""" @@ -243,6 +245,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = self.run_kernel(kernel_path, params) if result: @@ -276,6 +279,8 @@ def benchmark_problem_size( f"{structured_result['bandwidth_gb_s']:.2f} GB/s, " f"{structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -464,6 +469,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): ), "successful_runs": len(successful_results), "failed_runs": len(self.results) - len(successful_results), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "performance_summary": { "tflops_stats": { @@ -595,6 +603,29 @@ def main(): if args.json: benchmark.export_json(args.json, best_kernels) + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/batched_gemm/batched_gemm_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/batched_gemm/batched_gemm_benchmark.py index 8a67f56aad73..45acd8fffb08 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/batched_gemm/batched_gemm_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/batched_gemm/batched_gemm_benchmark.py @@ -16,6 +16,8 @@ def __init__(self, build_dir: str, verbose: bool = False): self.build_dir = Path(build_dir) self.verbose = verbose self.results = [] + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_batched_gemm_* executables in the build directory""" @@ -294,6 +296,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = self.run_kernel(kernel_path, params) if result: @@ -327,6 +330,8 @@ def benchmark_problem_size( print( f" {kernel_info['config_id']}: {structured_result['tflops']:.2f} TFLOPS, {structured_result['bandwidth_gb_s']:.2f} GB/s, {structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -517,6 +522,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): ), "successful_runs": len(successful_results), "failed_runs": len(self.results) - len(successful_results), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "performance_summary": { "tflops_stats": { @@ -669,6 +677,29 @@ def main(): if args.json: benchmark.export_json(args.json, best_kernels) + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py index 86422986464a..19e043d6f18d 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_abquant/gemm_abquant_benchmark.py @@ -15,6 +15,8 @@ def __init__(self, build_dir: str, verbose: bool = False): self.build_dir = Path(build_dir) self.verbose = verbose self.results = [] + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_gemm_abquant_* executables in the build directory.""" @@ -268,6 +270,7 @@ def benchmark_problem_size( "rotating_count": rotating_count, } + self.launch_attempted += 1 result = self.run_kernel(kernel_path, params) if result: @@ -302,6 +305,8 @@ def benchmark_problem_size( f"{structured_result['bandwidth_gb_s']:.2f} GB/s, " f"{structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -382,6 +387,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): "successful_runs": len( [r for r in self.results if r.get("tflops", 0) > 0] ), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "kernel_results": self.results, "best_kernels_by_problem": best_kernels or {}, @@ -470,6 +478,29 @@ def main(): if args.json: benchmark.export_json(args.json, best_kernels) + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py index 2f5534064318..fb105ab40d63 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_aquant/gemm_aquant_benchmark.py @@ -16,6 +16,8 @@ def __init__(self, build_dir: str, verbose: bool = False): self.build_dir = Path(build_dir) self.verbose = verbose self.results = [] + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_gemm_aquant_* executables in the build directory.""" @@ -275,6 +277,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = self.run_kernel(kernel_path, params) if result: @@ -309,6 +312,8 @@ def benchmark_problem_size( f"{structured_result['bandwidth_gb_s']:.2f} GB/s, " f"{structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -431,6 +436,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): "successful_runs": len( [r for r in self.results if r.get("tflops", 0) > 0] ), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "kernel_results": self.results, "best_kernels_by_problem": best_kernels or {}, @@ -540,6 +548,29 @@ def main(): if args.json: benchmark.export_json(args.json, best_kernels) + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py index b19d3702e2fc..2cfe9f0a7b74 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/block_scale_gemm/gemm_bquant/gemm_bquant_benchmark.py @@ -15,6 +15,8 @@ def __init__(self, build_dir: str, verbose: bool = False): self.build_dir = Path(build_dir) self.verbose = verbose self.results = [] + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_gemm_bquant_* executables in the build directory.""" @@ -244,6 +246,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = self.run_kernel(kernel_path, params) if result: @@ -278,6 +281,8 @@ def benchmark_problem_size( f"{structured_result['bandwidth_gb_s']:.2f} GB/s, " f"{structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -363,6 +368,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): "successful_runs": len( [r for r in self.results if r.get("tflops", 0) > 0] ), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "kernel_results": self.results, "best_kernels_by_problem": best_kernels or {}, @@ -457,6 +465,29 @@ def main(): if args.json: benchmark.export_json(args.json, best_kernels) + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_benchmark.py index b35390a1f98e..ce40410565bb 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_benchmark.py @@ -36,6 +36,8 @@ def __init__( self.verbose = verbose self.results = [] self.name = name + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_gemm_* executables in the build directory""" @@ -244,6 +246,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = benchmark_utils.run_kernel( self.build_dir, kernel_path, params, verbose=self.verbose ) @@ -278,6 +281,8 @@ def benchmark_problem_size( print( f" {kernel_info['config_id']}: {structured_result['tflops']:.2f} TFLOPS, {structured_result['bandwidth_gb_s']:.2f} GB/s, {structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json index 99902697f822..b021261c6010 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json @@ -88,5 +88,13 @@ ] } }, + "multi_abd_config": { + "num_a_tensors": 2, + "num_b_tensors": 2, + "num_d_tensors": 2, + "a_elementwise_op": "MultiDAdd", + "b_elementwise_op": "MultiDAdd", + "cde_elementwise_op": "MultiDAdd" + }, "k_block_per_cu": 1 } \ No newline at end of file diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/gemm_multi_abd_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/gemm_multi_abd_benchmark.py index 001e4a1d4f2b..82187ceb36db 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/gemm_multi_abd_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/gemm_multi_abd_benchmark.py @@ -148,7 +148,38 @@ def main(): # Export JSON if requested if args.json: - benchmark_utils.export_json(benchmark.results, args.json, best_kernels) + benchmark_utils.export_json( + benchmark.results, + args.json, + best_kernels, + launch_counts={ + "attempted": benchmark.launch_attempted, + "failed": benchmark.launch_failed, + }, + ) + + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py index 519644183776..7472a2e80a21 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_d/gemm_multi_d_benchmark.py @@ -140,7 +140,38 @@ def main(): # Export JSON if requested if args.json: - benchmark_utils.export_json(benchmark.results, args.json, best_kernels) + benchmark_utils.export_json( + benchmark.results, + args.json, + best_kernels, + launch_counts={ + "attempted": benchmark.launch_attempted, + "failed": benchmark.launch_failed, + }, + ) + + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py index 73ba1261a849..6fe1b07633ed 100755 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py @@ -140,7 +140,38 @@ def main(): # Export JSON if requested if args.json: - benchmark_utils.export_json(benchmark.results, args.json, best_kernels) + benchmark_utils.export_json( + benchmark.results, + args.json, + best_kernels, + launch_counts={ + "attempted": benchmark.launch_attempted, + "failed": benchmark.launch_failed, + }, + ) + + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py index a9c6e167fedd..f446366ea8e4 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm/grouped_gemm_benchmark.py @@ -16,6 +16,8 @@ def __init__(self, build_dir: str, verbose: bool = False): self.build_dir = Path(build_dir) self.verbose = verbose self.results = [] + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_grouped_gemm_* executables in the build directory""" @@ -308,6 +310,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = self.run_kernel(kernel_path, params) if result: @@ -341,6 +344,8 @@ def benchmark_problem_size( print( f" {kernel_info['config_id']}: {structured_result['tflops']:.2f} TFLOPS, {structured_result['bandwidth_gb_s']:.2f} GB/s, {structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -534,6 +539,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): ), "successful_runs": len(successful_results), "failed_runs": len(self.results) - len(successful_results), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "performance_summary": { "tflops_stats": { @@ -696,6 +704,29 @@ def main(): if args.json: benchmark.export_json(args.json, best_kernels) + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_rowcolquant/grouped_gemm_rowcolquant_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_rowcolquant/grouped_gemm_rowcolquant_benchmark.py index 6a4db06cd225..49c094e4b705 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_rowcolquant/grouped_gemm_rowcolquant_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_rowcolquant/grouped_gemm_rowcolquant_benchmark.py @@ -88,6 +88,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = benchmark_utils.run_kernel( self.build_dir, kernel_path, params, verbose=self.verbose ) @@ -122,6 +123,8 @@ def benchmark_problem_size( f" {kernel_info['config_id']}: {structured_result['tflops']:.2f} TFLOPS, " f"{structured_result['bandwidth_gb_s']:.2f} GB/s, {structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -301,6 +304,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): ), "successful_runs": len(successful_results), "failed_runs": len(self.results) - len(successful_results), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "performance_summary": { "tflops_stats": { @@ -474,9 +480,41 @@ def main(): benchmark_utils.export_csv(benchmark.results, args.csv) benchmark_utils.export_best_kernels(best_kernels, args.best) if args.json: - benchmark_utils.export_json(benchmark.results, args.json, best_kernels) + benchmark_utils.export_json( + benchmark.results, + args.json, + best_kernels, + launch_counts={ + "attempted": benchmark.launch_attempted, + "failed": benchmark.launch_failed, + }, + ) print("\nBenchmark complete!") + + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_tensorquant/grouped_gemm_tensorquant_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_tensorquant/grouped_gemm_tensorquant_benchmark.py index 2d3557454e0e..e638d0c11388 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_tensorquant/grouped_gemm_tensorquant_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/grouped_gemm_quant/grouped_gemm_tensorquant/grouped_gemm_tensorquant_benchmark.py @@ -86,6 +86,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = benchmark_utils.run_kernel( self.build_dir, kernel_path, params, verbose=self.verbose ) @@ -120,6 +121,8 @@ def benchmark_problem_size( f" {kernel_info['config_id']}: {structured_result['tflops']:.2f} TFLOPS, " f"{structured_result['bandwidth_gb_s']:.2f} GB/s, {structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -291,9 +294,41 @@ def main(): benchmark_utils.export_csv(benchmark.results, args.csv) benchmark_utils.export_best_kernels(best_kernels, args.best) if args.json: - benchmark_utils.export_json(benchmark.results, args.json, best_kernels) + benchmark_utils.export_json( + benchmark.results, + args.json, + best_kernels, + launch_counts={ + "attempted": benchmark.launch_attempted, + "failed": benchmark.launch_failed, + }, + ) print("\nBenchmark complete!") + + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm/mx_gemm/mx_gemm_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/mx_gemm/mx_gemm_benchmark.py index f69bc8328688..8d64ef8bae7d 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/mx_gemm/mx_gemm_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/mx_gemm/mx_gemm_benchmark.py @@ -181,7 +181,38 @@ def main(): benchmark_utils.export_best_kernels(best_kernels, args.best) if args.json: - benchmark_utils.export_json(benchmark.results, args.json, best_kernels) + benchmark_utils.export_json( + benchmark.results, + args.json, + best_kernels, + launch_counts={ + "attempted": benchmark.launch_attempted, + "failed": benchmark.launch_failed, + }, + ) + + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") return 0 diff --git a/projects/composablekernel/tile_engine/ops/gemm_streamk/CMakeLists.txt b/projects/composablekernel/tile_engine/ops/gemm_streamk/CMakeLists.txt index 02107de3fb8c..43348a46f360 100644 --- a/projects/composablekernel/tile_engine/ops/gemm_streamk/CMakeLists.txt +++ b/projects/composablekernel/tile_engine/ops/gemm_streamk/CMakeLists.txt @@ -231,9 +231,9 @@ message(VERBOSE "GEMM_STREAMK_DATATYPE: ${GEMM_STREAMK_DATATYPE}") message(VERBOSE "GEMM_STREAMK_LAYOUT: ${GEMM_STREAMK_LAYOUT}") message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") -# Filter GPU targets to only gfx90a, gfx942 +# Filter GPU targets to only gfx90a, gfx942, gfx950 set(GEMM_GPU_TARGETS_INDIVIDUAL "") -set(DESIRED_TARGETS "gfx90a;gfx942") # TODO: Add gfx950 when supported +set(DESIRED_TARGETS "gfx90a;gfx942;gfx950") foreach(target IN LISTS SUPPORTED_GPU_TARGETS) if(target IN_LIST DESIRED_TARGETS) @@ -244,7 +244,7 @@ endforeach() # Skip build if no matching targets found if(NOT GEMM_GPU_TARGETS_INDIVIDUAL) - message(WARNING "Skipping Tile Engine GEMM build: No supported GPU targets (gfx90a, gfx942) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") + message(WARNING "Skipping Tile Engine GEMM build: No supported GPU targets (gfx90a, gfx942, gfx950) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") else() message(VERBOSE "Building individual GEMM targets for GPU targets: ${GEMM_GPU_TARGETS_INDIVIDUAL}") diff --git a/projects/composablekernel/tile_engine/ops/gemm_streamk/gemm_streamk_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm_streamk/gemm_streamk_benchmark.py index ad8d9ff35c5a..7c83d4f1fae5 100644 --- a/projects/composablekernel/tile_engine/ops/gemm_streamk/gemm_streamk_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm_streamk/gemm_streamk_benchmark.py @@ -16,6 +16,8 @@ def __init__(self, build_dir: str, verbose: bool = False): self.build_dir = Path(build_dir) self.verbose = verbose self.results = [] + self.launch_attempted = 0 + self.launch_failed = 0 def discover_kernels(self) -> List[Path]: """Find all benchmark_gemm_streamk_* executables in the build directory""" @@ -307,6 +309,7 @@ def benchmark_problem_size( for kernel_path in kernels: kernel_info = self.extract_kernel_info(kernel_path) + self.launch_attempted += 1 result = self.run_kernel(kernel_path, params) if result: @@ -341,6 +344,8 @@ def benchmark_problem_size( print( f" {kernel_info['config_id']}: {structured_result['tflops']:.2f} TFLOPS, {structured_result['bandwidth_gb_s']:.2f} GB/s, {structured_result['time_ms']:.2f}ms" ) + else: + self.launch_failed += 1 return results @@ -528,6 +533,9 @@ def export_json(self, filename: str, best_kernels: Dict = None): ), "successful_runs": len(successful_results), "failed_runs": len(self.results) - len(successful_results), + "launches_attempted": self.launch_attempted, + "launches_succeeded": self.launch_attempted - self.launch_failed, + "launches_failed": self.launch_failed, }, "performance_summary": { "tflops_stats": { @@ -669,6 +677,29 @@ def main(): if args.json: benchmark.export_json(args.json, best_kernels) + # Exiting 0 after every launch failed would leave the CI lane green with no signal + attempted = benchmark.launch_attempted + failed = benchmark.launch_failed + succeeded = attempted - failed + print(f"Launches: {attempted} attempted, {succeeded} succeeded, {failed} failed") + + if attempted == 0: + print("No kernel launches were attempted - no kernels discovered") + return 1 + if not benchmark.results: + print("No benchmark results were collected") + return 1 + if failed > 0: + # A warning, not a failure: this sweep runs every generated config, and + # individual configs legitimately fail to launch (unsupported tile shape + # for the arch, workspace too large for the problem size). Making any + # single failure red would trade a permanently-green lane for a + # permanently-red one, which is no more informative. Total failure is + # already caught by the empty-result check above, and the per-run counts + # are in the JSON (launches_attempted/succeeded/failed) for anyone + # tracking the trend. + print(f"WARNING: {failed} of {attempted} kernel launch(es) failed") + return 0