The Kernel Hunt: Tracing the Marlin MoE Through SGLang's Source

Introduction

In the sprawling, multi-session effort to build a native C/C++/CUDA speculative decoding engine for the Kimi K2.6 model, there comes a moment where grand architectural plans meet the gritty reality of production code. Message [msg 12003] captures exactly that moment. The assistant has just received the user's directive to "go to the next phase" ([msg 12001]), and has spent the preceding message ([msg 12002]) reasoning through the enormous scope of Phase 3 — the integration of INT4 Marlin MoE GEMM operations, tensor parallelism across 8 GPUs, and the 548 GB weight load of K2.6. Now, in [msg 12003], the assistant receives the results of its first reconnaissance mission into SGLang's source code and must decide how to proceed.

This message is not about building anything. It is about finding something. The assistant is on a hunt for the actual kernel entry point — the precise function signature, weight format, and calling convention that will let it invoke the INT4 Marlin MoE grouped GEMM from its native engine. What it discovers reshapes its strategy and reveals the hidden complexity beneath SGLang's polished API surface.

The Context: A Pivot to Phase 3

To understand why this message exists, we must trace back through the session's arc. The assistant had just completed a monumental effort: building a complete native C/C++/CUDA DDTree inference engine from scratch, validating it against numpy golden references on actual 8× RTX PRO 6000 Blackwell hardware, and benchmarking it against SGLang's live service. The benchmarks were encouraging — the GPU tree builder achieved 6–474× speedup over SGLang's CPU implementation — but the honest analysis showed that at the target 1–10 stream concurrency, the real bottleneck was the INT4 Marlin MoE forward pass, not the tree build.

The user's command — "go to the next phase" — was unambiguous. The assistant responded in [msg 12002] with an extended reasoning trace that weighed the enormous scope of Phase 3 against what could realistically be achieved in a single session. The full end-to-end native K2.6 system requires INT4 MoE integration, weight loading and distribution across 8 GPUs via NCCL, flash-MLA for prefix attention, a Qwen3-style drafter, and support for K2.6-specific architectural details like YaRN RoPE and sigmoid routing. The assistant correctly identified this as "genuinely months of work" and chose to focus on the highest-leverage, highest-risk component: the INT4 Marlin MoE GEMM.

The plan was to explore what sgl_kernel exposes, how SGLang drives the Marlin MoE for K2.6, and then build a benchmark that validates correctness and measures throughput at realistic verify shapes. The bash command at the end of [msg 12002] was the first probe — checking torch.ops.sgl_kernel for marlin/moe ops and grepping SGLang's source for the relevant function signatures.

The Discovery: An Empty Namespace

The output that arrives in [msg 12003] is both revealing and slightly unsettling. The first probe — listing torch.ops.sgl_kernel for any ops containing "marlin", "moe", "wna16", or "gptq" — returns an empty list. The Marlin MoE kernel is simply not exposed as a top-level torch.ops operation. This is a significant finding: the kernel the assistant needs to call is hidden behind SGLang's own wrapper layer.

The second probe is more fruitful. Grepping SGLang's source reveals two critical pieces:

The actual entry point is fused_marlin_moe defined in sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py at line 28.
The underlying kernel is moe_wna16_marlin_gemm, imported from sglang.jit_kernel.moe_wna16_marlin.
There is also moe_sum_reduce imported from sgl_kernel. The K2.6 quantization config confirms the model uses compressed-tensors format with 4-bit weights, group size 32, and "pack-quantized" storage. This is the input knowledge the assistant now possesses: the kernel exists, it's a JIT kernel (likely Triton-based, given the jit_kernel module path), it's wrapped in a fused MoE function that handles routing, and the model's weights are stored in a specific compressed format. But the exact calling convention — the tensor shapes, dtypes, memory layouts, and the meaning of parameters like sorted_token, w13, w2, num_bits, and group — remains unknown.

The Reasoning: Two Paths Forward

