Marlin INT4 on Ada — how it works

TinyLoop's default INT4 GEMM kernel is a simple per-output-element scalar implementation (cuda/int2_gemm.cu::int4_gemm). It is correct and small and portable, but it caps at around 6.8 TFLOPS regardless of batch size because each output cell is computed by a single thread — bandwidth-bound on weight reads with no reuse across outputs.

On Ampere and Ada GPUs (SM 8.0 / 8.6 / 8.9) there is a much faster alternative: upstream Marlin (Frantar 2024, Apache 2.0), a tensor-core INT4 × FP16 mixed-precision GEMM that keeps weights INT4 in global memory, dequantizes per-tile into shared memory, and issues mma.sync instructions to the tensor cores. On an RTX 4090 it runs at 142 TFLOPS peak — roughly 86 % of the card's FP16 tensor-core peak (≈165 TFLOPS) and 21× faster than the naive kernel at production batch sizes.

This page explains what TinyLoop does to use Marlin unmodified on asymmetric GPTQ weights, how the runtime dispatches the right path automatically, and what the measured end-to-end story looks like.

:::note Hardware scope Upstream Marlin targets Ampere and Ada only. On Hopper (SM 9.0) the dispatcher short-circuits and falls back to the native INT4 path or the FP16 body cache. See the upstream README for the details; the TinyLoop wrapper calls kernels::marlin_available() before dispatching and the runtime toggle degrades gracefully when the kernel is unavailable. :::

The core mismatch: symmetric vs asymmetric

Upstream Marlin's kernel assumes symmetric offset-binary quantization:

W_marlin[k, n] = (q_int[k, n] − 8) · scale[g, n]

with q_int ∈ [0, 15] stored as 4-bit cells and scale per-group FP16. TinyLoop's GPTQ converter (tools/gptq_convert.py::gptq_quantize) produces asymmetric min-max weights:

W_tinyloop[k, n] = q_int[k, n] · scale[g, n] + zero[g, n]

with an extra per-group-per-column zero offset. Feeding TinyLoop weights directly into Marlin gives nonsense output because the kernel implicitly subtracts 8 · scale where TinyLoop intended to add zero. The residual per cell is

W_tinyloop − W_marlin = zero + 8 · scale   =:  δ[g, n]

and the matmul output difference collapses — because δ depends only on (g, n) and not on k — into a tiny per-group correction:

C_desired[m, n] = C_marlin[m, n] + Σ_g δ[g, n] · (Σ_{k in group g} A[m, k])
               = C_marlin[m, n] + (A_group_sum @ δ)[m, n]

The correction is a [M, groups] × [groups, N] matmul layered on top of Marlin's output. For our production shapes (M ≤ 256, groups ≤ 16, N ≤ 8192) it is ~1 % of Marlin's own compute, dominated by kernel-launch latency rather than compute.

This is what lets TinyLoop share upstream Marlin's vendored kernel file bit-for-bit — no zero-point support needed inside the kernel, no fork to maintain.

The pipeline, end to end

                  [load time, once per model]
.tinyloop file   →  QuantWeight { data, scale, zero, … }          // TinyLoop native
                  ↓                                               // (q, scale, zero)
                  ↓  tl.pack_model_for_marlin(model)               // tools/marlin_repack.py
                  ↓                                                //
                  ↓    marlin_pack_asymmetric(q_int, scale, zero)  //
                  ↓      ├─ reconstruct w_sym = (q_int − 8) · scale
                  ↓      ├─ call marlin_pack_symmetric
                  ↓      │    ├─ torch.round(w_sym / s) → INT4 codes
                  ↓      │    ├─ tile-reshape + permute + _PERM gather
                  ↓      │    └─ pack 8 × int4 → uint32 cells
                  ↓      └─ δ = zero + 8 · scale    // per group, per column
                  ↓
QuantWeight   ≣  { …,   marlin_B,   marlin_scales,   marlin_delta,   marlin_workspace }

                  [inference, every quant_gemm call]
