### Bug description
Chaining pre-norm transformer blocks (self-attention + cros…s-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
```python
"""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).