The assistant's reasoning in [msg 12003] reveals a genuine fork in the road. The first option is to "reverse-engineer the kernel call" — to read the full fused_marlin_moe.py source, understand every parameter, trace the weight loading path, and reconstruct the exact calling convention. This is the meticulous, bottom-up approach. It guarantees understanding but is time-consuming and risks getting lost in SGLang's abstraction layers.

The second option is more pragmatic: "benchmark K2.6's actual MoE layer directly by either instrumenting the running model or building a standalone benchmark using the fused_marlin_moe API with K2.6's exact dimensions." This approach treats the existing SGLang service as a black-box oracle. Instead of understanding every detail of the kernel interface, the assistant would hook into the live model's MoE forward pass, extract real weights, and measure throughput at the exact token counts that matter for speculative decoding: 1 token (autoregressive), 9 tokens (budget-8 verify), 90 tokens (10 streams × 9), and 330 tokens (10 streams × 33).

The assistant chooses the second path, but not before taking one more look at the source. The bash command it issues reads the first 140 lines of fused_marlin_moe.py and greps for the key patterns: torch.ops, def fused_marlin_moe, moe_wna16_marlin, sgl_kernel, gptq_marlin, w13, w2, sorted_token, expert, num_bits, group. This is a targeted reconnaissance — the assistant is looking for the specific function signature and parameter names that will tell it how to call the kernel.

Assumptions and Their Implications

The assistant makes several assumptions in this message, most of them reasonable but worth examining.

Assumption 1: The Marlin MoE kernel is the throughput bottleneck. This is well-supported by the benchmarks from the previous chunk, which showed the 1T MoE forward dominating step time at 1–10 stream concurrency. The assumption is correct for the current configuration, but it implicitly assumes that the INT4 Marlin GEMM is the only significant cost in the forward pass. In reality, the fused MoE also includes routing logic, activation functions (SwiGLU), and potentially shared-expert computation — all of which add overhead that a pure GEMM benchmark would miss.

Assumption 2: The kernel can be called independently from SGLang's service. The assistant plans to "extract real INT4 weights from the loaded model and time the grouped GEMM." This assumes that the weights are accessible as PyTorch tensors that can be copied and used outside the SGLang process. In practice, SGLang may hold weights in custom formats (e.g., repacked Marlin shards with specific quantization metadata) that are not trivially separable from the model's forward method. The compressed-tensors format with group_size=32 and pack-quantized storage may require careful handling to extract usable weight matrices.

Assumption 3: The fused_marlin_moe API is stable and well-documented enough to use standalone. The assistant is essentially trying to use SGLang's internal function outside of SGLang's model loading and execution pipeline. This is a form of API boundary crossing that may encounter unexpected dependencies — missing tensor initializations, unset configuration objects, or version mismatches between the JIT-compiled kernel and the calling environment.

Assumption 4: The kernel's behavior at benchmark shapes generalizes to the full engine. The assistant plans to benchmark at specific token counts (1, 9, 90, 330) that correspond to verify batch shapes. This assumes that the kernel's performance characteristics at these discrete points are representative of its behavior in the full speculative decoding loop, where token counts vary dynamically and the MoE is called with different routing patterns (different sets of active experts per token).

None of these assumptions are unreasonable — they are the standard working hypotheses of an engineer exploring an unfamiliar codebase. But they are worth noting because they shape the assistant's strategy and will need to be validated as the work progresses.

The Output Knowledge Created

Despite being a research-oriented message, [msg 12003] creates significant output knowledge:

The Marlin MoE is not a flat torch.ops operation. It is wrapped in SGLang's fused_marlin_moe function, which itself calls a JIT kernel (moe_wna16_marlin_gemm). This means the assistant cannot simply call torch.ops.sgl_kernel.marlin_moe(...) — it must either use the fused wrapper or replicate its logic.
The kernel lives in sglang.jit_kernel.moe_wna16_marlin. This is a Triton JIT kernel, which means it is compiled at runtime. The assistant will need to ensure the Triton compiler is available and compatible with the CUDA 13 / sm_120 environment on the PRO 6000 GPUs.
K2.6 uses compressed-tensors quantization with 4-bit weights and group size 32. This is the specific format the native engine must support. The "pack-quantized" format implies that multiple 4-bit values are packed into larger integer types (e.g., int32), which affects how weights are loaded and how the GEMM is called.
The MoE has a post-GEMM reduction step (moe_sum_reduce). The import of moe_sum_reduce from sgl_kernel reveals that the fused MoE does not just compute per-expert GEMMs — it also reduces the outputs across experts, likely weighted by routing probabilities. The native engine must replicate this reduction.
The function signature hints at specific parameters. The grep output reveals parameter names like sorted_token, w13, w2, num_bits, and group. These hint at the kernel's internal structure: w13 and w2 likely refer to the gate and up projections (SwiGLU uses two weight matrices per expert), sorted_token suggests the kernel expects tokens sorted by expert assignment, and num_bits/group are quantization parameters.

The Thinking Process: A Window into Engineering Strategy

The reasoning section of [msg 12003] is particularly revealing because it shows the assistant weighing two fundamentally different engineering strategies. The first — "reverse-engineer the kernel call" — is the thorough, academic approach. It would produce complete understanding but at the cost of time and the risk of getting lost in SGLang's abstraction layers. The second — "benchmark K2.6's actual MoE layer directly" — is the pragmatic, empirical approach. It would produce the key throughput numbers quickly, even if the understanding of the kernel interface remains incomplete.

The assistant's choice to pursue the second path, but with one more look at the source code, is a classic engineering compromise: gather just enough information to enable the empirical approach, without committing to a full reverse-engineering effort. The bash command reads only the first 140 lines of fused_marlin_moe.py — enough to get the function signature and key parameter names, but not enough to understand the full weight loading and execution pipeline.

This reveals a deeper strategic assumption: that the assistant can treat SGLang's MoE as a black box for benchmarking purposes, and only later integrate the kernel directly into the native engine. The benchmark will tell the assistant whether the INT4 Marlin MoE can deliver the throughput needed to beat the 138/517 tok/s baseline. If the answer is yes, the assistant can invest the effort to properly integrate the kernel. If the answer is no, the entire Phase 3 approach may need to be reconsidered.

The Broader Significance

Message [msg 12003] sits at a critical inflection point in the session. The assistant has completed the custom DDTree kernels and the MVP native engine (Phases 0–2). The user has asked for Phase 3. But Phase 3 requires integrating with SGLang's production inference stack — a stack that is complex, layered, and not designed for external consumption.

The message is fundamentally about the boundary between two systems: the assistant's clean, from-scratch native engine and SGLang's battle-tested but opaque production code. The Marlin MoE kernel is the point where these two systems must meet. The assistant's investigation in [msg 12003] is the first step in bridging that gap — finding the right API, understanding the weight format, and measuring the throughput that will determine whether the native engine can compete with SGLang's live service.

The empty torch.ops.sgl_kernel list is a small but significant discovery. It tells the assistant that the integration path is not a simple API call but a deeper engagement with SGLang's internal architecture. The kernel exists, but it is wrapped, JIT-compiled, and coupled to SGLang's weight loading and routing logic. Extracting it for standalone use will require either replicating that logic or finding a way to call SGLang's fused function directly.

Conclusion

Message [msg 12003] is a reconnaissance message — the assistant gathering intelligence before committing to a course of action. It reveals that the Marlin MoE kernel is not a simple, callable operation but a deeply embedded component of SGLang's inference stack, wrapped in fused functions and JIT compilation. The assistant's response — to read just enough source code to enable an empirical benchmark, rather than fully reverse-engineering the kernel — is a strategic choice that prioritizes actionable throughput data over complete understanding.

The message also reveals the assistant's assumptions about the integration path: that the kernel can be called independently, that the weights are accessible, and that benchmark results at specific token counts will generalize to the full speculative decoding loop. These assumptions will be tested in the messages that follow, as the assistant moves from investigation to implementation.