The previous Nemotron article separated four execution paths: Mamba2, attention, dense ReLU2 MLP, and routed ReLU2 MoE. This follow-up turns those paths around. For each one, what runs during inference? What must training save? How does the gradient return? Which tensors can be quantized without pretending that low-precision inference and low-precision training are the same problem?

Terminology: Nemotron names the model family. NVIDIA NeMo is the broader development and deployment ecosystem around models such as Nemotron. CKE is not reproducing the entire NeMo platform here. It is extracting the mathematical and numerical contracts required to execute and eventually train the model's kernels in generated C.

Part I: The Mathematical Prerequisites

Before starting Kernel 1, we need a small mathematical vocabulary. A derivative measures local sensitivity: if one value changes slightly, how strongly does another value respond? Backpropagation does not invent a new kind of mathematics. It applies familiar derivative rules to a circuit and carries the result backward. You do not need all of calculus to follow this article; the compact rulebook below is sufficient for the four kernel circuits.

A Derivative Is A Local Slope

If \(y=f(x)\), the derivative \(dy/dx\) answers: “If I nudge \(x\) a little at this point, approximately how much will \(y\) move?” In training, we ask the same question about loss: “If I change this activation or weight, how does the loss move?”

The Small Rulebook

Nemotron does not require every technique from a calculus textbook. Its kernels repeatedly use a compact set: constants and sums, powers, products, quotients, exponentials, logarithms, sigmoid, and the chain rule. Branching adds one bookkeeping rule: copy a gradient when the forward value splits, and add gradients when reverse paths meet.

The Chain Rule Is The Backpropagation Engine

Suppose \(x\) enters function \(f\), its result enters \(g\), and that result enters the loss. Forward evaluates left to right. Backward starts from \(dL/dL=1\) and multiplies each local derivative while moving right to left.

Five Labels Make Every Backward Kernel Readable

For one operation \(y=f(x)\), the backward kernel receives an upstream gradient \(dL/dy\). It computes the local derivative \(dy/dx\), multiplies them, and returns \(dL/dx\). “Upstream” means the gradient arrived from an operation that occurred later during forward execution.

Branches Copy Forward And Add Backward

A residual block sends \(x\) through an identity path and a kernel path. Forward adds the two outputs. Backward must add the gradient contributions from both uses of \(x\). Forgetting one path silently removes a learning signal.

Part II: Inference, Training Forward, And Backward Are Three Contracts

Inference is a forward evaluation with no parameter gradients. Prefill processes a sequence and creates the state needed for generation. Decode usually processes one new token, updates a KV cache or recurrent state, and discards most temporary activations. Training also runs a forward pass, but it must preserve or later recompute the values required by reverse-mode differentiation.

CKE's Nemotron v8 circuit already declares the four forward paths. It does not yet declare that all four have complete model-level backward coverage. The correct next step is to expose that gap operation by operation.

Setup: The Model Boundary, Embedding, Logits, And Loss

Token IDs first gather rows from an embedding matrix. After the hybrid layer stack, final RMSNorm and the language-model head produce vocabulary logits. Inference samples or selects from those logits. Teacher-forced training compares each position with its target token:

The cross-entropy boundary gives a compact starting gradient:

The LM-head backward accumulates an outer product into \(g_{W_{lm}}\) and sends a projected gradient into final RMSNorm and the last decoder block. Embedding backward is a scatter-add into the rows selected by the input token IDs. If input embeddings and output weights are tied, both uses contribute to one shared gradient buffer. A quantized LM head can reduce inference traffic, but training still needs a defined higher-precision gradient and master-weight policy.

Setup: The Shared Transformer Shell

Every Nemotron path sits inside a pre-normalized residual block:

Backward begins at the residual split. If the incoming gradient is \(g_y\), one branch passes directly to \(x\), while the other traverses the selected body:

The identity branch is simple. The normalized branch is not. RMSNorm backward needs the original input, learned scale, inverse RMS, and the exact reduction contract. CKE has explicit RMSNorm backward and ReLU2 backward maps. These shared primitives are the beginning of a trainable circuit, not proof that the complete hybrid model trains end to end.

