The MoE Throughput Lever: Quantifying the INT4 Marlin GEMM on Blackwell PRO 6000

Introduction

In the high-stakes world of large language model inference, every microsecond counts. When deploying a 1-trillion-parameter model like Kimi K2.6 across eight RTX PRO 6000 Blackwell GPUs, the difference between a system that delivers usable throughput and one that collapses under its own weight often comes down to understanding a single kernel's behavior under load. Message 12023 in this opencode session captures exactly such a moment: the assistant reports the completion of Phase 3's highest-value deliverable—a precise quantification of the INT4 Marlin Mixture-of-Experts (MoE) GEMM throughput characteristics on the target hardware.

This message is not a tool call, a code edit, or a command execution. It is a results summary and architectural analysis delivered to the user after a series of benchmarks were run and the live service was restored. It synthesizes raw benchmark numbers into a refined understanding of how the speculative decoding engine must be designed to maximize throughput, and it represents a critical decision point in the architecture of the custom DDTree inference engine being built from scratch.

The Context: Why This Measurement Matters

To understand why this message carries such weight, one must appreciate the broader project. The assistant and user are building a custom C/C++/CUDA inference engine (the "kdtree-engine") for Kimi K2.6, a 1-trillion-parameter Mixture-of-Experts model with 384 experts, using a speculative decoding technique called DFlash with DDTree (Draft Tree). The core idea of speculative decoding is to use a smaller, cheaper "drafter" model to propose multiple candidate next tokens in parallel (a "tree"), and then have the large target model verify all candidates in a single forward pass. The throughput win comes from the target model processing many tokens per step, amortizing its fixed costs.

The single most expensive operation in the target model's forward pass is the MoE GEMM—the matrix multiplication that routes each token through its top-8 activated experts out of 384 total. The K2.6 model uses INT4 quantization with Marlin kernels (moe_wna16_marlin_gemm via fused_marlin_moe), and understanding how this kernel behaves across different batch sizes (token counts) is essential to designing the verify step of the DDTree loop.

Prior to this message, the assistant had built and validated custom CUDA kernels for tree construction, tree-verify attention, and greedy tree acceptance. The MoE GEMM was the last major unknown—the component that would determine whether the speculative decoding engine could actually beat the autoregressive baseline. Without knowing the MoE's throughput curve, any design decisions about tree size, number of parallel streams, or batch configuration would be guesswork.

The Benchmark: What Was Measured

The assistant drove SGLang's real repack path for the fused_marlin_moe kernel at exact K2.6 dimensions: 384 experts (E=384), hidden dimension 7168, intermediate dimension 2048, top-8 routing, INT4 quantization with group size 32. The benchmark swept across token counts per layer (M) from 1 to 1024 on a single PRO 6000 GPU that was freed by briefly stopping the live DDTree service.

The resulting data is presented as a compact table in the message:

| M (tokens/layer) | ms/call | µs/token | regime | |---|---|---|---| | 1 (AR C=1) | 0.151 | 150.6 | ~8/384 experts active | | 9 (verify b8) | 1.18 | 131 | | | 90 (10×b8) | 5.19 | 57.7 | | | 256 | 7.02 | 27.4 | ~all experts active | | 512 | 7.01 | 13.7 | plateau | | 1024 | 7.02 | 6.85 | plateau |

This table is deceptively simple, but it reveals a profound architectural insight. The kernel exhibits two distinct regimes separated by a clear transition point around M=256 tokens per layer.

The Two Regimes: Weight-Streaming vs. Compute-Bound

The first regime, which the assistant labels weight-streaming-bound, occurs below approximately M=256 tokens per layer. In this region, the time per call grows roughly linearly with the number of active experts. At M=1 (a single token, autoregressive decode), only about 8 of the 384 experts are activated per layer (the top-8 routing), and the call takes 0.151 ms—150.6 µs per token. At M=9 (a small verify batch of 8 tokens), the time increases to 1.18 ms, or 131 µs/token. At M=90, it reaches 5.19 ms, or 57.7 µs/token.

What is happening here is that the GPU is spending most of its time streaming expert weights from memory into the compute units. Each expert's INT4 weights must be fetched, and the cost scales with how many experts are touched. Since the Marlin kernel processes tokens in groups and the expert routing is per-token, more tokens mean more distinct experts are activated, requiring more weight streaming.

