The Hidden Sorting Bug: Fixing Opportunistic Expert Activation on Blackwell GPUs

Introduction

In the high-stakes world of large language model inference optimization, the difference between a working optimization and a subtly broken one often comes down to a single assumption about tensor ordering. Message 1140 of this opencode session captures one such moment: the assistant discovers and fixes a correctness bug in the Opportunistic Expert Activation (OEA) optimization for the GLM-5-NVFP4 model running on 8× NVIDIA RTX PRO 6000 Blackwell GPUs. The bug was subtle — the OEA code assumed that the topk_ids tensor returned by the MoE routing kernel was sorted by score, when in fact it was not — and fixing it required understanding the internals of both the PyTorch torch.topk function and the SGLang MoE routing implementation.

This message is a masterclass in the kind of debugging that characterizes serious systems optimization work: a hypothesis about performance is formed, an optimization is implemented and tested, a discrepancy in results leads to deeper investigation, and a subtle assumption is uncovered and corrected. The fix itself is concise — a Python script that patches a single file on a remote server — but the reasoning behind it reveals a deep understanding of the model architecture, the inference framework, and the numerical properties of the routing mechanism.

The Context: Opportunistic Expert Activation

The GLM-5-NVFP4 model is a Mixture-of-Experts (MoE) architecture with 256 experts and top-8 routing. In MoE models, each token is routed to a subset of experts (in this case, 8 out of 256), and only those experts are computed. The key insight behind OEA is that at decode time, when multiple tokens are processed in a batch, many tokens may be routed to the same experts. If we can identify the most "popular" experts early, we can compute them first and potentially skip some of the less popular ones, improving throughput.

The OEA implementation works as follows: for each batch of tokens, it takes the top-k experts selected by the router, sorts them by score, picks the top-k0 experts (where k0 < k, e.g., k0=5 out of k=8), and uses those as the "base" set. The remaining k-k0 slots are filled from experts that are already in the base set for other tokens in the batch, reducing the total number of unique experts that need to be computed.

This optimization is gated by the environment variable SGLANG_OEA_K0 and was implemented in the file /root/sglang/python/sglang/srt/layers/moe/topk.py. The initial implementation had a critical flaw: it assumed that the topk_ids tensor — which contains the indices of the top-k experts selected for each token — was sorted by score, with the highest-scored expert at position 0 and the lowest at position k-1.

The Discovery: An Unsorted Assumption

The discovery of this bug was a textbook example of the scientific method in action. The assistant had benchmarked the OEA implementation and found a modest 5.7% throughput improvement at high concurrency (1024 concurrent requests). But the improvement was smaller than expected, and the assistant began questioning the implementation details.

In message 1137, the assistant writes: "I realize the OEA function has a subtle issue. The topk_ids from the fused kernel may NOT be sorted by score. The biased_grouped_topk_impl uses torch.topk with sorted=False when num_fused_shared_experts=0. This means my assumption that topk_ids[:, :k0] contains the highest-scored experts may be wrong!"

This is the crucial insight. The torch.topk function in PyTorch has a sorted parameter. When sorted=True, the returned indices are ordered by value in descending order. When sorted=False, the indices are returned in arbitrary order — they are guaranteed to be the top-k values, but their positions in the output tensor are not meaningful. The SGLang MoE routing code uses sorted=False for most paths, meaning that topk_ids[:, 0] is not necessarily the highest-scored expert, and topk_ids[:, :k0] is not necessarily the set of the k0 highest-scored experts.

The assistant confirmed this by grepping the source code for sorted= in message 1138, finding:

topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=False)[1]
group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
    sorted=(True if num_fused_shared_experts > 0 else False),
_, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)

The pattern is clear: sorted=False is the default, and sorted=True is only used when num_fused_shared_experts > 0. For the GLM-5 model with num_expert_group=1 and likely no fused shared experts, the routing path uses sorted=False.

The Fix: Sorting by Score

