diff --git a/docs/cli-design.md b/docs/cli-design.md index ed989997..34e7f63b 100644 --- a/docs/cli-design.md +++ b/docs/cli-design.md @@ -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` | @@ -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. diff --git a/ignition/bench/runner.py b/ignition/bench/runner.py index 365ba39d..80fba666 100644 --- a/ignition/bench/runner.py +++ b/ignition/bench/runner.py @@ -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) @@ -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])) @@ -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] @@ -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)] @@ -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") @@ -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: diff --git a/ignition/cli/bench.py b/ignition/cli/bench.py index b5006c7a..a2b0ac57 100644 --- a/ignition/cli/bench.py +++ b/ignition/cli/bench.py @@ -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", @@ -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") @@ -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): @@ -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): @@ -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") @@ -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", @@ -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) diff --git a/ignition/planning.py b/ignition/planning.py index 2132337a..d0569f07 100644 --- a/ignition/planning.py +++ b/ignition/planning.py @@ -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") diff --git a/ignition/preparation/selection.py b/ignition/preparation/selection.py index 3768e8e6..77dc09fe 100644 --- a/ignition/preparation/selection.py +++ b/ignition/preparation/selection.py @@ -14,6 +14,7 @@ class BenchmarkSelection: prefill_chunk: int = 0 speculation_study: bool = False study_widths: tuple[int, ...] = () + verification_prefixes: tuple[int, ...] = () def __post_init__(self): for name, values in (("batches", self.batches), ("contexts", self.contexts)): @@ -40,6 +41,17 @@ def __post_init__(self): or any(type(w) is not int or w not in (8, 13, 16) for w in self.study_widths) or self.width != max(self.study_widths)): raise ValueError("study widths must be increasing distinct widths in 8,13,16 with width equal to their maximum") + if self.verification_prefixes: + if (self.speculation_study or self.prefill_chunk or self.batches != (1,) + or not isinstance(self.verification_prefixes, tuple) + or tuple(sorted(set(self.verification_prefixes))) != self.verification_prefixes + or any(type(w) is not int or not 1 <= w <= self.width for w in self.verification_prefixes)): + raise ValueError("verification prefixes require batch one and distinct increasing prefixes within the prepared width") + + @property + def selected_verification_widths(self): + return tuple(sorted({1 if prefix == 1 else self.width + for prefix in self.verification_prefixes or (self.width,)})) @property def selected_study_widths(self): @@ -57,7 +69,8 @@ def max_context(self): def speculative_round_shapes(self): if self.prefill_chunk: return () - return tuple((batch, self.width, batch * self.width) for batch in self.batches) + return tuple((batch, 1 if prefix == 1 else self.width, batch * prefix) + for batch in self.batches for prefix in self.verification_prefixes or (self.width,)) @property def required_context(self): @@ -77,4 +90,5 @@ def physical_pages(self, page_tokens): def to_dict(self): return dict(batches=self.batches, contexts=self.contexts, width=self.width, prefill_chunk=self.prefill_chunk, speculation_study=self.speculation_study, study_widths=self.selected_study_widths, + verification_prefixes=self.verification_prefixes, speculative_round_shapes=self.speculative_round_shapes) diff --git a/ignition/recipes/kimi3/bootstrap.py b/ignition/recipes/kimi3/bootstrap.py index a692b639..5db59d30 100644 --- a/ignition/recipes/kimi3/bootstrap.py +++ b/ignition/recipes/kimi3/bootstrap.py @@ -140,7 +140,7 @@ def __post_init__(self): if (not self.drafting and not selection.prefill_chunk) or self.prompt_logprobs or selection.speculation_study: raise ValueError("Kimi benchmark preparation requires complete unscored text graphs") if selection.max_batch > self.max_active or (not selection.prefill_chunk - and selection.width not in self.speculative_widths): + and not set(selection.selected_verification_widths) <= set(self.speculative_widths)): raise ValueError("benchmark shapes exceed the configured request or width capacities") if selection.required_context > self.context: raise ValueError("benchmark workload and output allowance must fit the context") diff --git a/ignition/recipes/qwen38_dense/bootstrap.py b/ignition/recipes/qwen38_dense/bootstrap.py index c1eda47e..375c11b5 100644 --- a/ignition/recipes/qwen38_dense/bootstrap.py +++ b/ignition/recipes/qwen38_dense/bootstrap.py @@ -87,7 +87,8 @@ def __post_init__(self): if selection is not None: if (not self.mtp and not selection.prefill_chunk) or self.prompt_logprobs: raise ValueError("benchmark preparation requires speculative unscored text graphs") - if selection.max_batch > self.max_active or (not selection.prefill_chunk and not selection.speculation_study and selection.width not in self.speculative_widths): + if selection.max_batch > self.max_active or (not selection.prefill_chunk and not selection.speculation_study + and not set(selection.selected_verification_widths) <= set(self.speculative_widths)): raise ValueError("benchmark shapes exceed the configured request or width capacities") if selection.required_context > self.context: raise ValueError("benchmark workload and output allowance must fit the context") diff --git a/tests/bench/test_runner.py b/tests/bench/test_runner.py index 297a39df..13bbd16e 100644 --- a/tests/bench/test_runner.py +++ b/tests/bench/test_runner.py @@ -123,7 +123,7 @@ def test_missing_or_invalid_measurements_are_not_latency_estimates(times): summarize_case(BenchmarkCase(1, 8, 8, 512), [dict(seed=0, times_ms=times)]) -@pytest.mark.parametrize("arguments", [(3, 8, 24, 512), (1, 4, 4, 512), (2, 8, 8, 512), (1, 8, 8, 0)]) +@pytest.mark.parametrize("arguments", [(3, 8, 24, 512), (1, 0, 0, 512), (2, 8, 7, 512), (1, 8, 8, 0)]) def test_case_rejects_shapes_outside_complete_round_matrix(arguments): with pytest.raises(ValueError): BenchmarkCase(*arguments) @@ -258,3 +258,73 @@ def test_acceptance_observation_separates_verification_from_committed_clipping() sample = dict(seed=0, times_ms=[1.] * 20, acceptance=result) summary = summarize_case(BenchmarkCase(1, 8, 8, 8192), [sample]) assert summary["samples"][0]["acceptance"]["observations_per_request"] == 1 + + +def test_width_one_preserves_proposal_rng_and_advances_target_rng(): + from types import SimpleNamespace as NS + import torch + from ignition.bench.runner import _check_metadata + + metadata = NS(live_counts=torch.tensor([1, 1]), starts=torch.tensor([512]), + lengths=torch.tensor([1]), query_start_loc=torch.tensor([0, 1])) + owner = NS(generation=torch.tensor([2]), offset=torch.tensor([28]), + rng=NS(generation=torch.tensor([3]), target_offset=torch.tensor([4]))) + forward = NS(metadata=metadata, draft=None, owner=owner, + operations=(None, (), NS(draw_reservation=4), None)) + model = NS(benchmark_rounds={"round1": forward}) + step = NS(graph="round1", lanes=[NS(request=NS(slot=0, generation=3))], draft_tokens=0) + case = BenchmarkCase(1, 1, 1, 512) + _check_metadata(model, step, case, proposal_before=[(2, 28)]) + owner.offset.add_(4) + with pytest.raises(RuntimeError, match="without proposals changed"): + _check_metadata(model, step, case, proposal_before=[(2, 28)]) + owner.offset.sub_(4) + owner.rng.target_offset.zero_() + with pytest.raises(RuntimeError, match="draw reservations"): + _check_metadata(model, step, case, proposal_before=[(2, 28)]) + + +@pytest.mark.parametrize("prefix", [4, 6, 8]) +def test_packed_prefix_keeps_seven_proposals_and_packs_only_verified_rows(prefix): + from ignition.bench.runner import _round + from ignition.direct.scheduler import Lane, Scheduler + from tests.direct.test_scheduler import layout, records, request, uploads + + catalog = layout(capacity=1, max_active=1, widths=(1, 8), graphs={ + ("init", 1, 1): "init1", ("prefill", 1, 512): "prefill512", + ("speculative", 1, 8, prefix, 7): "round"}) + scheduler = Scheduler(catalog, prefix_cache=False) + member = request(1, [10] * 512, limit=64) + scheduler.submit(member) + step = scheduler.plan() + scheduler.complete(step, records(scheduler, step, {member: (1, 0, [], 0)})) + step = scheduler._step("prefill", 512, [Lane(member, 0, 512)]) + scheduler.complete(step, records(scheduler, step, {member: (1, 512, [42], 0)})) + step = _round(scheduler, BenchmarkCase(1, 8, prefix, 512)) + assert (step.width, step.draft_tokens, step.packed_tokens) == (8, 7, prefix) + packed = uploads(scheduler, step) + assert packed["query_lengths"] == {0: [prefix]} + assert step.lanes[0].tokens == prefix + + +@pytest.mark.parametrize("prefix", [4, 6, 8]) +def test_packed_metadata_retains_full_proposal_draw_reservation(prefix): + from types import SimpleNamespace as NS + import torch + from ignition.bench.runner import _check_metadata + + def metadata(rows): + return NS(live_counts=torch.tensor([rows, 1]), starts=torch.tensor([512]), + lengths=torch.tensor([rows]), query_start_loc=torch.tensor([0, rows])) + owner = NS(generation=torch.tensor([3]), offset=torch.tensor([28]), + rng=NS(generation=torch.tensor([3]), target_offset=torch.tensor([12]))) + forward = NS(metadata=metadata(prefix), draft=object(), draft_metadata=metadata(8), + draft_inputs=torch.empty((1, 8)), owner=owner, + operations=(None, (), NS(draw_reservation=12), None)) + model = NS(benchmark_rounds={"round": forward}) + step = NS(graph="round", lanes=[NS(request=NS(slot=0, generation=3))], draft_tokens=7) + _check_metadata(model, step, BenchmarkCase(1, 8, prefix, 512)) + if prefix < 8: + owner.offset.fill_(4 * (prefix - 1)) + with pytest.raises(RuntimeError, match="proposal RNG"): + _check_metadata(model, step, BenchmarkCase(1, 8, prefix, 512)) diff --git a/tests/cli/test_bench.py b/tests/cli/test_bench.py index 0dda48e5..25f0e3d4 100644 --- a/tests/cli/test_bench.py +++ b/tests/cli/test_bench.py @@ -58,6 +58,26 @@ def test_kimi_eight_request_round_with_autotuning(): assert args.max_concurrency == 8 and args.speculative_token_capacity == 64 +def test_verification_matrix_keeps_the_full_draft_block(): + value = settings("--verify-prefixes", "8,1,6,4", "--mtp-draft-tokens", "7") + assert value.verification_prefixes == (1, 4, 6, 8) + assert value.width == 8 + assert value.batches == (1,) + assert settings("--verify-prefixes", "4,6").width == 8 + + +@pytest.mark.parametrize("flags", [ + ("--verify-prefixes", "0,8"), ("--verify-prefixes", "4,9"), + ("--verify-prefixes", "1,8", "--batches", "2"), + ("--verify-prefixes", "1,8", "--prefill"), + ("--verify-prefixes", "1,8", "--speculation-study"), + ("--verify-prefixes", "1,8", "--mtp-draft-tokens", "3"), +]) +def test_invalid_verification_matrix(flags): + with pytest.raises(UsageError): + settings(*flags) + + def test_kimi_rejects_private_dflash_studies_before_preparation(tmp_path, capsys): config = tmp_path / "config.toml" config.write_text("") diff --git a/tests/preparation/test_benchmark_selection.py b/tests/preparation/test_benchmark_selection.py index 9a5611a1..3395028a 100644 --- a/tests/preparation/test_benchmark_selection.py +++ b/tests/preparation/test_benchmark_selection.py @@ -105,6 +105,19 @@ def test_invalid_selection(values): BenchmarkSelection(**values) +@pytest.mark.parametrize("values", [ + dict(verification_prefixes=(4, 1, 8)), dict(verification_prefixes=(1, 1, 8)), + dict(verification_prefixes=(1, 4), batches=(1, 2)), dict(verification_prefixes=(0, 8)), + dict(verification_prefixes=(1, 9)), dict(verification_prefixes=(True, 8)), + dict(verification_prefixes=(1, 8), prefill_chunk=4096, batches=(1,)), + dict(verification_prefixes=(1, 8), speculation_study=True, batches=(1,)), +]) +def test_invalid_serving_prefix_matrix(values): + values = {"batches": (1,), **values} + with pytest.raises(ValueError, match="verification prefixes"): + BenchmarkSelection(**values) + + def test_plan_retains_explicit_concurrency_and_token_budget(): plan = replace(_plan(), benchmark_selection=BenchmarkSelection()) assert plan.concurrency == 32 diff --git a/tests/recipes/test_kimi_benchmark.py b/tests/recipes/test_kimi_benchmark.py index c5f90aa4..a2192c07 100644 --- a/tests/recipes/test_kimi_benchmark.py +++ b/tests/recipes/test_kimi_benchmark.py @@ -32,6 +32,20 @@ def test_selected_kimi_round_preserves_serving_configuration(): assert len(ordinary.speculative_round_shapes) == 68 +def test_packed_prefix_matrix_preserves_k7_owner_and_backing(): + selection = BenchmarkSelection(batches=(1,), contexts=(512, 8192), + verification_prefixes=(1, 4, 6, 8)) + config = PreparationConfig(max_active=4, draft_checkpoint=Path("dflash2"), + speculative_token_capacity=32, benchmark_selection=selection) + assert config.draft_tokens == 7 + assert config.completion_token_capacity == 8 + assert config.speculative_round_shapes == ( + (1, 1, 1), (1, 8, 4), (1, 8, 6), (1, 8, 8)) + assert selection.physical_pages(2048) == 6 + with pytest.raises(ValueError, match="width capacities"): + replace(config, speculative_widths=(1,)) + + @pytest.mark.parametrize("changes,match", [ (dict(draft_checkpoint=None), "complete unscored"), (dict(prompt_logprobs=True), "complete unscored"),