Validating the INT4 W4A16 Path: A Milestone in Native DDTree Engine Development

In the sprawling, multi-session effort to build a native C/C++/CUDA speculative decoding engine for the Kimi K2.6 large language model, message [msg 12076] represents a quiet but critical inflection point. It is the moment when a weeks-long architectural bet — that the engine could be extended from a pure FP32 cuBLAS prototype to a genuine INT4 W4A16 group-quantized inference pipeline — is validated in a single, clean test run. The message itself is deceptively brief: a bash command that rebuilds the engine, followed by the output of two test executables confirming that both the autoregressive (AR) and DDTree speculative decoding paths produce token-exact results against the INT4 golden reference. But the weight of this moment can only be understood by tracing the chain of reasoning, engineering trade-offs, and architectural decisions that led to it.

The Context: From FP32 Prototype to Real Quantization

The native DDTree engine, housed in the kdtree-engine/ repository, had been built in phases. Phase 0 established the build infrastructure and a numpy reference model. Phase 1 delivered three custom CUDA kernels — a GPU best-first tree builder, a tree-verify MLA-absorb attention kernel, and a greedy tree-accept kernel — all validated against numpy references. Phase 2 produced a working MVP transformer engine in FP32 using cuBLAS GEMMs as placeholders for the eventual INT4 Marlin kernels. The engine could load model weights from a custom binary container format (KDTR), run the full DDTree speculative decode loop, and produce greedy-exact results matching the autoregressive baseline.

But the FP32 path was never the endgame. The Kimi K2.6 model, with its 548GB of weights, is stored in INT4 W4A16 format — 4-bit group-quantized weights with FP32 scales. Running the engine in FP32 meant it could demonstrate correctness but could never approach realistic memory footprints or throughput. The plan had always called for integrating INT4 Marlin MoE kernels, but the path to that integration was fraught with uncertainty.

The Fork in the Road: Marlin Extraction vs. Custom Kernel

In the messages leading up to [msg 12076], the assistant engaged in an extensive internal debate about how to bring INT4 quantization into the engine. The reasoning, captured in [msg 12063], reveals a careful weighing of options:

The most direct path would be to extract the marlin_moe_wna16 CUDA kernel from SGLang's sgl_kernel package and compile it directly into the engine. But this kernel lives inside a PyTorch custom op, wrapped in torch tensor machinery, and depends on the Cutlass library for its template metaprogramming. Extracting it would mean untangling a significant dependency surface, potentially pulling in Cutlass headers and JIT compilation infrastructure that the engine was deliberately designed to avoid.

A hybrid approach — running the real K2.6 MoE GEMM through Python's sgl_kernel while calling the DDTree kernels via ctypes — was considered but rejected as "essentially reimplementing the model in Python, which defeats the native engine goal."

The assistant then contemplated writing a simple INT4 dequant-plus-grouped-GEMM kernel from scratch. This would be self-contained, with no external dependencies, and could be validated against an FP32 reference. The downside was that it would duplicate Marlin's functionality and wouldn't match Marlin's performance, undermining the plan to use the engine for realistic throughput measurements.

After cycling through these options multiple times, the assistant committed to the custom kernel path with a clear rationale: "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. Performance can come from swapping in marlin later; the plan already separates correctness from optimization."

This was a deliberate architectural decision to decouple correctness from performance — a classic engineering strategy that prioritizes getting the right answer before optimizing for speed.

Building the INT4 Pipeline

With the decision made, the assistant executed a multi-step implementation plan across several files:

  1. ops.cuh / ops.cu ([msg 12064]-[msg 12066]): The W4A16 group-quantized GEMM kernel was added to the engine's CUDA operations. This kernel takes FP32 activations and multiplies them by INT4 packed weights with group scales, unpacking and dequantizing on the fly.
  2. model_ref.py ([msg 12067]): The numpy reference model was extended with quantize_w4a16 and dequantization functions, creating the specification that the engine's CUDA kernel must match exactly.
  3. gen_model_int4.py ([msg 12068]): A new bundle generator was written to produce INT4-quantized KDTR bundles, storing packed weights as .qw files and scales as .sc files alongside the existing FP32 weights for non-quantized layers (embeddings, MLA projections).
  4. model.h / model.cu ([msg 12070]-[msg 12075]): The engine's model loader was extended to recognize quantized weight files, load them into GPU memory, and dispatch through a unified linear() helper that routes to either the INT4 GEMM or the FP32 cuBLAS path depending on whether quantized weights exist for a given layer. A critical insight emerged during this work: the INT4 golden output would necessarily diverge from the FP32 golden output because quantization introduces numerical error. This meant the engine could not simply reproduce the FP32 reference — it needed to match the INT4 reference exactly. The assistant acknowledged this in [msg 12070]: "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."

The Validation: Message 12076

With all the pieces in place, message [msg 12076] executes the validation. The assistant takes three concrete actions:

First, it wires the INT4 bundle generation into the build scripts. Both run_tests_nvcc.sh and build_and_test.sh are modified with sed commands to add python3 python/gen_model_int4.py after the existing gen_model_ref.py invocation. This ensures that INT4 reference bundles are automatically regenerated as part of the build process, keeping them in sync with any changes to the model configuration.

Second, it rebuilds the engine locally using the nvcc build script. The build output confirms all five compilation targets succeeded: "[4/5] kernel unit tests + bench" and "[5/5] engine tests + demo (needs cublas)" completed, producing the build/ directory with the test executables.

Third, it runs the two critical tests against the INT4 model bundle model_int4_tiny.kdtr:

