[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:
co-authored by
Raiden-Makoto
parent
72cac88022
commit
9215da2515
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user