The Research Pivot: Investigating MoE Expert GEMM Optimization for Blackwell SM120

Introduction

In the course of a marathon optimization session for deploying trillion-parameter Mixture-of-Experts (MoE) language models on an 8× NVIDIA RTX PRO 6000 Blackwell GPU system, a pivotal moment arrived when the conversation took a sharp turn from communication optimization toward compute optimization. This article examines message index 2402 in the conversation — a pure research message that launched three parallel investigation subagents to explore the state of the art in optimizing tiny MoE expert GEMM (General Matrix Multiply) operations for NVIDIA's SM120 Blackwell architecture.

The message itself is concise, containing only a brief acknowledgment and three task tool calls:

Great question. Let me research what it would actually take to optimize those tiny MoE expert GEMMs for SM120 (Blackwell). This is the real bottleneck — let me dig deep.

>

[task] Research SM120 GEMM optimization — understanding current MoE expert GEMM kernels in vLLM for SM120 [task] Research Blackwell GEMM optimization — state of the art for small/skinny GEMM on Blackwell SM120 [task] Find local GEMM/MoE analysis files — search local project directory for existing analysis

Despite its brevity, this message represents a significant strategic pivot in the optimization campaign. It marks the moment when the assistant and user jointly recognized that the true bottleneck was not inter-GPU communication (allreduce) but rather the computational inefficiency of tiny matrix multiplications caused by aggressive tensor parallelism sharding. Understanding why this realization occurred, what assumptions it challenged, and what knowledge it produced is essential to grasping the broader narrative of deploying large MoE models on PCIe-only multi-GPU systems.

The Context: A Conversation About Bottlenecks

To understand message 2402, we must trace the conversation that led to it. The user had been engaged in an intensive multi-day effort to deploy and optimize several large MoE models — GLM-5, Kimi-K2.5 NVFP4, MiniMax-M2.5, and Kimi-K2.5 INT4 — on a system with 8× NVIDIA RTX PRO 6000 Blackwell GPUs connected solely via PCIe Gen5, with no NVLink interconnect. This hardware constraint meant that all inter-GPU communication (required by tensor parallelism) traversed the PCIe bus, creating a potential bottleneck.

Earlier in the conversation ([msg 2398]), the user had asked a perceptive question: "If allreduce is so slow it seems like it would make sense to run two inferences at the same time, one does allreduce while the other computes. Is that something vllm or sglang support?" This question revealed the user's intuition that overlapping communication with computation could hide allreduce latency — a classic optimization technique in distributed computing.

The assistant's response at [msg 2400] delivered a surprising finding after extensive research: allreduce was not the primary bottleneck. The assistant reported that all communication optimizations tested (MSCCLPP, Single Batch Overlap, allreduce fusion, even TP4+PP2 which halves allreduce data) yielded only 0-2% improvement. The real bottleneck, the assistant explained, was that with 256+ experts and TP=8, each expert GEMM was tiny (N=256) and couldn't saturate the GPU's compute units. The GPU was spending most of its time on memory-bandwidth-bound small matrix multiplies, not waiting on allreduce.

This revelation reframed the entire optimization problem. The user's response at [msg 2401] captured the mix of frustration and determination: "Aaah it's still the GEMMs! What would it take to properly optimize that for this GPU? Want to squeeze out every last FLOP."

The Message Itself: A Research-Only Pivot

Message 2402 is the assistant's response to this challenge. It is notable for what it does not contain: there are no bash commands, no file edits, no configuration changes, no model deployments. It is a pure research message, launching three parallel task subagents to gather information before any action is taken.

The structure is straightforward. The assistant acknowledges the question with enthusiasm ("Great question"), frames the investigation ("Let me research what it would actually take to optimize those tiny MoE expert GEMMs for SM120"), and then dispatches three research tasks in parallel. This is a pattern we see repeatedly in the conversation: when faced with a complex, multi-faceted problem, the assistant uses the task tool to spawn subagents that can independently research different aspects, then synthesizes the results.

The three research tasks cover distinct but complementary domains:

Task 1: "Research SM120 GEMM optimization" — This task investigates the current state of MoE expert GEMM kernels in vLLM for SM120 (Blackwell, compute capability 12.0). It specifically targets the problem of tiny GEMMs (e.g., N=256 for hidden_size=2048/8) caused by TP=8 sharding of 256+ experts. This is the most immediately practical task — it asks what kernel paths are currently used and what optimizations exist within the vLLM ecosystem.

Task 2: "Research Blackwell GEMM optimization" — This task broadens the scope to the state of the art for optimizing small/skinny GEMM operations on NVIDIA Blackwell GPUs generally. It searches for CUTLASS 4.x/CuTe support for SM120, Marlin kernels, persistent kernel approaches, L2 cache pinning, and column-major tile scheduling. This task recognizes that the solution may not come from vLLM-specific code but from fundamental GPU kernel optimization techniques.

