GEMM Memory Layout & Offset Calculation

GEMM vs GEMV

// Prefill: M tokens at once
gemm(activations[32, 896], weights[896, 4864], output[32, 4864])

// Decode: 1 token at a time
gemv(activation[1, 896], weights[896, 4864], output[1, 4864])

Layout Notation: NN, NT, TN, TT

The two-letter suffix indicates how matrices A and B are stored in memory:

Memory Layout Visualization

Normal vs Transposed Storage

// Memory layout for 896→4864 linear
W[0][0], W[0][1], W[0][2], ..., W[0][4863]  // Input 0 to all outputs
W[1][0], W[1][1], W[1][2], ..., W[1][4863]  // Input 1 to all outputs
...
W[895][0], W[895][1], ..., W[895][4863]    // Input 895 to all outputs

// Memory layout for 896→4864 linear
W[0][0], W[0][1], W[0][2], ..., W[0][895]   // All inputs → Output 0
W[1][0], W[1][1], W[1][2], ..., W[1][895]   // All inputs → Output 1
...
W[4863][0], W[4863][1], ..., W[4863][895]  // All inputs → Output 4863

Quantized Block Structure

Quantized formats store weights in blocks with a scale factor. The block structure affects offset calculations.

struct block_q8_0 {
    ggml_fp16_t d;      // 2 bytes: FP16 scale factor
    int8_t qs[32];      // 32 bytes: 32 quantized weights
};  // Total: 34 bytes for 32 weights

// Dequantization: weight[i] = qs[i] * d

Quantization Format Reference

Offset Calculation

row_offset = n × (K / block_size) × bytes_per_block

// Example: Q8_0 weights, K=896, accessing row 5
blocks_per_row = 896 / 32 = 28 blocks
row_bytes = 28 × 34 = 952 bytes
row_5_offset = 5 × 952 = 4760 bytes

size_t get_weight_row_offset(
    int n,              // output neuron index
    int K,              // input dimension
    int block_size,     // e.g., 32 for Q8_0
    int bytes_per_block // e.g., 34 for Q8_0
) {
    int blocks = K / block_size;
    return (size_t)n * blocks * bytes_per_block;
}

def get_weight_row_offset(
    n: int,              # output neuron index
    K: int,              # input dimension
    block_size: int,     # e.g., 32 for Q8_0
    bytes_per_block: int # e.g., 34 for Q8_0
) -> int:
    blocks = K // block_size
    return n * blocks * bytes_per_block

Row Bytes by Format (K=896)

* K-quants require K aligned to 256; 896 = 3.5 super-blocks (handled with partial blocks)

v7 Training: IR3 Memory + Parallel Execution Plan

What axis should GEMM split on?

Minimal execution-plan schema

{
  "schema": "ck.train.exec.v1",
  "runtime": {"threads": 12, "mode": "deterministic"},
  "ops": [
    {
      "op_id": 37,
      "phase": "forward",
      "kernel_id": "gemm_fwd_f32",
      "shape": {"m": 16, "n": 1024, "k": 1024},
      "dispatch": {
        "split_axis": "m",
        "tile_m": 4,
        "tile_n": 128,
        "threads": 12,
        "schedule": "static"
      },
      "reduction": {"type": "none"}
    },
    {
      "op_id": 109,
      "phase": "backward",
      "kernel_id": "gemm_backward_f32",
      "shape": {"m": 1024, "n": 1024, "k": 16},
      "dispatch": {"split_axis": "k", "threads": 12},
      "reduction": {
        "type": "sum",
        "partial_buffer": "tmp.partial.dw.L0.wq",
        "order": "fixed_tree"
      }
    }
  ]
}

Design rule: IR3 owns scheduling decisions, codegen emits fixed calls, kernels stay math-local.

Critical Bug: Row Stride Mismatch

// Generator and kernel agree
row_bytes = (K/32) × 34 = 952

Token 0: offset 0
Token 1: offset 952     ✓ Kernel reads here
Token 2: offset 1904    ✓ Kernel reads here

// Generator adds spurious +1
row_bytes = (K/32 + 1) × 34 = 986

Token 0: offset 0
Token 1: offset 986     ✗ Kernel expects 952!
Token 2: offset 1972    ✗ Kernel expects 1904!

Symptom: First token works fine, but all subsequent tokens produce NaN or garbage. The validation system in codegen catches this by verifying dimensions align to block sizes.

Kernel Naming Convention

C-Kernel-Engine kernel names follow this pattern:

gemm_nt_q5_0_q8_0
│    │  │     │
│    │  │     └─ Activation quantization (Q8_0)
│    │  └─────── Weight quantization (Q5_0)
│    └────────── Layout: A=Normal, B=Transposed
└─────────────── Operation type (GEMM)

Common Kernel Examples

Related Documentation