Nemotron architecture note

A hybrid model is not one block repeated many times. Mamba state, attention history, router decisions, and expert weights create different execution and memory contracts.

Nemotron kernel series: architecture to execution

Start with Nemotron Architecture From A C-Kernel-Engine Runtime Perspective. That report maps the complete hybrid runtime surface: attention, Mamba2 state, dense and routed feed-forward paths, checkpoint conversion, CKE lowering, generated C, and evidence gates. This article is the second step. It corrects the model-family history and then narrows the investigation to exact router mathematics, expert dispatch, and CPU memory behaviour.

The most useful lesson in Nemotron is not simply that Mamba can replace some attention layers. The deeper lesson is that a model family can evolve from dense hybrid blocks into routed hybrid MoE blocks while keeping the same broad architecture identity. A runtime must read the checkpoint and preserve the exact layer policy. It cannot infer the execution path from the name alone.

The model name selects a family. The configuration and weights determine the circuit that actually runs.

Terminology: Nemotron is NVIDIA's model family. NVIDIA NeMo is the broader framework and ecosystem used to develop, train, customize, and deploy model families such as Nemotron. The names are related, but they are not interchangeable. This series extracts the model's executable kernel contracts so they can be studied independently in CKE.

Comparison of original dense Nemotron-H and newer routed hybrid Nemotron variants
The original Nemotron-H and newer Nemotron MoE variants share hybrid Mamba and attention ideas, but they do not have the same feed-forward circuit.

The correction that matters

The original 2025 Nemotron-H models were not all routed MoE models. NVIDIA describes the 8B base model as 24 Mamba-2, 24 dense MLP, and four self-attention layers. The 56B base model uses 54 Mamba-2, 54 dense MLP, and ten self-attention layers.

Newer Nemotron 3 models extend the hybrid idea with routed experts. For example, NVIDIA documents Nemotron 3 Nano as a Mamba2-Transformer hybrid MoE with 128 routed experts plus one shared expert in each MoE layer and six selected experts per token.

CKE therefore needs both contracts: dense relu2 MLP layers and routed relu2 MoE layers.

One Family, Multiple Circuits

NVIDIA introduced Nemotron-H as a family of hybrid Mamba-Transformer models. The architecture replaces most self-attention layers with Mamba2 layers. During autoregressive decode, an attention layer repeatedly reads a KV cache that grows with context. A Mamba2 layer instead updates a fixed-size recurrent state. A small number of attention layers preserve direct access to token history while recurrent layers carry most of the sequence work.

That original design should be separated from the newer Nemotron 3 Nano design. Nemotron 3 adds sparse MoE layers. A router scores many experts but activates only a small subset for each token. The model can carry far more total parameters than it evaluates per token, but the runtime now has to solve sparse dispatch, expert locality, weighted accumulation, and load distribution.

Architecture surfaceOriginal Nemotron-HNewer hybrid MoE variantsPersistent runtime implication
Mamba2Present in most recurrent layersPresent in most recurrent layersConvolution history and selective-state buffers survive across decode calls.
AttentionA small number of attention layersA small number of GQA/attention layersOnly attention layers own KV cache; recurrent layers must not allocate fake KV state.
Feed-forwardDense ReLU2 MLPRouted ReLU2 experts plus a shared expertThe layer policy determines whether every row uses one dense weight set or selected expert weight sets.
RouterNot required for dense checkpointsScores, groups, selects, normalizes, and scales expertsOutputs integer expert IDs and floating-point routing weights for every token row.

This distinction is represented directly in CKE's Safetensors conversion tests. One fixture lowers four layer kinds as [mamba, attention, moe, mlp] and records a routed-ReLU2 policy only for the MoE layer. A separate dense fixture lowers [mamba, attention, mlp] and asserts that no MoE weights or router tensors appear. The tests are not merely checking filenames. They verify that checkpoint structure becomes a per-layer execution policy.

The Four Layer Kinds CKE Must Preserve

CKE's Nemotron-H v8 circuit does not encode one universal decoder body. It reads layer_kinds and expands each layer into a different operation sequence.

mamba:
  RMSNorm -> in projection -> split -> Conv1D/SiLU
  -> dt softplus -> selective scan -> gated RMSNorm
  -> output projection -> residual

attention:
  RMSNorm -> Q/K/V -> attention -> output projection -> residual

mlp:
  RMSNorm -> up projection -> ReLU2 -> down projection -> residual

moe:
  RMSNorm -> router -> group-limited top-k
  -> routed ReLU2 experts + shared expert -> residual

