Implement a numerically stable online-softmax attention pass (FlashAttention recurrence) i

Super-hard

Transformer Architecture Internals7 / 7

Chapter 017 / 7

Transformer Architecture Internals

Implement a numerically stable online-softmax attention pass (FlashAttention recurrence) i

Super-hard

Implement a numerically stable online-softmax attention pass (FlashAttention recurrence) in numpy that never materializes the $n\times n$ score matrix.

Target interface

flash_attention(Q, K, V, causal=True, bq=64, bk=64)

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.

solution.pyPython

Local draft

import numpy as np

def flash_attention(Q, K, V, causal=True, bq=64, bk=64):
    raise NotImplementedError

⌘/Ctrl↵to submit

AI review

No review yet

Submit your solution for a verdict, a score out of 100 and concrete feedback. Any correct approach passes.

Chapter 017 / 7

Transformer Architecture Internals

Implement a numerically stable online-softmax attention pass (FlashAttention recurrence) i

Super-hard

Implement a numerically stable online-softmax attention pass (FlashAttention recurrence) in numpy that never materializes the $n\times n$ score matrix.

Target interface

flash_attention(Q, K, V, causal=True, bq=64, bk=64)

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.