Task 3: "Find local GEMM/MoE analysis files" — This task searches the local project directory for any existing profiling data, analysis, or notes about GEMM performance, MoE expert kernel performance, or optimization strategies. This is a pragmatic task that acknowledges the project already contains extensive benchmarking data from previous model deployments (GLM-5, Kimi-K2.5 NVFP4, MiniMax-M2.5) that might contain relevant insights.

Why This Message Was Written: The Reasoning and Motivation

The motivation for message 2402 can be traced to a fundamental reframing of the optimization problem. Earlier in the conversation, the assistant and user had been operating under the assumption that allreduce was the primary bottleneck. This assumption seemed reasonable: the system had 8 GPUs connected only by PCIe Gen5, with no NVLink, and the Kimi-K2.5 model has 61 MLA (Multi-head Latent Attention) layers that all require allreduce operations across all 8 GPUs. The profiling data had shown allreduce consuming a significant portion of decode time.

However, the research at [msg 2400] revealed a more nuanced picture. While allreduce was measurable, it was not the dominant factor. The real issue was that tensor parallelism (TP=8) was sharding each expert's weight matrices into 8 pieces, making each piece too small to efficiently utilize the GPU's tensor cores. With 384 routed experts (top-8 activated per token) and a hidden size of 2048, each expert's GEMM had dimensions like M=token_count, K=2048/8=256, N=hidden_size. These tiny dimensions meant the GPU spent most of its time moving data rather than computing.

The user's reaction — "Want to squeeze out every last FLOP" — signaled a shift from trying to hide communication latency to trying to maximize compute efficiency. This is a fundamentally different optimization strategy. Instead of asking "how do we overlap allreduce with compute?" the question became "how do we make each individual matrix multiplication as fast as possible?"

Assumptions Made and Their Validity

Message 2402 makes several implicit assumptions, some of which turned out to be incorrect in light of later profiling data.

Assumption 1: The GEMMs are the primary bottleneck. The message states "This is the real bottleneck." This assumption was based on the earlier research showing that allreduce optimizations yielded only 0-2% improvement. However, later in the same segment (chunk 0), a comprehensive torch.profiler capture would reveal that AllReduce actually accounts for 51.5% of decode time (11.17ms per step) — not the GEMMs. This appears contradictory to the earlier finding, and the resolution lies in understanding that the Marlin W4A16 kernels (used for the INT4 quantized Kimi-K2.5) eliminated the dtype-cast overhead that had plagued earlier NVFP4 runs, thereby unmasking the PCIe allreduce bottleneck that had been hidden behind other inefficiencies.

Assumption 2: The solution lies in kernel optimization. By launching research into CUTLASS 4.x, Marlin kernels, and persistent fused kernels, the assistant implicitly assumes that the path to improvement is through better GPU kernel implementations. This assumption is reasonable — Blackwell (SM120) is a new architecture with significant differences from SM100 (datacenter Blackwell), and kernel support is still maturing. However, the later discovery that allreduce dominates 51.5% of decode time suggests that even perfect GEMM kernels would leave most of the performance on the table.

Assumption 3: The existing profiling data is sufficient. Task 3 searches for local analysis files, assuming that the extensive profiling from previous model deployments contains relevant insights. This assumption proved partially correct — the project did contain extensive data — but the critical insight (allreduce being 51.5% of decode time) came from new profiling, not from existing data.

Input Knowledge Required

To fully understand message 2402, one needs substantial background knowledge spanning multiple domains:

GPU Architecture Knowledge: Understanding what SM120 (Blackwell GeForce/RTX) is, how it differs from SM100 (datacenter Blackwell), and the implications for kernel development. The RTX PRO 6000 Blackwell uses SM120, which has different instruction set support and tensor core capabilities compared to the B200/GB200's SM100.

MoE Model Architecture: Understanding Mixture-of-Experts models, particularly the Kimi-K2.5 architecture with 384 routed experts, top-8 activation, MLA attention, and the implications of tensor parallelism on expert sharding. When TP=8 shards each expert across 8 GPUs, each GPU only holds 1/8 of each expert's weights, making the GEMM dimensions proportionally smaller.

