rocm: bound the INT8 lm_head GEMV dispatch by the kernel's max N - #1240
Open
zohaibm-amd wants to merge 1 commit into
Open
rocm: bound the INT8 lm_head GEMV dispatch by the kernel's max N#1240zohaibm-amd wants to merge 1 commit into
zohaibm-amd wants to merge 1 commit into
Conversation
The dispatch in hip_w8a16.py routed the lm_head GEMV to wvSplitK_int8 whenever the activation tile fit in LDS, admitting any batch N up to floor(32768 / K). The kernel's switch in csrc/rocm/skinny_gemms_w8a8.cu only implements N of 1 through 5 and throws on the default branch, so any K with floor(32768 / K) of 6 or more handed it a batch it cannot service: RuntimeError: Unsupported N value: 262144,5376,6 The throw does not unwind cleanly through _prefer_hipblaslt_for_logits, so it killed the whole EngineCore process rather than failing one request. Observed serving a 31B AWQ model with --dynamic-lm-head-quantization int8 as soon as six decode requests were scheduled together; identical with int8:g32. Smaller hidden sizes are more exposed, since the window is every N from 6 up to floor(32768 / K): 6 only at K=5376, but 6-16 at K=2048. Require N to be within what the kernel implements as well as fitting in LDS. Larger batches fall back to F.linear on the retained unquantized weights, which is already the designed path above the LDS bound, so the fallback is correctness-preserving and marginally more accurate than the INT8 path it replaces. The dispatch bound was introduced by 7847fce, which replaced a prior AssertionError with the LDS-capacity check without cross-checking it against the kernel's N switch. Changes: - N_BATCH in the existing test gains 5 and 6: the last N the kernel implements and the first that must fall back. It previously stopped at 4, which is why the gap went unnoticed. All three MK_SHAPES reach the kernel at N=6, so every one of them reproduces the throw without this fix. Signed-off-by: Zohaib Moti <zohaib.moti@amd.com> 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
eble-amd
approved these changes
Aug 31, 2026
eble-amd
left a comment
eble-amd
left a comment
There was a problem hiding this comment.
OK, but it would be nice to understand why support currently ends at 5. What if the ideal solution to this problem involves not just respecting the kernel's limitations in this layer, but also extending the kernel's capabilities?
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
2 participants
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.
Summary
--dynamic-lm-head-quantization int8(andint8:gN) kills the EngineCore process once the scheduler batches enough concurrent decode requests.hip_w8a16.pyroutes the lm_head GEMV towvSplitK_int8whenever the activation tile fits in LDS, admitting any batch N up tofloor(32768 / K). The kernel's switch incsrc/rocm/skinny_gemms_w8a8.cuimplements N of 1 through 5 andthrows on the default branch.Fix. Require N to be within what the kernel implements as well as fitting in LDS.
Test plan
N_BATCHintest_dynamic_int8_lm_head.pygains 5 and 6 — the last N the kernel implements and the first that must fall back. It previously stopped at 4, which is why the gap went unnoticed.Duplicate check
No open PR touches
hip_w8a16.pyorskinny_gemms_w8a8.cu. Neither file exists upstream invllm-project/vllm—--dynamic-lm-head-quantizationis a fork-specific feature (#888) — so no upstream PR can be duplicating it, and there is nothing upstream to cherry-pick.AI assistance
AI assistance (Claude Code) was used to locate the dispatch/kernel mismatch and prepare this change. The crash was reproduced and its traceback observed directly; the N bound was read from the kernel's switch statement rather than assumed.