The Golden Reference: Anchoring INT4 Quantization Correctness in a Custom CUDA Inference Engine

Message Overview

The subject message ([msg 12067]) is a brief but pivotal moment in a complex engineering session: the assistant edits model_ref.py to add numpy-based quantization and dequantization functions for INT4 W4A16 group-quantized weights. The full message reads:

## Agent Reasoning Now I'm setting up the numpy-based quantization reference implementation with a group size of 32, defining functions to pack 4-bit weights into int32 arrays with separate float32 scales, and then implementing the forward pass logic to use these quantized weights.

>

For the workflow, I'll create a function that takes the weight dictionary and quantizes specific layers (MLP gates, ups, downs, and all MoE expert variants), returning both a dequantized version for computing the golden output and a separate dictionary of packed weights and scales for the engine to store.

>

Writing quantization functions... Writing the engine loader... Now the numpy reference quantization (the spec the engine must match): [edit] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/model_ref.py Edit applied successfully.

On its surface, this is a simple file edit. But the reasoning section reveals a carefully considered design philosophy: the numpy implementation is not just a utility — it is the spec the engine must match. This message represents the moment when the assistant commits to a validation methodology that will govern the entire INT4 quantization integration.

Context: The Engine That Must Be Right

To understand why this message matters, we need to see where it fits in the larger project. The assistant has been building a native C/C++/CUDA inference engine called kdtree-engine for the Kimi K2.6 large language model, with a focus on speculative decoding via a custom algorithm called DDTree (Draft-Driven Tree). The engine is designed to be torch-free — a standalone C++ binary that links CUDA kernels directly, bypassing the Python runtime entirely.

In [msg 12063], the assistant laid out an extensive internal debate about how to approach Phase 3: integrating INT4 W4A16 quantized weights into the engine. The Kimi K2.6 model uses INT4 quantization for its MoE (Mixture-of-Experts) and MLP weights, making this step essential for realistic memory footprint and throughput. The assistant considered several approaches:

  1. Extracting the Marlin MoE kernel from sgl_kernel — rejected because the kernel is entangled with PyTorch's custom op machinery and Cutlass dependencies.
  2. A hybrid Python/C++ approach — rejected because it defeats the native engine goal.
  3. Writing a custom INT4 W4A16 GEMM kernel — chosen as the pragmatic path. The key insight in that earlier message was the assistant's realization that the real bottleneck is integration risk, not kernel performance. The assistant explicitly states: "Rather than risk the complexity of extracting marlin (cutlass dependencies, build integration), I'll write my own INT4 group-quantized GEMM kernel—it's self-contained, testable against a numpy reference, and proves the engine handles real quantization correctly." This decision frames everything that follows. The custom kernel won't match Marlin's peak performance, but it will prove the INT4 path works end-to-end. Marlin becomes "the documented peak-perf drop-in" — a future optimization, not a current dependency.

Why This Message Was Written

The subject message is the direct consequence of that architectural decision. If you're writing a custom INT4 GEMM kernel, you need two things:

  1. A specification of exactly what the kernel must compute.
  2. A ground-truth validator to compare against. The numpy reference implementation serves both roles simultaneously. It defines the quantization format (group size 32, symmetric per-group scaling, nibble packing into int32 arrays) and provides the dequantized forward pass that the CUDA kernel must match token-for-token. The assistant's reasoning reveals a sophisticated understanding of the validation challenge. Quantization introduces numerical error — the INT4 forward pass will not produce bit-identical results to the FP32 forward pass. But the goal is not bit-exactness; it's greedy-exactness: the argmax (most likely token) must agree. The assistant notes this explicitly in [msg 12063]: "For the kernel itself, I need to be careful about matching 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." This is a crucial engineering judgment. The assistant is not aiming for mathematical identity between the numpy reference and the CUDA kernel — that would be impossible given different floating-point evaluation orders. Instead, the target is functional equivalence: the engine must produce the same token selections as the reference, even if the raw logits differ in the 6th decimal place.

The Design Decisions Embedded in This Message

The message reveals several design decisions, some explicit and some implicit:

Group size 32. This matches Kimi K2.6's native quantization format, ensuring that the reference implementation is faithful to the production model's weight layout. It also means the engine can eventually load real K2.6 weights without reformatting.

Selective quantization. Only MLP gates, ups, downs, and all MoE expert variants are quantized. MLA (Multi-head Latent Attention) projections and embeddings remain in FP32. This mirrors K2.6's actual architecture and keeps validation tractable — the attention mechanism, which is critical for correctness, stays in full precision.

