class RMSNorm(nn.Module):
    def forward(self, x):
        return F.rms_norm(x, (x.size(-1),))

class MLP(nn.Module):
    def __init__(self, dim, mult=4):
        super().__init__()
        hidden = int(dim * mult)
        self.fc = nn.Linear(dim, hidden, bias=False)
        self.proj = nn.Linear(hidden, dim, bias=False)
    def forward(self, x):
        return self.proj(F.leaky_relu(self.fc(x), negative_slope=0.5).square())

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, num_kv_heads=4):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = dim // num_heads
        self.c_q = nn.Linear(dim, dim, bias=False)
        self.c_k = nn.Linear(dim, num_kv_heads * self.head_dim, bias=False)
        self.c_v = nn.Linear(dim, num_kv_heads * self.head_dim, bias=False)
        self.proj = nn.Linear(dim, dim, bias=False)
    def forward(self, x):
        B, T, D = x.shape
        q = self.c_q(x).reshape(B, T, self.num_heads, self.head_dim)
        k = self.c_k(x).reshape(B, T, self.num_kv_heads, self.head_dim)
        v = self.c_v(x).reshape(B, T, self.num_kv_heads, self.head_dim)
        q = F.rms_norm(q, (q.size(-1),))
        k = F.rms_norm(k, (k.size(-1),))
        q = q.transpose(1, 2)
        k = k.transpose(1, 2).repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
        v = v.transpose(1, 2).repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.proj(y.transpose(1, 2).reshape(B, T, D))

class Block(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.attn_norm = RMSNorm()
        self.mlp_norm = RMSNorm()
        self.attn = Attention(dim)
        self.mlp = MLP(dim)
    def forward(self, x):
        x = x + self.attn(self.attn_norm(x))
        x = x + self.mlp(self.mlp_norm(x))
        return x

class Model(nn.Module):
    def __init__(self, vocab_size=4096, dim=512, num_layers=11):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, dim)
        self.blocks = nn.ModuleList([Block(dim) for _ in range(num_layers)])
        self.norm = RMSNorm()
        self.head = nn.Linear(dim, vocab_size, bias=False)
    def forward(self, x, y):
        h = self.tok_emb(x)
        h = F.rms_norm(h, (h.size(-1),))
        for block in self.blocks:
            h = block(h)
        h = self.norm(h)
        logits = self.head(h)
        return F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))

model = Model().cuda().bfloat16()
compiled = torch.compile(model, dynamic=False, fullgraph=True)
x = torch.randint(0, 4096, (32, 2048), device='cuda')
y = torch.randint(0, 4096, (32, 2048), device='cuda')
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    compiled(x, y).backward()

# No loop, no accumulators, no workspace — one thread block per row
def triton_per_fused__fused_rms_norm_backward_8(
    in_out_ptr0, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, in_ptr5,
    out_ptr1, out_ptr2, xnumel, r0_numel, XBLOCK):
    # xnumel=98304, r0_numel=512, persistent R0_BLOCK=512
    # 7 loads → ~120 regs/thread, fits in 128 budget
    tmp0 = tl.load(in_ptr0 + ...)   # [98304, 512] rsqrt * Hessian
    tmp1 = tl.load(in_out_ptr0 + ...)  # [98304, 512] upstream grad
    tmp9 = tl.load(in_ptr1 + ...)   # [98304, 512] residual
    tmp10 = tl.load(in_ptr2 + ...)  # scalar: mix coefficient
    tmp20 = tl.load(in_ptr3 + ...)  # [98304, 1] rsqrt
    tmp25 = tl.load(in_ptr4 + ...)  # [512] norm weight 1
    tmp28 = tl.load(in_ptr5 + ...)  # [512] norm weight 2
    # ... 1 inner reduction (sum over 512), 3 stores

# RSPLIT loop with persistent accumulators and workspace memory
def triton_per_fused__fused_rms_norm_backward__to_copy_..._3(
    in_out_ptr0, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    in_ptr5, in_ptr6, in_ptr7, in_ptr8, in_ptr9,
    out_ptr0, out_ptr1, out_ptr3, out_ptr4, ws_ptr,
    xnumel, r0_numel, XBLOCK, RSPLIT_SIZE, NUM_STAGES):
    # 11 loads + 2 accumulators → ~200 regs/thread, EXCEEDS 128 budget
    accum0 = tl.full([R0_BLOCK], 0, tl.float32)  # [512] persists across loop
    accum1 = tl.full([R0_BLOCK], 0, tl.float32)  # [512] persists across loop
    for _ in tl.range(0, split_size, XBLOCK):
        tmp0 = tl.load(in_ptr0 + ...)     # [98304, 512]
        tmp1 = tl.load(in_ptr1 + ...)     # [98304, 512]
        tmp7 = tl.load(in_ptr2 + ...)     # [98304, 512]
        tmp13 = tl.load(in_ptr3 + ...)    # [98304, 512]
        tmp14 = tl.load(in_out_ptr0 + ...)# [98304, 512]
        tmp22 = tl.load(in_ptr4 + ...)    # scalar
        tmp32 = tl.load(in_ptr5 + ...)    # [98304, 1]
        tmp37 = tl.load(in_ptr6 + ...)    # [512]
        tmp40 = tl.load(in_ptr7 + ...)    # [512]
        tmp43 = tl.load(in_ptr8 + ...)    # [98304, 512]
        tmp46 = tl.load(in_ptr9 + ...)    # [98304, 512]
        # 3 inner reductions + 5 stores + 2 accumulator updates per iter
        # Spilled regs hit local memory EVERY iteration
    tl.store(ws_ptr + ..., accum0, ...)  # workspace for inter-block reduction
    tl.store(ws_ptr + ..., accum1, ...)

all_reads = {dep.name for all subnodes' reads if MemoryDep}
if len(all_reads) > max_reads:
    return False