docs: fix how-it-works inaccuracies, drop choose-a-kernel page, sweep stale claims#108
Conversation
- §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
|
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 ( Also heads up: both PRs add an 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. |
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.
3d7e459 to
2f4fc8f
Compare
Context
After v0.4.2 we reviewed the docs line-by-line against the kernel sources. Several pages had drifted:
how-it-works.htmlsaid the S tile lives in SRAM (it's in fp32 registers viatl.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.mdstill referencednum_consumer_groupswarp-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 themaxsimdispatch consolidation; its only code snippet raisedTypeErroras 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
[Unreleased].Validation
uv run --extra dev ruff check . && ruff format --check .cleanuv run --extra dev pytest: 262 passed, 311 skipped (CUDA auto-skip on CPU)Changes
Click here to expand
how-it-works.html· SRAM vs registers — §5 claimed the S tile "lives in SRAM";tl.dotaccumulates 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 comment —prune_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 thebq×bd×4term 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 overnum_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 overforward_configs()). The §7 lowmem bullet and push/pull figure now say tile-of-rows.how-it-works.html· §4 memory chart — recomputed frombytes(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, bakedq_startconstants under the unrolled copies, recompile-price line).choose-a-kernel.htmlalong with the README section, nav link, and preview asset; no dangling references remain.design.md— dropped thenum_consumer_groupswarp-spec claim (removed in v0.2.0 when Triton 3.5 dropped the API); long-query chunking yields one shared autotune entry (pinned bytest_compile_cache.py); SMEM prune budgets corrected to the usable 220/156 KiB (raw 228/164 minus Triton's 8 KiB scratch reserve).benchmarks.md) — theNq=16, Nd=32fused-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.Nd=128, Ld=1024.pytestmark_cuda/pytestmark_no_fp8module variables (every test carries its own marker); removed threefrom __future__ import annotationsimports;test_pylate_compat_warnings.pynow monkeypatchespackage_at_least → Falseto keep fake-pylate tests exercising the patch path once pylate ≥ 1.5.1 is in the lockfile; comment accuracy fixes intest_kd_pairs_fast.pyandtest_fp8.py.examples/basic.py— rerank demo passesnormalize=Trueso rawrandninputs don't fire the unnormalized-embeddings warning.