Unverified Commit 5264015d authored by Jinghui Zhang's avatar Jinghui Zhang Committed by GitHub
Browse files

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


Signed-off-by: default avatarJinghui Zhang <jinghuizhang0804@gmail.com>
parent 98ac0cb3
...@@ -279,7 +279,7 @@ def rocm_aiter_grouped_topk( ...@@ -279,7 +279,7 @@ def rocm_aiter_grouped_topk(
if e_score_correction_bias is not None: if e_score_correction_bias is not None:
torch.ops.vllm.rocm_aiter_biased_grouped_topk( torch.ops.vllm.rocm_aiter_biased_grouped_topk(
gating_output, gating_output,
e_score_correction_bias, e_score_correction_bias.to(gating_output.dtype),
topk_weights, topk_weights,
topk_ids, topk_ids,
num_expert_group, num_expert_group,
...@@ -409,15 +409,15 @@ def shuffle_weights( ...@@ -409,15 +409,15 @@ def shuffle_weights(
*tensors: torch.Tensor, layout: tuple[int, int] = (16, 16) *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
) -> tuple[torch.Tensor, ...]: ) -> tuple[torch.Tensor, ...]:
""" """
Applies shuffle_weight function from AITER to each Applies shuffle_weight function from AITER to each
input tensor and returns them. input tensor and returns them.
Rearranges (shuffles) the input tensor/s Rearranges (shuffles) the input tensor/s
into a specified block layout for optimized computation. into a specified block layout for optimized computation.
Args: Args:
*tensors: Variable number of torch.Tensor objects. *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. block sizes used to divide the tensors during shuffling.
Default is (16, 16). Default is (16, 16).
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment