diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6b7a708..6f294f2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,6 +13,8 @@ jobs: - run: npm ci - run: npm run typecheck - run: npm run format:check + - run: npm run build:cli + - run: node dist/cli/estimate.mjs chips --model k3 > /dev/null - run: npm run bench - run: npm test -- --exclude tests/bench.test.ts env: diff --git a/README.md b/README.md index fcf9976..00cc5f8 100644 --- a/README.md +++ b/README.md @@ -51,10 +51,34 @@ The explorer ranks sharding configurations per chip, and drills into any one wit fabric view (the chips on their real interconnect) and an execution trace (the op graph, colored by what each op is bound on). +## CLI + +The same engine as a command line, `estimate`, for scripting sweeps and pinning +configurations in other repos. Build it once (`npm run build:cli`, a single +node script at `dist/cli/estimate.mjs`, no runtime dependencies) and: + +``` +estimate models # the presets +estimate chips --model k3 # chips, with the least chips that hold the weights +estimate search k3 gb300-nvl72 32x1 --slo 20 # rank every sharding, streaming as it prices +estimate top k3 b300 8x1 --phase prefill # the same, silent +estimate explain k3 b300 8x1 --sizes TP=8,EP=8 # one sharding, fully worked, with engine flags +estimate sweep --model k3 --chips b300,gb300-nvl72 --machines 8x1,32x1 --slos none,20 --table +``` + +`--json` emits one row per line; `sweep` writes one file per cell and `report` +re-renders them. Two knobs correct the simulator where its defaults mislead, +both off unless asked for: `--state-slots N` (with `--spec-slots`, +`--state-dtype`, `--mem-fraction`) reserves the recurrent-state slots a hybrid +KV manager actually holds per sequence on linear-attention layers, and +`--scheduler dag` reads overlap off the op graph instead of the two `--overlap` +fractions. Every number is an estimate. + ## Layout ``` src/core engine (lowering, placement search, cost backends), model & hardware specs src/ui the explorer — leaderboard, fabric view, execution trace +src/cli the estimate command line (bundled by scripts/build-cli.mjs) tests property / fuzz tests ``` diff --git a/package-lock.json b/package-lock.json index a2d79a0..64f1acf 100644 --- a/package-lock.json +++ b/package-lock.json @@ -16,14 +16,22 @@ "react-dom": "^18.3.1", "recharts": "^2.12.7" }, + "bin": { + "estimate": "dist/cli/estimate.mjs" + }, "devDependencies": { + "@types/node": "^22.20.2", "@types/react": "^18.3.3", "@types/react-dom": "^18.3.0", "@vitejs/plugin-react": "^4.3.1", + "esbuild": "^0.21.5", "prettier": "^3.6.2", "typescript": "^5.5.3", "vite": "^5.3.4", "vitest": "^2.0.4" + }, + "engines": { + "node": ">=20" } }, "node_modules/@babel/code-frame": { @@ -1336,6 +1344,16 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/node": { + "version": "22.20.2", + "resolved": "https://registry.npmjs.org/@types/node/-/node-22.20.2.tgz", + "integrity": "sha512-xlvWf4Vs9n1PEVYwP1n4vvG07M6y8WgvJ2t0vbrWTmijsIHp1cS+uJ2kMIRdY3nHZK0nCYKrPeD171+SzF4/zw==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, "node_modules/@types/prop-types": { "version": "15.7.15", "resolved": "https://registry.npmjs.org/@types/prop-types/-/prop-types-15.7.15.tgz", @@ -2518,6 +2536,13 @@ "node": ">=14.17" } }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + }, "node_modules/update-browserslist-db": { "version": "1.2.3", "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz", diff --git a/package.json b/package.json index 42433d3..0c3e32c 100644 --- a/package.json +++ b/package.json @@ -8,11 +8,12 @@ "build": "tsc --noEmit && vite build", "test": "vitest run", "bench": "vitest run tests/bench.test.ts", - "search": "vite-node scripts/search.ts --", "typecheck": "tsc --noEmit", "preview": "vite preview", "format": "prettier --write .", - "format:check": "prettier --check ." + "format:check": "prettier --check .", + "build:cli": "node scripts/build-cli.mjs", + "cli": "node scripts/build-cli.mjs && node dist/cli/estimate.mjs" }, "prettier": { "printWidth": 100, @@ -32,12 +33,20 @@ "recharts": "^2.12.7" }, "devDependencies": { + "@types/node": "^22.20.2", "@types/react": "^18.3.3", "@types/react-dom": "^18.3.0", "@vitejs/plugin-react": "^4.3.1", + "esbuild": "^0.21.5", "prettier": "^3.6.2", "typescript": "^5.5.3", "vite": "^5.3.4", "vitest": "^2.0.4" + }, + "bin": { + "estimate": "dist/cli/estimate.mjs" + }, + "engines": { + "node": ">=20" } } diff --git a/scripts/build-cli.mjs b/scripts/build-cli.mjs new file mode 100644 index 0000000..621502f --- /dev/null +++ b/scripts/build-cli.mjs @@ -0,0 +1,29 @@ +// Bundle the CLI into one node script: the sources use extensionless +// imports under moduleResolution "bundler", which node cannot run as-is. +import { build } from 'esbuild'; +import { execSync } from 'node:child_process'; +import { mkdirSync } from 'node:fs'; + +const version = (() => { + try { + return execSync('git describe --tags --always --dirty', { stdio: ['ignore', 'pipe', 'ignore'] }) + .toString() + .trim(); + } catch { + return 'unknown'; + } +})(); + +mkdirSync('dist/cli', { recursive: true }); +await build({ + entryPoints: ['src/cli/main.ts'], + outfile: 'dist/cli/estimate.mjs', + bundle: true, + platform: 'node', + format: 'esm', + target: 'node20', + banner: { js: '#!/usr/bin/env node' }, + define: { __CLI_VERSION__: JSON.stringify(version) }, + logLevel: 'warning', +}); +console.error(`built dist/cli/estimate.mjs (${version})`); diff --git a/scripts/search.ts b/scripts/search.ts deleted file mode 100644 index 189589a..0000000 --- a/scripts/search.ts +++ /dev/null @@ -1,73 +0,0 @@ -import { searchShardings } from '../src/core/engine/optimizer/search'; -import { Mesh, roleSize } from '../src/core/engine/surface/deploy'; -import { ROLES, ShardingRole } from '../src/core/engine/sim/ir/sharding/roles'; -import { makeNaiveOpCostSumBackend } from '../src/core/engine/sim/cost/naiveOpCostSum'; -import { CHIPS_BY_ID } from '../src/core/hardware/chips'; -import { MODEL_PRESETS } from '../src/core/model/models'; - -// Streaming search CLI: one line per role-size tuple the moment it -// prices (the tuple's best placement and dispatch), new global bests -// marked with *. Usage: -// npm run search -- [prefill|decode] -// npm run search -- Kimi tpu-v5p 2x4x4 -// npm run search -- "LLaMA 3 70B" h100-sxm 8x2 decode (switched: domain x nodes) - -const [modelArg = 'Kimi', chipArg = 'tpu-v5p', machineArg = '2x4x4', phaseArg = 'prefill'] = - process.argv.slice(2); - -const model = MODEL_PRESETS.find((m) => m.name.startsWith(modelArg)); -const chip = CHIPS_BY_ID[chipArg]; -if (!model || !chip) throw new Error(`unknown model or chip`); - -const machine = chip.interconnect.topologies - ? chip.interconnect.topologies.find((t) => t.name === machineArg) - : (([domain, nodes = 1]) => ({ domain, nodes }))(machineArg.split('x').map(Number)); -if (!machine) throw new Error(`no ${machineArg} in the ${chipArg} catalog`); - -const fmt = (mesh: Mesh, f: (r: ShardingRole, size: number) => string) => - ROLES.filter((r) => roleSize(mesh, r) > 1) - .map((r) => f(r, roleSize(mesh, r))) - .join(' '); - -const chips = Array.isArray(machine) - ? 0 - : 'count' in machine - ? machine.count - : machine.domain * machine.nodes; - -const t0 = performance.now(); -let best = -Infinity; -let total = 0; -for (const step of searchShardings( - model, - chip, - machine, - { prefillLen: 4096, generateLen: 1024 }, - { - costBackend: makeNaiveOpCostSumBackend({ memoryOverlap: 0, commsOverlap: 0 }), - phase: - phaseArg === 'decode' - ? { kind: 'decode' } - : { kind: 'prefill', seqs: chips, mode: 'throughput' }, - }, -)) { - total = step.total; - const c = step.candidate; - if (!c) continue; - const isBest = c.score > best; - if (isBest) best = c.score; - const { mesh } = c.deployment; - console.log( - `[+${(performance.now() - t0).toFixed(0).padStart(6)}ms] ` + - `${step.done.toString().padStart(3)}/${step.total} ${isBest ? '*' : ' '} ` + - `${Math.round(c.score).toString().padStart(6)} tok/s/chip ` + - `${(fmt(mesh, (r, s) => `${r}=${s}`) || 'single chip').padEnd(30)} ` + - `best: ${c.batch !== undefined ? `batch=${c.batch} ` : ''}` + - // with EP = 1 there is nothing to route, so the dispatch is noise - `${roleSize(mesh, 'EP') > 1 ? `${c.deployment.moeDispatch} ` : ''}` + - `${fmt(mesh, (r) => `${r}[${mesh.roles[r].join('')}]`)}`, - ); -} -console.log( - `\n${total} tuples in ${((performance.now() - t0) / 1000).toFixed(1)}s, best ${Math.round(best)} tok/s/chip`, -); diff --git a/src/cli/args.ts b/src/cli/args.ts new file mode 100644 index 0000000..6949c17 --- /dev/null +++ b/src/cli/args.ts @@ -0,0 +1,134 @@ +import { parseArgs, type ParseArgsConfig } from 'node:util'; +import type { Scheduler } from '../core/engine/sim/cost/select'; +import { DEFAULT_SHARED, SharedOptions } from './context'; +import { CliError } from './resolve'; + +type OptionTable = NonNullable; +export type Values = Record; + +// the knobs every evaluating subcommand takes +export const SHARED_OPTIONS = { + 'prefill-len': { type: 'string' }, + 'gen-len': { type: 'string' }, + overlap: { type: 'string' }, + scheduler: { type: 'string' }, + 'cost-per-hour': { type: 'string' }, + 'mem-fraction': { type: 'string' }, + 'state-slots': { type: 'string' }, + 'state-dtype': { type: 'string' }, + 'spec-slots': { type: 'string' }, + json: { type: 'boolean' }, + help: { type: 'boolean', short: 'h' }, +} satisfies OptionTable; + +export const SHARED_HELP = ` +shared options + --prefill-len N prompt tokens per sequence (default 4096) + --gen-len N generated tokens per sequence (default 1024) + --overlap M,C naive scheduler: fraction of memory and comms traffic + hidden behind the widest stream (default 0.9,0.65) + --scheduler naive|dag dag reads overlap off the op graph and ignores --overlap + --cost-per-hour USD override the chip's rental price ($/chip-hour) + --mem-fraction F fraction of HBM the stack lets weights+cache use + (SGLang mem_fraction_static); applied as a smaller chip + --state-slots N recurrent-state slots reserved per sequence on + linear-attention layers (default 1; SGLang radix cache 5) + --state-dtype T capacity dtype of that state: fp32 (default) or bf16 + --spec-slots N extra state slots per sequence for speculative decoding + --json machine-readable output (one JSON object per line) +`.trimEnd(); + +export function parse( + argv: string[], + options: T, + allowPositionals = true, +): { values: Values; positionals: string[] } { + try { + const { values, positionals } = parseArgs({ + args: argv, + options, + allowPositionals, + strict: true, + }); + return { values: values as Values, positionals }; + } catch (err) { + throw new CliError(err instanceof Error ? err.message : String(err), 2); + } +} + +export function num(values: Values, name: string): number | undefined { + const v = values[name]; + if (v === undefined || typeof v === 'boolean') return undefined; + const n = Number(v); + if (!Number.isFinite(n)) throw new CliError(`--${name} expects a number, got "${v}"`, 2); + return n; +} + +export function str(values: Values, name: string): string | undefined { + const v = values[name]; + return typeof v === 'string' ? v : undefined; +} + +export function flag(values: Values, name: string): boolean { + return values[name] === true; +} + +export function oneOf( + values: Values, + name: string, + allowed: T, + dflt: T[number], +): T[number] { + const v = str(values, name) ?? dflt; + if (!allowed.includes(v)) + throw new CliError(`--${name} must be one of ${allowed.join('|')}, got "${v}"`, 2); + return v as T[number]; +} + +export function list(values: Values, name: string, dflt: string[] = []): string[] { + const v = str(values, name); + return v === undefined + ? dflt + : v + .split(',') + .map((s) => s.trim()) + .filter(Boolean); +} + +export function sharedOptions(values: Values): SharedOptions { + const overlapArg = str(values, 'overlap'); + const [memoryOverlap, commsOverlap] = overlapArg + ? overlapArg.split(',').map(Number) + : [DEFAULT_SHARED.overlap.memoryOverlap, DEFAULT_SHARED.overlap.commsOverlap]; + if ( + overlapArg && + ![memoryOverlap, commsOverlap].every((x) => Number.isFinite(x) && x >= 0 && x <= 1) + ) + throw new CliError(`--overlap expects two fractions "M,C", got "${overlapArg}"`, 2); + const scheduler: Scheduler = oneOf(values, 'scheduler', ['naive', 'dag'] as const, 'naive'); + + const statePool: SharedOptions['statePool'] = {}; + const slots = num(values, 'state-slots'); + if (slots !== undefined) statePool.slotsPerSeq = slots; + const spec = num(values, 'spec-slots'); + if (spec !== undefined) statePool.specSlots = spec; + const dtype = str(values, 'state-dtype'); + if (dtype !== undefined) { + const bytes = { fp32: 4, bf16: 2, fp16: 2, fp8: 1 }[dtype]; + if (!bytes) throw new CliError(`--state-dtype must be fp32|bf16|fp16|fp8, got "${dtype}"`, 2); + statePool.stateDtypeBytes = bytes; + } + + const memFraction = num(values, 'mem-fraction'); + if (memFraction !== undefined && !(memFraction > 0 && memFraction <= 1)) + throw new CliError(`--mem-fraction must be in (0, 1], got ${memFraction}`, 2); + + return { + prefillLen: num(values, 'prefill-len') ?? DEFAULT_SHARED.prefillLen, + genLen: num(values, 'gen-len') ?? DEFAULT_SHARED.genLen, + overlap: { memoryOverlap, commsOverlap, scheduler }, + costPerHour: num(values, 'cost-per-hour'), + memFraction, + statePool, + }; +} diff --git a/src/cli/commands/chips.ts b/src/cli/commands/chips.ts new file mode 100644 index 0000000..8f5de60 --- /dev/null +++ b/src/cli/commands/chips.ts @@ -0,0 +1,88 @@ +import { CHIPS, ChipSpec, peakFlops } from '../../core/hardware/chips'; +import { ModelSpec } from '../../core/model/models'; +import { weightBytesTotal } from '../../core/model/utils'; +import { fix, round, table } from '../format'; + +export interface ChipRow { + id: string; + name: string; + vendor: string; + hbmGb: number; + hbmTbps: number; + bf16Pf: number; + fp8Pf: number | null; + fp4Pf: number | null; + linkGbps: number; + domain: number; + // switched fabrics: nodes the scale-out tier reaches; ring fabrics: the + // slices sold + maxNodes: number | null; + slices: string[] | null; + costPerHour: number | null; + tdp: number | null; + // with --model: the least chips whose HBM holds the weights + minChipsForWeights: number | null; +} + +const pf = (chip: ChipSpec, d: 'bf16' | 'fp8' | 'fp4') => { + const v = peakFlops(chip, d) ?? (d === 'fp8' ? peakFlops(chip, 'mxfp8') : undefined); + return v === undefined ? null : round(v / 1e15, 2); +}; + +export function chipRows(model?: ModelSpec): ChipRow[] { + const w = model ? weightBytesTotal(model) : undefined; + return CHIPS.map((c) => ({ + id: c.id, + name: c.name, + vendor: c.vendor, + hbmGb: round(c.hbmCapacity / 1e9, 0), + hbmTbps: round(c.hbmBandwidth / 1e12, 2), + bf16Pf: pf(c, 'bf16')!, + fp8Pf: pf(c, 'fp8'), + fp4Pf: pf(c, 'fp4'), + linkGbps: round(c.interconnect.bandwidthPerChip / 1e9, 0), + domain: c.interconnect.domainSize, + maxNodes: c.interconnect.topologies ? null : (c.interconnect.scaleOut?.maxNodes ?? 1), + slices: c.interconnect.topologies ? c.interconnect.topologies.map((t) => t.name) : null, + costPerHour: c.costPerHour ?? null, + tdp: c.tdp ?? null, + minChipsForWeights: w === undefined ? null : Math.ceil(w / c.hbmCapacity), + })); +} + +export function renderChips(rows: ChipRow[]): string { + const withModel = rows.some((r) => r.minChipsForWeights !== null); + const head = [ + 'chip', + 'HBM GB', + 'TB/s', + 'bf16 PF', + 'fp8 PF', + 'fp4 PF', + 'link GB/s', + 'domain', + 'machines', + '$/hr', + 'W', + ...(withModel ? ['min chips'] : []), + ]; + const body = rows.map((r) => [ + r.id, + String(r.hbmGb), + fix(r.hbmTbps, 2), + fix(r.bf16Pf, 2), + fix(r.fp8Pf, 2), + fix(r.fp4Pf, 2), + String(r.linkGbps), + String(r.domain), + r.slices + ? `${r.slices.length} slices (${r.slices.slice(0, 3).join(',')}${r.slices.length > 3 ? ',...' : ''})` + : r.maxNodes === 1 + ? `${r.domain}x1` + : `${r.domain}x1 to ${r.domain}x${r.maxNodes}`, + fix(r.costPerHour, 2), + r.tdp === null ? '-' : String(r.tdp), + ...(withModel ? [r.minChipsForWeights === null ? '-' : String(r.minChipsForWeights)] : []), + ]); + return table(head, body, ['l', 'r', 'r', 'r', 'r', 'r', 'r', 'r', 'l', 'r', 'r', 'r']); +} diff --git a/src/cli/commands/explain.ts b/src/cli/commands/explain.ts new file mode 100644 index 0000000..7711338 --- /dev/null +++ b/src/cli/commands/explain.ts @@ -0,0 +1,193 @@ +import { + Candidate, + decodeContextParallels, + evaluateWithSearchOpts, + SearchPhase, +} from '../../core/engine/optimizer/search'; +import type { ResourceCostBackend } from '../../core/engine/sim/cost/types'; +import { ShardingRole } from '../../core/engine/sim/ir/sharding/roles'; +import { evaluateDecodeAtBatch } from '../../core/engine/sim/run/decode'; +import { Deployment, MoeDispatch, validateSizes } from '../../core/engine/surface/deploy'; +import { enumeratePlacements } from '../../core/engine/surface/placements'; +import { deployedAxes } from '../../core/hardware/topology'; +import { hasMoeLayers } from '../../core/model/utils'; +import { Context, evalOptions } from '../context'; +import { fix, int } from '../format'; +import { CliError } from '../resolve'; +import { candidateRow, EstimateRow } from '../row'; +import { SearchArgs, searchPhase } from './search'; + +export interface ExplainArgs extends SearchArgs { + // the roles given; the rest fill in from the machine (PP 1, DPA the + // remainder of the attention plane, EP the whole stage on MoE, ETP the + // remainder of the expert plane) + sizes: Partial>; + dcp?: number; + dispatch?: MoeDispatch; + // decode at exactly this batch instead of the policy's operating batch + batch?: number; +} + +export interface Explanation { + // the best placement/dispatch/DCP of the sizes, the way a search scores it + best: EstimateRow; + // every placement priced, best first + all: EstimateRow[]; + sizes: Record; + hints: string[]; +} + +export function fillSizes(ctx: Context, given: Partial>) { + const n = ctx.nChips; + const PP = given.PP ?? 1; + const stage = n / PP; + if (!Number.isInteger(stage)) throw new CliError(`PP=${PP} does not divide ${n} chips`, 2); + const TP = given.TP ?? (given.DPA ? stage / given.DPA : 1); + const DPA = given.DPA ?? stage / TP; + const moe = hasMoeLayers(ctx.model); + const EP = given.EP ?? (given.ETP ? stage / given.ETP : moe ? stage : 1); + const ETP = given.ETP ?? stage / EP; + const sizes = { PP, DPA, TP, EP, ETP }; + for (const [r, v] of Object.entries(sizes)) + if (!Number.isInteger(v) || v < 1) throw new CliError(`${r}=${v} is not a whole role size`, 2); + if (DPA * TP !== stage) + throw new CliError(`DPA=${DPA} x TP=${TP} is not the ${stage} chips of a stage`, 2); + if (EP * ETP !== stage) + throw new CliError(`EP=${EP} x ETP=${ETP} is not the ${stage} chips of a stage`, 2); + return sizes; +} + +export function explain(ctx: Context, a: ExplainArgs): Explanation { + const sizes = fillSizes(ctx, a.sizes); + const diags = validateSizes(ctx.model, { ...sizes, DCP: a.dcp }); + const errors = diags.filter((d) => d.severity === 'error'); + if (errors.length) throw new CliError(errors.map((d) => d.message).join('; '), 2); + + const axes = deployedAxes(ctx.chip.interconnect, ctx.machine); + const phase = searchPhase(ctx, a); + const opts = { ...evalOptions(ctx), phase }; + const dispatches: MoeDispatch[] = a.dispatch + ? [a.dispatch] + : hasMoeLayers(ctx.model) && sizes.EP > 1 + ? ['ring-of-experts', 'coalesced-a2a'] + : ['ring-of-experts']; + const dcps = a.dcp !== undefined ? [a.dcp] : decodeContextParallels(ctx.model, sizes.TP); + + const cands: Candidate[] = []; + for (const mesh of enumeratePlacements(axes, sizes)) + for (const moeDispatch of dispatches) + for (const decodeContextParallel of dcps) { + const deployment: Deployment = { chip: ctx.chip, mesh, moeDispatch, decodeContextParallel }; + const input = { model: ctx.model, deployment, workload: ctx.workload }; + const cand = + a.batch !== undefined && phase.kind === 'decode' + ? atBatch(input, a.batch, sizes.PP, opts) + : evaluateWithSearchOpts(input, opts); + if (cand) cands.push(cand); + } + if (!cands.length) + throw new CliError( + `no feasible placement of ${Object.entries(sizes) + .map(([r, v]) => `${r}=${v}`) + .join(' ')} on ${ctx.chip.id} ${ctx.machineName}` + + (a.slo ? ` at >= ${a.slo} tok/s/user` : '') + + (a.batch ? ` at batch ${a.batch}` : ''), + ); + cands.sort((x, y) => y.score - x.score); + const all = cands.map((c) => candidateRow(ctx, phase, c)); + return { best: all[0], all, sizes, hints: engineHints(sizes, all[0].dcp ?? 1) }; +} + +// a fixed decode batch, scored the way the search scores its operating one +function atBatch( + input: Parameters[0], + batch: number, + pp: number, + opts: Parameters>[1], +): Candidate | null { + const res = evaluateDecodeAtBatch(input, batch, pp, opts); + if (!res.ok) return null; + return { + deployment: input.deployment, + backend: opts.costBackend, + batch, + result: res, + score: res.tokPerSecPerChip, + }; +} + +// how the roles spell as engine flags (the mapping roles.ts documents) +export function engineHints(s: Record, dcp: number): string[] { + const world = s.TP * s.DPA; + const vllm = [`-tp ${s.TP}`, `-dp ${s.DPA}`]; + if (s.PP > 1) vllm.push(`-pp ${s.PP}`); + if (s.EP > 1) vllm.push('--enable-expert-parallel'); + const sglang = [`--tp ${world}`, `--dp ${s.DPA}`]; + if (s.DPA > 1) sglang.push('--enable-dp-attention'); + if (s.PP > 1) sglang.push(`--pp-size ${s.PP}`); + if (s.EP > 1) sglang.push(`--ep ${s.EP}`); + if (dcp > 1) sglang.push(`--decode-context-parallel-size ${dcp}`); + const out = [`vLLM: ${vllm.join(' ')}`, `SGLang: ${sglang.join(' ')}`]; + if (s.ETP > 1) out.push(`TRT-LLM: moe_tp ${s.ETP} (ETP has no vLLM/SGLang flag)`); + return out; +} + +export function renderExplanation(e: Explanation, phase: SearchPhase): string { + const r = e.best; + const lines: string[] = []; + const kv = (k: string, v: string) => lines.push(`${k.padEnd(22)} ${v}`); + kv('model', r.model); + kv('chip', `${r.chip} ${r.machine} (${r.nChips} chips)`); + kv('sharding', `${r.sharding}` + (r.dcp && r.dcp > 1 ? ` dcp=${r.dcp}` : '')); + kv('placement', `${r.placement}` + (r.moeDispatch ? `, ${r.moeDispatch}` : '')); + kv('phase', r.phase + (r.slo !== null ? ` at >= ${r.slo} tok/s/user` : '')); + if (r.weightGbPerChip !== null) { + kv('weights / chip', `${fix(r.weightGbPerChip, 1)} GB`); + kv( + 'kv / seq / chip', + `${fix(r.kvMbPerSeqPerChip, 1)} MB (paged ${fix(r.pagedKvMbPerSeqPerChip, 1)} MB + state ${fix(r.stateMbPerSeqPerChip, 1)} MB over ${r.stateSlotsPerSeq} slot${r.stateSlotsPerSeq === 1 ? '' : 's'})`, + ); + kv('max resident seqs', `${int(r.maxResidentSeqs)} on the machine`); + } + if (r.phase === 'decode') { + kv('batch', `${int(r.batch)} sequences`); + kv('step time', `${fix(r.stepTimeMs, 3)} ms`); + kv('tpot', `${fix(r.tpotMs, 3)} ms (${fix(r.tokPerSecPerUser, 1)} tok/s/user)`); + kv('tok/s/chip', `${int(r.tokPerSecPerChip)} (${int(r.tokPerSecMachine)} on the machine)`); + kv('mfu / mbu', `${fix((r.mfu ?? NaN) * 100, 1)}% / ${fix((r.mbu ?? NaN) * 100, 1)}%`); + } else { + kv('sequences', `${int(r.prefill?.batchSeqs)}`); + kv('pass latency', `${fix(r.prefill?.latencyMs, 3)} ms`); + kv('tok/s/chip', `${int(r.tokPerSecPerChip)} (${int(r.tokPerSecMachine)} on the machine)`); + } + if (r.busyMs) + kv( + 'bound by', + `${r.boundBy} (busy ms: compute ${fix(r.busyMs.compute, 3)}, memory ${fix(r.busyMs.memory, 3)}, comms ${fix(r.busyMs.comms, 3)})`, + ); + if (r.usdPerMtok !== null) kv('$/Mtok', `${fix(r.usdPerMtok, 4)} at $${r.costPerHour}/chip-hr`); + if (r.phase === 'decode' && r.prefill) + kv( + 'prefill (same config)', + `${int(r.prefill.tokPerSecPerChip)} tok/s/chip at ${r.prefill.batchSeqs} seqs, TTFT ${fix(r.prefill.ttftMs, 1)} ms`, + ); + kv( + 'scheduler', + r.scheduler + + (r.scheduler === 'naive' + ? ` (overlap ${r.overlap.memoryOverlap}/${r.overlap.commsOverlap})` + : ''), + ); + if (Object.keys(r.statePool).length) kv('state pool', JSON.stringify(r.statePool)); + for (const d of r.diagnostics) kv(d.severity === 'error' ? 'error' : d.severity, d.message); + lines.push('', ...e.hints); + if (e.all.length > 1) { + lines.push('', `${e.all.length} placements priced (${phase.kind}):`); + for (const x of e.all) + lines.push( + ` ${int(x.tokPerSecPerChip).padStart(6)} tok/s/chip ${x.placement}` + + (x.moeDispatch ? `, ${x.moeDispatch}` : ''), + ); + } + return lines.join('\n'); +} diff --git a/src/cli/commands/models.ts b/src/cli/commands/models.ts new file mode 100644 index 0000000..cb70954 --- /dev/null +++ b/src/cli/commands/models.ts @@ -0,0 +1,74 @@ +import { MODEL_PRESETS, ModelSpec } from '../../core/model/models'; +import { + activeParams, + layerCount, + minKvHeads, + moeExperts, + totalParams, + weightBytesTotal, +} from '../../core/model/utils'; +import { round, table } from '../format'; + +export interface ModelRow { + name: string; + paramsB: number; + activeB: number; + weightGb: number; + layers: number; + experts: number; + minKvHeads: number; + weights: string; + activations: string; + kv: string; +} + +const dtypes = (m: ModelSpec, which: 'weights' | 'activations') => { + const p = m.precision[which]; + const set = [...new Set([p.attention, p.denseMlp, p.routedExperts])]; + return set.join('/'); +}; + +export function modelRows(): ModelRow[] { + return MODEL_PRESETS.map((m) => ({ + name: m.name, + paramsB: round(totalParams(m) / 1e9, 1), + activeB: round(activeParams(m) / 1e9, 1), + weightGb: round(weightBytesTotal(m) / 1e9, 1), + layers: layerCount(m), + experts: moeExperts(m), + minKvHeads: minKvHeads(m), + weights: dtypes(m, 'weights'), + activations: dtypes(m, 'activations'), + kv: m.precision.kv, + })); +} + +export function renderModels(rows: ModelRow[]): string { + return table( + [ + 'model', + 'params B', + 'active B', + 'weights GB', + 'layers', + 'experts', + 'kv heads', + 'weights', + 'acts', + 'kv', + ], + rows.map((r) => [ + r.name, + String(r.paramsB), + String(r.activeB), + String(r.weightGb), + String(r.layers), + r.experts ? String(r.experts) : '-', + String(r.minKvHeads), + r.weights, + r.activations, + r.kv, + ]), + ['l', 'r', 'r', 'r', 'r', 'r', 'r', 'l', 'l', 'l'], + ); +} diff --git a/src/cli/commands/search.ts b/src/cli/commands/search.ts new file mode 100644 index 0000000..dd60d52 --- /dev/null +++ b/src/cli/commands/search.ts @@ -0,0 +1,146 @@ +import { + Candidate, + SearchPhase, + searchShardings, + searchTuples, +} from '../../core/engine/optimizer/search'; +import type { ResourceCostBackend } from '../../core/engine/sim/cost/types'; +import { Context, evalOptions } from '../context'; +import { fix, int, table } from '../format'; +import { CliError } from '../resolve'; +import { candidateRow, EstimateRow, shardingLabel } from '../row'; + +export interface SearchArgs { + phase: 'decode' | 'prefill'; + // decode: minimum tok/s/user; the batch backs off to meet it + slo?: number; + batching?: 'max' | 'b1'; + // prefill: a single-sequence TTFT pass instead of a full-machine batch + ttft?: boolean; + // prefill: sequences per throughput pass (default one per chip) + prefillSeqs?: number; + perDollar?: boolean; +} + +export function searchPhase(ctx: Context, a: SearchArgs): SearchPhase { + if (a.phase === 'decode') + return { kind: 'decode', policy: { sloTokPerSecPerUser: a.slo, batching: a.batching } }; + return a.ttft + ? { kind: 'prefill', seqs: 1, mode: 'ttft' } + : { kind: 'prefill', seqs: a.prefillSeqs ?? ctx.nChips, mode: 'throughput' }; +} + +export interface SearchStep { + done: number; + total: number; + candidate?: Candidate; + // is this tuple's candidate the best so far + best: boolean; + elapsedMs: number; +} + +export function searchSize(ctx: Context): number { + return searchTuples(ctx.model, ctx.nChips).length; +} + +// Every feasible tuple's best candidate, best first. onStep sees each +// tuple the moment it prices. +export function rankedCandidates( + ctx: Context, + a: SearchArgs, + onStep?: (s: SearchStep) => void, +): Candidate[] { + if (a.perDollar && ctx.chip.costPerHour === undefined) + throw new CliError(`${ctx.chip.id} has no price; pass --cost-per-hour to rank per dollar`, 2); + const phase = searchPhase(ctx, a); + const t0 = performance.now(); + const all: Candidate[] = []; + let best = -Infinity; + for (const step of searchShardings(ctx.model, ctx.chip, ctx.machine, ctx.workload, { + ...evalOptions(ctx), + phase, + rank: a.perDollar ? 'perDollar' : 'perChip', + })) { + const c = step.candidate; + const isBest = !!c && c.score > best; + if (c) { + all.push(c); + if (isBest) best = c.score; + } + onStep?.({ ...step, best: isBest, elapsedMs: performance.now() - t0 }); + } + return all.sort((x, y) => y.score - x.score); +} + +export function rankedRows( + ctx: Context, + a: SearchArgs, + top: number, + onStep?: (s: SearchStep) => void, +): EstimateRow[] { + const phase = searchPhase(ctx, a); + const cands = rankedCandidates(ctx, a, onStep); + return (top > 0 ? cands.slice(0, top) : cands).map((c) => candidateRow(ctx, phase, c)); +} + +// one streamed line per priced tuple, the scripts/search.ts format +export function stepLine(ctx: Context, s: SearchStep): string { + const c = s.candidate; + const head = + `[+${s.elapsedMs.toFixed(0).padStart(6)}ms] ` + + `${String(s.done).padStart(3)}/${s.total} ${s.best ? '*' : ' '} `; + if (!c) return head + 'infeasible'; + const label = shardingLabel(ctx.model, c.deployment.mesh); + const dcp = c.deployment.decodeContextParallel ?? 1; + return ( + head + + `${Math.round(c.result.tokPerSecPerChip).toString().padStart(6)} tok/s/chip ` + + `${label.padEnd(30)} ` + + (c.batch !== undefined ? `batch=${c.batch} ` : '') + + (dcp > 1 ? `dcp=${dcp} ` : '') + + (c.deployment.mesh.roles.EP.length ? `${c.deployment.moeDispatch} ` : '') + ); +} + +export function renderRows(rows: EstimateRow[], feasibleTotal: number): string { + if (!rows.length) return 'no feasible configuration'; + const decode = rows[0].phase === 'decode'; + const head = decode + ? ['#', 'tok/s/chip', 'tok/s/user', 'batch', 'residents', '$/Mtok', 'bound', 'sharding'] + : ['#', 'tok/s/chip', 'latency ms', 'seqs', '$/Mtok', 'bound', 'sharding']; + const body = rows.map((r, i) => { + const sharding = + `${r.sharding}` + + (r.dcp && r.dcp > 1 ? ` dcp=${r.dcp}` : '') + + (r.moeDispatch ? ` ${r.moeDispatch}` : ''); + return decode + ? [ + String(i + 1), + int(r.tokPerSecPerChip), + fix(r.tokPerSecPerUser, 1), + int(r.batch), + int(r.maxResidentSeqs), + fix(r.usdPerMtok, 3), + r.boundBy ?? '-', + sharding, + ] + : [ + String(i + 1), + int(r.tokPerSecPerChip), + fix(r.prefill?.latencyMs, 2), + int(r.prefill?.batchSeqs), + fix(r.usdPerMtok, 3), + r.boundBy ?? '-', + sharding, + ]; + }); + const align = decode + ? (['r', 'r', 'r', 'r', 'r', 'r', 'l', 'l'] as const) + : (['r', 'r', 'r', 'r', 'r', 'l', 'l'] as const); + const r0 = rows[0]; + const title = + `${r0.model} on ${r0.chip} ${r0.machine} (${r0.nChips} chips), ${r0.phase}` + + (r0.slo !== null ? ` at >= ${r0.slo} tok/s/user` : '') + + `: top ${rows.length} of ${feasibleTotal} feasible`; + return `${title}\n${table(head, body, [...align])}`; +} diff --git a/src/cli/commands/sweep.ts b/src/cli/commands/sweep.ts new file mode 100644 index 0000000..b0bc138 --- /dev/null +++ b/src/cli/commands/sweep.ts @@ -0,0 +1,182 @@ +import { Worker } from 'node:worker_threads'; +import { Context, makeContext, SharedOptions } from '../context'; +import { fix, int, table } from '../format'; +import { EstimateRow, infeasibleRow } from '../row'; +import { rankedRows, SearchArgs, searchPhase } from './search'; + +// one search: a model on a chip's machine, one phase, one SLO +export interface SweepCell { + model: string; + chip: string; + machine: string; + phase: 'decode' | 'prefill'; + slo: number | null; +} + +export interface SweepSpec { + model: string; + chips: string[]; + machines: string[]; + phases: ('decode' | 'prefill')[]; + // null = no SLO, batch to capacity + slos: (number | null)[]; +} + +export function sweepCells(s: SweepSpec): SweepCell[] { + const cells: SweepCell[] = []; + for (const chip of s.chips) + for (const machine of s.machines) + for (const phase of s.phases) + for (const slo of phase === 'decode' ? s.slos : [null]) + cells.push({ model: s.model, chip, machine, phase, slo }); + return cells; +} + +export function cellName(c: SweepCell): string { + return `${c.chip}-${c.machine}-${c.phase}-${c.slo ?? 'none'}`; +} + +// One cell's best configuration, or an infeasible row carrying the reason. +// Never throws: a sweep reports every cell. +export function runCell(cell: SweepCell, shared: SharedOptions): EstimateRow { + const a: SearchArgs = { phase: cell.phase, slo: cell.slo ?? undefined }; + let ctx: Context; + try { + ctx = makeContext(cell.model, cell.chip, cell.machine, shared); + } catch (err) { + return { + ...infeasibleRowFor(cell, shared), + error: err instanceof Error ? err.message : String(err), + }; + } + try { + const [row] = rankedRows(ctx, a, 1); + return row ?? infeasibleRow(ctx, searchPhase(ctx, a)); + } catch (err) { + return infeasibleRow( + ctx, + searchPhase(ctx, a), + err instanceof Error ? err.message : String(err), + ); + } +} + +// a row for a cell whose chip or machine did not even resolve +function infeasibleRowFor(cell: SweepCell, shared: SharedOptions): EstimateRow { + const ctx = { + model: { name: cell.model }, + chip: { id: cell.chip }, + machineName: cell.machine, + nChips: 0, + workload: { prefillLen: shared.prefillLen, generateLen: shared.genLen }, + opts: shared, + } as unknown as Context; + return infeasibleRow( + ctx, + cell.phase === 'decode' + ? { kind: 'decode', policy: { sloTokPerSecPerUser: cell.slo ?? undefined } } + : { kind: 'prefill', seqs: 0, mode: 'throughput' }, + ); +} + +export interface WorkerJob { + kind: 'sweep'; + cells: SweepCell[]; + shared: SharedOptions; +} + +// Run the cells, one worker thread per chip (a chip's warm reshard-plan +// caches stay with the thread that will be asked about it again), at most +// `jobs` at a time. jobs <= 1 or no worker script runs everything inline. +export async function runSweep( + cells: SweepCell[], + shared: SharedOptions, + jobs: number, + onRow: (row: EstimateRow) => void, + workerScript?: URL, +): Promise { + const rows: EstimateRow[] = []; + const emit = (row: EstimateRow) => { + rows.push(row); + onRow(row); + }; + if (jobs <= 1 || !workerScript) { + for (const cell of cells) emit(runCell(cell, shared)); + return rows; + } + + const byChip = new Map(); + for (const c of cells) byChip.set(c.chip, [...(byChip.get(c.chip) ?? []), c]); + const queue = [...byChip.values()]; + const runOne = (group: SweepCell[]) => + new Promise((resolve, reject) => { + const w = new Worker(workerScript, { + workerData: { kind: 'sweep', cells: group, shared } satisfies WorkerJob, + }); + w.on('message', (row: EstimateRow) => emit(row)); + w.on('error', reject); + w.on('exit', (code) => + code === 0 ? resolve() : reject(new Error(`sweep worker exited with ${code}`)), + ); + }); + const lanes = Array.from({ length: Math.min(jobs, queue.length) }, async () => { + for (let g = queue.shift(); g; g = queue.shift()) await runOne(g); + }); + await Promise.all(lanes); + return rows; +} + +// The grouped tables of a sweep: decode at each SLO, then prefill, best +// first within each, then the cells that had no feasible configuration. +export function renderReport(rows: EstimateRow[]): string { + const ok = rows.filter((r) => r.feasible); + const bad = rows.filter((r) => !r.feasible); + const out: string[] = []; + const section = (title: string, sub: EstimateRow[]) => { + if (!sub.length) return; + sub.sort((a, b) => (b.tokPerSecPerChip ?? 0) - (a.tokPerSecPerChip ?? 0)); + out.push( + `=== ${title} ===`, + table( + ['chip', 'machine', 'n', 'tok/s/chip', 'tok/s/user', 'batch', '$/Mtok', 'sharding'], + sub.map((r) => [ + r.chip, + r.machine, + String(r.nChips), + int(r.tokPerSecPerChip), + fix(r.tokPerSecPerUser, 1), + int(r.batch ?? r.prefill?.batchSeqs), + fix(r.usdPerMtok, 3), + `${r.sharding}` + + (r.dcp && r.dcp > 1 ? ` dcp=${r.dcp}` : '') + + (r.moeDispatch ? ` ${r.moeDispatch}` : ''), + ]), + ['l', 'l', 'r', 'r', 'r', 'r', 'r', 'l'], + ), + '', + ); + }; + section( + 'DECODE max throughput (no SLO)', + ok.filter((r) => r.phase === 'decode' && r.slo === null), + ); + const slos = [...new Set(ok.filter((r) => r.slo !== null).map((r) => r.slo!))].sort( + (a, b) => a - b, + ); + for (const slo of slos) + section( + `DECODE at >= ${slo} tok/s/user`, + ok.filter((r) => r.phase === 'decode' && r.slo === slo), + ); + section( + 'PREFILL throughput', + ok.filter((r) => r.phase === 'prefill'), + ); + out.push(`${rows.length} cells, ${ok.length} feasible, ${bad.length} infeasible`); + for (const r of bad) + out.push( + ` infeasible: ${r.chip} ${r.machine} ${r.phase} slo=${r.slo ?? 'none'}` + + (r.error ? `: ${r.error}` : ''), + ); + return out.join('\n'); +} diff --git a/src/cli/context.ts b/src/cli/context.ts new file mode 100644 index 0000000..2fbc05b --- /dev/null +++ b/src/cli/context.ts @@ -0,0 +1,72 @@ +import { makeCostBackend, OverlapOptions } from '../core/engine/sim/cost/select'; +import type { ResourceCostBackend } from '../core/engine/sim/cost/types'; +import type { EvalOptions, StatePoolOptions } from '../core/engine/surface/api'; +import { ChipSpec } from '../core/hardware/chips'; +import { Machine, machineName, machineSize, parseMachine } from '../core/hardware/machines'; +import { ModelSpec } from '../core/model/models'; +import { resolveChip, resolveModel } from './resolve'; + +// Every knob the subcommands share, as plain data (it crosses into sweep +// worker threads). +export interface SharedOptions { + prefillLen: number; + genLen: number; + overlap: OverlapOptions; + // chip overrides: rental price and the fraction of HBM the stack lets the + // cache have (SGLang's mem_fraction_static), applied as a smaller chip + costPerHour?: number; + memFraction?: number; + statePool: Partial; +} + +export const DEFAULT_SHARED: SharedOptions = { + prefillLen: 4096, + genLen: 1024, + overlap: { memoryOverlap: 0.9, commsOverlap: 0.65, scheduler: 'naive' }, + statePool: {}, +}; + +// One (model, chip, machine) the commands evaluate on, resolved. +export interface Context { + model: ModelSpec; + chip: ChipSpec; + machine: Machine; + machineName: string; + nChips: number; + workload: { prefillLen: number; generateLen: number }; + opts: SharedOptions; + backend: ResourceCostBackend; +} + +export function applyChipOverrides(chip: ChipSpec, opts: SharedOptions): ChipSpec { + return { + ...chip, + ...(opts.memFraction !== undefined ? { hbmCapacity: opts.memFraction * chip.hbmCapacity } : {}), + ...(opts.costPerHour !== undefined ? { costPerHour: opts.costPerHour } : {}), + }; +} + +export function makeContext( + modelArg: string, + chipArg: string, + machineArg: string, + opts: SharedOptions = DEFAULT_SHARED, +): Context { + const model = resolveModel(modelArg); + const chip = applyChipOverrides(resolveChip(chipArg), opts); + const machine = parseMachine(chip, machineArg); + return { + model, + chip, + machine, + machineName: machineName(machine), + nChips: machineSize(machine), + workload: { prefillLen: opts.prefillLen, generateLen: opts.genLen }, + opts, + backend: makeCostBackend(opts.overlap), + }; +} + +export function evalOptions(ctx: Context): EvalOptions { + return { costBackend: ctx.backend, statePool: ctx.opts.statePool }; +} diff --git a/src/cli/format.ts b/src/cli/format.ts new file mode 100644 index 0000000..d2cd110 --- /dev/null +++ b/src/cli/format.ts @@ -0,0 +1,31 @@ +// Fixed-width text tables and number cells for the terminal. + +export type Align = 'l' | 'r'; + +export function table(head: string[], rows: string[][], align: Align[] = []): string { + const widths = head.map((h, i) => Math.max(h.length, ...rows.map((r) => (r[i] ?? '').length))); + const line = (cells: string[]) => + cells + .map((c, i) => (align[i] === 'r' ? c.padStart(widths[i]) : c.padEnd(widths[i]))) + .join(' ') + .trimEnd(); + return [line(head), ...rows.map(line)].join('\n'); +} + +// a number cell, '-' for what a row does not have +export function fix(v: number | null | undefined, digits = 1): string { + return v === null || v === undefined || !Number.isFinite(v) ? '-' : v.toFixed(digits); +} + +export function int(v: number | null | undefined): string { + return v === null || v === undefined || !Number.isFinite(v) ? '-' : String(Math.round(v)); +} + +// JSON rows carry rounded numbers so a diff reads +export function round(v: number, digits: number): number { + return Number(v.toFixed(digits)); +} + +export function nullable(v: number | undefined, digits: number): number | null { + return v === undefined || !Number.isFinite(v) ? null : round(v, digits); +} diff --git a/src/cli/main.ts b/src/cli/main.ts new file mode 100644 index 0000000..b6b38cf --- /dev/null +++ b/src/cli/main.ts @@ -0,0 +1,296 @@ +import { isMainThread, parentPort, workerData } from 'node:worker_threads'; +import { readdirSync, readFileSync, statSync, writeFileSync, mkdirSync } from 'node:fs'; +import { join } from 'node:path'; +import { tmpdir } from 'node:os'; +import type { ShardingRole } from '../core/engine/sim/ir/sharding/roles'; +import type { MoeDispatch } from '../core/engine/surface/deploy'; +import { + flag, + list, + num, + oneOf, + parse, + SHARED_HELP, + SHARED_OPTIONS, + sharedOptions, + str, + Values, +} from './args'; +import { makeContext } from './context'; +import { CliError, resolveModel } from './resolve'; +import { VERSION } from './version'; +import { chipRows, renderChips } from './commands/chips'; +import { explain, renderExplanation } from './commands/explain'; +import { modelRows, renderModels } from './commands/models'; +import { rankedRows, renderRows, searchPhase, searchSize, stepLine } from './commands/search'; +import { cellName, renderReport, runCell, runSweep, sweepCells, WorkerJob } from './commands/sweep'; +import type { EstimateRow } from './row'; + +const USAGE = `estimate ${VERSION}: static roofline estimates of LLM serving (htdym) + +usage + estimate models [--json] + estimate chips [--model M] [--json] + estimate search M CHIP MACHINE [--phase decode|prefill] [--slo N] [--batching max|b1] + [--ttft] [--prefill-seqs N] [--top N] [--per-dollar] [--quiet] [--json] + estimate top M CHIP MACHINE [same as search; silent, --top 10] + estimate explain M CHIP MACHINE --sizes DPA=16,TP=2,EP=32 [--dcp N] [--dispatch D] + [--batch N | --slo N | --batching max|b1] [--phase decode|prefill] [--json] + estimate sweep --model M --chips a,b --machines 32x1,8x2 [--phases decode,prefill] + [--slos none,20,50] [--out DIR] [--jobs N] [--table] + estimate report + +M is a preset name, a unique prefix, or an alias (k3, dsv4flash, gptoss120b, ...). +CHIP is a chip id (estimate chips). MACHINE is x on a +switched fabric (32x1, 8x2) or a slice name on a ring fabric (4x4x4). + +Every number is an estimate from a static roofline model, not a measurement. +${SHARED_HELP} +`; + +const out = (s: string) => process.stdout.write(s + '\n'); +const err = (s: string) => process.stderr.write(s + '\n'); +const jsonl = (rows: EstimateRow[]) => rows.forEach((r) => out(JSON.stringify(r))); + +const SEARCH_OPTIONS = { + ...SHARED_OPTIONS, + phase: { type: 'string' }, + slo: { type: 'string' }, + batching: { type: 'string' }, + ttft: { type: 'boolean' }, + 'prefill-seqs': { type: 'string' }, + top: { type: 'string' }, + 'per-dollar': { type: 'boolean' }, + quiet: { type: 'boolean', short: 'q' }, +} as const; + +function searchArgs(values: Values) { + return { + phase: oneOf(values, 'phase', ['decode', 'prefill'] as const, 'decode'), + slo: num(values, 'slo'), + batching: str(values, 'batching') + ? oneOf(values, 'batching', ['max', 'b1'] as const, 'max') + : undefined, + ttft: flag(values, 'ttft'), + prefillSeqs: num(values, 'prefill-seqs'), + perDollar: flag(values, 'per-dollar'), + }; +} + +function needPositionals(p: string[], n: number, what: string) { + if (p.length < n) throw new CliError(`expected ${what}`, 2); +} + +function cmdModels(argv: string[]) { + const { values } = parse(argv, SHARED_OPTIONS, false); + const rows = modelRows(); + if (flag(values, 'json')) rows.forEach((r) => out(JSON.stringify(r))); + else out(renderModels(rows)); +} + +function cmdChips(argv: string[]) { + const { values } = parse(argv, { ...SHARED_OPTIONS, model: { type: 'string' } }, false); + const model = str(values, 'model'); + const rows = chipRows(model ? resolveModel(model) : undefined); + if (flag(values, 'json')) rows.forEach((r) => out(JSON.stringify(r))); + else out(renderChips(rows)); +} + +function cmdSearch(argv: string[], silent: boolean) { + const { values, positionals } = parse(argv, SEARCH_OPTIONS); + needPositionals(positionals, 3, 'MODEL CHIP MACHINE'); + const ctx = makeContext(positionals[0], positionals[1], positionals[2], sharedOptions(values)); + const a = searchArgs(values); + const top = num(values, 'top') ?? (silent ? 10 : 5); + const quiet = silent || flag(values, 'quiet'); + if (!quiet) + err( + `${ctx.model.name} on ${ctx.chip.id} ${ctx.machineName} (${ctx.nChips} chips): ` + + `${searchSize(ctx)} role-size tuples, ${a.phase}` + + (a.slo ? ` at >= ${a.slo} tok/s/user` : ''), + ); + let feasible = 0; + const rows = rankedRows(ctx, a, top, (s) => { + if (s.candidate) feasible++; + if (!quiet) err(stepLine(ctx, s)); + }); + if (flag(values, 'json')) jsonl(rows); + else out(renderRows(rows, feasible)); + if (!rows.length) throw new CliError('no feasible configuration', 1); +} + +function cmdExplain(argv: string[]) { + const { values, positionals } = parse(argv, { + ...SEARCH_OPTIONS, + sizes: { type: 'string' }, + dcp: { type: 'string' }, + dispatch: { type: 'string' }, + batch: { type: 'string' }, + }); + needPositionals(positionals, 3, 'MODEL CHIP MACHINE'); + const ctx = makeContext(positionals[0], positionals[1], positionals[2], sharedOptions(values)); + const sizes: Partial> = {}; + for (const part of list(values, 'sizes')) { + const m = /^(PP|DPA|TP|EP|ETP)=(\d+)$/i.exec(part); + if (!m) throw new CliError(`--sizes takes ROLE=N pairs (DPA=16,TP=2,EP=32), got "${part}"`, 2); + sizes[m[1].toUpperCase() as ShardingRole] = Number(m[2]); + } + const dispatch = str(values, 'dispatch'); + const dispatches = ['ring-of-experts', 'coalesced-a2a', 'expanded-a2a'] as const; + if (dispatch && !(dispatches as readonly string[]).includes(dispatch)) + throw new CliError(`--dispatch must be one of ${dispatches.join('|')}`, 2); + const a = { + ...searchArgs(values), + sizes, + dcp: num(values, 'dcp'), + dispatch: dispatch as MoeDispatch | undefined, + batch: num(values, 'batch'), + }; + const e = explain(ctx, a); + if (flag(values, 'json')) + out(JSON.stringify({ ...e.best, hints: e.hints, placements: e.all.length })); + else out(renderExplanation(e, searchPhase(ctx, a))); +} + +async function cmdSweep(argv: string[]) { + const { values } = parse( + argv, + { + ...SHARED_OPTIONS, + model: { type: 'string' }, + chips: { type: 'string' }, + machines: { type: 'string' }, + phases: { type: 'string' }, + slos: { type: 'string' }, + out: { type: 'string' }, + jobs: { type: 'string' }, + table: { type: 'boolean' }, + }, + false, + ); + const model = str(values, 'model'); + if (!model) throw new CliError('--model is required', 2); + const chips = list(values, 'chips'); + const machines = list(values, 'machines'); + if (!chips.length || !machines.length) + throw new CliError('--chips and --machines are required', 2); + const phases = list(values, 'phases', ['decode']).map((p) => { + if (p !== 'decode' && p !== 'prefill') throw new CliError(`--phases takes decode,prefill`, 2); + return p; + }); + const slos = list(values, 'slos', ['none']).map((s) => { + if (s === 'none') return null; + const n = Number(s); + if (!Number.isFinite(n) || n <= 0) + throw new CliError(`--slos takes none or tok/s/user numbers`, 2); + return n; + }); + const shared = sharedOptions(values); + const cells = sweepCells({ model, chips, machines, phases, slos }); + const outDir = str(values, 'out') ?? join(tmpdir(), `estimate-sweep-${Date.now()}`); + mkdirSync(outDir, { recursive: true }); + const jobs = num(values, 'jobs') ?? Math.min(chips.length, 4); + err(`${cells.length} cells over ${chips.length} chips, ${jobs} at a time, rows under ${outDir}`); + const t0 = performance.now(); + const rows = await runSweep( + cells, + shared, + jobs, + (row) => { + writeFileSync(join(outDir, `${cellName(rowCell(row))}.json`), JSON.stringify(row) + '\n'); + if (!flag(values, 'table')) out(JSON.stringify(row)); + err( + `[+${((performance.now() - t0) / 1000).toFixed(1).padStart(6)}s] ${cellName(rowCell(row))}: ` + + (row.feasible + ? `${Math.round(row.tokPerSecPerChip!)} tok/s/chip ${row.sharding}` + : `infeasible${row.error ? ` (${row.error})` : ''}`), + ); + }, + new URL(import.meta.url), + ); + if (flag(values, 'table')) out(renderReport(rows)); +} + +const rowCell = (r: EstimateRow) => ({ + model: r.model, + chip: r.chip, + machine: r.machine, + phase: r.phase, + slo: r.slo, +}); + +function cmdReport(argv: string[]) { + const { positionals } = parse(argv, SHARED_OPTIONS); + needPositionals(positionals, 1, 'a rows.jsonl file or a directory of row files'); + const rows: EstimateRow[] = []; + for (const p of positionals) { + const files = statSync(p).isDirectory() + ? readdirSync(p) + .filter((f) => f.endsWith('.json') || f.endsWith('.jsonl')) + .map((f) => join(p, f)) + : [p]; + for (const f of files) + for (const line of readFileSync(f, 'utf8').split('\n')) + if (line.trim().startsWith('{')) rows.push(JSON.parse(line) as EstimateRow); + } + out(renderReport(rows)); +} + +export async function main(argv: string[]): Promise { + const [cmd, ...rest] = argv; + try { + if (!cmd || cmd === '--help' || cmd === '-h' || cmd === 'help') { + out(USAGE); + return cmd ? 0 : 2; + } + if (cmd === '--version') { + out(VERSION); + return 0; + } + if (rest.includes('--help') || rest.includes('-h')) { + out(USAGE); + return 0; + } + switch (cmd) { + case 'models': + cmdModels(rest); + break; + case 'chips': + cmdChips(rest); + break; + case 'search': + cmdSearch(rest, false); + break; + case 'top': + cmdSearch(rest, true); + break; + case 'explain': + cmdExplain(rest); + break; + case 'sweep': + await cmdSweep(rest); + break; + case 'report': + cmdReport(rest); + break; + default: + throw new CliError(`unknown command "${cmd}"\n\n${USAGE}`, 2); + } + return 0; + } catch (e) { + if (e instanceof CliError) { + err(`estimate: ${e.message}`); + return e.exitCode; + } + throw e; + } +} + +if (!isMainThread && (workerData as WorkerJob | undefined)?.kind === 'sweep') { + const job = workerData as WorkerJob; + for (const cell of job.cells) parentPort!.postMessage(runCell(cell, job.shared)); +} else if (isMainThread) { + void main(process.argv.slice(2)).then((code) => { + process.exitCode = code; + }); +} diff --git a/src/cli/resolve.ts b/src/cli/resolve.ts new file mode 100644 index 0000000..faf09bf --- /dev/null +++ b/src/cli/resolve.ts @@ -0,0 +1,63 @@ +import { CHIPS, CHIPS_BY_ID, ChipSpec } from '../core/hardware/chips'; +import { MODEL_PRESETS, ModelSpec } from '../core/model/models'; + +// A failure the user can act on: printed as one line, no stack. +export class CliError extends Error { + constructor( + message: string, + public readonly exitCode: number = 1, + ) { + super(message); + } +} + +// short names people type, each to a unique preset-name prefix +const MODEL_ALIASES: Record = { + k3: 'Kimi K3', + kimik3: 'Kimi K3', + k2: 'Kimi K2.6', + 'k2.6': 'Kimi K2.6', + kimik2: 'Kimi K2.6', + dsv4flash: 'DeepSeek V4 Flash', + dsv4pro: 'DeepSeek V4 Pro', + glm53: 'GLM 5.3', + gptoss120b: 'gpt-oss-120b', + gptoss20b: 'gpt-oss-20b', + gemma31b: 'Gemma 4 31B', + gemma12b: 'Gemma 4 12B', + gemma26b: 'Gemma 4 26B', + llama8b: 'LLaMA 3 8B', + llama70b: 'LLaMA 3 70B', + llama405b: 'LLaMA 3.1 405B', + qwen35b: 'Qwen3.6 35B', + qwen27b: 'Qwen3.8 27B', + qwen4b: 'Qwen3 4B BF16', +}; + +const norm = (s: string) => s.toLowerCase().replace(/[\s_-]/g, ''); + +// exact preset name, an alias, or a unique case-insensitive prefix +export function resolveModel(arg: string): ModelSpec { + const exact = MODEL_PRESETS.find((m) => m.name === arg); + if (exact) return exact; + const target = MODEL_ALIASES[norm(arg)] ?? arg; + const hits = MODEL_PRESETS.filter((m) => norm(m.name).startsWith(norm(target))); + if (hits.length === 1) return hits[0]; + if (hits.length > 1) + throw new CliError(`model "${arg}" matches ${hits.map((m) => `"${m.name}"`).join(', ')}`, 2); + throw new CliError( + `unknown model "${arg}"; models: ${MODEL_PRESETS.map((m) => `"${m.name}"`).join(', ')}`, + 2, + ); +} + +// exact chip id or a unique prefix of one +export function resolveChip(arg: string): ChipSpec { + const exact = CHIPS_BY_ID[arg]; + if (exact) return exact; + const hits = CHIPS.filter((c) => c.id.startsWith(arg.toLowerCase())); + if (hits.length === 1) return hits[0]; + if (hits.length > 1) + throw new CliError(`chip "${arg}" matches ${hits.map((c) => c.id).join(', ')}`, 2); + throw new CliError(`unknown chip "${arg}"; chips: ${CHIPS.map((c) => c.id).join(', ')}`, 2); +} diff --git a/src/cli/row.ts b/src/cli/row.ts new file mode 100644 index 0000000..3a4a31c --- /dev/null +++ b/src/cli/row.ts @@ -0,0 +1,232 @@ +import { + enrichCandidate, + Boundedness, + boundBy, + ComponentTimes, +} from '../core/engine/optimizer/enrich'; +import type { Candidate, SearchPhase } from '../core/engine/optimizer/search'; +import type { Scheduler, OverlapOptions } from '../core/engine/sim/cost/select'; +import type { ResourceCostBackend } from '../core/engine/sim/cost/types'; +import { ROLES, ShardingRole } from '../core/engine/sim/ir/sharding/roles'; +import type { StatePoolOptions } from '../core/engine/surface/api'; +import { Diagnostic, Mesh, MoeDispatch, roleSize } from '../core/engine/surface/deploy'; +import { hasMoeLayers } from '../core/model/utils'; +import { ModelSpec } from '../core/model/models'; +import type { Context } from './context'; +import { nullable, round } from './format'; +import { VERSION } from './version'; + +// One estimated configuration as the CLI emits it (a JSON line, a table +// row). Field names follow the sweep output the tab fleet tooling reads. +export interface EstimateRow { + model: string; + chip: string; + machine: string; + nChips: number; + phase: 'decode' | 'prefill'; + slo: number | null; + feasible: boolean; + // "DPA=16 TP=2 EP=32": roles above 1, expert roles hidden on dense models + sharding: string | null; + sizes: Record | null; + dcp: number | null; + moeDispatch: MoeDispatch | null; + placement: string | null; + batch: number | null; + tokPerSecPerChip: number | null; + tokPerSecMachine: number | null; + tpotMs: number | null; + tokPerSecPerUser: number | null; + stepTimeMs: number | null; + boundBy: Boundedness | null; + busyMs: ComponentTimes | null; + visibleMs: ComponentTimes | null; + mfu: number | null; + mbu: number | null; + weightGbPerChip: number | null; + kvMbPerSeqPerChip: number | null; + pagedKvMbPerSeqPerChip: number | null; + stateMbPerSeqPerChip: number | null; + stateSlotsPerSeq: number | null; + maxResidentSeqs: number | null; + // the same deployment's prefill (for a decode row: a full-machine batch + // plus a single-sequence TTFT pass; for a prefill row: the searched pass) + prefill: { + tokPerSecPerChip: number; + ttftMs: number | null; + latencyMs: number | null; + batchSeqs: number; + boundBy: Boundedness; + } | null; + costPerHour: number | null; + usdPerMtok: number | null; + workload: { prefillLen: number; generateLen: number }; + overlap: OverlapOptions; + statePool: Partial; + scheduler: Scheduler; + diagnostics: Diagnostic[]; + error: string | null; + id: string | null; + version: string; +} + +// $/Mtok at one chip's rental rate and rate, the fleet tooling's own +// expression (cost_per_million_tokens in tab/clis/fleet/batch_floor.py) +export function usdPerMtok(costPerHour: number, tokPerSecPerChip: number): number { + return (costPerHour / (tokPerSecPerChip * 3600)) * 1e6; +} + +export function sizesOf(mesh: Mesh): Record { + return Object.fromEntries(ROLES.map((r) => [r, roleSize(mesh, r)])) as Record< + ShardingRole, + number + >; +} + +// the expert plane is structural noise on dense models, so it is hidden +export function shardingLabel(model: ModelSpec, mesh: Mesh): string { + const shown = ROLES.filter( + (r) => roleSize(mesh, r) > 1 && (hasMoeLayers(model) || (r !== 'EP' && r !== 'ETP')), + ); + return shown.map((r) => `${r}=${roleSize(mesh, r)}`).join(' ') || 'single-chip'; +} + +export function placementLabel(model: ModelSpec, mesh: Mesh, dcp: number): string { + const shown = ROLES.filter( + (r) => roleSize(mesh, r) > 1 && (hasMoeLayers(model) || (r !== 'EP' && r !== 'ETP')), + ); + return ( + shown.map((r) => `${r}[${mesh.roles[r].join('')}]`).join(' ') + + (dcp > 1 ? ` DCP=${dcp} of TP` : '') + ); +} + +export const sloOf = (phase: SearchPhase): number | null => + phase.kind === 'decode' ? (phase.policy?.sloTokPerSecPerUser ?? null) : null; + +function base(ctx: Context, phase: SearchPhase): EstimateRow { + return { + model: ctx.model.name, + chip: ctx.chip.id, + machine: ctx.machineName, + nChips: ctx.nChips, + phase: phase.kind, + slo: sloOf(phase), + feasible: false, + sharding: null, + sizes: null, + dcp: null, + moeDispatch: null, + placement: null, + batch: null, + tokPerSecPerChip: null, + tokPerSecMachine: null, + tpotMs: null, + tokPerSecPerUser: null, + stepTimeMs: null, + boundBy: null, + busyMs: null, + visibleMs: null, + mfu: null, + mbu: null, + weightGbPerChip: null, + kvMbPerSeqPerChip: null, + pagedKvMbPerSeqPerChip: null, + stateMbPerSeqPerChip: null, + stateSlotsPerSeq: null, + maxResidentSeqs: null, + prefill: null, + costPerHour: ctx.chip.costPerHour ?? null, + usdPerMtok: null, + workload: ctx.workload, + overlap: ctx.opts.overlap, + statePool: ctx.opts.statePool, + scheduler: ctx.opts.overlap.scheduler ?? 'naive', + diagnostics: [], + error: null, + id: null, + version: VERSION, + }; +} + +export function infeasibleRow(ctx: Context, phase: SearchPhase, error?: string): EstimateRow { + return { ...base(ctx, phase), error: error ?? null }; +} + +const ms = (c: ComponentTimes): ComponentTimes => ({ + compute: round(c.compute * 1e3, 4), + memory: round(c.memory * 1e3, 4), + comms: round(c.comms * 1e3, 4), +}); + +export function candidateRow( + ctx: Context, + phase: SearchPhase, + c: Candidate, +): EstimateRow { + const { mesh } = c.deployment; + const dcp = c.deployment.decodeContextParallel ?? 1; + const sizes = sizesOf(mesh); + const price = ctx.chip.costPerHour; + const rate = c.result.tokPerSecPerChip; + const row: EstimateRow = { + ...base(ctx, phase), + feasible: true, + sharding: shardingLabel(ctx.model, mesh), + sizes, + dcp, + moeDispatch: sizes.EP > 1 ? c.deployment.moeDispatch : null, + placement: placementLabel(ctx.model, mesh, dcp), + batch: c.batch ?? null, + tokPerSecPerChip: round(rate, 1), + tokPerSecMachine: round(rate * ctx.nChips, 0), + stepTimeMs: round(c.result.stepTime * 1e3, 3), + boundBy: boundBy(c.result.cost.busy), + busyMs: ms(c.result.cost.busy), + visibleMs: ms(c.result.cost.parts), + costPerHour: price ?? null, + usdPerMtok: price !== undefined ? round(usdPerMtok(price, rate), 4) : null, + diagnostics: c.result.diags, + id: `${ctx.chip.id}|${ctx.machineName}|${JSON.stringify(sizes)}|dcp${dcp}`, + }; + + if (!('tpot' in c.result)) { + // a searched prefill pass: its own latency is the TTFT only when it + // was a single-sequence pass + const single = phase.kind === 'prefill' && phase.mode === 'ttft'; + row.prefill = { + tokPerSecPerChip: round(rate, 1), + ttftMs: single ? round(c.result.latency * 1e3, 3) : null, + latencyMs: round(c.result.latency * 1e3, 3), + batchSeqs: phase.kind === 'prefill' ? phase.seqs : 0, + boundBy: boundBy(c.result.cost.busy), + }; + return row; + } + + const full = enrichCandidate(ctx.model, ctx.chip, ctx.chip.id, ctx.nChips, ctx.workload, c); + const dec = full?.decode; + const mem = full?.memory; + return { + ...row, + tpotMs: round(c.result.tpot * 1e3, 3), + tokPerSecPerUser: round(1 / c.result.tpot, 2), + mfu: nullable(dec?.mfu, 4), + mbu: nullable(dec?.mbu, 4), + weightGbPerChip: round(c.result.memory.weightBytesPerChip / 1e9, 2), + kvMbPerSeqPerChip: round(c.result.memory.kvBytesPerSeqPerChip / 1e6, 3), + pagedKvMbPerSeqPerChip: round(c.result.memory.pagedKvBytesPerSeqPerChip / 1e6, 3), + stateMbPerSeqPerChip: round(c.result.memory.stateBytesPerSeqPerChip / 1e6, 3), + stateSlotsPerSeq: c.result.memory.stateSlotsPerSeq, + maxResidentSeqs: mem?.maxResidentSeqs ?? null, + prefill: full?.prefill + ? { + tokPerSecPerChip: round(full.prefill.tokPerSecPerChip, 1), + ttftMs: round(full.prefill.ttft * 1e3, 3), + latencyMs: null, + batchSeqs: full.prefill.batchSeqs, + boundBy: full.prefill.boundBy, + } + : null, + }; +} diff --git a/src/cli/version.ts b/src/cli/version.ts new file mode 100644 index 0000000..ac6249a --- /dev/null +++ b/src/cli/version.ts @@ -0,0 +1,5 @@ +// Stamped by scripts/build-cli.mjs from git describe; 'dev' when the +// sources run unbundled (vitest imports the commands in-process). +declare const __CLI_VERSION__: string | undefined; + +export const VERSION: string = typeof __CLI_VERSION__ === 'string' ? __CLI_VERSION__ : 'dev'; diff --git a/src/core/engine/optimizer/enrich.ts b/src/core/engine/optimizer/enrich.ts new file mode 100644 index 0000000..83b5716 --- /dev/null +++ b/src/core/engine/optimizer/enrich.ts @@ -0,0 +1,184 @@ +import { matmulSeconds, roofline } from '../roofline'; +import { evaluateDecodeAtBatch } from '../sim/run/decode'; +import { evaluatePrefill } from '../sim/run/prefill'; +import { runnableOn } from '../sim/run/validate'; +import { ROLES, ShardingRole } from '../sim/ir/sharding/roles'; +import { Diagnostic, Mesh, MoeDispatch, roleSize } from '../surface/deploy'; +import type { ResourceCostBackend } from '../sim/cost/types'; +import type { HardwareResource } from '../surface/api'; +import { flopsPerDecodeToken, flopsPerPrefillToken, hasMoeLayers } from '../../model/utils'; +import { ChipSpec } from '../../hardware/chips'; +import { ModelSpec } from '../../model/models'; +import type { Candidate } from './search'; + +export type ComponentTimes = Record; +export type Boundedness = HardwareResource; + +// One searched configuration, evaluated for both phases: display-ready +// scalars plus the winning mesh, so a trace can be re-lowered from it. +export interface ConfigResult { + id: string; + chipId: string; + sizes: Partial>; + placement: string; + // how many of the TP ranks hold a sequence slice instead of a head slice + decodeContextParallel: number; + dispatch: MoeDispatch; + // the placement's resolved mesh + mesh: Mesh; + nChips: number; + workload: { prefillLen: number; generateLen: number }; + diagnostics: Diagnostic[]; + memory?: { + weightBytesPerChip: number; + // HBM left for KV after the weights land: what actually caps batch + kvSpaceBytesPerChip: number; + kvBytesPerSeqPerChip: number; + // the paged and reserved-state parts of kvBytesPerSeqPerChip, and how + // many state slots the latter reserves per sequence + pagedKvBytesPerSeqPerChip: number; + stateBytesPerSeqPerChip: number; + stateSlotsPerSeq: number; + // machine total (per-chip residency x DPA groups) + maxResidentSeqs: number; + }; + prefill?: { + tokPerSecPerChip: number; + ttft: number; + batchSeqs: number; + mfu: number; + fracOfCeiling: number; + boundBy: Boundedness; + components: ComponentTimes; + // the visible split of the phase time after overlap (sums to the time) + visible: ComponentTimes; + }; + decode?: { + tokPerSecPerChip: number; + tokPerSecPerUser: number; + tpot: number; + stepTime: number; + batchPerStage: number; + residentSeqs: number; + mfu: number; + // model bandwidth utilization: weight + KV bytes streamed per step over + // what the chip's peak HBM bandwidth could move in a step + mbu: number; + fracOfCeiling: number; + // operating tok/s/chip over this config's own B -> inf rate (KV gate + // off); low = throughput is KV-room-starved, not sharding-limited + batchSaturation?: number; + boundBy: Boundedness; + components: ComponentTimes; + // the visible split of the step time after overlap (sums to the time) + visible: ComponentTimes; + }; +} + +export function boundBy(b: ComponentTimes): Boundedness { + return b.compute >= b.memory && b.compute >= b.comms + ? 'compute' + : b.memory >= b.comms + ? 'memory' + : 'comms'; +} + +// A decode-search winner filled out into one config row: prefill is +// evaluated on the same deployment (throughput at a full-machine batch of +// sequences, plus a single-sequence pass for TTFT). The candidate's own +// backend prices both, so the row is the search's numbers, not a mirror. +export function enrichCandidate( + model: ModelSpec, + chip: ChipSpec, + key: string, + nChips: number, + workload: { prefillLen: number; generateLen: number }, + c: Candidate, + backend: ResourceCostBackend = c.backend, +): ConfigResult | undefined { + const { mesh } = c.deployment; + const input = { model, deployment: c.deployment, workload }; + const dec = c.result; + if (!('tpot' in dec)) return undefined; + + const runnableModel = runnableOn(model, chip); + const hw = roofline(model, chip, workload, chip.realizableFlopsFrac)!; + const dpa = roleSize(mesh, 'DPA'); + const pp = roleSize(mesh, 'PP'); + const ctxAvg = workload.prefillLen + workload.generateLen / 2; + const prefillBatchSeqs = dpa * Math.max(1, Math.ceil(hw.critTokens / workload.prefillLen)); + + const opts = { costBackend: backend }; + const pfThrough = evaluatePrefill(input, prefillBatchSeqs, 'throughput', opts); + const pfSingle = evaluatePrefill(input, 1, 'ttft', opts); + // the same config re-priced at an effectively infinite batch (KV gate + // off) — its own batch-scaling ceiling + const sat = evaluateDecodeAtBatch(input, 65536 * dpa, pp, { ...opts, ignoreKvCapacity: true }); + + // the expert plane is structural noise on dense models (never read by + // lowering), so hide it from the displayed sharding + const shown = ROLES.filter( + (r) => roleSize(mesh, r) > 1 && (hasMoeLayers(model) || (r !== 'EP' && r !== 'ETP')), + ); + const sizes = Object.fromEntries(shown.map((r) => [r, roleSize(mesh, r)])); + // DCP is not a role: it owns no dims, it re-spends TP's on the sequence. + // It still belongs in the identity and the label, since two rows can + // otherwise differ only by it. + const dcp = c.deployment.decodeContextParallel ?? 1; + const placement = + shown.map((r) => `${r}[${mesh.roles[r].join('')}]`).join(' ') + + (dcp > 1 ? ` DCP=${dcp} of TP` : ''); + + return { + id: `${key}|${JSON.stringify(sizes)}|dcp${dcp}`, + chipId: chip.id, + sizes, + decodeContextParallel: dcp, + placement, + dispatch: c.deployment.moeDispatch, + mesh, + nChips, + workload, + diagnostics: dec.diags, + memory: { + weightBytesPerChip: dec.memory.weightBytesPerChip, + kvSpaceBytesPerChip: Math.max(0, chip.hbmCapacity - dec.memory.weightBytesPerChip), + kvBytesPerSeqPerChip: dec.memory.kvBytesPerSeqPerChip, + pagedKvBytesPerSeqPerChip: dec.memory.pagedKvBytesPerSeqPerChip, + stateBytesPerSeqPerChip: dec.memory.stateBytesPerSeqPerChip, + stateSlotsPerSeq: dec.memory.stateSlotsPerSeq, + maxResidentSeqs: dec.memory.maxResidentSeqsPerChip * dpa, + }, + decode: { + tokPerSecPerChip: dec.tokPerSecPerChip, + tokPerSecPerUser: 1 / dec.tpot, + tpot: dec.tpot, + stepTime: dec.stepTime, + batchPerStage: c.batch!, + residentSeqs: c.batch! * pp, + mfu: + dec.tokPerSecPerChip * matmulSeconds(flopsPerDecodeToken(runnableModel, ctxAvg), chip, 1)!, + // like MFU, quoted against the datasheet peak, not the realizable fraction + mbu: (dec.traffic.weightBytes + dec.traffic.kvBytes) / dec.stepTime / chip.hbmBandwidth, + fracOfCeiling: dec.tokPerSecPerChip / hw.decodeCeilingOverlapped, + batchSaturation: sat.ok ? dec.tokPerSecPerChip / sat.tokPerSecPerChip : undefined, + boundBy: boundBy(dec.cost.busy), + components: dec.cost.busy, + visible: dec.cost.parts, + }, + prefill: pfThrough.ok + ? { + tokPerSecPerChip: pfThrough.tokPerSecPerChip, + ttft: pfSingle.ok ? pfSingle.latency : pfThrough.latency, + batchSeqs: prefillBatchSeqs, + mfu: + pfThrough.tokPerSecPerChip * + matmulSeconds(flopsPerPrefillToken(runnableModel, workload.prefillLen), chip, 1)!, + fracOfCeiling: pfThrough.tokPerSecPerChip / hw.prefillCeiling, + boundBy: boundBy(pfThrough.cost.busy), + components: pfThrough.cost.busy, + visible: pfThrough.cost.parts, + } + : undefined, + }; +} diff --git a/src/core/engine/optimizer/policy.ts b/src/core/engine/optimizer/policy.ts index 3f2bb5d..14bf7f6 100644 --- a/src/core/engine/optimizer/policy.ts +++ b/src/core/engine/optimizer/policy.ts @@ -36,6 +36,7 @@ export function operatingBatch( deployment, stages, workload.prefillLen + workload.generateLen, + opts.statePool, ); // batch sizes ONE of the pp microbatches const cap = Math.floor(memory.maxResidentSeqsPerChip / pp); diff --git a/src/core/engine/optimizer/search.ts b/src/core/engine/optimizer/search.ts index 0682e26..fce7261 100644 --- a/src/core/engine/optimizer/search.ts +++ b/src/core/engine/optimizer/search.ts @@ -25,6 +25,8 @@ export interface SearchOptions extends EvalOptions export interface Candidate { deployment: Deployment; + // the backend that priced it, so the candidate can be filled out further + backend: TBackend; // decode operating batch the policy chose (absent for prefill) batch?: number; result: Extract | PrefillEvaluation, { ok: true }>; @@ -132,8 +134,10 @@ export function* searchShardings( } } -// Evaluate one deployment, dropping infeasible results and computing its search score. -function evaluateWithSearchOpts( +// Evaluate one deployment, dropping infeasible results and computing its +// search score: what the search does per placement, for callers scoring +// one explicit deployment the same way. +export function evaluateWithSearchOpts( input: SimInput, opts: SearchOptions, ): Candidate | null { @@ -164,5 +168,5 @@ function evaluateWithSearchOpts( score = (score * 3600) / price; // tokens per dollar } - return { deployment: input.deployment, batch, result, score }; + return { deployment: input.deployment, backend: opts.costBackend, batch, result, score }; } diff --git a/src/core/engine/sim/cost/dagSchedule.ts b/src/core/engine/sim/cost/dagSchedule.ts new file mode 100644 index 0000000..bdeadaa --- /dev/null +++ b/src/core/engine/sim/cost/dagSchedule.ts @@ -0,0 +1,159 @@ +import { collectiveCost } from './helpers/collectives'; +import { naiveOpCost } from './helpers/naiveOpCost'; +import type { Deployment } from '../../surface/deploy'; +import type { ExpandedOp, OpId, Segment } from '../ir/ops'; +import type { HardwareResource } from '../../surface/api'; +import type { OpCost } from './helpers/naiveOpCost'; +import type { CostBackend, ResourceTraceCost } from './types'; + +const RESOURCES: readonly HardwareResource[] = ['compute', 'memory', 'comms']; + +// What set a trace's time: one of the three streams ran out of room, or +// the op graph's dependency chain did before any stream filled. +export type ScheduleBound = HardwareResource | 'deps'; + +export interface DagScheduleTraceCost extends ResourceTraceCost { + // per-resource work that ran under something else (busy - parts) + hidden: Record; + bound: ScheduleBound; +} + +export interface DagScheduleOptions { + // weight streaming is issued ahead of the ops that need it, so it never + // sits on the dependency chain: the memory stream's load still counts, + // the recurrence does not wait for it. Off puts every load back in + // line, the way a stack without prefetch runs. Default on. + prefetchWeights?: boolean; +} + +// Overlap read off the op graph instead of two hand-set fractions. Each +// op keeps the naive roofline price on its streams; a segment then runs as +// a software-pipelined loop over its repeats, whose period is the larger +// of the resource bound (the fullest stream) and the recurrence bound (the +// longest dependency chain through one iteration, each op as wide as its +// slowest stream). The first iteration also waits for its weights. The +// result is always within [max stream, sum of streams]: nothing hides more +// than a stream can absorb, and nothing serializes past the chain. +export function makeDagScheduleBackend(options: DagScheduleOptions = {}) { + const prefetch = options.prefetchWeights ?? true; + return ((deployment: Deployment) => { + return { + // collectives are priced exactly as the naive backend prices them, so + // the two share reshard plans + priceCollectiveHash: JSON.stringify(['naive-op-cost-sum', deployment.mesh.dims]), + priceCollective: (kind, over, input, elemBytes) => + collectiveCost(kind, over, input, elemBytes, deployment.mesh.dims), + priceTrace: (trace: Segment[]): DagScheduleTraceCost => { + const busy = zero(); + const parts = zero(); + const reason: Record = { ...zero(), deps: 0 }; + const busyPerOp = new Map>(); + + for (const s of trace) { + const costs = new Map(); + const load = zero(); + for (const op of s.ops) { + const c = naiveOpCost(op, deployment); + costs.set(op.id, c); + for (const r of RESOURCES) { + load[r] += c[r]; + busy[r] += c[r] * s.repeat; + } + busyPerOp.set(op.id, { + compute: c.compute * s.repeat, + memory: c.memory * s.repeat, + comms: c.comms * s.repeat, + }); + } + + const full = criticalPath(s.ops, costs, () => true); + const steady = prefetch + ? criticalPath(s.ops, costs, (op) => op.kind !== 'weight-load') + : full; + + // the first iteration: streams or the chain with its weights + charge(parts, reason, load, full, 1); + // every later one: streams or the chain, weights already there + charge(parts, reason, load, steady, s.repeat - 1); + } + + const time = parts.compute + parts.memory + parts.comms; + const hidden = zero(); + for (const r of RESOURCES) hidden[r] = Math.max(0, busy[r] - parts[r]); + const bound = (Object.keys(reason) as ScheduleBound[]).reduce((a, b) => + reason[b] > reason[a] ? b : a, + ); + return { time, busy, busyPerOp, parts, hidden, bound }; + }, + }; + }) satisfies CostBackend; +} + +const zero = (): Record => ({ compute: 0, memory: 0, comms: 0 }); + +// The longest dependency chain through the segment's ops, each op as long +// as its slowest stream. Deps outside the segment finished before it +// started. The chain's time is returned split by the resource each op on +// it was widest on, so the caller can attribute it. +interface Chain { + time: number; + parts: Record; +} + +function criticalPath( + ops: ExpandedOp[], + costs: Map, + include: (op: ExpandedOp) => boolean, +): Chain { + const byId = new Map(ops.filter(include).map((op) => [op.id, op])); + const memo = new Map(); + const empty = (): Chain => ({ time: 0, parts: zero() }); + + const finish = (id: OpId): Chain => { + const op = byId.get(id); + if (!op) return empty(); + const hit = memo.get(id); + if (hit) return hit; + + let best = empty(); + for (const dep of op.deps) { + const chain = finish(dep); + if (chain.time > best.time) best = chain; + } + const c = costs.get(id)!; + const widest = RESOURCES.reduce((a, b) => (c[b] > c[a] ? b : a)); + const out: Chain = { + time: best.time + c[widest], + parts: { ...best.parts, [widest]: best.parts[widest] + c[widest] }, + }; + memo.set(id, out); + return out; + }; + + let best = empty(); + for (const id of byId.keys()) { + const chain = finish(id); + if (chain.time > best.time) best = chain; + } + return best; +} + +// One iteration's period is the larger of the fullest stream and the +// chain: charge that many iterations to whichever it was. +function charge( + parts: Record, + reason: Record, + load: Record, + chain: Chain, + iterations: number, +): void { + if (iterations <= 0) return; + const fullest = RESOURCES.reduce((a, b) => (load[b] > load[a] ? b : a)); + if (load[fullest] >= chain.time) { + parts[fullest] += load[fullest] * iterations; + reason[fullest] += load[fullest] * iterations; + } else { + for (const r of RESOURCES) parts[r] += chain.parts[r] * iterations; + reason.deps += chain.time * iterations; + } +} diff --git a/src/core/engine/sim/cost/naiveOpCostSum.ts b/src/core/engine/sim/cost/naiveOpCostSum.ts index 40e8082..4d618c1 100644 --- a/src/core/engine/sim/cost/naiveOpCostSum.ts +++ b/src/core/engine/sim/cost/naiveOpCostSum.ts @@ -5,16 +5,11 @@ import { naiveOverlapBreakdown } from './helpers/naiveOverlap'; import type { OpId, Segment } from '../ir/ops'; import type { HardwareResource } from '../../surface/api'; -import type { CostBackend, TraceCost } from './types'; +import type { CostBackend, ResourceTraceCost } from './types'; export type { CostBackend, TraceCost } from './types'; -export interface NaiveOpCostSumTraceCost extends TraceCost { - // per-resource busy sums before overlap - busy: Record; - // per-resource cost sums after overlap - parts: Record; -} +export type NaiveOpCostSumTraceCost = ResourceTraceCost; export function makeNaiveOpCostSumBackend(options: { memoryOverlap: number; diff --git a/src/core/engine/sim/cost/select.ts b/src/core/engine/sim/cost/select.ts new file mode 100644 index 0000000..0d21c4e --- /dev/null +++ b/src/core/engine/sim/cost/select.ts @@ -0,0 +1,18 @@ +import { makeDagScheduleBackend } from './dagSchedule'; +import { makeNaiveOpCostSumBackend } from './naiveOpCostSum'; +import type { ResourceCostBackend } from './types'; + +// 'naive' hides fixed fractions of the memory and comms streams behind the +// widest one (the two overlap constants); 'dag' reads the overlap off the +// op graph and ignores the constants. +export type Scheduler = 'naive' | 'dag'; + +export interface OverlapOptions { + memoryOverlap: number; + commsOverlap: number; + scheduler?: Scheduler; +} + +export function makeCostBackend(o: OverlapOptions): ResourceCostBackend { + return o.scheduler === 'dag' ? makeDagScheduleBackend() : makeNaiveOpCostSumBackend(o); +} diff --git a/src/core/engine/sim/cost/types.ts b/src/core/engine/sim/cost/types.ts index 7fff08b..e99b7f3 100644 --- a/src/core/engine/sim/cost/types.ts +++ b/src/core/engine/sim/cost/types.ts @@ -16,6 +16,17 @@ export interface TraceCost { busyPerOp?: ReadonlyMap>; } +// A trace price that also says how long each resource was busy (before +// any overlap) and how the wall-clock time splits between them (after). +// Both naive and scheduled backends answer this, so anything that shows +// a bound-by or a stream bar can take either. +export interface ResourceTraceCost extends TraceCost { + // per-resource busy sums before overlap + busy: Record; + // additive visible parts after overlap, they sum to time + parts: Record; +} + // A backend bound to one evaluation's context. priceCollective() // is used by the reshard expander to price candidate collectives. export interface BoundBackend { @@ -39,6 +50,11 @@ export interface BoundBackend { // A backend owns all physics: bind it to a deployment to price anything. export type CostBackend = (deployment: Deployment) => BoundBackend; +// A backend whose trace prices carry the per-resource breakdown. +export type ResourceCostBackend = ( + deployment: Deployment, +) => BoundBackend & { priceTrace(trace: Segment[]): ResourceTraceCost }; + // The price a given backend's priceTrace returns: the backend-specific // cost breakdown of one stage trace. export type TracePriceOf = ReturnType< diff --git a/src/core/engine/sim/run/decode.ts b/src/core/engine/sim/run/decode.ts index c57dbc9..335ed66 100644 --- a/src/core/engine/sim/run/decode.ts +++ b/src/core/engine/sim/run/decode.ts @@ -36,6 +36,7 @@ export function evaluateDecodeAtBatch( deployment, stages, workload.prefillLen + workload.generateLen, + opts.statePool, ); // each DPA group holds batch/dpa of every resident microbatch's KV if (!opts.ignoreKvCapacity && (batch / dpa) * microbatches > memory.maxResidentSeqsPerChip) { diff --git a/src/core/engine/sim/run/memory.ts b/src/core/engine/sim/run/memory.ts index e770cd0..e64a42b 100644 --- a/src/core/engine/sim/run/memory.ts +++ b/src/core/engine/sim/run/memory.ts @@ -8,30 +8,43 @@ import { blockRouterParams, blockRoutedExpertParams, blockSharedExpertParams, + blockStateBytes, } from '../../../model/block'; import { chipsPerStage, dcpSize, kvFraction, Deployment, roleSize } from '../../surface/deploy'; import type { Stage } from '../lowering/stages'; -import type { MemoryFootprint } from '../../surface/api'; +import { DEFAULT_STATE_POOL, type MemoryFootprint, type StatePoolOptions } from '../../surface/api'; import { DTYPE_BYTES } from '../../../model/dtype'; // Exact per-chip memory footprint from the ordered stage map. Weights are // counted at the stored format on every chip: they stay packed, with a // widening kernel assumed where none ships (validate warns there). +// +// A sequence's bytes split into the paged cache and the recurrent state of +// linear-attention blocks. The state is reserved statePool slots at a time: +// one pool, so the resident count is where free HBM runs out under +// paged + slots * state per sequence (which is also the fixed point a +// two-pool manager sized "so both fill together" lands on). export function memoryFootprint( m: ModelSpec, d: Deployment, stages: Stage[], fullLen: number, + statePool: Partial = {}, ): MemoryFootprint { + const pool = { ...DEFAULT_STATE_POOL, ...statePool }; + const slots = pool.slotsPerSeq + pool.specSlots; const tp = roleSize(d.mesh, 'TP'); const B = (c: keyof typeof m.precision.weights) => DTYPE_BYTES[m.precision.weights[c]]; let worstWeights = 0; let worstKv = 0; + let worstPaged = 0; + let worstState = 0; let residents = Infinity; for (const s of stages) { let weights = 0; - let kvPerSeq = 0; + let pagedPerSeq = 0; + let statePerSlot = 0; for (const g of s.groups) for (const { block: b, count } of g.pattern) { const n = g.repeat * count; @@ -48,18 +61,25 @@ export function memoryFootprint( weights += (n * blockRoutedExpertParams(m, b).total * B('routedExperts')) / chipsPerStage(d); - kvPerSeq += - n * - blockKvBytes(b, DTYPE_BYTES[m.precision.kv], fullLen, 'store') * - kvFraction(blockKvHeads(b), tp, dcpSize(d)); + const share = kvFraction(blockKvHeads(b), tp, dcpSize(d)); + if (b.attn.kind === 'linear') + statePerSlot += n * blockStateBytes(b, pool.stateDtypeBytes) * share; + else + pagedPerSeq += n * blockKvBytes(b, DTYPE_BYTES[m.precision.kv], fullLen, 'store') * share; } const emb = (m.vocab * m.modelDim * B('embeddings')) / tp; if (s.hasEmbedding) weights += emb; if (s.hasUnembedding && !(m.tiedEmbeddings && s.hasEmbedding)) weights += emb; + const statePerSeq = slots * statePerSlot; + const kvPerSeq = pagedPerSeq + statePerSeq; worstWeights = Math.max(worstWeights, weights); - worstKv = Math.max(worstKv, kvPerSeq); + if (kvPerSeq > worstKv) { + worstKv = kvPerSeq; + worstPaged = pagedPerSeq; + worstState = statePerSeq; + } const free = d.chip.hbmCapacity - weights; residents = Math.min( @@ -70,6 +90,9 @@ export function memoryFootprint( return { weightBytesPerChip: worstWeights, kvBytesPerSeqPerChip: worstKv, + pagedKvBytesPerSeqPerChip: worstPaged, + stateBytesPerSeqPerChip: worstState, + stateSlotsPerSeq: slots, maxResidentSeqsPerChip: residents, }; } diff --git a/src/core/engine/sim/run/prefill.ts b/src/core/engine/sim/run/prefill.ts index c9b0828..e2db8ff 100644 --- a/src/core/engine/sim/run/prefill.ts +++ b/src/core/engine/sim/run/prefill.ts @@ -31,7 +31,7 @@ export function evaluatePrefill( const pp = roleSize(deployment.mesh, 'PP'); const stages = partitionIntoStages(model, pp); - const memory = memoryFootprint(model, deployment, stages, T); + const memory = memoryFootprint(model, deployment, stages, T, opts.statePool); // Throughput keeps one microbatch on every pipeline stage. TTFT runs // one request through the stages in sequence, so it has no PP multiplier. const residentSeqsPerChip = mode === 'throughput' ? (seqs / dpa) * pp : seqs; diff --git a/src/core/engine/surface/api.ts b/src/core/engine/surface/api.ts index 418bd63..a64c0e0 100644 --- a/src/core/engine/surface/api.ts +++ b/src/core/engine/surface/api.ts @@ -18,11 +18,34 @@ export interface SimInput { }; } +// How a serving stack provisions the recurrent state of linear-attention +// blocks. The simulator's default charges one state per resident sequence, +// which is the pure-physics floor; real hybrid KV managers keep several +// per request (SGLang: 3 with the radix cache off, 4 with extra_buffer_lazy, +// 5 with extra_buffer) plus one per speculative draft token, and on a model +// like Kimi K3 that pool, not the paged cache, caps concurrency. +export interface StatePoolOptions { + // state slots reserved per resident sequence (1 = one working state) + slotsPerSeq: number; + // extra slots per sequence for speculative decoding intermediates + // (DSPARK: draft block + 1) + specSlots: number; + // bytes per state element for capacity, when the stack stores the state + // narrower than the model's stateBytes (2 = bf16). Capacity only: the + // HBM traffic the step streams keeps the model's dtype. + stateDtypeBytes?: number; +} + +export const DEFAULT_STATE_POOL: StatePoolOptions = { slotsPerSeq: 1, specSlots: 0 }; + export interface EvalOptions { // Skip the KV-residency feasibility gate: evaluate the step as if the // batch fit. Only for B_inf saturation diagnostics (batchSaturation's own // ceiling), never for reported operating points. ignoreKvCapacity?: boolean; + // reserved recurrent-state slots for linear-attention blocks; unset means + // DEFAULT_STATE_POOL (one slot, nothing speculative, the model's dtype) + statePool?: Partial; // bound per evaluation, then prices traces and candidate collectives costBackend: TBackend; } @@ -53,8 +76,15 @@ export interface BaseEvaluation { export interface MemoryFootprint { // resident weight bytes on the heaviest chip weightBytesPerChip: number; - // KV bytes one full-length sequence costs its group's chips (worst stage) + // bytes one full-length sequence costs its group's chips (worst stage): + // paged cache plus every reserved state slot, what actually divides HBM kvBytesPerSeqPerChip: number; + // the paged (growing) part of that: MLA latents, GQA heads, windows + pagedKvBytesPerSeqPerChip: number; + // the recurrent-state part, all reserved slots included + stateBytesPerSeqPerChip: number; + // slots that state part reserves per sequence (slotsPerSeq + specSlots) + stateSlotsPerSeq: number; // sequences one chip's free HBM holds KV for (worst stage). Each DPA // group holds its own sequences, so the machine total is dpa times this. maxResidentSeqsPerChip: number; diff --git a/src/core/hardware/chips.ts b/src/core/hardware/chips.ts index 6d977fb..4b50424 100644 --- a/src/core/hardware/chips.ts +++ b/src/core/hardware/chips.ts @@ -313,6 +313,39 @@ export const CHIPS: ChipSpec[] = [ costPerHour: 5.4, tdp: 1400, }, + { + // ESTIMATED ENTRY (Morph, 2026-09): no vendor datasheet is encoded here + // yet. Every number below is derived: B300 SXM formats stepped up the way + // GB200 steps up B200 (bf16 2250 -> 2500), so fp8/mxfp8 5000e12 and + // fp4/mxfp4/nvfp4 15000e12 (NVIDIA quotes the rack at ~1.1 EF dense FP4 + // over 72 GPUs -> ~15 PF each); HBM 288 GB with 270e9 usable and 8e12 B/s + // as on the B300 entry; NVLink5 900e9 one-way at 2e-6 s, domain 64 of 72 + // for the same reason as GB200; costPerHour 6.75 = B300's 5.4 x the + // GB200-over-B200 NVL72 premium (1.25); tdp 1650 = ~3.6 kW superchip + // minus ~300 W of Grace, over its 2 GPUs. Update from the datasheet and + // from the actual rack rental rate. + id: 'gb300-nvl72', + mmaShapes: BLACKWELL_ULTRA_MMA, + name: 'GB300 NVL72', + vendor: 'NVIDIA', + formats: { + bf16: 2500e12, + fp8: 5000e12, + mxfp8: 5000e12, + fp4: 15000e12, + mxfp4: 15000e12, + nvfp4: 15000e12, + int8: 'kernel-widened', + int4: 'kernel-widened', + }, + hbmCapacity: 270e9, + hbmBandwidth: 8e12, + interconnect: { bandwidthPerChip: 900e9, latency: 2e-6, domainSize: 64 }, + realizableFlopsFrac: 0.75, + realizableHbmBwFrac: 0.85, + costPerHour: 6.75, + tdp: 1650, + }, { id: 'vr100-nvl72', mmaShapes: RUBIN_MMA, diff --git a/src/ui/machines.ts b/src/core/hardware/machines.ts similarity index 51% rename from src/ui/machines.ts rename to src/core/hardware/machines.ts index 7894a9b..d141aab 100644 --- a/src/ui/machines.ts +++ b/src/core/hardware/machines.ts @@ -1,6 +1,49 @@ -import { ChipSpec } from '../core/hardware/chips'; -import { SliceTopology } from '../core/hardware/topology'; -import type { UiChip } from './results'; +import { ChipSpec } from './chips'; +import { SliceTopology } from './topology'; + +// A chip with the machine it deploys on: a ring-fabric slice name or a +// switched-fabric node count, both defaulting in machineOf when unset. +export interface ChipOnMachine extends ChipSpec { + slice?: string; + nodes?: number; +} + +// The machine a search runs on, as searchShardings takes it. +export type Machine = SliceTopology | { domain: number; nodes: number }; + +// A machine spelled on a command line: a catalog slice name on a ring +// fabric ("4x4x4"), or "x" ("8x2", "32x1") on a +// switched fabric, where a bare count means one node. +export function parseMachine(chip: ChipSpec, spec: string): Machine { + const catalog = chip.interconnect.topologies; + if (catalog) { + const slice = catalog.find((s) => s.name === spec); + if (!slice) + throw new Error( + `${chip.id} has no ${spec} slice; it offers ${catalog.map((s) => s.name).join(', ')}`, + ); + return slice; + } + const m = /^(\d+)(?:x(\d+))?$/.exec(spec); + if (!m) throw new Error(`machine ${spec} is not x`); + const domain = Number(m[1]); + const nodes = m[2] === undefined ? 1 : Number(m[2]); + if (domain < 1 || domain > chip.interconnect.domainSize) + throw new Error(`${chip.id} domains hold up to ${chip.interconnect.domainSize} chips`); + const maxNodes = chip.interconnect.scaleOut?.maxNodes ?? 1; + if (nodes < 1 || nodes > maxNodes) + throw new Error(`${chip.id} scales out to ${maxNodes} node${maxNodes > 1 ? 's' : ''}`); + return { domain, nodes }; +} + +// The name parseMachine accepts for a machine, for labels and file names. +export function machineName(m: Machine): string { + return 'dims' in m ? m.name : `${m.domain}x${m.nodes}`; +} + +export function machineSize(m: Machine): number { + return 'dims' in m ? m.count : m.domain * m.nodes; +} // Most scale-out nodes selectable on a switched fabric (1 = no scale-out). export function maxNodesOf(chip: ChipSpec): number { @@ -14,7 +57,7 @@ export function slicesOf(chip: ChipSpec): SliceTopology[] | undefined { // Every machine a chip can be deployed on, as chips the rest of the UI can // treat independently: one per slice, or one per scale-out node count. -export function machineVariants(chip: UiChip): UiChip[] { +export function machineVariants(chip: ChipOnMachine): ChipOnMachine[] { const slices = slicesOf(chip); return slices ? slices.map((s) => ({ ...chip, slice: s.name })) @@ -23,7 +66,7 @@ export function machineVariants(chip: UiChip): UiChip[] { // The offered slice closest to a host count's worth of chips (ties go to // the larger slice). Sliced-fabric chips only. -function sliceAtHosts(chip: UiChip, hosts: number): SliceTopology { +function sliceAtHosts(chip: ChipOnMachine, hosts: number): SliceTopology { const target = hosts * chip.interconnect.chipsPerHost!; return slicesOf(chip)!.reduce((a, b) => { const da = Math.abs(a.count - target); @@ -36,19 +79,19 @@ function sliceAtHosts(chip: UiChip, hosts: number): SliceTopology { // scale-out nodes on a switched fabric, or the slice closest to the hosts' // worth of chips on a ring fabric. Fixed-size machines (NVL72, Trainium2) // resolve to their one machine. -export function machineAtNodes(chip: UiChip, nodes: number): UiChip { +export function machineAtNodes(chip: ChipOnMachine, nodes: number): ChipOnMachine { if (!slicesOf(chip)) return { ...chip, nodes: Math.min(nodes, maxNodesOf(chip)) }; return { ...chip, slice: sliceAtHosts(chip, nodes).name }; } // Identifies a chip *on a machine*: the key groups are held under, so a // machine sweep's variants of one chip stay separate rows. -export function machineKey(chip: UiChip): string { +export function machineKey(chip: ChipOnMachine): string { const m = machineOf(chip); return `${chip.id}|${'dims' in m ? m.name : m.nodes}`; } -export function machineOf(chip: UiChip): SliceTopology | { domain: number; nodes: number } { +export function machineOf(chip: ChipOnMachine): Machine { const catalog = chip.interconnect.topologies; // unset machine picks default to one node/host of chips if (!catalog) @@ -63,13 +106,12 @@ export function machineOf(chip: UiChip): SliceTopology | { domain: number; nodes // empty when the chip offers no scale-out at all: the pill would just say // "1 node" on every row -export function machineLabel(chip: UiChip): string { +export function machineLabel(chip: ChipOnMachine): string { const m = machineOf(chip); if ('dims' in m) return m.name; return maxNodesOf(chip) > 1 ? `${m.nodes} node${m.nodes > 1 ? 's' : ''}` : ''; } -export function machineChips(chip: UiChip): number { - const m = machineOf(chip); - return 'dims' in m ? m.count : m.domain * m.nodes; +export function machineChips(chip: ChipOnMachine): number { + return machineSize(machineOf(chip)); } diff --git a/src/core/model/block/index.ts b/src/core/model/block/index.ts index aba3237..5e56f6b 100644 --- a/src/core/model/block/index.ts +++ b/src/core/model/block/index.ts @@ -242,11 +242,7 @@ export function blockKvBytesParts( len: number, access: 'store' | 'read', ): KvByteParts { - if (b.attn.kind === 'linear') - return { - core: b.attn.valueHeads * b.attn.headDim * b.attn.valueHeadDim * b.attn.stateBytes, - indexer: 0, - }; + if (b.attn.kind === 'linear') return { core: blockStateBytes(b), indexer: 0 }; if (b.attn.kind === 'mla') { const { dc, dRope, dsa } = b.attn; const latentLen = access === 'read' && dsa ? Math.min(len, dsa.topk) : len; @@ -277,6 +273,17 @@ export function blockKvBytesParts( }; } +// Bytes of recurrent state one sequence holds for one linear-attention +// block (0 for every other kind). stateBytes overrides the model's state +// dtype, for a serving stack that keeps the state narrower than the model +// says (SGLang's --mamba-ssm-dtype bf16). +export function blockStateBytes(b: BlockSpec, stateBytes?: number): number { + if (b.attn.kind !== 'linear') return 0; + return ( + b.attn.valueHeads * b.attn.headDim * b.attn.valueHeadDim * (stateBytes ?? b.attn.stateBytes) + ); +} + export function blockKvBytes( b: BlockSpec, kvBytes: number, diff --git a/src/ui/App.tsx b/src/ui/App.tsx index dcc4f69..3ab4e76 100644 --- a/src/ui/App.tsx +++ b/src/ui/App.tsx @@ -24,7 +24,7 @@ import { machineKey, machineLabel, machineVariants, -} from './machines'; +} from '../core/hardware/machines'; import './app.css'; function SailResearchLogo() { @@ -115,8 +115,12 @@ export function App() { // stable identity so the baseline hook's change-detection key isn't // recomputed on unrelated renders const overlap = useMemo( - () => ({ memoryOverlap: sweep.memoryOverlap, commsOverlap: sweep.commsOverlap }), - [sweep.memoryOverlap, sweep.commsOverlap], + () => ({ + memoryOverlap: sweep.memoryOverlap, + commsOverlap: sweep.commsOverlap, + scheduler: sweep.scheduler, + }), + [sweep.memoryOverlap, sweep.commsOverlap, sweep.scheduler], ); const baseline = useHmvpBaseline(model, workload, overlap, chipsById[H100_ID]); @@ -202,6 +206,7 @@ export function App() { workload, sweep.memoryOverlap, sweep.commsOverlap, + sweep.scheduler, ]); const prev = lastSearched.current; const stale = enabled.filter((c) => { diff --git a/src/ui/components/DetailsPanel.tsx b/src/ui/components/DetailsPanel.tsx index 354feae..6c47527 100644 --- a/src/ui/components/DetailsPanel.tsx +++ b/src/ui/components/DetailsPanel.tsx @@ -95,7 +95,12 @@ export function DetailsPanel({ result: r, group, model, overlap, basis, onClose
- 1 ? 'stage' : 'batch'} /> + 1 ? 'stage' : 'batch'} + />
)} @@ -147,7 +152,7 @@ export function DetailsPanel({ result: r, group, model, overlap, basis, onClose
- +
)} @@ -297,14 +302,38 @@ function HbmBar({ r, cap }: { r: UiResult; cap: number }) { const COMPONENT_LABEL = { compute: 'Compute', memory: 'Memory', comms: 'Comms' } as const; +function scheduledBreakdown(busy: ComponentTimes, parts: ComponentTimes) { + const keys = ['compute', 'memory', 'comms'] as const; + const by = keys.reduce((x, y) => (parts[y] > parts[x] ? y : x)); + const hidden: ComponentTimes = { + compute: Math.max(0, busy.compute - parts.compute), + memory: Math.max(0, busy.memory - parts.memory), + comms: Math.max(0, busy.comms - parts.comms), + }; + return { parts, pool: { by, time: parts[by] }, hidden }; +} + /** * Phase time as a stacked bar: each solid segment is the visible time a * component contributes under the tuned overlap, the pool winner first. * Work hidden behind the pool is hatched inside the pool's span — * concurrent, adding no wall-clock time. */ -function StepTimeBar({ c, a, per }: { c: ComponentTimes; a: UiOverlap; per: string }) { - const { parts, pool, hidden } = naiveOverlapBreakdown(c, a); +function StepTimeBar({ + c, + visible, + a, + per, +}: { + c: ComponentTimes; + visible: ComponentTimes; + a: UiOverlap; + per: string; +}) { + // the fractions have a closed-form breakdown; a scheduled trace reports + // its visible split directly, and everything past it ran hidden + const { parts, pool, hidden } = + a.scheduler === 'dag' ? scheduledBreakdown(c, visible) : naiveOverlapBreakdown(c, a); const total = parts.compute + parts.memory + parts.comms; if (total <= 0) return null; const order = [ diff --git a/src/ui/components/OpDagView.tsx b/src/ui/components/OpDagView.tsx index b00f4ae..dfb8bd1 100644 --- a/src/ui/components/OpDagView.tsx +++ b/src/ui/components/OpDagView.tsx @@ -12,7 +12,7 @@ import { } from '@xyflow/react'; import '@xyflow/react/dist/style.css'; import type { ModelSpec } from '../../core/model/models'; -import { makeNaiveOpCostSumBackend } from '../../core/engine/sim/cost/naiveOpCostSum'; +import { makeCostBackend } from '../../core/engine/sim/cost/select'; import { evaluateDecodeAtBatch } from '../../core/engine/sim/run/decode'; import { evaluatePrefill } from '../../core/engine/sim/run/prefill'; import type { ExpandedOp, OpId } from '../../core/engine/sim/ir/ops'; @@ -357,9 +357,10 @@ export function OpDagView({ const { run, backend } = useMemo(() => { const deployment: Deployment = { chip, mesh: result.mesh, moeDispatch: result.dispatch }; const input = { model, deployment, workload: result.workload }; - const costBackend = makeNaiveOpCostSumBackend({ + const costBackend = makeCostBackend({ memoryOverlap: overlap.memoryOverlap, commsOverlap: overlap.commsOverlap, + scheduler: overlap.scheduler, }); const pp = result.sizes.PP ?? 1; const ev = @@ -374,7 +375,7 @@ export function OpDagView({ // same context the run bound, so this shares its reshard-plan cache backend: costBackend(deployment), }; - }, [result, model, chip, overlap.memoryOverlap, overlap.commsOverlap, phase]); + }, [result, model, chip, overlap.memoryOverlap, overlap.commsOverlap, overlap.scheduler, phase]); const stages = run?.perStageTrace ?? []; const stageCount = stages.length; diff --git a/src/ui/components/Sidebar.tsx b/src/ui/components/Sidebar.tsx index 507aa6c..4e626d5 100644 --- a/src/ui/components/Sidebar.tsx +++ b/src/ui/components/Sidebar.tsx @@ -5,9 +5,16 @@ import { kvBytesPerSeq } from '../../core/model/utils'; import { MODEL_PRESETS, ModelSpec } from '../../core/model/models'; import { DTYPE_BYTES } from '../../core/model/dtype'; import { fmtBytes } from '../format'; -import { machineAtNodes, machineLabel, machineOf, maxNodesOf, slicesOf } from '../machines'; +import { + machineAtNodes, + machineLabel, + machineOf, + maxNodesOf, + slicesOf, +} from '../../core/hardware/machines'; import { CostBasis, H100_ID, relPriceOf, relPriceToDollars } from '../pricing'; import { UiChip, UiWorkload } from '../results'; +import type { Scheduler } from '../../core/engine/sim/cost/select'; import { VendorLogo } from './VendorLogo'; export interface SweepControls { @@ -19,6 +26,8 @@ export interface SweepControls { memoryOverlap: number; /** fraction of collective traffic hidden behind compute */ commsOverlap: number; + /** how the streams overlap: the two fractions above, or the op graph */ + scheduler?: Scheduler; } interface Props { @@ -365,6 +374,26 @@ export function Sidebar(p: Props) {

Overlap

+
+ + Source + + ⓘ + + + +
>; - placement: string; - // how many of the TP ranks hold a sequence slice instead of a head slice - decodeContextParallel: number; - dispatch: MoeDispatch; - // the placement's resolved mesh, for re-lowering the trace in the details DAG - mesh: Mesh; - nChips: number; - workload: { prefillLen: number; generateLen: number }; - diagnostics: Diagnostic[]; - memory?: { - weightBytesPerChip: number; - // HBM left for KV after the weights land: what actually caps batch - kvSpaceBytesPerChip: number; - kvBytesPerSeqPerChip: number; - // machine total (per-chip residency x DPA groups) - maxResidentSeqs: number; - }; - // cost-efficiency for the whole workload vs the HMVP baseline; filled at - // render time from the live prices and baseline, never by the worker +// A streamed configuration row: the engine's result plus the efficiency +// figures filled at render time from the live prices and baseline, never +// by the worker. +export interface UiResult extends ConfigResult { + // cost-efficiency for the whole workload vs the HMVP baseline requestEff?: number; - prefill?: { - tokPerSecPerChip: number; - ttft: number; - batchSeqs: number; - mfu: number; - fracOfCeiling: number; - boundBy: Boundedness; - components: ComponentTimes; - // (rate ÷ relative price) over the HMVP's rate; filled at render time - // from the live prices and baseline, never by the worker - eff?: number; - // rate over the HMVP's rate — pure speed, price not included - relRate?: number; - }; - decode?: { - tokPerSecPerChip: number; - tokPerSecPerUser: number; - tpot: number; - stepTime: number; - batchPerStage: number; - residentSeqs: number; - mfu: number; - // model bandwidth utilization: weight + KV bytes streamed per step over - // what the chip's peak HBM bandwidth could move in a step - mbu: number; - fracOfCeiling: number; - // operating tok/s/chip over this config's own B -> inf rate (KV gate - // off); low = throughput is KV-room-starved, not sharding-limited - batchSaturation?: number; - boundBy: Boundedness; - components: ComponentTimes; - // (rate ÷ relative price) over the HMVP's rate; filled at render time - // from the live prices and baseline, never by the worker + prefill?: ConfigResult['prefill'] & { + // (rate ÷ relative price) over the HMVP's rate eff?: number; // rate over the HMVP's rate — pure speed, price not included relRate?: number; }; + decode?: ConfigResult['decode'] & { eff?: number; relRate?: number }; } // One chip's slot in the leaderboard: its fixed machine, the hardware diff --git a/src/ui/search.worker.ts b/src/ui/search.worker.ts index d302f83..9d741a9 100644 --- a/src/ui/search.worker.ts +++ b/src/ui/search.worker.ts @@ -1,17 +1,10 @@ -import { Candidate, searchShardings, searchTuples } from '../core/engine/optimizer/search'; +import { searchShardings, searchTuples } from '../core/engine/optimizer/search'; import { ServingPolicy } from '../core/engine/optimizer/policy'; -import { matmulSeconds, roofline } from '../core/engine/roofline'; -import { makeNaiveOpCostSumBackend } from '../core/engine/sim/cost/naiveOpCostSum'; -import { evaluateDecodeAtBatch } from '../core/engine/sim/run/decode'; -import { evaluatePrefill } from '../core/engine/sim/run/prefill'; -import { runnableOn } from '../core/engine/sim/run/validate'; -import { ROLES } from '../core/engine/sim/ir/sharding/roles'; -import { roleSize } from '../core/engine/surface/deploy'; -import { flopsPerDecodeToken, flopsPerPrefillToken, hasMoeLayers } from '../core/model/utils'; -import { ChipSpec } from '../core/hardware/chips'; -import { MODEL_PRESETS, ModelSpec } from '../core/model/models'; -import { machineChips, machineKey, machineLabel, machineOf } from './machines'; -import { Boundedness, ComponentTimes, UiResult } from './results'; +import { roofline } from '../core/engine/roofline'; +import { makeCostBackend } from '../core/engine/sim/cost/select'; +import { enrichCandidate } from '../core/engine/optimizer/enrich'; +import { MODEL_PRESETS } from '../core/model/models'; +import { machineChips, machineKey, machineLabel, machineOf } from '../core/hardware/machines'; import type { SearchRequest, SearchUpdate } from './searchClient'; // the client pins one worker per chip, so this instance only ever sees one @@ -29,7 +22,7 @@ async function run(req: SearchRequest) { const model = MODEL_PRESETS.find((m) => m.name === req.modelName); if (!model) return; - const backend = makeNaiveOpCostSumBackend(req.overlap); + const backend = makeCostBackend(req.overlap); const policy: ServingPolicy = { batching: req.workload.batching, sloTokPerSecPerUser: req.workload.sloTokPerSecPerUser || undefined, @@ -75,7 +68,7 @@ async function run(req: SearchRequest) { done: step.done, row: step.candidate && - toRow(model, chip, group.key, nChips, workload, backend, step.candidate), + enrichCandidate(model, chip, group.key, nChips, workload, step.candidate), }); await new Promise((r) => setTimeout(r)); } @@ -84,106 +77,3 @@ async function run(req: SearchRequest) { } } } - -type Backend = ReturnType; - -// A decode-search winner filled out into the UI's config row: prefill is -// evaluated on the same deployment (throughput at a full-machine batch of -// sequences, plus a single-sequence pass for TTFT). -function toRow( - model: ModelSpec, - chip: ChipSpec, - key: string, - nChips: number, - workload: { prefillLen: number; generateLen: number }, - backend: Backend, - c: Candidate, -): UiResult | undefined { - const { mesh } = c.deployment; - const input = { model, deployment: c.deployment, workload }; - const dec = c.result; - if (!('tpot' in dec)) return undefined; - - const runnableModel = runnableOn(model, chip); - const hw = roofline(model, chip, workload, chip.realizableFlopsFrac)!; - const dpa = roleSize(mesh, 'DPA'); - const pp = roleSize(mesh, 'PP'); - const ctxAvg = workload.prefillLen + workload.generateLen / 2; - const prefillBatchSeqs = dpa * Math.max(1, Math.ceil(hw.critTokens / workload.prefillLen)); - - const opts = { costBackend: backend }; - const pfThrough = evaluatePrefill(input, prefillBatchSeqs, 'throughput', opts); - const pfSingle = evaluatePrefill(input, 1, 'ttft', opts); - // the same config re-priced at an effectively infinite batch (KV gate - // off) — its own batch-scaling ceiling - const sat = evaluateDecodeAtBatch(input, 65536 * dpa, pp, { ...opts, ignoreKvCapacity: true }); - - const bound = (b: ComponentTimes): Boundedness => - b.compute >= b.memory && b.compute >= b.comms - ? 'compute' - : b.memory >= b.comms - ? 'memory' - : 'comms'; - - // the expert plane is structural noise on dense models (never read by - // lowering), so hide it from the displayed sharding - const shown = ROLES.filter( - (r) => roleSize(mesh, r) > 1 && (hasMoeLayers(model) || (r !== 'EP' && r !== 'ETP')), - ); - const sizes = Object.fromEntries(shown.map((r) => [r, roleSize(mesh, r)])); - // DCP is not a role: it owns no dims, it re-spends TP's on the sequence. - // It still belongs in the identity and the label, since two rows can - // otherwise differ only by it. - const dcp = c.deployment.decodeContextParallel ?? 1; - const placement = - shown.map((r) => `${r}[${mesh.roles[r].join('')}]`).join(' ') + - (dcp > 1 ? ` DCP=${dcp} of TP` : ''); - - return { - id: `${key}|${JSON.stringify(sizes)}|dcp${dcp}`, - chipId: chip.id, - sizes, - decodeContextParallel: dcp, - placement, - dispatch: c.deployment.moeDispatch, - mesh, - nChips, - workload, - diagnostics: dec.diags, - memory: { - weightBytesPerChip: dec.memory.weightBytesPerChip, - kvSpaceBytesPerChip: Math.max(0, chip.hbmCapacity - dec.memory.weightBytesPerChip), - kvBytesPerSeqPerChip: dec.memory.kvBytesPerSeqPerChip, - maxResidentSeqs: dec.memory.maxResidentSeqsPerChip * dpa, - }, - decode: { - tokPerSecPerChip: dec.tokPerSecPerChip, - tokPerSecPerUser: 1 / dec.tpot, - tpot: dec.tpot, - stepTime: dec.stepTime, - batchPerStage: c.batch!, - residentSeqs: c.batch! * pp, - mfu: - dec.tokPerSecPerChip * matmulSeconds(flopsPerDecodeToken(runnableModel, ctxAvg), chip, 1)!, - // like MFU, quoted against the datasheet peak, not the realizable fraction - mbu: (dec.traffic.weightBytes + dec.traffic.kvBytes) / dec.stepTime / chip.hbmBandwidth, - fracOfCeiling: dec.tokPerSecPerChip / hw.decodeCeilingOverlapped, - batchSaturation: sat.ok ? dec.tokPerSecPerChip / sat.tokPerSecPerChip : undefined, - boundBy: bound(dec.cost.busy), - components: dec.cost.busy, - }, - prefill: pfThrough.ok - ? { - tokPerSecPerChip: pfThrough.tokPerSecPerChip, - ttft: pfSingle.ok ? pfSingle.latency : pfThrough.latency, - batchSeqs: prefillBatchSeqs, - mfu: - pfThrough.tokPerSecPerChip * - matmulSeconds(flopsPerPrefillToken(runnableModel, workload.prefillLen), chip, 1)!, - fracOfCeiling: pfThrough.tokPerSecPerChip / hw.prefillCeiling, - boundBy: bound(pfThrough.cost.busy), - components: pfThrough.cost.busy, - } - : undefined, - }; -} diff --git a/src/ui/searchClient.ts b/src/ui/searchClient.ts index 71c9f7e..367a908 100644 --- a/src/ui/searchClient.ts +++ b/src/ui/searchClient.ts @@ -38,7 +38,7 @@ export type SearchUpdate = { id: number; key: string } & ( export function makeSearchClient(onUpdate: (u: SearchUpdate) => void) { const workers = new Map(); const buffer: SearchUpdate[] = []; - let flush = 0; + let flush: ReturnType | 0 = 0; const onMessage = (e: MessageEvent) => { buffer.push(e.data); if (!flush) diff --git a/tests/cli.test.ts b/tests/cli.test.ts new file mode 100644 index 0000000..f94b8dc --- /dev/null +++ b/tests/cli.test.ts @@ -0,0 +1,218 @@ +import { expect, test } from 'vitest'; +import { CHIPS_BY_ID } from '../src/core/hardware/chips'; +import { parseMachine } from '../src/core/hardware/machines'; +import { sharedOptions } from '../src/cli/args'; +import { DEFAULT_SHARED, makeContext } from '../src/cli/context'; +import { CliError, resolveChip, resolveModel } from '../src/cli/resolve'; +import { usdPerMtok } from '../src/cli/row'; +import { chipRows } from '../src/cli/commands/chips'; +import { modelRows } from '../src/cli/commands/models'; +import { explain, fillSizes, engineHints } from '../src/cli/commands/explain'; +import { rankedRows, renderRows } from '../src/cli/commands/search'; +import { renderReport, runSweep, sweepCells } from '../src/cli/commands/sweep'; + +test('models resolve by name, alias, or unique prefix', () => { + expect(resolveModel('Kimi K3 MXFP4/MXFP8').name).toBe('Kimi K3 MXFP4/MXFP8'); + expect(resolveModel('k3').name).toBe('Kimi K3 MXFP4/MXFP8'); + expect(resolveModel('kimi k3').name).toBe('Kimi K3 MXFP4/MXFP8'); + expect(resolveModel('gpt-oss-120b').name).toBe('gpt-oss-120b MXFP4/BF16'); + expect(resolveModel('dsv4flash').name).toBe('DeepSeek V4 Flash MXFP4/FP8'); + // "Kimi" alone is K2.6 or K3 + expect(() => resolveModel('kimi')).toThrow(CliError); + expect(() => resolveModel('nonesuch')).toThrow(/unknown model/); + expect(resolveChip('gb300-nvl72').id).toBe('gb300-nvl72'); + expect(resolveChip('gb300').id).toBe('gb300-nvl72'); + expect(() => resolveChip('h')).toThrow(/matches/); +}); + +test('machines parse as domain x nodes on switched fabrics and as slices on rings', () => { + const gb300 = CHIPS_BY_ID['gb300-nvl72']; + expect(parseMachine(gb300, '32x1')).toEqual({ domain: 32, nodes: 1 }); + expect(parseMachine(gb300, '64')).toEqual({ domain: 64, nodes: 1 }); + expect(() => parseMachine(gb300, '72x1')).toThrow(/up to 64/); + expect(() => parseMachine(gb300, '32x2')).toThrow(/scales out to 1 node/); + const h100 = CHIPS_BY_ID['h100-sxm']; + expect(parseMachine(h100, '8x2')).toEqual({ domain: 8, nodes: 2 }); + const v5p = CHIPS_BY_ID['tpu-v5p']; + expect(parseMachine(v5p, '2x2x2')).toMatchObject({ name: '2x2x2', count: 8 }); + expect(() => parseMachine(v5p, '8x1')).toThrow(/no 8x1 slice/); + expect(makeContext('k3', 'gb300', '32x1').nChips).toBe(32); +}); + +test('shared flags map onto the engine knobs', () => { + const o = sharedOptions({ + 'state-slots': '5', + 'spec-slots': '8', + 'state-dtype': 'bf16', + 'mem-fraction': '0.9', + scheduler: 'dag', + overlap: '0.5,0.25', + 'cost-per-hour': '7', + }); + expect(o.statePool).toEqual({ slotsPerSeq: 5, specSlots: 8, stateDtypeBytes: 2 }); + expect(o.memFraction).toBe(0.9); + expect(o.overlap).toEqual({ memoryOverlap: 0.5, commsOverlap: 0.25, scheduler: 'dag' }); + expect(o.costPerHour).toBe(7); + expect(sharedOptions({})).toEqual(DEFAULT_SHARED); + const ctx = makeContext('k3', 'b300', '8x1', o); + expect(ctx.chip.hbmCapacity).toBeCloseTo(0.9 * CHIPS_BY_ID['b300'].hbmCapacity, 0); + expect(ctx.chip.costPerHour).toBe(7); + expect(() => sharedOptions({ scheduler: 'magic' })).toThrow(/scheduler/); + expect(() => sharedOptions({ overlap: '2,0' })).toThrow(/overlap/); +}); + +test('the chip table carries the estimated GB300 entry', () => { + const k3 = resolveModel('k3'); + const row = chipRows(k3).find((r) => r.id === 'gb300-nvl72')!; + expect(row).toMatchObject({ hbmGb: 270, hbmTbps: 8, costPerHour: 6.75, domain: 64 }); + expect(row.minChipsForWeights).toBe(6); + expect(modelRows().find((m) => m.name.startsWith('Kimi K3'))).toMatchObject({ + layers: 93, + experts: 896, + minKvHeads: 1, + }); +}); + +test('$/Mtok is the fleet tooling’s cost_per_million_tokens', () => { + // cost_per_million_tokens(gpu_hourly, gpus, tok_s) = gpu_hourly*gpus / (tok_s*3600) * 1e6 + expect(usdPerMtok(6.75, 4371)).toBeCloseTo((6.75 / (4371 * 3600)) * 1e6, 12); +}); + +test('explain fills the roles the way the search sizes them', () => { + const ctx = makeContext('k3', 'gb300', '32x1'); + expect(fillSizes(ctx, { DPA: 16, TP: 2, EP: 32 })).toEqual({ + PP: 1, + DPA: 16, + TP: 2, + EP: 32, + ETP: 1, + }); + expect(fillSizes(ctx, { TP: 4 })).toEqual({ PP: 1, DPA: 8, TP: 4, EP: 32, ETP: 1 }); + expect(fillSizes(ctx, { PP: 2, TP: 8, ETP: 2 })).toEqual({ PP: 2, DPA: 2, TP: 8, EP: 8, ETP: 2 }); + expect(() => fillSizes(ctx, { TP: 3 })).toThrow(CliError); + expect(engineHints({ PP: 1, DPA: 16, TP: 2, EP: 32, ETP: 1 }, 2)).toEqual([ + 'vLLM: -tp 2 -dp 16 --enable-expert-parallel', + 'SGLang: --tp 32 --dp 16 --enable-dp-attention --ep 32 --decode-context-parallel-size 2', + ]); +}); + +test('the state pool caps explain’s batch near the measured B300 admission', () => { + const ctx = makeContext('k3', 'b300', '8x1', { + ...DEFAULT_SHARED, + prefillLen: 8192, + genLen: 1024, + memFraction: 0.9, + statePool: { slotsPerSeq: 5 }, + }); + const e = explain(ctx, { phase: 'decode', sizes: { TP: 8, EP: 8 }, dcp: 1 }); + expect(e.best.maxResidentSeqs).toBe(117); + expect(e.best.batch).toBe(117); + expect(e.best.stateSlotsPerSeq).toBe(5); + expect(e.best.kvMbPerSeqPerChip).toBeCloseTo( + e.best.pagedKvMbPerSeqPerChip! + e.best.stateMbPerSeqPerChip!, + 2, + ); +}); + +test('the pinned GB300 search: K3 on 32 chips at 20 tok/s/user lands on DPA=16 TP=2 EP=32 DCP=2', () => { + const ctx = makeContext('k3', 'gb300-nvl72', '32x1'); + const [top] = rankedRows(ctx, { phase: 'decode', slo: 20 }, 1); + expect(top.sizes).toEqual({ PP: 1, DPA: 16, TP: 2, EP: 32, ETP: 1 }); + expect(top.dcp).toBe(2); + expect(top.moeDispatch).toBe('coalesced-a2a'); + expect(top.sharding).toBe('DPA=16 TP=2 EP=32'); + expect(top.tokPerSecPerUser).toBeGreaterThanOrEqual(20); + expect(top.tokPerSecPerUser).toBeLessThan(21); + expect(top.usdPerMtok).toBeCloseTo(usdPerMtok(6.75, top.tokPerSecPerChip!), 3); + expect(top.batch).toBeLessThanOrEqual(top.maxResidentSeqs!); + expect(top.prefill?.ttftMs).toBeGreaterThan(0); + expect(renderRows([top], 1)).toContain('DPA=16 TP=2 EP=32 dcp=2 coalesced-a2a'); + + // the row's key set is the contract sweep consumers read + expect(Object.keys(top).sort()).toEqual( + [ + 'batch', + 'boundBy', + 'busyMs', + 'chip', + 'costPerHour', + 'dcp', + 'diagnostics', + 'error', + 'feasible', + 'id', + 'kvMbPerSeqPerChip', + 'machine', + 'maxResidentSeqs', + 'mbu', + 'mfu', + 'model', + 'moeDispatch', + 'nChips', + 'overlap', + 'pagedKvMbPerSeqPerChip', + 'phase', + 'placement', + 'prefill', + 'scheduler', + 'sharding', + 'sizes', + 'slo', + 'stateMbPerSeqPerChip', + 'stateSlotsPerSeq', + 'statePool', + 'stepTimeMs', + 'tokPerSecMachine', + 'tokPerSecPerChip', + 'tokPerSecPerUser', + 'tpotMs', + 'usdPerMtok', + 'version', + 'visibleMs', + 'weightGbPerChip', + 'workload', + ].sort(), + ); + + // explain on the winner's sizes reproduces the winner + const e = explain(ctx, { phase: 'decode', slo: 20, sizes: { DPA: 16, TP: 2, EP: 32 } }); + expect(e.best.batch).toBe(top.batch); + expect(e.best.tokPerSecPerChip).toBe(top.tokPerSecPerChip); + expect(e.best.placement).toBe(top.placement); +}, 180_000); + +test('a sweep reports every cell and the report groups them', async () => { + const cells = sweepCells({ + model: 'k3', + chips: ['b300', 'nonesuch'], + machines: ['8x1'], + phases: ['decode', 'prefill'], + slos: [null, 50], + }); + expect(cells.map((c) => `${c.chip}/${c.phase}/${c.slo}`)).toEqual([ + 'b300/decode/null', + 'b300/decode/50', + 'b300/prefill/null', + 'nonesuch/decode/null', + 'nonesuch/decode/50', + 'nonesuch/prefill/null', + ]); + const seen: string[] = []; + const rows = await runSweep(cells, DEFAULT_SHARED, 1, (r) => seen.push(r.chip)); + expect(seen.length).toBe(6); + const ok = rows.filter((r) => r.feasible); + expect(ok.map((r) => `${r.phase}/${r.slo}`)).toEqual([ + 'decode/null', + 'decode/50', + 'prefill/null', + ]); + const bad = rows.filter((r) => !r.feasible); + expect(bad).toHaveLength(3); + expect(bad[0].error).toMatch(/unknown chip/); + const at50 = ok.find((r) => r.slo === 50)!; + expect(at50.tokPerSecPerUser).toBeGreaterThanOrEqual(50); + const report = renderReport(rows); + expect(report).toContain('DECODE at >= 50 tok/s/user'); + expect(report).toContain('PREFILL throughput'); + expect(report).toContain('3 infeasible'); +}, 120_000); diff --git a/tests/dagSchedule.test.ts b/tests/dagSchedule.test.ts new file mode 100644 index 0000000..c85a901 --- /dev/null +++ b/tests/dagSchedule.test.ts @@ -0,0 +1,232 @@ +import { expect, test } from 'vitest'; +import { makeDagScheduleBackend } from '../src/core/engine/sim/cost/dagSchedule'; +import { makeNaiveOpCostSumBackend } from '../src/core/engine/sim/cost/naiveOpCostSum'; +import { makeCostBackend } from '../src/core/engine/sim/cost/select'; +import { naiveOpCost } from '../src/core/engine/sim/cost/helpers/naiveOpCost'; +import type { ExpandedOp, OpId, Segment } from '../src/core/engine/sim/ir/ops'; +import { tt } from '../src/core/engine/sim/ir/tensors'; +import { evaluateDecodeAtBatch } from '../src/core/engine/sim/run/decode'; +import { evaluatePrefill } from '../src/core/engine/sim/run/prefill'; +import { searchShardings } from '../src/core/engine/optimizer/search'; +import { Deployment, makeMesh } from '../src/core/engine/surface/deploy'; +import { CHIPS_BY_ID } from '../src/core/hardware/chips'; +import { deployedAxes } from '../src/core/hardware/topology'; +import { MODEL_PRESETS } from '../src/core/model/models'; + +const h100 = CHIPS_BY_ID['h100-sxm']; +const naive = makeNaiveOpCostSumBackend({ memoryOverlap: 0.9, commsOverlap: 0.65 }); +const serial = makeNaiveOpCostSumBackend({ memoryOverlap: 0, commsOverlap: 0 }); +const dag = makeDagScheduleBackend(); + +function deploymentOn(chip = h100, tp = 1): Deployment { + const axes = deployedAxes(chip.interconnect, { domain: tp, nodes: 1 }); + const names = tp > 1 ? axes.map((a) => a.name) : []; + return { + chip, + mesh: makeMesh(axes, { DPA: [], TP: names, EP: names, ETP: [], PP: [] }), + moeDispatch: 'ring-of-experts', + }; +} + +const id = (s: string) => s as OpId; +const RES = ['compute', 'memory', 'comms'] as const; +const sum = (c: Record<(typeof RES)[number], number>) => c.compute + c.memory + c.comms; +const max = (c: Record<(typeof RES)[number], number>) => Math.max(c.compute, c.memory, c.comms); + +// a decode-shaped gemm: few rows, so it streams more than it multiplies +const gemm = (name: string, deps: string[], m: number, k: number, n: number): ExpandedOp => ({ + id: id(name), + label: name, + deps: deps.map(id), + kind: 'gemm', + x: tt([m, k]), + w: tt([k, n]), + out: tt([m, n]), + dtype: 'bf16', +}); +const weightLoad = (name: string, k: number, n: number): ExpandedOp => ({ + id: id(name), + label: name, + deps: [], + kind: 'weight-load', + out: tt([k, n]), + dtype: 'bf16', + loadFraction: 1, +}); + +test('an op and an independent weight stream overlap fully', () => { + const d = deploymentOn(); + const ops = [gemm('g', [], 4096, 8192, 8192), weightLoad('w', 8192, 8192)]; + const trace: Segment[] = [{ label: 's', ops, repeat: 1 }]; + const cost = dag(d).priceTrace(trace); + const g = naiveOpCost(ops[0], d); + const w = naiveOpCost(ops[1], d); + // the streams overlap: the fullest one sets the time + expect(cost.time).toBeCloseTo(Math.max(g.compute, g.memory + w.memory), 12); + expect(cost.busy).toEqual(naive(d).priceTrace(trace).busy); + expect(sum(cost.parts)).toBeCloseTo(cost.time, 12); +}); + +test('a dependent chain through a collective serializes exactly', () => { + const d = deploymentOn(h100, 2); + const axis = d.mesh.dims[0].name; + const g1 = gemm('g1', [], 64, 8192, 8192); + const ar: ExpandedOp = { + id: id('ar'), + label: 'ar', + deps: [id('g1')], + kind: 'collective', + variant: 'all-reduce', + axes: [axis], + x: tt([64, 8192], [[], []], [axis]), + out: tt([64, 8192]), + dtype: 'bf16', + }; + const g2 = gemm('g2', ['ar'], 64, 8192, 8192); + const trace: Segment[] = [{ label: 's', ops: [g1, ar, g2], repeat: 1 }]; + const cost = dag(d).priceTrace(trace); + const c1 = naiveOpCost(g1, d); + const c2 = naiveOpCost(g2, d); + const cc = naiveOpCost(ar, d); + expect(cost.time).toBeCloseTo(max(c1) + cc.comms + max(c2), 12); + expect(cost.bound).toBe('deps'); + expect(sum(cost.parts)).toBeCloseTo(cost.time, 12); +}); + +test('repeats pipeline: the first iteration pays its weights, the rest do not', () => { + const d = deploymentOn(); + const w = weightLoad('w', 8192, 8192); + const g: ExpandedOp = { ...gemm('g', ['w'], 64, 8192, 8192) }; + const one = dag(d).priceTrace([{ label: 's', ops: [w, g], repeat: 1 }]); + const n = 10; + const many = dag(d).priceTrace([{ label: 's', ops: [w, g], repeat: n }]); + const cw = naiveOpCost(w, d); + const cg = naiveOpCost(g, d); + // one iteration: the load then the gemm, in that order + expect(one.time).toBeCloseTo(cw.memory + max(cg), 12); + // then each later one is bounded by its fullest stream, the memory stream + // that still carries the prefetched weights + const steady = Math.max(cg.compute, cg.memory + cw.memory, max(cg)); + expect(many.time).toBeCloseTo(one.time + (n - 1) * steady, 12); + expect(many.time).toBeLessThan(n * one.time); + + // without prefetch every iteration waits for its weights + const noPrefetch = makeDagScheduleBackend({ prefetchWeights: false })(d); + expect(noPrefetch.priceTrace([{ label: 's', ops: [w, g], repeat: n }]).time).toBeCloseTo( + n * one.time, + 12, + ); +}); + +test('busy and busyPerOp are the naive backend’s, op for op', () => { + const gptoss = MODEL_PRESETS.find((m) => m.name.startsWith('gpt-oss-120b'))!; + const k3 = MODEL_PRESETS.find((m) => m.name.startsWith('Kimi K3'))!; + const h200 = CHIPS_BY_ID['h200-sxm']; + const b300 = CHIPS_BY_ID['b300']; + const cases = [ + evaluateDecodeAtBatch( + { + model: gptoss, + deployment: deploymentOn(h200), + workload: { prefillLen: 4096, generateLen: 1024 }, + }, + 64, + 1, + { costBackend: dag, ignoreKvCapacity: true }, + ), + evaluatePrefill( + { + model: k3, + deployment: deploymentOn(b300, 8), + workload: { prefillLen: 4096, generateLen: 1024 }, + }, + 1, + 'throughput', + { costBackend: dag, ignoreKvCapacity: true }, + ), + ]; + for (const r of cases) { + if (!r.ok) throw new Error(r.diags.map((x) => x.message).join('; ')); + for (const trace of r.perStageTrace) { + const d = cases[0] === r ? deploymentOn(h200) : deploymentOn(b300, 8); + const a = naive(d).priceTrace(trace); + const b = dag(d).priceTrace(trace); + expect(b.busy).toEqual(a.busy); + expect([...b.busyPerOp!.entries()]).toEqual([...a.busyPerOp!.entries()]); + } + } +}); + +test('every preset schedules between its widest stream and its serial sum', () => { + const chip = CHIPS_BY_ID['b300']; + const d = deploymentOn(chip, 8); + const workload = { prefillLen: 2048, generateLen: 512 }; + for (const model of MODEL_PRESETS) { + for (const phase of ['decode', 'prefill'] as const) { + const r = + phase === 'decode' + ? evaluateDecodeAtBatch({ model, deployment: d, workload }, 32, 1, { + costBackend: dag, + ignoreKvCapacity: true, + }) + : evaluatePrefill({ model, deployment: d, workload }, 1, 'throughput', { + costBackend: dag, + ignoreKvCapacity: true, + }); + if (!r.ok) continue; + const c = r.cost; + expect(c.time, `${model.name} ${phase}`).toBeGreaterThanOrEqual(max(c.busy) * (1 - 1e-9)); + expect(c.time, `${model.name} ${phase}`).toBeLessThanOrEqual(sum(c.busy) * (1 + 1e-9)); + expect(sum(c.parts)).toBeCloseTo(c.time, 9); + for (const k of RES) expect(c.hidden[k]).toBeCloseTo(c.busy[k] - c.parts[k], 9); + } + } +}); + +test('gpt-oss-120b decode on one H200 at batch 64 lands under the measured TPOT', () => { + const gptoss = MODEL_PRESETS.find((m) => m.name.startsWith('gpt-oss-120b'))!; + const d = deploymentOn(CHIPS_BY_ID['h200-sxm']); + const input = { model: gptoss, deployment: d, workload: { prefillLen: 4096, generateLen: 1024 } }; + const r = evaluateDecodeAtBatch(input, 64, 1, { costBackend: dag, ignoreKvCapacity: true }); + if (!r.ok) throw new Error('anchor did not evaluate'); + // vLLM on this cell measures ~22.5 ms per token; the serial sum overshoots + // it and a fully hidden memory stream undershoots it, the schedule sits + // between: past the memory floor, under the measurement + expect(r.stepTime).toBeLessThan(22.5e-3); + expect(r.stepTime).toBeGreaterThan(r.cost.busy.memory); + const s = evaluateDecodeAtBatch(input, 64, 1, { costBackend: serial, ignoreKvCapacity: true }); + if (!s.ok) throw new Error('anchor did not evaluate'); + expect(r.stepTime).toBeLessThan(s.stepTime); +}); + +test('makeCostBackend picks the scheduler', () => { + const d = deploymentOn(); + const trace: Segment[] = [ + { label: 's', ops: [gemm('g', [], 4096, 8192, 8192), weightLoad('w', 8192, 8192)], repeat: 1 }, + ]; + const o = { memoryOverlap: 0.9, commsOverlap: 0.65 }; + expect(makeCostBackend(o)(d).priceTrace(trace).time).toBe(naive(d).priceTrace(trace).time); + expect(makeCostBackend({ ...o, scheduler: 'dag' })(d).priceTrace(trace).time).toBe( + dag(d).priceTrace(trace).time, + ); + // the dag backend shares the naive backend's reshard plans + expect(dag(d).priceCollectiveHash).toBe(naive(d).priceCollectiveHash); +}); + +test('on K3 across a GB300 NVL72 the schedule runs slower than the constants, not faster', () => { + const k3 = MODEL_PRESETS.find((m) => m.name.startsWith('Kimi K3'))!; + const chip = CHIPS_BY_ID['gb300-nvl72']; + const workload = { prefillLen: 4096, generateLen: 1024 }; + const best = (backend: typeof dag | typeof naive) => { + let top = 0; + for (const step of searchShardings(k3, chip, { domain: 32, nodes: 1 }, workload, { + costBackend: backend, + phase: { kind: 'decode', policy: { sloTokPerSecPerUser: 20 } }, + })) + if (step.candidate) top = Math.max(top, step.candidate.score); + return top; + }; + const ratio = best(dag) / best(naive); + expect(ratio).toBeGreaterThan(0.7); + expect(ratio).toBeLessThanOrEqual(1); +}, 180_000); diff --git a/tests/state-slots.test.ts b/tests/state-slots.test.ts new file mode 100644 index 0000000..d757f3a --- /dev/null +++ b/tests/state-slots.test.ts @@ -0,0 +1,158 @@ +import { expect, test } from 'vitest'; +import { MODEL_PRESETS } from '../src/core/model/models'; +import { CHIPS_BY_ID, ChipSpec } from '../src/core/hardware/chips'; +import { deployedAxes } from '../src/core/hardware/topology'; +import { Deployment, makeMesh, roleSize } from '../src/core/engine/surface/deploy'; +import { memoryFootprint } from '../src/core/engine/sim/run/memory'; +import { partitionIntoStages } from '../src/core/engine/sim/lowering/stages'; +import { evaluateDecodeAtBatch } from '../src/core/engine/sim/run/decode'; +import { operatingBatch } from '../src/core/engine/optimizer/policy'; +import { searchShardings } from '../src/core/engine/optimizer/search'; +import { makeNaiveOpCostSumBackend } from '../src/core/engine/sim/cost/naiveOpCostSum'; +import { blockStateBytes } from '../src/core/model/block'; +import { DEFAULT_STATE_POOL, StatePoolOptions } from '../src/core/engine/surface/api'; + +// The anchor cell: Kimi K3 (24 MLA + 69 KDA layers) on one 8x B300 node at +// TP=8, ISL 8192 / OSL 1024, fp8 latents, fp32 recurrent state. SGLang on +// this exact cell admits 101 concurrent requests with radix cache +// (extra_buffer, 5 state slots per request) and mem_fraction_static 0.9, +// and 68 with DSPARK speculative decoding on top (draft block + 1 = 8 +// more slots). The simulator's physics floor is one slot per sequence. +const k3 = MODEL_PRESETS.find((m) => m.name.startsWith('Kimi K3'))!; +const b300 = CHIPS_BY_ID['b300']; +const axes = deployedAxes(b300.interconnect, { domain: 8, nodes: 1 }); +const names = axes.map((a) => a.name); +const stages = partitionIntoStages(k3, 1); +const workload = { prefillLen: 8192, generateLen: 1024 }; +const fullLen = workload.prefillLen + workload.generateLen; +const backend = makeNaiveOpCostSumBackend({ memoryOverlap: 0.9, commsOverlap: 0.65 }); + +const tp8 = (chip: ChipSpec = b300, dcp = 1): Deployment => ({ + chip, + mesh: makeMesh(axes, { DPA: [], TP: names, EP: names, ETP: [], PP: [] }), + moeDispatch: 'ring-of-experts', + decodeContextParallel: dcp, +}); +const footprint = (pool: Partial, chip = b300, dcp = 1) => + memoryFootprint(k3, tp8(chip, dcp), stages, fullLen, pool); +// SGLang's mem_fraction_static is a chip-level headroom, not a per-sequence +// cost, so it is modelled as the chip having that much HBM +const atFraction = (f: number): ChipSpec => ({ ...b300, hbmCapacity: f * b300.hbmCapacity }); + +const kdaLayers = 69; +const kdaBlock = k3.blocks[0].pattern[0].block; +const statePerSlotPerChip = (kdaLayers * blockStateBytes(kdaBlock)) / 8; + +test('the state pool defaults to one slot and reproduces the old footprint', () => { + expect(DEFAULT_STATE_POOL).toEqual({ slotsPerSeq: 1, specSlots: 0 }); + const before = memoryFootprint(k3, tp8(), stages, fullLen); + const explicit = footprint({}); + expect(explicit).toEqual(before); + expect(before.stateSlotsPerSeq).toBe(1); + expect(before.kvBytesPerSeqPerChip).toBe( + before.pagedKvBytesPerSeqPerChip + before.stateBytesPerSeqPerChip, + ); + expect(before.stateBytesPerSeqPerChip).toBeCloseTo(statePerSlotPerChip, 0); + expect(before.maxResidentSeqsPerChip).toBe(407); +}); + +test('reserved slots multiply the state part only, by exact byte arithmetic', () => { + const one = footprint({}); + const five = footprint({ slotsPerSeq: 5 }); + expect(five.stateSlotsPerSeq).toBe(5); + expect(five.pagedKvBytesPerSeqPerChip).toBe(one.pagedKvBytesPerSeqPerChip); + expect(five.stateBytesPerSeqPerChip).toBeCloseTo(5 * one.stateBytesPerSeqPerChip, 0); + expect(five.kvBytesPerSeqPerChip).toBeCloseTo( + one.pagedKvBytesPerSeqPerChip + 5 * one.stateBytesPerSeqPerChip, + 0, + ); + expect(five.weightBytesPerChip).toBe(one.weightBytesPerChip); + expect(five.maxResidentSeqsPerChip).toBe( + Math.floor((b300.hbmCapacity - five.weightBytesPerChip) / five.kvBytesPerSeqPerChip), + ); + expect(five.maxResidentSeqsPerChip).toBe(185); + + // spec slots stack on top of the working slots + const spec = footprint({ slotsPerSeq: 5, specSlots: 8 }); + expect(spec.stateSlotsPerSeq).toBe(13); + expect(spec.stateBytesPerSeqPerChip).toBeCloseTo(13 * one.stateBytesPerSeqPerChip, 0); +}); + +test('a narrower state dtype halves the state part and leaves the paged part alone', () => { + const fp32 = footprint({ slotsPerSeq: 5 }); + const bf16 = footprint({ slotsPerSeq: 5, stateDtypeBytes: 2 }); + expect(bf16.stateBytesPerSeqPerChip).toBeCloseTo(fp32.stateBytesPerSeqPerChip / 2, 0); + expect(bf16.pagedKvBytesPerSeqPerChip).toBe(fp32.pagedKvBytesPerSeqPerChip); + expect(bf16.maxResidentSeqsPerChip).toBeGreaterThan(fp32.maxResidentSeqsPerChip); +}); + +test('DCP shards the paged cache and cannot touch the reserved state', () => { + const dcp1 = footprint({ slotsPerSeq: 5 }); + const dcp8 = footprint({ slotsPerSeq: 5 }, b300, 8); + expect(dcp8.pagedKvBytesPerSeqPerChip).toBeCloseTo(dcp1.pagedKvBytesPerSeqPerChip / 8, 0); + expect(dcp8.stateBytesPerSeqPerChip).toBe(dcp1.stateBytesPerSeqPerChip); +}); + +test('five slots plus 0.9 static fraction lands near the measured admission', () => { + const measuredNoSpec = 101; + const measuredDspark = 68; + + // one slot per sequence overshoots the measured ceiling by more than 2x: + // this is the bug the knob exists for + expect(footprint({}, atFraction(0.9)).maxResidentSeqsPerChip).toBeGreaterThan(2 * measuredNoSpec); + + const noSpec = footprint({ slotsPerSeq: 5 }, atFraction(0.9)).maxResidentSeqsPerChip; + expect(noSpec).toBe(117); + expect(Math.abs(noSpec - measuredNoSpec) / measuredNoSpec).toBeLessThan(0.25); + + // DSPARK's extra slots are charged per sequence the same way, which is a + // stricter accounting than the stack's (it shares a pool between the + // working and draft states): direction and a 30% band, not a hit + const spec = footprint({ slotsPerSeq: 5, specSlots: 8 }, atFraction(0.9)).maxResidentSeqsPerChip; + expect(spec).toBeLessThan(noSpec); + expect(Math.abs(spec - measuredDspark) / measuredDspark).toBeLessThan(0.3); +}); + +test('the pool changes capacity only: the priced step is byte-identical', () => { + const input = { model: k3, deployment: tp8(), workload }; + const base = { costBackend: backend, ignoreKvCapacity: true }; + const a = evaluateDecodeAtBatch(input, 64, 1, base); + const b = evaluateDecodeAtBatch(input, 64, 1, { ...base, statePool: { slotsPerSeq: 5 } }); + if (!a.ok || !b.ok) throw new Error('anchor cell did not evaluate'); + expect(b.stepTime).toBe(a.stepTime); + expect(b.traffic).toEqual(a.traffic); + expect(b.cost.busy).toEqual(a.cost.busy); + expect(b.memory.maxResidentSeqsPerChip).toBe(185); + expect(a.memory.maxResidentSeqsPerChip).toBe(407); +}); + +test('the KV gate and the operating batch see the reserved slots', () => { + const input = { model: k3, deployment: tp8(), workload }; + const opts = { costBackend: backend, statePool: { slotsPerSeq: 5 } }; + + expect(evaluateDecodeAtBatch(input, 185, 1, opts).ok).toBe(true); + const over = evaluateDecodeAtBatch(input, 400, 1, opts); + expect(over.ok).toBe(false); + expect(over.diags.map((d) => d.code)).toContain('kv-no-room'); + expect(evaluateDecodeAtBatch(input, 400, 1, { costBackend: backend }).ok).toBe(true); + + expect(operatingBatch(input, { batching: 'max' }, opts)).toBe(185); + expect(operatingBatch(input, { batching: 'max' }, { costBackend: backend })).toBe(407); +}); + +test('a search carries the pool into every candidate it yields', () => { + const gen = searchShardings(k3, b300, { domain: 8, nodes: 1 }, workload, { + costBackend: backend, + statePool: { slotsPerSeq: 5 }, + phase: { kind: 'decode', policy: { batching: 'max' } }, + }); + // the first tuples (TP=1) cannot hold K3's replicated weights on one + // B300, so take the first tuple that priced + let cand; + for (const step of gen) if ((cand = step.candidate)) break; + if (!cand) throw new Error('no feasible tuple on the anchor node'); + if (!('memory' in cand.result)) throw new Error('decode search yielded a prefill result'); + expect(cand.result.memory.stateSlotsPerSeq).toBe(5); + const dpa = roleSize(cand.deployment.mesh, 'DPA'); + expect(cand.batch).toBeLessThanOrEqual(cand.result.memory.maxResidentSeqsPerChip * dpa); +}); diff --git a/tsconfig.json b/tsconfig.json index 7fdbbde..00e40ce 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -12,7 +12,7 @@ "skipLibCheck": true, "isolatedModules": true, "noEmit": true, - "types": ["vite/client"] + "types": ["vite/client", "node"] }, "include": ["src", "tests"] }