Kernel 1: Mamba2 Inference And Backpropagation

Forward And Inference: State Instead Of A Growing KV Cache

The CKE Mamba2 forward circuit is:

RMSNorm -> input projection -> split(gate, x/B/C, dt)
-> causal Conv1D + SiLU -> dt softplus
-> selective scan -> gated group RMSNorm
-> output projection -> residual add

The short causal convolution is repeated multiply-and-add, but that compact equation hides an important idea. Digital systems first represent a changing signal as an ordered sequence of samples. A one-dimensional convolution aligns a short weight vector with a local sample window, multiplies matching positions, adds the products, and writes one output. The window then shifts by one sample and the same weights are reused. For a causal filter, the output at time or token position \(t\) may read the present and past, but never a future position.

Why Convolution Appears In Signals, CNNs, And Mamba

Convolution predates neural networks. In signal processing, a finite impulse response filter describes how a system combines the latest sample with a short history. Carefully chosen coefficients can smooth noise, emphasize rapid changes, isolate bands, or detect a known local pattern. The convolution theorem adds another view: convolution in the time domain corresponds to multiplication in the frequency domain after a Fourier transform. Convolution is not itself a Fourier transform, but the relationship explains why a local filter can suppress or amplify frequency components.

Convolutional neural networks made those filter coefficients learnable. Instead of an engineer manually selecting an edge detector or smoothing filter, training discovers local patterns useful for the task. In an image, the filter moves across rows and columns. In audio, sensor telemetry, or token sequences, a one-dimensional filter moves through sampled time or sequence position. The reusable insight is locality: nearby samples often form meaningful short patterns before a larger system reasons over longer distances.

What Conv1D Gives The Mamba Block

Mamba's short depthwise causal convolution supplies the first tier of sequence memory: immediate local context. It cheaply mixes neighboring projected token features before the selective state update. That local stage can recognize order-sensitive fragments, transitions, and short motifs. The state-space recurrence then decides how strongly that prepared evidence should alter a fixed-size memory carried across the sequence. In a hybrid Nemotron layer stack, this lower-memory path complements the layers that still perform full attention.

This division of labour is important. Full attention provides addressable memory: a query can directly compare itself with stored token representations. Mamba combines two older mathematical ideas for the layers between those global-memory points: a local causal filter for nearby structure and a dynamical state-space system for compressed longer-range history. Conv1D prevents the recurrent update from seeing each token position as an isolated sample, while the recurrent state avoids making the short convolution window responsible for remembering the entire context.

During backward, every output gradient contributes to the input values in its window, to the convolution weights that touched those values, and to the bias. Overlapping windows cause several gradient contributions to meet at one input location.

The CKE decode kernel applies SiLU to the convolution sum. Its backward therefore first uses the SiLU derivative shown later in the gating diagram, then sends that resulting gradient through the convolution products and overlapping windows.

State Space Before Mamba: The Engineering Equation

State space is a standard language in control systems, robotics, state estimation, mechanical dynamics, and electrical engineering. The central idea is that a complicated system can often be advanced through time if we preserve a compact vector containing the information needed for its future evolution. For a drone, that state may contain position, velocity, attitude, and angular rate. For an electrical circuit, it may contain capacitor voltages and inductor currents. For a robot arm, it may contain joint positions and velocities.

In continuous time, \(x(t)\) is the hidden state, \(u(t)\) is an external input, and \(y(t)\) is the observable output. Matrix \(A\) describes how existing state evolves, \(B\) describes how the input changes the state, \(C\) reads state into an output, and \(D\) allows an immediate input-to-output path. After choosing a sampling interval \(\Delta t\), the system becomes a discrete recurrence:

This does not always mean “predict the future” in the statistical sense. It means propagate the modeled state from sample \(k\) to sample \(k+1\). A controller uses the state to choose an input. An observer or Kalman filter uses dynamics plus noisy measurements to estimate state that cannot be measured directly. The Kalman filter adds uncertainty models and a predict-correct cycle; Mamba does not secretly run a Kalman filter, but both rely on the principle that a compact latent state can summarize information carried across time.

