[Perf] Make merge attention context count a runtime argument (#48739)

Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
This commit is contained in:
liminfei-amd
2026-07-27 13:24:42 +00:00
committed by GitHub
parent 96fa3f42c9
commit a89015c6df
@@ -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,