[BUG] Deterministic numerical corruption compounding across chained transformer blocks on GPU (bf16)

Bug description

Chaining pre-norm transformer blocks (self-attention + cross-attention + SwiGLU FFN, bf16, 2-token sequence) in a single graph produces deterministic numerical corruption that compounds with depth, while every block verifies in isolation (cos 0.9999 vs a PyTorch reference) and 1–2-block prefixes are near-clean. This is independent of the ops.rms_norm bug reported in modular/modular#6883 — the repro below uses a manually composed fp32 RMSNorm throughout.

With the random-weight repro below (expected cos ~0.9999 at every depth; a torch-bf16-vs-fp32 oracle stays at 0.9999 through 13 such blocks):

depth 1: cos vs torch = 0.999986
depth 2: cos vs torch = 0.973392
depth 3: cos vs torch = 0.969994
depth 4: cos vs torch = 0.881777

With a real bf16 checkpoint's weights in the same structure the divergence is catastrophic: a 13-block chain forks at block 2 (cos drops from 0.9999 to ~0.21 in one block) and produces NaN by block 9 — while the identical blocks pass individually, and per-layer taps show the torch bf16 reference tracking a fp32 oracle at 0.9999 the whole way. Swapping implementation details (fp32 vs bf16 attention scores, ops.softmax vs manual exp/sum, ops.split vs slice_tensor) does not change the failure, which points at compilation (fusion/scheduling/memory planning) rather than any individual op.

The corruption is deterministic across runs and unchanged under MODULAR_DEBUG_DEVICE_ALLOCATOR=uninitialized-poison (unlike the rms_norm issue).

Found while porting a custom transformer model; the repro is self-contained (random weights).

Repro

"""Minimal repro: deterministic corruption/NaN when chaining >=3 transformer blocks.

Structure per block (pre-norm, all bf16 weights, random init):
    x = x + o_proj(SDPA_self(rms(x)))          # 2-token sequence, 12 heads x 128
    x = x + xo_proj(SDPA_cross(rms(x), K, V))  # 6-token memory, K/V precomputed
    x = x + down(silu(gate(rms(x))) * up(rms(x)))
where rms() is a manually composed fp32 RMSNorm (cast->mean->rsqrt->mul->cast)
to avoid the separate builtin ops.rms_norm bug (see repro_rms_norm.py), and
attention is explicit matmul + fp32 softmax.

Result vs a torch bf16 reference on identical inputs/weights (this random
init; expected ~0.9999 at every depth — torch bf16 vs a fp32 oracle stays
at 0.9999 through 13 such blocks):
    depth 1: cos 0.999986   depth 2: 0.973392   depth 3: 0.969994   depth 4: 0.881777
With a real bf16 checkpoint's weights in the same structure the divergence is
catastrophic (cos ~0.1-0.2 from depth 3, NaN by depth 9) while every 1-2-block
prefix and every block in isolation verifies at cos 0.9999. Output is
deterministic across runs and unchanged under
MODULAR_DEBUG_DEVICE_ALLOCATOR=uninitialized-poison.

Observed on MAX 26.6.0.dev2026081105, H100. Reproduces identically on the
supported CUDA 13 path (cuda-compat 580.178.04) and the
MODULAR_NVPTX_COMPILER_PATH fallback, ruling out driver/ptxas.

Usage: python repro_block_chain.py [depth ...]   (default: 1 2 3 4)
"""

import sys

import torch
import torch.nn.functional as F

from max.driver import Accelerator, CPU, Buffer
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType, Weight, ops

B, T, D, HEADS = 1, 2, 1536, 12
HD = D // HEADS
INNER = 4 * D
MEM_T, MEM_D = 6, 1024
SCALE = HD ** -0.5
BF16, F32, GPU = DType.bfloat16, DType.float32, DeviceRef.GPU()

torch.manual_seed(0)
NB = 4
reg: dict[str, torch.Tensor] = {}
for i in range(NB):
    for n, shape in [("qkv", (3 * D, D)), ("ao", (D, D)), ("xq", (D, D)), ("xo", (D, D)),
                     ("xk", (D, MEM_D)), ("xv", (D, MEM_D)),
                     ("gate", (INNER, D)), ("up", (INNER, D)), ("down", (D, INNER))]:
        reg[f"b{i}.{n}_t"] = (torch.randn(shape, dtype=torch.bfloat16) * 0.02).T.contiguous()
    for n in ("n1", "nx", "n2"):
        reg[f"b{i}.{n}"] = torch.randn(D, dtype=torch.bfloat16).abs() * 0.2 + 0.9
x_in = torch.randn(B, T, D, dtype=torch.bfloat16)
mem_in = torch.randn(B, MEM_T, MEM_D, dtype=torch.bfloat16)


def build(depth: int) -> Graph:
    ws = {n: Weight(n, BF16, tuple(t.shape), GPU) for n, t in reg.items()}

    def w(g, n):
        return g.add_weight(ws[n])

    def rms(g, x, wn):
        x32 = ops.cast(x, F32)
        xn = x32 * ops.rsqrt(ops.mean(x32 * x32, axis=-1) + 1e-5)
        return ops.cast(xn, BF16) * w(g, wn)

    def attend(q, k, v):
        q = ops.permute(ops.reshape(q, (B, T, HEADS, HD)), [0, 2, 1, 3])
        s = ops.matmul(ops.cast(q, F32), ops.permute(ops.cast(k, F32), [0, 1, 3, 2]))
        p = ops.cast(ops.softmax(s * SCALE, axis=-1), BF16)
        return ops.reshape(ops.permute(ops.matmul(p, v), [0, 2, 1, 3]), (B, T, D))

    with Graph(f"chain{depth}", input_types=[
        TensorType(BF16, (B, T, D), device=GPU),
        TensorType(BF16, (B, MEM_T, MEM_D), device=GPU),
    ]) as g:
        x, mem = (i.tensor for i in g.inputs)
        for i in range(depth):
            k = ops.permute(ops.reshape(ops.matmul(mem, w(g, f"b{i}.xk_t")), (B, MEM_T, HEADS, HD)), [0, 2, 1, 3])
            v = ops.permute(ops.reshape(ops.matmul(mem, w(g, f"b{i}.xv_t")), (B, MEM_T, HEADS, HD)), [0, 2, 1, 3])
            h = rms(g, x, f"b{i}.n1")
            qkv = ops.matmul(h, w(g, f"b{i}.qkv_t"))
            q, sk, sv = ops.split(qkv, [D, D, D], axis=2)
            sk = ops.permute(ops.reshape(sk, (B, T, HEADS, HD)), [0, 2, 1, 3])
            sv = ops.permute(ops.reshape(sv, (B, T, HEADS, HD)), [0, 2, 1, 3])
            x = x + ops.matmul(attend(q, sk, sv), w(g, f"b{i}.ao_t"))
            h = rms(g, x, f"b{i}.nx")
            x = x + ops.matmul(attend(ops.matmul(h, w(g, f"b{i}.xq_t")), k, v), w(g, f"b{i}.xo_t"))
            h = rms(g, x, f"b{i}.n2")
            x = x + ops.matmul(ops.silu(ops.matmul(h, w(g, f"b{i}.gate_t"))) * ops.matmul(h, w(g, f"b{i}.up_t")),
                               w(g, f"b{i}.down_t"))
        g.output(x)
    return g


def torch_ref(depth: int) -> torch.Tensor:
    x = x_in.clone()
    for i in range(depth):
        wt = lambda n: reg[f"b{i}.{n}_t"].T
        k, v = (F.linear(mem_in, wt(n)).view(B, MEM_T, HEADS, HD).transpose(1, 2) for n in ("xk", "xv"))
        h = F.rms_norm(x, (D,), reg[f"b{i}.n1"], 1e-5)
        q, sk, sv = (y.view(B, T, HEADS, HD).transpose(1, 2) for y in F.linear(h, wt("qkv")).chunk(3, dim=-1))
        a = F.scaled_dot_product_attention(q, sk, sv, scale=SCALE)
        x = x + F.linear(a.transpose(1, 2).reshape(B, T, D), wt("ao"))
        h = F.rms_norm(x, (D,), reg[f"b{i}.nx"], 1e-5)
        q = F.linear(h, wt("xq")).view(B, T, HEADS, HD).transpose(1, 2)
        a = F.scaled_dot_product_attention(q, k, v, scale=SCALE)
        x = x + F.linear(a.transpose(1, 2).reshape(B, T, D), wt("xo"))
        h = F.rms_norm(x, (D,), reg[f"b{i}.n2"], 1e-5)
        x = x + F.linear(F.silu(F.linear(h, wt("gate"))) * F.linear(h, wt("up")), wt("down"))
    return x


def cos(a: torch.Tensor, b: torch.Tensor) -> float:
    a, b = a.double().flatten(), b.double().flatten()
    return float((a @ b) / (a.norm() * b.norm()))


def main() -> None:
    depths = [int(d) for d in sys.argv[1:]] or [1, 2, 3, 4]
    device = Accelerator()
    session = InferenceSession(devices=[device])
    for depth in depths:
        model = session.load(build(depth), weights_registry=reg)
        got = torch.from_dlpack(
            model.execute(Buffer.from_dlpack(x_in).to(device), Buffer.from_dlpack(mem_in).to(device))[0].to(CPU())
        )
        print(f"depth {depth}: cos vs torch = {cos(got, torch_ref(depth)):.6f}")


if __name__ == "__main__":
    main()

System information

• MAX 26.6.0.dev2026081105 (pip nightly)
• NVIDIA H100 80GB HBM3, Ubuntu 22.04.5
• Reproduces identically on two driver setups: kernel driver 570.172.08 + cuda-compat-13-0 580.178.04 forward-compat libs (native CUDA 13 path), and the MODULAR_NVPTX_COMPILER_PATH=/usr/local/cuda-12.8/bin/ptxas fallback — so driver/ptxas are ruled out.
• Compile cache cleared between runs (rm -rf ~/.cache/modular).
• torch 2.11.0+cu128 (reference computation + DLPack buffer staging only).

Start by running the embedded repro with depths 1–4 and compare the compiled output with torch_ref using the reported cosine values. Trace the graph through compilation, fusion, scheduling, and memory planning, since the issue reports the corruption survives operator substitutions and allocator poisoning. Done means chained bf16 blocks match the reference across depth without divergence or NaNs.