MIGraphX: derive static seq-pad length from the attention mask - #66
Open
aditya-dl wants to merge 1 commit into
Open
MIGraphX: derive static seq-pad length from the attention mask#66aditya-dl wants to merge 1 commit into
aditya-dl wants to merge 1 commit into
Conversation
Static seq-padding took its target length from the static_pad_seq_len provider option. For OGA callers that value is config.search.max_length, read while the session is being created. A caller that sets max_length the normal way, via GeneratorParams::set_search_options, does so after the model is built, so the value never reaches the EP and we fall back to whatever the model declares as its context length. On the certification models that is 131072, so the prefill was padded and compiled for the full context instead of the few thousand tokens actually in use. Take the length from the attention mask instead. The caller sizes that mask from the length it is really generating with, and hands it to us with the data, so it cannot be stale. It is also the same number the caller used for its KV cache, so deriving from it keeps the padded token axis and the caller's buffers consistent by construction rather than by both sides happening to agree. The provider option and ORT_MIGRAPHX_STATIC_PAD_SEQ_LEN still work. An explicit env value overrides the mask, and models with no attention_mask input keep the configured length, so behaviour is unchanged wherever the mask is unavailable. Verified with MIGRAPHX_TRACE_COMPILE, comparing the configured length against the compiled input_ids extent: harness config mask compiled python, no overlay 131072 19 19 modelbench, no overlay 131072 145 145 EPCert, unmodified 131072 4096 4096 EPCert previously needed ORT_MIGRAPHX_STATIC_PAD_SEQ_LEN=4096 to avoid running out of memory on the first prefill; it now completes with nothing set. A run with the length already correct in the config produces byte-identical tokens.
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.
Port of #65 to
main. Identical change; the release branch could not be merged directly because of unrelated commits in its history.Problem
Static seq-padding took its target from the
static_pad_seq_lenprovider option. For OGA that isconfig.search.max_length, read while the session is created. A caller that setsmax_lengththe normal way (GeneratorParams::set_search_options) does so after the model is built, so it never reaches the EP and we pad to the model's declared context length.EPCert hits this: 131072 instead of 4096, out of memory on the first prefill. It needed
ORT_MIGRAPHX_STATIC_PAD_SEQ_LEN=4096to run.Fix
Take the length from
attention_maskinstead. It arrives with the data so it cannot be stale, and it is the same length the caller sized its KV cache to. No OGA or caller change, which matters because EPCert is not ours to patch.ORT_MIGRAPHX_STATIC_PAD_SEQ_LENstill overrides. Models with no mask keep the configured length. Non-LLM paths never enablestatic_pad_seq, so nothing changes there.