Skip to content

[ROCM] Optimize the fused_topk_bias to use aiter instead of fallback torch ops. #36253

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

Merged
gshtras merged 4 commits into from
Mar 9, 2026
Merged

[ROCM] Optimize the fused_topk_bias to use aiter instead of fallback torch ops. #36253

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c12f25e
add minimax m2.5 opt
benenzhu Mar 6, 2026
98d714c
change code according to gemini
benenzhu Mar 6, 2026
913d44c
Merge branch 'main' into rocm-minimax-opt
benenzhu Mar 6, 2026
302fdee
Merge branch 'main' into rocm-minimax-opt
benenzhu Mar 7, 2026
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
36 changes: 36 additions & 0 deletions vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,42 @@ def fused_topk_bias(
return topk_weights, topk_ids
else:
raise ValueError(f"Unsupported scoring function: {scoring_func}")
elif rocm_aiter_ops.is_fused_moe_enabled() and scoring_func == "sigmoid":
# aiter only have sigmoid bias grouped topk now.
M = hidden_states.size(0)
num_experts = gating_output.shape[-1]
_AITER_MAX_EXPERTS_PER_GROUP = 32
num_expert_group = max(
1, -(-num_experts // _AITER_MAX_EXPERTS_PER_GROUP)
) # ceil division
while num_experts % num_expert_group != 0:
num_expert_group += 1

@gemini-code-assist gemini-code-assist Bot Mar 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

high

The calculation of num_expert_group is performed on every forward pass. The while loop can be inefficient if num_experts is a large number with no small divisors (e.g., a prime number). For instance, if num_experts is a prime like 59, this loop will iterate many times on every call. This could introduce a performance overhead in what is supposed to be a fast path.

Since num_experts is constant for a given model, this value should be computed once and cached. The best approach would be to compute it during model initialization and pass it to this function. As a self-contained alternative, you can cache the result within this function using functools.lru_cache as shown in the suggestion.

        import functools

        @functools.lru_cache(maxsize=8)
        def _get_num_expert_group(num_experts: int) -> int:
            _AITER_MAX_EXPERTS_PER_GROUP = 32
            g = max(1, -(-num_experts // _AITER_MAX_EXPERTS_PER_GROUP))
            while num_experts % g != 0:
                g += 1
            return g

        num_expert_group = _get_num_expert_group(num_experts)

@benenzhu benenzhu Mar 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in new commit, added the LRU cache to the function.

topk_weights = torch.empty(
M, topk, dtype=torch.float32, device=hidden_states.device
)
assert num_experts % num_expert_group == 0, (
f"{num_experts=} not divisible by {num_expert_group=}"
)
assert num_experts // num_expert_group <= _AITER_MAX_EXPERTS_PER_GROUP, (
f"group size {num_experts // num_expert_group} \
exceeds limit {_AITER_MAX_EXPERTS_PER_GROUP}"
)
topk_ids = torch.empty(
M,
topk,
dtype=torch.int32 if indices_type is None else indices_type,
device=hidden_states.device,
)
rocm_aiter_ops.biased_grouped_topk(
gating_output,
e_score_correction_bias.to(gating_output.dtype),
topk_weights,
topk_ids,
num_expert_group=num_expert_group,
topk_group=num_expert_group,
need_renorm=renormalize,
)

@gemini-code-assist gemini-code-assist Bot Mar 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

security-high high

The code added in the PR introduces a fast path for ROCm using the aiter library's biased_grouped_topk kernel. However, there is a potential mismatch between the allocated size of the output tensors (topk_weights and topk_ids) and the number of elements the kernel is instructed to write. The output tensors are allocated with topk columns (lines 121-123 and 131-136), but the kernel is called with num_expert_group and topk_group=num_expert_group (lines 142-143). In grouped top-k kernels like aiter's, the total number of elements written to the output is typically num_expert_group * topk_group. In this code, that total is num_expert_group * num_expert_group. There is no check to ensure that topk == num_expert_group * num_expert_group. If topk is less than num_expert_group**2, the kernel will write past the end of the allocated buffers, leading to GPU memory corruption. For example, for a model with 128 experts and topk=1, num_expert_group would be 4, and the kernel would attempt to write 16 elements into a buffer of size 1. This can lead to data integrity issues, crashes, or potential information leakage in a multi-user environment.

@benenzhu benenzhu Mar 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in new commit. Added checks and fallback to original path when checks don't meet.

return topk_weights, topk_ids

n_routed_experts = gating_output.shape[-1]
if scoring_func == "softmax":
Expand Down