[Perf] Make merge attention context count a runtime argument (#48739)
Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
liminfei-amd committed
Jul 27, 2026 at 21:24 UTC
a89015c6df8eeb37a843b717c97a5be1355de83d
1 file changed
+2
-2
vllm/v1/attention/ops/triton_merge_attn_states.py
+2
-2
@@ -158,10 +158,10 @@ def merge_attn_states(
158
prefix_head_stride,
159
output_head_stride,
160
output_scale,
161
+ prefill_tokens_with_context,
162
head_size,
163
padded_head_size,
164
output_lse is not None,
164
- prefill_tokens_with_context,
165
output_scale is not None,
166
)
167
@@ -177,10 +177,10 @@ def merge_attn_states_kernel(
177
prefix_head_stride,
178
output_head_stride,
179
output_scale, # scale tensor or None
180
+ prefill_tokens_with_context,
181
HEAD_SIZE: tl.constexpr,
182
PADDED_HEAD_SIZE: tl.constexpr,
183
OUTPUT_LSE: tl.constexpr,
183
- prefill_tokens_with_context: tl.constexpr,
184
USE_FP8: tl.constexpr,
185
FP8_MIN: tl.constexpr = float8_info.min,
186
FP8_MAX: tl.constexpr = float8_info.max,