quant_gemm(w, X, Y, BT, N, K):
    if marlin_dispatch_enabled()
       && w.has_marlin_pack()
       && marlin_available()
       && BT ≥ 16 && K % 128 == 0 && N % 128 == 0:
        gs = (w.group_size == K) ? -1 : w.group_size      // per-row → per-col
        rc = marlin_mm_asymmetric(                         // include/kernels.h
            X, w.marlin_B, Y, w.marlin_scales, w.marlin_delta,
            w.marlin_workspace, BT, N, K, gs
        )                                                  // two kernels:
                                                            //   1) marlin_int4_tc_gemm
                                                            //   2) marlin_asymm_correction
        if rc == 0:  return
    // fall through to native INT4 / FP16 body cache path

Three moving parts:

marlin_pack_asymmetric in tools/marlin_repack.py — Python, runs on GPU via torch at load time. Reuses the symmetric packer so we don't re-implement the _PERM / _SCALE_PERM permutations. Returns a (B_packed, S_packed, δ) triple.
marlin_asymm_correction in cuda/marlin_asymm_correction.cu — two small CUDA kernels with a static cached scratch buffer:
- a_group_sum_kernel reduces A[M, K] in groupsize-sized chunks down to a_sum[M, groups]. One block per (m, g) cell, 128-thread tree reduction.
- correction_accum_kernel computes C[m, n] += Σ_g a_sum[m, g] · δ[g, n] with one thread per (m, n) output cell.
The scratch buffer grows monotonically across calls (doubling), so steady-state has no cudaMalloc overhead.
marlin_mm_asymmetric in the same file — wraps Marlin's symmetric GEMM + the correction into one call. Zeros C first, returns Marlin's rc so the dispatcher can fall back on CALL_IF gaps.

How the dispatcher chooses

src/inference.cpp::quant_gemm is the single entry point every block calls for its attn_qkv / attn_out / mlp_down GEMMs (and the equivalent inside quant_swiglu for mlp_gate / mlp_up — with a caveat, below). The branch order is:

FP16 body cache (w.fp16_cache != nullptr) — memory-for-speed trade; wins on Hopper where HBM bandwidth is abundant.
Marlin + correction — when all six preconditions hold:
- w.bits == 4 && !exact_int4
- marlin_dispatch_enabled() — TINYLOOP_USE_MARLIN=1 or tl.set_marlin_dispatch(True)
- w.has_marlin_pack() — pack ran at load time
- kernels::marlin_available() — Ampere/Ada at runtime
- BT ≥ 16 — M-axis CALL_IF floor for non-per-col shapes
- K % 128 == 0 && N % 128 == 0 — Marlin tile size multiples
Cross-iter stochastic rounding paths — frac4 or sign_bit channels if the .tinyloop file carries them and TINYLOOP_CROSS_ITER_ROUND=1 is set.
Native INT4 — int4_gemm, the fallback that always works.

The (w.group_size == K) ? -1 : w.group_size translation in step 2 is worth calling out. TinyLoop encodes per-row / per-column quant as group_size == K so there's one group of size K. Upstream Marlin's CALL_IF table doesn't instantiate a template for the group_blocks = K / thread_k case when group_blocks > 8; it uses the groupsize == -1 sentinel for per-column instead. That one-line translation was the last bug standing between the pack and a working e2e run — the first version silently fell back to native on every call because Marlin returned rc = 2.

Why there are two kernels, not one fused kernel

The initial fused design (one kernel per output tile, reading A in-shared + reducing + accumulating correction + writing C) was slower than the two-kernel design at production shapes. Reason: each output N-tile redundantly reduced A[m, :] for every one of the N / TILE_N n-tiles, burning ~N / TILE_N× more work on the reduction stage than the two-pass design where a_sum is computed once.

The current design:

Two kernel launches per Marlin call (for the correction), ≈ 5 µs each on SM_89.
Static cached scratch so cudaMalloc cost amortizes to zero.
Compute is sub-microsecond — launch latency dominates.

Further halving is possible via CUDA Graphs replay (bake the Marlin + correction pair into a single graph launch per shape). Deferred because the 31–71 % relative overhead it adds over Marlin's 11–66 µs is still much cheaper than the 6.8 TFLOPS naive ceiling we'd otherwise be stuck at.

