From a89015c6df8eeb37a843b717c97a5be1355de83d Mon Sep 17 00:00:00 2001 From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:24:42 +0800 Subject: [PATCH] [Perf] Make merge attention context count a runtime argument (#48739) Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> --- vllm/v1/attention/ops/triton_merge_attn_states.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/ops/triton_merge_attn_states.py b/vllm/v1/attention/ops/triton_merge_attn_states.py index ca06c2970b..e8f90efce9 100644 --- a/vllm/v1/attention/ops/triton_merge_attn_states.py +++ b/vllm/v1/attention/ops/triton_merge_attn_states.py @@ -158,10 +158,10 @@ def merge_attn_states( prefix_head_stride, output_head_stride, output_scale, + prefill_tokens_with_context, head_size, padded_head_size, output_lse is not None, - prefill_tokens_with_context, output_scale is not None, ) @@ -177,10 +177,10 @@ def merge_attn_states_kernel( prefix_head_stride, output_head_stride, output_scale, # scale tensor or None + prefill_tokens_with_context, HEAD_SIZE: tl.constexpr, PADDED_HEAD_SIZE: tl.constexpr, OUTPUT_LSE: tl.constexpr, - prefill_tokens_with_context: tl.constexpr, USE_FP8: tl.constexpr, FP8_MIN: tl.constexpr = float8_info.min, FP8_MAX: tl.constexpr = float8_info.max,