Quantization Formats: Understanding the difference between NVFP4 (NVIDIA's FP4 format) and INT4 (compressed-tensors INT4 with group_size=32). The INT4 format uses Marlin kernels for W4A16 computation, which have different performance characteristics than the NVFP4 kernels used earlier.

vLLM/SGLang Architecture: Understanding the different kernel backends (Triton, FlashAttention, FlashInfer, CUTLASS) and which ones support which GPU architectures. For SM120, only Triton-based MLA works, which constrains optimization options.

CUDA/CUTLASS Ecosystem: Understanding the CUTLASS 4.x template library for GEMM operations, CuTe (CUDA Templates), and the Marlin sparse matrix kernel format.

Output Knowledge Created

Message 2402 itself does not produce concrete results — it is a research dispatch. The output knowledge is created by the three subagents that execute in response to these task calls. However, the message establishes the framework for what knowledge is sought:

1. Current kernel paths used by vLLM for MoE expert computation on SM120 — identifying exactly which code paths are executed for the Kimi-K2.5 INT4 model's expert GEMMs.
2. State-of-the-art techniques for small GEMM optimization on Blackwell — including CUTLASS 4.x support, Marlin kernel applicability, persistent kernel approaches, L2 cache pinning strategies, and column-major tile scheduling.
3. Existing project profiling data — inventorying all the analysis files, benchmark results, and optimization notes already accumulated in the project directory. This knowledge would then inform the next phase of the optimization campaign, which (as the chunk summary reveals) would involve a comprehensive three-phase benchmarking plan: macro-level throughput/latency tests, micro-benchmarks of individual GEMM operations, NCCL AllReduce burst measurements, and a full torch.profiler capture.

The Thinking Process Visible in Reasoning

The assistant's reasoning in message 2402 reveals several important cognitive patterns:

Problem Decomposition: The assistant immediately breaks the broad question ("optimize for this GPU") into three distinct research streams: what vLLM currently does (task 1), what the state of the art is (task 2), and what we already know (task 3). This decomposition is characteristic of effective technical problem-solving — rather than trying to answer the question in one shot, the assistant identifies the sub-questions that need to be answered first.

Parallel Exploration: By launching all three tasks simultaneously, the assistant exploits the subagent parallelism available in the tool architecture. This is a deliberate strategy: the three research directions are independent (they don't depend on each other's results), so they can be explored concurrently. This minimizes the number of conversation rounds needed to gather the information.

Acknowledging Uncertainty: The assistant doesn't pretend to know the answer. Instead, it explicitly frames the message as research ("Let me research what it would actually take"). This intellectual honesty is important in the context of a complex, novel problem — Blackwell SM120 is a new architecture, and the behavior of MoE expert GEMMs at these dimensions is not well-documented.

Reframing the Problem: The most significant cognitive move is the reframing from "how do we overlap allreduce with compute?" to "how do we make these tiny GEMMs faster?" This reframing was enabled by the research at [msg 2400] which showed that allreduce optimizations had negligible impact. The assistant is now operating under a new hypothesis: that the bottleneck is computational, not communication.

Mistakes and Incorrect Assumptions

The most significant "mistake" in message 2402 is the assumption that GEMMs are the primary bottleneck. As the later profiling would reveal, AllReduce actually accounts for 51.5% of decode time. However, this isn't really a mistake in the traditional sense — it's a hypothesis that was being tested. The assistant explicitly says "Let me research" and "let me dig deep," indicating that the conclusion is provisional.

The apparent contradiction between "allreduce is only 5-8% of time" (from [msg 2400]) and "allreduce is 51.5% of decode time" (from later profiling) deserves examination. The resolution likely lies in the difference between the NVFP4 model (where dtype-casting overhead dominated) and the INT4 model (where Marlin kernels eliminated that overhead, revealing the allreduce bottleneck underneath). Message 2402 was written before this distinction was fully understood.

Another potential issue is the assumption that kernel optimization within the existing vLLM framework is the right approach. The alternative — Expert Parallelism (EP), which eliminates MoE allreduce entirely by placing complete experts on each GPU — was mentioned in [msg 2400] but not pursued here. The user had concerns about EP's all-to-all communication on PCIe without NVLink, which would face serialization issues at higher batch sizes. Message 2402 implicitly accepts this constraint by focusing on GEMM optimization within the TP=8 framework.

Conclusion

Message 2402 represents a strategic pivot point in a complex optimization campaign. It marks the transition from communication optimization to compute optimization, from trying to hide allreduce latency to trying to maximize GEMM throughput. While the assumption that GEMMs are the primary bottleneck would later be challenged by more detailed profiling, the message's value lies in its systematic approach to problem decomposition and parallel research.

The message also illustrates a key pattern in the conversation: the use of parallel research subagents to gather information before committing to a course of action. This pattern — research first, act second — is particularly appropriate for novel problems where the solution space is poorly understood.

For readers following the broader narrative, message 2402 is the moment when the optimization campaign entered its most intensive phase. The subsequent profiling would reveal the true nature of the bottleneck (allreduce at 51.5%), leading to yet another pivot — this time toward speculative decoding as a software-only optimization path. But that pivot would not have been possible without the foundational research initiated in this message.