Deep Dive: LLM Concepts

Position vs Dimension

Both are involved, but differently:

angle(pos, dim) = position * theta_i

Dim pair 0: rotates FAST (local patterns)
Dim pair 31: rotates SLOW (global patterns)

Why Rotation Works

When you compute Q_m · K_n after rotation:

Q_m · K_n = f(Q, K, m - n)

Only the angle difference matters!

• Pos 5 → Pos 3: angle diff = 2*theta
• Pos 100 → Pos 98: angle diff = 2*theta

Same relative distance = same attention behavior.

Split-Half / NEOX RoPE

This layout is common in GPT-NeoX-style implementations. In this engine's current templates, Qwen and Gemma follow this family.

(0, d/2), (1, d/2+1), (2, d/2+2), ...

The vector is split into two halves, then each index in the first half rotates against the matching index in the second half.

Practical upside: it fits naturally with split-vector implementations and long-standing NEOX-style kernels. Practical requirement: it is only correct for checkpoints trained with that layout.

Before training: how would an ML engineer choose?

If you are designing a model family from scratch, this choice is mostly about ecosystem fit and implementation coherence, not a proven universal quality advantage.

• Choose pairwise if you want a Llama-style family that should match llama.cpp, common Llama checkpoints, and adjacent even/odd rotation math.
• Choose split-half if you are building on a GPT-NeoX-style stack or existing kernels, loaders, and training code that already assume first-half/second-half pairing.
• Choose the stack you can support end-to-end: training code, checkpoint export, inference kernels, parity tooling, and debugging all need to agree.

For most teams, the practical decision is: pick the RoPE layout of the ecosystem you want to join. Bigger modeling choices usually matter more for quality than pairwise-versus-split-half by itself.

The Problem

Standard attention materializes the full score matrix:

scores = Q @ K.T  # [N, N] matrix!

N = 2048:  16 MB
N = 8192:  256 MB
N = 32768: 4 GB
N = 131072: 64 GB  # Won't fit in GPU!

Memory is O(N²), which limits context length.

Windowed Causal Mask

For token i:
  attend j in [max(0, i - W + 1), i]

scores[i][j] = -inf if j > i
scores[i][j] = -inf if i - j >= W
        

Compute: O(N * W) vs O(N²). Memory stays O(N * D).

LayerNorm (Original)

mean = sum(x) / n
var = sum((x - mean)²) / n
y = gamma * (x - mean) / sqrt(var + eps) + beta
        

4 operations: mean, center, variance, normalize

Learnable: gamma (scale) and beta (shift)

What CK Has to Preserve

The hard part is not only implementing the math once. The runtime has to preserve state ownership and layout contracts across different layer types:

• Attention: Q/K/V projection, RoPE or no-RoPE policy, KV-cache placement, and output projection.
• Mamba2 / SSM: input projection split, convolution state, dt softplus, recurrent selective state update, and gated normalization.
• MoE: group-limited routing, top-k expert selection, ReLU2 expert execution, and weighted accumulation.

Decode Contract

projected -> split(gate, hidden_bc, dt)
conv_state = roll_append(conv_state, hidden_bc)
dt = softplus(dt + dt_bias)
state = state * exp(dt * A) + dt * B * x
y = C dot state + D * x
out = grouped_rmsnorm(y * silu(gate))
        

The important part is state ownership. CK exposes each step as a named kernel so the runtime can decide where state lives, how it is streamed, and which future optimized kernel replaces the scalar reference.

Group-Limited Top-K Router

scores = sigmoid(router_logits + correction_bias)
groups = partition(experts, n_group)
active_groups = topk(group_scores, topk_group)
experts = topk(scores inside active_groups, top_k)
weights = normalize(scores[experts])
        

The router is not just a matrix multiply. It includes bias correction, group selection, top-k selection, optional normalization, and a scaling factor before expert dispatch.

ReLU and ReLU2

ReLU(x)  = max(0, x)
ReLU2(x) = max(0, x)^2

d ReLU2 / dx = 2x if x > 0 else 0
        

ReLU2 keeps the same negative-value mask as ReLU, but positive activations grow quadratically. The backward kernel must use the original input value, not only a binary mask.

GELU (GPT-2, BERT)

GELU(x) = x * Φ(x)
        ≈ x * sigmoid(1.702 * x)
        

Smooth approximation of ReLU that allows small negative values.

The Frontend Contract

im2patch:        [C,H,W] -> [N, C*P*P]    (N = (H/P)*(W/P))
patch embed:     [N, C*P*P] -> [N, d]     (dual proj + stream merge + bias)
spatial_merge_2x2: group 4 patches -> [N/4, 4*d]
+ tiled 2D pos:  position_embeddings_add_tiled_2d + 2D RoPE
        

Each patch is a P×P tile across all channels, flattened into one row. spatial_merge_2x2 then groups a 2×2 block of patches into a single token using tile_order_index_2d, which is why the encoder works on a coarser grid than the raw patch grid.

1. Multi-section (M-)RoPE — the rotation path

It is the same RoPE math, just applied to two indices. The head dimension is split into sections: one slice of dims is rotated by an angle derived from the row θ(y), another slice by an angle derived from the column θ(x).

text   :  angle = t * base^(-2i/d)          // one index
vision :  y-dims rotate by  y * base^(-2i/d)
          x-dims rotate by  x * base^(-2i/d) // two indices
        

vision_position_ids_2d_merge builds the position buffer as four flattened streams [y | x | y | x], emitted in the same 2×2 merged-tile order the patches use, so a token's coordinates line up with its embedding row. After rotation, attention can recover both vertical and horizontal distance between patches — true 2D relative position.