These paths cannot share one memory plan. Attention owns layer-specific KV cache. Mamba owns recurrent state and convolution history. A dense MLP streams one pair of projection matrices. An MoE layer may touch different expert matrices for adjacent token rows. The tensor shapes can all end at the same hidden width while the storage lifetimes and access patterns remain fundamentally different.

Four Nemotron layer kinds lowered into distinct CKE operation sequences
The residual stream is common, but each layer kind expands into a different circuit and state policy.

The Router Is A Deterministic Algorithm

“Choose the best experts” is not a sufficient router specification. The implementation must answer which score is used for selection, whether experts are grouped, how groups are ranked, whether selected weights are renormalized, how ties are resolved, and when scaling is applied.

For one token row with E experts, CKE's group-limited router receives sigmoid probabilities:

\[s_e = \sigma(z_e), \qquad e \in \{0,\ldots,E-1\}\]

An optional correction bias affects the selection score:

\[c_e = s_e + b_e\]

Experts are divided into G equal groups. The score of a group is the sum of its two largest corrected expert scores. The router keeps the best T_G groups, masks every other group, and then selects the global top K experts from the surviving corrected scores.

\[q_g = \operatorname{Top2Sum}\{c_e : e \in g\}\]
\[\mathcal{G}^{*}=\operatorname{TopK}(q, T_G), \qquad\mathcal{E}^{*}=\operatorname{TopK}\{c_e : e \in \mathcal{G}^{*}\}, K\]

The subtle step comes next. Correction bias chooses the experts, but the output weights are gathered from the original sigmoid scores s. If configured, those weights are normalized and then multiplied by the routed scaling factor:

\[w_e = \alpha\frac{s_e}{\sum_{j\in\mathcal{E}^{*}}s_j + 10^{-20}},\qquad e\in\mathcal{E}^{*}\]

This is exactly the kind of detail that can produce plausible but numerically different output. Using corrected scores as expert weights, ranking groups by their maximum instead of top-two sum, or allowing unstable tie ordering changes the selected execution path.

Step-by-step group-limited top-k Nemotron router mathematics
Selection and weighting use related but distinct scores. The correction bias affects choice; original probabilities determine mixture weights.

The C Kernel Contract

The current C reference kernel makes the interface explicit:

scores          float[R][E]
correction_bias float[E]       // optional
indices         int[R][K]
weights         float[R][K]

parameters:
  n_group, topk_group, norm_topk_prob, routed_scaling_factor

It also gives equal scores a deterministic expert-index tie break. That detail matters for reproducibility: two implementations can have the same values but select different experts if their top-k ordering is unstable.

Routing Is Cheap; Expert Memory Is The Hard Part

Once the router returns expert IDs and weights, every active expert performs a dense ReLU2 MLP:

\[u_{r,e}=W^{up}_e h_r\]
\[a_{r,e}=\max(0,u_{r,e})^2\]
\[y_r=\sum_{e\in\mathcal{E}^{*}_r}w_{r,e}W^{down}_e a_{r,e} + y^{shared}_r\]

The arithmetic is familiar GEMV/GEMM work. The scheduling is not. If token row 0 chooses experts 3 and 41 while row 1 chooses experts 7 and 99, a naive row-by-row loop jumps between distant weight regions. A model may perform fewer FLOPs than a dense network while creating worse cache and memory behavior.

Comparison of token-order and expert-batched MoE execution on CPU memory
Expert batching groups rows by selected expert so each expert's weights can be reused before the runtime moves to another expert.

A CPU implementation therefore has several possible execution schedules:

  1. Token-major reference: for each row, execute its selected experts immediately. This is simple and useful for correctness.
  2. Expert-major batching: build a dispatch list, group token rows by expert, run a denser expert batch, then scatter weighted outputs back.
  3. Thread-local expert queues: reduce synchronization while preserving some expert weight locality.
  4. NUMA-aware ownership: place expert weights and their worker threads together so routed rows do not repeatedly cross sockets.
  5. Distributed expert parallelism: place experts on different nodes and exchange routed activations rather than replicating every expert everywhere.

Each optimization changes the reduction order and temporary buffers. CKE's reference-first approach is valuable here: the optimized scheduler still has to reproduce the reference expert IDs, routing weights, and accumulated output within a declared numerical contract.

What CKE Has Actually Established

CKE has more than a diagram, but less than a finished high-performance Nemotron runtime. The distinction should remain visible.

