The Bridge Between FP32 and INT4: Building the Test Infrastructure for a Native DDTree Engine

In the course of developing a native C/C++/CUDA speculative decoding engine for the Kimi K2.6 model, one seemingly small message ([msg 12068]) marks a pivotal architectural transition: the creation of gen_model_int4.py, a Python script that generates INT4-quantized test bundles. On its surface, this is a simple file write — the assistant writes a script and moves on. But beneath this brief action lies a carefully reasoned decision about how to validate one of the most critical components of the entire engine: the W4A16 group-quantized GEMM path that will ultimately determine whether the engine can run the real K2.6 model at all.

The Subject Message: A Script Is Born

The message itself is brief and direct:

## Agent Reasoning Now I'm writing gen_model_int4.py to generate INT4 quantized bundles with non-quantized fp32 weights, quantized weights stored as .qw and .sc files, and golden tokens computed using dequantized weights, reusing the configuration from the reference model generator. Now the INT4 bundle generator: [write] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/gen_model_int4.py Wrote file successfully.

The assistant's reasoning block succinctly describes the script's purpose: it generates INT4 quantized bundles that include non-quantized FP32 weights alongside quantized weights stored as .qw (packed 4-bit values) and .sc (float32 group scales), with golden tokens computed using dequantized weights. The script reuses the configuration from the existing reference model generator, ensuring consistency between the FP32 and INT4 test suites. The file is written successfully — a single tool call, a single action.

But this simplicity is deceptive. The message is the culmination of an extended reasoning chain that spans multiple previous messages, weighing architectural trade-offs, assessing integration risk, and ultimately committing to a specific validation strategy. To understand why this particular file exists and why it takes the form it does, we must trace the reasoning that led to it.

The Fork in the Road

To understand why this message matters, we must trace the reasoning that led to it. The story begins with the assistant standing at a fork in the road ([msg 12063]). CT200 — an 8× RTX PRO 6000 Blackwell machine — had been moved to a production VLAN and restarted. The SGLang DDTree service was inactive, leaving all eight GPUs completely idle. The user had chosen to use this free capacity for "the next engine phase" rather than restoring the production service.

But what was the next engine phase? The assistant's reasoning in [msg 12063] reveals an extended internal debate, weighing three distinct paths:

Path one: Extract the real Marlin MoE kernel from sgl_kernel and link it into the native C++ engine. This was architecturally pure — the engine would use the same high-performance INT4 GEMM that SGLang itself uses. But it carried serious risk: the Marlin kernel is entangled with Cutlass, PyTorch custom op registration, and a complex build system. Extracting it meant either pulling in a massive dependency surface or reverse-engineering the raw CUDA launcher.

Path two: Write a custom W4A16 group-quantized GEMM kernel from scratch, self-contained within the engine's own CUDA code. This would not match Marlin's peak performance, but it would prove the INT4 path end-to-end with full control over correctness. Marlin could remain a documented performance upgrade for later.

Path three: A hybrid approach — run the real K2.6 MoE GEMM through sgl_kernel's Marlin op in Python, paired with the DDTree kernels called via ctypes. This would produce realistic performance numbers quickly, but it defeated the entire purpose of building a native engine.

The assistant chose Path two. The reasoning is instructive: "Performance can come from swapping in marlin later; the plan already separates correctness from optimization." This is a classic engineering trade-off — prioritize correctness and self-containment over peak performance, establishing a known-correct baseline that can be optimized later without changing the validation surface.

The Validation Pipeline

With the W4A16 kernel committed to ops.cuh and ops.cu ([msg 12064], [msg 12065], [msg 12066]) and the quantization reference functions added to model_ref.py ([msg 12067]), the assistant faced a new problem: how to validate that the engine's CUDA kernel produces the same results as the numpy reference, given that quantization introduces approximation error.

The answer is gen_model_int4.py — the file written in [msg 12068]. This script performs several critical functions:

  1. Weight quantization: It takes the existing FP32 model weights and quantizes specific layers — MoE expert weights, MLP gate/up/down projections — to W4A16 format with group size 32, matching K2.6's actual quantization scheme.
  2. Packing: The 4-bit values are packed two-per-byte into a compact binary representation, demonstrating genuine 4-bit compression rather than storing unpacked int8 values.
  3. Golden reference generation: Crucially, the script computes "golden tokens" using dequantized weights — weights that have been quantized and then reconstructed back to FP32. This means the golden output accounts for quantization error, so the engine must match the quantized behavior exactly, not the original FP32 behavior.
  4. Bundle output: The quantized weights, scales, and golden tokens are written to .kdtr bundle files that the engine's loader can read, using the same binary container format (KDTR) established earlier in the project. The script reuses the model configuration from the existing reference model generator, ensuring consistency between the FP32 and INT4 test suites. It writes two files: .qw for the packed 4-bit weights and .sc for the float32 group scales.

Assumptions Embedded in the Design

The gen_model_int4.py script encodes several assumptions that deserve scrutiny:

Group size of 32: This matches K2.6's actual quantization scheme, but it also imposes divisibility constraints on the model dimensions. The assistant explicitly noted the need to verify these constraints across all model sizes before committing to the format.

Which layers to quantize: Only the MoE expert weights and dense MLP layers (the SwiGLU weights where Marlin would apply) are quantized. The MLA projections and embeddings remain in FP32. This is faithful to K2.6's approach, but it means the engine must handle a mixed-precision model — some weights quantized, others not — and dispatch to the correct GEMM path for each layer.

Packing order: The 4-bit values are packed two-per-byte in a specific order that the CUDA kernel must unpack identically. Any mismatch between the numpy packing and the CUDA unpacking would cause silent correctness failures that might not be caught by simple token comparisons.

Dequantization arithmetic: The assistant recognized this as a critical risk: "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." The golden tokens are computed with numpy's floating-point semantics; the CUDA kernel must reproduce those semantics closely enough that argmax selects the same token.

Why This Message Matters

The gen_model_int4.py script is the bridge between two worlds. On one side is the FP32 reference model — a numpy-based golden implementation that defines correct behavior but bears no resemblance to how the real K2.6 model runs. On the other side is the native engine — a CUDA-accelerated C++ program that must eventually load and run the actual 548GB of INT4 Marlin weights across 8 GPUs with tensor parallelism.

This script creates the intermediate test infrastructure that makes it possible to validate the engine's INT4 path before tackling the enormous complexity of real weight loading, TP-8 sharding, and NCCL communication. It is a deliberate scoping decision: prove correctness on small synthetic bundles first, then scale up to real weights.

The assistant's reasoning reveals a sophisticated understanding of validation strategy. Rather than attempting to build the entire K2.6 inference stack in one shot — which the assistant correctly identified as "genuinely multi-week work" — the work is decomposed into testable increments. The W4A16 kernel is validated against a numpy reference. The kernel integration is validated through the engine's linear() dispatch. The full engine is validated through end-to-end greedy-exact tests. Each layer of validation depends on the test infrastructure that gen_model_int4.py provides.

In the very next message ([msg 12069]), the assistant runs the script and verifies that the INT4 golden tokens diverge from the FP32 golden tokens — as expected, since quantization introduces error. The divergence is not a bug; it is the precise behavior that the engine must match. The test infrastructure is working as designed.