Skip to content

[Bugfix] Fix hash topk dtype mismatch #43425

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wangyicong52 wants to merge 1 commit into from
Closed

[Bugfix] Fix hash topk dtype mismatch #43425

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions vllm/_custom_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2444,6 +2444,13 @@ def topk_hash_softplus_sqrt(
input_tokens: torch.Tensor | None = None,
hash_indices_table: torch.Tensor | None = None,
) -> None:
if hash_indices_table is not None:
assert input_tokens is not None
if input_tokens.dtype != topk_indices.dtype:
input_tokens = input_tokens.to(dtype=topk_indices.dtype)
if hash_indices_table.dtype != topk_indices.dtype:
hash_indices_table = hash_indices_table.to(dtype=topk_indices.dtype)

@jeejeelee jeejeelee May 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

QQ: Will these changes launch the type conversion kernel?

@wangyicong52 wangyicong52 May 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes. The current version is a quick and minimal correctness fix, and it may launch CUDA type conversion kernels when the dtypes do not match. If we want to avoid that entirely, the more complete solution would be to change the CUDA op / kernel contract so that topk_indices and the hash routing metadata (input_ids / tid2eid) do not have to use the same dtype. That should avoid the conversion kernels, but it is a larger kernel/dispatch change.

Without changing the CUDA kernel, a better way would be to move the dtype alignment upstream, so the static hash_indices_table is aligned once when the MoE backend determines the required topk_indices_dtype, instead of casting it before every custom op call.

Which way do you think is better? I can work on either direction next to fix this bug. @jeejeelee

@jeejeelee jeejeelee Jun 11, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for missing this PR. It looks like #41183 has fixed this bug

@wangyicong52 wangyicong52 Jun 11, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, I see. Thanks anyway!


torch.ops._moe_C.topk_softplus_sqrt(
topk_weights,
topk_indices,
Expand Down
Loading