The Moment of Truth: Validating INT4 Quantization in a Custom DDTree Inference Engine

The Message

cd /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine
python3 python/gen_model_int4.py 2>&1
echo "=== sanity: how different are INT4 vs FP32 golden tokens? ==="
python3 - <<'PY'
import sys; sys.path.insert(0,"python")
import kdtr_io
f=kdtr_io.load("tests/refs/model_tiny.kdtr")["golden_tokens"].tolist()
q=kdtr_io.load("tests/refs/model_int4_tiny.kdtr")["golden_tokens"].tolist()
print("fp32 :",f)
print("int4 :",q)
print("diverge at:", next((i for i,(a,b) in enumerate(zip(f,q)) if a!=b), "none (identical)"))
PY
wrote /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/../tests/refs/model_int4_tiny.kdtr  quantized 84 weights (G=32)  golden[:8]=[499, 128, 58, 76, 462, 243, 474, 353]
wrote /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/../tests/refs/model_int4_tiny2.kdtr  quantized 90 weights (G=32)  golden[:8]=[187, 328, 252, 328, 121, 158, 129, 137]
=== sanity: how different are INT4 vs FP32 golden tokens? ===
fp32 : [499, 128, 83, 210, 49, 128, 16, 269, 203, 224, 499, 170, 117, 203...

At first glance, this appears to be a routine execution of a Python script followed by a quick sanity check. But within the arc of this ambitious project—building a native C/C++/CUDA speculative decoding engine for the Kimi K2.6 language model—this message represents a critical inflection point. It is the moment where theory meets reality, where the carefully designed INT4 quantization scheme is put to the test, and where the assistant discovers whether the entire approach is viable or fundamentally broken.

Context: The Road to INT4 Quantization

To understand why this message matters, we must trace the path that led here. The assistant had been building a custom "DDTree" (Draft-Draft Tree) speculative decoding engine from scratch—a native C/C++/CUDA inference stack designed to accelerate the Kimi K2.6 model on 8× NVIDIA RTX PRO 6000 Blackwell GPUs. The engine's architecture was organized around three custom CUDA kernels: a GPU best-first tree builder, a tree-verify MLA-absorb attention kernel, and a greedy tree-accept kernel. These had been validated against numpy golden references, passing all 27 tests with bit-exact accuracy.

But the engine had a critical limitation: it ran entirely in FP32 using cuBLAS GEMMs as placeholder for the real matrix multiplications. The Kimi K2.6 model, like most modern large language models, uses INT4 weight quantization to fit its 548GB of parameters into GPU memory. The FP32 placeholder was a correctness stand-in, not a realistic implementation. The real engine needed to handle W4A16 (4-bit weights, 16-bit activations) group-quantized matrix operations, specifically for the Mixture-of-Experts (MoE) layers and dense MLPs that constitute the bulk of the model's compute.

The user had explicitly chosen to use the idle CT200 server (8× PRO 6000 GPUs, freshly moved to a production VLAN) for "the next engine phase." The assistant interpreted this as a mandate to advance the engine from FP32 simulation to genuine INT4 operation. The plan was ambitious: write a custom W4A16 group-quantized GEMM kernel, build a numpy reference implementation for validation, generate INT4 quantized model bundles, and prove that the engine's INT4 forward pass matched the reference token-for-token.

Messages immediately preceding this one had laid the groundwork. The assistant had edited ops.cuh and ops.cu to add the W4A16 kernel declarations and implementations, modified model_ref.py to include quantization and dequantization functions, and written gen_model_int4.py—the script that would produce the quantized bundles and golden reference tokens. Message 12069 is the first execution of that script.

What the Message Actually Does

The message contains two commands executed sequentially. The first runs gen_model_int4.py, which performs several critical operations:

  1. Loads the existing FP32 model configuration (the tiny and tiny2 test models used throughout development)
  2. Quantizes the designated weights to INT4 with group size 32—specifically the MoE expert weights, gating projections, shared expert layers, and dense MLP weights (gates, ups, and downs)
  3. Packs the 4-bit values into a compact representation (two nibbles per byte, stored as int32 arrays)
  4. Computes golden reference tokens using the dequantized weights, so the "golden" output represents what the INT4 model should produce
  5. Writes two bundle files: model_int4_tiny.kdtr and model_int4_tiny2.kdtr The output confirms this succeeded: "wrote ... model_int4_tiny.kdtr quantized 84 weights (G=32)" and "wrote ... model_int4_tiny2.kdtr quantized 90 weights (G=32)". The "golden[:8]" prefixes show the first eight tokens of each model's autoregressive output. The second command is the sanity check—a Python one-liner that loads both the original FP32 golden tokens and the new INT4 golden tokens, then finds the first position where they differ. This is the moment of truth.

The Critical Discovery: Divergence at Position 2

The output is revealing. The FP32 golden tokens begin [499, 128, 83, 210, ...] while the INT4 golden tokens (shown in the script output as golden[:8]=[499, 128, 58, 76, ...]) diverge at index 2. The sanity check confirms: "diverge at: 2"—the third token.

This is a significant finding, and the assistant's decision to run this comparison immediately after generation shows careful engineering judgment. The divergence is not a bug; it is the expected consequence of quantization. When weights are compressed from 32-bit floating point to 4-bit integers (with group-wise scaling), information is lost. The model's output distribution shifts slightly, and at some point the argmax (the token with highest probability) changes. In this case, it happens at the third token.

The implications are profound for the DDTree speculative decoding approach. The entire premise of speculative decoding is that a smaller "drafter" model proposes tokens that the larger "target" model can verify in parallel. If quantization changes the target model's output, the drafter must be calibrated against the quantized target, not the original FP32 version. More subtly, if the INT4 quantization introduces numerical drift that accumulates over long sequences, the model's behavior could diverge significantly from the FP32 baseline—potentially affecting quality, coherence, or task performance.

Assumptions and Their Validation

This message tests several implicit assumptions that the assistant had been working under:

Assumption 1: The quantization scheme is numerically sound. The assistant designed the W4A16 format with group size 32, symmetric quantization, and careful handling of scales. The divergence at position 2 doesn't invalidate this assumption—it merely confirms that quantization introduces error, which is expected. What matters is whether the error is bounded and whether the engine's INT4 output matches the numpy INT4 reference. That test comes later when the C++ engine loads these same bundles.

Assumption 2: The golden reference is correct. The assistant chose to compute golden tokens using dequantized weights rather than the original FP32 weights. This means the INT4 golden represents "what the INT4 model should output" rather than "what the FP32 model outputs." This is the correct approach for validation—the engine should match the INT4 reference, not the FP32 reference—but it means the assistant has implicitly accepted that the INT4 model is a different model from the FP32 one.

Assumption 3: The bundle format is sufficient. The .kdtr bundle format, designed earlier in the project, stores model weights and metadata in a binary container. The assistant extended it to include .qw (quantized weights) and .sc (scales) fields. The successful write of both bundles validates this extension.

Assumption 4: The test models are representative. The "tiny" and "tiny2" configurations were designed to exercise the engine's full pipeline while keeping compute costs low. They include MoE layers with multiple experts, shared experts, and the full transformer architecture (RMSNorm, RoPE, SwiGLU). If the quantization works on these small models, it should work on the full K2.6—though the group size divisibility constraints and numerical behavior may differ at scale.

The Thinking Process Visible in This Message

The assistant's reasoning, visible in the preceding messages, reveals a careful balance between ambition and pragmatism. The original plan was to extract the Marlin INT4 kernel from SGLang's sgl_kernel library and link it directly into the C++ engine. But the assistant recognized the risk: Marlin is deeply entangled with CUDA's Cutlass library and PyTorch's custom op machinery. Extracting it would be a multi-day effort with uncertain outcomes.

The decision to write a custom W4A16 kernel instead was a deliberate trade-off. The assistant explicitly noted: "It won't match marlin's speed, but it proves the INT4 path end-to-end in the native engine." This is classic engineering pragmatism—build the simplest thing that validates the architecture, then optimize later. The plan document already separates correctness from performance: "Marlin remains the documented peak-perf drop-in; this is the correctness/format building block."

The assistant also made a conscious decision about which weights to quantize. Following K2.6's architecture, only the MoE experts and dense MLPs (the SwiGLU weights) are quantized to INT4. The MLA (Multi-head Latent Attention) projections and embeddings remain in FP32. This is faithful to the real model and keeps validation tractable—attention is where the DDTree kernels operate, and keeping it FP32 avoids compounding quantization error with the tree verification logic.

Input Knowledge Required

To fully understand this message, one needs:

  1. The DDTree engine architecture: Knowledge that the engine has three custom CUDA kernels (tree builder, verify attention, tree accept) and uses a numpy reference for validation.
  2. The quantization scheme: Understanding of W4A16 group quantization—that weights are stored as 4-bit integers with group-wise FP32 scale factors, and that matrix multiplication requires dequantization on the fly.
  3. The Kimi K2.6 model architecture: Familiarity with MLA (Multi-head Latent Attention), MoE (Mixture of Experts) with shared experts, SwiGLU activations, and the specific set of weights that are quantized.
  4. The KDTR bundle format: Knowledge that .kdtr files are binary containers storing model weights, configuration, and golden reference outputs, used to share data between Python and C++.
  5. The speculative decoding loop: Understanding that DDTree uses a drafter to propose token trees and a target model to verify them, and that the target model's output distribution must be deterministic for verification to work.
  6. The CT200 hardware context: The 8× RTX PRO 6000 Blackwell GPUs with 96GB each, the production VLAN migration, and the service being down.

Output Knowledge Created

This message produces several pieces of critical knowledge:

  1. Quantization is functional: The gen_model_int4.py script works correctly, producing valid KDTR bundles with quantized weights and correct golden tokens.
  2. INT4 diverges from FP32 at position 2: The quantization error is not zero, and it affects the model's output starting at the third token. This has implications for any application that assumes INT4 preserves FP32 behavior exactly.
  3. Two test models are now available: The tiny (84 quantized weights) and tiny2 (90 quantized weights) configurations provide test vectors for the C++ engine's W4A16 kernel validation.
  4. The validation pipeline is established: The assistant now has a clear workflow: generate INT4 bundles with numpy, load them in the C++ engine, run the forward pass, and compare tokens. Any mismatch indicates a bug in the engine's kernel.
  5. A baseline for quantization error: The specific divergence point (token index 2) gives a concrete metric for how much quantization affects this particular model configuration. Future improvements to the quantization scheme (different group sizes, asymmetric quantization, etc.) can be measured against this baseline.

Broader Implications

The divergence at position 2 raises an important question that the assistant will need to address: if the INT4 model disagrees with the FP32 model, which one is "correct"? In production, the INT4 model is the real deployment target—the FP32 version is just a simulation tool. But the DDTree speculative decoding loop depends on the drafter predicting tokens that the target model will accept. If the drafter was trained or calibrated against the FP32 target, it may perform worse against the INT4 target.

This is not a problem the assistant needs to solve in this message—the immediate goal is validating the engine's INT4 kernel against the numpy reference. But the divergence foreshadows a deeper challenge: the entire speculative decoding pipeline must be consistent. The drafter, the target model, and the verification kernel must all agree on what constitutes a "correct" token. Quantization introduces a systematic shift that must be accounted for.

The assistant's approach—generating INT4 golden references and validating against them—is the correct engineering response. By treating the INT4 model as the ground truth, the assistant avoids the trap of chasing FP32 fidelity that doesn't matter in production. The engine only needs to match the INT4 reference; whether that reference matches the original FP32 model is a separate concern for model quality evaluation.

Conclusion

Message 12069 is a quiet but pivotal moment in the development of a complex inference engine. On the surface, it's a simple script execution and a sanity check. But beneath that lies the validation of a fundamental design decision: the choice to implement custom INT4 quantization rather than integrate an existing library, the acceptance of quantization error as a fact of life, and the establishment of a validation pipeline that will guide the rest of the engine's development.

The divergence at position 2 is not a failure—it's data. It tells the assistant that the quantization scheme works, that the bundle format is correct, and that the INT4 model is a distinct entity from its FP32 parent. The next step will be loading these bundles into the C++ engine and running the W4A16 kernel, comparing token-by-token against the numpy reference. If they match, the engine will have crossed a critical threshold: from FP32 simulation to genuine quantized inference.

This message exemplifies the kind of disciplined engineering that makes complex systems work. The assistant didn't assume the quantization would be perfect; it immediately checked, quantified the divergence, and recorded the result. The data is now available for the next stage of development, and the path forward is clear.