Photographic Memory And Compressed Memory

Attention and state-space layers solve the same broad architectural problem: make information from earlier positions useful at the current position. They are not the same mathematical operator, but they are complementary memory mechanisms. Full attention acts like the model's photographic memory: it stores key and value representations for previous positions, and a new query can directly address that history to form a fresh weighted mixture. “Photographic” is an engineering analogy rather than literal perfect recall, but it captures the practical implication: previous positions remain separately available instead of being collapsed into one summary. The cost is that stored KV history grows with context and each attention layer must read that history.

A state-space layer supplies compressed working memory. It does not preserve every previous token as an individually addressable record. It repeatedly folds the current input into a fixed-size hidden state. The output reads from that state, and the next token updates it again. This creates linear recurrent execution and bounded state memory, but it also creates an information bottleneck: details not preserved in the state cannot later be retrieved from that Mamba layer by directly addressing the original token.

The important hybrid-design variable is therefore not “attention or memory.” Both paths are memory. It is how many layers require context-length-dependent photographic memory through individually addressable KV history and how many layers can operate with bounded compressed state. Reducing the number of full-attention layers reduces total KV-cache growth while retaining selected layers where global token-to-token lookup remains available.

From Classical State Space To A Selective Mamba State

The classical discrete equation already resembles a recurrent neural network: carry old state through \(A_d\), inject the current input through \(B_d\), and expose an output through \(C\) and \(D\). In a language model, the projected token features take the role of the input, Mamba's hidden tensor takes the role of state, and the layer output returns information to the model's residual stream.

Mamba's important change is selectivity. A fixed linear time-invariant system applies the same dynamics at every sample. Mamba lets the current token influence quantities such as the time step \(\Delta_t\), the write vector \(B_t\), and the read vector \(C_t\). The model can therefore learn to preserve, ignore, overwrite, or expose different information depending on the token. The hardware-aware selective scan evaluates this recurrent structure efficiently across a sequence.

The projection creates several streams. The hidden stream passes through a short causal convolution. The time-step stream is constrained by softplus. The selective scan then updates a fixed recurrent state for every time step:

CKE then clamps the transformed time step into a configured interval. Inside the interval, the local derivative is the sigmoid shown above. Outside it, the clamp has saturated and its derivative is zero; at the exact boundary, the implementation must declare a convention. This apparently small branch determines whether a time-step parameter can still receive a learning signal.

Before adding heads and state dimensions, one scalar recurrence contains the essential idea. The new state is one product carrying old memory plus another product writing the new input. The product rule immediately tells us how the gradient divides among prior state, decay, input, and write strength.

After the scan, a learned gate controls how much signal passes. The gate itself is SiLU: \(\operatorname{SiLU}(g)=g\sigma(g)\). Backward first splits at the final product, then applies the product, sigmoid, and chain rules inside SiLU.

During prefill, CKE's selective scan walks the complete sequence and returns both \(u_{1:T}\) and the final state. During decode, the state-update kernel evaluates the same recurrence for one step. The unit test requires full-sequence scan and repeated one-token decode to agree.

Backpropagation: What Mamba2 Backward Must Do

Reverse-mode differentiation walks the recurrence from \(T\) back to 1. The gradient arriving at a state combines the current output path with the future recurrent path:

In the indexed equations below, \(t\) is the token position, \(h\) is a Mamba head, \(d\) is a channel within that head, \(s\) is a state coordinate, and \(g\) is a parameter-sharing group. The function \(g(h)\) returns the group that owns head \(h\). Therefore, a shared \(B_{t,g,s}\) or \(C_{t,g,s}\) receives the sum of every head and channel assigned to that group.

That state gradient then produces contributions for the input-dependent state parameters:

Gradients also pass through \(\bar A_t=\exp(\Delta_t A)\), the softplus transform, causal convolution, gating and both projections. The exact tensor broadcasting depends on heads, groups, head dimension and state dimension. That is why a generic statement such as “differentiate the scan” is not a kernel contract.

Kernel 2: Attention Inference And Backpropagation

Forward And Inference: Save Probabilities Or Recompute Them

For an attention layer:

Inference prefill writes K and V into the layer's cache. Decode appends one position and computes one query row against cached keys and values. Training does not mutate a decode cache. It needs enough information to recover the softmax probabilities and the Q/K/V inputs for backward.

Backpropagation: Reverse Value Mixing, Softmax, And Q/K/V

Let \(S=QK^{\mathsf T}/\sqrt{d_h}+M\) be one row of attention scores and \(P=\operatorname{softmax}(S)\). Reversing \(O=PV\) first separates the value path from the probability path:

Softmax couples every probability in a row through the shared denominator. For row coordinates \(i\) and \(j\), its Jacobian entry is \(\partial P_i/\partial S_j=P_i(\delta_{ij}-P_j)\). Multiplying that Jacobian by the arriving probability gradient gives:

The vectorized row operation and the following score-matrix gradients are therefore:

CKE's GQA attention backward map makes the storage contract visible: it receives the upstream gradient, forward Q/K/V, and saved attention weights, and produces gradients for Q, K, V and score space. Projection backward then accumulates \(g_{W_Q},g_{W_K},g_{W_V},g_{W_O}\).

Nemotron's split RoPE is another local operation around Q and K. Forward rotates coordinate pairs; backward applies the transpose of that rotation to the Q and K gradients before their projection backward kernels run. In grouped-query attention, several query heads share one K/V head, so their K and V gradient contributions must be added into the same shared destination.

Kernel 3: Dense ReLU2 Inference And Backpropagation

Forward And Inference: Project, Activate, Project

A dense ReLU2 MLP is:

Backpropagation: The Smallest Complete Reverse Path

Its reverse path is direct:

The important numerical boundary is the point \(u=0\), where the implementation's derivative convention must be fixed. The CKE forward and backward maps preserve the input to ReLU2 so the reverse kernel can apply the same branch used by the forward kernel.

Kernel 4: Routed ReLU2 MoE Inference And Backpropagation

Forward And Inference: Select, Execute, Weight, And Combine

For selected expert set \(\mathcal E_r^*\), CKE evaluates:

Start with only two selected experts. The output is a weighted sum. The product rule gives a gradient for each floating mixture weight and each expert output; the sum rule adds both expert input gradients where they meet.

Backpropagation: Differentiate Mixture Weights And Active Experts

The selected expert path has two gradient branches:

Each active expert then runs the dense ReLU2 reverse equations. Gradients from all selected and shared experts add into \(g_{x_r}\). CKE's routed expert backward kernel explicitly returns d_hidden, d_routing_weights, d_expert_up, and d_expert_down.

Backpropagation Boundary: Hard Top-K Has No Ordinary Derivative

The router's integer expert IDs do not have an ordinary derivative. A small score change may leave the selected set unchanged, then abruptly replace one expert. Training normally treats dispatch as fixed for the selected path and differentiates the selected mixture weights, while auxiliary balancing objectives and the model's training recipe shape router learning. CKE has a deterministic forward router, but it does not yet declare a complete Nemotron router-training contract or auxiliary-loss implementation. That must be matched to the exact model recipe before claiming training parity.

Other systems can choose surrogate estimators for this discrete boundary. A straight-through estimator can copy a gradient across the hard decision, Gumbel-softmax can replace selection with a temperature-controlled differentiable relaxation, and score-function methods such as REINFORCE can estimate gradients from sampled routing outcomes. These methods optimize different surrogate objectives and have different variance, bias, and numerical contracts. Their existence does not make hard top-k itself differentiable, and CKE must not silently select one. Nemotron training parity requires reproducing the exact router objective, selected-weight normalization, auxiliary losses, and gradient convention used by the reference recipe.

Where Quantization Fits

