The 32-Bit Trap: how a 4 GB tensor becomes a wild pointer

0The crash

A megakernel is one kernel launch that runs an entire transformer forward pass (every layer, every op) across all SMs, using grid-wide barriers instead of separate kernel launches. Ours decoded Qwen2.5-0.5B at 163 tok/s on an RTX 4050, 95% of the DRAM roofline, 3.4× HF eager, with output matching HuggingFace token-for-token. Porting it to Qwen2.5-14B on an A100-80GB should have been a matter of changing config constants. Instead:

[2] compiling megakernel...
    compiled in 247.7s (161 grid barriers)
[gpu-health] XID: NVRM: Xid 31, MMU Fault: ENGINE GRAPHICS GPC5 ...
             faulted @ 0x2ac1_d4010000. Fault is of type
             FAULT_PDE ACCESS_TYPE_VIRT_READ
cutlass.base_dsl.common.DSLCudaRuntimeError:
error: CUDA_ERROR_ILLEGAL_ADDRESS (error code: 700)

Note what is not here: no line number, no kernel frame, no op name. The fault surfaced at the host-side self.compiled(*self.args) call, the launch itself, because an illegal address on the GPU only becomes visible to the host at the next synchronization point. Every debugging instinct that starts from the traceback is useless here. That is the first lesson: a GPU fault tells you that a pointer was wild, never which pointer.

The root cause, once found, was one sentence:

Root cause

The DSL computes tensor byte offsets in 32-bit arithmetic. Any single tensor whose byte span exceeds 2³² (4.29 GB) wraps its addressing. Layers past the boundary read through a wrapped, wild, pointer.

To see why that happens, and why it only bites at 14B, you have to follow a weight from the checkpoint file all the way into the kernel. That journey is where the trap is set.

1How a weight gets to the GPU

HuggingFace → torch → one giant stacked tensor

The engine starts life as ordinary HuggingFace. Nothing about this part is exotic, and that matters: the bug is born here, in completely innocent-looking torch code:

from transformers import AutoModelForCausalLM, AutoTokenizer

hf = AutoModelForCausalLM.from_pretrained(model_id, dtype=load_dt).eval()
sd = {k: v.detach().clone() for k, v in hf.named_parameters()}

def st(key):  # stack one weight across all layers
    return torch.stack([sd[key.format(i)] for i in range(NUM_LAYERS)])

W = {
    "q":    st("model.layers.{}.self_attn.q_proj.weight"),
    "gate": st("model.layers.{}.mlp.gate_proj.weight"),
    "up":   st("model.layers.{}.mlp.up_proj.weight"),
    "down": st("model.layers.{}.mlp.down_proj.weight"),
    ...
}
del hf  # free the reference model

torch.stack turns 48 per-layer matrices [N, K] into one contiguous [L, N, K] tensor, and the engine does this once per weight type, so one layer's seven matrices become seven stacked tensors. This is the right thing to do for a megakernel (one pointer, one layout, the layer index is just the outermost coordinate), and it is also exactly where the oversized tensor is created:

w_gate = w_up = [48, 13824, 5120] bf16   =  6.79 GB each
w_down        = [48, 5120, 13824] bf16   =  6.79 GB

layer stride in BYTES = 13824 · 5120 · 2 = 141,557,760

At 0.5B the largest stacked tensor is [24, 4864, 896] bf16 = 209 MB, 27× under the limit. Nothing was ever going to go wrong there, which is why the port looked safe. The code didn't change; the magnitudes did.

The hand-off: from_dlpack

CuTe DSL doesn't copy your weights. It takes them by reference through the DLPack protocol, a standard C struct that wraps any framework's tensor as (data pointer, shape, strides, dtype). One call per tensor, once, at setup:

from cutlass.cute.runtime import from_dlpack

ct = lambda t: from_dlpack(t, assumed_align=16)
self.args = (
    ct(self.mTok), ct(self.mPos), ct(W["embed"]), ct(self.rope_t),
    ct(W["q"]), ct(W["gate_lo"]), ct(W["gate_hi"]), ...
    ct(self.gKC), ct(self.gVC), ct(self.gBar), ct(self.mLogits),
    self.stream,
)
self.compiled = cute.compile(Qwen25Megakernel(), *self.args, n_layers=nl)

