[ROCm] Allow DSA indexer native decode above next_n=2 - #13
Open
amd-sriram wants to merge 1 commit into
Open
amd-sriram wants to merge 1 commit into
amd-sriram wants to merge 1 commit into
Conversation
Gate it on VLLM_ROCM_INDEXER_NATIVE_DECODE_MAX_NEXT_N; the default of 2 keeps the current flattening behaviour. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@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.
Problem
Under MTP the DSA indexer only passes
next_nQ rows per request to the pagedMQA-logits kernel when
_supports_native_decode(next_n)is true. On ROCm thatreturns
next_n in (1, 2), sonum_speculative_tokens >= 2falls into_use_flattening, and_prepare_decode_tensorsexpands every draft token intoits own batch row. The kernel is then called at
next_n = 1withB = num_decode_tokens, and each row re-reads the whole KV sweep for its ownsequence instead of reusing a fetched page across the request's draft rows.
The aiter paged kernels do not have that limit. Gluon's
deepgemm_fp8_paged_mqa_logitsreadsnext_noffq_fp8.size(), andtop_k_per_row_decodeindexesrowIdx / next_nwith no upper bound.Change
_supports_native_decodetakes a ROCm branch that compares against a newVLLM_ROCM_INDEXER_NATIVE_DECODE_MAX_NEXT_N. The default of2reproducesnext_n in (1, 2)exactly, sincenext_n >= 1always, so nothing changesunless the cap is raised.
Not yet verified
other
next_nassumptions in the ROCm decode path.get_cudagraph_supportfromALWAYStoUNIFORM_BATCH, because that tier is chosen by_use_flattening. That needsto hold under capture before any perf number from this is meaningful.
The env var exists so both arms fit in one image while we measure. If the
native path holds up, the upstream version should drop the flag and return
next_n <= 8directly.