The Golden Reference: Building a NumPy Specification for a Native CUDA Inference Engine

Introduction

In the sprawling, multi-month effort to deploy Kimi K2.6 with DFlash speculative decoding on 8× RTX PRO 6000 Blackwell GPUs, there comes a quiet but pivotal moment: the creation of a NumPy reference model. Message [msg 11920] in this opencode session captures that moment — a brief dispatch from the assistant that, on its surface, appears to be a simple file write. But beneath the surface, this message represents a critical architectural decision that shapes the entire validation strategy for a custom C/C++/CUDA inference engine.

The message reads:

## Agent Reasoning I'm setting up the reference implementation in NumPy with MLA parameterized directly in absorb form using w_kc and w_vc as parameters, keeping everything in float32 to match the engine's precision. Starting with the numpy reference model (the spec + golden source). MLA parameterized directly in absorb form (w_kc/w_vc), MoE with top-k softmax routing + shared expert, SwiGLU: [write] /home/theuser/glm-kimi-sm120-rtx6000bw/kdtree-engine/python/model_ref.py Wrote file successfully.

This is the first concrete implementation step of Phase 2 — the native engine — following the user's directive in [msg 11918] to "continue and don't stop until we have a working MVP C hyper-optimized inference stack." To understand why this particular message matters, we must trace the reasoning that led to it, the decisions embedded in its brief text, and the assumptions that underpin the entire validation architecture.

The Context: From CUDA Kernels to a Complete Engine

The session leading up to [msg 11920] had already accomplished significant work. In [msg 11917], the assistant had completed and committed the kdtree-engine/ repository, containing three validated custom CUDA kernels — a GPU best-first tree builder, a tree-verify MLA-absorb attention kernel with visibility masking, and a greedy tree-accept kernel — all passing 27/27 tests bit-exact against NumPy references. These kernels were the highest-risk, most novel components of the DDTree speculative decoding stack, proven on real sm_120 silicon (the RTX 5070 Ti, sharing architecture with the target PRO 6000 GPUs).

But kernels alone do not make an inference engine. The user's command was unambiguous: build a working MVP. The assistant's extensive reasoning in [msg 11919] reveals the strategic dilemma. A full K2.6 native engine would require weight loading for a 548GB model, Marlin INT4 quantization extraction, TP-8 NCCL parallelism, MLA attention forward passes, MoE routing, and a dozen other components — genuinely months of work. The alternative was a self-contained MVP on a small model that exercises the entire stack, proving correctness end-to-end, with the remaining K2.6-specific work documented as mechanical scale-up.

The assistant chose the latter, and [msg 11920] is the first execution of that choice. The NumPy reference model is the cornerstone of the validation strategy: a golden specification that defines the exact mathematical behavior the C++ engine must reproduce.

Why NumPy? The Reference-as-Spec Approach

The decision to write the reference model in NumPy rather than PyTorch is revealing. The assistant's reasoning notes that PyTorch is not available locally, but the choice runs deeper. NumPy offers several advantages for this role:

First, purity of specification. NumPy is a mathematical library, not a deep learning framework. A NumPy implementation of a transformer forward pass is closer to executable pseudocode than to production code. It makes no assumptions about GPU memory layouts, kernel fusion, or batching strategies. It is the "what" without any hint of the "how" — exactly what a specification should be.

Second, precision control. The assistant explicitly chooses float32 for both the NumPy reference and the engine, noting that this avoids "BF16 numerics complicating reference matching." This is a deliberate scoping decision: prove correctness in FP32 first, then optimize to BF16 or INT4 as a later phase. The reference must be a faithful oracle, and float32 provides a straightforward path to bit-exact or near-bit-exact agreement.

Third, determinism and reproducibility. NumPy operations with a fixed seed produce identical results across runs and platforms. The reference model generates "golden" token sequences and per-step logits that the C++ engine must match exactly. This is the same philosophy used in hardware verification and compiler testing: define the correct output once, then test every implementation against it.

The assistant's approach mirrors the strategy used in the earlier kernel validation (Phase 1), where each CUDA kernel was tested against a NumPy reference. Now that pattern scales from individual kernels to the entire engine.

The Architectural Decisions Embedded in the Reference

The message's reasoning text reveals several specific architectural choices that define the reference model's structure:

MLA in absorb form with w_kc/w_vc parameterization: The Multi-head Latent Attention (MLA) mechanism used in DeepSeek V2/V3 and Kimi K2.6 is notoriously complex, involving low-rank KV projections that are later "absorbed" into the query and output projections for computational efficiency. The assistant chooses to parameterize MLA directly in the absorbed form, using w_kc (key components) and w_vc (value components) as the actual model parameters rather than deriving them from a larger kv_b_proj matrix. This is a deliberate simplification for the test model — it matches the computation the engine will actually perform, avoiding unnecessary indirection in the reference.

MoE with top-k softmax routing and shared expert: The Mixture-of-Experts configuration mirrors the K2.6 architecture but at a tiny scale (8 experts, top-2, with the first layer as a dense MLP). The shared expert is a feature of DeepSeek's MoE design that the assistant preserves in the reference.

SwiGLU activation: The choice of SwiGLU (swish-gated linear unit) over simpler activations like ReLU or GELU is faithful to the modern LLM architecture that K2.6 uses.

Row-major weight convention: The assistant defines all projections as y = x @ W where W is [in, out] — the row-major convention natural to C/C++ and cuBLAS. This contrasts with PyTorch's column-major convention (W @ x), and getting this right is essential for the cuBLAS wrapper to produce correct results.

These decisions are not arbitrary. Each one reflects a mapping between the NumPy reference and the eventual C++ implementation. The reference is designed not just to produce correct tokens, but to produce them through the same mathematical pathway the engine will use, making the comparison meaningful.

The Validation Strategy: Token-Exact Greedy Output

The ultimate goal of the reference model is to prove what the assistant calls "the critical invariant": that DDTree speculative greedy output matches autoregressive greedy output token-for-token. This is the fundamental correctness property of any speculative decoding system — the draft-and-verify process must never change the output distribution.

To prove this, the assistant needs:

  1. A reference autoregressive implementation that produces a golden token sequence
  2. A reference DDTree implementation that produces a token sequence using draft-verify-accept
  3. Proof that both sequences are identical The NumPy reference provides both. The engine will be validated against it in stages: first the autoregressive forward pass (prefill + decode logits), then the full DDTree loop. The assistant's plan, detailed in [msg 11919], includes generating "golden AR greedy tokens/logits" from the reference, which the engine must reproduce exactly.

Assumptions and Their Implications

The reference model makes several assumptions that are worth examining:

The drafter is a placeholder. The assistant explicitly notes that for the MVP, the drafter can be trivial — even random candidates — because the goal is to validate the tree-build and verification loop, not optimize acceptance rates. The real DFlash drafter is a "drop-in replacement with the same interface." This is a pragmatic assumption that separates correctness validation from performance optimization, but it means the MVP cannot demonstrate the throughput benefits of speculative decoding. The acceptance rate with a random drafter will be near zero.

FP32 precision is sufficient for correctness proof. The assistant assumes that FP32 agreement between reference and engine implies correctness, and that later precision changes (BF16, INT4) are purely mechanical. This is reasonable but not trivial — quantization can introduce numerical edge cases that change token selection, especially for low-probability tokens near the decision boundary.

cuBLAS GEMM is a faithful placeholder for Marlin INT4. The engine uses cuBLAS FP32 matrix multiplications where the production system would use Marlin INT4 kernels. The assistant treats this as a correctness-safe substitution because the reference is also FP32. This is valid for proving the control flow and architecture, but it means the engine's performance characteristics (memory bandwidth, compute throughput) will differ dramatically from the final system.

Single-sequence batch is sufficient for MVP. The assistant explicitly notes "starting with a single-sequence batch for the MVP" and marks batching as a future phase. This is a reasonable scope constraint, but it means the engine cannot yet exploit the parallelism that makes GPU inference efficient.

Input Knowledge Required

To fully understand this message, the reader needs knowledge of:

Output Knowledge Created

This message produces:

  1. The reference model file (model_ref.py): A complete NumPy implementation of a DeepSeekV3/Kimi-style MLA+MoE transformer, serving as the golden specification for the C++ engine.
  2. A validated architectural design: The specific parameterization choices (absorb-form MLA, w_kc/w_vc, row-major weights, top-k softmax MoE) that define the mathematical contract between reference and engine.
  3. A reproducibility foundation: The reference model generates deterministic golden outputs that can be used to validate every subsequent change to the engine, from primitive operations to the full DDTree loop. This output knowledge is the foundation for everything that follows. Without the NumPy reference, the engine's correctness would be impossible to verify — the assistant would be building in the dark, unable to distinguish between a correct implementation and one that merely appears to work.

The Thinking Process: From Abstract Plan to Concrete Code

The reasoning in [msg 11920] is notably terse compared to the extensive deliberation in [msg 11919]. This is characteristic of the transition from planning to execution. In [msg 11919], the assistant explored multiple options (full K2.6 engine vs. small-model MVP), weighed trade-offs (months of work vs. achievable validation), and designed the architecture (MLA absorb form, FP32, cuBLAS GEMM placeholder). By [msg 11920], those decisions are settled, and the focus shifts to implementation.

The reasoning reveals a clear hierarchy of concerns:

  1. Precision strategy: Float32 throughout, avoiding BF16 complications
  2. Parameterization: MLA in absorb form with w_kc/w_vc as direct parameters
  3. Architecture components: MoE with top-k routing, shared expert, SwiGLU
  4. Implementation priority: Reference first, then engine primitives, then integration This ordering reflects a "spec-first" development methodology: define the correct behavior in a high-level, debuggable language before implementing it in a high-performance one. It is the same approach used in hardware design (where a functional specification precedes RTL implementation) and in compiler construction (where a reference interpreter precedes the optimizing compiler).

Conclusion

Message [msg 11920] is a quiet but foundational moment in the construction of a custom inference engine. It is the point where planning becomes code, where architectural decisions become executable specifications, and where the abstract goal of "a working MVP C hyper-optimized inference stack" begins to take concrete form. The NumPy reference model it creates is not merely a testing tool — it is the mathematical constitution against which the entire engine will be judged. Every kernel, every GEMM call, every attention computation will be validated against this golden standard.

The message's brevity belies its importance. In a single file write, the assistant establishes the entire validation framework for one of the most complex software engineering projects in modern ML infrastructure: a custom, hyper-optimized inference engine for a state-of-the-art language model, running on cutting-edge Blackwell hardware, with speculative decoding that must produce exactly the same tokens as the autoregressive baseline. The NumPy reference is the anchor that makes this audacious goal achievable.