The Pivot to Real Hardware: Diagnosing Throughput and Building the INT4 MoE Benchmark for Kimi K2.6

Introduction

In the long arc of building a native C/C++/CUDA speculative decoding engine for the Kimi K2.6 language model, there comes a moment when theory must meet hardware. Message 12006 of this opencode conversation captures exactly that transition. After dozens of rounds spent building custom CUDA kernels, validating bit-exact tree operations, and deploying a live SGLang DDTree service on an 8× RTX PRO 6000 Blackwell machine (codename CT200), the assistant faces a critical fork: it must decide how to make meaningful progress on Phase 3—the integration of the INT4 Marlin MoE (Mixture-of-Experts) GEMM that will determine whether the native engine can actually outperform the existing SGLang baseline.

This message is not a flashy milestone. It is a quiet, technical deliberation—a reasoning trace that reveals how a skilled engineer navigates the treacherous waters between what is theoretically possible and what is practically achievable in a single session. The assistant writes a benchmark script, but more importantly, it performs a rigorous mental calculation of memory budgets, weight formats, and kernel signatures that determines the entire trajectory of the next phase of development.

The Context: Where We Stand

To understand message 12006, one must first understand the landscape. The project has been building a native DDTree inference engine for Kimi K2.6, a 548 GB Mixture-of-Experts model with 384 experts, INT4 quantization via the compressed-tensors format, and a unique MLA (Multi-head Latent Attention) architecture. The engine lives in a repository called kdtree-engine/ and is written in C/C++/CUDA, organized into phases:

The Reasoning Process: A Mental Architecture Walkthrough

Message 12006 opens with the assistant's internal reasoning block, labeled "Agent Reasoning." This is not a polished document; it is a stream of consciousness that reveals how the assistant works through a complex technical problem.

Step 1: Understanding the Weight Format

The assistant begins by reasoning about the Marlin quantization weight format. The key insight is that Marlin uses a repacked layout where the weight matrices are transformed from a logical GPTQ format (INT4 values packed 8 per int32) into a hardware-friendly layout optimized for shared memory access patterns.

For the w1 weight matrix (the first GEMM in the MoE feed-forward network), the shape is:

w1: [E, K//16, 2N * (num_bits // 2)]

Where E is the number of experts (384 for K2.6), K is the hidden dimension (7168), and N relates to the intermediate dimension. The //16 factor comes from Marlin's tiling strategy—it processes 16 rows of the weight matrix at a time. The 2N * (num_bits // 2) accounts for the fact that w1 in a SwiGLU MLP actually stores two weight matrices (gate and up projections) fused together, with each INT4 value packed into half a byte.

For w2 (the down projection), the shape is reversed:

w2: [E, N//16, K * (num_bits // 2)]

The scale tensors follow their own grouping patterns based on group_size=32. Each group of 32 values along the K dimension shares a scale factor.

This reasoning is critical because getting the shapes wrong means the kernel will crash or silently produce garbage. The assistant cannot simply guess—it must verify against SGLang's actual implementation.

Step 2: The Repack Pipeline

The assistant realizes it cannot construct Marlin weights from scratch. The format is too fiddly, with permutation tensors, group indices, and sorting indices that must be exactly right. Instead, it plans to use SGLang's own gptq_marlin_moe_repack function, which takes GPTQ-formatted INT4 weights and transforms them into Marlin layout.

The repack function signature, discovered in earlier messages, is:

def gptq_marlin_moe_repack(
    b_q_weight: torch.Tensor,
    perm: torch.Tensor,
    size_k: int,
    size_n: int,
    num_bits: int,
) -> torch.Tensor:

For the no-activation-order case (which K2.6 uses), the permutation tensor is empty—the repack function treats this as "no permutation applied."

The assistant notes that for a timing benchmark, the actual weight values don't matter—the kernel does the same amount of work regardless of the data. What matters is that the shapes are correct and the kernel runs without errors. As long as the output has the right shape and contains finite values, the timing is representative.

Step 3: The Memory Calculation

This is where the assistant's reasoning becomes most concrete—and most revealing. It calculates the memory required for the full E=384 benchmark:

Step 4: The Decision to Write the Benchmark Script

The final action in the message is writing the benchmark script to /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/bench_marlin_moe.py. This script will:

  1. Construct INT4 GPTQ-packed weights for all 384 experts at K2.6 dimensions
  2. Repack them via SGLang's gptq_marlin_moe_repack
  3. Call the fused_marlin_moe kernel across a sweep of token counts (1, 9, 90, 330, etc.)
  4. Report throughput in tokens per second The script is the tangible output of this message—the artifact that will answer whether the native engine's approach (batching multiple token sequences through the INT4 MoE GEMM) can beat the SGLang baseline.

Assumptions and Potential Pitfalls

The assistant makes several assumptions that deserve scrutiny:

Assumption 1: Weight Values Don't Matter for Timing

The assistant asserts that "the kernel does the same amount of work regardless of weight values." This is mostly true for GEMM operations—the computation is data-independent. However, there are edge cases:

Assumption 2: The Repack Produces Valid Weights

The assistant assumes that calling gptq_marlin_moe_repack with random GPTQ-formatted weights will produce valid Marlin weights that the fused_marlin_moe kernel will accept. This is a reasonable assumption—the repack function is designed to transform any valid GPTQ weights into Marlin format. However, if there are any invariants in the GPTQ format that random data violates (e.g., specific bit patterns in the permutation tensor), the repack might produce garbage or crash.

The assistant hedges this by noting that "the asserts guard layout"—the kernel's internal assertions will catch shape mismatches. But asserts won't catch semantically invalid weights that happen to have the right shape.

Assumption 3: Stopping the Service is Safe

The assistant decides to "briefly stop the service" to free GPU memory. This assumes:

Assumption 4: The Benchmark Answers the Right Question

The benchmark measures raw INT4 MoE GEMM throughput at various token batch sizes. The assistant assumes this directly answers whether "the native engine can beat the baseline." In reality, the full system has many other components—attention, routing, KV cache management, NCCL communication for tensor parallelism—that contribute to end-to-end latency. A fast MoE GEMM is necessary but not sufficient for overall performance.

Input Knowledge Required

To fully understand message 12006, the reader needs knowledge of:

  1. Mixture-of-Experts (MoE) architecture: How K2.6 uses 384 experts with top-8 routing, and the w1/w2 weight matrix structure in a SwiGLU MLP.
  2. INT4 quantization and Marlin format: How 4-bit weights are packed (8 values per int32), the group-wise scaling with group_size=32, and Marlin's specific tiling layout (processing 16 rows at a time).
  3. GPTQ vs Marlin repack: The transformation from logical GPTQ format (where weights are stored in a straightforward packed layout) to Marlin's optimized format (which rearranges data for efficient shared memory access).
  4. CUDA GPU memory hierarchy: Understanding why 96 GB GPUs might have only 9–10 GB free when a large model is loaded, and why peak memory during repack (~16.8 GB) matters.
  5. The broader project context: Phases 0–2 of the native engine, the live SGLang service baseline (138–517 tok/s), and the user's directive to proceed to Phase 3.
  6. Speculative decoding: How DDTree works with a budget (C=1, C=4, etc.) and why the verify batch size (number of tokens verified per step) determines the throughput leverage.

Output Knowledge Created

Message 12006 produces several forms of knowledge:

Explicit Output

The benchmark script bench_marlin_moe.py is the primary artifact. It will produce throughput numbers for the INT4 MoE GEMM at K2.6 shapes across token batch sizes, directly answering whether batching multiple token sequences through the MoE is compute-bound enough to yield a throughput gain.

Implicit Output

The reasoning trace itself is valuable knowledge. It documents:

Knowledge for Future Decisions

The benchmark results (produced in subsequent messages) will inform:

The Broader Significance

Message 12006 exemplifies a pattern that recurs throughout engineering: the moment when abstract reasoning meets concrete constraints. The assistant has spent hours reasoning about kernel signatures, weight formats, and repack pipelines. But the final decision—stop the service, run the benchmark at full scale—is driven by a simple memory calculation: 16.8 GB peak exceeds 9 GB available, so the service must pause.

This is the essence of systems engineering. No amount of theoretical analysis can substitute for running on real hardware with real constraints. The assistant could have spent more time trying to construct a reduced-scale benchmark that fits within the available memory, but that would introduce uncertainty about whether the results scale to the full model. Instead, it chooses the more disruptive but more definitive path: full-scale measurement on the real hardware.

The message also reveals the assistant's deep understanding of the quantization pipeline. It doesn't just copy shapes from documentation—it reasons through why w1 has shape [E, K//16, 2N*(num_bits//2)] and w2 has shape [E, N//16, K*(num_bits//2)]. It understands that the //16 comes from Marlin's tiling, the 2N comes from SwiGLU's fused gate/up projections, and the num_bits//2 comes from packing 4-bit values into bytes. This level of understanding is necessary because the documentation is often incomplete or scattered across multiple source files.

Conclusion

Message 12006 is a turning point in the project. It marks the transition from analysis to implementation, from theory to hardware. The assistant has diagnosed the throughput bottleneck (INT4 MoE GEMM), identified the key measurement (throughput vs. token batch size), calculated the resource requirements (16.8 GB peak), and made the operational decision (stop the service, run at full scale).

The benchmark script written in this message will produce the numbers that determine whether the entire Phase 3 approach is viable. If the INT4 MoE GEMM shows strong throughput scaling with token batch size, the native engine can beat the SGLang baseline by batching verify tokens more efficiently. If not, the project needs a different strategy—perhaps focusing on the attention kernel or the drafter quality instead.

In the broader narrative of this opencode session, message 12006 is where the assistant stops planning and starts building. It is the moment when the abstract architecture diagrams become concrete code running on real GPUs. And it is a masterclass in how to think about performance optimization: identify the bottleneck, measure it at realistic scales, and make decisions based on data rather than intuition.