[profiler] Emit non-contiguous V operand stride in ViT attention profiler scope - #1071
Open
parthash0804 wants to merge 1 commit into
Open
[profiler] Emit non-contiguous V operand stride in ViT attention profiler scope#1071parthash0804 wants to merge 1 commit into
parthash0804 wants to merge 1 commit into
Conversation
…pe text
Add an optional v_token_stride argument to create_attention_profiler_scope
and wire it at all ViT encoder attention sites, for every selectable
mm-encoder backend:
- TRITON_ATTN (MMEncoderAttention._forward_triton, triton_attn_wrapper)
- FLASH_ATTN / ROCM_AITER_FA (MMEncoderAttention._forward_fa,
flash_attn_maxseqlen_wrapper)
- TORCH_SDPA (MMEncoderAttention._forward_sdpa)
The stride is appended to the profiler scope label as "Vstride=<n>" only
when it is *not* the natural contiguous value (num_kv_heads * head_size) --
i.e. the real ViT layout where V is a non-contiguous slice of the fused QKV
projection (token stride mult*H*D, e.g. 3*H*D). Natural strides are omitted
to keep the label short.
This surfaces the V operand memory layout directly in profiles regardless of
which ViT backend is selected, which carries a measurable latency impact on
some GPUs (e.g. Strix Point cold cache). Decoder / KV-cache attention paths
and all other callers are unchanged since the new arg defaults to None.
Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Parth Ashwin Jain <parthash@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
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.
Summary
The ViT encoder attention is fed a non-contiguous V — a slice of the fused QKV projection with token stride
mult*H*D(e.g.3*H*D), vsH*Dfor a contiguous V. This operand layout has a measurable latency impact on some GPUs (notably Strix Point /gfx1150when the cache is cold), but it was previously invisible in profiles and could only be confirmed by monkeypatchingcontext_attention_fwdat runtime. This PR surfaces the V operand's token stride directly in the attention profiler scope label, for every selectable ViT-encoder backend, so the layout can be read straight from a profile (and from thevllm-bench.pyattentiontable) without any instrumentation.
Changes
create_attention_profiler_scope(vllm/v1/utils.py) gains an optionalv_token_strideargument.Vstride=<n>only when the stride is not the natural contiguous value (num_kv_heads * head_size). Natural/contiguous strides are omitted to keep the label short, so the tag only appears for the fused-QKV ViT layout that actually matters.--mm-encoder-attn-backend:MMEncoderAttention._forward_triton,triton_attn_wrapperMMEncoderAttention._forward_fa,flash_attn_maxseqlen_wrapperMMEncoderAttention._forward_sdpatriton_attn,rocm_attn,rocm_aiter_unified_attn,rocm_aiter_fa) and all other callers are unchanged — the new arg defaults toNone, and their V (pagedvalue_cache) has the natural stride anyway.Result