From 28efb4395cfcedcba9974f94c2724b8cb2eae55e Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 17 Oct 2025 19:32:38 +0800 Subject: [PATCH] multiout at every level --- test/test_rangeify.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/test/test_rangeify.py b/test/test_rangeify.py index 2737bb714d..6c7331488c 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -52,22 +52,25 @@ class TestPcontig(unittest.TestCase): target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize() GlobalCounters.reset() - attn = q.scaled_dot_product_attention(k, v) + #print("****\n\n\n\n\n") + attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward() attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1) out = attn_output(attn) loss = (out - target).square().mean() loss.backward() + ret = [out, Tensor.stack(q.grad, k.grad, v.grad)] + Tensor.realize(*ret) + return ret + Tensor.realize(out, q.grad, k.grad, v.grad) - return q,k,v + return out, q.grad, k.grad, v.grad with Context(PCONTIG=2, DEBUG=2): - q,k,v = fa_bw() - grads = q.grad, k.grad, v.grad + grads= fa_bw() print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS") with Context(DEBUG=2): - q,k,v = fa_bw() - cmp_grads = q.grad, k.grad, v.grad + cmp_grads = fa_bw() print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS") with Context(DEBUG=0):