The Weight of Precision: Building Native INT4 Quantization for a Custom Speculative Decoding Engine

Introduction

In the sprawling, multi-month endeavor to build a custom inference engine for the Kimi K2.6 language model, message [msg 12087] stands as a quiet but decisive milestone. It is not the flashy moment of a first deployment or the drama of a throughput regression diagnosed — those moments have their own articles. Instead, this message is the careful, methodical announcement that a foundational piece of the architecture has been locked down: the native engine now genuinely supports INT4 W4A16 group-quantized weights, matching the exact format that Kimi K2.6 ships in.

The message, written by the AI assistant to the user, reads as a status summary after a burst of implementation work. It opens with a bold claim — "the engine is now genuinely INT4 W4A16 group-quantized" — and proceeds to catalog what shipped, how it was validated, and what remains. On its surface, it is a progress report. But beneath that surface lies a dense web of engineering reasoning, trade-off decisions, validation philosophy, and architectural foresight that reveals how serious inference engine development actually works in practice.

This article examines that single message in depth: the context that made it necessary, the decisions embedded in its technical claims, the assumptions it makes about hardware and correctness, and the knowledge it both consumes and produces. To understand this message is to understand a critical juncture in the lifecycle of a custom ML inference engine — the moment when a prototype stops being a prototype and starts becoming a real system.

The Context That Produced This Message

To grasp why message [msg 12087] matters, one must understand the broader project it belongs to. The assistant and user have been building a native C/C++/CUDA inference engine called kdtree-engine — a from-scratch implementation designed to run speculative decoding (specifically, a "draft tree" or DDTree approach) for the Kimi K2.6 large language model on NVIDIA RTX PRO 6000 Blackwell GPUs.

This is not a trivial undertaking. Modern LLM inference engines like vLLM and SGLang are the products of years of engineering by teams of specialists. Building one from scratch means re-implementing every layer: the transformer architecture (MLA attention, MoE routing, RMSNorm, RoPE, SwiGLU), the KV cache management, the CUDA kernels for attention and matrix multiplication, and — crucially — the speculative decoding loop that makes DDTree interesting.

The project had already progressed through several phases before this message. The assistant had built and validated a full FP32 engine with MLA, MoE, and the DDTree loop. It had created numpy reference implementations to serve as golden test vectors. It had deployed the engine to an 8× PRO 6000 Blackwell machine (the "CT200 box") and run end-to-end validation. But FP32 inference is not production-ready for a 200B+ parameter model like Kimi K2.6 — the memory and compute requirements are prohibitive. The model ships with INT4 group-quantized weights, and any serious engine must support that format.

This message announces that the INT4 path is now complete and validated. It is the culmination of a focused implementation sprint that spanned roughly a dozen preceding messages ([msg 12067] through [msg 12086]), during which the assistant wrote quantization reference code, implemented the w4a16_gemm CUDA kernel, wired the engine's linear layer dispatch, generated INT4 test bundles, and validated everything on both a local RTX 5070 Ti and the target PRO 6000 hardware.

What Actually Shipped: Three Layers of Engineering

The message organizes its deliverables into three bullet points, but each of those bullets represents a significant engineering achievement.

The Kernel: w4a16_gemm

The first deliverable is the w4a16_gemm kernel in src/engine/ops.cu. This is a CUDA kernel that performs matrix multiplication where the activations are FP32 (the "W4A16" notation means weights are 4-bit, activations are 16-bit/FP32) and the weights are packed as INT4 values — 8 nibbles per int32, stored along the input dimension. The quantization is symmetric (value = nibble − 8, meaning the representable range is [-8, 7]), with per-group scales applied at group size 32.

The choice of group size 32 is not arbitrary. The message explicitly notes it "matching K2.6" — the actual Kimi K2.6 model uses this group size for its quantized weights. This is a critical design decision: the engine is not implementing a general-purpose quantization framework; it is implementing the exact format the target model uses. This means the engine can load real K2.6 weight checkpoints without any format conversion, which is essential for production deployment.

The kernel itself is described as "naive" — a word that appears both in this message and in the assistant's earlier reasoning. This is an honest admission. The current w4a16_gemm is a straightforward implementation that gets correctness right but is not optimized for peak throughput. The assistant has a clear plan: replace it with the Marlin kernel, which was separately benchmarked at ~67% of HBM peak bandwidth. But correctness must come before optimization, and this message marks the correctness milestone.

The Reference and Bundle Pipeline

The second deliverable is the Python-side infrastructure: model_ref.quantize_w4a16 (the quantization algorithm in the numpy reference implementation) and gen_model_int4.py (the script that produces INT4 test bundles). These bundles consist of .qw files (packed INT4 weights as int32 arrays) and .sc files (per-group scales as float32 arrays), along with an INT4-dequantized golden token sequence for validation.