End-to-end result on RTX 4090

tests/test_marlin_e2e.py (a CTest target) loads a real .tinyloop INT4 model, A/B's Model.score() on a 1 K-token sequence with the dispatch flag off and on, and verifies the two paths agree.

Measured on model_2b_int4_nosvd.tinyloop (per-row INT4, D=2048, FFN=8192, pre=2, loops=8):

Path	Model.score() PPL (1 K tokens)	Notes
Native INT4 (`int4_gemm`)	8.9755	Baseline
Marlin + asymmetric correction	8.9157	Δ 0.0597 = 0.67 % PPL drift

15 weights packed (2 pre × 5 linears + 1 loop × 5 linears)
Marlin VRAM overhead: ~101 MB
All 15 weights hit the Marlin branch on every forward pass, zero fallbacks
First dispatch logged on stderr: [marlin] first dispatch ok: BT=1024 N=6144 K=2048 gs=-1

The 0.67 % drift is below the FP16-accumulation noise floor for a full multi-layer × 8-loop-iter forward pass. Both runs are valid samples from the same quantized model; neither is "correct" relative to the other, analogous to how TINYLOOP_EXPERIMENTAL_FP16_BODY=1 produces a different-but-valid sample from the INT4 reference.

Parity validation chain

Four CTest targets cover the pipeline at increasing levels of integration. All PASS on SM 8.9:

Target	Layer	Shapes	Check
`tinyloop.test_marlin_parity`	symmetric packer	3	Marlin vs FP16 ref, symmetric only
`tinyloop.test_marlin_asymmetric_parity`	Python correction reference	4	`apply_marlin_correction` matches `A @ w_fq`
`tinyloop.test_marlin_asymm_correction_kernel`	CUDA correction kernel	5	CUDA matches Python ref within 1–2 ULPs
`tinyloop.test_marlin_mm_asymmetric`	combined C++ wrapper	7	`marlin_mm_asymmetric` on shapes up to M=64 N=2048 K=8192 (64 groups)
`tinyloop.test_marlin_e2e`	full runtime	—	`Model.score()` parity with dispatch flag on vs off

All five tests skip cleanly (exit 77) when their preconditions are unmet — no model file, Marlin not linked, pybind11 missing, etc. — so CI on non-GPU or non-Ada runners stays green.

Gaps and follow-ups

quant_swiglu is not yet Marlin-dispatched. The fused SwiGLU dispatch (mlp_gate + mlp_up in one kernel) uses its own path that still takes the naive INT4 kernel. Separate follow-up — would need either a Marlin pair dispatch or an unfused fallback mode.
Head projection (final vocab-size GEMM) is not packed by pack_model_for_marlin — the vocab dimension (V = 50257) does not divide 128 cleanly, so the precondition fails and it would always fall back anyway.
Hopper (SM 9.0). The dispatcher gate is in place but there is no backend. The cleanest path would be vendoring gptq-marlin (the vLLM fork), which has both a WGMMA path for Hopper and native asymmetric-scale support — but that is a significantly bigger vendoring effort (~several thousand template-heavy lines) than the single-file upstream Marlin. Not on today's roadmap.
CUDA Graphs for the correction kernel pair, to eliminate the ~10 µs launch-latency floor on small-BT decode. Optional; current overhead is already acceptable.

References

Upstream Marlin: github.com/IST-DASLab/marlin — kernel source is vendored at cuda/marlin/ with its Apache 2.0 license.
tools/marlin_repack.py — the Python-side pack + correction reference, plus pack_model_for_marlin driver.
cuda/marlin_wrapper.cu — our namespaced C++ entry point (kernels::marlin_int4_tc_gemm).
cuda/marlin_asymm_correction.cu — the correction kernel pair and combined wrapper.
Benchmark Comparison — the measured throughput and accuracy numbers.
Environment Variables — the TINYLOOP_USE_MARLIN=1 toggle documentation.

The core mismatch: symmetric vs asymmetric​

The pipeline, end to end​

How the dispatcher chooses​

Why there are two kernels, not one fused kernel​

End-to-end result on RTX 4090​

Parity validation chain​

Gaps and follow-ups​

References​