The Unsorted Top-K: A Critical Assumption Uncovered in MoE Routing Optimization

Introduction

In the course of an intensive optimization campaign for deploying the GLM-5-NVFP4 mixture-of-experts (MoE) model across 8 NVIDIA RTX PRO 6000 Blackwell GPUs, a single assistant message (message index 1139) captures a pivotal moment of debugging insight. The message is deceptively short — a two-sentence analysis followed by a single bash command — but it represents the culmination of a multi-step reasoning chain that had been building over the preceding dozen messages. In it, the assistant realizes that a core assumption underlying its newly implemented "Opportunistic Expert Activation" (OEA) optimization is incorrect, potentially invalidating the optimization's correctness and calling into question the benchmark results already collected. This article unpacks that moment: the reasoning that led to the discovery, the assumptions that were broken, the technical context required to understand the stakes, and the knowledge produced by the correction.

The Context: What Is OEA and Why Does Sorting Matter?

To understand message 1139, one must first understand what the assistant was trying to accomplish. The GLM-5-NVFP4 model is a Mixture-of-Experts (MoE) transformer with 256 experts per MoE layer. During inference, each token is routed to a subset of experts — specifically, the top-k experts by router score, where k is typically 8 for this model. The computational cost of an MoE layer scales with the number of unique experts activated across the batch, because each unique expert requires loading its weights into SRAM and performing a GEMM (general matrix multiply) operation.

The OEA optimization, implemented in the messages immediately preceding 1139 (messages 1110–1125), introduced a clever idea: at decode time, when the batch contains multiple tokens that may each activate different experts, why not opportunistically skip the lowest-scored experts for each token? If a token's 8th-best expert has a very low score, perhaps the model can get away with only computing the top 5 (k0=5) experts, reducing the total number of unique experts that need to be computed across the batch. The reduction in unique experts directly reduces the GEMM workload, potentially improving throughput.

The OEA implementation worked by intercepting the router's output — specifically the topk_ids tensor, which contains the indices of the top-k experts for each token — and truncating it to only the first k0 columns. The critical assumption, stated explicitly in the code comment at line 934 of topk.py, was that topk_ids is sorted by score, meaning the first k0 entries are guaranteed to be the k0 highest-scored experts. If this assumption holds, then topk_ids[:, :k0] correctly selects the best k0 experts, and the optimization is sound.

The Discovery: A Chain of Doubt

The chain of reasoning that led to message 1139 began several messages earlier. In message 1122, the assistant benchmarked OEA at concurrency 1024 and observed a 5.7% improvement in output throughput and a 25% improvement in peak output. These were promising but not transformative numbers. In message 1123, the assistant started questioning implementation details — specifically whether the router_logits tensor passed to OEA was raw logits or sigmoid scores, and whether the OEA code path was actually being triggered for the GLM-5 model's specific routing configuration (num_expert_group=1, likely using the flashinfer fused kernel path).

Then in message 1137, a deeper concern surfaced. The assistant recalled that biased_grouped_topk_impl — the routing function used by GLM-5 — calls torch.topk with sorted=False when num_fused_shared_experts=0. This is a performance optimization: sorting the top-k results is unnecessary for the routing computation itself (the model only needs the indices and weights, not their order), so the fused kernel skips the sorting step to save cycles. But OEA's correctness depends on that ordering. If the top-k results are unsorted, then topk_ids[:, :k0] is not the k0 highest-scored experts — it's an arbitrary subset of the top-k, determined by whatever order the kernel happened to produce.

Message 1138 confirmed the suspicion. The assistant grepped for sorted= across the entire topk.py file and found multiple occurrences of sorted=False:

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

Only the num_fused_shared_experts > 0 path uses sorted=True. For the standard path used by GLM-5, sorting is disabled.

Message 1139: The Moment of Clarity

With the evidence gathered, message 1139 delivers the conclusion with characteristic concision:

sorted=False for most paths. This means topk_ids[:, :k0] is NOT guaranteed to be the k0 highest-scored experts — it's an arbitrary subset of the top-k. This is a correctness issue in our OEA implementation.

The assistant then immediately moves to fix the issue, running a bash command to locate the exact line where the assumption was embedded in the code:

Let me fix the OEA function to explicitly sort the top-k by score first: [bash] ssh root@10.1.230.174 'grep -n "# Phase 1: baseline = top-k0 per token (topk_ids are sorted by score)" /root/sglang/python/sglang/srt/layers/moe/topk.py' 934: # Phase 1: baseline = top-k0 per token (topk_ids are sorted by score)

The command finds the comment at line 934 that explicitly documents the now-broken assumption. The irony is palpable: the code comment itself asserts "topk_ids are sorted by score" — a statement that the assistant now knows to be false.

The Reasoning Process: What Makes This Message Significant

Message 1139 is significant not for its length but for what it represents: the moment when a chain of evidence converges into a definitive conclusion. The assistant's thinking process, visible across the preceding messages, follows a classic debugging pattern:

  1. Observation: OEA shows modest improvement (5.7%), but the mechanism is unclear. Is OEA actually doing what we think it's doing?
  2. Hypothesis generation: Perhaps the router logits are pre-sigmoid (raw) rather than post-sigmoid (scores). Perhaps OEA isn't being called at all for this model's routing path.
  3. Investigation: The assistant checks whether OEA is loaded (it is — 10 occurrences in the log), then checks the sigmoid issue and fixes it.
  4. Deeper questioning: But wait — even with correct scores, is the topk_ids tensor sorted? The routing kernel uses sorted=False for performance.
  5. Confirmation: The assistant greps the codebase and finds multiple sorted=False calls.
  6. Conclusion (message 1139): The assumption is broken. The optimization as implemented is incorrect. This pattern — moving from surface-level benchmarking to progressively deeper implementation scrutiny — is characteristic of the assistant's methodology throughout the session. Each optimization is tested, questioned, and either validated or discarded based on evidence, not speculation.

The Assumptions at Stake

Message 1139 exposes several layers of assumptions that were baked into the OEA implementation:

Assumption 1: torch.topk returns sorted results by default. This is actually the default behavior of torch.topk — when called without the sorted argument, it defaults to sorted=True. The assistant reasonably assumed that the fused routing kernels in sglang would either use the default or explicitly request sorted output. But the codebase's custom kernels and routing implementations explicitly pass sorted=False for performance, overriding the default.

Assumption 2: The code comment at line 934 is accurate. The comment "topk_ids are sorted by score" was either written based on the same incorrect assumption, or was left stale after a code change that switched to unsorted output. In either case, it's a reminder that comments can mislead as easily as they can inform.

Assumption 3: The OEA optimization's correctness is independent of the routing kernel's internals. The assistant designed OEA as a post-processing step on the router's output, assuming a stable interface contract (sorted top-k). But the router's internal implementation violated that contract for performance reasons, and the OEA code never verified its input.

Assumption 4: The benchmark results at concurrency 1024 reflect genuine OEA benefit. If OEA was selecting arbitrary subsets of experts rather than the top-k0, then the 5.7% throughput improvement might be partly or entirely spurious — perhaps the model was simply computing fewer experts on average (a real computational saving) but with degraded output quality (since the wrong experts were being selected). The assistant had not yet measured output quality (e.g., perplexity or task accuracy) to validate that OEA preserved model fidelity.

Input Knowledge Required to Understand This Message

A reader needs several layers of context to fully grasp message 1139:

MoE routing fundamentals: Understanding that each token in an MoE layer is routed to a subset of experts via a learned router, that the router produces scores (often via sigmoid or softmax over logits), and that the top-k experts by score are selected for computation.

The torch.topk API: Knowing that torch.topk has a sorted parameter that controls whether the returned indices are sorted by descending value. The default is sorted=True, but the codebase overrides this.

