The Final Wire: How a Single Edit Integrated INT4 Quantization into a Native DDTree Inference Engine

Subject message: [assistant] Now route the MoE/MLP GEMMs through linear(): [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/engine/model.cu Edit applied successfully.

At first glance, message <msg id=12074> appears almost trivial: a single sentence of intent followed by a confirmation that a file edit succeeded. There is no diff shown, no fanfare, no test output. Yet this message represents the culmination of a carefully orchestrated, multi-step integration effort spanning over a dozen preceding messages. It is the moment when a substantial body of preparatory work — a custom W4A16 group-quantized GEMM kernel, a numpy reference quantization pipeline, INT4 golden test bundles, and a new weight-loading and dispatch infrastructure — finally gets wired into the live forward pass of a native C/C++/CUDA speculative decoding engine. This article unpacks why that single edit mattered, what decisions led to it, and what it reveals about the architecture of the engine being built.

The Context: Building a Native DDTree Engine for Kimi K2.6

To understand message <msg id=12074>, one must understand the project it belongs to. The assistant has been building a native C/C++/CUDA inference engine for the Kimi K2.6 large language model, organized under a repository called kdtree-engine/. The engine implements a form of speculative decoding called "DDTree" (Draft-and-Diff Tree), where a smaller drafter model proposes candidate token sequences that a larger target model verifies in parallel. The assistant had already delivered three custom CUDA kernels — a GPU best-first tree builder, a tree-verify MLA-absorb attention kernel, and a greedy tree-accept kernel — and validated them against numpy references.

However, the engine's matrix computations were still running in FP32 via cuBLAS GEMMs. Real K2.6 inference uses INT4 W4A16 group-quantized weights (4-bit weights, 16-bit activations, grouped along the input dimension with shared scales). The FP32 path was functionally correct but bore no resemblance to the actual memory footprint or arithmetic of a production deployment. The assistant's next phase — Phase 3 of the engine roadmap — was to make the engine genuinely INT4-quantized.

The Multi-Step Integration Chain

Message <msg id=12074> is the final link in a chain of preparatory messages. Understanding it requires tracing that chain backward:

1. Kernel creation ([msg 12064], [msg 12065], [msg 12066]): The assistant wrote a W4A16 group-quantized GEMM kernel in ops.cuh and ops.cu. This kernel takes FP32 activations, INT4 packed weights (two 4-bit values per byte), and FP32 group scales, and produces FP32 outputs. It is the computational core of the quantization path.
2. Numpy reference quantization ([msg 12067]): The assistant added quantize_w4a16 and dequantization functions to model_ref.py, establishing the exact arithmetic that the CUDA kernel must match. This creates a spec: the engine's INT4 forward pass must produce the same outputs as the numpy dequantized forward pass, token-for-token.
3. INT4 bundle generator ([msg 12068]): A new script gen_model_int4.py was written to produce .kdtr bundle files containing quantized weights (.qw), scales (.sc), and golden tokens computed from the dequantized reference. This gives the engine a concrete test target.
4. Bundle generation and validation ([msg 12069]): The assistant ran the generator and confirmed that INT4 golden tokens diverge from FP32 golden tokens at token position 2 — proving quantization is actually being applied and that the INT4 golden is a distinct, meaningful reference point.
5. Weight loading and dispatch in model.h ([msg 12070]): The assistant added fields for quantized weight maps and scale maps to the model structure, along with a linear() function that checks whether a weight name exists in the quantized map and dispatches to either the W4A16 kernel or the FP32 cuBLAS path.
6. Implementation in model.cu ([msg 12071], [msg 12072], [msg 12073]): The assistant added the linear() method body, included <cstring> for string comparison, and prepared the infrastructure.
7. The final wire ([msg 12074]): Route the MoE/MLP GEMMs through linear(). Each step depends on the previous one. The kernel must exist before it can be called. The reference must exist before correctness can be checked. The loader must exist before weights can be dispatched. Message <msg id=12074> is the last domino: it replaces the direct gemm() calls in the MoE expert routing and MLP forward pass with calls to the new linear() dispatcher.

What the Edit Actually Changes

While the message does not show a diff, the architecture is clear from the surrounding context. The engine's layer_forward function (or equivalent) contains code paths for:

• MoE (Mixture of Experts): For each token, the router selects top-k experts. Each expert's gate, up, and down projections are matrix multiplications. Previously, these were FP32 cuBLAS gemm() calls. After the edit, they go through linear("expert.{i}.gate"), which checks if qw_expert_{i}_gate exists in the quantized weight map. If so, it calls the W4A16 kernel; otherwise, it falls back to FP32.
• MLP (dense feed-forward): The shared MLP layers (gate, up, down projections) follow the same pattern. Previously FP32 gemm; now routed through linear(). The key architectural insight is the fallback: linear() is not an either/or replacement. It is a unified dispatcher that silently selects the correct path based on what weights are available. This means the same engine binary can run FP32 bundles, INT4 bundles, or hybrid bundles where only some layers are quantized. This design choice — a single dispatch point rather than separate code paths — reduces maintenance burden and makes testing incremental: you can validate INT4 on a single layer while keeping everything else in FP32.

The Reasoning: Why This Approach Was Chosen

The assistant's reasoning, visible in the agent thinking blocks of preceding messages, reveals a deliberate architectural philosophy. The assistant considered several alternatives before settling on the linear() dispatcher pattern:

Option 1: Separate code paths. The assistant could have written a separate int4_forward() function alongside the existing fp32_forward(), with the caller choosing which to invoke. This would be simpler to implement but would duplicate control flow and make it harder to mix quantization levels across layers.

Option 2: Marlin extraction. The assistant considered extracting the marlin_moe_wna16 kernel from sgl_kernel (SGLang's kernel library) and linking it directly into the native engine. This would give production-grade INT4 performance but was deemed too risky due to Cutlass dependencies, JIT compilation machinery, and the difficulty of disentangling the raw CUDA kernel from the PyTorch op registration layer.

Option 3: Custom INT4 kernel with dispatcher. The assistant ultimately chose to write a simple, self-contained W4A16 GEMM kernel and wrap it in a linear() dispatcher. This approach is explicitly positioned as a correctness building block, not a performance peak. The plan documentation (from earlier in the session) notes that Marlin remains "the documented peak-perf drop-in" — the custom kernel proves the INT4 path works end-to-end, and Marlin can replace it later without changing the dispatch logic.

This is a classic engineering tradeoff: build the simplest thing that proves the architecture, then optimize. The linear() dispatcher is the interface that makes this substitution transparent.

Assumptions Embedded in the Design

The linear() dispatcher and the surrounding INT4 integration rest on several assumptions, some explicit and some implicit:

Assumption 1: Weight names are a reliable quantization key. The dispatcher checks for the existence of qw_{name} in the quantized weight map. This assumes a naming convention where quantized weights are stored with a .qw suffix and scales with .sc, and that the engine can infer which weights should be quantized from the bundle metadata alone. If a bundle contains a .qw entry for a weight that the engine expects in FP32, the dispatcher correctly uses the INT4 path. But if a bundle omits a .qw entry for a weight that was supposed to be quantized, the engine silently falls back to FP32 — which would produce wrong outputs if the golden tokens were computed with quantization. This is a silent correctness risk.

Assumption 2: Group size 32 matches K2.6. The quantization uses a group size of 32 (32 input columns share one scale per output column). This matches the K2.6 format, but the assistant does not verify this against actual K2.6 weights — the test bundles are synthetic. If real K2.6 weights use a different group size (e.g., 64 or 128), the kernel would need adjustment.

Assumption 3: FP32 fallback is safe. The dispatcher falls back to FP32 cuBLAS when no quantized weights are found. This is correct for testing but would be incorrect in production if the model is supposed to run entirely in INT4 — the fallback would consume 4× the memory and produce different numerical outputs. The assistant treats this as a testing convenience, but it could mask configuration errors.

Assumption 4: The W4A16 kernel's numerical accuracy is sufficient for greedy decoding. The INT4 golden tokens diverge from FP32 at position 2, meaning quantization introduces error. The engine is validated against the INT4 golden (not the FP32 golden), so it is self-consistent. But the assumption is that this quantization error does not change the semantic quality of the output — a claim that can only be verified with real model weights and downstream task evaluation.

Input Knowledge Required

To understand message <msg id=12074>, a reader needs knowledge of:

• The kdtree-engine project: Its architecture (C++/CUDA engine with numpy reference), its DDTree speculative decoding loop, and its three custom kernels.
• The INT4 quantization pipeline: How weights are packed (2 per byte), how group quantization works (shared scales along the input dimension), and how the .qw/.sc bundle format stores quantized weights.
• The MoE and MLP structure of Kimi K2.6: The router, expert gating, shared expert, SwiGLU activation, and the specific weight names that need quantization.
• The linear() dispatcher pattern: A function that abstracts away the quantization decision, allowing the forward pass to be written once and work with any weight format.
• The preceding edit chain: Messages <msg id=12064> through <msg id=12073> each contribute a piece of the puzzle; message <msg id=12074> is meaningless without them.

Output Knowledge Created

Message <msg id=12074> itself produces no visible output — no test results, no benchmark numbers, no golden token comparisons. Its output is purely structural: the engine's forward pass now routes through the quantization-aware dispatcher. The knowledge it creates is:

• The engine is now INT4-capable. Before this edit, the engine could only run FP32. After it, any bundle with .qw/.sc weights will use the W4A16 kernel automatically.
• The dispatcher architecture is validated. The linear() function is now exercised in the actual forward pass, not just tested in isolation. Any bugs in the dispatch logic (e.g., wrong weight name matching, incorrect tensor dimensions) would surface immediately.
• The integration chain is complete. All the preparatory work — kernel, reference, bundle generator, loader — is now connected to the execution path. The next step is to run the engine against the INT4 golden bundles and verify token-for-token correctness.

What Comes Next

The message's placement in the conversation is telling. It is followed by the assistant running tests to verify that the INT4 engine output matches the golden reference. Those tests are the payoff: if they pass, the INT4 path is proven correct. If they fail, the assistant must debug the dispatcher, the kernel, or the weight loading.

The broader narrative is one of incremental, test-driven infrastructure building. Each message adds a small piece, but the pieces are carefully ordered so that each can be validated before the next is added. Message <msg id=12074> is the moment when the pieces stop being independent and start being a system.

Conclusion

A single sentence — "Now route the MoE/MLP GEMMs through linear()" — belies the depth of engineering judgment it represents. The assistant chose a dispatcher pattern over separate code paths, a custom kernel over Marlin extraction, and a correctness-first approach over premature optimization. It built a chain of dependencies — kernel → reference → bundle → loader → dispatcher — and then connected the final link. The edit itself is trivial; the architecture it completes is not. Message <msg id=12074> is a reminder that in complex systems, the most important changes are often the ones that look the smallest.