Packed storage. The assistant decides to pack two 4-bit values per byte rather than storing them unpacked as int8. This is more work but "demonstrates genuine 4-bit compression." The assistant is thinking about the end-to-end system: the engine must eventually handle real model weights that are packed, so the reference should validate the packing and unpacking logic.

Dual return values. The quantization function returns both a dequantized version (for computing golden outputs) and a separate dictionary of packed weights and scales (for the engine to load). This cleanly separates the validation path from the engine path, allowing the assistant to verify that the engine's packed-weight forward pass matches the reference's dequantized forward pass.

Assumptions and Potential Pitfalls

The assistant makes several assumptions that are worth examining:

The numpy reference is authoritative. This is a reasonable engineering practice — the reference implementation defines the spec. But it means any bugs in the quantization logic propagate to both the reference and the engine. The assistant mitigates this by keeping the quantization simple (symmetric, per-group) and by planning to validate against the FP32 golden output as a sanity check.

Group size 32 is always feasible. The assistant notes the need to "verify the divisibility constraints for the quantization scheme across all model sizes." This is a real concern — if a model dimension is not divisible by 32, the quantization scheme breaks. The assistant is aware of this and plans to check it, but the message doesn't show the verification.

The dequantization arithmetic can be made identical between numpy and CUDA. This is harder than it sounds. NumPy evaluates expressions left-to-right with full FP32 precision, while CUDA kernels may fuse multiply-add operations or use different intermediate precision. The assistant's plan to use "identical fp32 accumulation" is sound in principle but requires careful implementation.

The engine's INT4 path will be validated against the numpy reference, not against the FP32 path. This is a deliberate choice — the assistant could validate that the INT4 engine matches the INT4 reference, and separately validate that the INT4 reference matches the FP32 reference (within quantization error). But the message doesn't discuss the second validation step explicitly.

Input Knowledge Required

To understand this message, one needs:

Output Knowledge Created

This message produces:

  1. Quantization functions in model_ref.py. These define the exact arithmetic for packing 4-bit weights, computing per-group scales, and dequantizing during the forward pass.
  2. A validation spec. The numpy reference becomes the authoritative definition of what the CUDA W4A16 kernel must compute. Any future changes to the kernel must be validated against this reference.
  3. A bundle format extension. The quantization functions produce packed weights and scales that the engine loader must read. This extends the KDTR bundle format to support quantized weight storage.

The Thinking Process

The reasoning section reveals a methodical, layered thought process:

  1. Goal identification: "set up the numpy-based quantization reference implementation with a group size of 32."
  2. Format design: "pack 4-bit weights into int32 arrays with separate float32 scales."
  3. Scope definition: "quantizes specific layers (MLP gates, ups, downs, and all MoE expert variants)."
  4. Workflow design: "returning both a dequantized version for computing the golden output and a separate dictionary of packed weights and scales for the engine to store."
  5. Execution: "Writing quantization functions... Writing the engine loader..." The assistant is thinking in terms of contracts: the quantization function defines a contract between the numpy reference and the CUDA engine. The packed weights and scales are the serialized form of that contract. The dequantized forward pass is the validation oracle. This is classic systems engineering: define the interface first, implement both sides independently, then verify they agree. The assistant is building a test harness as much as a quantization implementation.

Conclusion

Message [msg 12067] is a small edit that carries significant weight. It represents the transition from planning to implementation in a complex engineering effort — adding INT4 quantization to a custom CUDA inference engine. The assistant's decision to build a numpy reference implementation as the authoritative spec reflects a disciplined approach to correctness in a domain where numerical errors can silently corrupt model outputs.

The message also illustrates a key principle of the kdtree-engine project: every CUDA kernel is validated against a numpy golden reference before it is trusted. This methodology, established in earlier phases for the DDTree kernels (build, verify, accept), is now being extended to the quantization path. The assistant is not just writing code — it is building a verification framework that will catch regressions and ensure the engine remains correct as it evolves from FP32 placeholders to real INT4 weights.

The deeper story here is about engineering risk management. Rather than attempting the high-risk path of extracting Marlin from sgl_kernel's tangled dependency graph, the assistant chose a lower-risk path: write a custom kernel, validate it against a numpy reference, and defer peak-performance optimization to a later phase. This message is where that choice becomes concrete, as the first lines of the reference implementation are written.