The whole lifecycle, in one picture:

  safetensors on disk
        │  AutoModelForCausalLM.from_pretrained
        ▼
  torch tensors (one [N,K] matrix per layer)         host RAM / VRAM
        │  torch.stack + .to("cuda").contiguous() + .to(bf16)
        ▼
  ONE contiguous [L, N, K] tensor per matrix         ← the 5.66 GB giant is born
        │  from_dlpack(t, assumed_align=16)          ← DLPack: zero-copy hand-off
        ▼
  cute.Tensor = base pointer + layout(shape, stride) ← ptr is 64-bit; INDEX MATH is the risk
        │  cute.compile(...)  →  trace → MLIR → LLVM/NVVM → PTX → ptxas → cubin
        ▼
  self.compiled(*self.args)  →  ONE launch           ← base ptr passed as kernel arg;
                                                       offsets recomputed per access
                                                       on device

Three properties of this hand-off matter for the trap:

• The base pointer is always 64-bit. The GPU is a 64-bit address space; ld.global takes a 64-bit address. The pointer is not what truncates.
• The offset is recomputed on device, per access. t[(layer, n, k)] is not a Python index; it's IR that becomes offset = layer·s0 + n·s1 + k·s2, and the bit-width of that arithmetic is the compiler's choice. That choice is the trap.
• Our shapes are constexpr. The megakernel traces with fixed model dimensions, so strides are baked into the IR as constants. The multiply layer · 141,557,760 is emitted as plain 32-bit integer arithmetic, not as the 64-bit dynamic-stride path.

2How CuTe DSL addresses memory

A cute.Tensor is two things and only two things: a base pointer, and a layout, a function from coordinates to offsets. Indexing is layout evaluation:

t[(layer, n, k)]  ──crd2idx──►  offset_elems = layer·s0 + n·s1 + k·s2
                  ──scale───►   offset_bytes = offset_elems · itemsize
                  ──GEP─────►   address      = base_ptr + offset_bytes
                  ──load────►   ld.global  [address]

Following it down to PTX

Each step above is a real IR instruction you can inspect. The Python indexing lowers to cute.crd2idx in the CuTe dialect, then to plain LLVM integer ops, then to NVPTX:

# MLIR (LLVM dialect), the healthy version
%off   = llvm.mul %layer, %stride0   : i64      # coordinate · stride
%addr  = llvm.getelementptr %base[%off] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>
%val   = llvm.load %addr : !llvm.ptr<1>

# PTX: what ptxas sees
mad.wide.s64   %rd1, %r1, %rd2, %rd3;    # 64-bit: base + offset·size
ld.global.b64  %rd4, [%rd1];

The load itself is always 64-bit (mad.wide takes a 64-bit base). The danger is upstream: if the multiply is emitted in 32-bit, the offset wraps before it ever reaches the 64-bit add. And that is not hypothetical: NVIDIA's own bug tracker shows the exact broken pattern. From cutlass#3208 (cute.domain_offset truncating in CuTe DSL 4.5.0):

%13 = llvm.mul  %12, %5        : i64          # offset computed in i64, fine
%14 = llvm.trunc %13          : i64 to i32    # ← then NARROWED to 32 bits
%15 = llvm.sext %14           : i32 to i64    # ← sign-extended back, damage done
%16 = llvm.getelementptr %10[%15] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>

In that reproducer the offset was exactly 2³², the truncation wrapped it to zero, and a kernel meant to write the far end of a 4 GB buffer wrote out[0] instead, silently. Our megakernel hit the same wall from the other side: offsets past 2³² wrapped to nonzero garbage, and the MMU caught the wild read.

So the picture to keep in your head¹:

base pointer          64-bit   ALWAYS safe: the GPU is a 64-bit address space
coordinate (i,j,k)    32-bit   fine: no dimension is 4 billion long
stride                baked    dangerous when stride·coord can exceed 2^32
offset = Σ coord·stride       ← THE TRAP LIVES HERE, in the width of this sum
address = base + offset        64-bit GEP: receives already-wrapped garbage

Nothing in this chain is visible from Python. The tensor reports its true shape; the pointer is correct; only the intermediate arithmetic lies. That is why the kernel can be "provably correct" and still fault.