EvidenceWhat existsWhat it establishesWhat it does not establish
Architecture inspectionNemotron-H configuration and Safetensors-family classification testsMamba, attention, dense MLP, router, routed expert, and shared-expert tensors can be identified.It does not prove complete model output parity.
ConversionDense and routed Safetensors-to-BUMP fixtures with zero unmapped source tensorsCheckpoint tensors become explicit CKE names and per-layer policies.Small synthetic fixtures are not full public checkpoint validation.
Mamba referencesFP32 in-projection split, Conv1D decode, dt softplus, state update, sequence scan, and gated RMSNorm testsIndividual recurrent contracts can be compared with PyTorch.The scalar path is not the final SIMD/chunked implementation.
RouterPyTorch parity for Nano-like 128-expert top-6 routing and additional configurationsExpert indices match exactly and weights match at 1e-6 absolute tolerance.A router microbenchmark is not full MoE-layer throughput.
ExpertsFP32 forward/backward references and selected quantized forward variantsReLU2 expert math is isolated behind kernel maps.Expert batching, cache locality, NUMA placement, and end-to-end scaling remain optimization work.

The router parity test includes a Nano-like case with 128 experts, top-6 selection, eight groups, and three selected groups. It requires exact integer indices and FP32 weights within 1e-6. This is strong primitive evidence because a wrong expert ID cannot be hidden by a small floating-point tolerance.

Current status

Validated primitives: architecture inspection, dense/routed conversion fixtures, reference Mamba2 kernels, group-limited routing, and ReLU2 expert kernels.

Still required: full checkpoint conversion and execution, layer-by-layer PyTorch attribution, teacher-forced generation parity, production scheduling, profiler evidence, and measured model-level throughput.

Why This Architecture Is Especially Interesting On CPUs

Hybrid recurrent MoE models challenge the simplistic CPU-versus-GPU argument. Most Mamba2 layers avoid a context-length-dependent KV-cache read during decode. MoE activates only a fraction of total expert weights per token. Those properties reduce some work, but they also make locality and scheduling more important.

A CPU server has large addressable memory, deep caches, many cores, NUMA topology, and mature thread-placement controls. That can be useful if the runtime batches rows by expert, places state deliberately, and keeps the selected weights close to execution. It can be terrible if every token triggers random weight streaming across sockets. The architecture does not automatically favor a CPU or GPU. The execution schedule determines whether the hardware is used coherently.

This is where VTune and Advisor become more useful than theoretical peak arithmetic. The next measurements need to answer:

  • How much time is spent in router selection versus active expert projection?
  • Do adjacent rows select enough common experts to reuse weights?
  • Does expert-major batching improve LLC hit rate and memory bandwidth utilization?
  • At what row count does dispatch-list construction pay for itself?
  • Should experts be partitioned by thread, NUMA node, or physical machine?
  • How do quantized expert weights change the compute-to-bandwidth balance?

The Next Honest Gates

  1. Run a real supported Nemotron checkpoint through the converter and publish the conversion audit.
  2. Export matching PyTorch hidden states after Mamba, attention, dense MLP, router, routed expert, and residual boundaries.
  3. Locate the first full-model divergence rather than relying only on primitive tests.
  4. Freeze teacher-forced parity before changing expert scheduling.
  5. Measure token-major versus expert-major execution on AVX2 and Xeon AVX-512/AMX hardware.
  6. Add cache, bandwidth, branch, thread-placement, and power evidence beside throughput.
  7. Only then claim an optimized Nemotron CPU path.

Conclusion

Nemotron is valuable to CKE because it forces the engine to stop pretending a model is a homogeneous stack. Attention layers carry token history. Mamba2 layers carry recurrent state. Dense MLP layers stream one weight set. MoE layers make a discrete routing decision and then stream a sparse subset of expert weights.

The router is the hinge between model mathematics and systems behavior. Its indices determine which memory is read, which threads receive work, which NUMA node becomes busy, and eventually which machine receives an activation. That makes router correctness, expert locality, and distributed placement one connected problem.

The earlier Nemotron architecture report mapped the full hybrid surface. This follow-up narrows the focus to the routed path and corrects the architecture history. CKE now has the reference contracts needed to investigate that path honestly. The next step is not another support badge. It is full-model parity followed by measured CPU scheduling evidence.

Continue through the circuit

The next article, Nemotron Kernels From Inference To Backpropagation, reverses these four execution paths. It derives the gradients for Mamba2, attention, dense ReLU2, and routed ReLU2 experts; separates differentiable mixture weights from hard top-k choices; and identifies where quantized inference ends and training precision begins.

Primary Sources And CKE Artifacts