[BugFix][AMD][Deepseek] fix a dtype mismatch error for deepseek running on AMD (#23864)

Signed-off-by: Jinghui Zhang <jinghuizhang0804@gmail.com>
This commit is contained in:
Jinghui Zhang 2025-08-28 22:54:12 -07:00 committed by GitHub
parent 98ac0cb32d
commit 5264015d74
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -279,7 +279,7 @@ def rocm_aiter_grouped_topk(
if e_score_correction_bias is not None:
torch.ops.vllm.rocm_aiter_biased_grouped_topk(
gating_output,
e_score_correction_bias,
e_score_correction_bias.to(gating_output.dtype),
topk_weights,
topk_ids,
num_expert_group,
@ -409,15 +409,15 @@ def shuffle_weights(
*tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
) -> tuple[torch.Tensor, ...]:
"""
Applies shuffle_weight function from AITER to each
Applies shuffle_weight function from AITER to each
input tensor and returns them.
Rearranges (shuffles) the input tensor/s
into a specified block layout for optimized computation.
Args:
*tensors: Variable number of torch.Tensor objects.
layout: A pair of integers specifying the
layout: A pair of integers specifying the
block sizes used to divide the tensors during shuffling.
Default is (16, 16).