Skip to content

[ROCm] Allow DSA indexer native decode above next_n=2 - #13

Open
amd-sriram wants to merge 1 commit into
mainfrom
rocm/indexer-native-decode-nextn
Open

amd-sriram wants to merge 1 commit into
mainfrom
rocm/indexer-native-decode-nextn

Conversation

@amd-sriram

Copy link
Copy Markdown
Owner

Problem

Under MTP the DSA indexer only passes next_n Q rows per request to the paged
MQA-logits kernel when _supports_native_decode(next_n) is true. On ROCm that
returns next_n in (1, 2), so num_speculative_tokens >= 2 falls into
_use_flattening, and _prepare_decode_tensors expands every draft token into
its own batch row. The kernel is then called at next_n = 1 with
B = num_decode_tokens, and each row re-reads the whole KV sweep for its own
sequence 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_logits reads next_n off q_fp8.size(), and
top_k_per_row_decode indexes rowIdx / next_n with no upper bound.

Change

_supports_native_decode takes a ROCm branch that compares against a new
VLLM_ROCM_INDEXER_NATIVE_DECODE_MAX_NEXT_N. The default of 2 reproduces
next_n in (1, 2) exactly, since next_n >= 1 always, so nothing changes
unless the cap is raised.

Not yet verified

  • No GPU run. This has only been reasoned through the call path and checked for
    other next_n assumptions in the ROCm decode path.
  • Raising the cap flips get_cudagraph_support from ALWAYS to
    UNIFORM_BATCH, because that tier is chosen by _use_flattening. That needs
    to hold under capture before any perf number from this is meaningful.
  • Accuracy under MTP with the cap raised is unmeasured.

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 <= 8 directly.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant