The Quantization Bridge: Wiring INT4 Precision into a Native Inference Engine

Introduction

In the sprawling, multi-session effort to build a native C/C++/CUDA speculative decoding engine for the Kimi K2.6 language model, message [msg 12070] represents a pivotal moment of integration. It is the point at which the assistant, having just validated that its INT4 quantization pipeline produces numerically distinct (and correct) outputs from the FP32 baseline, commits to wiring the engine's model loading and linear dispatch logic to consume those quantized weights. This message is not flashy—it contains no new kernel launches, no benchmark results, no dramatic performance revelations. It is an infrastructure message, a "glue" message, and its importance lies precisely in that unglamorous role: without it, the custom W4A16 GEMM kernel written in the preceding messages would remain an orphaned piece of code, never called, never tested end-to-end, never proving that the engine can actually run on the 4-bit quantized weights that real K2.6 deployment demands.

To understand why this message was written, one must appreciate the architecture being constructed. The kdtree-engine repository, built over the course of segment 65, is a from-scratch native inference engine for DeepSeekV3/Kimi-style transformer architectures. It implements Multi-head Latent Attention (MLA), Mixture-of-Experts (MoE) routing with shared experts, RMSNorm, NeoX rotary position embeddings, SwiGLU activations, and—most critically—the full DFlash speculative decoding loop with three custom CUDA kernels: a GPU best-first tree builder, a tree-verify MLA-absorb attention kernel, and a greedy tree-accept kernel. The engine had been validated at FP32 precision using cuBLAS GEMMs as placeholders for the eventual INT4 Marlin kernels. The next phase—Phase 3—was to make the engine genuinely INT4 W4A16 group-quantized, matching the weight format that K2.6 uses in production.

The Context: A Chain of Dependencies

The message sits at the end of a carefully constructed dependency chain. In [msg 12063], the assistant laid out the plan: write a self-contained INT4 W4A16 group-quantized GEMM kernel, add quantization and dequantization functions to the numpy reference model, generate INT4 golden bundles, then wire the engine to load and dispatch through the new quantized path. The assistant explicitly chose to write a custom INT4 kernel rather than extract Marlin from sgl_kernel, reasoning that the extraction risk (Cutlass dependencies, build integration complexity, torch runtime entanglement) was too high for a single session. The pragmatic choice was correctness first, performance later—a decision that shapes everything that follows.

Messages [msg 12064] through [msg 12067] implemented the W4A16 GEMM kernel in ops.cuh and ops.cu. Message [msg 12067] added the numpy quantization reference functions (quantize_w4a16, dequantize_w4a16, INT4 forward pass) to model_ref.py. Message [msg 12068] created gen_model_int4.py, the INT4 bundle generator that produces .kdtr files containing both the original FP32 weights and the quantized weights stored as .qw (packed int4) and .sc (float32 scales) files, along with golden tokens computed using the dequantized weights.

Then came the critical validation in [msg 12069]. The assistant ran the INT4 bundle generator, which produced two quantized model bundles (model_int4_tiny.kdtr and model_int4_tiny2.kdtr), and then compared the INT4 golden tokens against the FP32 golden tokens:

