Skip to content

docs: fix how-it-works inaccuracies, drop choose-a-kernel page, sweep stale claims#108

Merged
h-tonywu merged 7 commits into
mainfrom
fix-documentation
Jun 10, 2026
Merged

docs: fix how-it-works inaccuracies, drop choose-a-kernel page, sweep stale claims#108
h-tonywu merged 7 commits into
mainfrom
fix-documentation

Conversation

@tonywu71

@tonywu71 tonywu71 commented Jun 9, 2026

Copy link
Copy Markdown
Collaborator

Context

After v0.4.2 we reviewed the docs line-by-line against the kernel sources. Several pages had drifted: how-it-works.html said the S tile lives in SRAM (it's in fp32 registers via tl.dot), described a backward-autotuning scheme predating the block-tiled lowmem grad-D kernel, and had a §4 memory chart whose axis coordinates didn't match the formula. design.md still referenced num_consumer_groups warp-spec configs removed in v0.2.0 and quoted SMEM prune budgets before Triton's 8 KiB scratch reserve was subtracted. The kernel picker (choose-a-kernel.html) predated the maxsim dispatch consolidation; its only code snippet raised TypeError as written.

This PR fixes every wrong claim and removes the picker. The constexpr appendix was also tightened for clarity: each paragraph now makes one point, and the figure states both sides of the trade-off.

Note

  • CHANGELOG history entries mentioning the kernel picker are left as-is (they record past releases); the removal is logged under [Unreleased].
  • Measured benchmark numbers are untouched; only the prose around them changed.

Validation

  • uv run --extra dev ruff check . && ruff format --check . clean
  • uv run --extra dev pytest: 262 passed, 311 skipped (CUDA auto-skip on CPU)
  • SVG coordinates in both reworked figures (§4 memory chart, constexpr loop figure) were recomputed from the formulas, not eyeballed

Changes

Click here to expand
  • how-it-works.html · SRAM vs registers — §5 claimed the S tile "lives in SRAM"; tl.dot accumulates S in fp32 registers, with SRAM holding only the Q/D operand tiles. Prose, animation legend, and step-log strings now agree.
  • _autotune.py · S tile budget commentprune_forward's comment said "the fp32 S accumulator stays in one slot" without saying where, implying SRAM. Updated to clarify S lives in registers and the bq×bd×4 term is a conservative upper bound, consistent with the above.
  • how-it-works.html · backward autotuning — the section now covers both regimes: grad-Q/unified kernels tune over num_warps/num_stages (BWD_CONFIGS); the lowmem grad-D kernel is block-tiled like the forward (one program per slab × BLOCK_D-row doc-tile, tuned over forward_configs()). The §7 lowmem bullet and push/pull figure now say tile-of-rows.
  • how-it-works.html · §4 memory chart — recomputed from bytes(L) = 64·64·L²·2 (fp16): even x-ticks 0–2048, y-axis 0–40 GB at 5.5 px/GB, curve points every 128 tokens (34.4 GB at L=2048). The "80 GB H100" line moved to a legend note.
  • how-it-works.html · smaller fixes — Appendix A register bandwidth corrected to ~80 TB/s; §1 now reads "64 queries × 64 documents (4,096 scored pairs)"; Appendix C accounts for the Lq > 4096 pass-through; §6 credits online-max exactness to associativity/commutativity rather than idempotence.
  • how-it-works.html · constexpr appendix — prose restructured so each paragraph makes one point; the figure now shows both sides of the trade-off (one-binary-for-any-Ld on the loop row, baked q_start constants under the unrolled copies, recompile-price line).
  • Removed choose-a-kernel.html along with the README section, nav link, and preview asset; no dangling references remain.
  • design.md — dropped the num_consumer_groups warp-spec claim (removed in v0.2.0 when Triton 3.5 dropped the API); long-query chunking yields one shared autotune entry (pinned by test_compile_cache.py); SMEM prune budgets corrected to the usable 220/156 KiB (raw 228/164 minus Triton's 8 KiB scratch reserve).
  • Memory tables (README, benchmarks.md) — the Nq=16, Nd=32 fused-fwd cell is 2 KB (Nq·Nd·4 B), not 64 KB; the naive column is now described as the measured allocator peak, which sits above the S-tensor formula because of fp32-cast operand copies.
  • README fused D-head headline — now 0.94–4.5×: two small LateOn shapes come in at 0.94–0.95×, with wins of 1.4× and up from Nd=128, Ld=1024.
  • Test hygiene — deleted dead pytestmark_cuda/pytestmark_no_fp8 module variables (every test carries its own marker); removed three from __future__ import annotations imports; test_pylate_compat_warnings.py now monkeypatches package_at_least → False to keep fake-pylate tests exercising the patch path once pylate ≥ 1.5.1 is in the lockfile; comment accuracy fixes in test_kd_pairs_fast.py and test_fp8.py.
  • examples/basic.py — rerank demo passes normalize=True so raw randn inputs don't fire the unnormalized-embeddings warning.

tonywu71 added 4 commits June 9, 2026 21:28
- §5 + animation: the S tile and running max accumulate in fp32
  registers, not SRAM; SRAM holds only the Q/D operand tiles
- §7 / backward autotuning: describe both regimes — per-row kernels
  tuned over num_warps/num_stages vs the block-tiled lowmem grad_D
  kernel reusing the forward config pool; lowmem programs own a
  BLOCK_D-row tile, not a single row
- Appendix A: register bandwidth ~80 TB/s so the table matches the
  'four times faster than SRAM' prose
- §4 chart: recompute the curve from bytes(L) = 64·64·L²·2 with honest
  axes (0–2048 tokens, 0–40 GB); replace the off-scale 80 GB line with
  a legend note (fp16 OOMs near L ≈ 3125 on 80 GB)
- §1: '64 pairs' → 64 queries × 64 documents (4,096 scored pairs)
- Appendix C: note the Lq > 4096 pass-through path; §6: online max is
  exact by associativity/commutativity, not idempotence
The interactive kernel picker predates the API consolidation: with
maxsim dispatching on layout and the patchers / native PyLate &
colpali-engine backends covering the framework paths, a decision tree
over many entry points no longer reflects the library (and its only
code snippet was broken). Remove the page, its README section, nav
link, and preview asset; note the removal in the changelog.
- design.md: drop the num_consumer_groups warp-spec claim (configs lost
  those kwargs in v0.2.0 when Triton 3.5 removed the API); long-query
  chunking collapses onto exactly one autotune entry (mask flag is out
  of the key, pinned by test_compile_cache.py); SMEM prune budgets are
  the usable 220/156 KiB (raw 228/164 minus the 8 KiB Triton reserve)
- README + benchmarks.md memory tables: the Nq=16,Nd=32 fused-fwd cell
  is 2 KB (Nq·Nd·4 B), not 64 KB; clarify that the naive column is the
  measured allocator peak, which exceeds the S-tensor formula because
  fp32-cast operand copies coexist with it
- README: honest fused D-head range 0.94-4.5× (two small LateOn shapes
  are slightly slower; ≥1.4× from Nd=128, Ld=1024 up)
…etup

- delete dead pytestmark_cuda / pytestmark_no_fp8 module variables
  (every test already carries its own marker; the misnamed variables
  never affected collection)
- drop 'from __future__ import annotations' from three test files
  (repo targets Python >= 3.10)
- test_kd_pairs_fast.py: neutral wording for the forward-parity
  comment; the backward comment now points at the real parity bar
  (test_backward.py, rel < 1e-2) instead of test_forward.py
- test_fp8.py: comment matches the 5% assert
- test_pylate_compat_warnings.py: monkeypatch package_at_least to False
  so the fake-pylate warning tests keep exercising the patch path once
  the installed pylate reaches 1.5.1 (native-LIK no-op guard reads dist
  metadata, not sys.modules)
- examples/basic.py: pass normalize=True in the rerank demo so raw
  randn inputs don't fire the unnormalized-embeddings UserWarning
@tonywu71 tonywu71 self-assigned this Jun 9, 2026
@tonywu71 tonywu71 added the documentation Improvements or additions to documentation label Jun 9, 2026
@tonywu71 tonywu71 requested a review from h-aurelien-lac June 9, 2026 19:34
@h-aurelien-lac

Copy link
Copy Markdown
Collaborator

Nice cleanup — I recomputed the memory table and the §4 chart numbers and they all check out. Lowering the fused D-head headline to 0.94× is the kind of honesty we never see, appreciated.

One thing to reconcile with #109 though: this PR says the S tile never touches SRAM (tl.dot accumulates straight into registers), but #109's new _smem_bytes() comments say "the fp32 S accumulator stays in one slot" and counts bq·bd·4 bytes of SMEM for it. One of the two needs rewording.

Also heads up: both PRs add an ## [Unreleased] heading at the same spot in the CHANGELOG, so whichever merges second gets a conflict.

Minor: the strict Conventional Commits + no-AI-attribution rules in AGENTS.md are policy changes, not doc drift — probably worth a quick team ack rather than slipping them into a docs PR.

tonywu71 added 3 commits June 10, 2026 09:37
The section opened with an unexplained metaphor, used 'pipelines'
without saying what pipelining buys, and packed the unroll argument,
its cost, and the bucketing fix into one paragraph. Split it so each
paragraph makes one point and define terms on first use.
The figure showed how the two compilations differ but not why it
matters: add the one-binary-for-any-Ld chip to the loop row, the baked
q_start constants under the unrolled copies, and the recompile price
line, so the figure states both sides of the bargain.
…ve bound

The prune_forward comment said "the fp32 S accumulator stays in one
slot" without saying where — implying SRAM, which contradicts the
how-it-works.html claim that tl.dot accumulates S straight into the
register file. Both are correct: the kernel never writes S to SRAM, but
the budget formula includes one bq×bd×4-byte slot as a conservative
over-estimate to avoid OOM on configs where register pressure might
otherwise spill.
@tonywu71 tonywu71 force-pushed the fix-documentation branch from 3d7e459 to 2f4fc8f Compare June 10, 2026 07:38
@h-tonywu h-tonywu merged commit 57a05c1 into main Jun 10, 2026
8 checks passed
@h-tonywu h-tonywu deleted the fix-documentation branch June 10, 2026 07:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants