Muon is one of the more interesting optimizer ideas showing up in modern LLM training discussions because it does not merely change a learning-rate schedule or add another element-wise moment buffer. It changes the geometry of the update for hidden-layer weight matrices. AdamW looks at each parameter mostly as an element with its own first and second moment statistics. Muon looks at a 2D matrix update and asks whether the update direction can be made more orthogonal, balanced, and matrix-aware.

C-Kernel-Engine does not currently ship a Muon kernel. That is important to say clearly. This post is therefore a research and implementation target: what Muon is, what the update rule does, why it differs from AdamW, and what a clean CKE implementation path would look like. scope This is not an inference kernel. Muon is a training optimizer. It matters when CKE updates weights, not when CKE only runs a fixed checkpoint.

Why Another Optimizer?

Post 34 built the optimizer path from SGD to Momentum to AdamW. That sequence is still the right baseline. SGD gives the simplest possible update. Momentum remembers a smoothed direction. Adam adds per-parameter adaptive scaling. AdamW fixes weight decay by decoupling it from the gradient update.

AdamW became the default because it is robust, easy to tune relative to older methods, and works across many parameter types. But AdamW pays for that robustness with optimizer state. For every weight it usually stores a first moment and a second moment. At scale, optimizer state becomes a major memory object.

Muon asks a different question. Many important neural network weights are not isolated scalars. They are 2D matrices: projection weights, MLP weights, expert weights, and other hidden-layer transforms. If the parameter is a matrix, maybe the update should respect matrix geometry instead of treating the entries as unrelated scalar slots.

The diagram is not saying one optimizer is universally better. It is showing the shape of the update each optimizer believes in. All three methods start from the same object: a gradient computed by backpropagation. The difference is what happens before that gradient becomes a weight update.

This is the core mental shift. SGD asks, “which direction does the gradient point?” AdamW asks, “how should each coordinate be scaled based on first and second moments?” Muon asks, “if this parameter is a matrix, should the update itself be shaped like a matrix operator?” That is why Muon belongs in a kernel-engineering series. The optimizer is no longer just a vectorized element-wise pass. The optimizer now contains matrix multiplications, temporary buffers, dtype decisions, and layout-sensitive performance behavior.

The Core Muon Step

The PyTorch implementation makes the high-level structure clear. For each 2D parameter matrix with gradient G, Muon keeps a momentum buffer B. It computes a momentum or Nesterov-style update, orthogonalizes that update using a Newton-Schulz iteration, applies decoupled weight decay, adjusts the learning rate based on matrix shape, and subtracts the resulting update from the parameter.

for W in two_dimensional_hidden_layer_weights:
    G = W.grad

    B = momentum * B + G
    U = G + momentum * B if nesterov else B

    O = newton_schulz_orthogonalize(U, steps=5)

    adjusted_lr = adjust_lr(lr, W.shape)
    W *= (1.0 - lr * weight_decay)
    W -= adjusted_lr * O

Why Only 2D Parameters?

This point matters. Muon is not meant to replace every optimizer path in the model. The local PyTorch implementation rejects non-2D parameters. The documentation notes that other parameters, such as biases and embeddings, should be optimized by a standard method such as AdamW.

This is a natural fit for transformer hidden layers. The most expensive trainable objects in the model are matrix-shaped: W_Q, W_K, W_V, W_O, MLP up/down/gate projections, MoE expert matrices, and the final LM head. But embeddings, biases, scalar norm weights, and router quirks may need a different optimizer route. 2D Muon is best understood as a matrix optimizer for hidden-layer weights, not a universal replacement for every parameter tensor.

The Newton-Schulz Part

The Newton-Schulz loop is where Muon becomes interesting for kernel engineering. The update matrix is normalized, optionally transposed so the smaller dimension is handled conveniently, and then repeatedly transformed using matrix products. The local PyTorch implementation uses default coefficients approximately:

At a high level, the loop looks like this:

def zeropower_via_newton_schulz(G, a, b, c, steps, eps):
    X = bf16(G)

    if X.rows > X.cols:
        X = X.T

    X = X / max(norm(X), eps)

    for _ in range(steps):
        A = X @ X.T
        P = b * A + c * (A @ A)
        X = a * X + P @ X

    if original_was_transposed:
        X = X.T

    return X

From a kernel perspective, that loop is not mysterious. It is a sequence of matrix multiplies plus scalar combinations. But it is very different from AdamW. AdamW is mostly element-wise vector math. Muon introduces real matrix work inside the optimizer step.

The exact implementation details matter. The common Muon implementation first normalizes the update matrix so the iteration is stable. It may transpose the matrix so the iteration handles the smaller side more efficiently. Then each Newton-Schulz step builds products such as X Xᵀ, A², and P X. For square matrices this is real cubic-ish matrix work. For rectangular transformer matrices the cost depends heavily on which side is smaller and how the implementation chooses its orientation.

Why This Is Interesting For C-Kernel-Engine

CKE already treats training as explicit kernels, saved tensors, generated backward paths, and optimizer updates. Post 34 covered the current optimizer surface: SGD momentum, AdamW, gradient norm clipping, and fused multi-tensor update paths. Muon would add a new kind of optimizer kernel: a matrix-update optimizer instead of a pure element-wise optimizer.

That makes it a good next implementation target because it naturally exercises the parts CKE cares about: matrix layout, scratch buffers, BF16/FP32 behavior, GEMM reuse, deterministic update contracts, optimizer routing, and parity against a known reference.

Proposed CKE Kernel Surface

