The Quiet Wiring: How a Single Edit Confirmed INT4 Quantized Inference in a Native DDTree Engine
Message: [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/engine/model.cu Edit applied successfully.
At first glance, the subject message — message index 12075 in the conversation — is almost invisible. It is a one-line tool confirmation: the assistant reports that an edit to a file called model.cu was applied successfully. There is no reasoning block, no dramatic announcement, no fanfare. Yet this message is the keystone of a much larger architectural effort. It represents the moment when the native C++/CUDA DDTree inference engine for Kimi K2.6 crossed a critical threshold: the transition from a purely FP32 cuBLAS-based prototype to a mixed-precision engine capable of W4A16 group-quantized inference, the weight format used by the real production model.
To understand why this single edit matters, one must trace the chain of decisions and implementations that led to it — a chain that spans dozens of messages and multiple days of work across a cluster of 8× RTX PRO 6000 Blackwell GPUs.
The Context: Building a Native Inference Engine
The kdtree-engine project was born from a practical need: the Kimi K2.6 model, a large Mixture-of-Experts (MoE) language model with speculative decoding via a DDTree (Dynamic Draft Tree) drafter, was being deployed on a cluster of 8× RTX PRO 6000 GPUs. The SGLang inference framework provided a baseline, but the team wanted a native, optimized engine that could integrate custom CUDA kernels for the DDTree speculative decoding loop — kernels that the assistant had already written and validated in earlier phases of the project.
By the time we reach message 12075, the engine already had:
- A complete FP32 forward pass using cuBLAS GEMMs for all linear layers
- Three custom CUDA kernels: a GPU best-first tree builder, a tree-verify MLA-absorb attention kernel, and a greedy tree-accept kernel
- A full DDTree speculative decode loop validated against numpy golden references
- A binary container format (KDTR) for sharing test data between Python and C++ But the engine was running entirely in FP32. The real Kimi K2.6 model uses INT4 Marlin-format weights — W4A16 group quantization with a group size of 32. Running FP32 meant the engine could not load real model weights, could not measure realistic memory bandwidth utilization, and could not serve as a production path. Phase 3 of the project was to bridge this gap.
The Decision Fork: Marlin Extraction vs. Custom Kernel
Message 12063 reveals the assistant's extensive deliberation about how to add INT4 support. Three paths were considered:
- Extract the Marlin MoE kernel from sgl_kernel — the ideal path for peak performance, but fraught with risk. The Marlin kernel is deeply entangled with the Cutlass template library and PyTorch's custom op registration machinery. Extracting a clean, callable C++ launcher would require significant reverse-engineering and build-system surgery.
- Write a custom INT4 W4A16 GEMM kernel — self-contained, testable, and under the assistant's full control. It would not match Marlin's performance, but it would prove the INT4 path end-to-end in the native engine. Performance could be improved later by swapping in Marlin as a drop-in replacement.
- Build a hybrid Python benchmark — load real K2.6 weights through Python, call the Marlin op for MoE, and use ctypes to invoke the DDTree kernels. This would measure realistic performance but would not advance the native engine itself. The assistant chose path 2: write a custom W4A16 group-quantized GEMM kernel directly in the engine. The reasoning was pragmatic and disciplined: "Marlin remains the documented peak-perf drop-in; this is the correctness/format building block." The goal was not to achieve peak throughput immediately, but to prove that the engine could handle real INT4 quantization, load quantized weights, and produce correct tokens — all within the native C++/CUDA framework.
The Implementation Chain
The work unfolded across several messages, each building on the last:
- Message 12064–12066: The W4A16 GEMM kernel was written into
ops.cuhandops.cu. This kernel takes FP32 activations, INT4 packed weights (two 4-bit values per byte), and FP32 group scales, and produces FP32 outputs. The kernel had to match the dequantization arithmetic exactly with the numpy reference to ensure bit-exact agreement on argmax. - Message 12067: The numpy reference model (
model_ref.py) was extended withquantize_w4a16()and dequantization functions, establishing the mathematical specification that the CUDA kernel must match. - Message 12068–12069: The INT4 bundle generator (
gen_model_int4.py) was written and executed. It quantized 84–90 weights per model configuration with group size 32, produced INT4 golden tokens, and saved.qw(packed int4 weights) and.sc(float32 scales) files alongside the existing KDTR bundles. The output confirmed that INT4 quantization introduced real numerical divergence from FP32 — the golden tokens diverged at position 2 — establishing the INT4 golden as a separate, well-defined reference target. - Message 12070–12071: The engine's model loader (
model.handmodel.cu) was extended with fields for quantized weight maps and scale maps, plus alinear()dispatcher function. The loader was modified to recognize.qwand.scfiles and load them into GPU memory alongside the FP32 weights. - Message 12072–12073: A missing
<cstring>header was added tomodel.cufor string comparison functions, and thelinear()method was implemented to check whether quantized weights exist for a given layer name and dispatch to either the W4A16 kernel or the FP32 cuBLAS fallback. - Message 12074: The assistant announced the final step: "Now route the MoE/MLP GEMMs through
linear():" — replacing the direct cuBLAScublasSgemmcalls in the MoE and MLP forward paths with calls to the new dispatcher. - Message 12075 (the subject): The edit was applied successfully. The routing was complete.
What the Edit Actually Did
The edit in message 12075 modified the layer_forward function in model.cu to replace direct cuBLAS GEMM calls — which had been hard-coded to operate on FP32 weights — with calls to the new linear() method. The linear() method, in turn, checks a map of quantized weight names. If the requested layer has a .qw entry, it calls the W4A16 kernel with the packed int4 weights and float32 scales. If not, it falls back to the original FP32 cuBLAS path.
This is a textbook example of the Strategy pattern applied at the GPU kernel level: the caller (layer_forward) does not need to know whether a weight is quantized. It simply asks linear() to perform y = Wx, and the dispatcher handles the rest. The MoE experts, the shared expert, the gate projection, and the MLP layers all route through this single point, meaning a single edit touched every quantized layer in the model.
The Validation
The proof that the edit worked came in the very next message (msg 12076). The assistant built the engine locally on a 5070 Ti GPU and ran both the autoregressive (AR) and DDTree tests against the INT4 tiny model:
golden AR : 499 128 58 76 462 243 474 353 74 410 401 211 211 474 150 386 ...
PASS model_ar tokens=24/24 exact max_abs_logit_diff=8.225e-06
ddtree : 499 128 58 76 462 243 474 353 74 410 401 211 211 474 150 386 ...
DDTree stats: steps=3 committed=25 avg_accept=8.00 tokens/step (block_size=8)
PASS model_ddtree ...
Both the AR forward pass and the full DDTree speculative decode loop produced tokens identical to the INT4 golden reference, with a maximum absolute logit difference of 8.225e-06 — essentially bit-exact at FP32 precision. The INT4 path was validated end-to-end.
Assumptions and Their Verification
The assistant made several assumptions during this work, all of which were validated by the test results:
- The W4A16 kernel's dequantization arithmetic matches numpy exactly. This was critical because any discrepancy in the dequantization formula (e.g., how scales are applied, how nibbles are unpacked) would produce different floating-point values and potentially different argmax decisions. The 8.225e-06 max logit difference confirmed near-bit-exact agreement.
- Group size 32 provides sufficient fidelity for greedy decoding. The INT4 golden tokens diverged from FP32 at position 2, but the engine's output matched the INT4 golden exactly. The quantization error, while real, was consistent and predictable.
- The
linear()dispatch overhead is negligible. The dispatcher performs a string map lookup per layer call. In the context of a GEMM operation that processes millions of floating-point operations, this overhead is essentially zero. - The MoE and MLP layers are the only layers that need quantization. The assistant explicitly scoped the INT4 work to the MoE experts, shared expert, gate projection, and MLP SwiGLU weights — the layers where Marlin quantization would apply in the real K2.6 model. MLA projections and embeddings remained in FP32, consistent with the production model's approach.
The Broader Significance
Message 12075 is, in a sense, the most boring and the most important message in this sequence. It is boring because it is just a tool confirmation — no reasoning, no explanation, no drama. It is important because it marks the moment when the native engine stopped being a FP32 toy and became a mixed-precision inference engine capable of loading and running quantized weights.
The edit itself is small — a few lines of C++ that replace direct GEMM calls with a dispatcher. But those few lines represent the culmination of:
- A deliberate architectural decision to write a custom INT4 kernel rather than extract Marlin
- The implementation and validation of that kernel against a numpy reference
- The extension of the model loader to handle quantized weight files
- The creation of INT4 golden reference bundles
- The careful matching of dequantization arithmetic between Python and CUDA Without this edit, the engine could not load real K2.6 weights. With it, the path is clear: load the real Marlin-quantized weights, measure the throughput, and then optimize by swapping in the Marlin kernel itself. The edit is the hinge between correctness and performance.
Input Knowledge Required
To understand this message, one must know:
- The kdtree-engine project structure:
model.handmodel.cudefine the Model class with weight maps, layer forward passes, and the DDTree speculative decode loop - The W4A16 kernel in
ops.cu: a custom GEMM that unpacks INT4 nibbles, applies group scales, and accumulates in FP32 - The KDTR bundle format:
.qwfor packed int4 weights,.scfor float32 scales, and the convention that quantized weights are stored alongside FP32 weights - The Kimi K2.6 model architecture: MoE with shared expert, SwiGLU MLP, and the specific layers that use INT4 quantization
Output Knowledge Created
This message produced:
- A validated INT4 inference path in the native DDTree engine, proven correct against golden references on two model configurations
- The architectural pattern for dispatching between quantized and full-precision paths, which can be extended to other quantization formats (e.g., FP8, INT8)
- A clear separation between correctness (custom W4A16 kernel) and performance (future Marlin swap-in), documented in the project plans
- A working demonstration that the engine's DDTree speculative decode loop functions correctly with quantized weights, preserving the greedy-exact invariant
Conclusion
Message 12075 is a testament to the fact that in complex engineering work, the most consequential moments are often the quietest. A single edit confirmation — "Edit applied successfully" — marked the completion of a multi-step integration that transformed the kdtree-engine from a FP32 prototype into a mixed-precision inference engine. The edit itself was small, but the thinking, the kernel implementation, the reference validation, and the careful architectural decisions that enabled it were anything but. The engine could now run INT4 quantized models, and the path to running the real Kimi K2.6 at scale was open.