The second regime, the compute-bound plateau, begins at M≈256 and extends to at least M=1024. In this region, the time per call is essentially flat at ~7.02 ms, while the per-token cost drops from 27.4 µs at M=256 to 6.85 µs at M=1024—a 4× improvement. What has happened is that all 384 experts' weights are now being streamed once per layer, and additional tokens are nearly free because the computation (the matrix multiplies) can be overlapped with the already-in-flight weight data. The kernel has become compute-bound rather than memory-bound.

This plateau behavior is the key insight. It means that once the verify batch size (M) exceeds approximately 256, the MoE cost per layer becomes a fixed overhead regardless of how many tokens are being verified. Adding more tokens to the batch does not increase the MoE latency. This is exactly the property that makes speculative decoding with large trees attractive: you can commit many tokens per step while paying the same MoE cost as for a much smaller batch.

The Architectural Implication: Refining the Engine Design

The assistant immediately translates this raw benchmark data into concrete design guidance. The verify batch size M is the product of two parameters: the number of parallel streams (how many independent draft sequences are verified simultaneously) and the sequence length per stream (the number of tokens committed in each stream, which is the draft budget plus the accepted tokens).

The sweet spot, the assistant concludes, is big trees (budget 32–64) combined with 8–16 streams. For example, a budget of 32 tokens per stream with 8 streams yields M≈264, just entering the plateau. With 16 streams, M≈528, solidly in the plateau. This combination allows the engine to maximize the number of committed tokens per step while keeping the MoE cost flat.

A critical caveat is noted: C=1 (a single stream, which is the autoregressive baseline) cannot reach the plateau. A single stream with any reasonable budget stays in the weight-streaming-bound regime. This means that the speculative decoding engine's throughput advantage over autoregressive decoding comes from two sources: (1) committing many tokens per sparse forward pass (the traditional speculative decoding win), and (2) reaching the MoE plateau where additional tokens cost nearly nothing (the new insight from this benchmark).

The assistant also documents a crucial limitation: these are single-GPU numbers running the full MoE (all 384 experts) on one card. In the production deployment with tensor parallelism across 8 GPUs (TP-8), each GPU holds only one-eighth of the expert weights. The per-GPU time would scale down by roughly a factor of 8, but the shape of the curve—the transition from weight-streaming to compute-bound at a threshold M—remains the same. The plateau threshold in absolute tokens might shift slightly, but the qualitative insight holds.

Where Phase 3 Stands: The Broader Picture

The message then steps back to summarize where the entire Phase 3 effort stands. Two dominant decode-cost components have now been characterized on the real target hardware:

  1. MoE GEMM: The reusable moe_wna16_marlin_gemm, with the plateau curve quantified above—this is the engine's throughput path.
  2. Tree build/accept/verify_attn: These custom CUDA kernels have been validated and benchmarked. The GPU tree builder achieves 6–474× speedup over SGLang's CPU implementation. The tree build and accept kernels are microsecond-cheap. The verify attention kernel is functional but acknowledged as naive—the production path will use FlashMLA for the long prefix attention and reserve the custom kernel only for the small tree tail.
  3. Baseline to beat: The live service achieves 138 tokens/second at C=1 (autoregressive) and 517 tokens/second at C=10 (speculative decoding with 10 streams). These are the numbers the native engine must exceed. The remaining work—integrating the Marlin MoE into the engine, loading the 548 GB of INT4 weights (by mmap'ing SGLang's repacked shards), sharding with NCCL for TP-8, swapping the verify attention prefix to FlashMLA, and porting the real DFlash drafter—is characterized as "integration plumbing" with no architectural risk. The decode loop, attention, tree kernels, cache logic, and now the MoE GEMM lever are all proven.

The Thinking Process: What This Message Reveals

This message is a window into the assistant's engineering reasoning. Several aspects of the thinking process are visible:

Prioritization: The assistant explicitly identifies the INT4 Marlin MoE GEMM as "the single component that determines whether the native engine can beat the baseline (and the biggest unknown)." This is a clear statement of risk assessment—the assistant is systematically attacking the highest-uncertainty items first.

Experimental discipline: The benchmark was run on a freed PRO 6000 GPU by briefly stopping the live service. The baseline was captured beforehand, and the service was restored and verified healthy afterward. This shows careful experimental hygiene—no data was lost, and the production service was not compromised.

Caveat documentation: The assistant is meticulous about documenting limitations. The TP-8 caveat ("single-GPU full-MoE numbers; TP-8 shards experts → per-GPU ≈ ÷8") is explicitly noted. The verify attention kernel is acknowledged as naive with a clear path to the production solution (FlashMLA). This intellectual honesty is crucial for maintaining trust in the analysis.