The fix in message 1140 is elegant and minimal. Instead of changing the MoE routing code to use sorted=True (which would add computational overhead and might change behavior elsewhere), the assistant modifies the OEA function to explicitly sort the top-k experts by score before selecting the top-k0.

The fix script does two things:

Sorts the top-k indices by score: It computes sigmoid scores from the raw router logits, gathers the scores for the selected top-k experts using torch.gather, sorts the indices by score in descending order using argsort, and then takes the first k0 from the sorted list. This guarantees that base_ids contains the k0 highest-scored experts, regardless of whether the original topk_ids was sorted.
Removes a duplicate scores computation: The earlier fix in message 1124 had added a line scores = router_logits.float().sigmoid() inside the OEA function. But the same line was also present in Phase 2 of the function (added by the same fix). The script removes the duplicate, keeping only the Phase 1 computation. The key code change is:

# Before (buggy):
base_ids = topk_ids[:, :k0].to(torch.int64)  # (B, k0)

# After (fixed):
scores = router_logits.float().sigmoid()  # (B, N) -- sigmoid scores
topk_scores = scores.gather(1, topk_ids.to(torch.int64))  # (B, k)
sorted_indices = topk_scores.argsort(dim=-1, descending=True)  # (B, k)
topk_ids_sorted = topk_ids.gather(1, sorted_indices)  # (B, k) sorted by score
base_ids = topk_ids_sorted[:, :k0].to(torch.int64)  # (B, k0) — guaranteed highest-scored

This fix is computationally efficient: it adds only a gather, an argsort, and another gather — all on the small top-k tensor (shape B×k, where k=8), not on the full router logits (shape B×256). The sigmoid scores are computed once and reused.

Why Sigmoid, Not Raw Logits?

An important design decision visible in this fix is the use of sigmoid scores rather than raw logits for sorting. The router logits are raw, unbounded values produced by the gating network. The sigmoid function maps these to probabilities in (0, 1). Since sigmoid is a monotonic function, the ordering is preserved — sorting by sigmoid scores gives the same result as sorting by raw logits.

However, using sigmoid scores is more correct for the subsequent weight gathering in Phase 2 of the OEA function, where scores are used to compute weighted sums. The raw logits are not probabilities and would give incorrect results if used directly as weights. The sigmoid transformation ensures that the scores used for weight gathering are proper probabilities, matching the behavior of the original routing code.

The assistant had already identified this issue earlier (message 1123-1124) and added the sigmoid transformation. The current fix builds on that by also using the sigmoid scores for sorting.

Assumptions and Their Consequences

The original OEA implementation made several assumptions, some of which were correct and some incorrect:

Correct assumption: Sigmoid is monotonic, so sorting by sigmoid scores is equivalent to sorting by raw logits. This is mathematically sound.

Correct assumption: The top-k experts selected by the router are the correct set of experts to consider. The OEA function does not change which experts are selected — it only changes the order in which they are processed.

Incorrect assumption: topk_ids is sorted by score. This was the bug. The torch.topk function with sorted=False returns indices in arbitrary order, meaning that topk_ids[:, :k0] could contain any subset of the top-k experts, not necessarily the highest-scored ones.

Incorrect assumption (earlier, already fixed): Raw logits can be used as scores for weight gathering. This was fixed in message 1124 by adding the sigmoid transformation.

The consequence of the sorting bug was that OEA might select a suboptimal set of experts as the "base" set. Instead of picking the k0 highest-scored experts (which are the most important for each token), it might pick a random subset of the top-k. This would reduce the effectiveness of the optimization because the base set would be less representative of each token's preferences, leading to more unique experts overall and less opportunity for reuse.

Interestingly, the benchmarks showed a 5.7% improvement even with the buggy implementation. This suggests that even a random subset of the top-k experts provides some benefit — the optimization was working, just not at full efficiency. The fix should improve the effectiveness of OEA, particularly at medium concurrency where expert overlap is partial.

The Thinking Process: A Window into Debugging

The reasoning visible in the messages leading up to this fix reveals a systematic debugging process:

