| # 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) |
|
|
| 1. 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. |
| 2. 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. |
| 3. Section numbering offset between M1 and M2/M3 (see above). |
| 4. 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. |
| 5. Series subtitle wording still differs ("Batch M1/M2", M3 has descriptive text). |
| 6. 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). |
|
|
