Skip to content

Benchmark packed verification prefixes with a fixed proposal count #244

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
voipmonitor wants to merge 1 commit into from
Closed

Benchmark packed verification prefixes with a fixed proposal count #244

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion docs/cli-design.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ the timed graph starts.
| `--batches LIST` | Distinct comma-separated request batches selected from `1,2,4,8`; prefill requires `1` | `1,2,4`; `1` with `--prefill` or `--speculation-study` |
| `--contexts LIST` | Distinct consumed context lengths before each measured graph, multiples of 128; prefill also allows zero | `512,2048,8192,16384`; `0` with `--prefill`; `8192` with `--speculation-study` |
| `--verify-width N` | Target tokens per request, including the pending token; `8,16`, or `8,13,16` with `--speculation-study` | `8` |
| `--verify-prefixes LIST` | C1 packed prefixes measured from shared context, within `--verify-width`; proposal count remains fixed above prefix one. Cannot accompany `--prefill` or `--speculation-study` | unset |
| `--study-widths LIST` | Prepare selected widths from `8,13,16` together and measure each from shared prefilled state; replaces scalar width/depth flags | unset |
| `--speculation-study` | Measure private single-request speculative graphs without native completion publication or export | off |
| `--warmup N` | Positive count of untimed restored rounds per case and seed | `5` |
Expand All @@ -241,7 +242,25 @@ that capacity from the concurrency requested through configuration or flags.
`--no-mtp` applies only with `--prefill`. Mixed prefill and prompt-logprob graphs
are unsupported.

Kimi K3.1 Pro supports verification width eight. Its selected catalog contains
For Kimi K3.1 Pro, compare packed verification prefixes without loading weights again:

```sh
.venv/bin/ignition bench kimi-k3.1-pro --gpus 0-15 \
--mtp-draft-tokens 7 \
--batches 1 --max-concurrency 1 --verify-prefixes 1,4,6,8 --contexts 8192 \
--speculative-token-capacity 8 --prefill-chunk-size 4096 --online-mxfp8 \
--autotune-effort default --output /tmp/kimi31-verification-prefixes.json
```

Each prefix restores the same request, cache and sampling state before replay.
For prefixes four, six and eight, the learned draft generates and samples all
seven proposals, then verifies the selected prefix. The graph
retains logical width eight while physical target rows are four, six or eight.
Prefix one uses the width-one serving graph: it omits proposal generation but
updates draft context, so it is not a target-only model measurement. JSON records
the executed draft-forward size and acceptance separately from replay latency.

Kimi K3.1 Pro supports verification widths through eight. Its selected catalog contains
the requested speculative rounds and the initialization and prefill graphs
needed to establish their contexts. Vision weights retain their configured
residency, while image graphs are omitted from this text benchmark.
Expand Down
95 changes: 62 additions & 33 deletions ignition/bench/runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,10 @@ class BenchmarkCase:
context_tokens: int

def __post_init__(self):
if (self.batch not in (1, 2, 4, 8) or self.width not in (8, 16)
or self.packed_tokens != self.batch * self.width or self.context_tokens <= 0):
raise ValueError("round benchmark requires B=1,2,4,8, W=8,16, P=B*W and positive context")
if (self.batch not in (1, 2, 4, 8) or type(self.width) is not int or not 1 <= self.width <= 16
or not self.batch <= self.packed_tokens <= self.batch * self.width
or self.packed_tokens % self.batch or self.context_tokens <= 0):
raise ValueError("round benchmark requires B=1,2,4,8, W=1 through 16, equal positive prefixes within W and positive context")


@dataclass(frozen=True)
Expand Down Expand Up @@ -173,7 +174,7 @@ def _round(scheduler, case):
if len(matches) != 1:
raise ValueError(f"prepared catalog must contain exactly one complete round for {case}")
shape, graph = matches[0]
lanes = [Lane(request, request.consumed, case.width) for request in scheduler.active]
lanes = [Lane(request, request.consumed, case.packed_tokens // case.batch) for request in scheduler.active]
return scheduler._commit(Step("speculative", graph, case.batch, case.width, lanes, 0,
packed_tokens=case.packed_tokens, draft_tokens=shape[4]))

Expand Down Expand Up @@ -206,7 +207,13 @@ def _check_request_state(model, step, records):
raise RuntimeError(f"request {request.ident} device state disagrees with graph completion")


def _check_metadata(model, step, case):
def _proposal_counters(forward, step):
owner = forward.owner
return [(int(owner.generation[lane.request.slot].item()),
int(owner.offset[lane.request.slot].item())) for lane in step.lanes]


def _check_metadata(model, step, case, *, proposal_before=None):
if isinstance(case, PrefillCase):
metadata = model.benchmark_prefills[step.graph]
if (metadata.live_counts.cpu().tolist() != [case.chunk_tokens, 1]
Expand All @@ -216,8 +223,10 @@ def _check_metadata(model, step, case):
raise RuntimeError("prefill metadata does not describe the exact benchmark chunk")
return
forward = model.benchmark_rounds[step.graph]
for role, metadata, width in (("target", forward.metadata, case.width),
("draft", forward.draft_metadata, forward.draft_inputs.shape[1])):
stages = [("target", forward.metadata, case.packed_tokens // case.batch)]
if forward.draft is not None:
stages.append(("draft", forward.draft_metadata, forward.draft_inputs.shape[1]))
for role, metadata, width in stages:
expected_starts = [case.context_tokens] * case.batch
expected_lengths = [width] * case.batch
expected_offsets = [row * width for row in range(case.batch + 1)]
Expand All @@ -228,11 +237,14 @@ def _check_metadata(model, step, case):
raise RuntimeError(f"{role} metadata does not describe the complete benchmark round")
owner = forward.owner
verification = forward.operations[2]
if forward.draft is None and _proposal_counters(forward, step) != proposal_before:
raise RuntimeError("a round without proposals changed draft RNG state")
for lane in step.lanes:
slot, generation = lane.request.slot, lane.request.generation
if (int(owner.generation[slot].item()) != generation
or int(owner.offset[slot].item()) != 4 * step.draft_tokens
or int(owner.rng.generation[slot].item()) != generation
if forward.draft is not None and (int(owner.generation[slot].item()) != generation
or int(owner.offset[slot].item()) != 4 * step.draft_tokens):
raise RuntimeError("proposal RNG did not consume its declared draw reservation")
if (int(owner.rng.generation[slot].item()) != generation
or int(owner.rng.target_offset[slot].item()) != verification.draw_reservation):
raise RuntimeError("speculative sampling streams did not consume their declared draw reservations")

Expand Down Expand Up @@ -314,34 +326,51 @@ def result(status, failure=None):
return measurement_result(cases, measurements, warmup=warmup, repetitions=repetitions,
seeds=seeds, status=status, failure=failure)

groups = {}
for index, case in enumerate(cases):
key = (case.batch, case.context_tokens, getattr(case, "chunk_tokens", 0))
groups.setdefault(key, []).append(index)
generation = int(model.benchmark_requests.generation.max().item())
try:
for case, samples in zip(cases, measurements):
for indices in groups.values():
# Every width starts from the same real prefix, request budget and RNG state.
setup = max((cases[index] for index in indices), key=lambda case: getattr(case, "width", 0))
for seed in seeds:
times = []
samples.append(dict(seed=seed, times_ms=times))
generation += 1
baseline = _prefill(model, runner, layout, case, tokenizer, seed, generation)
baseline = _prefill(model, runner, layout, setup, tokenizer, seed, generation)
runner.snapshot()
for iteration in range(-1, warmup + repetitions):
scheduler = deepcopy(baseline, {id(layout): layout})
runner.restore()
step = _prefill_step(scheduler, case) if isinstance(case, PrefillCase) else _round(scheduler, case)
elapsed = runner.run(_message(scheduler, step))
if iteration >= warmup:
times.append(elapsed / 1e6)
records = _complete(model, runner, scheduler, step, elapsed)
_check_request_state(model, step, records)
if isinstance(case, PrefillCase) and any(
lane.request.consumed != case.context_tokens + case.chunk_tokens
or len(lane.request.generated) != 1 for lane in step.lanes):
raise RuntimeError("prefill replay did not consume the complete chunk and sample its pending token")
if iteration == -1:
_check_metadata(model, step, case)
if not isinstance(case, PrefillCase):
samples[-1]["acceptance"] = acceptance_observation(model.benchmark_rounds[step.graph])
if progress is not None:
progress(result("running"))
for index in indices:
case, samples = cases[index], measurements[index]
times = []
samples.append(dict(seed=seed, times_ms=times))
for iteration in range(-1, warmup + repetitions):
scheduler = deepcopy(baseline, {id(layout): layout})
runner.restore()
step = _prefill_step(scheduler, case) if isinstance(case, PrefillCase) else _round(scheduler, case)
proposal_before = None
if iteration == -1 and not isinstance(case, PrefillCase):
proposal_before = _proposal_counters(model.benchmark_rounds[step.graph], step)
elapsed = runner.run(_message(scheduler, step))
if iteration >= warmup:
times.append(elapsed / 1e6)
records = _complete(model, runner, scheduler, step, elapsed)
_check_request_state(model, step, records)
if isinstance(case, PrefillCase) and any(
lane.request.consumed != case.context_tokens + case.chunk_tokens
or len(lane.request.generated) != 1 for lane in step.lanes):
raise RuntimeError("prefill replay did not consume the complete chunk and sample its pending token")
if iteration == -1:
_check_metadata(model, step, case, proposal_before=proposal_before)
if not isinstance(case, PrefillCase):
forward = model.benchmark_rounds[step.graph]
samples[-1]["acceptance"] = acceptance_observation(forward)
samples[-1]["draft_forward_tokens_per_request"] = (
forward.draft_inputs.shape[1] if forward.draft is not None else 0)
samples[-1]["draft_context_updated"] = forward.context is not None
samples[-1]["proposal_tokens_per_request"] = step.draft_tokens
samples[-1]["target_tokens_per_request"] = case.packed_tokens // case.batch
if progress is not None:
progress(result("running"))
runner.restore()
runner.release_snapshot()
except BaseException as error:
Expand Down
36 changes: 28 additions & 8 deletions ignition/cli/bench.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ def arguments(parser):
help="preceding consumed tokens, multiples of 128 (default: 512,2048,8192,16384; --prefill: 0)")
options.add_argument("--verify-width", type=int, metavar="N",
help="target tokens per request including pending token: 8 or 16; 8,13,16 with --speculation-study")
options.add_argument("--verify-prefixes", type=_integers, metavar="LIST",
help="C1 packed verification prefixes, e.g. 1,4,6,8; retains the full proposal count for prefixes above one")
options.add_argument("--warmup", type=int, default=5, metavar="N",
help="positive untimed restored replays per case and seed (default: 5)")
options.add_argument("--repetitions", type=int, default=20, metavar="N",
Expand All @@ -64,17 +66,23 @@ class Settings:
prefill: bool = False
speculation_study: bool = False
study_widths: tuple[int, ...] = ()
verification_prefixes: tuple[int, ...] = ()

@classmethod
def from_args(cls, args):
study = getattr(args, "speculation_study", False)
batches = args.batches if args.batches is not None else ((1,) if args.prefill or study else (1, 2, 4))
batches = args.batches if args.batches is not None else ((1,) if args.prefill or study or args.verify_prefixes else (1, 2, 4))
contexts = args.contexts if args.contexts is not None else ((0,) if args.prefill else (8192,) if study else (512, 2048, 8192, 16384))
if args.prefill and args.verify_width is not None:
raise UsageError("--verify-width does not apply to --prefill")
if args.prefill and batches != (1,):
raise UsageError("--prefill requires --batches 1")
study_widths = getattr(args, "study_widths", None)
verification_prefixes = getattr(args, "verify_prefixes", None)
if verification_prefixes is not None:
if args.prefill or study or batches != (1,):
raise UsageError("--verify-prefixes requires --batches 1 without --prefill or --speculation-study")
verification_prefixes = tuple(sorted(verification_prefixes))
if study_widths is not None:
if not study or any(w not in (8, 13, 16) for w in study_widths):
raise UsageError("--study-widths requires --speculation-study and widths from 8,13,16")
Expand All @@ -90,6 +98,8 @@ def from_args(cls, args):
raise UsageError("--speculation-study requires --batches 1 and --verify-width 8,13,16 without --prefill")
if not args.prefill and not study and width not in (8, 16):
raise UsageError("ignition bench requires --verify-width 8 or 16 and matching --mtp-draft-tokens 7 or 15")
if verification_prefixes and (verification_prefixes[0] < 1 or verification_prefixes[-1] > width):
raise UsageError("--verify-prefixes must be positive and no larger than --verify-width")
if any(value not in (1, 2, 4, 8) for value in batches):
raise UsageError("--batches must select request counts from 1,2,4,8")
if any(value < (0 if args.prefill else 128) or value % 128 for value in contexts):
Expand All @@ -107,7 +117,8 @@ def from_args(cls, args):
if args.official_source is not None:
raise UsageError("--official-source applies to core publishing; ignition bench exports no core")
return cls(tuple(sorted(batches)), tuple(sorted(contexts)), 8 if args.prefill else width,
args.warmup, args.repetitions, args.seeds, args.prefill, study, study_widths or ())
args.warmup, args.repetitions, args.seeds, args.prefill, study, study_widths or (),
verification_prefixes or ())


def table(result):
Expand Down Expand Up @@ -161,7 +172,7 @@ def run(args):
plan = replace(plan, benchmark_selection=BenchmarkSelection(
batches=settings.batches, contexts=settings.contexts, width=settings.width,
prefill_chunk=plan.prefill_chunk if settings.prefill else 0, speculation_study=settings.speculation_study,
study_widths=settings.study_widths))
study_widths=settings.study_widths, verification_prefixes=settings.verification_prefixes))
print(json.dumps({"plan": plan.to_dict()}), file=sys.stderr, flush=True)
progress = _progress(args.verbose)
hub_cache = configuration.directory("hub_cache")
Expand Down Expand Up @@ -221,6 +232,7 @@ def _measure(model, plan, configuration, settings, progress, *, compile_forwards
details["configuration"] = {"model": model.model_id, "checkpoint_directory": str(snapshot),
"repository": model.checkpoint.repository, "revision": model.checkpoint.revision,
"compiled_forwards": compile_forwards, "torch_version": torch.__version__,
"verification_prefixes": settings.verification_prefixes, "logical_verification_width": settings.width,
"requested_max_concurrency": plan.concurrency, "prepared_request_capacity": max(settings.batches),
"preparation": {key: preparation[key] for key in (
"tensor_parallel", "prefill_chunk", "stream_overlap", "context_parallel",
Expand Down Expand Up @@ -256,11 +268,19 @@ def record(result):
tokenizer_asset = root.server_config["assets"]["tokenizer"]
tokenizer = (Tokenizer.from_str(tokenizer_asset["text"]) if "text" in tokenizer_asset
else Tokenizer.from_file(tokenizer_asset["path"]))
scenarios = (tuple(PrefillCase(chunk_tokens=plan.prefill_chunk, context_tokens=context)
for context in settings.contexts) if settings.prefill else
tuple((StudyCase if settings.speculation_study else BenchmarkCase)(batch=batch, width=width, packed_tokens=batch * width,
context_tokens=context) for batch in settings.batches for context in settings.contexts
for width in settings.study_widths or (settings.width,)))
if settings.prefill:
scenarios = tuple(PrefillCase(chunk_tokens=plan.prefill_chunk, context_tokens=context)
for context in settings.contexts)
elif settings.speculation_study:
scenarios = tuple(StudyCase(batch=batch, width=width, packed_tokens=batch * width,
context_tokens=context)
for batch in settings.batches for context in settings.contexts
for width in settings.study_widths or (settings.width,))
else:
scenarios = tuple(BenchmarkCase(batch=batch, width=width, packed_tokens=packed,
context_tokens=context)
for batch, width, packed in plan.benchmark_selection.speculative_round_shapes
for context in settings.contexts)
execute = run_study if settings.speculation_study else run_benchmark
result = execute(prepared, tokenizer, scenarios=scenarios, warmup=settings.warmup,
repetitions=settings.repetitions, seeds=settings.seeds, progress=record)
Expand Down
4 changes: 3 additions & 1 deletion ignition/planning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,9 @@ def __post_init__(self):
raise UsageError("benchmark graph preparation requires Qwen3.8-27B on one GPU or Kimi K3.1 Pro on sixteen GPUs")
if selection.speculation_study and self.recipe.name != "qwen38_dense":
raise UsageError("speculation studies require Qwen3.8-27B DFlash2")
if not selection.prefill_chunk and (self.speculation is None or (8 if selection.speculation_study else selection.width) not in self.speculation.widths):
required_widths = (8,) if selection.speculation_study else selection.selected_verification_widths
if not selection.prefill_chunk and (self.speculation is None
or not set(required_widths) <= set(self.speculation.widths)):
raise UsageError("benchmark rounds require the selected speculative width")
if selection.max_batch > self.concurrency:
raise UsageError("benchmark batches exceed --max-concurrency")
Expand Down
Loading
Loading