diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index 1c9ee2bf5bfb..d0eae8f72814 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -529,6 +529,7 @@ def _forward_sdpa( num_kv_heads=self.num_kv_heads, dtype=query.dtype, is_causal=False, # ViT attention is non-causal + v_token_stride=value.stride(1), ): output = vit_torch_sdpa_wrapper( q=query, @@ -582,6 +583,7 @@ def _forward_fa( num_kv_heads=self.num_kv_heads, dtype=query.dtype, is_causal=False, # ViT attention is non-causal + v_token_stride=value.stride(1), ): output = vit_flash_attn_wrapper( q=query, @@ -633,6 +635,7 @@ def _forward_triton( num_kv_heads=self.num_kv_heads, dtype=query.dtype, is_causal=False, # ViT attention is non-causal + v_token_stride=value.stride(1), ): output = vit_triton_attn_wrapper( q=query, diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index c1274377b6e1..391b788a25b9 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -71,6 +71,7 @@ def flash_attn_maxseqlen_wrapper( num_kv_heads=num_kv_heads, dtype=q.dtype, is_causal=False, + v_token_stride=v.stride(0), ): output = flash_attn_varlen_func( q, @@ -170,6 +171,7 @@ def triton_attn_wrapper( num_kv_heads=num_kv_heads, dtype=q.dtype, is_causal=False, + v_token_stride=v.stride(0), ): context_attention_fwd( q, diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index b089e551cbea..7ee12f56c456 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -707,6 +707,7 @@ def create_attention_profiler_scope( num_kv_heads: int, dtype: torch.dtype, is_causal: bool, + v_token_stride: int | None = None, ) -> AbstractContextManager: """Create a profiler scope for attention operations with metadata. @@ -720,6 +721,13 @@ def create_attention_profiler_scope( num_kv_heads: Number of KV heads (for GQA/MQA) dtype: Data type of the tensors is_causal: Whether causal masking is applied + v_token_stride: Optional per-token stride of the V operand. Only + surfaced in the label 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). This carries a measurable latency impact on + some GPUs (e.g. Strix Point cold cache); natural strides are omitted + to keep the scope label short. Returns: Context manager for the profiler scope @@ -733,6 +741,13 @@ def create_attention_profiler_scope( f"H={num_heads} D={head_size} KV_H={num_kv_heads} " f"DT={dtype_str} BE={backend_name}" ) + # Only surface the V operand token stride when it is *not* the natural + # contiguous value (num_kv_heads * head_size). The real ViT layout feeds a + # non-contiguous V sliced from the fused QKV projection (token stride + # mult*H*D), which is what we want visible in the profile; natural strides + # are omitted to keep the label short. + if v_token_stride is not None and v_token_stride != num_kv_heads * head_size: + scope_name += f" Vstride={v_token_stride}" return record_function_or_nullcontext(scope_name)