GStar Assignment 2: Build Kimi Delta Attention

How can a model remember a long context with a fixed-size state? In this assignment, you will build the forward kernels step by step: start with linear attention, add forgetting and prediction-error correction, then combine your work into chunked Kimi Delta Attention (KDA).

Like Assignment 1, the starter code supplies the GPU plumbing. You implement the masked arithmetic in PyTorch and Triton. Only the forward pass is required.

Start here

1. Read the overview or open the introduction slides.
2. Follow Problems 1–9 in order, reading each PDF before editing its code.
3. Implement the STUDENT IMPLEMENTATION REQUIRED regions and run the checker.

Learning path

Problem	Implement	Reading
P1	Sequential and chunked linear attention in PyTorch	PDF
P2	Recurrent linear attention in Triton	PDF
P3	Chunked linear attention	PDF
P4	Forgetting: gated linear attention	PDF
P5	Prediction-error correction: DeltaNet	PDF
P6	Chunked DeltaNet and the triangular inverse	PDF
P7	Combine forgetting and correction: Gated DeltaNet	PDF
P8	Recurrent KDA and stable scores	PDF
P9	Assemble the complete chunked KDA forward pass	PDF

Later problems reuse your earlier kernels. In particular, P9 combines the P4 prefix scan, P6 triangular inverse, and P8 score builder. P4 uses a supplied score baseline until you replace it with your P8 implementation.

Setup

Use Python 3.10+. P1 runs on CPU, including macOS. P2–P9 require Linux, Triton, and an NVIDIA CUDA GPU.

cd GStar-Assignment-2-KDA
python -m pip install -r requirements.txt
python autograder.py --check-templates

Implement and check

Replace each NotImplementedError or tl.static_assert(False, ...) in the student regions. Keep the function signatures and region markers intact. The unfinished starter code is expected to fail numerical checks.

# Start with P1 on CPU.
python autograder.py --p1 --device cpu

# On a CUDA GPU, check each problem as you finish it.
python autograder.py --p2   # use --p3 through --p9 for later problems

# Check the complete assignment, then benchmark KDA after correctness passes.
python autograder.py --all --dtype bfloat16
python autograder.py --p9 --benchmark

The checker compares outputs and final states against independent references. Its numerical tolerances are provisional; final grading criteria come from the course staff. Use the required Triton kernels and reuse earlier student work; do not call the checker references from your implementations.

For shapes, precision, and kernel conventions, consult the implementation reference. The formulas omit query scaling; use scale=1.0 to reproduce the written examples.

Submit

Submit problem_1.py through problem_9.py with the explanation comments requested in each problem. See the verification and submission guide for details.