Quantization is not one switch applied to the entire graph. It is a storage and arithmetic policy for particular tensor edges. The most attractive targets are the large projection and expert matrices because they dominate bytes moved from memory.

For quantization-aware training, the forward graph can insert fake quantization or low-precision simulation while backward uses a defined estimator and higher-precision accumulators. For post-training quantization, training remains unchanged and calibration or error measurement selects a storage format. These are different products and require different evidence.

What Exists In CKE Today

How We Would Validate The Complete Path

1. Freeze one tiny Nemotron configuration and deterministic input batch.
2. Export PyTorch values after every forward boundary: normalization, projections, Conv1D, scan, attention, router, experts and residuals.
3. Export gradients at the same boundaries after one scalar loss.
4. Use finite differences on tiny tensors to catch an implementation agreeing with itself.
5. Require exact router indices before comparing selected mixture-weight gradients.
6. Validate FP32 first, then introduce BF16 or quantized storage one edge at a time.
7. Measure memory, recomputation and speed only after numerical gates pass.

Conclusion

Nemotron is not merely a collection of unusual forward kernels. It is a useful test of whether a runtime can preserve several kinds of mathematical state in both directions. Attention moves gradients through a normalized all-pairs interaction. Mamba2 moves them backward through a recurrence. ReLU2 exposes a simple activation boundary. MoE adds discrete dispatch, selected mixture weights and sparse parameter updates.

Quantization makes the distinction sharper. Quantized inference can reduce the weight traffic that dominates CPU execution. Training still needs stable gradients, reduction contracts and usually higher-precision accumulators. A credible engine must state exactly where precision changes, what is saved, what is recomputed, and what reference decides correctness.

Primary Sources And CKE Artifacts

• Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• NVIDIA ADLR: Nemotron-H architecture
• Hugging Face Transformers: Nemotron-H model contract
• NVIDIA Nemotron 3 Nano model card
• CKE Nemotron v8 circuit
• CKE Mamba2 reference kernels
• CKE routed expert forward and backward kernels
• CKE Mamba2 PyTorch-reference tests
• CKE ReLU2 expert forward/backward tests

Continue Through CKE And The Model Families

This article links equations to the exact CKE circuit, kernel maps, C kernels, and parity tests because the implementation should remain inspectable. The related notes below provide several deliberate entry paths rather than treating this Nemotron investigation as an isolated post.

Start With C-Kernel-Engine

• What Is The C-Kernel-Engine? introduces the complete path from model structure and numerical contracts to generated C and CPU execution.
• Templates Are Circuit Maps explains how CKE makes model-family differences explicit before lowering begins.
• V8 IR Pipeline Codegen follows the architecture through IR, memory planning, kernel binding, and generated C.
• C-Kernel-Engine on GitHub contains the circuits, kernel maps, implementation, tests, and current reports referenced throughout these articles.

Follow The Qwen Model Family

• Qwen3 From A C Runtime Perspective shows how a dense decoder becomes an explicit sequence of normalization, attention, projection, activation, and residual kernels.
• How Qwen3-VL Vision Works extends the same method into patch embeddings, a vision encoder, multimodal projection, and language-model integration.

Deepen The Memory And Kernel Mathematics

• Attention: The Core Of The Transformer derives Q/K/V, softmax, value mixing, backward propagation, and the role of addressable photographic memory.
• KV Cache Memory explains the context-length-dependent storage that hybrid recurrent models are designed to reduce across most layers.
• DeltaNet And SSM compares two ways of carrying bounded learned state through a sequence.
• K-Quants Deep Dive explains how low-bit storage becomes concrete block layouts, scales, mixed dot products, and CPU kernels.

Read The Nemotron Series In Order

1. Nemotron Architecture From A CKE Runtime Perspective: the complete hybrid architecture and runtime surface.
2. Nemotron Is Not One Circuit: exact layer policy, deterministic routing, expert execution, and CPU memory behaviour.
3. Nemotron Kernels From Inference To Backpropagation: the forward equations, reverse equations, precision boundaries, source code, and current evidence gates.