perf(dsv4): adaptive BLOCK_K for csa_translate_pack#1464
Merged
Conversation
The previous BLOCK_K = min(64, next_pow2(index_topk)) fixes the per-program tile at <=64 slots, so a large decode batch launches an enormous grid of tiny workgroups (e.g. T=256, index_topk=2048 -> 8192 workgroups). On gfx1250 the per-workgroup dispatch overhead grows steeply with workgroup count, so these tiny tiles are dominated by launch/scheduling cost rather than useful work. Grow the per-program tile while the grid still has enough workgroups to fill the GPU: keep the workgroup count (T * cdiv(index_topk, BLOCK_K)) at or above _CSA_TARGET_WG (512). Large-T decode then gets fat tiles that amortize launch overhead, while small-T stays fine-grained for occupancy. Tile floor stays at 64 slots and is capped at 1024. Measured up to 3.2x on gfx1250 (11.3 -> 3.5 us) and up to 1.5x on gfx950 for large decode batches; small cases are unchanged. Kernel body and numerics are untouched -- this only retunes the launch tiling.
…e writes Mark the config-fixed stride/dim/cache_size args of _swa_write_kernel and _update_compressor_states_kernel as tl.constexpr so the backend folds the ring-buffer address arithmetic to constants and specializes pos%cache_size (pow-2 -> bitmask). ~1.04-1.11x on both kernels, bit-exact. Also fix update_compressor_states_reference to skip sentinel (-1) plan rows like the kernel does (Python -1%N wraps to N-1, writing a garbage ring row).
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
1 participant
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.
What
Adaptive
BLOCK_Ktiling for the DSV4csa_translate_packTriton kernel.The old launch used
BLOCK_K = min(64, next_pow2(index_topk)), fixing theper-program tile at ≤64 slots. A large decode batch therefore launches an
enormous grid of tiny workgroups (e.g.
T=256, index_topk=2048→ 8192workgroups). On gfx1250 the per-workgroup dispatch overhead grows steeply with
workgroup count, so those tiny tiles are dominated by launch/scheduling cost.
This grows the per-program tile while the grid still has enough workgroups to
fill the GPU — keeping
T * cdiv(index_topk, BLOCK_K)at or above_CSA_TARGET_WG(512). Large-T decode gets fat tiles that amortize launchoverhead; small-T stays fine-grained for occupancy. Tile floor stays 64, capped
at 1024.
Correctness
Kernel body and numerics are untouched — this only retunes the launch tiling.
Verified 10/10 correctness against the reference on gfx950.
Performance