Why Did My Attention Layer Allocate So Much Memory?

Lesson

Replace explicit scores, causal masking, softmax and values with scaled_dot_product_attention. Preserve scale=1/sqrt(64), dropout=0 and causal self-attention. Do not pass a second incompatible mask. Float32 inputs and backend selection are deliberate: FlashAttention is not guaranteed. BACKEND_OPERATOR_EVIDENCE lists observed attention operators from a separate diagnostic capture; an empty list means no attention operator was observed, not a guessed backend. Compare peak allocation, elapsed time and numerical tolerance at sequence lengths 128, 256 and 512. Diagnostic profiling is outside the judged interval.

Measurement contract: the real read-only driver defines WORKLOAD. Two warmups precede five paired trials with alternating baseline/candidate order. Each trial uses fresh deterministic inputs. Synchronized host time includes solve and completion, excludes input construction and reference verification. METRICS reports every sample, median and IQR; memory includes persistent buffers and allocator state. One additional solve is isolated in judged for Trace. No speedup is promised. Run the unchanged starter first.

What you will learn

• Identify the bottleneck from the appropriate evidence.
• Preserve the output contract on fresh inputs.
• Compare paired measurements and numerical error.

Starter code

    return baseline(case)