3The arithmetic of the fault

The byte offset of layer L inside the stacked w_gate tensor is L · 141,557,760. Where does that cross 2³²?

2^32 ÷ 141,557,760 = 30.34   →   layer 31 is the first bad one

Forty layers, and the last nine of them read through a wrapped pointer. But "the offset wraps" is a theory until you can predict a boundary with it. So instead of reading code, we bisected scale: MEGA_NUM_LAYERS runs only the first N of the 40 layers (the output is garbage; this is a crash test, not a correctness test).

Read the table twice; it eliminates every competing theory:

• 2³¹ (2.15 GB) is ruled out. The 24-layer run works with a 3.26 GB offset, which already exceeds 2³¹. So it is not a signed-int32 problem.
• The boundary sits between 3.26 and 4.39 GB. The only power of two there is 2³² = 4.29 GB, and 4.29 GB ÷ 141.5 MB = 30.34 predicts the first failing layer is 31, consistent with 24 passing and 32 failing.
• Everything else was already cleared: 0.5B ran on the same A100 (arch and CTA count fine); 14B with 1 layer ran (every single-layer op fine at 14B dims); the L40S failure was a separate SMEM-capacity issue, not this.

Bisection over inspection was the right call for a reason worth memorizing: scale-dependent bugs are invisible in the source. The same line of kernel code is correct at 24 layers and wrong at 32. No amount of staring at the line reveals that. An hour of inspection found nothing; three bisect runs found the exact boundary.²

4Two fixes that could not work

Two plausible fixes failed first, and both failures teach something more durable than the fix itself.

4.1. The int8 local repro: a false negative

Renting A100s to test a theory is slow and not free. The cheap idea: reproduce the overflow locally on a 6 GB laptop GPU by using int8: same element count, one byte per element, no big tensor required:

L, N, K = 36, 13824, 5120     # 2.55 B elements, int8 = 2.55 GB
# one thread per layer reads w[layer, N-1, K-1]

All 36 layers read back correctly; theory "disproved," an evening wasted. Except the experiment was testing the wrong boundary. With int8, byte offset equals element offset, so the test topped out at 2.55 GB of byte offset and never crossed 4.29 GB. It disproved a 2³¹-element boundary while the real 2³²-byte boundary sat untested above it.

Lesson

When testing an addressing limit, make sure the test actually reaches the byte magnitude of the real workload, not just the element count.

4.2. Reshaping to 4D: mathematically incapable of helping

Second idea: keep every partial index product small by splitting the layer axis into two axes, so no intermediate multiply gets big:

W[k] = W[k].view(2, 20, N, K)              # host
w_all[(layer // 20, layer % 20, n, k)]     # kernel

The partial products do shrink. But the final offset is unchanged; it cannot change, because it must address the same byte of the same buffer:

layer 31 via 3D:  31 · (N·K)               = 2.19 B elements
layer 31 via 4D:   1 · (20·N·K) + 11 · (N·K) = 2.19 B elements   # identical

If the truncation happens when the final offset is materialized, no reshape can save you. Only shrinking the tensor's own byte span works.

Lesson

Reshaping changes index arithmetic, never the resulting address.

5The fix that worked: split the tensor

If the disease is "one tensor's span exceeds 2³² bytes," the cure is to not have such a tensor. Split the oversized ones along the layer dimension into two halves, each under the limit. The elegant part: in a megakernel the layer index is a compile-time constant (cutlass.Constexpr), the loop over layers is unrolled at trace time, so selecting the half costs zero runtime instructions. It is a Python-level if evaluated while tracing:

# host side: AFTER the bf16 cast (order matters, see below)
HL = NUM_LAYERS // 2
for k in ["gate", "up", "down"]:
    W[k + "_lo"], W[k + "_hi"] = W[k][:HL], W[k][HL:]
    del W[k]   # slices keep the storage alive; avoids double-counting bytes

# kernel side: layer is constexpr, so this is resolved at TRACE time
HLh: cutlass.Constexpr = NUM_LAYERS // 2
if cutlass.const_expr(layer < HLh):
    mlp_gate_silu_mc(sN, w_gate_lo, w_up_lo, gG, tid, cta, layer, ...)
else:
    mlp_gate_silu_mc(sN, w_gate_hi, w_up_hi, gG, tid, cta, layer - HLh, ...)

Only three tensors needed splitting. q/k/v/o at 2.05 GB are fine precisely because the boundary is 2³² and not 2³¹, which is why pinning down the exact power of two in §3 mattered before touching code. And the result:

compiled in 247.7s (161 grid barriers)
DECODE: 26.42 ms/token → 37.9 tok/s (951 GB/s achieved), all 40 layers, A100-80GB

The ordering trap inside the fix

One sharp edge, hit on the first attempt: if you split before the bf16 cast, the halves are fp32 views of the pre-cast tensor while W["gate"] gets rebound to a new bf16 tensor. The kernel then tries to store fp32 into a bf16 SMEM tile and you get a type-mismatch error that points nowhere near the real mistake. Cast first, then slice.

And a mechanical one: going from 3 weight args to 6 required edits in four places: the kernel(...) signature, the self.kernel(...) call, the __call__(...) signature, and the host-side self.args tuple. Miss one and you get a silent argument-order shift that surfaces as a bizarre type error somewhere unrelated. Argument marshalling errors never point at the real line; treat any "impossible" type error after an arg-list change as a marshalling bug first.

6You are in good company

This is not an obscure edge we wandered into alone; it is one of the most-reported classes of CuTe DSL bugs, and it has bitten essentially every serious user of the DSL, including NVIDIA's own examples and Tri Dao's kernel library:

Same disease, three different organs: baked constexpr strides (us), dynamic-stride defaults (#2312), and a narrowing in the lowering itself (#3208). The lesson generalizes: in CuTe DSL, treat every offset as guilty of being 32-bit until you have proven the width of the arithmetic that produced it.

The official knobs

Note the trade-off NVIDIA names explicitly: 64-bit addressing costs registers and extra instructions on the hottest path in the kernel, which is why 32-bit was the default and why "just make everything i64" is not the plan. The megakernel's split sidesteps the trade-off entirely: offsets stay small and fast.

7Checklist: porting a DSL kernel to a bigger model

1. Compute every tensor's byte span. numel · itemsize > 2**32 must be split. Do this on paper before you rent anything.
2. Verify per-lane dim coverage for any HEAD_DIM-dependent code: HEAD_DIM // 32 dims per lane, never a hardcoded 2. (At HEAD_DIM=128 a hardcoded pair silently skips dims 64..127 of every head: in-bounds, no crash, wrong output.)
3. Check SMEM against the target arch, not your dev GPU. sm_89 = 99 KB, sm_80 = 164 KB, sm_90 = 227 KB. The 14B kernel needs 142 KB/CTA and cannot run on sm_89 at all.
4. Bisect scale before reading code. MEGA_NUM_LAYERS found the boundary in three runs; inspection found nothing in an hour.
5. Make the cheap test reach the real magnitude: the int8 repro (§4.1) tested element count, not byte magnitude, and lied to us.
6. Validate numerics against HF after it stops crashing. "It runs" is not "it's right"; a silently-wrong attention dim hid behind the crash the whole time.

8Fault-signature reference

References

1. cutlass#2312: "CuteDSL example hits IMA for large tensors due to strides in Int32." Known limitation, 64-bit strides in 4.3.0.dev0, use_32bit_stride knob.
2. cutlass#3208: "CuTe DSL 4.5: cute.domain_offset truncates large dynamic offsets through i32." Found via quack test failures (Tri Dao's DSL kernel library); fixed in 4.5.1. Source of the llvm.trunc MLIR quoted in §2.
3. cutlass#2886: Blackwell fmha_bwd workspace overflows int32 spec.
4. flashinfer#3392: widening cute-DSL norm kernels to int64 for inputs > 2³¹ elements.
5. flashinfer#3230: widening GDN decode pool indices to Int64 against element-offset overflow.
6. CUTLASS docs: Integration with Frameworks: from_dlpack(tensor, assumed_align, use_32bit_stride) semantics and the make_ptr / cute.make_tensor DLPack bypass.
7. Local: .agents/the_32bit_trap.md (full debugging log), megakernels/models/qwen25/engine.py (the pipeline in §1), megakernels/README.md (results table).