The OEA optimization: Understanding that OEA truncates the top-k expert list to k0 < k, and that this truncation only makes sense if the first k0 entries are the highest-scored ones.

The sglang codebase structure: Knowing that topk.py in sglang/srt/layers/moe/ contains the routing implementations, and that the GLM-5 model uses biased_grouped_topk_impl with num_expert_group=1 and num_fused_shared_experts=0.

The hardware context: The optimization is running on NVIDIA RTX PRO 6000 Blackwell GPUs (SM120 architecture), where the small per-expert GEMMs are the primary bottleneck, making expert count reduction a valuable optimization target.

Output Knowledge Created by This Message

Message 1139 produces several forms of knowledge:

Direct knowledge: The topk_ids tensor from the MoE router is unsorted for the GLM-5 model's routing path, invalidating the assumption that topk_ids[:, :k0] selects the top-k0 experts.

Methodological knowledge: The message demonstrates a reproducible debugging methodology — when an optimization's mechanism is unclear, trace the data flow through the codebase rather than relying on benchmark numbers alone. The assistant could have accepted the 5.7% improvement at face value, but instead chose to verify the implementation's correctness.

Code archaeology: The grep results reveal that sorted=False is used in multiple routing paths (scores_for_choice, tmp_scores, group_scores), suggesting this is a deliberate design choice across the codebase, not an isolated oversight. Only the num_fused_shared_experts &gt; 0 path uses sorted=True.

Documentation gap: The comment at line 934 ("topk_ids are sorted by score") is now known to be inaccurate, creating a documentation fix opportunity.

Future work: The assistant now needs to either (a) add an explicit sorting step before the OEA truncation, (b) modify the routing kernel to return sorted results, or (c) redesign OEA to work with unsorted inputs (e.g., by using the scores tensor directly rather than relying on index ordering).

The Broader Significance

Message 1139 is a microcosm of the entire optimization campaign documented in this session. The campaign is characterized by a relentless cycle of: implement → benchmark → question → investigate → discover flaw → fix → re-benchmark. Each cycle produces not just performance improvements but also deeper understanding of the system's internals.

The OEA story is particularly instructive because the optimization was partially working even with the bug — the 5.7% improvement at concurrency 1024 was real, but it was achieved with incorrect expert selection. This is a dangerous regime for optimization work: a bug that produces positive results can easily escape detection if the implementer doesn't dig deep enough. The assistant's decision to question the results rather than celebrate them is what separates rigorous engineering from lucky hacking.

The message also highlights a tension between performance and correctness that pervades ML systems engineering. The routing kernels use sorted=False because sorting adds overhead — and for the routing computation itself, ordering is irrelevant. But OEA, as a downstream consumer of the routing output, implicitly depends on that ordering. This is a classic interface contract violation: the producer (routing kernel) optimized for its own performance without documenting the change in output semantics, and the consumer (OEA) assumed a contract that no longer held.

Conclusion

Message 1139 captures the moment when an incorrect assumption meets hard evidence. In just two sentences and a bash command, the assistant articulates a correctness bug in its OEA optimization, traces it to the sorted=False parameter in torch.topk calls, and begins the fix process. The message is the culmination of a multi-step reasoning chain that moved from benchmark numbers to implementation scrutiny to code-level verification. It demonstrates that in systems optimization, the most valuable output is not always a throughput number — sometimes it's the knowledge that a number you trusted was built on sand.

The fix that follows from this message — adding explicit sorting before the OEA truncation, or redesigning the approach — will determine whether OEA is a genuine optimization or a bug that happened to improve throughput by accident. Either way, the knowledge produced by message 1139 is essential: without it, the OEA results would have been incorporated into the findings document as valid, and future work would have built on a flawed foundation. The message is a testament to the value of intellectual honesty in engineering work, and a reminder that the most important debugging tool is the willingness to question one's own assumptions.