This pipeline embodies a crucial engineering philosophy: test against a golden reference, not against yourself. The assistant generates the INT4 golden tokens using a numpy implementation of the dequantized forward pass, then compares the CUDA engine's output against those golden tokens. If they match token-for-token, the engine is correct. If they diverge, there is a bug. This is the same approach used for the FP32 path, and it provides an unambiguous correctness criterion.

The fact that the INT4 golden differs from the FP32 golden (diverging at token 2) is treated as a feature, not a bug. It confirms that quantization is actually being applied — that the INT4 path is not accidentally falling through to FP32. This is the kind of sanity check that separates careful engineering from cargo-cult coding.

The Engine Dispatch: Model::linear

The third deliverable is the integration into the engine itself. The Model::linear method now dispatches to either the W4A16 path or the FP32 cuBLAS path depending on whether the weights for a given layer are quantized. The MoE experts, shared expert, and dense MLP layers run through the INT4 path; the MLA projections and the language model head remain in FP32.

This selective dispatch is another design decision worth examining. Why not quantize everything? Because some layers are more sensitive to quantization error than others. The attention projections (MLA) and the final LM head operate on activations that are critical for output quality, and quantizing them to INT4 could introduce unacceptable degradation. The MoE and MLP layers, which dominate the compute budget, are more tolerant of quantization. This is consistent with how production systems like vLLM and TensorRT-LLM handle quantization — they typically keep attention in FP16/FP32 while quantizing the feed-forward layers.

The result is a single engine that transparently handles both FP32 and INT4 bundles. The same binary, the same code paths, just different weight files. This is elegant engineering: the complexity of quantization is hidden behind a clean abstraction.

Validation: The Gold Standard of Greedy-Exact Matching

The message reports validation results on the target hardware — PRO 6000 Blackwell GPUs with sm_120 compute capability. The key claim is:

Engine INT4 AR == INT4 golden and INT4 DDTree == INT4 golden, token-exact (greedy-exact) across both model configs.

This is a strong claim, and it is worth understanding what it means and why it matters.

"Greedy-exact" means that when the engine runs greedy decoding (always picking the most likely token), it produces exactly the same token sequence as the reference implementation. For autoregressive (AR) decoding, this is a straightforward comparison: run the engine step by step, compare each output token to the golden sequence. For DDTree speculative decoding, the claim is more interesting: the speculative decoding loop, which uses draft trees and verification, must produce exactly the same tokens as if it had run autoregressively. This is the fundamental correctness invariant of speculative decoding — it should be mathematically equivalent to AR decoding, just faster.

The message reports that this invariant holds for INT4, just as it held for FP32. The 8× reduction in target forwards (the number of full model forward passes) is preserved. The speculative decoding is working correctly on quantized weights.

The message also reports that "the INT4 golden differs from FP32 (diverges at token 2) — confirming real quantization, not a no-op." This is a subtle but important point. If the INT4 output were identical to the FP32 output, it would mean the quantization was not actually affecting the computation — perhaps the weights were being dequantized internally and the INT4 path was never really exercised. The divergence confirms that the quantization is real and that the engine is faithfully reproducing the quantized computation.

Assumptions Embedded in the Message

Every engineering message carries assumptions, and this one is no exception. Some are explicit, but others are implicit in the claims being made.

Hardware assumption: sm_120 Blackwell GPUs. The engine is built for and validated on NVIDIA Blackwell architecture (RTX PRO 6000, compute capability sm_120). The CUDA kernels are compiled for this target. The Marlin kernel that will eventually replace the naive w4a16_gemm was benchmarked on this same hardware. The message assumes that the target deployment platform is Blackwell, and that the engine will not need to run on older architectures.

Weight format assumption: K2.6's exact quantization scheme. The engine implements symmetric INT4 quantization with group size 32, value = nibble − 8. This matches Kimi K2.6's format. The message assumes this format is stable and will not change. If the model were updated with a different quantization scheme (e.g., asymmetric quantization, different group size, different packing order), the engine would need updates.

Correctness criterion: greedy-exact matching. The message treats token-for-token agreement with the golden reference as the definition of correctness. This is appropriate for greedy decoding, but it does not guarantee correctness for sampling-based decoding (temperature, top-k, top-p). The engine may need additional validation for non-greedy modes.

Performance assumption: Marlin is the right replacement. The message states that Marlin is the "documented drop-in" for the naive kernel, benchmarked at ~67% of HBM peak. This assumes that Marlin is compatible with the weight format, that it can be integrated into the engine's linear dispatch, and that it will deliver the expected speedup on Blackwell. These are reasonable assumptions given the assistant's earlier work, but they remain assumptions until validated.

