[AMD] Fix no-op dtype cast in _topk_ids_logical_to_physical_dynamic on HIP (#28074)

Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-06-20 11:07:22 -07:00
parent ff1fc1fbdf
commit d6d06cdc17
2 changed files with 157 additions and 1 deletions
@@ -111,6 +111,7 @@ def _topk_ids_logical_to_physical_dynamic(
    topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo]
 ) -> torch.Tensor:
    topk_ids_original_shape = topk_ids.shape
+    original_dtype = topk_ids.dtype
    device = topk_ids.device
    topk_ids = topk_ids.flatten()

@@ -120,7 +121,7 @@ def _topk_ids_logical_to_physical_dynamic(
    )
    topk_ids = info.partial_logical_to_all_physical_map[topk_ids, chosen_dispatch_index]
    if _is_hip:
-        topk_ids = topk_ids.to(topk_ids.dtype)
+        topk_ids = topk_ids.to(original_dtype)

    topk_ids = topk_ids.view(topk_ids_original_shape)
    return topk_ids