Translation from data to design: The assistant doesn't just report numbers; it immediately derives architectural guidance. The plateau observation is connected to specific recommendations about tree budgets (32–64) and stream counts (8–16). This is the mark of an engineer who understands that benchmarks are only valuable insofar as they inform design decisions.

Confidence calibration: The assistant's language shifts from exploratory ("the biggest unknown") to declarative ("The result proves the central thesis"). This reflects a genuine update in confidence based on experimental evidence. The uncertainty has been resolved.

Assumptions and Limitations

Several assumptions underlie this analysis, some explicit and some implicit:

The plateau threshold is model-dependent: The M≈256 transition point was measured on K2.6 with its specific dimensions (E=384, hidden 7168, inter 2048, top-8). A different model with different expert counts, hidden dimensions, or routing strategies would have a different threshold. The qualitative insight generalizes, but the quantitative numbers are specific to K2.6.

Single-GPU numbers don't directly translate to TP-8: The assistant acknowledges this but doesn't attempt to predict the TP-8 curve. The per-GPU cost will be lower because each GPU handles fewer experts, but the plateau behavior might shift. The true TP-8 characterization remains to be measured.

The benchmark measures raw kernel performance, not end-to-end throughput: The MoE GEMM is one component of a complex pipeline that includes attention, routing, normalization, and the drafter forward pass. The plateau analysis shows that MoE cost can be amortized, but the other components may introduce their own scaling behaviors.

The weight-streaming model assumes a particular memory hierarchy: The analysis implicitly assumes that expert weights are in GPU HBM and that the bottleneck is HBM bandwidth. This is correct for the PRO 6000 with its 96 GB of HBM per GPU, but the analysis would need revision for different memory architectures.

Input Knowledge Required

To fully understand this message, the reader needs:

  1. Knowledge of Mixture-of-Experts architectures: Understanding that K2.6 has 384 experts with top-8 routing, and that each token activates only a subset of experts.
  2. Knowledge of speculative decoding: Understanding the DDTree approach where a drafter proposes a tree of candidates and the target model verifies them in parallel, with the verify batch size being the product of streams and sequence length.
  3. Knowledge of GPU memory hierarchy: Understanding the distinction between weight-streaming-bound (memory bandwidth limited) and compute-bound (arithmetic limited) regimes, and why the transition occurs when all weights have been loaded once.
  4. Knowledge of INT4 quantization and Marlin kernels: Understanding that Marlin is an efficient GPU kernel for INT4 matrix multiplication that achieves near-ideal throughput by careful tiling and asynchronous data movement.
  5. Knowledge of tensor parallelism: Understanding that TP-8 means each GPU holds one-eighth of the expert weights, reducing per-GPU memory and computation.

Output Knowledge Created

This message creates several forms of knowledge:

  1. A quantified throughput curve for the K2.6 MoE GEMM on Blackwell PRO 6000: This is the primary empirical result—a precise measurement of how the kernel behaves across token counts.
  2. A design rule for the DDTree engine: The sweet spot of budget 32–64 with 8–16 streams is a concrete architectural decision derived from the data.
  3. A risk assessment of the remaining work: The assistant's characterization of the remaining integration as "no architectural risk" is a knowledge claim about project status—it tells the user that the hard unknowns have been resolved.
  4. A documented baseline to beat: The 138 tok/s (C=1) and 517 tok/s (C=10) numbers serve as the target for the native engine.
  5. A methodology for future characterization: The approach of isolating individual kernel components, benchmarking them on the target hardware, and deriving architectural guidance from the results is a reusable pattern.

Conclusion

Message 12023 is a masterclass in engineering communication. It takes a complex benchmark result—the throughput curve of an INT4 Marlin MoE GEMM on Blackwell GPUs—and distills it into a clear architectural insight: the MoE cost plateaus at M≈256, and the DDTree engine should be designed to operate in this plateau region with big trees and multiple streams. The message is simultaneously a progress report, a technical analysis, a design document, and a risk assessment. It demonstrates the value of systematic measurement and the power of translating raw data into actionable design guidance.

The assistant's disciplined approach—identifying the highest-uncertainty component, designing a focused experiment, running it cleanly on the target hardware, documenting caveats, and immediately deriving architectural implications—is a model for how to make progress on complex systems. The message leaves the reader with confidence that the remaining work, while substantial, is well-understood and architecturally sound. The critical unknowns have been resolved, and the path forward is clear.