Style alignment: align/notes and align/notebooks vs Flash-Attention-101

Reference: /apps/development/rlhf/reference-repos/Flash-Attention-101 (notes/0000-0008, 009, 010; README.md; scratch-pad.ipynb). Provenance caveat: the reference PDFs are handwritten notes whose title page (0000, p.1) is signed "Flash Attention from first principles, Umar Jamil". The repo (README clone URL viai957/Flash-Attention-101) packages them with the user's own README, notebooks and kernels. The "house style" below is therefore the handwritten-lecture style plus the repo's layout, not a typeset template. Where typeset equivalents are needed I mapped them (see "Mapping").

1. Reference style rules (page citations are PDF page numbers)

Layout and naming

R1. Numbered, one-topic-per-file notes named NNNN - Title Case.pdf (0001 - Multi Head Attention, 0003 - Online Softmax, ...). 0000 is a Table of Contents (0000 p.2 "Topics" list, p.3 "Prerequisites"). Numbering is not perfectly uniform (009, 010 have three digits).
R2. Length 2-9 pages per note (0002: 4, 0003: 6, 0004: 9, 0008: 9, 009/010: 2). One idea per page.
R3. Notebook: a single scratch-pad.ipynb that goes naive first (cells 0-9 "Conventional Softmax": loop for max, loop for normaliser, loop for division, then one consolidated typed function), then the improved version ("Safe Softmax" cell 10 on: single fused loop), then block/grid scaffolding for the Triton kernel. Markdown cell = pseudocode, code cell = direct transcription, then a consolidated function with a docstring. No asserts; checks are by printing. Kernel code lives in triton/ and cuda/; README maps folder -> role.
R4. README: Overview, Folder Structure, "Understanding ...", "Implementation Details" in the order PyTorch naive -> Triton -> CUDA, Benchmarking, Setup.

Page structure and tone

R5. Question-first opening. Each note starts with a bare problem statement as a question, then the pseudocode that fails: 0003 p.1 "The problem statement is: can we find a way to compute the softmax of a vector without going through the vector 3 times, but also preventing the exponential from exploding?"; 0004 p.2 "WHY SHOULD WE CARE?"; 0002 p.1 "But there's a problem! If the values of the vector are large, the exponential will explode!"; 0008 p.1 states the object ("Derivative: scalar input, scalar output") first.
R6. Motivation -> failing naive version -> fix -> proof. 0003: problem (p.1) -> try fusing (p.2) -> step-by-step trace on the same tiny vector x = [3, 2, 5, 1] (pp.2-4) -> "Can we prove this algorithm is correct? Let's do it!" (p.5) -> induction with base case and step, each closed by a marker (p.6).
R7. One running toy example reused on every page (x = [3,2,5,1]; the 8x8 blocked Q,K,V of 0004), with explicit shapes written next to every object ((N,d), (d,N), "ORIGINAL (8,128) / BLOCK (4,128)", 0002 p.1, 0004 pp.2-5).
R8. Pseudocode appears next to the derivation, is restated in full each time it is modified, and is labelled "Pseudocode" / "New pseudocode" (0003 pp.1, 4).
R9. Figures carry the argument: grid pictures of blocked matrices, computation graph for autograd (0008 pp.6-7), memory-hierarchy figure (0000 p.1). Colour has a fixed role: blue = given/derivation, red = warning or the open question, purple = annotation/shape, green = "correct" / the question to prove.
R10. Warnings are explicit and loud: "WRONG!" box with the reason (0004 p.5), "The l3 we computed is wrong! Here it is wrong!" (0003 p.3). A "THE IDEA" boxed summary precedes the fix (0004 p.6).
R11. Tone: second person plural, conversational ("Let's restore it... with a twist!", "Let's multiply by V", "Lots of patience", 0000 p.3). Short sentences, no citations inside notes; the only external reference is a pasted algorithm (FlashAttention-2 Algorithm 1, 0004 p.9) annotated by hand.
R12. Complexity is stated per step (time and memory reads, 0002 p.3).
R13. Notation: m_i (running max), l_i (normaliser), O_i, S = QK^T, P = softmax(S), blocks Q_i, K_j, V_j; softmax* for the unnormalised block softmax (0004 p.4); upstream/local/downstream gradient vocabulary (0008 p.9).

2. Mapping to the typeset notes (what "same style" means here)

Reference	Typeset equivalent used
R1 file naming	already `NNNN - Title.pdf`; kept
R5 question first	`problem` box titled "The problem" whose first line is `\question{...}`
R6 motivation -> math -> proof	Notation, Derivation (proofs inline), unchanged content
R7 shapes and a running toy	the notes already carry toy numbers in `\ver{...}` lines and notation tables; not altered
R10 warnings	`What breaks` box (all 24)
R9 "THE IDEA" boxed summary	`Boxed result` box (all 24)
(flow to next page)	`Next` box (all 24)

3. Gaps found (before the change)

Compared 0107 (M1), 0203 (M2), 0304 (M3) and then grepped all 24.

Three series had three different style files and three different skeletons:

Header: M1 \hdr (no "Note" word, grey series line), M2/M3 \notehead ("Note NNNN -- ...").
Opening: M1 grey "The problem this solves" box; M2 numbered section "The problem this solves"; M3 numbered section "The problem this concept solves". None opened with an explicit question (R5), although many posed one inside the paragraph.
Result: M1 yellow "Result"; M2 blue box; M3 blue box, section called "Boxed final results" (M3) vs "Boxed result" (M1 box with no section, M2).
Failure modes: M1 box "What breaks" under a section of the same name; M2 box with title override; M3 plain section (0301 titled "What breaks (failure modes and the assumption each violates)").
Next: M1 blue box; M2/M3 bold inline line under "Link to the next note" (M2) or none (M3).
References: M2 had \section{References} followed by thebibliography (duplicate heading).
Page: M1 1in margins, 11pt, colour links; M2 2.2cm; M3 2.2cm with Unicode/Lean listings; M2 10pt, M1/M3 11pt.
Macro clashes between the files: \ind (bracketed one-argument in M1, bare symbol in M2), \Pr (operator vs blackboard P), \nextnote (environment vs command).

Not present in any of the 24 (and not added, because it would be new content): hand-drawn style figures (R9), a "WRONG!" worked-trace page (R10), complexity-per-step tables (R12), a Table of Contents note (R1 0000), a prerequisites page.

4. What I changed

Style infrastructure

notes/src/common.sty replaces common-m1.sty, common-m2.sty, common-m3.sty (the three are deleted; originals backed up in /home/vignesh/.claude/jobs/b952890a/tmp/notes_style/backup/). It is the union of the three: geometry 2.2cm, 11pt documents, hidelinks, Cerulean result box, BrickRed breaks box, grey problem box, blue Next box, Lean listings and the Unicode declarations, all shared macros.
\ind is the bare symbol everywhere; the bracketed M1 form is now \indb{...} (renamed in 0101-0110 only; same rendering as before). \Pr is \mathbb{P} in all notes (0301, 0302, 0304, 0305 used the default operator "Pr"; only the glyph changes).
notes/src/build_all.sh rebuilds all 24 (or a subset: ./build_all.sh 0107 0203), takes the PDF file name from the existing ../NNNN - *.pdf. build.sh and build_m1.sh still work unchanged.

All 24 .tex files

\usepackage{common}, \documentclass[11pt].
One header, \notehead{NNNN}{title}{subtitle} (M1's \hdr converted; subtitle text kept: series line plus companion notebook path).
One opening: a grey "The problem" box whose first line is a new one-sentence \question{...} (24 questions, each only restating the question the existing paragraph already poses; no new claim, no number), followed by the original paragraph untouched. For M2/M3 the old numbered section was converted into this box.
One tail: Boxed result section, What breaks section with the red box (M2/M3 explicit title override dropped, the default is the same text), Next box, References (M2/M3 duplicate heading removed).
Section numbers: the box does not take a number in M1 (Notation = 1, Derivation = 2). In M2/M3 the old problem section was section 1, so I inserted \stepcounter{section} after the box so Notation stays 2 and Derivation 3. Reason: many text and notebook cross-references use those numbers (for example 0301 "\S3.3", 0302 notebook "section 3.4", 0304 notebook "section 3", 0203 notebook "section 3", 0102 cited as "Sec.~2.5" from 0105). Renumbering needs a coordinated edit of notes and notebooks and was not done. Consequence: M1 numbers sections from 1, M2/M3 from 2.

Rebuild

24 of 24 PDFs rebuilt with latexmk, 0 failures, 0 undefined references/citations. Page counts: 0103 5->4, 0201 5->6, 0202 5->6, 0204 5->6, 0205 5->6, 0206 5->6, 0207 5->6, 0302 6->7, 0304 6->7, all others unchanged. Text diff of old vs new PDFs (word multiset) shows only header, section-number, heading, box-title and the 24 added The question. lines; no equation, number or verified-result line changed.
notes/src/build/m1/ (stale build artefacts of the old M1 script) was deleted; notes/src/build/all/ is the new build dir.

Notebooks

Not changed. They already follow the reference's naive -> efficient progression with asserts and verification tables (0107 section 7 "Naive vs efficient implementation"; 0203 sections 1-3 naive exact/float vs stable form; 0304 "A1 NAIVE ... A2 EFFICIENT"). The reference notebook has no asserts and a single scratch-pad.ipynb; ours are per-note numbered NNNN-slug.ipynb, which I kept (the numbered naming matches R1 and the .tex/PDF linkage). Nothing was re-executed; no Lean notebook touched.

5. Not aligned (left deliberately)

Handwritten diagram/colour language (R9, R10), worked trace on one toy vector (R7), per-step complexity (R12): would be new content. A reasonable follow-up is a \begin{warning} "WRONG!" box and a 4-element trace per note, authored by hand per topic.
0000 Table of Contents note and a prerequisites page: not requested; a generated 0000 - Table of Contents.pdf from the 24 titles is mechanical if wanted.
Section numbering offset between M1 and M2/M3 (see above).
M3 notes keep their \paragraph{Numerical verification (notebook ...)} block while M1 uses inline \ver{...} tags and M2 mixes both; both carry the same information, and merging them touches content layout.
Series subtitle wording still differs ("Batch M1/M2", M3 has descriptive text).
Notebook per-note files vs reference single scratch-pad: kept.

6. Possible content issues (not edited)

notebooks/0101-autoregressive-likelihood-sft.ipynb cell 2 cites "(note, Sec. 3.1)" for the proper-distribution lemma; in the M1 numbering the lemma sits under Section 2 (Derivation, "Step 2 (Lemma 1)"), so this reference looks stale. Not verified against the notebook text beyond the grep.
No mathematical claim or number was found to be wrong during this pass (this was a style pass; the content was not re-derived).