fp32 : [499, 128, 83, 210, 49, 128, 16, 269, 203, 224, 499, 170, 117, 203...
int4 : [499, 128, 58, 76, 462, 243, 474, 353]
diverge at: 2

The divergence at token index 2 confirmed that quantization was actually being applied—the INT4 forward pass produces different outputs from FP32, as expected due to the information loss from 4-bit quantization. This was not a bug; it was the desired behavior. The INT4 golden output is now a separate, well-defined reference point. The engine's job is to reproduce that INT4 golden exactly, not the FP32 golden.

The Message Itself: Reasoning and Action

Message [msg 12070] opens with the assistant's reasoning block, which reveals the thinking process behind the subsequent edit:

INT4 quantization introduces real numerical error, so the INT4 golden output is now a separate, well-defined reference point. The engine needs to reproduce that INT4 golden exactly to verify the dequantization and matrix multiplication are correct, and the divergence from FP32 confirms quantization is actually being applied. Now I need to update the model code to load the quantized weights from the .qw file as integers and the scales from .sc as floats.

This reasoning encapsulates several key insights. First, the assistant recognizes that the INT4 golden is the authoritative reference—the engine must match this, not the FP32 golden. This is a subtle but important point: the INT4 forward pass introduces quantization error, but that error is deterministic and reproducible. The engine's W4A16 kernel must produce bit-identical (or at least argmax-identical) results to the numpy reference's INT4 forward pass. Second, the divergence from FP32 is not a problem to be fixed but a feature to be confirmed—it proves the quantization pipeline is actually quantizing. Third, the assistant correctly identifies the next concrete steps: load .qw as integers, load .sc as floats, and implement a dispatch mechanism.

The reasoning then describes the implementation plan in more detail:

I'm adding fields to store the quantized weights and scales in separate maps, plus a group size parameter from the metadata. The load function needs to distinguish between quantized weight files (ending in .qw), scale files (.sc), and regular float weights, allocating GPU memory and copying each type appropriately. Then I'm implementing a linear layer dispatcher that checks whether a weight is quantized and calls the appropriate INT4 GEMM kernel if so.

This reveals the assistant's mental model of the engine's architecture. The model loader currently reads FP32 weights from the KDTR bundle and copies them to GPU memory. To support quantized weights, it needs to:

  1. Add new storage fields: std::unordered_map<std::string, float*> qweight for the packed int4 weights, std::unordered_map<std::string, float*> scales for the group scales, and an int group_size parameter.
  2. Extend the loading logic to check file extensions: .qw files contain int4-packed integers, .sc files contain float32 scales, and unadorned weight names contain float32 tensors.
  3. Add a linear() dispatcher function that checks whether a weight name exists in the quantized maps and, if so, calls the W4A16 GEMM kernel instead of the FP32 cuBLAS GEMM. The assistant then executes the edit:
INT4 quantization is real (diverges from FP32 at token 2, as expected) and the INT4 golden is now a well-defined target. Now wire the engine to load .qw/.sc and dispatch a linear() that uses W4A16 for the MoE/MLP weights: [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/src/engine/model.h Edit applied successfully.

The edit is to model.h, the engine's model header file. This is the central file that defines the model structure, weight storage, and forward pass logic. By adding the quantized weight loading and dispatch here, the assistant is integrating the W4A16 path into the engine's core computation graph.

Design Decisions and Trade-offs

Several design decisions are implicit in this message. The choice to store quantized weights in separate maps (.qw and .sc) rather than modifying the existing weight map or using a unified weight structure reflects a pragmatic separation of concerns. The FP32 path remains untouched; the INT4 path is additive. This minimizes the risk of breaking the existing, validated FP32 engine while adding the new capability.

The decision to dispatch at the linear() level rather than at a higher layer (e.g., at the MoE routing or MLP level) is architecturally significant. It means that any linear layer in the model can potentially use quantized weights, but the dispatch is invisible to the calling code—the MoE router, the SwiGLU activation, the shared expert projection all just call linear() and get the appropriate precision automatically. This is a clean abstraction that keeps the forward pass logic readable.

The assumption that the INT4 golden must be reproduced exactly is worth examining. In practice, floating-point accumulation order between numpy (which uses Python's float32, potentially with different associativity) and CUDA (which uses IEEE 754 float32 with its own instruction-level parallelism) can introduce tiny differences even when the same arithmetic is specified. The assistant is aware of this—in earlier reasoning ([msg 12063]), it noted that the kernel must match the dequantization arithmetic "exactly between numpy and CUDA—using identical fp32 accumulation so the results align closely enough for argmax to work despite tiny floating-point differences." The goal is not bit-exact match of every floating-point value, but argmax-identical token selection, which is a slightly looser constraint that accommodates the inherent non-associativity of floating-point arithmetic.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

Output Knowledge Created

This message produces:

  1. An edited model.h that adds quantized weight storage fields, loading logic for .qw/.sc files, and a linear() dispatcher.
  2. A validated design pattern for integrating quantized weights into the engine: separate maps, file-extension-based loading, transparent dispatch.
  3. A clear contract between the Python reference and the C++ engine: the INT4 golden is the authoritative target, and the engine must reproduce it.

The Thinking Process

The most revealing part of this message is the reasoning block, which shows the assistant's thought process in real time. The assistant begins by reflecting on the implications of the INT4/FP32 divergence: "INT4 quantization introduces real numerical error, so the INT4 golden output is now a separate, well-defined reference point." This is not merely stating a fact; it is actively constructing the validation strategy. The assistant is deciding that the INT4 golden, not the FP32 golden, is the ground truth for the quantized engine.

The next sentence—"The engine needs to reproduce that INT4 golden exactly to verify the dequantization and matrix multiplication are correct"—reveals the assistant's mental model of correctness. The engine is not being validated against "the correct answer" in some absolute sense; it is being validated against a specific numerical pipeline. If the numpy reference and the CUDA engine implement the same arithmetic (same group size, same quantization formula, same dequantization formula, same accumulation), they should produce the same tokens. This is a closed-world validation: the reference defines the truth.

The final sentence of the reasoning—"and the divergence from FP32 confirms quantization is actually being applied"—shows the assistant interpreting the earlier test results. The divergence at token 2 is not a bug report; it is a confirmation signal. The quantization pipeline is working. The assistant can now proceed with confidence.

Conclusion

Message [msg 12070] is a quintessential "bridge" message in a complex engineering effort. It connects the validated quantization pipeline (Python reference + bundle generator) to the engine's runtime (model loading + kernel dispatch). It is not the most technically impressive message in the session—that honor belongs to the custom CUDA kernels or the full engine validation—but it is the message that makes all that work cohere into a functional system. Without it, the W4A16 kernel would never be called, the INT4 bundles would never be loaded, and the engine would remain forever FP32. With it, the engine takes its first step toward running on the actual quantized weights that K2.6 uses in production, proving that the architecture can support real-world deployment.