Two Bucket Sidebar

A side-by-side memory bucket sidebar — SRAM tile on top, HBM buffer bars below — that contrasts a naive attention pass (which materialises the full n × n scores matrix) against a flash pass (which tiles Q / K / V through SRAM and never writes the scores matrix at all). Toggle the mode and the sequence length; the HBM bars rescale and the quadratic Scores bar visibly collapses out of view under flash mode.

Sequence lengthn = 128
SRAM(fast, small)

current tile (4×4)

HBM(slow, big)
total: 96 KB
Q16 KB
K16 KB
V16 KB
O16 KB
Scores (n²)32 KB

Naive attention materialises the full 128×128 scores matrix in HBM — 32 KB. That grows quadratically with sequence length.

Customize
Mode
Shape
128
64

Installation

npx shadcn@latest add https://craftbits.dev/r/two-bucket-sidebar.json

Usage

import { TwoBucketSidebar } from "@craft-bits/viz/two-bucket-sidebar";
 
<TwoBucketSidebar />

Drive mode + sequence length from outside:

<TwoBucketSidebar
  mode={mode}
  onModeChange={setMode}
  n={n}
  onNChange={setN}
/>

Custom buffer set with a KV-cache bar that only shows under flash:

<TwoBucketSidebar
  buffers={[
    { key: "Q", label: "Q" },
    { key: "K", label: "K" },
    { key: "V", label: "V" },
    { key: "O", label: "O", color: "var(--cb-success)" },
    { key: "KV cache", label: "KV cache", color: "var(--cb-info)", onlyOnMode: "flash" },
    { key: "scores", label: "Scores (n²)", color: "var(--cb-error)", onlyOnMode: "naive" },
  ]}
/>

Override the byte calculator (fp32, custom head dim):

<TwoBucketSidebar
  d={128}
  elemBytes={4}
  bytesFor={(key, { n, d, elemBytes }) =>
    key === "scores" ? n * n * elemBytes : n * d * elemBytes
  }
/>

Understanding the component

  1. Two stacked buckets. The SRAM mini heat grid sits on top of the HBM buffer bars. The mode toggle decides whether SRAM is lit (flash) or dark (naive), and whether the Scores bar joins the HBM stack.
  2. Bar widths are normalised per render. Each bar's width is max(0.04, bytes / maxBytes) where maxBytes is the largest visible bar in the current (mode, n) configuration. Adding or removing the Scores bar rescales every other bar in lock-step.
  3. AnimatePresence on the Scores bar. Switching from naive → flash collapses the Scores bar's height to zero and fades it out; the remaining bars expand to fill the freed visual space. Reduced-motion users get an instant swap.
  4. Phase machine. The root carries data-phase="naive" | "flash" for styling hooks in surrounding UI.
  5. Pluggable byte calculation. Pass bytesFor to override the default Flash Attention math (n·d·elemBytes for vector buffers, n·n·elemBytes for anything keyed on "scores"). d and elemBytes are passed through to the calculator so the simple default uses them too.
  6. Reduced motion. Under prefers-reduced-motion: reduce the mode-pill colour transition, the bar-width spring, the SRAM tile fades, and the Scores collapse all reduce to instant.

Props

PropTypeDefaultDescription
mode / defaultModeTwoBucketSidebarMode"naive"Controlled / uncontrolled compute mode.
onModeChange(mode) => voidFires when the mode changes.
n / defaultNnumber128Controlled / uncontrolled sequence length.
onNChange(n) => voidFires when the sequence length changes.
nOptionsreadonly number[][64, 128, 256, 512]Sequence-length pills.
dnumber64Per-head dimension used by defaultBytesFor.
elemBytesnumber2Bytes per scalar (fp16 = 2, fp32 = 4).
tileSizenumber4Side length of the SRAM mini tile grid.
buffersreadonly TwoBucketSidebarBufferBar[]flash-attention defaultsHBM buffer bars in render order.
bytesFor(key, { mode, n, d, elemBytes }) => numberflash-attention mathResolve a buffer's byte size.
formatBytes(bytes) => string1024-aware "B / KB / MB"Override the bytes formatter.
renderNarration(args) => ReactNodenaive vs flash copyOverride the bottom narration.
transitionTransitionSPRINGS.snapOverride the spring used for bar fills, tile fades, and the Scores collapse.

Accessibility

  • The mode and n selectors are role="radiogroup" blocks with aria-labels; each pill is role="radio" with aria-checked. Colour is always paired with the data-state, focus ring, and a text label so it is never the only signal.
  • The SRAM mini heat grid is aria-hidden="true" — it is purely decorative and mirrors the mode pill that already carries the screen-reader-readable label.
  • The HBM "total" byte count and the bottom narration are aria-live="polite" regions so updates announce naturally as the mode or n changes.
  • Every pill clears a ≥ 32 × 32 px hit area (32 × 56 for the mode pills, 32 × 40 for the n pills) — meets the Fitts target-size bar.
  • Motion respects prefers-reduced-motion: reduce — the pill colour transition, the SRAM tile fades, the bar-width spring, and the Scores collapse all reduce to instant.
  • Numbers throughout the sidebar use tabular-nums so byte readouts and n values do not shift width as they update.

Credits

  • Extracted from: craftingattention (app/src/lessons/primitives/viz/TwoBucketSidebar.tsx). The source imported TogglePill from the lesson UI kit, hand-rolled the bar-fill spring as SPRINGS.snappy (not present on the library motion module), and hard-coded its colour vocabulary to var(--color-accent-400) / var(--color-warn-400) / var(--color-success-400) / var(--color-fail-500) / var(--color-ink-200) / var(--color-surface-elevated). The viz extract strips every lesson-only import, remaps every inline colour to the --cb-accent / --cb-success / --cb-error / --cb-warning / --cb-fg-* / --cb-bg-* token vocabulary, swaps SPRINGS.snappy for the canonical SPRINGS.snap, generalises the five hard-coded buffers (Q, K, V, O, Scores) into a buffers prop with onlyOnMode visibility, lifts the byte math to a bytesFor callback (with flash-attention defaults), exposes controlled+uncontrolled mode and n props so an outer scrubber can drive the sidebar, and rebuilds the mode and sequence-length pills as role="radio" controls with explicit radiogroups, focus rings, ≥ 32 × 32 hit areas, and reduced-motion-aware animations.