The user's priorities. The message offers two paths forward: restart the DDTree service for production, or continue toward the next Phase-3 piece (Marlin integration or the K2.6 weight loader). This assumes the user has the bandwidth and interest to continue development rather than shifting to a different task.

Potential Mistakes and Incorrect Assumptions

While the message is well-reasoned, there are a few points worth scrutinizing.

The "naive" kernel may have hidden issues. The message describes the current w4a16_gemm as naive, but the validation only covers small test models (tiny and tiny2 configurations with hidden size 192, 6 layers, 6 heads). Real K2.6 dimensions are vastly larger (hidden size 7168, 61 layers, etc.). The naive kernel may have correctness issues at scale that don't manifest in the tiny test configurations — for example, register pressure, shared memory bank conflicts, or integer overflow in index calculations. The validation on small models is necessary but not sufficient.

The greedy-exact invariant may not hold for all inputs. The test models use random weights and a fixed prompt. Real model weights have structure (e.g., many weights near zero, specific activation patterns) that could expose numerical edge cases not triggered by random weights. The assistant's earlier work included validation against actual K2.6 weights in some contexts, but the INT4 validation appears to use only synthetic test models.

The divergence at token 2 is a weak sanity check. The message treats the INT4-vs-FP32 divergence at token 2 as confirmation of real quantization. This is correct as far as it goes, but it does not quantify the magnitude of the divergence. A single token difference could be the result of a tiny numerical difference that happens to cross a decision boundary, or it could indicate significant quality degradation. The message does not report perplexity or any other quality metric on the quantized model. For production deployment, one would want to know that the quantized model's output quality is acceptable, not just that it differs from FP32.

The assumption that Marlin is a drop-in may be optimistic. Marlin is a sophisticated kernel with specific requirements for weight layout, shared memory usage, and thread block geometry. Integrating it into the engine's existing dispatch mechanism may require nontrivial changes to the weight loading code, the linear layer abstraction, and the memory management. The message's characterization of Marlin as a "documented drop-in" reflects the plan, not the reality.

Knowledge Required and Knowledge Created

To fully understand this message, a reader needs:

The Thinking Process Visible in the Message

The message is a summary, not a reasoning trace, but it reveals the assistant's thinking through its structure and emphasis.

First, there is a clear prioritization of correctness over performance. The message leads with what was shipped and validated, not with benchmark numbers. The naive kernel is acknowledged as slow, but that is presented as a known limitation with a planned fix, not as a failure. This reflects a mature engineering approach: get the right answer first, then make it fast.

Second, there is a strong emphasis on reproducible validation. The message reports exact test outcomes ("token-exact," "diverges at token 2," "full suite green") rather than vague claims of "it works." This is the language of someone who has been burned by false positives and has learned to demand concrete evidence.

Third, there is architectural foresight in the dispatch design. By making Model::linear handle both FP32 and INT4 transparently, the assistant has created an engine that can load either format without code changes. This is not strictly necessary for the current phase, but it anticipates future needs — for example, running FP32 for debugging or running mixed-precision configurations.

Fourth, there is honest scoping in the "what remains" section. The message does not claim the INT4 work is complete; it explicitly identifies the Marlin kernel swap and the real weight layout as remaining work. This builds trust with the user and sets clear expectations.

Finally, there is situational awareness in the offer to restart the DDTree service. The assistant knows the user's priorities have shifted (the user asked to use the idle box for development rather than keeping the service running), and it offers to pivot back to production if needed. This reflects an understanding that the project has both development and operational dimensions.

Conclusion: A Milestone, Not a Destination

Message [msg 12087] is a milestone marker in a long engineering journey. It announces that the kdtree-engine can now handle the weight format that Kimi K2.6 actually ships in — a critical capability for any engine that aspires to run this model in production. The INT4 W4A16 path is validated, the greedy-exact invariant holds, and the test infrastructure is in place.

But the message is also honest about what remains. The naive kernel needs to be replaced with Marlin for acceptable performance. The real K2.6 weight layout needs to be supported. The engine needs TP-8 (tensor parallelism across 8 GPUs) to handle the full model. And the DDTree service, currently down, needs to be restarted when the user is ready.

What makes this message interesting as a subject of analysis is not the technical details alone — it is the window it provides into the engineering mindset. The careful separation of correctness from performance. The insistence on golden references and token-exact matching. The architectural decisions that anticipate future needs. The honest acknowledgment of what is still to be done.

In a world where AI development is often characterized by hype and premature claims of completion, this message is a refreshing example of measured, evidence-based progress reporting. It does not claim the engine is done. It claims that a specific, well-defined piece of work is complete and validated. And it lays out, with clarity and precision, what comes next.