A C implementation should not begin as one giant monolithic function. It should expose the algorithmic pieces separately so each piece can be tested:

typedef struct {
    float lr;
    float weight_decay;
    float momentum;
    float eps;
    float ns_a;
    float ns_b;
    float ns_c;
    int   ns_steps;
    int   nesterov;
    int   adjust_lr_mode;
} ck_muon_config_t;

void ck_muon_update_matrix_f32(
    float *weight,             // [rows, cols]
    const float *grad,          // [rows, cols]
    float *momentum_buffer,     // [rows, cols]
    float *scratch,             // workspace for X, A, A2, P
    int rows,
    int cols,
    const ck_muon_config_t *cfg);

The scratch pointer matters. Newton-Schulz needs temporary matrices. If the optimizer allocates inside every matrix update, the implementation will become noisy, slow, and hard to reason about. CKE should allocate a training scratch arena once, plan the largest required temporary shape, and reuse it across parameters.

AdamW State vs Muon State

AdamW carries first and second moment buffers. Muon carries a momentum buffer and computes an orthogonalized update. This does not mean Muon is automatically cheaper in wall-clock time. It shifts cost from element-wise state math toward matrix products during the optimizer step.

The practical trade-off is this: AdamW is memory-state heavy and element-wise compute friendly. Muon is matrix-update heavy and may be attractive when the training dynamics improvement is worth the optimizer-side GEMM work. trade-off Muon is not “free speed.” It is a different optimizer geometry. The right benchmark is convergence per token, memory footprint, and total training cost, not only optimizer kernel latency.

Shape-Based Learning Rate Adjustment

Muon implementations usually adjust the learning rate based on matrix shape. The local PyTorch implementation includes two modes. The original mode scales by a ratio involving the rectangular shape. Another mode attempts to match AdamW-style RMS behavior.

This matters for CKE because the optimizer cannot just read a flat parameter array and blindly apply one update rule. It needs the parameter manifest. The manifest must know whether a tensor is 2D, whether it should use Muon, what shape adjustment applies, and whether a fallback optimizer should own that parameter.

Where Muon Fits In The Training Pipeline

Muon sits exactly where AdamW sits: after forward, loss, backward, gradient accumulation, optional gradient clipping, and before zeroing gradients for the next step. It does not change the chain rule. It changes how the final gradient buffer modifies the weights.

for step in training_steps {
    ck_forward(model, batch);
    ck_cross_entropy_loss(model, targets);
    ck_backward(model);

    ck_gradient_allreduce_or_accumulate(model);
    ck_clip_grad_norm_if_needed(model);

    // Existing path:
    // ck_adamw_update_all(model);

    // Proposed mixed optimizer path:
    ck_optimizer_route_update(model, ADAMW_FOR_NON_2D, MUON_FOR_2D);

    ck_zero_grad(model);
}

The mixed optimizer route is the real systems problem. If Muon only supports hidden-layer matrices, then the optimizer pass must split parameters into groups: Muon-owned matrices, AdamW-owned embeddings/bias/norm parameters, and possibly special cases like MoE routers.

This is where a generated-runtime project like CKE has an advantage if it is disciplined. The optimizer router should not be a pile of string checks like “if the tensor name contains wq, use Muon.” It should be a property of the model manifest and lowered training graph. Each tensor should carry enough metadata to answer: is it trainable, is it 2D, is it a hidden-layer matrix, is it sparse-like, does it need weight decay, does it use Muon, and what fallback owns it if Muon does not.

What To Benchmark

A good Muon benchmark is not “does the optimizer step run fast once?” That is too shallow. The optimizer matters because it changes the training trajectory. A serious CKE benchmark should track:

How This Differs From “AdamW But Faster”

A common mistake is to evaluate Muon as if it were trying to be a faster AdamW kernel. That is not the right comparison. AdamW is usually cheap per step relative to the rest of training because the update is mostly element-wise. Muon deliberately adds matrix computation to the optimizer step. The bet is that the changed update geometry can improve the training trajectory enough to justify that extra work.

That means a fair benchmark needs at least three axes: loss per token, loss per second, and total memory/runtime cost. If Muon reaches the same loss with fewer tokens but each step costs more, the question becomes whether the net training run is cheaper. If Muon uses less persistent optimizer state for matrix weights but needs more scratch workspace during the step, the question becomes whether the memory plan is better for the actual machine. This is why optimizer work belongs beside systems work, not above it.

Why Muon Belongs After The AdamW Post

The previous optimizer post explained the default training path. Muon is the right follow-up because it asks a deeper question: should the optimizer know the tensor is a matrix? AdamW says every element gets adaptive scalar statistics. Muon says the update for a matrix can be shaped as a matrix.

That is exactly the kind of topic that belongs in this series. The point of the series is not to memorize formulas. The point is to connect math, tensor shape, kernel implementation, memory, and training behavior. Muon touches all of those layers.

Summary

Muon is best understood as a matrix-aware optimizer for hidden-layer weights. It keeps momentum, optionally uses Nesterov, transforms the update through Newton-Schulz orthogonalization, applies decoupled weight decay, adjusts learning rate based on matrix shape, and updates the weight matrix.

For C-Kernel-Engine, the key is not to chase hype. The key is to translate the optimizer into contracts: which parameters use it, which tensors are saved, which scratch buffers are required, which GEMM paths are reused, and which parity tests prove the update is correct.

That makes Muon a useful next research target. It is not part of CKE yet, but it fits the project’s direction: build the training stack from math to kernels, make every step inspectable, and only optimize after correctness is nailed down.