Repository navigation
[Enh] Add CompiledFunction.max_blocks_per_cu() - #1259
Merged
Merged
Conversation
Return the max resident blocks per CU of a compiled kernel via hipModuleOccupancyMaxActiveBlocksPerMultiprocessor, so callers can size grids to full waves. A runtime wrapper next to mgpuModuleLoad loads the code object into a temporary module, so nothing is launched and it works before the first launch. Block size and dynamic smem are read from the compiled module's gpu.launch_func.
akii96
added a commit
to akii96/aiter
that referenced
this pull request
Oct 10, 2026
Replace the measured occupancy table, the IR metadata parser and the HIP generator script with CompiledFunction.max_blocks_per_cu() (ROCm/FlyDSL#1259), compiled with meta tensors once per launcher. Falls back to the cu_oversub rule on FlyDSL < 0.4.0 and on a first call inside graph capture. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96
added a commit
to akii96/aiter
that referenced
this pull request
Oct 10, 2026
Replace the measured occupancy table, the IR metadata parser and the HIP generator script with CompiledFunction.max_blocks_per_cu() (ROCm/FlyDSL#1259), compiled with meta tensors once per launcher. Falls back to the cu_oversub rule on a first call inside graph capture. Requires FlyDSL >= 0.4.0. Do not merge before aiter's flydsl pin is bumped to it. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96
added a commit
to akii96/aiter
that referenced
this pull request
Oct 15, 2026
Replace the measured occupancy table, the IR metadata parser and the HIP generator script with CompiledFunction.max_blocks_per_cu() (ROCm/FlyDSL#1259), compiled with meta tensors once per launcher. Falls back to the cu_oversub rule on a first call inside graph capture. Requires FlyDSL >= 0.4.0. Do not merge before aiter's flydsl pin is bumped to it. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
coderfeli
pushed a commit
to ROCm/aiter
that referenced
this pull request
Oct 25, 2026
#5921) * [FlyDSL] fp8_mqa_logits: derive the split target from kernel occupancy _auto_num_splits targets a fixed multiple of the CU count, but the number of blocks a CU actually holds is a property of the compiled kernel, not of the device. On gfx942 the mfma_r2_w4 variant is 8 blocks/CU at num_heads=16 and 2 at num_heads=128, so one cu_oversub constant is correct for one kernel instance and wrong for the rest: at 16x CU it over-splits the low-occupancy kernels (19 splits where 7 is optimal) and under-splits the high-occupancy ones (3 where 8 is optimal). Replace the constant with the effective CU count, cu_count * occupancy, and pick the split that minimizes a wave-quantized cost model: cost(s) = ceil(grid_x * s / effective_cus) * (ceil(W / s) + overhead) Occupancy comes from the kernel itself. FlyDSL JITs, so there is no hipFuncAttributes handle before the artifact exists, but the artifact carries the AMDHSA metadata that determines occupancy; kernel_occupancy reads it and applies the hardware allocation rules. That result equals hipModuleOccupancyMaxActiveBlocksPerMultiprocessor on all 216 buildable gfx942 configurations and all 308 on gfx950. A measured table covers the first launch of a shape, which necessarily asks before any artifact exists, and a conservative default covers anything else; both fall on the under-estimate side, which under-splits rather than oversubscribes. The two arches allocate AGPRs differently and the model is parameterized for it: CDNA3 shares one register file between VGPRs and AGPRs, CDNA4 splits them. Summing them on gfx950 mispredicts the bkv256 variants (vgpr=452/agpr=200), which the sum puts over the 512-VGPR file at occupancy 1 where the hardware reports 2. Every gfx942 kernel in this op reports agpr=0, so the distinction changes nothing there. gfx950 keeps the fixed-oversubscription path. Its occupancy is verified, but its own measurements do not justify switching: its auto-selected variants span only two occupancy values against six on gfx942, its optima sit 8-24x past one effective wave (load imbalance across causal windows, which this model does not represent), and 32 of 35 shapes have four or more split choices within 3% of optimal against a ~2% noise floor. Rationale recorded in _SPLIT_POLICIES. Perf, gfx942/MI325X, over the shapes vLLM issues for chunked long-context indexer prefill (seq_len pinned near 1024 by VLLM_SPARSE_INDEXER_MAX_LOGITS_MB, seq_len_kv sweeping 16k-128k, causal), on top of #4963: DeepSeek-V4.1-Flash (index_n_heads=64): 1.25x, and 1.14x over a flat 16x GLM-5.3-Flash (index_n_heads=32): 1.19x, and 1.07x over a flat 16x #4963 is roughly neutral on DSv4.1 across that family, so nearly all of its 1.24x combined gain comes from this change; on GLM the two compose to 1.74x. On aiter's own --model-shapes grid (block-diagonal batch_causal) the same change is 1.08x over main with no shape regressing beyond its measured spread. block_overhead_tiles is fitted over 86 shapes spanning three regimes -- vLLM long-context prefill, block-diagonal batch_causal, and wide-window causal -- scored with the regimes weighted equally rather than by sample count, since counts reflect how much was measured rather than what matters. Fitted on wide causal alone it lands at 3, which then costs up to 1.21x on the narrow windows production actually issues, because per-block fixed cost is a far larger share of the work at 32 tiles than at 1024. The joint optimum is 8 under equal, vllm-weighted and production-only weightings alike, and the worst case falls from 1.21x to 1.07x. Output is bit-identical to both prior policies at every in-window position on every shape; splitting a row's window across grid.y is pure parallelism. aiter's own sweep passes 1590/1590 cases. Adds op_tests/flydsl/gen_fp8_mqa_logits_occupancy.py, which regenerates the table and re-checks the model against HIP, and 25 unit tests for the selection logic, the occupancy fallback chain, and the per-arch register-file rules. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * [FlyDSL] fp8_mqa_logits: address review on occupancy-aware splits Trim comments, cache the occupancy on the launcher and the split choice, read the public CompiledArtifact.ir, load HIP through load_hip_runtime, and add a test that reads the compiled kernel's metadata against the measured occupancy table. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * [FlyDSL] fp8_mqa_logits: take occupancy from FlyDSL max_blocks_per_cu() Replace the measured occupancy table, the IR metadata parser and the HIP generator script with CompiledFunction.max_blocks_per_cu() (ROCm/FlyDSL#1259), compiled with meta tensors once per launcher. Falls back to the cu_oversub rule on a first call inside graph capture. Requires FlyDSL >= 0.4.0. Do not merge before aiter's flydsl pin is bumped to it. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * [FlyDSL] fp8_mqa_logits: tidy split-policy docs and occupancy test Name the split search cap (_MAX_SPLITS), drop stale numbers from the _auto_num_splits docstring, and have the occupancy test call _max_blocks_per_cu directly instead of monkeypatching _auto_num_splits. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * [FlyDSL] fp8_mqa_logits: drop the split-heuristic unit test from op_tests --------- Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Some callers size their grid to fill the GPU, for example aiter
fp8_mqa_logitswhen it picks a KV split count (ROCm/aiter#5921). For that they need to know how many blocks of a kernel fit on one CU. This depends on each compiled kernel's register and LDS usage. Today a caller can only guess it, keep a hand-measured table, or parse FlyDSL's IR and load the binary into HIP itself.Change
Adds
CompiledFunction.max_blocks_per_cu():The value comes from HIP's
hipModuleOccupancyMaxActiveBlocksPerMultiprocessor. A small wrapper next tomgpuModuleLoadloads the kernel into a temporary module, asks HIP, and unloads it, using the same runtime library as launches. Nothing is launched, so it works before the first launch.The kernel binary, block size and shared memory size are read from the compiled module, which is what the disk cache keeps, so cached kernels work too. The launcher must launch exactly one kernel, with constant block and shared memory sizes.
Test
tests/unit/test_max_blocks_per_cu.pycompiles an empty kernel with 0 and with 48 KiB of shared memory and checks that the second fits fewer blocks. Tested on MI300X. It also gives the expected value for each of 72 aiterfp8_mqa_logitskernel configurations.Also verified on MI350 (gfx950): matches HIP on 40 block-size and shared-memory configurations.