Hypothesis formation: The assistant hypothesizes that OEA should improve throughput by reducing unique expert computations.
Implementation: OEA is implemented and tested.
Benchmarking: Results show a modest 5.7% improvement at high concurrency, but the peak throughput improvement is 25%.
Questioning assumptions: The assistant wonders why the improvement isn't larger and begins examining the implementation details.
Code inspection: The assistant greps the source code to check the sorted parameter of torch.topk.
Bug identification: The sorted=False pattern is confirmed, and the bug is identified.
Fix design: The assistant designs a fix that sorts the top-k by score before selecting the top-k0.
Implementation: The fix is applied via a Python script executed on the remote server. This is classic debugging methodology: form a hypothesis, test it, examine discrepancies, dig deeper into the implementation, find the root cause, and fix it. The assistant doesn't just apply a band-aid — it understands the root cause and designs a correct, efficient fix.

Input Knowledge Required

To understand this message, one needs knowledge of:

Mixture-of-Experts (MoE) architecture: Understanding that tokens are routed to a subset of experts based on router logits, and that top-k routing selects the k experts with the highest scores.
PyTorch's torch.topk function: Knowing that the sorted parameter controls whether the output indices are ordered by value. This is a subtle API detail that many users might overlook.
Sigmoid function and its properties: Understanding that sigmoid is monotonic (preserving order) and maps logits to probabilities.
SGLang inference framework: Knowing the code structure, particularly the MoE routing implementation in topk.py.
The GLM-5-NVFP4 model architecture: Knowing that it has 256 experts with top-8 routing, and that num_expert_group=1 means the grouped routing degenerates to standard top-k.
The OEA optimization: Understanding the concept of selecting a subset of experts to reduce unique expert computations.
SSH and remote execution: The fix is applied via SSH to a remote server, requiring knowledge of the command structure and file paths.

Output Knowledge Created

This message creates several pieces of knowledge:

The bug itself: The OEA implementation had a correctness bug where it assumed topk_ids was sorted by score. This is now documented and fixed.
The fix: A Python script that patches the OEA function to explicitly sort top-k experts by score before selecting the top-k0 subset.
A reusable debugging pattern: The method of checking sorted= parameter usage in torch.topk calls is a general technique for identifying similar bugs in other codebases.
Documentation of assumptions: The fix clarifies that the OEA function now correctly handles unsorted top-k indices, and the code comments explain why the sorting is necessary.
Performance implications: The fix should improve OEA's effectiveness, particularly at medium concurrency where expert overlap is partial and the selection of the base set matters most.

The Broader Significance

This bug fix is more than just a correction — it's an example of the kind of deep, systematic thinking that characterizes high-quality systems optimization work. The assistant didn't just accept the benchmark results and move on. It questioned why the improvement was modest, dug into the implementation details, found a subtle assumption violation, and fixed it.

The fix also demonstrates the importance of understanding the numerical properties of the operations involved. The assistant knew that sigmoid is monotonic, that sorting by sigmoid scores preserves the ordering of logits, and that using proper probability scores for weight gathering is essential for correctness. This kind of mathematical understanding is what separates a hack from a robust optimization.

In the broader context of the GLM-5 optimization campaign, this fix represents a refinement of one of the more creative optimization ideas — OEA. While the initial results were modest, the fix should improve the optimization's effectiveness, and the lessons learned about tensor ordering assumptions will inform future optimization work.

Conclusion

Message 1140 captures a moment of insight and correction in a complex optimization campaign. The assistant discovered that a seemingly innocuous assumption — that torch.topk returns sorted indices — was incorrect, and that this assumption was undermining the OEA optimization. The fix is elegant, efficient, and correct, and the reasoning behind it reveals a deep understanding of the MoE architecture, the PyTorch API, and the inference framework.

This message is a reminder that in systems optimization, the devil is in the details. A single assumption about tensor ordering can make the difference between an optimization that works and one that works correctly. The assistant's systematic approach — hypothesize, implement, test, question, inspect, fix — is a model for how to do this kind of work well.