[AMD] [GLM5] skip redundant -inf pre-fill of HIP indexer MQA-logits (#28757)

Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
This commit is contained in:
Raiden Makoto
2026-06-24 23:21:39 -07:00
committed by GitHub
co-authored by Raiden-Makoto
parent 72cac88022
commit 9215da2515
2 changed files with 132 additions and 1 deletions
@@ -870,8 +870,18 @@ class Indexer(MultiPlatformOp):
from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
kv, scale = kv_fp8
# Match the CUDA deep_gemm path (clean_logits=False): the topk
# transform masks invalid positions via ks/ke/lengths, so the
# -inf pre-fill of the [tokens x seq_len_kv] logits buffer is
# redundant and grows quadratically with context length.
logits = fp8_mqa_logits(
q_fp8[:q_offset], kv, scale, weights[:q_offset], ks, ke
q_fp8[:q_offset],
kv,
scale,
weights[:q_offset],
ks,
ke,
clean_logits=False,
)
else:
logits = deep_gemm.fp8_mqa_logits(
@@ -915,6 +925,7 @@ class Indexer(MultiPlatformOp):
from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
kv, scale = kv_fp8
# clean_logits=False: topk transform handles masking (see above)
logits_chunk = fp8_mqa_logits(
q_fp8[start:end],
kv,
@@ -922,6 +933,7 @@ class Indexer(MultiPlatformOp):
weights[start:end],
ks[start:end],
ke[start:end],
clean_logits=False,
)
else:
logits_chunk = deep_gemm.fp8_mqa_logits(