The Sigmoid Fix: A Single-Line Correction in an MoE Routing Optimization

In the course of an intensive optimization campaign targeting the GLM-5-NVFP4 large language model running on eight NVIDIA RTX PRO 6000 Blackwell GPUs, an assistant implementing a novel decode-time routing optimization called Opportunistic Expert Activation (OEA) paused mid-benchmark to correct a subtle but important bug. The message in question, message [msg 1124], is deceptively simple — a single sed command that changes one line of Python code:

ssh root@10.1.230.174 'sed -i "s/    scores = router_logits.float()  # (B, N)/    scores = router_logits.float().sigmoid()  # (B, N) — sigmoid scores for proper weight gathering/" /root/sglang/python/sglang/srt/layers/moe/topk.py'

This one-line edit transforms router_logits.float() into router_logits.float().sigmoid(), adding a sigmoid activation to convert raw logits into proper probability scores. Behind this tiny change lies a chain of careful reasoning, a discovered bug in a freshly implemented optimization, and a deeper understanding of how Mixture-of-Experts routing works under the hood.

The Context: Implementing OEA

The story begins with the assistant's broader mission: maximizing inference throughput for GLM-5-NVFP4, a Mixture-of-Experts (MoE) model with 256 experts and top-8 routing. The core bottleneck, as established across many prior experiments, was the small per-expert GEMM operations on Blackwell SM120 GPUs, where each expert's matrix multiply is too small to fully utilize the GPU's compute capacity. One promising optimization was Opportunistic Expert Activation (OEA), a technique described in a recent paper (arXiv:2511.02237) that reduces the number of unique experts loaded per decode batch by "piggybacking" tokens onto experts that are already being used by other tokens in the same batch.

The assistant had implemented OEA as a patch to SGLang's topk.py file in message [msg 1106], inserting a new opportunistic_expert_activation() function and a call site within the select_experts routine. The implementation was careful and thorough: it preserved the top-k0 baseline experts per token, then filled remaining slots with experts already selected by other tokens in the batch, gathering weights and renormalizing. The function was gated by the SGLANG_OEA_K0 environment variable and skipped for small batches.

Discovering the Bug

After implementing OEA, the assistant ran a series of benchmarks to measure its impact. The results at concurrency 1024 showed a modest 5.7% improvement in output throughput (1,607.61 tok/s vs 1,520.55 tok/s baseline) and a more impressive 25% improvement in peak output tokens per second (4,012 vs 3,210). These were encouraging numbers, but the assistant's attention was drawn to a deeper question.

In message [msg 1123], the assistant explicitly reasoned about the correctness of the implementation:

"The key question is whether router_logits passed to OEA is pre-sigmoid (raw logits) or post-sigmoid (scores). Looking at the code: select_experts receives router_logits which is raw logits, and the sigmoid is applied inside biased_grouped_topk_impl or the fused kernels. My OEA function uses router_logits.float() for scoring — but these are raw logits, not sigmoid scores. For sorting/ranking this should still work (sigmoid is monotonic, so ordering is preserved), but the renormalization needs to use sigmoid scores instead of raw logits."

This is a critical insight. In MoE routing, the router network produces logits — unbounded real numbers — which are then passed through a sigmoid (or softmax) activation to produce probabilities in the range [0, 1]. The top-k selection is based on these probabilities, and the selected weights are used in the weighted sum of expert outputs. The OEA function was using raw logits for two purposes: (1) ranking experts to decide which ones to piggyback, and (2) gathering weights for renormalization after modifying the expert selection.

For purpose (1), using raw logits is harmless because sigmoid is a monotonic function — the order of experts by logit value is identical to the order by sigmoid probability. But for purpose (2), using raw logits is incorrect. The renormalization step divides each selected weight by the sum of all selected weights, and raw logits can be arbitrarily large or small, producing nonsensical normalized weights. The sigmoid squashes values into [0, 1], giving proper probability-like weights for the renormalization.

The Fix

The fix itself is trivial in terms of code change — adding .sigmoid() to a tensor operation — but it represents an important correction. The assistant used a sed in-place substitution over SSH to modify the remote file, targeting line 942 of topk.py (as confirmed by the grep -n command in message [msg 1123]). The comment was also updated to document the change: "sigmoid scores for proper weight gathering."

In the subsequent message [msg 1125], the assistant verified the correctness of the fix: "Good. Now the weights gathered from scores will be proper sigmoid probabilities, matching what the original routing code produces. The ranking is the same either way (sigmoid is monotonic), so the expert selection logic is unchanged."

Why This Matters

This single-line fix illuminates several important aspects of the optimization process. First, it demonstrates the importance of understanding the data flow through a complex system. The select_experts function receives raw logits, but the downstream fused kernels (like FlashInfer's biased_grouped_topk_impl) apply sigmoid internally. The OEA function, intercepting the routing mid-stream, needed to replicate that transformation. A less careful implementer might have missed this detail entirely, producing incorrect weight renormalization that could silently degrade model quality.

Second, it shows how benchmarking and analysis feed back into correctness. The assistant didn't just implement OEA and declare victory; it scrutinized the benchmark results, questioned whether the implementation was correct, traced the code path, and found the bug. This is the hallmark of rigorous engineering: every optimization is not just tested for performance but audited for correctness.

Third, it reveals an interesting property of the OEA algorithm itself. Even with the bug fix, the assistant's subsequent benchmarks (as described in the chunk summary) showed that OEA produced "near-zero average throughput improvement on random data." The peak throughput improved ~5% at high concurrency, but the average was flat. The reason was that on random data with uniform expert routing, nearly all 256 experts are already active at high batch sizes — there's no opportunity for piggybacking. OEA's benefit would only manifest on workloads with non-uniform expert routing patterns, such as conversational or structured generation tasks.

The Broader Pattern

This message is a microcosm of the entire optimization campaign documented in Segment 9. The assistant operates in a tight loop of hypothesis, implementation, measurement, analysis, and refinement. Each optimization idea — whether it's OEA, Expert Parallelism, Piecewise CUDA Graphs, or MSCCLPP allreduce — is implemented cleanly, benchmarked against a baseline, analyzed for root causes, and either adopted or discarded based on real data. The OEA sigmoid fix is a perfect example of the "analysis" phase feeding back into "implementation": the benchmark results prompted a correctness audit, which revealed a bug, which was fixed before further testing.

The article also highlights the value of understanding the mathematical properties of the operations involved. The assistant's reasoning that "sigmoid is monotonic, so ordering is preserved" is a concise application of a fundamental property of monotonic functions, allowing the fix to be applied with confidence that the expert selection logic remains unchanged while only the weight computation is corrected.

Conclusion

Message [msg 1124] is a single sed command — seven words of Python code changed. But it encapsulates a complete cycle of engineering rigor: implement, benchmark, question, trace, correct, verify. The sigmoid fix ensured that OEA's weight renormalization used proper probability scores rather than raw logits, preventing a subtle correctness bug that could have compromised model quality. It also demonstrates that even in a high-throughput optimization campaign focused on pushing tokens-per-second numbers, the assistant never lost sight of correctness — a crucial discipline when modifying the internals of a production inference engine.