=== INT4 AR + DDTree (local 5070 Ti) ===
golden          AR : 499 128 58 76 462 243 474 353 74 410 401 211 211 474 150 386 128 26 197 236 472 197 32 32 
PASS  model_ar  tokens=24/24 exact  max_abs_logit_diff=8.225e-06
ddtree : 499 128 58 76 462 243 474 353 74 410 401 211 211 474 150 386 128 26 197 236 472 197 32 32 
DDTree stats: steps=3 committed=25 avg_accept=8.00 tokens/step (block_size=8)
PASS  model_ddtree  AR==golden and DDTREE==golden (24 tokens), greedy-exact

The autoregressive test confirms that the engine's INT4 forward pass produces exactly the same 24 tokens as the numpy INT4 golden reference, with a maximum absolute logit difference of only 8.225e-06 — well within the tolerance for greedy decoding where only the argmax matters.

The DDTree test confirms that the speculative decoding loop, using the custom tree-build, tree-verify, and tree-accept kernels, produces the same 24 tokens as the autoregressive path. The statistics show 3 steps committed 25 tokens (one extra from the tree), with an average acceptance rate of 8.00 tokens per step at block_size=8 — the theoretical maximum for a perfectly predictable sequence.

What This Message Proves

The significance of this validation extends beyond the immediate test results. It proves several things simultaneously:

  1. The W4A16 GEMM kernel is correct. The dequantization arithmetic in CUDA matches the numpy reference closely enough that the argmax (and therefore the token output) is identical. This is non-trivial: floating-point accumulation order differences between numpy (Python, typically row-major) and CUDA (parallel, with warp-level reductions) can produce small numerical discrepancies. The 8.225e-06 max logit difference confirms the implementation is faithful.
  2. The bundle format extension works. The .qw/.sc convention for storing quantized weights alongside FP32 weights, with the engine falling back to FP32 for non-quantized layers, is functional. The loader correctly distinguishes between weight types and populates the appropriate GPU memory.
  3. The dispatch mechanism is correct. The linear() helper routes MoE and MLP GEMMs through the INT4 path while keeping embeddings and attention projections in FP32, and the combined forward pass produces the correct output.
  4. The DDTree loop is quantization-agnostic. The speculative decoding kernels (tree builder, verify attention, tree accept) operate on the hidden states produced by the model forward pass. They don't care whether those hidden states came from INT4 or FP32 weights — they just need the model to produce the correct logits. This test confirms that property.
  5. The local development workflow is sound. The assistant chose to validate on a local RTX 5070 Ti rather than the production 8× RTX PRO 6000 Blackwell box (CT200). This allowed fast iteration — the build and test cycle completes in seconds rather than requiring a deployment to the remote machine. The fact that the test passes locally gives high confidence it will pass on the target hardware.

Assumptions and Their Validity

The validation rests on several assumptions that deserve examination:

The INT4 golden reference is the correct target. The assistant assumes that matching the INT4 reference is sufficient to prove the engine handles real quantization correctly. This is valid because the INT4 reference was generated by the same quantization algorithm that would be applied to real K2.6 weights. If the engine matches the reference, it will match any weights quantized with the same scheme.

Group size 32 is representative. The test uses group size 32, matching the K2.6 configuration. The divisibility constraints are satisfied for the test model dimensions, and the same kernel will work for any dimensions that are multiples of 32. This is a reasonable assumption given that K2.6's hidden dimensions are designed to be compatible with group-quantized formats.

The tiny model is representative. The test models (model_int4_tiny and model_int4_tiny2) have small dimensions (e.g., hidden_size=256, intermediate_size=768, 4 experts). The assumption is that correctness on small models implies correctness on full-scale models, since the kernel is data-parallel and doesn't change behavior with size. This is standard practice in CUDA kernel development.

Floating-point agreement is sufficient. The test checks for token-exact matches, which only requires the argmax to agree. The 8.225e-06 max logit difference suggests the numerical error is well-controlled, but the assistant implicitly assumes this level of agreement will hold at larger scales. This is a reasonable assumption given that the error comes from floating-point accumulation order, which doesn't systematically grow with problem size.

The Broader Trajectory

Message [msg 12076] closes a major chapter in the engine's development. The INT4 W4A16 path is now validated as correct, opening the door to the next phase: loading real K2.6 weights, setting up tensor parallelism across 8 GPUs, and measuring end-to-end throughput with the speculative decoding loop operating on genuinely quantized weights.

The assistant's todo list from [msg 12063] shows what remains: deploying to the CT200 box, running the INT4 tests on the 8× PRO 6000 hardware, and then tackling the "real K2.6 end-to-end" integration — weight loading, TP-8 sharding, NCCL communication, and the full inference pipeline. The Marlin kernel swap for peak performance is documented as a future optimization, deferred until the correctness path is proven.

But for this single message, the focus is narrower and more satisfying: the engine now handles INT4 quantization, and it does so correctly. The 24 tokens of output, identical across AR and DDTree paths, are a small but decisive signal that the architecture is sound. The custom kernel bet paid off.

Conclusion

Message [msg 12076] is a validation message in the truest sense — it confirms that a complex, multi-file implementation effort produced the correct result. The assistant's reasoning process, visible across the preceding messages, shows a methodical approach to a difficult engineering problem: weighing integration strategies, decoupling correctness from performance, building reference implementations, and incrementally wiring the pieces together. The final test output, with its clean "PASS" lines and precise numerical agreement, is the payoff for that careful work. It transforms the engine from an FP32 prototype into a genuinely quantized inference pipeline, ready for the next challenge of full-scale K2.6 deployment.