From efad7b50532d03c718bd6b8e4c20f662cf258e31 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 02:13:35 -0700 Subject: [PATCH 01/19] perf: defer DDP progress readback until training completes --- code/labs/train_distributed/optimized_ddp.py | 24 +++++--- .../optimized_ddp_multigpu.py | 24 +++++--- .../training_utils/deferred_metrics.py | 48 +++++++++++++++ .../test_benchmark_hygiene_regressions.py | 7 +++ code/tests/test_deferred_training_progress.py | 61 +++++++++++++++++++ 5 files changed, 148 insertions(+), 16 deletions(-) create mode 100644 code/labs/train_distributed/training_utils/deferred_metrics.py create mode 100644 code/tests/test_deferred_training_progress.py diff --git a/code/labs/train_distributed/optimized_ddp.py b/code/labs/train_distributed/optimized_ddp.py index 2d3997464..7469fc1a6 100644 --- a/code/labs/train_distributed/optimized_ddp.py +++ b/code/labs/train_distributed/optimized_ddp.py @@ -26,6 +26,7 @@ gradient_sync_context, validate_gradient_accumulation, ) +from labs.train_distributed.training_utils.deferred_metrics import DeferredTrainingProgress from labs.train_distributed.training_utils.torchrun_harness import TorchrunScriptBenchmark from labs.train_distributed.training_utils.utils import ( build_dataloader, @@ -109,8 +110,12 @@ def main(): num_steps = min(args.steps, len(dataloader)) accumulation_plan = build_gradient_accumulation_plan(num_steps, args.grad_accum) total_tokens = 0 + progress = ( + DeferredTrainingProgress(num_steps=num_steps, interval=10, device=device) + if is_main and num_steps > 0 + else None + ) start_time = perf_counter() - loss_value_buffer = torch.empty(1, dtype=torch.float64, device=device) completed_steps = 0 final_batch = None @@ -140,19 +145,22 @@ def main(): total_tokens += batch["input_ids"].numel() if step % 10 == 0 and is_main: - loss_value_buffer[0].copy_(loss.detach()) - loss_value = float(loss_value_buffer.detach().cpu()[0]) - print( - f"[optimized-ddp] step {step}/{num_steps} " - f"loss={loss_value:.4f} " - f"tokens/step={batch['input_ids'].numel():,}" - ) + if progress is None: + raise RuntimeError("Training progress buffer was not initialized") + progress.record(step=step, loss=loss, tokens=batch["input_ids"].numel()) torch.cuda.synchronize(device) total_time = perf_counter() - start_time if completed_steps <= 0: raise RuntimeError("DDP training completed no optimization steps") if is_main: + if progress is None: + raise RuntimeError("Training progress buffer was not initialized") + for sample in progress.read(): + print( + f"[optimized-ddp] step {sample.step}/{num_steps} " + f"loss={sample.loss:.4f} tokens/step={sample.tokens:,}" + ) toks_sec = total_tokens / total_time if total_time > 0 else 0.0 effective_bs = args.batch_size * args.grad_accum * world_size print( diff --git a/code/labs/train_distributed/optimized_ddp_multigpu.py b/code/labs/train_distributed/optimized_ddp_multigpu.py index 1767b5ea9..1dc14add2 100644 --- a/code/labs/train_distributed/optimized_ddp_multigpu.py +++ b/code/labs/train_distributed/optimized_ddp_multigpu.py @@ -34,6 +34,7 @@ gradient_sync_context, validate_gradient_accumulation, ) +from labs.train_distributed.training_utils.deferred_metrics import DeferredTrainingProgress from labs.train_distributed.training_utils.torchrun_harness import TorchrunScriptBenchmark from labs.train_distributed.training_utils.utils import ( build_dataloader, @@ -121,8 +122,12 @@ def main(): num_steps = min(args.steps, len(dataloader)) accumulation_plan = build_gradient_accumulation_plan(num_steps, args.grad_accum) total_tokens = 0 + progress = ( + DeferredTrainingProgress(num_steps=num_steps, interval=10, device=device) + if is_main and num_steps > 0 + else None + ) start_time = perf_counter() - loss_value_buffer = torch.empty(1, dtype=torch.float64, device=device) completed_steps = 0 final_batch = None @@ -156,19 +161,22 @@ def main(): total_tokens += batch["input_ids"].numel() if step % 10 == 0 and is_main: - loss_value_buffer[0].copy_(loss.detach()) - loss_value = float(loss_value_buffer.detach().cpu()[0]) - print( - f"[optimized-ddp] step {step}/{num_steps} " - f"loss={loss_value:.4f} " - f"tokens/step={batch['input_ids'].numel():,}" - ) + if progress is None: + raise RuntimeError("Training progress buffer was not initialized") + progress.record(step=step, loss=loss, tokens=batch["input_ids"].numel()) torch.cuda.synchronize(device) total_time = perf_counter() - start_time if completed_steps <= 0: raise RuntimeError("DDP training completed no optimization steps") if is_main: + if progress is None: + raise RuntimeError("Training progress buffer was not initialized") + for sample in progress.read(): + print( + f"[optimized-ddp] step {sample.step}/{num_steps} " + f"loss={sample.loss:.4f} tokens/step={sample.tokens:,}" + ) toks_sec = total_tokens / total_time if total_time > 0 else 0.0 effective_bs = args.batch_size * args.grad_accum * world_size print( diff --git a/code/labs/train_distributed/training_utils/deferred_metrics.py b/code/labs/train_distributed/training_utils/deferred_metrics.py new file mode 100644 index 000000000..4a2daf7e7 --- /dev/null +++ b/code/labs/train_distributed/training_utils/deferred_metrics.py @@ -0,0 +1,48 @@ +"""Retain sampled training losses without synchronizing every progress update.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class TrainingProgress: + step: int + loss: float + tokens: int + + +class DeferredTrainingProgress: + """Copy detached losses on device; transfer recorded samples once at the end. + + Construct before the measured loop and read after its final synchronization. + Report end-to-end process time separately from the training-loop timer. + """ + + def __init__(self, *, num_steps: int, interval: int, device: torch.device) -> None: + if num_steps < 1 or interval < 1: + raise ValueError("num_steps and interval must be positive") + self.interval = interval + self.num_steps = num_steps + self._losses = torch.empty( + (num_steps + interval - 1) // interval, device=device, dtype=torch.float64 + ) + self._samples: list[tuple[int, int]] = [] + + def record(self, *, step: int, loss: torch.Tensor, tokens: int) -> None: + if step != len(self._samples) * self.interval or step >= self.num_steps: + raise ValueError("Progress samples must follow the declared step interval") + if loss.numel() != 1 or loss.device != self._losses.device: + raise ValueError("Loss must be a scalar on the progress buffer device") + with torch.no_grad(): + self._losses[len(self._samples)].copy_(loss.detach().reshape(())) + self._samples.append((step, tokens)) + + def read(self) -> list[TrainingProgress]: + values = self._losses[: len(self._samples)].detach().cpu().tolist() + return [ + TrainingProgress(step=step, loss=value, tokens=tokens) + for (step, tokens), value in zip(self._samples, values, strict=True) + ] diff --git a/code/tests/test_benchmark_hygiene_regressions.py b/code/tests/test_benchmark_hygiene_regressions.py index 2058a2500..88814e5b1 100644 --- a/code/tests/test_benchmark_hygiene_regressions.py +++ b/code/tests/test_benchmark_hygiene_regressions.py @@ -15318,6 +15318,13 @@ def test_train_distributed_optimized_wrappers_log_detached_loss_values() -> None assert "loss.item()" not in source assert "float(loss.detach())" not in source + if relative in ("optimized_ddp.py", "optimized_ddp_multigpu.py"): + assert "DeferredTrainingProgress(num_steps=num_steps, interval=10, device=device)" in source + assert "progress.record(step=step, loss=loss, tokens=batch[\"input_ids\"].numel())" in source + assert "for sample in progress.read():" in source + assert "loss={sample.loss:.4f}" in source + assert source.index("total_time =") < source.index("for sample in progress.read():") + continue assert "loss_value_buffer = torch.empty(1, dtype=torch.float64" in source assert "loss_value_buffer[0].copy_(loss.detach())" in source assert "loss_value = float(loss_value_buffer.detach().cpu()[0])" in source diff --git a/code/tests/test_deferred_training_progress.py b/code/tests/test_deferred_training_progress.py new file mode 100644 index 000000000..039fc587f --- /dev/null +++ b/code/tests/test_deferred_training_progress.py @@ -0,0 +1,61 @@ +"""Actual sampled-loss storage, autograd independence, and CUDA transfer checks.""" + +import pytest +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from labs.train_distributed.training_utils.deferred_metrics import DeferredTrainingProgress + + +def test_retains_each_sample_after_source_changes_without_autograd_history(): + progress = DeferredTrainingProgress(num_steps=21, interval=10, device=torch.device("cpu")) + value = torch.tensor(1.25, requires_grad=True) + for step, expected in [(0, 1.25), (10, 2.5), (20, 3.75)]: + with torch.no_grad(): + value.fill_(expected) + progress.record(step=step, loss=value, tokens=128 + step) + with torch.no_grad(): + value.fill_(100) + rows = progress.read() + assert [(row.step, row.loss, row.tokens) for row in rows] == [ + (0, 1.25, 128), (10, 2.5, 138), (20, 3.75, 148) + ] + assert value.grad is None + + +def test_partial_training_does_not_read_unwritten_buffer_entries(): + progress = DeferredTrainingProgress(num_steps=100, interval=10, device=torch.device("cpu")) + assert progress.read() == [] + progress.record(step=0, loss=torch.tensor(0.5), tokens=64) + assert [row.loss for row in progress.read()] == [0.5] + with pytest.raises(ValueError, match="declared step interval"): + progress.record(step=20, loss=torch.tensor(0.6), tokens=64) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Actual CUDA device required") +def test_cuda_records_never_transfer_loss_to_host_until_read(): + class Transfers(TorchDispatchMode): + def __init__(self): + super().__init__() + self.host_copies = 0 + self.scalar_reads = 0 + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + options = kwargs or {} + if func is torch.ops.aten._to_copy.default and options.get("device") == torch.device("cpu"): + self.host_copies += 1 + if func is torch.ops.aten._local_scalar_dense.default: + self.scalar_reads += 1 + return func(*args, **options) + + device = torch.device("cuda", 0) + progress = DeferredTrainingProgress(num_steps=20, interval=10, device=device) + losses = torch.tensor([1.25, 2.5], device=device) + trace = Transfers() + with trace: + progress.record(step=0, loss=losses[0], tokens=128) + progress.record(step=10, loss=losses[1], tokens=128) + assert trace.host_copies == 0 and trace.scalar_reads == 0 + rows = progress.read() + assert trace.host_copies == 1 and trace.scalar_reads == 0 + assert [row.loss for row in rows] == [1.25, 2.5] From 30446ffe9606f096997f09bb4d8205cdcb09a1cc Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 02:15:54 -0700 Subject: [PATCH 02/19] feat: add explicit execution placement and destination write audits --- code/core/harness/execution_audit.py | 655 ++++++++++++++++++ ...stributed_protection_receipt_test_utils.py | 132 ++++ code/tests/test_anti_cheat_edge_cases.py | 53 +- code/tests/test_anti_cheat_protections.py | 14 +- code/tests/test_execution_audit_guards.py | 269 +++++++ 5 files changed, 1108 insertions(+), 15 deletions(-) create mode 100644 code/core/harness/execution_audit.py create mode 100644 code/tests/distributed_protection_receipt_test_utils.py create mode 100644 code/tests/test_execution_audit_guards.py diff --git a/code/core/harness/execution_audit.py b/code/core/harness/execution_audit.py new file mode 100644 index 000000000..5ab8077b4 --- /dev/null +++ b/code/core/harness/execution_audit.py @@ -0,0 +1,655 @@ +"""Explicit, out-of-timing execution audits for benchmark callbacks. + +This module intentionally does not participate in normal benchmark timing. A +dispatcher mode observes tensor operands and results for one audited invocation, +and an optional destination guard poisons declared floating-point output buffers +before that invocation to prove that every logical element was overwritten. + +The placement audit covers PyTorch dispatcher operations executed on the current +thread. It cannot see arbitrary work inside custom extensions, background +threads, host callbacks, or external processes. Destination poisoning proves +write coverage only for the exact buffers supplied by the caller; it is not a +general uninitialized-memory or allocation-provenance detector. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import sys +from collections import Counter +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +PLACEMENT_SCOPE = ( + "PyTorch dispatcher-visible tensor operations on the audited invocation's current thread" +) +WRITE_COVERAGE_SCOPE = ( + "exact declared contiguous floating-point or complex destinations poisoned before the " + "audited invocation" +) + + +@dataclass(frozen=True, eq=False) +class HostTensorAllowance: + """Allow one exact CPU tensor only for explicitly named dispatcher operations.""" + + label: str + tensor: torch.Tensor + operations: tuple[str, ...] + + def __post_init__(self) -> None: + if not self.label.strip(): + raise ValueError("host tensor allowance label must be non-empty") + if not isinstance(self.tensor, torch.Tensor): + raise TypeError("host tensor allowance must reference a torch.Tensor") + if self.tensor.device.type != "cpu": + raise ValueError("host tensor allowances apply only to exact CPU tensor identities") + if not self.operations or any(not operation.strip() for operation in self.operations): + raise ValueError("host tensor allowance operations must be non-empty strings") + + +@dataclass(frozen=True) +class TensorExtentEvidence: + """Observed tensor identity-independent extent and placement evidence.""" + + path: str + device: str + dtype: str + shape: tuple[int, ...] + stride: tuple[int, ...] + numel: int + nbytes: int + matches_expected_device: bool + allowed_host_tensor: bool + + def to_dict(self) -> dict[str, Any]: + return { + "path": self.path, + "device": self.device, + "dtype": self.dtype, + "shape": list(self.shape), + "stride": list(self.stride), + "numel": self.numel, + "nbytes": self.nbytes, + "matches_expected_device": self.matches_expected_device, + "allowed_host_tensor": self.allowed_host_tensor, + } + + +@dataclass(frozen=True) +class OperationEvidence: + """Bounded evidence for one dispatcher operation.""" + + operator: str + tensors: tuple[TensorExtentEvidence, ...] + mismatched_paths: tuple[str, ...] + + def to_dict(self) -> dict[str, Any]: + return { + "operator": self.operator, + "tensors": [tensor.to_dict() for tensor in self.tensors], + "mismatched_paths": list(self.mismatched_paths), + } + + +@dataclass(frozen=True) +class OperationPlacementResult: + """Result of checking every dispatcher-visible tensor in one invocation.""" + + expected_device: str + operations_seen: int + operator_counts: tuple[tuple[str, int], ...] + operation_evidence: tuple[OperationEvidence, ...] + operation_evidence_truncated: bool + violations_seen: int + violation_evidence: tuple[OperationEvidence, ...] + violation_evidence_truncated: bool + + @property + def passed(self) -> bool: + return self.violations_seen == 0 + + def to_dict(self) -> dict[str, Any]: + return { + "passed": self.passed, + "scope": PLACEMENT_SCOPE, + "expected_device": self.expected_device, + "operations_seen": self.operations_seen, + "operator_counts": dict(self.operator_counts), + "operation_evidence": [item.to_dict() for item in self.operation_evidence], + "operation_evidence_truncated": self.operation_evidence_truncated, + "violations_seen": self.violations_seen, + "violation_evidence": [item.to_dict() for item in self.violation_evidence], + "violation_evidence_truncated": self.violation_evidence_truncated, + } + + +def _iter_tensor_paths(value: Any, path: str): + if isinstance(value, torch.Tensor): + yield path, value + return + if isinstance(value, Mapping): + for key, item in value.items(): + yield from _iter_tensor_paths(item, f"{path}[{key!r}]") + return + if isinstance(value, list | tuple): + for index, item in enumerate(value): + yield from _iter_tensor_paths(item, f"{path}[{index}]") + + +class TensorOperationPlacementAudit(TorchDispatchMode): + """Observe tensor devices for PyTorch operations on the current thread.""" + + def __init__( + self, + expected_device: str | torch.device, + *, + allowed_host_tensors: Sequence[HostTensorAllowance] = (), + evidence_limit: int = 64, + ) -> None: + super().__init__() + if isinstance(evidence_limit, bool) or evidence_limit <= 0: + raise ValueError("evidence_limit must be a positive integer") + self.expected_device = torch.device(expected_device) + self.evidence_limit = evidence_limit + self._allowances = tuple(allowed_host_tensors) + self._allowed_operations_by_identity: dict[int, set[str]] = {} + for allowance in self._allowances: + if not isinstance(allowance, HostTensorAllowance): + raise TypeError("allowed_host_tensors must contain HostTensorAllowance instances") + self._allowed_operations_by_identity.setdefault(id(allowance.tensor), set()).update( + allowance.operations + ) + self._operations_seen = 0 + self._operator_counts: Counter[str] = Counter() + self._operation_evidence: list[OperationEvidence] = [] + self._violations_seen = 0 + self._violation_evidence: list[OperationEvidence] = [] + + def _matches_expected_device(self, actual: torch.device) -> bool: + if actual.type != self.expected_device.type: + return False + if self.expected_device.index is None: + return True + return actual.index == self.expected_device.index + + def _is_allowed_host_tensor(self, tensor: torch.Tensor, operator: str) -> bool: + if tensor.device.type != "cpu": + return False + return operator in self._allowed_operations_by_identity.get(id(tensor), set()) + + def _tensor_evidence( + self, + *, + path: str, + tensor: torch.Tensor, + operator: str, + ) -> TensorExtentEvidence: + matches = self._matches_expected_device(tensor.device) + allowed = not matches and self._is_allowed_host_tensor(tensor, operator) + return TensorExtentEvidence( + path=path, + device=str(tensor.device), + dtype=str(tensor.dtype).removeprefix("torch."), + shape=tuple(tensor.shape), + stride=tuple(tensor.stride()), + numel=tensor.numel(), + nbytes=tensor.numel() * tensor.element_size(), + matches_expected_device=matches, + allowed_host_tensor=allowed, + ) + + def __torch_dispatch__( + self, + func: Any, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Any: + del types + actual_kwargs = kwargs or {} + result = func(*args, **actual_kwargs) + operator = str(func) + self._operations_seen += 1 + self._operator_counts[operator] += 1 + + observed = [ + *self._collect_tensor_evidence(args, "args", operator), + *self._collect_tensor_evidence(actual_kwargs, "kwargs", operator), + *self._collect_tensor_evidence(result, "output", operator), + ] + mismatched_paths = tuple( + item.path + for item in observed + if not item.matches_expected_device and not item.allowed_host_tensor + ) + evidence = OperationEvidence( + operator=operator, + tensors=tuple(observed), + mismatched_paths=mismatched_paths, + ) + if len(self._operation_evidence) < self.evidence_limit: + self._operation_evidence.append(evidence) + if mismatched_paths: + self._violations_seen += 1 + if len(self._violation_evidence) < self.evidence_limit: + self._violation_evidence.append(evidence) + return result + + def _collect_tensor_evidence( + self, + value: Any, + path: str, + operator: str, + ) -> list[TensorExtentEvidence]: + return [ + self._tensor_evidence(path=item_path, tensor=tensor, operator=operator) + for item_path, tensor in _iter_tensor_paths(value, path) + ] + + def result(self) -> OperationPlacementResult: + """Return immutable bounded evidence collected so far.""" + + return OperationPlacementResult( + expected_device=str(self.expected_device), + operations_seen=self._operations_seen, + operator_counts=tuple(sorted(self._operator_counts.items())), + operation_evidence=tuple(self._operation_evidence), + operation_evidence_truncated=self._operations_seen > len(self._operation_evidence), + violations_seen=self._violations_seen, + violation_evidence=tuple(self._violation_evidence), + violation_evidence_truncated=self._violations_seen > len(self._violation_evidence), + ) + + +@dataclass(frozen=True) +class DestinationWriteCoverageEvidence: + """Write-coverage result for one exact declared destination tensor.""" + + name: str + device: str + dtype: str + shape: tuple[int, ...] + stride: tuple[int, ...] + storage_offset: int + numel: int + nbytes: int + unwritten_elements: int + first_unwritten_flat_indices: tuple[int, ...] + + @property + def passed(self) -> bool: + return self.unwritten_elements == 0 + + def to_dict(self) -> dict[str, Any]: + return { + "passed": self.passed, + "scope": WRITE_COVERAGE_SCOPE, + "name": self.name, + "device": self.device, + "dtype": self.dtype, + "shape": list(self.shape), + "stride": list(self.stride), + "storage_offset": self.storage_offset, + "numel": self.numel, + "nbytes": self.nbytes, + "unwritten_elements": self.unwritten_elements, + "first_unwritten_flat_indices": list(self.first_unwritten_flat_indices), + } + + +class DestinationWriteCoverageGuard: + """Poison exact declared output buffers and check that all elements change.""" + + def __init__(self, destinations: Mapping[str, torch.Tensor]) -> None: + self._destinations = dict(destinations) + if not self._destinations: + raise ValueError("at least one destination tensor is required") + seen_storage: dict[tuple[str, int], str] = {} + for name, tensor in self._destinations.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError("destination names must be non-empty strings") + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"destination {name!r} must be a torch.Tensor") + if tensor.device.type == "meta": + raise ValueError(f"destination {name!r} cannot be a meta tensor") + if tensor.numel() == 0: + raise ValueError(f"destination {name!r} is empty; write coverage is unobservable") + if not (tensor.is_floating_point() or tensor.is_complex()): + raise TypeError( + f"destination {name!r} must be floating-point or complex for NaN poisoning" + ) + if not tensor.is_contiguous(): + raise ValueError( + f"destination {name!r} must be contiguous for exact logical write coverage" + ) + storage_key = (str(tensor.device), tensor.untyped_storage().data_ptr()) + previous_name = seen_storage.get(storage_key) + if previous_name is not None: + raise ValueError( + f"destinations {previous_name!r} and {name!r} share storage; " + "audit them separately" + ) + seen_storage[storage_key] = name + self._poisoned = False + + def _synchronize_cuda_destinations(self) -> None: + devices = { + tensor.device for tensor in self._destinations.values() if tensor.device.type == "cuda" + } + for device in sorted(devices, key=str): + torch.cuda.synchronize(device) + + def poison(self) -> None: + """Fill destinations with NaNs before the audited invocation.""" + + if self._poisoned: + raise RuntimeError("destination write-coverage guard has already been poisoned") + with torch.no_grad(): + for tensor in self._destinations.values(): + tensor.fill_(float("nan")) + self._synchronize_cuda_destinations() + self._poisoned = True + + def inspect(self) -> tuple[DestinationWriteCoverageEvidence, ...]: + """Synchronize CUDA and report any poison that survived the invocation.""" + + if not self._poisoned: + raise RuntimeError("poison() must be called before inspect()") + self._synchronize_cuda_destinations() + results: list[DestinationWriteCoverageEvidence] = [] + for name, tensor in self._destinations.items(): + unwritten_mask = torch.isnan(tensor) + unwritten_elements = int(unwritten_mask.count_nonzero().item()) + first_indices: tuple[int, ...] = () + if unwritten_elements: + first_indices = tuple( + int(index) + for index in unwritten_mask.flatten().nonzero()[:8].flatten().tolist() + ) + results.append( + DestinationWriteCoverageEvidence( + name=name, + device=str(tensor.device), + dtype=str(tensor.dtype).removeprefix("torch."), + shape=tuple(tensor.shape), + stride=tuple(tensor.stride()), + storage_offset=int(tensor.storage_offset()), + numel=tensor.numel(), + nbytes=tensor.numel() * tensor.element_size(), + unwritten_elements=unwritten_elements, + first_unwritten_flat_indices=first_indices, + ) + ) + return tuple(results) + + +@dataclass(frozen=True) +class ExecutionAuditResult: + """Combined placement and declared-destination evidence.""" + + placement: OperationPlacementResult + destinations: tuple[DestinationWriteCoverageEvidence, ...] + destination_identity_errors: tuple[str, ...] = () + + @property + def passed(self) -> bool: + return ( + self.placement.passed + and all(destination.passed for destination in self.destinations) + and not self.destination_identity_errors + ) + + def raise_for_failure(self) -> None: + if self.passed: + return + diagnostics: list[str] = [] + if not self.placement.passed: + first = self.placement.violation_evidence[0] + diagnostics.append( + f"{self.placement.violations_seen} operation placement violation(s); " + f"first={first.operator} paths={list(first.mismatched_paths)}" + ) + for destination in self.destinations: + if not destination.passed: + diagnostics.append( + f"destination {destination.name!r} retained poison in " + f"{destination.unwritten_elements}/{destination.numel} elements" + ) + diagnostics.extend(self.destination_identity_errors) + raise RuntimeError("EXECUTION AUDIT FAILED: " + " | ".join(diagnostics)) + + def to_dict(self) -> dict[str, Any]: + return { + "schema": "aisp.execution-audit.v1", + "passed": self.passed, + "placement": self.placement.to_dict(), + "destinations": [destination.to_dict() for destination in self.destinations], + "destination_identity_errors": list(self.destination_identity_errors), + "limits": [ + "Placement covers dispatcher-visible tensor operands/results on the current thread.", + "Custom extension internals, background work, host callbacks, and external processes " + "are outside this audit.", + "Destination checks prove overwrite coverage only for exact declared buffers.", + "Destination checks do not establish general allocation or uninitialized-memory " + "provenance.", + ], + } + + +def audit_callable_once( + callback: Callable[[], Any], + *, + expected_device: str | torch.device, + destinations: Mapping[str, torch.Tensor] | None = None, + allowed_host_tensors: Sequence[HostTensorAllowance] = (), + evidence_limit: int = 64, +) -> ExecutionAuditResult: + """Audit one callback invocation outside benchmark timing.""" + + if not callable(callback): + raise TypeError("callback must be callable") + write_guard = DestinationWriteCoverageGuard(destinations) if destinations else None + if write_guard is not None: + write_guard.poison() + placement_audit = TensorOperationPlacementAudit( + expected_device, + allowed_host_tensors=allowed_host_tensors, + evidence_limit=evidence_limit, + ) + with placement_audit: + callback() + destination_results = write_guard.inspect() if write_guard is not None else () + return ExecutionAuditResult( + placement=placement_audit.result(), + destinations=destination_results, + ) + + +def _resolve_attribute(instance: Any, attribute_path: str) -> Any: + if not attribute_path or any(part == "" for part in attribute_path.split(".")): + raise ValueError("attribute paths must be non-empty dotted names") + value = instance + for part in attribute_path.split("."): + if part.startswith("__"): + raise ValueError("dunder attribute paths are not supported") + value = getattr(value, part) + return value + + +def _load_target_module(path: Path) -> Any: + spec = importlib.util.spec_from_file_location("_aisp_execution_audit_target", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load benchmark module from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Run one fresh benchmark invocation under explicit execution audits." + ) + parser.add_argument("benchmark_path", type=Path) + parser.add_argument( + "--factory", + default="get_benchmark", + help="No-argument factory function or class name in the benchmark module.", + ) + parser.add_argument("--expected-device", required=True) + parser.add_argument( + "--destination", + action="append", + default=[], + metavar="ATTRIBUTE", + help="Exact preallocated tensor attribute to poison and verify; may be repeated.", + ) + parser.add_argument( + "--allow-host-tensor", + action="append", + default=[], + metavar="ATTRIBUTE=OPERATOR", + help=( + "Allow one exact CPU tensor attribute for one exact dispatcher operator; " + "repeat for additional operators." + ), + ) + parser.add_argument( + "--target-arg", + action="append", + default=[], + help="Argument exposed to the target module as sys.argv; may be repeated.", + ) + parser.add_argument("--evidence-limit", type=int, default=64) + return parser + + +def _host_allowances(instance: Any, specs: Sequence[str]) -> tuple[HostTensorAllowance, ...]: + operations_by_attribute: dict[str, list[str]] = {} + for spec in specs: + attribute, separator, operator = spec.partition("=") + if not separator or not attribute.strip() or not operator.strip(): + raise ValueError("--allow-host-tensor requires ATTRIBUTE=OPERATOR") + operations_by_attribute.setdefault(attribute, []).append(operator) + return tuple( + HostTensorAllowance( + label=attribute, + tensor=_resolve_attribute(instance, attribute), + operations=tuple(operations), + ) + for attribute, operations in operations_by_attribute.items() + ) + + +def _audit_fresh_benchmark(args: argparse.Namespace) -> tuple[ExecutionAuditResult, dict[str, Any]]: + benchmark_path = args.benchmark_path.expanduser().resolve() + if not benchmark_path.is_file(): + raise FileNotFoundError(f"benchmark module does not exist: {benchmark_path}") + + original_argv = sys.argv + benchmark = None + teardown_error: Exception | None = None + try: + sys.argv = [str(benchmark_path), *args.target_arg] + module = _load_target_module(benchmark_path) + factory = getattr(module, args.factory) + if not callable(factory): + raise TypeError(f"target symbol {args.factory!r} is not callable") + benchmark = factory() + for method_name in ("setup", "benchmark_fn", "teardown"): + if not callable(getattr(benchmark, method_name, None)): + raise TypeError(f"fresh benchmark must implement callable {method_name}()") + benchmark.setup() + + destinations: dict[str, torch.Tensor] = {} + for attribute in args.destination: + if attribute in destinations: + raise ValueError(f"duplicate destination attribute: {attribute}") + destinations[attribute] = _resolve_attribute(benchmark, attribute) + original_destination_identities = { + attribute: id(tensor) for attribute, tensor in destinations.items() + } + result = audit_callable_once( + benchmark.benchmark_fn, + expected_device=args.expected_device, + destinations=destinations, + allowed_host_tensors=_host_allowances(benchmark, args.allow_host_tensor), + evidence_limit=args.evidence_limit, + ) + identity_errors = tuple( + f"destination attribute {attribute!r} no longer references the declared tensor" + for attribute, original_identity in original_destination_identities.items() + if id(_resolve_attribute(benchmark, attribute)) != original_identity + ) + result = replace(result, destination_identity_errors=identity_errors) + metadata = { + "benchmark_path": str(benchmark_path), + "factory": args.factory, + "process_id": os.getpid(), + "fresh_instance": True, + "normal_timing_lifecycle_modified": False, + } + return result, metadata + finally: + if benchmark is not None and callable(getattr(benchmark, "teardown", None)): + try: + benchmark.teardown() + except Exception as error: # preserve teardown evidence after the primary audit + teardown_error = error + sys.argv = original_argv + if teardown_error is not None: + raise RuntimeError( + f"fresh benchmark teardown failed: {type(teardown_error).__name__}: " + f"{teardown_error}" + ) from teardown_error + + +def main(argv: Sequence[str] | None = None) -> int: + """CLI entry point for a fresh, standalone audited invocation.""" + + args = _build_parser().parse_args(argv) + try: + result, metadata = _audit_fresh_benchmark(args) + except Exception as error: + print( + json.dumps( + { + "schema": "aisp.execution-audit.v1", + "passed": False, + "error": f"{type(error).__name__}: {error}", + }, + sort_keys=True, + ) + ) + return 1 + payload = result.to_dict() + payload.update(metadata) + print(json.dumps(payload, sort_keys=True)) + return 0 if result.passed else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = [ + "DestinationWriteCoverageEvidence", + "DestinationWriteCoverageGuard", + "ExecutionAuditResult", + "HostTensorAllowance", + "OperationEvidence", + "OperationPlacementResult", + "TensorExtentEvidence", + "TensorOperationPlacementAudit", + "audit_callable_once", + "main", +] diff --git a/code/tests/distributed_protection_receipt_test_utils.py b/code/tests/distributed_protection_receipt_test_utils.py new file mode 100644 index 000000000..32903dc0d --- /dev/null +++ b/code/tests/distributed_protection_receipt_test_utils.py @@ -0,0 +1,132 @@ +"""Shared real contract controls for retained distributed anti-cheat test IDs.""" + +from __future__ import annotations + +from dataclasses import replace + +from core.benchmark.distributed_work_contract import ( + BARRIER_BEFORE_TIMED_CLOSE, + DECLARED_ALGORITHM_EVIDENCE, + WAIT_FOR_ASYNC_BEFORE_TIMED_CLOSE, + DistributedRankWorkReceipt, + validate_distributed_work_receipts, +) +from core.benchmark.verification import DistributedTopology, compare_topologies + + +def _topology(**overrides: object) -> DistributedTopology: + values = { + "world_size": 2, + "ranks": [0, 1], + "shards": 2, + "per_rank_batch_size": 2, + "collective_type": "all_reduce", + "collective_algorithm": "ring", + "gradient_bucket_bytes": 4096, + "barrier_policy": BARRIER_BEFORE_TIMED_CLOSE, + "async_completion_policy": WAIT_FOR_ASYNC_BEFORE_TIMED_CLOSE, + } + values.update(overrides) + return DistributedTopology(**values) + + +def _receipt(rank: int, **overrides: object) -> DistributedRankWorkReceipt: + values = { + "rank": rank, + "world_size": 2, + "backend": "gloo", + "collective_type": "all_reduce", + "declared_collective_algorithm": "ring", + "gradient_bucket_bytes": 4096, + "barrier_policy": BARRIER_BEFORE_TIMED_CLOSE, + "async_completion_policy": WAIT_FOR_ASYNC_BEFORE_TIMED_CLOSE, + "timed_region_start_ns": 100, + "collective_launch_ns": (110,), + "collective_completion_ns": (130,), + "barrier_entry_ns": 140, + "barrier_completion_ns": 150, + "timed_region_close_ns": 160, + } + values.update(overrides) + return DistributedRankWorkReceipt(**values) + + +def _assert_clean_receipts() -> tuple[DistributedRankWorkReceipt, ...]: + receipts = (_receipt(0), _receipt(1)) + validation = validate_distributed_work_receipts( + _topology(), + receipts, + expected_backend="gloo", + ) + assert validation.passed, validation.errors + assert validation.collective_algorithm_evidence == DECLARED_ALGORITHM_EVIDENCE + return receipts + + +def assert_collective_algorithm_declaration_controls() -> None: + """Check declared ring/tree parity without claiming runtime algorithm inspection.""" + + baseline = _topology() + assert compare_topologies(baseline, replace(baseline)) == (True, None) + passed, differences = compare_topologies( + baseline, + replace(baseline, collective_algorithm="tree"), + ) + assert not passed + assert differences == "Collective algorithm mismatch: ring vs tree" + + receipts = list(_assert_clean_receipts()) + receipts[1] = replace(receipts[1], declared_collective_algorithm="tree") + validation = validate_distributed_work_receipts(_topology(), receipts) + assert not validation.passed + assert any("declared_collective_algorithm mismatch" in error for error in validation.errors) + + +def assert_gradient_bucket_declaration_controls() -> None: + """Check bucket-byte parity in topology declarations and registered receipts.""" + + baseline = _topology() + assert compare_topologies(baseline, replace(baseline)) == (True, None) + passed, differences = compare_topologies( + baseline, + replace(baseline, gradient_bucket_bytes=8192), + ) + assert not passed + assert differences == "Gradient bucket bytes mismatch: 4096 vs 8192" + + receipts = list(_assert_clean_receipts()) + receipts[1] = replace(receipts[1], gradient_bucket_bytes=8192) + validation = validate_distributed_work_receipts(_topology(), receipts) + assert not validation.passed + assert any("gradient_bucket_bytes mismatch" in error for error in validation.errors) + + +def assert_barrier_completion_receipt_controls() -> None: + """Check that a registered final barrier completes before timed-region close.""" + + receipts = list(_assert_clean_receipts()) + receipts[1] = replace(receipts[1], barrier_completion_ns=170) + validation = validate_distributed_work_receipts(_topology(), receipts) + assert not validation.passed + assert any( + "barrier must enter and complete before timed-region close" in error + for error in validation.errors + ) + + +def assert_async_completion_receipt_controls() -> None: + """Check that every registered async collective completes before the barrier.""" + + receipts = list(_assert_clean_receipts()) + receipts[1] = replace(receipts[1], collective_completion_ns=()) + validation = validate_distributed_work_receipts(_topology(), receipts) + assert not validation.passed + assert any("asynchronous collectives incomplete" in error for error in validation.errors) + + +__all__ = [ + "assert_async_completion_receipt_controls", + "assert_barrier_completion_receipt_controls", + "assert_collective_algorithm_declaration_controls", + "assert_gradient_bucket_declaration_controls", +] diff --git a/code/tests/test_anti_cheat_edge_cases.py b/code/tests/test_anti_cheat_edge_cases.py index 40d159d61..2313b0d90 100644 --- a/code/tests/test_anti_cheat_edge_cases.py +++ b/code/tests/test_anti_cheat_edge_cases.py @@ -16,6 +16,13 @@ from core.benchmark.verification import InputSignature, PrecisionFlags, ToleranceSpec from core.benchmark.verify_runner import VerifyConfig +from core.harness.execution_audit import audit_callable_once +from tests.distributed_protection_receipt_test_utils import ( + assert_async_completion_receipt_controls, + assert_barrier_completion_receipt_controls, + assert_collective_algorithm_declaration_controls, + assert_gradient_bucket_declaration_controls, +) from tests.evaluation_contract_test_utils import assert_evaluation_contract_controls from tests.protection_test_utils import ( TensorWork, assert_comparison_controls, assert_compile_cache_reset, @@ -307,12 +314,34 @@ def test_train_test_overlap_single_sample(self, tmp_path): class TestLocationEdgeCases: def test_cpu_spillover_single_op_on_cpu(self): - 'Requirement remains open; this retained test ID is not passing coverage.' - pytest.skip('Missing production protection: no per-operation CPU spillover detector is implemented; wall-time measurement alone does not identify execution placement') + """A real CPU tensor op is visible when the audited invocation expects CUDA.""" + value = torch.arange(12, dtype=torch.float32).reshape(3, 4) + result = audit_callable_once(lambda: value.square(), expected_device="cuda") + assert not result.passed + assert result.placement.operations_seen == 1 + violation = result.placement.violation_evidence[0] + assert violation.operator == "aten.pow.Tensor_Scalar" + assert {tensor.device for tensor in violation.tensors} == {"cpu"} + assert {tensor.shape for tensor in violation.tensors} == {(3, 4)} + assert {tensor.numel for tensor in violation.tensors} == {12} def test_cpu_spillover_data_dependent_branch(self): - 'Requirement remains open; this retained test ID is not passing coverage.' - pytest.skip('Missing production protection: no per-operation CPU spillover detector is implemented; wall-time measurement alone does not identify execution placement') + """Placement evidence follows the branch that actually executes.""" + cpu_value = torch.arange(8, dtype=torch.float32) + + def execute_branch(use_cpu: bool) -> None: + if use_cpu: + torch.relu(cpu_value) + + clean = audit_callable_once(lambda: execute_branch(False), expected_device="cuda") + assert clean.passed + assert clean.placement.operations_seen == 0 + + violation = audit_callable_once(lambda: execute_branch(True), expected_device="cuda") + assert not violation.passed + evidence = violation.placement.violation_evidence[0] + assert evidence.operator == "aten.relu.default" + assert any(tensor.shape == (8,) and tensor.device == "cpu" for tensor in evidence.tensors) def test_setup_precomputation_cached_result(self, runner): assert check_jitter(runner, 'real') == (True, None) @@ -533,20 +562,20 @@ def test_collective_shortcircuit_single_element(self, runner): assert_comparison_controls(runner, expected, torch.tensor([1.0])) def test_topology_mismatch_ring_vs_tree(self): - 'Requirement remains open; this retained test ID is not passing coverage.' - pytest.skip('Missing production protection: DistributedTopology does not encode ring versus tree algorithm, so it cannot enforce this claimed policy') + """Declared algorithm parity is enforced; runtime selection still needs profiling.""" + assert_collective_algorithm_declaration_controls() def test_barrier_timing_straggler(self): - 'Requirement remains open; this retained test ID is not passing coverage.' - pytest.skip('Missing production protection: no rank barrier timing detector is implemented') + """Registered receipts reject a final barrier completing after timed close.""" + assert_barrier_completion_receipt_controls() def test_gradient_bucketing_different_sizes(self): - 'Requirement remains open; this retained test ID is not passing coverage.' - pytest.skip('Missing production protection: no gradient bucket-size parity field or detector is implemented') + """Declared and observed receipt bucket bytes must match.""" + assert_gradient_bucket_declaration_controls() def test_async_gradient_overlap(self): - 'Requirement remains open; this retained test ID is not passing coverage.' - pytest.skip('Missing production protection: no async gradient completion timing detector is implemented') + """Registered async work must complete before the final barrier.""" + assert_async_completion_receipt_controls() def test_pipeline_bubble_microbatch_count(self): 'Declared per-rank workload comparison, not timing or detecting pipeline bubbles.' diff --git a/code/tests/test_anti_cheat_protections.py b/code/tests/test_anti_cheat_protections.py index bc58b624a..4a13fdd68 100644 --- a/code/tests/test_anti_cheat_protections.py +++ b/code/tests/test_anti_cheat_protections.py @@ -19,6 +19,11 @@ import pytest import torch +from tests.distributed_protection_receipt_test_utils import ( + assert_async_completion_receipt_controls, + assert_barrier_completion_receipt_controls, + assert_gradient_bucket_declaration_controls, +) from tests.evaluation_contract_test_utils import assert_evaluation_contract_controls from tests.protection_test_utils import ( @@ -1224,13 +1229,16 @@ def test_collective_short_circuit_detection(self): assert result.error_message == "RANK SKIPPING: Missing outputs from ranks [1]" def test_barrier_timing_protection(self): - pytest.skip('Missing production protection: no rank-barrier timing detector') + """Registered receipts require final-barrier completion before timed close.""" + assert_barrier_completion_receipt_controls() def test_gradient_bucketing_mismatch_detection(self): - pytest.skip('Missing production protection: no gradient bucket-size parity field or detector') + """Declared and observed receipt bucket bytes must match.""" + assert_gradient_bucket_declaration_controls() def test_async_gradient_timing(self): - pytest.skip('Missing production protection: no asynchronous gradient completion timing detector') + """Registered async work must complete before the final barrier.""" + assert_async_completion_receipt_controls() @requires_cuda def test_pipeline_bubble_tracking(self): diff --git a/code/tests/test_execution_audit_guards.py b/code/tests/test_execution_audit_guards.py new file mode 100644 index 000000000..bd584b979 --- /dev/null +++ b/code/tests/test_execution_audit_guards.py @@ -0,0 +1,269 @@ +"""Real CPU and capability-gated CUDA controls for explicit execution audits.""" + +from __future__ import annotations + +import json +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +import torch + +from core.harness.execution_audit import ( + DestinationWriteCoverageGuard, + HostTensorAllowance, + audit_callable_once, +) + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="real CUDA execution-audit integration requires a device", +) + + +def test_cpu_operation_audit_retains_positive_and_negative_extent_evidence() -> None: + value = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + clean = audit_callable_once(lambda: value.square(), expected_device="cpu") + assert clean.passed + assert clean.placement.operations_seen == 1 + clean_record = clean.placement.operation_evidence[0] + assert clean_record.operator == "aten.pow.Tensor_Scalar" + assert clean_record.mismatched_paths == () + assert {(item.device, item.shape, item.numel) for item in clean_record.tensors} == { + ("cpu", (3, 4), 12) + } + + violation = audit_callable_once(lambda: value.square(), expected_device="cuda") + assert not violation.passed + assert violation.placement.violations_seen == 1 + violation_record = violation.placement.violation_evidence[0] + assert violation_record.operator == "aten.pow.Tensor_Scalar" + assert violation_record.mismatched_paths == ("args[0]", "output") + assert {(item.device, item.shape, item.numel) for item in violation_record.tensors} == { + ("cpu", (3, 4), 12) + } + with pytest.raises(RuntimeError, match="EXECUTION AUDIT FAILED"): + violation.raise_for_failure() + + +def test_host_tensor_allowance_requires_exact_identity_and_operator_scope() -> None: + declared_scalar = torch.tensor(7.0) + other_scalar = torch.tensor(8.0) + allowance = HostTensorAllowance( + label="declared_scalar", + tensor=declared_scalar, + operations=("aten._local_scalar_dense.default",), + ) + + clean = audit_callable_once( + declared_scalar.item, + expected_device="cuda", + allowed_host_tensors=(allowance,), + ) + assert clean.passed + assert clean.placement.operations_seen == 1 + tensor_evidence = clean.placement.operation_evidence[0].tensors + assert len(tensor_evidence) == 1 + assert tensor_evidence[0].allowed_host_tensor + + wrong_identity = audit_callable_once( + other_scalar.item, + expected_device="cuda", + allowed_host_tensors=(allowance,), + ) + assert not wrong_identity.passed + assert not wrong_identity.placement.violation_evidence[0].tensors[0].allowed_host_tensor + + wrong_operator = audit_callable_once( + declared_scalar.neg, + expected_device="cuda", + allowed_host_tensors=(allowance,), + ) + assert not wrong_operator.passed + assert wrong_operator.placement.violation_evidence[0].operator == "aten.neg.default" + + +def test_declared_destination_write_coverage_accepts_full_write_and_rejects_partial_write() -> None: + source = torch.arange(8, dtype=torch.float32) + full_destination = torch.empty_like(source) + full = audit_callable_once( + lambda: full_destination.copy_(source), + expected_device="cpu", + destinations={"output": full_destination}, + ) + assert full.passed + assert len(full.destinations) == 1 + full_evidence = full.destinations[0] + assert full_evidence.shape == (8,) + assert full_evidence.numel == 8 + assert full_evidence.nbytes == 32 + assert full_evidence.unwritten_elements == 0 + torch.testing.assert_close(full_destination, source, rtol=0, atol=0) + + partial_destination = torch.empty_like(source) + partial = audit_callable_once( + lambda: partial_destination[:3].copy_(source[:3]), + expected_device="cpu", + destinations={"output": partial_destination}, + ) + assert partial.placement.passed + assert not partial.passed + partial_evidence = partial.destinations[0] + assert partial_evidence.unwritten_elements == 5 + assert partial_evidence.first_unwritten_flat_indices == (3, 4, 5, 6, 7) + + +@pytest.mark.parametrize( + ("tensor", "exception", "diagnostic"), + [ + (torch.empty(8, dtype=torch.int64), TypeError, "floating-point or complex"), + (torch.empty(2, 4).T, ValueError, "contiguous"), + (torch.empty(0), ValueError, "empty"), + ], +) +def test_destination_write_coverage_refuses_unprovable_tensor_contracts( + tensor: torch.Tensor, + exception: type[Exception], + diagnostic: str, +) -> None: + with pytest.raises(exception, match=diagnostic): + DestinationWriteCoverageGuard({"output": tensor}) + + +def _write_cli_benchmark(path: Path, *, mode: str) -> None: + write_statement = { + "full": "torch.mul(self.input, 2, out=self.output)", + "partial": "self.output[:3].copy_(self.input[:3] * 2)", + "reassigned": ( + "torch.mul(self.input, 2, out=self.output)\n" + " self.output = self.input * 3" + ), + }[mode] + path.write_text( + textwrap.dedent( + f""" + import torch + + class AuditedCpuBenchmark: + def setup(self): + self.input = torch.arange(6, dtype=torch.float32) + self.output = torch.empty_like(self.input) + + def benchmark_fn(self): + {write_statement} + + def teardown(self): + pass + + def get_benchmark(): + return AuditedCpuBenchmark() + """ + ).strip() + + "\n", + encoding="utf-8", + ) + + +@pytest.mark.parametrize( + ("mode", "expected_returncode", "unwritten_elements", "identity_errors"), + [ + ("full", 0, 0, []), + ("partial", 2, 3, []), + ( + "reassigned", + 2, + 0, + ["destination attribute 'output' no longer references the declared tensor"], + ), + ], +) +def test_fresh_benchmark_cli_runs_real_out_of_timing_audit( + tmp_path: Path, + mode: str, + expected_returncode: int, + unwritten_elements: int, + identity_errors: list[str], +) -> None: + benchmark_path = tmp_path / "audited_cpu_benchmark.py" + _write_cli_benchmark(benchmark_path, mode=mode) + completed = subprocess.run( + [ + sys.executable, + "-m", + "core.harness.execution_audit", + str(benchmark_path), + "--expected-device", + "cpu", + "--destination", + "output", + ], + cwd=Path(__file__).resolve().parents[1], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + assert completed.returncode == expected_returncode, completed.stderr + payload = json.loads(completed.stdout) + assert payload["schema"] == "aisp.execution-audit.v1" + assert payload["fresh_instance"] is True + assert payload["normal_timing_lifecycle_modified"] is False + assert payload["placement"]["operations_seen"] > 0 + assert payload["destinations"][0]["numel"] == 6 + assert payload["destinations"][0]["unwritten_elements"] == unwritten_elements + assert payload["destination_identity_errors"] == identity_errors + assert payload["passed"] is (mode == "full") + + +@requires_cuda +def test_cuda_operation_audit_detects_real_cpu_spillover() -> None: + device = torch.device("cuda", torch.cuda.current_device()) + cuda_value = torch.arange(16, dtype=torch.float32, device=device) + cpu_value = torch.arange(5, dtype=torch.float32) + + clean = audit_callable_once(lambda: cuda_value.square(), expected_device=device) + assert clean.passed + assert clean.placement.operations_seen == 1 + assert {item.device for item in clean.placement.operation_evidence[0].tensors} == {str(device)} + + def spill_to_cpu() -> None: + cuda_value.square() + cpu_value.square() + + violation = audit_callable_once(spill_to_cpu, expected_device=device) + assert not violation.passed + assert violation.placement.operations_seen == 2 + assert violation.placement.violations_seen == 1 + record = violation.placement.violation_evidence[0] + assert record.operator == "aten.pow.Tensor_Scalar" + assert {(item.device, item.shape, item.numel) for item in record.tensors} == {("cpu", (5,), 5)} + + +@requires_cuda +def test_cuda_destination_write_coverage_detects_real_partial_write() -> None: + device = torch.device("cuda", torch.cuda.current_device()) + source = torch.arange(16, dtype=torch.float32, device=device) + + full_destination = torch.empty_like(source) + full = audit_callable_once( + lambda: full_destination.copy_(source), + expected_device=device, + destinations={"output": full_destination}, + ) + assert full.passed + assert full.destinations[0].unwritten_elements == 0 + + partial_destination = torch.empty_like(source) + partial = audit_callable_once( + lambda: partial_destination[:7].copy_(source[:7]), + expected_device=device, + destinations={"output": partial_destination}, + ) + assert partial.placement.passed + assert not partial.passed + assert partial.destinations[0].shape == (16,) + assert partial.destinations[0].unwritten_elements == 9 + assert partial.destinations[0].first_unwritten_flat_indices == tuple(range(7, 15)) From ce4681dd63740ca10a66ea3627d83a6e2d9e7a1b Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 02:20:47 -0700 Subject: [PATCH 03/19] fix: freeze independent KV and Ozaki arithmetic requirements --- code/core/scripts/refresh_readmes.py | 27 +- .../ACCURACY_REQUIREMENTS.md | 69 +++++ code/labs/kv_cache_compression/README.md | 29 +- code/labs/kv_cache_compression/accuracy.py | 120 +++++++- .../kv_cache_compression/accuracy_policy.json | 45 +++ .../calibrate_accuracy.py | 96 ++++++- .../kv_cache_compression/qualify_accuracy.py | 123 +++++++++ .../ozaki_scheme/ACCURACY_REQUIREMENTS.md | 120 ++++++++ code/labs/ozaki_scheme/README.md | 32 ++- code/labs/ozaki_scheme/accuracy.h | 36 +++ code/labs/ozaki_scheme/accuracy_policy.json | 44 +++ code/labs/ozaki_scheme/accuracy_policy.py | 127 ++++++++- code/labs/ozaki_scheme/lab_utils.py | 22 +- .../labs/ozaki_scheme/ozaki_scheme_common.cuh | 137 ++++++++- code/labs/ozaki_scheme/qualify_accuracy.py | 134 +++++++++ ...est_accuracy_requirements_followthrough.py | 260 ++++++++++++++++++ code/tests/test_refresh_readmes.py | 5 +- 17 files changed, 1340 insertions(+), 86 deletions(-) create mode 100644 code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md create mode 100644 code/labs/kv_cache_compression/accuracy_policy.json create mode 100644 code/labs/kv_cache_compression/qualify_accuracy.py create mode 100644 code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md create mode 100644 code/labs/ozaki_scheme/accuracy_policy.json create mode 100644 code/labs/ozaki_scheme/qualify_accuracy.py create mode 100644 code/tests/test_accuracy_requirements_followthrough.py diff --git a/code/core/scripts/refresh_readmes.py b/code/core/scripts/refresh_readmes.py index 7a46b8b9a..33db0119e 100644 --- a/code/core/scripts/refresh_readmes.py +++ b/code/core/scripts/refresh_readmes.py @@ -4982,49 +4982,50 @@ def lab_entry( ), ), MarkdownSection( - 'Accuracy gate: target calibration pending', + 'Accuracy gate: requirements defined, target qualification pending', dedent( """\ Every token, head, and channel in both K and V is checked against an independent PyTorch BF16 projection reference using the original weights and inputs. The reference bypasses Transformer Engine's GEMMs and packing. Checks reject shape mismatches, non-finite values and aliased reference storage; relative L2 and maximum error normalized by reference magnitude avoid signed-checksum cancellation. Verification then snapshots the full cache for the harness pair comparison. - No workload accuracy bound has been calibrated. An accepted benchmark run requires `AISP_KV_CACHE_ACCURACY_POLICY` pointing to a JSON file with `schema_version: 1`, separate `fp8` and `nvfp4` objects, and the fields `relative_l2`, `normalized_max_abs`, `pairwise_rtol`, `pairwise_atol`. The first three must be finite and in `[0,1)`; the last must be finite and nonnegative. Bounds are deliberately not supplied here. Configuring bounds is not evidence that they are appropriate or that this workload passes them. + [`ACCURACY_REQUIREMENTS.md`](ACCURACY_REQUIREMENTS.md) records the reviewed, predeclared arithmetic policy and exact B200 qualification matrix. Full-cache relative-L2 and normalized-maximum ceilings are `2^-4` for E4M3 FP8 and `2^-2` for E2M1 NVFP4. These are format-scale engineering limits, not attention, model, task, or application-quality guarantees. `accuracy.py` rejects any configured limit above the checked-in source ceiling. - Collect measurements on the actual CUDA/Transformer Engine host before reviewing a policy: + The policy requires nominal, unseen holdout, alternating-sign, and sparse-outlier receipts for both variants. Collect each on the actual CUDA/Transformer Engine host without accepting a benchmark result: ```bash - python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --seed 42 --output /tmp/kv-fp8-seed42.json - python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --seed 42 --output /tmp/kv-nvfp4-seed42.json + python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --cohort nominal --seed 2026 --output /tmp/kv-fp8-nominal-2026.json + python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --cohort nominal --seed 2026 --output /tmp/kv-nvfp4-nominal-2026.json ``` - Repeat with other fixed seeds and preserve the hardware, software and workload metadata. These commands collect error metrics only; they do not accept output or claim a speedup. After independent accuracy review, run: + `qualify_accuracy.py` requires the complete matrix and retains every failure reason. These commands collect error metrics only; they do not accept output or claim a speedup. After the matrix passes, select the same policy for the ordinary pair: ```bash - AISP_KV_CACHE_ACCURACY_POLICY=/absolute/path/reviewed-policy.json python -m cli.aisp bench run --targets labs/kv_cache_compression:kv_cache --profile minimal + AISP_KV_CACHE_ACCURACY_POLICY="$PWD/labs/kv_cache_compression/accuracy_policy.json" python -m cli.aisp bench run --targets labs/kv_cache_compression:kv_cache --profile minimal ``` - The historical 6066.040/5897.083 ms measurements used the old permissive verifier and are not evidence for the revised accuracy contract or cache compression. Fresh GPU accuracy, memory and performance measurements remain pending. + Historical calibration and timing receipts remain diagnostics; they were not used to widen these ceilings and do not establish qualification or cache compression. Fresh B200 accuracy and performance measurements remain pending. """ ), ), ], goals=[ 'Compare FP8 and NVFP4 projection GEMMs with the same BF16 KV cache storage.', - 'Measure full-cache numerical error before reviewing any accuracy policy.', + 'Qualify predeclared format-scale arithmetic ceilings across nominal, holdout, and edge cohorts.', 'Keep allocated storage bytes separate from compute precision and latency.', ], contents=[ ('`baseline_kv_cache.py`, `optimized_kv_cache_nvfp4.py`', 'FP8/NVFP4 compute benchmark pair with BF16 cache storage.'), ('`kv_cache_common.py`', 'Shared attention workload and cache allocation.'), - ('`accuracy.py`, `calibrate_accuracy.py`', 'Independent full-cache reference, explicit policy, and measurement-only driver.'), + ('`accuracy.py`, `accuracy_policy.json`, `calibrate_accuracy.py`, `qualify_accuracy.py`', 'Independent full-cache reference, source-bounded policy, measurement-only driver, and retained-receipt qualifier.'), + ('`ACCURACY_REQUIREMENTS.md`', 'Threshold rationale, claim boundary, and exact serial B200 qualification plan.'), ], run=RunSection( - commands=['python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --seed 42 --output /tmp/kv-fp8-seed42.json', 'python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --seed 42 --output /tmp/kv-nvfp4-seed42.json'], - notes=['These collect error metrics without accepting an accuracy threshold. Accepted benchmark runs require the separately reviewed policy described above.'], + commands=['python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --cohort nominal --seed 2026 --output /tmp/kv-fp8-nominal-2026.json', 'python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --cohort nominal --seed 2026 --output /tmp/kv-nvfp4-nominal-2026.json'], + notes=['These collect measurement-only errors. Run the complete matrix in `ACCURACY_REQUIREMENTS.md` and pass `qualify_accuracy.py` before selecting the checked-in policy for an accepted benchmark.'], ), run_heading='Collecting Accuracy Measurements', run_intro='Run on the actual CUDA/Transformer Engine host, preserving target and workload metadata.', validation=[ - 'Require an independently reviewed accuracy policy and full-output comparisons before accepting timing.', + 'Require the checked-in source-bounded policy, complete nominal/holdout/edge receipt matrix, and full-output comparisons before accepting timing.', 'Reject zeros, corruption, non-finite values, aliasing, and shape mismatches using the independent reference.', 'Verify allocated cache storage bytes and the BF16-relative compression ratio of 1.0.', ], diff --git a/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md b/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md new file mode 100644 index 000000000..a47b8bb6c --- /dev/null +++ b/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md @@ -0,0 +1,69 @@ +# KV projection arithmetic requirements + +## Decision fixed before target runs + +The reviewed policy is [`accuracy_policy.json`](accuracy_policy.json). It compares every +BF16 K/V cache element with the existing unquantized BF16 PyTorch projection path, +which bypasses Transformer Engine quantization and packing. Shape, dtype, finite-value, +storage-alias, relative-L2, and maximum-error checks all remain mandatory. + +The source ceilings are: + +| Variant | Full-cache relative L2 | Maximum error / maximum reference | Pairwise rtol | Pairwise atol | +| --- | ---: | ---: | ---: | ---: | +| Delayed-scaling FP8 E4M3 | `0.0625` (`2^-4`) | `0.0625` (`2^-4`) | `0.25` | `0.0625` | +| NVFP4 E2M1 | `0.25` (`2^-2`) | `0.25` (`2^-2`) | `0.25` | `0.0625` | + +E4M3 stores three fraction bits, making half the spacing within a normal binade +`2^-4`; E2M1 stores one fraction bit, making the analogous quantity `2^-2`. +Those representation-scale quantities define the independent full-cache engineering +ceilings. NVFP4 also uses a per-16-element E4M3 block scale and a global FP32 scale, +as described in the [Transformer Engine NVFP4 documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/features/low_precision_training/nvfp4/nvfp4.html). +The pairwise allowance is secondary: each arm must first pass its own full-cache +reference check. + +These limits were fixed from the declared formats before the new candidate runs. +Prior calibration errors remain diagnostics and were not used to widen a bound. +Changing the JSON above a source ceiling is rejected by `load_accuracy_policy()`; +a deliberate ceiling change therefore requires a reviewed source change. + +## Required qualification matrix + +Both variants must pass one nominal seed, two unseen holdout seeds, an alternating-sign +edge distribution, and a sparse two-channel outlier distribution. All cohorts preserve +the production batch, hidden dimension, sequence lengths, cache shape, and stored dtype. +`calibrate_accuracy` writes measurement-only receipts, and `qualify_accuracy` rejects a +missing, duplicate, mismatched, non-finite, or over-budget receipt while retaining every +reason in its summary. + +Run these commands serially on the requested B200 from `code/`: + +```bash +accuracy_out=/tmp/ai-perf-followthrough-20260908-private/accuracy +mkdir -p "$accuracy_out" + +for variant in fp8 nvfp4; do + python -m labs.kv_cache_compression.calibrate_accuracy --variant "$variant" --cohort nominal --seed 2026 --output "$accuracy_out/kv-receipt-$variant-nominal-2026.json" + python -m labs.kv_cache_compression.calibrate_accuracy --variant "$variant" --cohort holdout --seed 2027 --output "$accuracy_out/kv-receipt-$variant-holdout-2027.json" + python -m labs.kv_cache_compression.calibrate_accuracy --variant "$variant" --cohort holdout --seed 2029 --output "$accuracy_out/kv-receipt-$variant-holdout-2029.json" + python -m labs.kv_cache_compression.calibrate_accuracy --variant "$variant" --cohort alternating --seed 2039 --output "$accuracy_out/kv-receipt-$variant-alternating-2039.json" + python -m labs.kv_cache_compression.calibrate_accuracy --variant "$variant" --cohort sparse_outlier --seed 2053 --output "$accuracy_out/kv-receipt-$variant-sparse_outlier-2053.json" +done + +python -m labs.kv_cache_compression.qualify_accuracy \ + --policy labs/kv_cache_compression/accuracy_policy.json \ + --output "$accuracy_out/kv-qualification.json" \ + "$accuracy_out"/kv-receipt-*.json +``` + +Only after `kv-qualification.json` says `qualified_arithmetic_gate`, run the ordinary +pair with the same checked-in policy: + +```bash +AISP_KV_CACHE_ACCURACY_POLICY="$PWD/labs/kv_cache_compression/accuracy_policy.json" \ + python -m cli.aisp bench run --targets labs/kv_cache_compression:kv_cache --profile minimal +``` + +Passing establishes this lab's arithmetic full-cache requirement on the recorded +hardware and software stack. It does not establish attention quality, model quality, +task quality, or application quality, and it does not mean the BF16 cache is compressed. diff --git a/code/labs/kv_cache_compression/README.md b/code/labs/kv_cache_compression/README.md index d5f0fe49a..fec485e81 100644 --- a/code/labs/kv_cache_compression/README.md +++ b/code/labs/kv_cache_compression/README.md @@ -8,29 +8,29 @@ Both variants use batch 8, hidden dimension 16384, 64 heads, 4096 prefill tokens The FP8 recipe is `DelayedScaling`; it is not MXFP8 block scaling. The NVFP4 recipe uses supported `NVFP4BlockScaling()` defaults. Both retain identical unquantized BF16 parameter representations while Transformer Engine autocast chooses the low-precision GEMMs. -## Accuracy gate: target calibration pending +## Accuracy gate: requirements defined, target qualification pending Every token, head, and channel in both K and V is checked against an independent PyTorch BF16 projection reference using the original weights and inputs. The reference bypasses Transformer Engine's GEMMs and packing. Checks reject shape mismatches, non-finite values and aliased reference storage; relative L2 and maximum error normalized by reference magnitude avoid signed-checksum cancellation. Verification then snapshots the full cache for the harness pair comparison. -No workload accuracy bound has been calibrated. An accepted benchmark run requires `AISP_KV_CACHE_ACCURACY_POLICY` pointing to a JSON file with `schema_version: 1`, separate `fp8` and `nvfp4` objects, and the fields `relative_l2`, `normalized_max_abs`, `pairwise_rtol`, `pairwise_atol`. The first three must be finite and in `[0,1)`; the last must be finite and nonnegative. Bounds are deliberately not supplied here. Configuring bounds is not evidence that they are appropriate or that this workload passes them. +[`ACCURACY_REQUIREMENTS.md`](ACCURACY_REQUIREMENTS.md) records the reviewed, predeclared arithmetic policy and exact B200 qualification matrix. Full-cache relative-L2 and normalized-maximum ceilings are `2^-4` for E4M3 FP8 and `2^-2` for E2M1 NVFP4. These are format-scale engineering limits, not attention, model, task, or application-quality guarantees. `accuracy.py` rejects any configured limit above the checked-in source ceiling. -Collect measurements on the actual CUDA/Transformer Engine host before reviewing a policy: +The policy requires nominal, unseen holdout, alternating-sign, and sparse-outlier receipts for both variants. Collect each on the actual CUDA/Transformer Engine host without accepting a benchmark result: ```bash -python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --seed 42 --output /tmp/kv-fp8-seed42.json -python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --seed 42 --output /tmp/kv-nvfp4-seed42.json +python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --cohort nominal --seed 2026 --output /tmp/kv-fp8-nominal-2026.json +python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --cohort nominal --seed 2026 --output /tmp/kv-nvfp4-nominal-2026.json ``` -Repeat with other fixed seeds and preserve the hardware, software and workload metadata. These commands collect error metrics only; they do not accept output or claim a speedup. After independent accuracy review, run: +`qualify_accuracy.py` requires the complete matrix and retains every failure reason. These commands collect error metrics only; they do not accept output or claim a speedup. After the matrix passes, select the same policy for the ordinary pair: ```bash -AISP_KV_CACHE_ACCURACY_POLICY=/absolute/path/reviewed-policy.json python -m cli.aisp bench run --targets labs/kv_cache_compression:kv_cache --profile minimal +AISP_KV_CACHE_ACCURACY_POLICY="$PWD/labs/kv_cache_compression/accuracy_policy.json" python -m cli.aisp bench run --targets labs/kv_cache_compression:kv_cache --profile minimal ``` -The historical 6066.040/5897.083 ms measurements used the old permissive verifier and are not evidence for the revised accuracy contract or cache compression. Fresh GPU accuracy, memory and performance measurements remain pending. +Historical calibration and timing receipts remain diagnostics; they were not used to widen these ceilings and do not establish qualification or cache compression. Fresh B200 accuracy and performance measurements remain pending. ## Learning Goals - Compare FP8 and NVFP4 projection GEMMs with the same BF16 KV cache storage. -- Measure full-cache numerical error before reviewing any accuracy policy. +- Qualify predeclared format-scale arithmetic ceilings across nominal, holdout, and edge cohorts. - Keep allocated storage bytes separate from compute precision and latency. ## Directory Layout @@ -38,18 +38,19 @@ The historical 6066.040/5897.083 ms measurements used the old permissive verifie | --- | --- | | `baseline_kv_cache.py`, `optimized_kv_cache_nvfp4.py` | FP8/NVFP4 compute benchmark pair with BF16 cache storage. | | `kv_cache_common.py` | Shared attention workload and cache allocation. | -| `accuracy.py`, `calibrate_accuracy.py` | Independent full-cache reference, explicit policy, and measurement-only driver. | +| `accuracy.py`, `accuracy_policy.json`, `calibrate_accuracy.py`, `qualify_accuracy.py` | Independent full-cache reference, source-bounded policy, measurement-only driver, and retained-receipt qualifier. | +| `ACCURACY_REQUIREMENTS.md` | Threshold rationale, claim boundary, and exact serial B200 qualification plan. | ## Collecting Accuracy Measurements Run on the actual CUDA/Transformer Engine host, preserving target and workload metadata. ```bash -python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --seed 42 --output /tmp/kv-fp8-seed42.json -python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --seed 42 --output /tmp/kv-nvfp4-seed42.json +python -m labs.kv_cache_compression.calibrate_accuracy --variant fp8 --cohort nominal --seed 2026 --output /tmp/kv-fp8-nominal-2026.json +python -m labs.kv_cache_compression.calibrate_accuracy --variant nvfp4 --cohort nominal --seed 2026 --output /tmp/kv-nvfp4-nominal-2026.json ``` -- These collect error metrics without accepting an accuracy threshold. Accepted benchmark runs require the separately reviewed policy described above. +- These collect measurement-only errors. Run the complete matrix in `ACCURACY_REQUIREMENTS.md` and pass `qualify_accuracy.py` before selecting the checked-in policy for an accepted benchmark. ## Validation Checklist -- Require an independently reviewed accuracy policy and full-output comparisons before accepting timing. +- Require the checked-in source-bounded policy, complete nominal/holdout/edge receipt matrix, and full-output comparisons before accepting timing. - Reject zeros, corruption, non-finite values, aliasing, and shape mismatches using the independent reference. - Verify allocated cache storage bytes and the BF16-relative compression ratio of 1.0. diff --git a/code/labs/kv_cache_compression/accuracy.py b/code/labs/kv_cache_compression/accuracy.py index 22f5f3d26..85f762cc3 100644 --- a/code/labs/kv_cache_compression/accuracy.py +++ b/code/labs/kv_cache_compression/accuracy.py @@ -1,7 +1,8 @@ """Full-cache accuracy checks and independent unquantized BF16 reference. -No quantization acceptance threshold has been calibrated for this workload. A -missing policy is an error, not permission to use a permissive default. +The checked-in policy defines conservative arithmetic requirements from the +operand formats. A caller must still select that policy explicitly: merely +collecting candidate errors never changes an acceptance bound. """ from __future__ import annotations @@ -18,6 +19,27 @@ from labs.kv_cache_compression.kv_cache_common import KVCache +REFERENCE_ID = "pytorch-unquantized-bf16-full-cache-v1" +POLICY_ID = "kv-cache-projection-format-ceilings-v1" +DEFAULT_POLICY_PATH = Path(__file__).with_name("accuracy_policy.json") +WORKLOAD = { + "batch_size": 8, + "hidden_dim": 16384, + "num_heads": 64, + "prefill_seq": 4096, + "decode_seq": 128, + "decode_steps": 128, + "storage_dtype": "bfloat16", +} +QUALIFICATION_VARIANTS = ["fp8", "nvfp4"] +QUALIFICATION_RECEIPTS = [ + {"cohort": "nominal", "seeds": [2026]}, + {"cohort": "holdout", "seeds": [2027, 2029]}, + {"cohort": "alternating", "seeds": [2039]}, + {"cohort": "sparse_outlier", "seeds": [2053]}, +] + + @dataclass(frozen=True) class AccuracyLimits: relative_l2: float @@ -34,18 +56,98 @@ def __post_init__(self): raise ValueError("pairwise_atol must be finite and nonnegative") +# These ceilings are set from the quantized operand representations, before +# candidate execution. E4M3 has three stored fraction bits, so half an ULP at +# a normal binade is 2^-4. E2M1 has one stored fraction bit, so the analogous +# bound is 2^-2. The full-cache aggregate and global-maximum requirements use +# those format-scale bounds. The pairwise check compares FP8 and NVFP4 caches +# only after each arm passes its independent reference check, so it uses the +# coarser E2M1 ceiling plus one E4M3-scale absolute allowance near zero. +ENGINEERING_CEILINGS = { + "fp8": AccuracyLimits( + relative_l2=2.0**-4, + normalized_max_abs=2.0**-4, + pairwise_rtol=2.0**-2, + pairwise_atol=2.0**-4, + ), + "nvfp4": AccuracyLimits( + relative_l2=2.0**-2, + normalized_max_abs=2.0**-2, + pairwise_rtol=2.0**-2, + pairwise_atol=2.0**-4, + ), +} + + +def _limits_from_item(item: dict) -> AccuracyLimits: + try: + return AccuracyLimits(**{ + name: float(item[name]) + for name in ("relative_l2", "normalized_max_abs", "pairwise_rtol", "pairwise_atol") + }) + except KeyError as exc: + raise ValueError(f"KV accuracy policy missing {exc.args[0]}") from exc + + +def load_accuracy_policy(path: Path) -> dict: + """Load and validate the reviewed policy contract, without running a candidate.""" + policy = json.loads(path.read_text()) + schema_version = policy.get("schema_version") + if schema_version == 1: + # Preserve exact synthetic fixture policies. A nonzero legacy policy has + # no reference/workload identity and cannot accept a quantized run. + for variant in ("fp8", "nvfp4"): + if variant not in policy: + continue + limits = _limits_from_item(policy[variant]) + if any(getattr(limits, name) != 0 for name in ( + "relative_l2", "normalized_max_abs", "pairwise_rtol", "pairwise_atol" + )): + raise ValueError("schema_version=1 is permitted only for exact-zero test policies") + return policy + if schema_version != 2: + raise ValueError("KV accuracy policy requires schema_version=2") + if policy.get("policy_id") != POLICY_ID: + raise ValueError(f"KV accuracy policy requires policy_id={POLICY_ID}") + if policy.get("reference", {}).get("id") != REFERENCE_ID: + raise ValueError(f"KV accuracy policy requires reference.id={REFERENCE_ID}") + if policy.get("workload") != WORKLOAD: + raise ValueError("KV accuracy policy workload does not match the benchmark contract") + qualification = policy.get("qualification", {}) + if (qualification.get("variants") != QUALIFICATION_VARIANTS or + qualification.get("required_receipts") != QUALIFICATION_RECEIPTS): + raise ValueError("KV accuracy policy qualification matrix does not match the source contract") + variants = policy.get("variants") + if not isinstance(variants, dict): + raise ValueError("KV accuracy policy requires a variants object") + for variant, ceiling in ENGINEERING_CEILINGS.items(): + if variant not in variants: + raise ValueError(f"KV accuracy policy missing variant {variant}") + limits = _limits_from_item(variants[variant]) + for name in ("relative_l2", "normalized_max_abs", "pairwise_rtol", "pairwise_atol"): + if getattr(limits, name) > getattr(ceiling, name): + raise ValueError( + f"{variant}.{name} exceeds the source-defined engineering ceiling " + f"{getattr(ceiling, name):.8g}" + ) + return policy + + def load_accuracy_limits(variant: str) -> AccuracyLimits: path = os.environ.get("AISP_KV_CACHE_ACCURACY_POLICY") if not path: raise RuntimeError( - "KV compute accuracy is uncalibrated: AISP_KV_CACHE_ACCURACY_POLICY is required. " - "Use python -m labs.kv_cache_compression.calibrate_accuracy to collect errors; " - "a configured policy alone is not measured accuracy evidence." + "KV compute accuracy is uncalibrated because no policy is selected: " + "AISP_KV_CACHE_ACCURACY_POLICY is required. " + f"The reviewed repository policy is {DEFAULT_POLICY_PATH}. A configured policy alone " + "is not target-hardware accuracy evidence." ) - policy = json.loads(Path(path).read_text()) - if policy.get("schema_version") != 1: - raise ValueError("KV accuracy policy requires schema_version=1") - return AccuracyLimits(**policy[variant]) + policy = load_accuracy_policy(Path(path)) + try: + item = policy[variant] if policy["schema_version"] == 1 else policy["variants"][variant] + except KeyError as exc: + raise ValueError(f"KV accuracy policy missing variant {variant}") from exc + return _limits_from_item(item) def reference_cache(model, groups, cache: KVCache) -> KVCache: diff --git a/code/labs/kv_cache_compression/accuracy_policy.json b/code/labs/kv_cache_compression/accuracy_policy.json new file mode 100644 index 000000000..004a162c7 --- /dev/null +++ b/code/labs/kv_cache_compression/accuracy_policy.json @@ -0,0 +1,45 @@ +{ + "schema_version": 2, + "policy_id": "kv-cache-projection-format-ceilings-v1", + "reference": { + "id": "pytorch-unquantized-bf16-full-cache-v1", + "description": "Original BF16 inputs and weights through PyTorch layer_norm and linear, bypassing Transformer Engine quantization and packing; every stored K/V element is compared." + }, + "workload": { + "batch_size": 8, + "hidden_dim": 16384, + "num_heads": 64, + "prefill_seq": 4096, + "decode_seq": 128, + "decode_steps": 128, + "storage_dtype": "bfloat16" + }, + "variants": { + "fp8": { + "operand_format": "E4M3 forward operands with delayed per-tensor scaling", + "relative_l2": 0.0625, + "normalized_max_abs": 0.0625, + "pairwise_rtol": 0.25, + "pairwise_atol": 0.0625, + "rationale": "The independent full-cache aggregate and global-maximum ceilings are one E4M3 half-ULP at a normal binade (2^-4)." + }, + "nvfp4": { + "operand_format": "E2M1 values with per-16-element E4M3 block scale and FP32 tensor scale", + "relative_l2": 0.25, + "normalized_max_abs": 0.25, + "pairwise_rtol": 0.25, + "pairwise_atol": 0.0625, + "rationale": "The independent full-cache aggregate and global-maximum ceilings are one E2M1 half-ULP at a normal binade (2^-2); finite and overflow checks remain mandatory." + } + }, + "qualification": { + "required_receipts": [ + {"cohort": "nominal", "seeds": [2026]}, + {"cohort": "holdout", "seeds": [2027, 2029]}, + {"cohort": "alternating", "seeds": [2039]}, + {"cohort": "sparse_outlier", "seeds": [2053]} + ], + "variants": ["fp8", "nvfp4"], + "claim_boundary": "Passing establishes this lab's arithmetic full-cache gate on the recorded stack. It does not establish attention, model, task, or application quality." + } +} diff --git a/code/labs/kv_cache_compression/calibrate_accuracy.py b/code/labs/kv_cache_compression/calibrate_accuracy.py index 036908c68..c3b3ba76d 100644 --- a/code/labs/kv_cache_compression/calibrate_accuracy.py +++ b/code/labs/kv_cache_compression/calibrate_accuracy.py @@ -5,41 +5,107 @@ import argparse import importlib.metadata import json +import subprocess +import traceback from pathlib import Path import torch +from labs.kv_cache_compression.accuracy import REFERENCE_ID, WORKLOAD from labs.kv_cache_compression.baseline_kv_cache import BaselineKVCacheBenchmark from labs.kv_cache_compression.optimized_kv_cache_nvfp4 import OptimizedKVCacheNVFP4Benchmark +COHORTS = ("nominal", "holdout", "alternating", "sparse_outlier") + + +def _source_commit() -> str: + completed = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parents[3], + check=True, + capture_output=True, + text=True, + ) + return completed.stdout.strip() + + +def _apply_input_cohort(benchmark, cohort: str) -> None: + """Apply a deterministic distribution change without changing workload shape.""" + if cohort in ("nominal", "holdout"): + return + tensors = benchmark.prefill_inputs + benchmark.decode_inputs + with torch.no_grad(): + if cohort == "alternating": + feature = torch.ones(benchmark.hidden_dim, dtype=benchmark.tensor_dtype, device=benchmark.device) + feature[1::2] = -1 + for tensor in tensors: + tensor.copy_(feature) + tensor[:, 1::2].neg_() + elif cohort == "sparse_outlier": + for tensor in tensors: + tensor.zero_() + tensor[..., 0] = 1 + tensor[..., 1] = -1 + else: # pragma: no cover - argparse owns the public choices. + raise ValueError(f"Unknown KV accuracy cohort: {cohort}") + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--variant", choices=("fp8", "nvfp4"), required=True) parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--cohort", choices=COHORTS, default="nominal") parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("Accuracy calibration requires actual CUDA/Transformer Engine hardware") - torch.manual_seed(args.seed) - benchmark = (BaselineKVCacheBenchmark() if args.variant == "fp8" - else OptimizedKVCacheNVFP4Benchmark()) - recipe = benchmark.fp8_recipe if args.variant == "fp8" else benchmark.nvfp4_recipe + receipt = { + "schema_version": 2, + "status": "failure_not_accepted", + "variant": args.variant, + "cohort": args.cohort, + "seed": args.seed, + "reference_id": REFERENCE_ID, + "git_commit": _source_commit(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "workload": {key: value for key, value in WORKLOAD.items() if key != "storage_dtype"}, + "thresholds": None, + } + benchmark = None try: + if not torch.cuda.is_available(): + raise RuntimeError("Accuracy calibration requires actual CUDA/Transformer Engine hardware") + torch.manual_seed(args.seed) + benchmark = (BaselineKVCacheBenchmark() if args.variant == "fp8" + else OptimizedKVCacheNVFP4Benchmark()) + recipe = benchmark.fp8_recipe if args.variant == "fp8" else benchmark.nvfp4_recipe benchmark._setup_with_recipe(recipe, require_accuracy_policy=False) + _apply_input_cohort(benchmark, args.cohort) + if args.cohort not in ("nominal", "holdout"): + if recipe.delayed(): + benchmark._calibrate_fp8(recipe) + benchmark._warmup_runtime(recipe) benchmark.benchmark_fn() metrics = benchmark.measure_accuracy() - args.output.write_text(json.dumps({ - "schema_version": 1, "status": "measurement_only_not_accepted", - "variant": args.variant, "seed": args.seed, "torch": torch.__version__, + receipt.update({ + "status": "measurement_only_not_accepted", "transformer_engine": importlib.metadata.version("transformer_engine"), - "gpu": torch.cuda.get_device_name(), "metrics": metrics, - "workload": {key: getattr(benchmark, key) for key in ( - "batch_size", "hidden_dim", "num_heads", "prefill_seq", "decode_seq", "decode_steps")}, - "thresholds": None, - }, indent=2) + "\n") + "gpu": torch.cuda.get_device_name(), + "compute_capability": list(torch.cuda.get_device_capability()), + "metrics": metrics, + }) + args.output.write_text(json.dumps(receipt, indent=2) + "\n") + except Exception as exc: + receipt.update({ + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(), + }) + args.output.write_text(json.dumps(receipt, indent=2) + "\n") + raise finally: - benchmark.teardown() + if benchmark is not None: + benchmark.teardown() if __name__ == "__main__": diff --git a/code/labs/kv_cache_compression/qualify_accuracy.py b/code/labs/kv_cache_compression/qualify_accuracy.py new file mode 100644 index 000000000..2849a5547 --- /dev/null +++ b/code/labs/kv_cache_compression/qualify_accuracy.py @@ -0,0 +1,123 @@ +"""Qualify KV accuracy receipts against the predeclared policy and cohorts. + +This command never runs a benchmark. It consumes retained measurement-only +receipts, writes every rejection reason, and exits nonzero unless the complete +nominal, holdout, and edge matrix passes the source-defined ceilings. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +from labs.kv_cache_compression.accuracy import ( + DEFAULT_POLICY_PATH, + REFERENCE_ID, + WORKLOAD, + _limits_from_item, + load_accuracy_policy, +) + + +def required_cases(policy: dict) -> set[tuple[str, str, int]]: + cases: set[tuple[str, str, int]] = set() + qualification = policy["qualification"] + for variant in qualification["variants"]: + for group in qualification["required_receipts"]: + for seed in group["seeds"]: + cases.add((variant, group["cohort"], int(seed))) + return cases + + +def assess_receipts(policy: dict, receipts: list[dict]) -> dict: + """Return a durable pass/fail summary without discarding malformed receipts.""" + required = required_cases(policy) + seen: set[tuple[str, str, int]] = set() + failures: list[str] = [] + receipt_results: list[dict] = [] + expected_provenance: tuple | None = None + variants = policy["variants"] + expected_workload = {key: value for key, value in WORKLOAD.items() if key != "storage_dtype"} + + for index, receipt in enumerate(receipts): + try: + seed = int(receipt.get("seed", -1)) + except (TypeError, ValueError): + seed = -1 + key = (str(receipt.get("variant")), str(receipt.get("cohort")), seed) + reasons: list[str] = [] + if key not in required: + reasons.append("receipt is not a required policy case") + if key in seen: + reasons.append("duplicate policy case") + if receipt.get("schema_version") != 2: + reasons.append("receipt schema_version is not 2") + if receipt.get("status") != "measurement_only_not_accepted": + reasons.append("receipt is not measurement-only source evidence") + if receipt.get("reference_id") != REFERENCE_ID: + reasons.append("reference identity mismatch") + if receipt.get("workload") != expected_workload: + reasons.append("workload mismatch") + provenance = tuple(receipt.get(name) for name in ( + "git_commit", "torch", "cuda", "transformer_engine", "gpu", "compute_capability" + )) + if any(value in (None, "", []) for value in provenance): + reasons.append("hardware/software provenance is incomplete") + elif expected_provenance is None: + expected_provenance = provenance + elif provenance != expected_provenance: + reasons.append("hardware/software provenance differs across receipts") + metrics = receipt.get("metrics") + if key[0] in variants and isinstance(metrics, dict): + limits = _limits_from_item(variants[key[0]]) + for tensor in ("cache_k", "cache_v"): + for metric_name in ("relative_l2", "normalized_max_abs"): + name = f"{tensor}.{metric_name}" + value = metrics.get(name) + limit = getattr(limits, metric_name) + if not isinstance(value, (int, float)) or not math.isfinite(float(value)): + reasons.append(f"{name} is missing or non-finite") + elif float(value) > limit: + reasons.append(f"{name}={float(value):.8g} exceeds {limit:.8g}") + else: + reasons.append("variant or metrics are invalid") + if key in required and key not in seen: + seen.add(key) + if reasons: + failures.extend(f"receipt[{index}] {key}: {reason}" for reason in reasons) + receipt_results.append({"case": list(key), "passed": not reasons, "reasons": reasons}) + + for key in sorted(required - seen): + failures.append(f"missing required receipt: {key}") + return { + "schema_version": 1, + "policy_id": policy["policy_id"], + "status": "qualified_arithmetic_gate" if not failures else "failed_arithmetic_gate", + "required_case_count": len(required), + "passing_case_count": sum(item["passed"] for item in receipt_results), + "failures": failures, + "receipts": receipt_results, + "claim_boundary": policy["qualification"]["claim_boundary"], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("receipts", type=Path, nargs="+") + parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY_PATH) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + policy = load_accuracy_policy(args.policy) + if policy.get("schema_version") != 2: + raise ValueError("Receipt qualification requires the reviewed schema_version=2 policy") + receipts = [json.loads(path.read_text()) for path in args.receipts] + summary = assess_receipts(policy, receipts) + args.output.write_text(json.dumps(summary, indent=2) + "\n") + if summary["status"] != "qualified_arithmetic_gate": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md b/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md new file mode 100644 index 000000000..5579ae3a0 --- /dev/null +++ b/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md @@ -0,0 +1,120 @@ +# Ozaki arithmetic requirements + +## Decision fixed before target runs + +The reviewed policy is [`accuracy_policy.json`](accuracy_policy.json). Production-size +candidate arrays are compared element by element with a separately executed native-FP64 +cuBLAS result. Small edge cohorts additionally use `reference_gemm_long_double()`, a +CPU row-major GEMM with wider-than-FP64 accumulation. That path refuses to run when the +host aliases `long double` to FP64 or when a request exceeds 50 million FMAs. + +The source ceilings are: + +| Variant | Relative L2 | Maximum error / maximum reference | +| --- | ---: | ---: | +| Dynamic, max 16 bits, offset `-56` | `0.03125` (`2^-5`) | `0.0625` (`2^-4`) | +| Fixed 12-bit | `0.000244140625` (`2^-12`) | `0.00048828125` (`2^-11`) | + +NVIDIA documents that dynamic mantissa control targets native-FP64 accuracy at the +default precision, but a mantissa-bit offset explicitly trades accuracy for performance; +fixed control likewise does not guarantee native-FP64 accuracy. See the +[cuBLAS fixed-point emulation documentation](https://docs.nvidia.com/cuda/cublas/index.html#fixed-point). +Because this lab deliberately uses `dynamic_offset=-56`, the dynamic thresholds are +an explicit five-bit aggregate and four-bit worst-output engineering floor. Fixed-12 +must remain within one `2^-12` grid step in aggregate and two steps at the worst output. +These are acceptance decisions for the declared lab and cohorts, not forward-error +theorems for arbitrary or ill-conditioned matrices. + +The checksum tolerances are secondary harness bounds. They follow from the full-array +relative-L2 ceiling, Cauchy-Schwarz, and the declared uniform input bound: +`atol = relative_l2 * m * n * k * input_scale^2`; `rtol` is zero. Candidate results +cannot pass on the checksum alone because the executable first gates the complete array. + +These limits were fixed before new candidate execution. Prior measurement-only logs +remain diagnostics and were not used to widen a bound. `load_accuracy_policy()` rejects +any JSON value above the source-defined ceilings. + +## Required qualification matrix + +Both dynamic and fixed variants must pass the 4096-cubed nominal seed, two unseen +4096-cubed holdouts, and two small rectangular edges. The alternating and dynamic-range +edges use the independent CPU long-double reference. The executable prints the input, +algorithm, and reference identities so `qualify_accuracy` can reject substitutions. + +From `code/labs/ozaki_scheme/` on the requested B200, build once: + +```bash +make ARCH=sm_100 all +accuracy_out=/tmp/ai-perf-followthrough-20260908-private/accuracy +mkdir -p "$accuracy_out" +``` + +The measurement-only binary intentionally exits `2`. This helper preserves the log and +accepts only that explicit disposition: + +```bash +measure_only() { + output=$1 + shift + set +e + "$@" >"$output" 2>&1 + status=$? + set -e + if [ "$status" -ne 2 ]; then + return "$status" + fi +} +``` + +Run each candidate and cohort serially: + +```bash +for variant in dynamic fixed; do + if [ "$variant" = dynamic ]; then + binary=./optimized_ozaki_scheme_dynamic_sm100 + variant_args=(--dynamic-max-bits 16 --dynamic-offset -56) + else + binary=./optimized_ozaki_scheme_fixed_sm100 + variant_args=(--fixed-bits 12) + fi + + for seed in 2026 2027 2029; do + cohort=holdout + if [ "$seed" -eq 2026 ]; then cohort=nominal; fi + measure_only "$accuracy_out/ozaki-$variant-$cohort-$seed.log" "$binary" \ + --m 4096 --n 4096 --k 4096 --warmup 3 --iters 10 --seed "$seed" \ + --input-scale 0.001 --input-pattern uniform --reference-mode native_fp64 \ + --emulation-strategy eager "${variant_args[@]}" --accuracy-measure-only + done + + measure_only "$accuracy_out/ozaki-$variant-alternating_edge-2039.log" "$binary" \ + --m 17 --n 19 --k 23 --warmup 1 --iters 1 --seed 2039 --input-scale 1 \ + --input-pattern alternating --reference-mode cpu_long_double \ + --emulation-strategy eager "${variant_args[@]}" --accuracy-measure-only + + measure_only "$accuracy_out/ozaki-$variant-dynamic_range_edge-2053.log" "$binary" \ + --m 31 --n 29 --k 37 --warmup 1 --iters 1 --seed 2053 --input-scale 1 \ + --input-pattern dynamic_range --reference-mode cpu_long_double \ + --emulation-strategy eager "${variant_args[@]}" --accuracy-measure-only +done +``` + +Qualify the complete set from `code/`: + +```bash +python -m labs.ozaki_scheme.qualify_accuracy \ + --policy labs/ozaki_scheme/accuracy_policy.json \ + --output /tmp/ai-perf-followthrough-20260908-private/accuracy/ozaki-qualification.json \ + /tmp/ai-perf-followthrough-20260908-private/accuracy/ozaki-*.log +``` + +Only after the summary says `qualified_arithmetic_gate`, run the ordinary pair with +the same checked-in policy: + +```bash +AISP_OZAKI_ACCURACY_POLICY="$PWD/labs/ozaki_scheme/accuracy_policy.json" \ + python -m cli.aisp bench run --targets labs/ozaki_scheme --profile minimal +``` + +Passing establishes the declared arithmetic GEMM requirement for this stack and these +cohorts. It does not establish solver, model, task, or application quality. diff --git a/code/labs/ozaki_scheme/README.md b/code/labs/ozaki_scheme/README.md index 2959b5fff..770844985 100644 --- a/code/labs/ozaki_scheme/README.md +++ b/code/labs/ozaki_scheme/README.md @@ -9,7 +9,8 @@ The implemented runtime paths are: - Ozaki-I style dynamic retained-bit control through cuBLAS FP64 emulation - Ozaki-I style fixed retained-bit control through cuBLAS FP64 emulation -The lab includes narrative-check drivers, but their optimized runs require an explicitly reviewed accuracy policy. The following claims remain target-validation gates: +The lab includes narrative-check drivers and a checked-in, source-bounded arithmetic +policy. The following claims still require fresh target qualification: - controllable accuracy - adaptive retained-bit behavior @@ -43,21 +44,34 @@ These earlier numbers predate the full-array verifier. The old signed-checksum c The dynamic path's reported `3e-6` maximum error is material relative to this small output scale. The old tolerance does not establish FP64-equivalent accuracy. -## Accuracy gate: target calibration pending +## Accuracy gate: requirements defined, target qualification pending -The CUDA executable now checks every element against a separately allocated native FP64 reference before returning an accepted timing. The host comparator in `accuracy.h` rejects non-finite values and aliased reference storage and reports relative L2 and maximum absolute error normalized by the largest reference magnitude. Both accuracy limits must be explicitly configured, finite, and in `[0,1)`. The Python checksum comparison is secondary and defaults to exact equality; it cannot replace the full-array check. +The CUDA executable checks every element against a separately executed native FP64 +reference before returning an accepted timing. Small edge cohorts can instead use the +CPU long-double GEMM in `accuracy.h`, which fails closed when `long double` is not wider +than FP64. The comparator rejects non-finite values and overlapping reference storage, +then reports relative L2 and maximum absolute error normalized by the largest reference +magnitude. The Python checksum comparison remains secondary and cannot replace the +full-array check. -No numerical threshold has been calibrated or approved here. Without limits, optimized binaries fail before GPU allocation. Python wrappers and runners accept `AISP_OZAKI_ACCURACY_POLICY`, a JSON file with `schema_version: 1` and separate `dynamic`/`fixed` objects containing `relative_l2`, `normalized_max_abs`, `checksum_rtol`, and `checksum_atol`. The first three fields must be finite and in `[0,1)`; the last must be finite and nonnegative. A configured policy is not measured accuracy evidence. +[`ACCURACY_REQUIREMENTS.md`](ACCURACY_REQUIREMENTS.md) records the predeclared policy, +rationale, nominal and holdout seeds, independent-reference edge cohorts, exact B200 +commands, and claim boundary. The dynamic path is capped at relative L2 `2^-5` and +normalized maximum error `2^-4`; fixed-12 is capped at `2^-12` and `2^-11`. These are +lab arithmetic requirements, not FP64-equivalence or application-quality guarantees. +`accuracy_policy.py` rejects any configured limit above the checked-in source ceiling. -After compiling on the target, collect errors without accepting a benchmark result (binary suffix depends on the target): +Select the checked-in policy for ordinary accepted runs: ```bash -make -C labs/ozaki_scheme all -labs/ozaki_scheme/optimized_ozaki_scheme_dynamic_sm100 --m 4096 --n 4096 --k 4096 --seed 2026 --input-scale 0.001 --dynamic-max-bits 16 --dynamic-offset -56 --accuracy-measure-only -labs/ozaki_scheme/optimized_ozaki_scheme_fixed_sm100 --m 4096 --n 4096 --k 4096 --seed 2026 --input-scale 0.001 --fixed-bits 12 --accuracy-measure-only +AISP_OZAKI_ACCURACY_POLICY="$PWD/labs/ozaki_scheme/accuracy_policy.json" \ + python -m cli.aisp bench run --targets labs/ozaki_scheme --profile minimal ``` -Measurement-only runs exit **2** with `ACCURACY_STATUS: MEASUREMENT_ONLY_NOT_ACCEPTED`; they omit `TIME_MS`, TFLOPS and the verifier checksum. Sweep seeds, retained-bit settings, and input scales before selecting workload-specific bounds. Repeated controlled GPU runs, numerical calibration and performance acceptance remain open gates. +This command is acceptable only after the required measurement-only matrix qualifies. +Those runs still exit **2** with `ACCURACY_STATUS: MEASUREMENT_ONLY_NOT_ACCEPTED` and +remain retained evidence if a cohort fails. Fresh B200 qualification and repeated +performance acceptance remain open gates. ## Why This Lab Exists The motivating story from the slides is that low-precision tensor-core hardware keeps getting faster while native FP64 throughput improves much more slowly, so accurate FP64-equivalent matrix multiplication increasingly wants an emulation story instead of a brute-force FP64 story. diff --git a/code/labs/ozaki_scheme/accuracy.h b/code/labs/ozaki_scheme/accuracy.h index fa435f482..fd2d84993 100644 --- a/code/labs/ozaki_scheme/accuracy.h +++ b/code/labs/ozaki_scheme/accuracy.h @@ -7,6 +7,7 @@ #include #include #include +#include namespace ozaki_scheme { struct AccuracyMetrics { @@ -16,6 +17,41 @@ struct AccuracyMetrics { double normalized_max_abs = 0; }; +inline std::vector reference_gemm_long_double( + const double* a, + const double* b, + std::size_t m, + std::size_t n, + std::size_t k) { + if (!a || !b || !m || !n || !k || + m > std::numeric_limits::max() / k || + k > std::numeric_limits::max() / n || + m > std::numeric_limits::max() / n) { + throw std::runtime_error("CPU long-double reference requires valid nonempty matrix dimensions"); + } + if (std::numeric_limits::digits <= std::numeric_limits::digits) { + throw std::runtime_error("CPU long-double reference requires precision wider than FP64"); + } + constexpr std::size_t kMaxReferenceFmas = 50'000'000; + if (m > kMaxReferenceFmas / n || m * n > kMaxReferenceFmas / k) { + throw std::runtime_error("CPU long-double reference is limited to 50,000,000 FMAs"); + } + std::vector result(m * n); + for (std::size_t row = 0; row < m; ++row) { + for (std::size_t column = 0; column < n; ++column) { + long double sum = 0; + for (std::size_t inner = 0; inner < k; ++inner) { + sum = std::fma( + static_cast(a[row * k + inner]), + static_cast(b[inner * n + column]), + sum); + } + result[row * n + column] = static_cast(sum); + } + } + return result; +} + inline AccuracyMetrics measure_accuracy(const double* actual, const double* reference, std::size_t count) { if (!count || !actual || !reference || count > std::numeric_limits::max() / sizeof(double)) { throw std::runtime_error("Accuracy requires nonempty independent candidate/reference arrays"); diff --git a/code/labs/ozaki_scheme/accuracy_policy.json b/code/labs/ozaki_scheme/accuracy_policy.json new file mode 100644 index 000000000..a6084b3f2 --- /dev/null +++ b/code/labs/ozaki_scheme/accuracy_policy.json @@ -0,0 +1,44 @@ +{ + "schema_version": 2, + "policy_id": "ozaki-fp64-emulation-arithmetic-ceilings-v1", + "reference": { + "id": "native-fp64-full-plus-cpu-long-double-edge-v1", + "description": "Every production output is compared with a separate native-FP64 cuBLAS result; small rectangular, cancellation, and dynamic-range cohorts are also compared with a CPU long-double dot-product implementation." + }, + "workload": { + "m": 4096, + "n": 4096, + "k": 4096, + "input_scale": 0.001, + "emulation_strategy": "eager", + "dynamic_max_bits": 16, + "dynamic_offset": -56, + "fixed_bits": 12 + }, + "variants": { + "dynamic": { + "relative_l2": 0.03125, + "normalized_max_abs": 0.0625, + "checksum_rtol": 0.0, + "checksum_atol": 2147.483648, + "rationale": "The -56 offset deliberately gives up the native-FP64 accuracy guarantee. Require at least a five-bit aggregate output-quality floor (2^-5) and a four-bit worst-output floor (2^-4). These are engineering acceptance limits, not a theorem about arbitrary matrices." + }, + "fixed": { + "relative_l2": 0.000244140625, + "normalized_max_abs": 0.00048828125, + "checksum_rtol": 0.0, + "checksum_atol": 16.777216, + "rationale": "Fixed-12 must stay within one 2^-12 input-grid step in aggregate and two steps at the worst output. These are engineering acceptance limits, not a forward-error guarantee for ill-conditioned products." + } + }, + "qualification": { + "required_receipts": [ + {"cohort": "nominal", "seeds": [2026], "m": 4096, "n": 4096, "k": 4096, "input_scale": 0.001, "input_pattern": "uniform", "reference_mode": "native_fp64"}, + {"cohort": "holdout", "seeds": [2027, 2029], "m": 4096, "n": 4096, "k": 4096, "input_scale": 0.001, "input_pattern": "uniform", "reference_mode": "native_fp64"}, + {"cohort": "alternating_edge", "seeds": [2039], "m": 17, "n": 19, "k": 23, "input_scale": 1.0, "input_pattern": "alternating", "reference_mode": "cpu_long_double"}, + {"cohort": "dynamic_range_edge", "seeds": [2053], "m": 31, "n": 29, "k": 37, "input_scale": 1.0, "input_pattern": "dynamic_range", "reference_mode": "cpu_long_double"} + ], + "variants": ["dynamic", "fixed"], + "claim_boundary": "Passing establishes this lab's arithmetic GEMM gate on the recorded stack and cohorts. It does not establish application, solver, or model quality for arbitrary matrices." + } +} diff --git a/code/labs/ozaki_scheme/accuracy_policy.py b/code/labs/ozaki_scheme/accuracy_policy.py index 42ff54565..ae17445d6 100644 --- a/code/labs/ozaki_scheme/accuracy_policy.py +++ b/code/labs/ozaki_scheme/accuracy_policy.py @@ -1,4 +1,4 @@ -"""Explicit, externally reviewed Ozaki bounds; configuration is not calibration evidence.""" +"""Source-defined Ozaki arithmetic bounds; configuration is not run evidence.""" import json import math @@ -6,22 +6,119 @@ from pathlib import Path +POLICY_ID = "ozaki-fp64-emulation-arithmetic-ceilings-v1" +REFERENCE_ID = "native-fp64-full-plus-cpu-long-double-edge-v1" +DEFAULT_POLICY_PATH = Path(__file__).with_name("accuracy_policy.json") +WORKLOAD = { + "m": 4096, + "n": 4096, + "k": 4096, + "input_scale": 0.001, + "emulation_strategy": "eager", + "dynamic_max_bits": 16, + "dynamic_offset": -56, + "fixed_bits": 12, +} +QUALIFICATION_VARIANTS = ["dynamic", "fixed"] +QUALIFICATION_RECEIPTS = [ + {"cohort": "nominal", "seeds": [2026], "m": 4096, "n": 4096, "k": 4096, + "input_scale": 0.001, "input_pattern": "uniform", "reference_mode": "native_fp64"}, + {"cohort": "holdout", "seeds": [2027, 2029], "m": 4096, "n": 4096, "k": 4096, + "input_scale": 0.001, "input_pattern": "uniform", "reference_mode": "native_fp64"}, + {"cohort": "alternating_edge", "seeds": [2039], "m": 17, "n": 19, "k": 23, + "input_scale": 1.0, "input_pattern": "alternating", "reference_mode": "cpu_long_double"}, + {"cohort": "dynamic_range_edge", "seeds": [2053], "m": 31, "n": 29, "k": 37, + "input_scale": 1.0, "input_pattern": "dynamic_range", "reference_mode": "cpu_long_double"}, +] + +# The dynamic policy is an explicit five-bit output-quality floor for the +# deliberately accuracy-reducing -56 offset. It is a conservative engineering +# requirement, not a cuBLAS guarantee. Fixed-12 is held to one 2^-12 grid step +# in aggregate and two steps at the worst output. Checksum tolerances are only +# secondary harness bounds, derived with Cauchy-Schwarz from the full-array L2 +# ceiling and |a|,|b| <= input_scale: limit * m*n*k*input_scale^2. +ENGINEERING_CEILINGS = { + "dynamic": { + "relative_l2": 2.0**-5, + "normalized_max_abs": 2.0**-4, + "checksum_rtol": 0.0, + "checksum_atol": (2.0**-5) * 4096**3 * 0.001**2, + }, + "fixed": { + "relative_l2": 2.0**-12, + "normalized_max_abs": 2.0**-11, + "checksum_rtol": 0.0, + "checksum_atol": (2.0**-12) * 4096**3 * 0.001**2, + }, +} + + +def _limits_from_item(item: dict) -> dict[str, float]: + result = {} + for name in ("relative_l2", "normalized_max_abs", "checksum_rtol", "checksum_atol"): + try: + value = float(item[name]) + except KeyError as exc: + raise ValueError(f"Ozaki accuracy policy missing {name}") from exc + if not math.isfinite(value) or value < 0 or (name != "checksum_atol" and value >= 1): + interval = "nonnegative" if name == "checksum_atol" else "in [0,1)" + raise ValueError(f"{name} must be finite and {interval}") + result[name] = value + return result + + +def load_accuracy_policy(path: Path) -> dict: + policy = json.loads(path.read_text()) + schema_version = policy.get("schema_version") + if schema_version == 1: + # Exact-zero fixture policies remain useful for CPU comparator tests, + # but legacy documents cannot carry a nonzero benchmark acceptance bar. + for variant in ("dynamic", "fixed"): + if variant not in policy: + continue + if any(_limits_from_item(policy[variant]).values()): + raise ValueError("schema_version=1 is permitted only for exact-zero test policies") + return policy + if schema_version != 2: + raise ValueError("Ozaki accuracy policy requires schema_version=2") + if policy.get("policy_id") != POLICY_ID: + raise ValueError(f"Ozaki accuracy policy requires policy_id={POLICY_ID}") + if policy.get("reference", {}).get("id") != REFERENCE_ID: + raise ValueError(f"Ozaki accuracy policy requires reference.id={REFERENCE_ID}") + if policy.get("workload") != WORKLOAD: + raise ValueError("Ozaki accuracy policy workload does not match the benchmark contract") + qualification = policy.get("qualification", {}) + if (qualification.get("variants") != QUALIFICATION_VARIANTS or + qualification.get("required_receipts") != QUALIFICATION_RECEIPTS): + raise ValueError("Ozaki accuracy policy qualification matrix does not match the source contract") + variants = policy.get("variants") + if not isinstance(variants, dict): + raise ValueError("Ozaki accuracy policy requires a variants object") + for variant, ceilings in ENGINEERING_CEILINGS.items(): + if variant not in variants: + raise ValueError(f"Ozaki accuracy policy missing variant {variant}") + limits = _limits_from_item(variants[variant]) + for name, ceiling in ceilings.items(): + if limits[name] > ceiling: + raise ValueError( + f"{variant}.{name} exceeds the source-defined engineering ceiling {ceiling:.8g}" + ) + return policy + + def configured_accuracy(variant: str) -> tuple[list[str], tuple[float, float]]: path = os.environ.get("AISP_OZAKI_ACCURACY_POLICY") if not path: # The binary rejects emulation without bounds before allocating/running. return [], (0.0, 0.0) - policy = json.loads(Path(path).read_text()) - if policy.get("schema_version") != 1: - raise ValueError("Ozaki accuracy policy requires schema_version=1") - item = policy[variant] - for name in ("relative_l2", "normalized_max_abs", "checksum_rtol"): - value = float(item[name]) - if not math.isfinite(value) or not 0 <= value < 1: - raise ValueError(f"{name} must be finite and in [0,1)") - atol = float(item["checksum_atol"]) - if not math.isfinite(atol) or atol < 0: - raise ValueError("checksum_atol must be finite and nonnegative") - return (["--relative-l2-limit", str(item["relative_l2"]), - "--normalized-max-abs-limit", str(item["normalized_max_abs"])], - (float(item["checksum_rtol"]), atol)) + policy = load_accuracy_policy(Path(path)) + try: + item = policy[variant] if policy["schema_version"] == 1 else policy["variants"][variant] + except KeyError as exc: + raise ValueError(f"Ozaki accuracy policy missing variant {variant}") from exc + limits = _limits_from_item(item) + relative_arg = item["relative_l2"] if policy["schema_version"] == 1 else limits["relative_l2"] + normalized_arg = item["normalized_max_abs"] if policy["schema_version"] == 1 else limits["normalized_max_abs"] + return (["--relative-l2-limit", str(relative_arg), + "--normalized-max-abs-limit", str(normalized_arg)], + (limits["checksum_rtol"], limits["checksum_atol"])) diff --git a/code/labs/ozaki_scheme/lab_utils.py b/code/labs/ozaki_scheme/lab_utils.py index b130ef4c4..348a3d946 100644 --- a/code/labs/ozaki_scheme/lab_utils.py +++ b/code/labs/ozaki_scheme/lab_utils.py @@ -19,6 +19,21 @@ _METRIC_PATTERNS: Final[dict[str, re.Pattern[str]]] = { "variant": re.compile(r"VARIANT:\s*(\S+)"), + "m": re.compile(r"M:\s*(\d+)"), + "n": re.compile(r"N:\s*(\d+)"), + "k": re.compile(r"K:\s*(\d+)"), + "seed": re.compile(r"SEED:\s*(-?\d+)"), + "input_scale": re.compile(r"INPUT_SCALE:\s*([0-9.eE+-]+)"), + "input_pattern": re.compile(r"INPUT_PATTERN:\s*(\S+)"), + "reference_mode": re.compile(r"REFERENCE_MODE:\s*(\S+)"), + "gpu_name": re.compile(r"GPU_NAME:\s*(.+)"), + "compute_capability": re.compile(r"COMPUTE_CAPABILITY:\s*(\S+)"), + "cuda_runtime_version": re.compile(r"CUDA_RUNTIME_VERSION:\s*(\d+)"), + "cublas_version": re.compile(r"CUBLAS_VERSION:\s*(\d+)"), + "accuracy_status": re.compile(r"ACCURACY_STATUS:\s*(\S+)"), + "dynamic_max_bits": re.compile(r"DYNAMIC_MAX_BITS:\s*(\d+)"), + "dynamic_offset": re.compile(r"DYNAMIC_OFFSET:\s*(-?\d+)"), + "fixed_bits": re.compile(r"FIXED_BITS:\s*(\d+)"), "time_ms": re.compile(r"TIME_MS:\s*([0-9.eE+-]+)"), "tflops": re.compile(r"TFLOPS:\s*([0-9.eE+-]+)"), "retained_bits": re.compile(r"RETAINED_BITS:\s*(-?\d+)"), @@ -43,9 +58,12 @@ def parse_metrics(stdout: str) -> dict[str, MetricValue]: if not match: continue value = match.group(1) - if key in {"variant", "emulation_strategy"}: + if key in {"variant", "emulation_strategy", "input_pattern", "reference_mode", "accuracy_status", + "gpu_name", "compute_capability"}: metrics[key] = value - elif key in {"retained_bits", "emulation_used"}: + elif key in {"m", "n", "k", "seed", "dynamic_max_bits", "dynamic_offset", "fixed_bits", + "cuda_runtime_version", "cublas_version", + "retained_bits", "emulation_used"}: metrics[key] = int(value) else: metrics[key] = float(value) diff --git a/code/labs/ozaki_scheme/ozaki_scheme_common.cuh b/code/labs/ozaki_scheme/ozaki_scheme_common.cuh index b473faac2..a8a615258 100644 --- a/code/labs/ozaki_scheme/ozaki_scheme_common.cuh +++ b/code/labs/ozaki_scheme/ozaki_scheme_common.cuh @@ -34,6 +34,17 @@ enum class EmulationStrategy { kEager, }; +enum class ReferenceMode { + kNativeFp64, + kCpuLongDouble, +}; + +enum class InputPattern { + kUniform, + kAlternating, + kDynamicRange, +}; + struct Options { int m = 4096; int n = 4096; @@ -50,6 +61,8 @@ struct Options { bool accuracy_measure_only = false; std::size_t workspace_bytes = 64ull << 20; EmulationStrategy emulation_strategy = EmulationStrategy::kEager; + ReferenceMode reference_mode = ReferenceMode::kNativeFp64; + InputPattern input_pattern = InputPattern::kUniform; }; struct Metrics { @@ -62,6 +75,11 @@ struct Metrics { double normalized_max_abs_error = 0.0; int retained_bits = -1; int emulation_used = 0; + int cuda_runtime_version = 0; + int cublas_version = 0; + int compute_capability_major = 0; + int compute_capability_minor = 0; + std::string gpu_name; }; inline const char* variant_name(Variant variant) { @@ -88,6 +106,28 @@ inline const char* emulation_strategy_name(EmulationStrategy strategy) { return "unknown"; } +inline const char* reference_mode_name(ReferenceMode mode) { + switch (mode) { + case ReferenceMode::kNativeFp64: + return "native_fp64"; + case ReferenceMode::kCpuLongDouble: + return "cpu_long_double"; + } + return "unknown"; +} + +inline const char* input_pattern_name(InputPattern pattern) { + switch (pattern) { + case InputPattern::kUniform: + return "uniform"; + case InputPattern::kAlternating: + return "alternating"; + case InputPattern::kDynamicRange: + return "dynamic_range"; + } + return "unknown"; +} + inline void check_cuda(cudaError_t status, const char* expr, const char* file, int line) { if (status != cudaSuccess) { std::ostringstream oss; @@ -147,6 +187,8 @@ inline void print_usage(const char* program) { << " --iters Timed matmuls averaged into TIME_MS (default 10)\n" << " --seed RNG seed for deterministic inputs (default 2026)\n" << " --input-scale Uniform input scale (default 0.001)\n" + << " --input-pattern uniform|alternating|dynamic_range (default uniform)\n" + << " --reference-mode native_fp64|cpu_long_double (default native_fp64)\n" << " --dynamic-max-bits Max retained bits for dynamic Ozaki (default 16)\n" << " --dynamic-offset Dynamic mantissa bias (default -56)\n" << " --fixed-bits Retained bits for fixed Ozaki (default 12)\n" @@ -172,6 +214,31 @@ inline EmulationStrategy parse_emulation_strategy(const std::string& raw) { " (expected default|performant|eager)"); } +inline ReferenceMode parse_reference_mode(const std::string& raw) { + if (raw == "native_fp64") { + return ReferenceMode::kNativeFp64; + } + if (raw == "cpu_long_double") { + return ReferenceMode::kCpuLongDouble; + } + throw std::runtime_error(std::string("Invalid value for --reference-mode: ") + raw + + " (expected native_fp64|cpu_long_double)"); +} + +inline InputPattern parse_input_pattern(const std::string& raw) { + if (raw == "uniform") { + return InputPattern::kUniform; + } + if (raw == "alternating") { + return InputPattern::kAlternating; + } + if (raw == "dynamic_range") { + return InputPattern::kDynamicRange; + } + throw std::runtime_error(std::string("Invalid value for --input-pattern: ") + raw + + " (expected uniform|alternating|dynamic_range)"); +} + inline Options parse_args(int argc, char** argv) { Options options; for (int i = 1; i < argc; ++i) { @@ -208,6 +275,10 @@ inline Options parse_args(int argc, char** argv) { parse_numeric_arg(value, &options.fixed_bits, "--fixed-bits"); } else if (arg == "--emulation-strategy") { options.emulation_strategy = parse_emulation_strategy(value); + } else if (arg == "--reference-mode") { + options.reference_mode = parse_reference_mode(value); + } else if (arg == "--input-pattern") { + options.input_pattern = parse_input_pattern(value); } else if (arg == "--input-scale") { parse_numeric_arg(value, &options.input_scale, "--input-scale"); } else if (arg == "--relative-l2-limit") { @@ -368,11 +439,39 @@ inline void launch_matmul( CUBLAS_GEMM_DEFAULT_TENSOR_OP)); } -inline void fill_host_matrix(std::vector* data, int seed, double scale) { +inline void fill_host_matrix( + std::vector* data, + std::size_t rows, + std::size_t columns, + int seed, + double scale, + InputPattern pattern, + bool right_operand) { + if (data->size() != rows * columns) { + throw std::runtime_error("Input pattern shape does not match allocated matrix"); + } std::mt19937_64 rng(static_cast(seed)); std::uniform_real_distribution dist(-scale, scale); - for (double& value : *data) { - value = dist(rng); + if (pattern == InputPattern::kUniform) { + for (double& value : *data) { + value = dist(rng); + } + return; + } + for (std::size_t row = 0; row < rows; ++row) { + for (std::size_t column = 0; column < columns; ++column) { + const std::size_t index = row * columns + column; + const bool negative = ((row + column + static_cast(seed) + + (right_operand ? 1u : 0u)) & 1u) != 0; + double magnitude = scale; + if (pattern == InputPattern::kAlternating) { + magnitude *= 1.0 - static_cast((index + static_cast(seed)) % 17) / 64.0; + } else { + const int exponent = -static_cast((index + static_cast(seed)) % 13); + magnitude = std::ldexp(scale, exponent); + } + (*data)[index] = negative ? -magnitude : magnitude; + } } } @@ -399,8 +498,10 @@ inline Metrics benchmark_variant(Variant variant, const Options& options) { std::vector h_a(a_elements); std::vector h_b(b_elements); - fill_host_matrix(&h_a, options.seed, options.input_scale); - fill_host_matrix(&h_b, options.seed + 17, options.input_scale); + fill_host_matrix(&h_a, options.m, options.k, options.seed, options.input_scale, + options.input_pattern, false); + fill_host_matrix(&h_b, options.k, options.n, options.seed + 17, options.input_scale, + options.input_pattern, true); double* d_a = nullptr; double* d_b = nullptr; @@ -426,7 +527,7 @@ inline Metrics benchmark_variant(Variant variant, const Options& options) { state = create_handle_state(variant, options, stream, workspace); - if (variant != Variant::kNative) { + if (variant != Variant::kNative && options.reference_mode == ReferenceMode::kNativeFp64) { OZAKI_CHECK_CUDA(cudaMalloc(&d_ref, c_bytes)); ref_state = create_handle_state(Variant::kNative, options, stream, workspace); launch_matmul( @@ -457,6 +558,11 @@ inline Metrics benchmark_variant(Variant variant, const Options& options) { OZAKI_CHECK_CUDA(cudaEventElapsedTime(&elapsed_ms, start, stop)); Metrics metrics; + metrics.gpu_name = props.name; + metrics.compute_capability_major = props.major; + metrics.compute_capability_minor = props.minor; + OZAKI_CHECK_CUDA(cudaRuntimeGetVersion(&metrics.cuda_runtime_version)); + OZAKI_CHECK_CUBLAS(cublasGetVersion(state.handle, &metrics.cublas_version)); metrics.time_ms = static_cast(elapsed_ms) / static_cast(options.iters); metrics.tflops = (2.0 * static_cast(options.m) * options.n * options.k) / (metrics.time_ms * 1.0e9); @@ -474,8 +580,14 @@ inline Metrics benchmark_variant(Variant variant, const Options& options) { } if (variant != Variant::kNative) { - std::vector h_ref(c_elements); - OZAKI_CHECK_CUDA(cudaMemcpy(h_ref.data(), d_ref, c_bytes, cudaMemcpyDeviceToHost)); + std::vector h_ref; + if (options.reference_mode == ReferenceMode::kNativeFp64) { + h_ref.resize(c_elements); + OZAKI_CHECK_CUDA(cudaMemcpy(h_ref.data(), d_ref, c_bytes, cudaMemcpyDeviceToHost)); + } else { + h_ref = reference_gemm_long_double( + h_a.data(), h_b.data(), options.m, options.n, options.k); + } const AccuracyMetrics accuracy = measure_accuracy(h_c.data(), h_ref.data(), c_elements); metrics.max_abs_error = accuracy.max_abs_error; metrics.mean_abs_error = accuracy.mean_abs_error; @@ -513,6 +625,15 @@ inline void print_metrics(Variant variant, const Options& options, const Metrics std::cout << "K: " << options.k << "\n"; std::cout << "WARMUP: " << options.warmup << "\n"; std::cout << "ITERS: " << options.iters << "\n"; + std::cout << "SEED: " << options.seed << "\n"; + std::cout << "INPUT_SCALE: " << options.input_scale << "\n"; + std::cout << "INPUT_PATTERN: " << input_pattern_name(options.input_pattern) << "\n"; + std::cout << "REFERENCE_MODE: " << reference_mode_name(options.reference_mode) << "\n"; + std::cout << "GPU_NAME: " << metrics.gpu_name << "\n"; + std::cout << "COMPUTE_CAPABILITY: " << metrics.compute_capability_major << "." + << metrics.compute_capability_minor << "\n"; + std::cout << "CUDA_RUNTIME_VERSION: " << metrics.cuda_runtime_version << "\n"; + std::cout << "CUBLAS_VERSION: " << metrics.cublas_version << "\n"; if (variant == Variant::kDynamic) { std::cout << "DYNAMIC_MAX_BITS: " << options.dynamic_max_bits << "\n"; std::cout << "DYNAMIC_OFFSET: " << options.dynamic_offset << "\n"; diff --git a/code/labs/ozaki_scheme/qualify_accuracy.py b/code/labs/ozaki_scheme/qualify_accuracy.py new file mode 100644 index 000000000..4b1396cf8 --- /dev/null +++ b/code/labs/ozaki_scheme/qualify_accuracy.py @@ -0,0 +1,134 @@ +"""Qualify retained Ozaki measurement logs against the predeclared policy. + +The CUDA executables intentionally return 2 for measurement-only runs. Capture +their stdout as individual logs, then use this CPU-only command to require the +complete nominal, holdout, and independent long-double edge matrix. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +from labs.ozaki_scheme.accuracy_policy import ( + DEFAULT_POLICY_PATH, + WORKLOAD, + _limits_from_item, + load_accuracy_policy, +) +from labs.ozaki_scheme.lab_utils import parse_metrics + + +def _required_cases(policy: dict) -> dict[tuple[str, str, int], dict]: + cases = {} + for variant in policy["qualification"]["variants"]: + for group in policy["qualification"]["required_receipts"]: + for seed in group["seeds"]: + key = (variant, group["cohort"], int(seed)) + cases[key] = {**group, "variant": variant, "seed": int(seed)} + return cases + + +def _matching_case(required: dict[tuple[str, str, int], dict], metrics: dict) -> tuple[str, str, int] | None: + variant = str(metrics.get("variant", "")).removeprefix("ozaki_") + for key, case in required.items(): + if key[0] != variant or key[2] != metrics.get("seed"): + continue + if all(metrics.get(name) == case[name] for name in ( + "m", "n", "k", "input_scale", "input_pattern", "reference_mode" + )): + return key + return None + + +def assess_logs(policy: dict, log_texts: list[str]) -> dict: + required = _required_cases(policy) + seen: set[tuple[str, str, int]] = set() + failures: list[str] = [] + results: list[dict] = [] + expected_provenance: tuple | None = None + + for index, text in enumerate(log_texts): + metrics = parse_metrics(text) + key = _matching_case(required, metrics) + reasons: list[str] = [] + if key is None: + reasons.append("log does not match a required policy case") + elif key in seen: + reasons.append("duplicate policy case") + if metrics.get("accuracy_status") != "MEASUREMENT_ONLY_NOT_ACCEPTED": + reasons.append("log is not retained measurement-only evidence") + if metrics.get("emulation_used") != 1 or int(metrics.get("retained_bits", -1)) < 0: + reasons.append("cuBLAS did not report active fixed-point emulation") + provenance = tuple(metrics.get(name) for name in ( + "gpu_name", "compute_capability", "cuda_runtime_version", "cublas_version" + )) + if any(value in (None, "") for value in provenance): + reasons.append("GPU/CUDA/cuBLAS provenance is incomplete") + elif expected_provenance is None: + expected_provenance = provenance + elif provenance != expected_provenance: + reasons.append("GPU/CUDA/cuBLAS provenance differs across logs") + variant = str(metrics.get("variant", "")).removeprefix("ozaki_") + if variant in policy["variants"]: + limits = _limits_from_item(policy["variants"][variant]) + for metric_name, limit_name in ( + ("relative_l2_error", "relative_l2"), + ("normalized_max_abs_error", "normalized_max_abs"), + ): + value = metrics.get(metric_name) + if not isinstance(value, (int, float)) or not math.isfinite(float(value)): + reasons.append(f"{metric_name} is missing or non-finite") + elif float(value) > limits[limit_name]: + reasons.append(f"{metric_name}={float(value):.8g} exceeds {limits[limit_name]:.8g}") + if variant == "dynamic" and ( + metrics.get("dynamic_max_bits") != WORKLOAD["dynamic_max_bits"] or + metrics.get("dynamic_offset") != WORKLOAD["dynamic_offset"] + ): + reasons.append("dynamic mantissa configuration mismatch") + if variant == "fixed" and metrics.get("fixed_bits") != WORKLOAD["fixed_bits"]: + reasons.append("fixed mantissa configuration mismatch") + if metrics.get("emulation_strategy") != WORKLOAD["emulation_strategy"]: + reasons.append("emulation strategy mismatch") + else: + reasons.append("unknown Ozaki variant") + if key is not None and key not in seen: + seen.add(key) + if reasons: + failures.extend(f"log[{index}] {key}: {reason}" for reason in reasons) + results.append({"case": list(key) if key else None, "passed": not reasons, + "reasons": reasons, "metrics": metrics}) + + for key in sorted(set(required) - seen): + failures.append(f"missing required log: {key}") + return { + "schema_version": 1, + "policy_id": policy["policy_id"], + "status": "qualified_arithmetic_gate" if not failures else "failed_arithmetic_gate", + "required_case_count": len(required), + "passing_case_count": sum(item["passed"] for item in results), + "failures": failures, + "logs": results, + "claim_boundary": policy["qualification"]["claim_boundary"], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("logs", type=Path, nargs="+") + parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY_PATH) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + policy = load_accuracy_policy(args.policy) + if policy.get("schema_version") != 2: + raise ValueError("Log qualification requires the reviewed schema_version=2 policy") + summary = assess_logs(policy, [path.read_text() for path in args.logs]) + args.output.write_text(json.dumps(summary, indent=2) + "\n") + if summary["status"] != "qualified_arithmetic_gate": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/code/tests/test_accuracy_requirements_followthrough.py b/code/tests/test_accuracy_requirements_followthrough.py new file mode 100644 index 000000000..15f743ae1 --- /dev/null +++ b/code/tests/test_accuracy_requirements_followthrough.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +import json +import shutil +import subprocess +import sys +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +def test_kv_checked_in_policy_uses_format_ceiling_and_rejects_widening(tmp_path: Path) -> None: + from labs.kv_cache_compression.accuracy import ( + DEFAULT_POLICY_PATH, + ENGINEERING_CEILINGS, + load_accuracy_policy, + ) + + policy = load_accuracy_policy(DEFAULT_POLICY_PATH) + assert policy["variants"]["fp8"]["relative_l2"] == 2**-4 + assert policy["variants"]["nvfp4"]["relative_l2"] == 2**-2 + + widened = deepcopy(policy) + widened["variants"]["fp8"]["relative_l2"] = ENGINEERING_CEILINGS["fp8"].relative_l2 + 0.001 + path = tmp_path / "widened.json" + path.write_text(json.dumps(widened)) + with pytest.raises(ValueError, match="exceeds the source-defined engineering ceiling"): + load_accuracy_policy(path) + + +@pytest.mark.parametrize("variant,corruption", [("fp8", 0.07), ("nvfp4", 0.26)]) +def test_kv_format_requirements_reject_meaningful_full_cache_corruption( + variant: str, corruption: float +) -> None: + from labs.kv_cache_compression.accuracy import ENGINEERING_CEILINGS, assert_cache_accuracy + from labs.kv_cache_compression.kv_cache_common import KVCache + + reference = KVCache(torch.ones(2, 4, 2, 8), -torch.ones(2, 4, 2, 8)) + actual = KVCache(reference.cache_k.clone() + corruption, reference.cache_v.clone() - corruption) + with pytest.raises(AssertionError, match="KV cache accuracy failed"): + assert_cache_accuracy(actual, reference, ENGINEERING_CEILINGS[variant]) + + +def test_kv_edge_cohorts_preserve_shape_and_create_stress_patterns() -> None: + from labs.kv_cache_compression.calibrate_accuracy import _apply_input_cohort + + alternating = SimpleNamespace( + hidden_dim=8, + tensor_dtype=torch.bfloat16, + device=torch.device("cpu"), + prefill_inputs=[torch.empty(2, 3, 8, dtype=torch.bfloat16)], + decode_inputs=[torch.empty(2, 2, 8, dtype=torch.bfloat16)], + ) + shapes = [tensor.shape for tensor in alternating.prefill_inputs + alternating.decode_inputs] + _apply_input_cohort(alternating, "alternating") + assert [tensor.shape for tensor in alternating.prefill_inputs + alternating.decode_inputs] == shapes + torch.testing.assert_close( + alternating.prefill_inputs[0][0, 0], + -alternating.prefill_inputs[0][0, 1], + rtol=0, + atol=0, + ) + + _apply_input_cohort(alternating, "sparse_outlier") + for tensor in alternating.prefill_inputs + alternating.decode_inputs: + assert torch.count_nonzero(tensor[..., 2:]) == 0 + assert torch.all(tensor[..., 0] == 1) + assert torch.all(tensor[..., 1] == -1) + + +def test_kv_calibration_retains_unsupported_host_failure(tmp_path: Path) -> None: + if torch.cuda.is_available(): + pytest.skip("This control exercises the unsupported CPU-host receipt") + output = tmp_path / "failure.json" + completed = subprocess.run( + [ + sys.executable, + "-m", + "labs.kv_cache_compression.calibrate_accuracy", + "--variant", + "fp8", + "--cohort", + "nominal", + "--seed", + "2026", + "--output", + str(output), + ], + cwd=Path(__file__).parents[1], + capture_output=True, + text=True, + ) + assert completed.returncode != 0 + receipt = json.loads(output.read_text()) + assert receipt["status"] == "failure_not_accepted" + assert receipt["error_type"] == "RuntimeError" + assert "actual CUDA" in receipt["error"] + + +def _kv_receipts(policy: dict) -> list[dict]: + from labs.kv_cache_compression.accuracy import REFERENCE_ID, WORKLOAD + from labs.kv_cache_compression.qualify_accuracy import required_cases + + workload = {key: value for key, value in WORKLOAD.items() if key != "storage_dtype"} + return [{ + "schema_version": 2, + "status": "measurement_only_not_accepted", + "variant": variant, + "cohort": cohort, + "seed": seed, + "reference_id": REFERENCE_ID, + "git_commit": "0123456789abcdef", + "torch": "2.9.1", + "cuda": "13.0", + "transformer_engine": "2.18.0", + "gpu": "NVIDIA B200", + "compute_capability": [10, 0], + "workload": workload, + "metrics": { + "cache_k.relative_l2": 0.0, + "cache_k.normalized_max_abs": 0.0, + "cache_v.relative_l2": 0.0, + "cache_v.normalized_max_abs": 0.0, + }, + } for variant, cohort, seed in sorted(required_cases(policy))] + + +def test_kv_qualification_requires_all_holdouts_and_retains_failure() -> None: + from labs.kv_cache_compression.accuracy import DEFAULT_POLICY_PATH, load_accuracy_policy + from labs.kv_cache_compression.qualify_accuracy import assess_receipts + + policy = load_accuracy_policy(DEFAULT_POLICY_PATH) + receipts = _kv_receipts(policy) + assert assess_receipts(policy, receipts)["status"] == "qualified_arithmetic_gate" + + damaged = deepcopy(receipts) + damaged[-1]["metrics"]["cache_v.normalized_max_abs"] = 0.5 + result = assess_receipts(policy, damaged) + assert result["status"] == "failed_arithmetic_gate" + assert any("cache_v.normalized_max_abs" in failure for failure in result["failures"]) + + +def test_ozaki_checked_in_policy_rejects_widening(tmp_path: Path) -> None: + from labs.ozaki_scheme.accuracy_policy import DEFAULT_POLICY_PATH, load_accuracy_policy + + policy = load_accuracy_policy(DEFAULT_POLICY_PATH) + assert policy["variants"]["dynamic"]["relative_l2"] == 2**-5 + assert policy["variants"]["fixed"]["relative_l2"] == 2**-12 + widened = deepcopy(policy) + widened["variants"]["fixed"]["normalized_max_abs"] = 0.001 + path = tmp_path / "widened.json" + path.write_text(json.dumps(widened)) + with pytest.raises(ValueError, match="exceeds the source-defined engineering ceiling"): + load_accuracy_policy(path) + + +@pytest.fixture(scope="module") +def ozaki_long_double_reference_probe(tmp_path_factory: pytest.TempPathFactory) -> Path: + compiler = shutil.which("c++") + if not compiler: + pytest.skip("C++ compiler unavailable") + directory = tmp_path_factory.mktemp("ozaki_long_double_reference") + source = directory / "probe.cpp" + source.write_text(r''' +#include "accuracy.h" +#include +#include +int main() { + const double a[] = {1, 2, 3, 4, 5, 6}; + const double b[] = {7, 8, 9, 10, 11, 12}; + const double expected[] = {58, 64, 139, 154}; + std::vector reference; + try { + reference = ozaki_scheme::reference_gemm_long_double(a, b, 2, 2, 3); + } catch (const std::exception& error) { + std::cout << error.what() << "\n"; + return 4; + } + const auto exact = ozaki_scheme::measure_accuracy(reference.data(), expected, 4); + if (exact.relative_l2 != 0 || exact.normalized_max_abs != 0) return 2; + auto corrupt = reference; + corrupt.back() += 1; + const auto damaged = ozaki_scheme::measure_accuracy(corrupt.data(), reference.data(), 4); + try { + ozaki_scheme::assert_accuracy(damaged, 1.0 / 4096.0, 1.0 / 2048.0); + } catch (const std::exception&) { + std::cout << "corruption rejected\n"; + return 0; + } + return 3; +} +''') + completed = subprocess.run( + [compiler, "-std=c++17", "-Wall", "-Wextra", "-pedantic", "-I", + str(Path(__file__).parents[1] / "labs" / "ozaki_scheme"), str(source), "-o", str(directory / "probe")], + capture_output=True, + text=True, + ) + assert completed.returncode == 0, completed.stderr + return directory / "probe" + + +def test_ozaki_cpu_long_double_reference_and_corruption_control( + ozaki_long_double_reference_probe: Path, +) -> None: + completed = subprocess.run([str(ozaki_long_double_reference_probe)], capture_output=True, text=True) + assert completed.returncode in (0, 4), completed.stdout + completed.stderr + if completed.returncode == 0: + assert completed.stdout.strip() == "corruption rejected" + else: + # arm64 Darwin aliases long double to FP64. The reference must fail + # closed there; x86_64 B200 qualification executes the wider oracle. + assert completed.stdout.strip() == "CPU long-double reference requires precision wider than FP64" + + +def _ozaki_log(case: dict, *, relative_l2: float = 0.0) -> str: + variant = case["variant"] + knob = ( + "DYNAMIC_MAX_BITS: 16\nDYNAMIC_OFFSET: -56" + if variant == "dynamic" + else "FIXED_BITS: 12" + ) + return f"""VARIANT: ozaki_{variant} +M: {case['m']} +N: {case['n']} +K: {case['k']} +SEED: {case['seed']} +INPUT_SCALE: {case['input_scale']} +INPUT_PATTERN: {case['input_pattern']} +REFERENCE_MODE: {case['reference_mode']} +GPU_NAME: NVIDIA B200 +COMPUTE_CAPABILITY: 10.0 +CUDA_RUNTIME_VERSION: 13000 +CUBLAS_VERSION: 130000 +{knob} +EMULATION_STRATEGY: eager +EMULATION_USED: 1 +RETAINED_BITS: 12 +RELATIVE_L2_ERROR: {relative_l2} +NORMALIZED_MAX_ABS_ERROR: 0 +ACCURACY_STATUS: MEASUREMENT_ONLY_NOT_ACCEPTED +""" + + +def test_ozaki_qualification_requires_independent_edges_and_rejects_corruption() -> None: + from labs.ozaki_scheme.accuracy_policy import DEFAULT_POLICY_PATH, load_accuracy_policy + from labs.ozaki_scheme.qualify_accuracy import _required_cases, assess_logs + + policy = load_accuracy_policy(DEFAULT_POLICY_PATH) + cases = list(_required_cases(policy).values()) + logs = [_ozaki_log(case) for case in cases] + assert assess_logs(policy, logs)["status"] == "qualified_arithmetic_gate" + + logs[-1] = _ozaki_log(cases[-1], relative_l2=0.1) + result = assess_logs(policy, logs) + assert result["status"] == "failed_arithmetic_gate" + assert any("relative_l2_error" in failure for failure in result["failures"]) diff --git a/code/tests/test_refresh_readmes.py b/code/tests/test_refresh_readmes.py index afae4b7e0..bc02c20b8 100644 --- a/code/tests/test_refresh_readmes.py +++ b/code/tests/test_refresh_readmes.py @@ -210,9 +210,12 @@ def test_ch10_and_priority_labs_render_custom_evidence_sections() -> None: _assert_evidence_sections(markdown) assert "## Storage and workload" in kv_cache_compression_markdown - assert "## Accuracy gate: target calibration pending" in kv_cache_compression_markdown + assert "## Accuracy gate: requirements defined, target qualification pending" in kv_cache_compression_markdown assert "neither path compresses the KV cache" in kv_cache_compression_markdown assert "AISP_KV_CACHE_ACCURACY_POLICY" in kv_cache_compression_markdown + assert "2^-4" in kv_cache_compression_markdown + assert "2^-2" in kv_cache_compression_markdown + assert "ACCURACY_REQUIREMENTS.md" in kv_cache_compression_markdown def test_all_generator_readmes_match_generated_content() -> None: From 0de2f31d723fd786036253e3430470c238de94ad Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 02:31:52 -0700 Subject: [PATCH 04/19] fix: keep minimal Nsight captures to the requested five metrics --- code/cli/aisp.py | 2 +- code/core/profiling/nsight_automation.py | 22 +++++++++++++--------- code/core/profiling/profiler_config.py | 9 ++++++--- code/docs/api-reference.md | 2 +- code/mcp/mcp_server.py | 2 +- code/tests/test_cli_profile_ncu.py | 4 ++-- code/tests/test_ncu_app_range_commands.py | 3 ++- code/tests/test_ncu_command_builder.py | 15 +++++++++++---- 8 files changed, 37 insertions(+), 22 deletions(-) diff --git a/code/cli/aisp.py b/code/cli/aisp.py index 3980f72a4..2d9a6ae5e 100644 --- a/code/cli/aisp.py +++ b/code/cli/aisp.py @@ -696,7 +696,7 @@ def profile_ncu( metric_set: str = typer.Option( "full", "--metric-set", - help="NCU set: full, roofline, minimal, speed-of-light, basic", + help="NCU metrics: minimal selects exactly five; full, roofline, speed-of-light, basic select NVIDIA sections", show_default=True, ), replay_mode: str = typer.Option( diff --git a/code/core/profiling/nsight_automation.py b/code/core/profiling/nsight_automation.py index 606ff7dd0..a3d63506c 100644 --- a/code/core/profiling/nsight_automation.py +++ b/code/core/profiling/nsight_automation.py @@ -189,16 +189,18 @@ def _available_ncu_sets(self) -> set[str]: return set() def _resolve_ncu_set(self, metric_set: str) -> str: - """Resolve user-facing metric-set aliases to an installed NCU --set value.""" + """Resolve section aliases, or the exact repository minimal metric list.""" metric_set_norm = str(metric_set or "").strip().lower() + if metric_set_norm == "minimal": + # A section set adds hundreds of counters on recent NCU releases, + # even when --metrics is also supplied. Minimal means five metrics. + self._last_resolved_ncu_set = "minimal" + return "minimal" alias_candidates = { "full": ["full"], "roofline": ["roofline"], # Nsight versions vary: some expose speed-of-light, others expose basic. "speed-of-light": ["speed-of-light", "basic"], - # Prefer `basic` first for minimal runs; it is substantially lower - # overhead while still providing SpeedOfLight-derived signals. - "minimal": ["basic", "speed-of-light"], "basic": ["basic", "speed-of-light"], } if metric_set_norm not in alias_candidates: @@ -772,6 +774,8 @@ def build_ncu_command( raise ValueError(f"Unsupported workload_type: {workload_type}") metrics = self.METRIC_SETS[workload_type] ncu_set = self._resolve_ncu_set(metric_set) + if ncu_set == 'minimal': + metrics = MINIMAL_METRICS replay_mode = validate_ncu_replay_mode(replay_mode) nvtx_filters = ( list(dict.fromkeys(str(tag).strip() for tag in nvtx_includes or [] if str(tag).strip())) @@ -779,7 +783,7 @@ def build_ncu_command( else self._normalize_nvtx_includes(nvtx_includes) ) if replay_mode == 'app-range': - if ncu_set != 'basic': + if ncu_set not in {'minimal', 'basic'}: raise ValueError("app-range requires metric_set='minimal' or 'basic'.") metrics = MINIMAL_METRICS validate_ncu_app_range_capture( @@ -792,9 +796,9 @@ def build_ncu_command( profile_from_start=profile_from_start, ) ncu_cmd = ['ncu'] - # --set adds section metrics beyond --metrics. Keep app-range at its - # exact validated metric list so collectives do not need extra replays. - if replay_mode != 'app-range': + # --set adds section metrics beyond --metrics. Keep minimal and + # app-range at the exact requested list to bound replay overhead. + if replay_mode != 'app-range' and ncu_set != 'minimal': ncu_cmd.extend(['--set', ncu_set]) ncu_cmd.extend([ '--target-processes', 'all', @@ -804,7 +808,7 @@ def build_ncu_command( if replay_mode: ncu_cmd.extend(['--replay-mode', replay_mode]) # Only add custom metrics when using the full set; other sets bring their own. - if metrics and (ncu_set == 'full' or replay_mode == 'app-range'): + if metrics and (ncu_set in {'full', 'minimal'} or replay_mode == 'app-range'): ncu_cmd.extend(['--metrics', ",".join(metrics)]) if kernel_filter: if kernel_name_base: diff --git a/code/core/profiling/profiler_config.py b/code/core/profiling/profiler_config.py index ccaa60291..881f02358 100644 --- a/code/core/profiling/profiler_config.py +++ b/code/core/profiling/profiler_config.py @@ -585,9 +585,12 @@ def get_ncu_command_for_target( "--clock-control", "none", ] - # Section sets add metrics to --metrics. Application-range replay must - # collect exactly its validated five metrics, avoiding extra NCCL passes. - if requested_replay_mode != "app-range": + # Section sets add metrics to --metrics. Minimal captures must retain + # their requested small list, rather than silently adding basic sections. + explicit_minimal_metrics = metric_set_norm == "minimal" or ( + metric_set_norm == "auto" and ncu_set == "basic" + ) + if requested_replay_mode != "app-range" and not explicit_minimal_metrics: cmd.extend(["--set", ncu_set]) cmd.extend(["--metrics", ",".join(metrics)]) diff --git a/code/docs/api-reference.md b/code/docs/api-reference.md index b8f7a6c6d..cd2aa55c5 100644 --- a/code/docs/api-reference.md +++ b/code/docs/api-reference.md @@ -168,7 +168,7 @@ Profiling with Nsight Systems, Nsight Compute, and torch.profiler. **MCP profiling captures include metrics JSON:** `profile_nsys` returns `nsys_metrics`, `profile_ncu` returns `ncu_metrics`, `profile_torch` returns `torch_metrics` (and `report` alias), and `profile_hta` includes `nsys_metrics`. Use these payloads to analyze regressions and bottleneck shifts. -**Targeted NCU capture (CLI + MCP):** `profile_ncu` supports kernel scoping (`kernel_filter`, optional `kernel_name_base`) plus NVTX gating (`nvtx_include`, `profile_from_start='off'`) to isolate specific kernels and avoid setup-noise captures. Captures now fail loudly when NCU profiles zero kernels or collects zero metrics, and `metric_set='minimal'` auto-resolves to `speed-of-light` or `basic` depending on Nsight Compute version. `compare_ncu` flags rank-only kernel symbol alignment as low-confidence for tuning (advisory), applies kernel-family alias matching, and returns a concrete `staged_pair_dir` for stable follow-up diffs. +**Targeted NCU capture (CLI + MCP):** `profile_ncu` supports kernel scoping (`kernel_filter`, optional `kernel_name_base`) plus NVTX gating (`nvtx_include`, `profile_from_start='off'`) to isolate specific kernels and avoid setup-noise captures. Captures fail loudly when NCU profiles zero kernels or collects zero metrics. `metric_set='minimal'` requests exactly the repository's five metrics without an additional NVIDIA section set, keeping replay overhead bounded; choose `basic` explicitly for NVIDIA's broader basic sections. `compare_ncu` flags rank-only kernel symbol alignment as low-confidence for tuning (advisory), applies kernel-family alias matching, and returns a concrete `staged_pair_dir` for stable follow-up diffs. **Pair-health + manifests (CLI + MCP):** profile comparisons now emit `pair_health` metadata (presence/absence of baseline/optimized NSYS/NCU pairs), and comparison staging writes `pair_manifest.json` so downstream automation can consume deterministic pair context without re-discovering files. diff --git a/code/mcp/mcp_server.py b/code/mcp/mcp_server.py index fe3519e0d..2d525737e 100644 --- a/code/mcp/mcp_server.py +++ b/code/mcp/mcp_server.py @@ -7415,7 +7415,7 @@ def _execute_capture_with_job(): "metric_set selects the NCU --set; workload_type picks custom metrics (only when metric_set=full). " "workload_type: memory_bound (default, fast), compute_bound, tensor_core. " "🕐 SLOW (varies). WORKFLOW: profile_kernels → profile_ncu. NOT FOR: Timeline (use profile_nsys). " - "DEFAULTS: metric_set='full' by default; use metric_set='minimal' (auto-resolves to speed-of-light/basic by Nsight version) for routine baseline/optimized compares; " + "DEFAULTS: metric_set='full' by default; use metric_set='minimal' (exactly five metrics without additional NVIDIA sections) for routine baseline/optimized compares; " "use metric_set='roofline' for bound analysis; use metric_set='full' for deep dives. " "COMPARE: compare_ncu auto-pairs baseline/optimized across subdirectories; pass pair if multiple pairs exist. " "Use launch_skip/launch_count to limit captures on many-launch benchmarks (e.g., 4096 batches). " diff --git a/code/tests/test_cli_profile_ncu.py b/code/tests/test_cli_profile_ncu.py index eefd91b75..f078c3c82 100644 --- a/code/tests/test_cli_profile_ncu.py +++ b/code/tests/test_cli_profile_ncu.py @@ -36,7 +36,7 @@ def profile_ncu(self, **kwargs: object) -> Path: type(self).calls.append(kwargs) output_name = str(kwargs.get("output_name", "profile_ncu")) metric_set = str(kwargs.get("metric_set", "full")) - resolved_metric_set = "basic" if metric_set == "minimal" else metric_set + resolved_metric_set = metric_set output_path = self.output_root / f"{output_name}.ncu-rep" output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text("fake ncu report", encoding="utf-8") @@ -107,7 +107,7 @@ def test_cli_profile_ncu_minimal_metric_set(fake_binary: Path, tmp_path: Path) - assert call["metric_set"] == "minimal" assert call["timeout_seconds"] == 60 assert "NCU report:" in result.stdout - assert "Metric set: minimal (resolved: basic)" in result.stdout + assert "Metric set: minimal (resolved: minimal)" in result.stdout def test_cli_profile_ncu_launch_limiting(fake_binary: Path, tmp_path: Path) -> None: diff --git a/code/tests/test_ncu_app_range_commands.py b/code/tests/test_ncu_app_range_commands.py index 6b83f32e9..f6a158f2e 100644 --- a/code/tests/test_ncu_app_range_commands.py +++ b/code/tests/test_ncu_app_range_commands.py @@ -48,7 +48,8 @@ def test_existing_lower_level_replay_policies(mode: str, honor: bool, expected: config = ProfilerConfig(ncu_replay_mode=mode, honor_replay_mode_in_minimal=honor) command = config.get_ncu_command_for_target("out", ["program"], nvtx_includes=[RANGE_NAME]) assert command[command.index("--replay-mode") + 1] == expected - assert command[command.index("--set") + 1] == "basic" + assert "--set" not in command + assert command[command.index("--metrics") + 1].split(",") == MINIMAL_METRICS @pytest.mark.parametrize("mode", NCU_REPLAY_MODES) diff --git a/code/tests/test_ncu_command_builder.py b/code/tests/test_ncu_command_builder.py index 18a414b2e..ab2564f19 100644 --- a/code/tests/test_ncu_command_builder.py +++ b/code/tests/test_ncu_command_builder.py @@ -3,6 +3,7 @@ import pytest from core.profiling.nsight_automation import NsightAutomation +from core.profiling.profiler_config import MINIMAL_METRICS def _build_cmd(**kwargs) -> list[str]: @@ -14,10 +15,16 @@ def _build_cmd(**kwargs) -> list[str]: ) -def test_ncu_command_minimal_uses_supported_lightweight_set(): - cmd = _build_cmd(metric_set="minimal") - set_idx = cmd.index("--set") - assert cmd[set_idx + 1] in {"speed-of-light", "basic"} +@pytest.mark.parametrize("replay_mode", ["kernel", "application"]) +def test_ncu_command_minimal_uses_exact_metrics_without_extra_sections(replay_mode): + cmd = _build_cmd(metric_set="minimal", replay_mode=replay_mode) + assert "--set" not in cmd + assert cmd[cmd.index("--metrics") + 1].split(",") == MINIMAL_METRICS + + +def test_explicit_basic_still_selects_nvidia_sections(): + cmd = _build_cmd(metric_set="basic") + assert cmd[cmd.index("--set") + 1] in {"basic", "speed-of-light"} assert "--metrics" not in cmd From f69a6c08dddcead28efa6ccb79098eceac65c2e1 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 02:37:01 -0700 Subject: [PATCH 05/19] perf: reuse serving engines and remove redundant 1P1D barriers --- .../cache_aware_disagg_inference/README.md | 1 + .../cache_aware_disagg_multigpu_common.py | 115 +++- code/labs/dynamic_router/README.md | 2 +- .../dynamic_router/baseline_dual_pool_vllm.py | 21 +- .../baseline_dynamic_router_vllm.py | 21 +- .../optimized_dual_pool_vllm.py | 21 +- .../optimized_dynamic_router_vllm.py | 21 +- code/labs/dynamic_router/vllm_runner.py | 583 ++++++++++++++++-- ..._serving_followthrough_engine_lifecycle.py | 272 ++++++++ 9 files changed, 997 insertions(+), 60 deletions(-) create mode 100644 code/tests/test_serving_followthrough_engine_lifecycle.py diff --git a/code/labs/cache_aware_disagg_inference/README.md b/code/labs/cache_aware_disagg_inference/README.md index 35c71820b..72e423596 100644 --- a/code/labs/cache_aware_disagg_inference/README.md +++ b/code/labs/cache_aware_disagg_inference/README.md @@ -44,6 +44,7 @@ python -m cli.aisp bench run --targets labs/cache_aware_disagg_inference --profi ## Notes - With two GPUs, the distributed target has one prefill and one decode rank. Both placement policies select the same decode rank, so this topology cannot demonstrate a reduction in migrations between decode ranks. September 7, 2026 repeated ABBA measurements on two B200s passed every full 2,048-element output comparison but found no speedup: median 14.331648 ms baseline and 14.499441 ms optimized (0.988428x; eight observations per arm, four fresh seeds). Standard deviations were 0.470046 and 0.517233 ms. These portable, unlocked observations did not reproduce the earlier single live-input run's 1.64306x. Both-arm Nsight traces are retained. Use at least two decode ranks to investigate migration benefits; this topology cannot establish that mechanism. +- Multi-GPU results now expose `cache_aware.decode_rank_count`, `cache_aware.affinity_opportunity_count`, and `cache_aware.affinity_placement_distinguishable`. A direct 1P1D run is classified as a comparison surface. The optimized 1P1D path removes redundant global barriers between blocking point-to-point transfers while preserving the final per-request drain; `cache_aware.direct_1p1d_sync_fast_path` and `cache_aware.global_barriers_avoided_per_request` identify that separate synchronization mechanism. Its performance effect remains unmeasured until a repeated B200 rerun with full-output checks and retained traces. - This lab is intentionally a logical reproduction of the scheduler/caching story, not a full serving engine. - Treat single-GPU `cache_aware_disagg` as a locality-comparison benchmark with a local comparison contract. The stable value on one GPU is the cache hit rate, KV transfer volume, and worker affinity improvement; the timed delta is recorded, but it is not a trustworthy headline speed gate on this host. - Judge the single-GPU target by cache hit rate, KV transfer volume, and worker affinity before raw wall-clock speedup. diff --git a/code/labs/cache_aware_disagg_inference/cache_aware_disagg_multigpu_common.py b/code/labs/cache_aware_disagg_inference/cache_aware_disagg_multigpu_common.py index a502d5ade..9b3efe4d1 100644 --- a/code/labs/cache_aware_disagg_inference/cache_aware_disagg_multigpu_common.py +++ b/code/labs/cache_aware_disagg_inference/cache_aware_disagg_multigpu_common.py @@ -235,6 +235,62 @@ def _choose_decode_rank( return prefill_ranks + ((plan.global_request_idx + stage_idx) % decode_ranks) +def _affinity_opportunity_count( + plans: Sequence[DistributedRequestPlan], + *, + prefill_ranks: int, + decode_ranks: int, +) -> int: + """Count placements where sticky and round-robin select different ranks.""" + count = 0 + for plan in plans: + for stage_idx in range(plan.warm_chunks, plan.total_chunks + 1): + baseline_rank = _choose_decode_rank( + plan, + stage_idx, + affinity_mode=DecodeAffinityMode.ROUND_ROBIN, + prefill_ranks=prefill_ranks, + decode_ranks=decode_ranks, + ) + sticky_rank = _choose_decode_rank( + plan, + stage_idx, + affinity_mode=DecodeAffinityMode.STICKY, + prefill_ranks=prefill_ranks, + decode_ranks=decode_ranks, + ) + count += int(baseline_rank != sticky_rank) + return count + + +def _use_direct_1p1d_sync_fast_path( + *, + affinity_mode: DecodeAffinityMode, + world_size: int, + prefill_ranks: int, + decode_ranks: int, +) -> bool: + """Use blocking point-to-point ordering instead of redundant global barriers.""" + return ( + affinity_mode == DecodeAffinityMode.STICKY + and world_size == 2 + and prefill_ranks == 1 + and decode_ranks == 1 + ) + + +def _direct_1p1d_barriers_avoided_per_request( + plans: Sequence[DistributedRequestPlan], +) -> float: + if not plans: + return 0.0 + barriers = sum( + (2 * (plan.total_chunks - plan.warm_chunks)) + 1 + for plan in plans + ) + return float(barriers) / float(len(plans)) + + def _build_request_plans( cfg: CacheAwareDisaggMultiGPUConfig, *, @@ -470,6 +526,17 @@ def _run_torchrun_worker( } plans = _build_request_plans(cfg, prefill_ranks=prefill_ranks) + affinity_opportunities = _affinity_opportunity_count( + plans, + prefill_ranks=prefill_ranks, + decode_ranks=decode_ranks, + ) + direct_1p1d_sync_fast_path = _use_direct_1p1d_sync_fast_path( + affinity_mode=affinity_mode, + world_size=world_size, + prefill_ranks=prefill_ranks, + decode_ranks=decode_ranks, + ) warm_cache_store: Dict[int, torch.Tensor] = {} prefill_seed_store: Dict[int, torch.Tensor] = {} @@ -595,7 +662,8 @@ def run_iteration( active_caches=active_caches, metrics=local_metrics, ) - _sync_and_barrier(device) + if not direct_1p1d_sync_fast_path: + _sync_and_barrier(device) if rank == plan.prefill_rank: chunk_kv, seed = model.prefill(chunks[chunk_idx]) @@ -615,7 +683,8 @@ def run_iteration( kv_buffers=kv_buffers, allow_allocation=False, ) - _sync_and_barrier(device) + if not direct_1p1d_sync_fast_path: + _sync_and_barrier(device) current_owner = target_rank current_cache_len += _chunk_length(cfg, chunk_idx) @@ -644,7 +713,8 @@ def run_iteration( active_caches=active_caches, metrics=local_metrics, ) - _sync_and_barrier(device) + if not direct_1p1d_sync_fast_path: + _sync_and_barrier(device) if rank == plan.prefill_rank: if seed is None: @@ -765,6 +835,19 @@ def run_iteration( "cache_aware.wall_tokens_per_second": ( total_generated_tokens * (int(iters) / elapsed_s) ), + "cache_aware.decode_rank_count": float(decode_ranks), + "cache_aware.affinity_opportunity_count": float(affinity_opportunities), + "cache_aware.affinity_placement_distinguishable": float( + affinity_opportunities > 0 + ), + "cache_aware.direct_1p1d_sync_fast_path": float( + direct_1p1d_sync_fast_path + ), + "cache_aware.global_barriers_avoided_per_request": ( + _direct_1p1d_barriers_avoided_per_request(plans) + if direct_1p1d_sync_fast_path + else 0.0 + ), } _write_metrics_sidecar( label=label, @@ -1350,6 +1433,22 @@ def benchmark_fn(self) -> None: ) custom_metrics["cache_aware.peer_handoffs"] = metrics["peer_handoffs"] custom_metrics["cache_aware.shared_reload_mb"] = metrics["shared_reload_bytes"] / 1e6 + affinity_opportunities = _affinity_opportunity_count( + self._request_plans, + prefill_ranks=self._resolved_prefill_ranks, + decode_ranks=self._resolved_decode_ranks, + ) + custom_metrics["cache_aware.decode_rank_count"] = float( + self._resolved_decode_ranks + ) + custom_metrics["cache_aware.affinity_opportunity_count"] = float( + affinity_opportunities + ) + custom_metrics["cache_aware.affinity_placement_distinguishable"] = float( + affinity_opportunities > 0 + ) + custom_metrics["cache_aware.direct_1p1d_sync_fast_path"] = 0.0 + custom_metrics["cache_aware.global_barriers_avoided_per_request"] = 0.0 def capture_verification_payload(self) -> None: if not self._outputs_ready or self._verify_prompt is None: @@ -1440,6 +1539,16 @@ def get_config(self) -> BenchmarkConfig: ncu_replay_mode_override=True, ) + def get_optimization_goal(self) -> str: + """Treat a one-decode-rank run as a topology/control comparison.""" + world_size = self._resolved_world_size or _world_size_hint() + prefill_ranks = self._resolved_prefill_ranks or _hint_prefill_ranks( + world_size, + self.cfg.prefill_ranks, + ) + decode_ranks = world_size - prefill_ranks + return "comparison" if decode_ranks < 2 else "speed" + def get_workload_metadata(self) -> Optional[WorkloadMetadata]: return self._workload_metadata diff --git a/code/labs/dynamic_router/README.md b/code/labs/dynamic_router/README.md index c1df19c30..7be5978ba 100644 --- a/code/labs/dynamic_router/README.md +++ b/code/labs/dynamic_router/README.md @@ -37,7 +37,7 @@ python -m cli.aisp bench run --targets labs/dynamic_router --profile minimal ## Notes - The dual-pool policies produce different batch shapes. On the pinned vLLM 0.16 stack with GPT-OSS-20B, the default backend produced different greedy tokens for identical prompts, including across repeated optimized runs. The explicit batch-invariant Triton configuration above matched all 1,734 output elements on 2×B200. Apply the same backend and environment to both arms; the option does not change the default backend for other workloads. Other models and stacks still require their own correctness check. -- Harness latency includes constructing both model engines on every benchmark invocation. Treat it as startup plus request processing, rather than steady-state routing throughput. Prefix caching is disabled so every request processes its full declared prompt. +- The vLLM benchmarks construct each model engine once in `setup()`, execute the harness-required five full-workload warmups, and reuse the idle engines for exactly three steady-state iterations. Every invocation clears completed-request bookkeeping, uses a fresh request-id generation, and still verifies every generated token. Custom metrics report engine startup, warmup request processing, and steady-state request processing separately. Teardown emits a `vllm_engine_lifecycle` JSON record with teardown and end-to-end wall time, so moving engine construction outside the steady-state timer cannot be presented as an end-to-end speedup. Prefix caching remains disabled so every request processes its full declared prompt. - The harness prepares live CPU prompt IDs during setup. The topology-aware runner requires this input; standalone entrypoints create default prompts before calling it. Conversion to the Python token lists required by vLLM remains part of request admission, and GPU-resident prompt inputs fail explicitly before conversion. - `driver.py` accepts knobs such as `--prefill-gpus`, `--decode-gpus`, and `--migration-budget` to stress different regimes. - vLLM integration now takes flags (`--model`, `--prefill-gpus`, `--decode-gpus`, etc.) plus locally available tokenizer/model weights. diff --git a/code/labs/dynamic_router/baseline_dual_pool_vllm.py b/code/labs/dynamic_router/baseline_dual_pool_vllm.py index def3fc190..b4205b2d5 100644 --- a/code/labs/dynamic_router/baseline_dual_pool_vllm.py +++ b/code/labs/dynamic_router/baseline_dual_pool_vllm.py @@ -42,6 +42,7 @@ def __init__(self) -> None: vllm_runner._CLI_ARGS ) self._topology = None + self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False request_count = len(self._prompt_lengths) self.register_workload_metadata( @@ -58,17 +59,26 @@ def setup(self) -> None: self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + self._engine_session = vllm_runner.create_dual_pool_vllm_session( + "shared", + topology_snapshot=self._topology, + cli_args=vllm_runner._CLI_ARGS, + warmup_runs=vllm_runner.WARMUP_ITERATIONS, + ) def benchmark_fn(self) -> None: if self._mode_input is None or int(self._mode_input[0]) != 0: raise RuntimeError("setup() must initialize shared-pool routing mode") if self._prompt_token_ids is None: raise RuntimeError("setup() must initialize live prompt-token input") + if self._engine_session is None: + raise RuntimeError("setup() must initialize reusable vLLM engines") self._summary = run_dual_pool_vllm_with_topology( "shared", topology_snapshot=self._topology, cli_args=vllm_runner._CLI_ARGS, prompt_token_ids=self._prompt_token_ids, + engine_session=self._engine_session, ) self._summary_ready = True @@ -94,6 +104,9 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + if self._engine_session is not None: + self._engine_session.close() + self._engine_session = None self.output = None self._metric_values = None self._metric_output_buffer = None @@ -104,7 +117,13 @@ def teardown(self) -> None: super().teardown() def get_config(self) -> Optional[BenchmarkConfig]: - return BenchmarkConfig(iterations=1, warmup=5, multi_gpu_required=True) + return BenchmarkConfig( + iterations=vllm_runner.STEADY_STATE_ITERATIONS, + warmup=vllm_runner.WARMUP_ITERATIONS, + adaptive_iterations=False, + timing_method="wall_clock", + multi_gpu_required=True, + ) def get_custom_metrics(self) -> Optional[Dict[str, float]]: return self._summary or None diff --git a/code/labs/dynamic_router/baseline_dynamic_router_vllm.py b/code/labs/dynamic_router/baseline_dynamic_router_vllm.py index 5b4066a15..45a02b252 100644 --- a/code/labs/dynamic_router/baseline_dynamic_router_vllm.py +++ b/code/labs/dynamic_router/baseline_dynamic_router_vllm.py @@ -40,6 +40,7 @@ def __init__(self) -> None: self._prompt_token_ids: Optional[torch.Tensor] = None self._prompt_lengths = vllm_runner.routing_prompt_lengths(vllm_runner._CLI_ARGS) self._topology = None + self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False request_count = len(self._prompt_lengths) self.register_workload_metadata( @@ -55,17 +56,26 @@ def setup(self) -> None: self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + self._engine_session = vllm_runner.create_vllm_routing_session( + "baseline", + topology_snapshot=self._topology, + cli_args=vllm_runner._CLI_ARGS, + warmup_runs=vllm_runner.WARMUP_ITERATIONS, + ) def benchmark_fn(self) -> None: if self._mode_input is None or int(self._mode_input[0]) != 0: raise RuntimeError("setup() must initialize baseline routing mode") if self._prompt_token_ids is None: raise RuntimeError("setup() must initialize live prompt-token input") + if self._engine_session is None: + raise RuntimeError("setup() must initialize reusable vLLM engines") self._summary = run_vllm_routing_with_topology( "baseline", topology_snapshot=self._topology, cli_args=vllm_runner._CLI_ARGS, prompt_token_ids=self._prompt_token_ids, + engine_session=self._engine_session, ) self._summary_ready = True @@ -91,6 +101,9 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + if self._engine_session is not None: + self._engine_session.close() + self._engine_session = None self.output = None self._metric_values = None self._metric_output_buffer = None @@ -101,7 +114,13 @@ def teardown(self) -> None: super().teardown() def get_config(self) -> Optional[BenchmarkConfig]: - return BenchmarkConfig(iterations=1, warmup=5, multi_gpu_required=True) + return BenchmarkConfig( + iterations=vllm_runner.STEADY_STATE_ITERATIONS, + warmup=vllm_runner.WARMUP_ITERATIONS, + adaptive_iterations=False, + timing_method="wall_clock", + multi_gpu_required=True, + ) def get_custom_metrics(self) -> Optional[Dict[str, float]]: return self._summary or None diff --git a/code/labs/dynamic_router/optimized_dual_pool_vllm.py b/code/labs/dynamic_router/optimized_dual_pool_vllm.py index 71ce3e981..1dc4e0e67 100644 --- a/code/labs/dynamic_router/optimized_dual_pool_vllm.py +++ b/code/labs/dynamic_router/optimized_dual_pool_vllm.py @@ -42,6 +42,7 @@ def __init__(self) -> None: vllm_runner._CLI_ARGS ) self._topology = None + self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False request_count = len(self._prompt_lengths) self.register_workload_metadata( @@ -58,17 +59,26 @@ def setup(self) -> None: self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + self._engine_session = vllm_runner.create_dual_pool_vllm_session( + "dual", + topology_snapshot=self._topology, + cli_args=vllm_runner._CLI_ARGS, + warmup_runs=vllm_runner.WARMUP_ITERATIONS, + ) def benchmark_fn(self) -> None: if self._mode_input is None or int(self._mode_input[0]) != 1: raise RuntimeError("setup() must initialize dual-pool routing mode") if self._prompt_token_ids is None: raise RuntimeError("setup() must initialize live prompt-token input") + if self._engine_session is None: + raise RuntimeError("setup() must initialize reusable vLLM engines") self._summary = run_dual_pool_vllm_with_topology( "dual", topology_snapshot=self._topology, cli_args=vllm_runner._CLI_ARGS, prompt_token_ids=self._prompt_token_ids, + engine_session=self._engine_session, ) self._summary_ready = True @@ -94,6 +104,9 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + if self._engine_session is not None: + self._engine_session.close() + self._engine_session = None self.output = None self._metric_values = None self._metric_output_buffer = None @@ -104,7 +117,13 @@ def teardown(self) -> None: super().teardown() def get_config(self) -> Optional[BenchmarkConfig]: - return BenchmarkConfig(iterations=1, warmup=5, multi_gpu_required=True) + return BenchmarkConfig( + iterations=vllm_runner.STEADY_STATE_ITERATIONS, + warmup=vllm_runner.WARMUP_ITERATIONS, + adaptive_iterations=False, + timing_method="wall_clock", + multi_gpu_required=True, + ) def get_custom_metrics(self) -> Optional[Dict[str, float]]: return self._summary or None diff --git a/code/labs/dynamic_router/optimized_dynamic_router_vllm.py b/code/labs/dynamic_router/optimized_dynamic_router_vllm.py index ec635dea9..887f6f7e2 100644 --- a/code/labs/dynamic_router/optimized_dynamic_router_vllm.py +++ b/code/labs/dynamic_router/optimized_dynamic_router_vllm.py @@ -40,6 +40,7 @@ def __init__(self) -> None: self._prompt_token_ids: Optional[torch.Tensor] = None self._prompt_lengths = vllm_runner.routing_prompt_lengths(vllm_runner._CLI_ARGS) self._topology = None + self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False request_count = len(self._prompt_lengths) self.register_workload_metadata( @@ -55,17 +56,26 @@ def setup(self) -> None: self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + self._engine_session = vllm_runner.create_vllm_routing_session( + "optimized", + topology_snapshot=self._topology, + cli_args=vllm_runner._CLI_ARGS, + warmup_runs=vllm_runner.WARMUP_ITERATIONS, + ) def benchmark_fn(self) -> None: if self._mode_input is None or int(self._mode_input[0]) != 2: raise RuntimeError("setup() must initialize optimized routing mode") if self._prompt_token_ids is None: raise RuntimeError("setup() must initialize live prompt-token input") + if self._engine_session is None: + raise RuntimeError("setup() must initialize reusable vLLM engines") self._summary = run_vllm_routing_with_topology( "optimized", topology_snapshot=self._topology, cli_args=vllm_runner._CLI_ARGS, prompt_token_ids=self._prompt_token_ids, + engine_session=self._engine_session, ) self._summary_ready = True @@ -91,6 +101,9 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + if self._engine_session is not None: + self._engine_session.close() + self._engine_session = None self.output = None self._metric_values = None self._metric_output_buffer = None @@ -101,7 +114,13 @@ def teardown(self) -> None: super().teardown() def get_config(self) -> Optional[BenchmarkConfig]: - return BenchmarkConfig(iterations=1, warmup=5, multi_gpu_required=True) + return BenchmarkConfig( + iterations=vllm_runner.STEADY_STATE_ITERATIONS, + warmup=vllm_runner.WARMUP_ITERATIONS, + adaptive_iterations=False, + timing_method="wall_clock", + multi_gpu_required=True, + ) def get_custom_metrics(self) -> Optional[Dict[str, float]]: return self._summary or None diff --git a/code/labs/dynamic_router/vllm_runner.py b/code/labs/dynamic_router/vllm_runner.py index ff6184052..21d8271eb 100644 --- a/code/labs/dynamic_router/vllm_runner.py +++ b/code/labs/dynamic_router/vllm_runner.py @@ -16,6 +16,7 @@ import time from contextlib import redirect_stdout from dataclasses import dataclass +from functools import wraps from typing import Dict, List, Optional, Sequence, Set, Tuple import torch @@ -49,6 +50,8 @@ def _skip(reason: str) -> None: _EXPECTED_TORCH_VERSION = _SERVING_STACK_PINS.torch_version _EXPECTED_VLLM_DIST_VERSION = _SERVING_STACK_PINS.vllm_version _EXPECTED_FLASHINFER_DIST_VERSION = _SERVING_STACK_PINS.flashinfer_version +WARMUP_ITERATIONS = 5 +STEADY_STATE_ITERATIONS = 3 def _is_vllm_abi_mismatch_error(exc: BaseException) -> bool: @@ -335,6 +338,7 @@ def __init__(self, gpu_id: str, device_index: int, model_id: str, *, attention_b print(captured, file=sys.stderr) self._inflight: Dict[str, _RequestRuntime] = {} self._completed_output_token_ids: Dict[str, Tuple[int, ...]] = {} + self._closed = False def add_request( self, @@ -434,6 +438,50 @@ def _consume_request_outputs(self, outputs, observed_at: float) -> Tuple[List[st def queue_depth(self) -> int: return self.engine.get_num_unfinished_requests() + def reset_request_state(self) -> None: + """Prepare an idle engine for another exact-output workload.""" + if self._closed: + raise RuntimeError(f"vLLM engine {self.gpu_id} is already closed") + unfinished = self.engine.get_num_unfinished_requests() + if unfinished or self._inflight: + raise RuntimeError( + f"Cannot reuse vLLM engine {self.gpu_id} with unfinished requests: " + f"engine={unfinished}, tracked={len(self._inflight)}" + ) + self._completed_output_token_ids.clear() + + def close(self, *, force: bool = False) -> None: + """Shut down the pinned vLLM EngineCore after all requests drain.""" + if self._closed: + return + unfinished = self.engine.get_num_unfinished_requests() + if unfinished or self._inflight: + if not force: + raise RuntimeError( + f"Cannot close vLLM engine {self.gpu_id} with unfinished requests: " + f"engine={unfinished}, tracked={len(self._inflight)}" + ) + errors: List[str] = [] + if force and self._inflight: + abort_request = getattr(self.engine, "abort_request", None) + if callable(abort_request): + request_ids = list(self._inflight) + try: + abort_request(request_ids) + except Exception as exc: + errors.append(f"abort {request_ids}: {exc}") + self._inflight.clear() + core_client = getattr(self.engine, "engine_core", None) + shutdown = getattr(core_client, "shutdown", None) + if callable(shutdown): + try: + shutdown() + except Exception as exc: + errors.append(f"EngineCore shutdown: {exc}") + self._closed = True + if errors: + raise RuntimeError("; ".join(errors)) + def snapshot_metrics(self, ttft_ema: Optional[float], tpot_ema: float) -> Dict[str, float]: mem_free_gb = 0.0 if torch.cuda.is_available(): @@ -494,6 +542,7 @@ def __init__(self, gpu_id: str, device_index: int, model_id: str, *, attention_b _skip("EngineCore.step_fn is unavailable; update vLLM to V1 or disable --use-v1-core-loop.") self._inflight: Dict[str, _RequestRuntime] = {} self._completed_output_token_ids: Dict[str, Tuple[int, ...]] = {} + self._closed = False def step(self, now: Optional[float] = None) -> Tuple[List[str], List[Tuple[str, float]], int]: outputs_dict, executed = self._core.step_fn() @@ -541,6 +590,246 @@ class _GPUHandle: numa_node: Optional[int] = None +class VllmEngineSession: + """Own reusable vLLM engines and retain their full lifecycle timings.""" + + def __init__( + self, + *, + workload_kind: str, + mode: str, + handles: Sequence[_GPUHandle], + model_id: str, + attention_backend: Optional[str], + wrapper_cls: type[_VllmWrapper], + warmup_runs: int, + ) -> None: + if warmup_runs < 0: + raise ValueError("warmup_runs must be non-negative") + self.workload_kind = workload_kind + self.mode = mode + self.handles = tuple(handles) + self.model_id = model_id + self.attention_backend = attention_backend + self.warmup_runs = int(warmup_runs) + self._created_at = time.perf_counter() + self._closed = False + self._run_count = 0 + self._active_run: Optional[Tuple[str, float]] = None + self._phase_durations_ms: Dict[str, List[float]] = { + "warmup": [], + "steady_state": [], + } + self._failed_runs: List[Dict[str, object]] = [] + self._primary_failure: Optional[BaseException] = None + self._teardown_ms: Optional[float] = None + self._end_to_end_ms: Optional[float] = None + self.engine_startup_ms = 0.0 + self.engines: Dict[str, _VllmWrapper] = {} + + startup_start = time.perf_counter() + try: + for handle in self.handles: + self.engines[handle.gpu_id] = wrapper_cls( + handle.gpu_id, + handle.device_index, + model_id, + attention_backend=attention_backend, + ) + except BaseException as exc: + self.engine_startup_ms = (time.perf_counter() - startup_start) * 1000.0 + self._primary_failure = exc + cleanup_start = time.perf_counter() + cleanup_errors: List[str] = [] + for engine in self.engines.values(): + try: + engine.close(force=True) + except Exception as cleanup_exc: + cleanup_errors.append(f"{engine.gpu_id}: {cleanup_exc}") + self._teardown_ms = (time.perf_counter() - cleanup_start) * 1000.0 + self._end_to_end_ms = (time.perf_counter() - self._created_at) * 1000.0 + self._closed = True + self._emit_lifecycle("startup_failed", cleanup_errors) + if cleanup_errors and hasattr(exc, "add_note"): + exc.add_note( + "vLLM partial-startup cleanup errors: " + "; ".join(cleanup_errors) + ) + raise + self.engine_startup_ms = (time.perf_counter() - startup_start) * 1000.0 + + def validate_layout( + self, + *, + workload_kind: str, + mode: str, + handles: Sequence[_GPUHandle], + model_id: str, + attention_backend: Optional[str], + ) -> None: + if self._closed: + raise RuntimeError("vLLM engine session is closed") + expected = ( + workload_kind, + mode, + tuple(handles), + model_id, + attention_backend, + ) + actual = ( + self.workload_kind, + self.mode, + self.handles, + self.model_id, + self.attention_backend, + ) + if actual != expected: + raise RuntimeError( + "Reusable vLLM engine session does not match the requested workload layout" + ) + + def begin_run(self) -> Tuple[str, str]: + if self._active_run is not None: + raise RuntimeError("vLLM engine session already has an active run") + for engine in self.engines.values(): + engine.reset_request_state() + phase = "warmup" if self._run_count < self.warmup_runs else "steady_state" + request_prefix = f"session-{self._run_count:04d}-" + self._active_run = (phase, time.perf_counter()) + return phase, request_prefix + + def finish_run(self, phase: str) -> None: + active = self._active_run + if active is None or active[0] != phase: + raise RuntimeError("vLLM engine session run phase is inconsistent") + elapsed_ms = (time.perf_counter() - active[1]) * 1000.0 + self._phase_durations_ms[phase].append(elapsed_ms) + self._run_count += 1 + self._active_run = None + + def abort_run( + self, + exc: BaseException, + *, + retain_inactive_failure: bool = False, + ) -> None: + """Release an active run lease while retaining its primary failure.""" + active = self._active_run + if active is None: + if retain_inactive_failure and self._primary_failure is None: + self._primary_failure = exc + return + if self._primary_failure is None: + self._primary_failure = exc + self._failed_runs.append( + { + "phase": active[0], + "elapsed_ms": (time.perf_counter() - active[1]) * 1000.0, + "error_type": type(exc).__name__, + "error": str(exc), + } + ) + self._active_run = None + + def lifecycle_metrics(self) -> Dict[str, float]: + warmup_samples = self._phase_durations_ms["warmup"] + steady_samples = self._phase_durations_ms["steady_state"] + metrics = { + "lifecycle.setup_engine_startup_ms": self.engine_startup_ms, + "lifecycle.engine_startup_ms": self.engine_startup_ms, + "lifecycle.engine_count": float(len(self.engines)), + "lifecycle.warmup_runs": float(len(warmup_samples)), + "lifecycle.warmup_request_processing_ms_total": float(sum(warmup_samples)), + "lifecycle.warmup_request_processing_ms_mean": ( + float(sum(warmup_samples) / len(warmup_samples)) if warmup_samples else 0.0 + ), + "lifecycle.steady_state_runs": float(len(steady_samples)), + "lifecycle.steady_state_request_processing_ms_total": float(sum(steady_samples)), + "lifecycle.steady_state_request_processing_ms_mean": ( + float(sum(steady_samples) / len(steady_samples)) if steady_samples else 0.0 + ), + "lifecycle.steady_state_request_processing_ms_last": ( + float(steady_samples[-1]) if steady_samples else 0.0 + ), + "lifecycle.engine_reuse_count": float(max(self._run_count - 1, 0)), + "lifecycle.failed_runs": float(len(self._failed_runs)), + "lifecycle.elapsed_before_teardown_ms": ( + time.perf_counter() - self._created_at + ) + * 1000.0, + } + if self._teardown_ms is not None: + metrics["lifecycle.engine_teardown_ms"] = self._teardown_ms + if self._end_to_end_ms is not None: + metrics["lifecycle.end_to_end_ms"] = self._end_to_end_ms + return metrics + + def _emit_lifecycle(self, disposition: str, errors: Sequence[str]) -> None: + failure = self._primary_failure + print( + json.dumps( + { + "event": "vllm_engine_lifecycle", + "disposition": disposition, + "workload_kind": self.workload_kind, + "mode": self.mode, + "setup_engine_startup_ms": self.engine_startup_ms, + "engine_startup_ms": self.engine_startup_ms, + "warmup_request_processing_ms": self._phase_durations_ms["warmup"], + "steady_state_request_processing_ms": self._phase_durations_ms[ + "steady_state" + ], + "failed_runs": self._failed_runs, + "engine_teardown_ms": self._teardown_ms, + "end_to_end_ms": self._end_to_end_ms, + "request_state_reset_per_run": True, + "primary_failure": ( + {"type": type(failure).__name__, "message": str(failure)} + if failure is not None + else None + ), + "shutdown_errors": list(errors), + }, + sort_keys=True, + ), + file=sys.stderr, + flush=True, + ) + + def close(self, *, preserve_primary_error: bool = False) -> List[str]: + if self._closed: + return [] + if self._active_run is not None: + if self._primary_failure is None: + raise RuntimeError("Cannot close vLLM engine session during an active run") + self.abort_run(self._primary_failure) + teardown_start = time.perf_counter() + errors: List[str] = [] + for engine in self.engines.values(): + try: + engine.close(force=self._primary_failure is not None) + except Exception as exc: + errors.append(f"{engine.gpu_id}: {exc}") + self._teardown_ms = (time.perf_counter() - teardown_start) * 1000.0 + self._end_to_end_ms = (time.perf_counter() - self._created_at) * 1000.0 + self._closed = True + if self._primary_failure is not None: + disposition = "failed_run" + else: + disposition = "completed" + if errors: + disposition += "_with_teardown_errors" + self._emit_lifecycle(disposition, errors) + if errors and self._primary_failure is not None and hasattr( + self._primary_failure, "add_note" + ): + self._primary_failure.add_note( + "vLLM engine teardown errors: " + "; ".join(errors) + ) + if errors and not preserve_primary_error and self._primary_failure is None: + raise RuntimeError("vLLM engine teardown failed: " + "; ".join(errors)) + return errors + + def _parse_device_list(raw: Optional[str], default: str, max_device: int) -> List[int]: raw = raw or default ids: List[int] = [] @@ -613,6 +902,199 @@ def _build_handles( return handles +def _require_vllm_host(*, workload_label: str, minimum_gpus: int) -> int: + if not torch.cuda.is_available(): + _skip(f"CUDA is required for {workload_label}.") + total_gpus = torch.cuda.device_count() + if total_gpus < minimum_gpus: + _skip(f"{workload_label} requires at least {minimum_gpus} GPUs.") + return total_gpus + + +def _routing_session_layout( + *, + mode: str, + topology_snapshot: TopologySnapshot, + cli_args: argparse.Namespace, +) -> Tuple[str, List[_GPUHandle], str, Optional[str], type[_VllmWrapper]]: + total_gpus = _require_vllm_host( + workload_label="vLLM routing demo", + minimum_gpus=2, + ) + model_id = cli_args.model + if not model_id: + _skip("Pass --model to run vLLM demo.") + _assert_vllm_runtime_ready() + decode_ids = _parse_device_list(cli_args.decode_gpus, "0,1", total_gpus) + if not decode_ids: + decode_ids = list(range(min(2, total_gpus))) + handles = _build_handles( + "shared", + decode_ids, + decode_ids, + gpu_numa=topology_snapshot.gpu_numa, + ) + return ( + mode, + handles, + model_id, + getattr(cli_args, "attention_backend", None), + _VllmWrapper, + ) + + +def _dual_pool_session_layout( + *, + mode: str, + topology_snapshot: TopologySnapshot, + cli_args: argparse.Namespace, +) -> Tuple[str, List[_GPUHandle], str, Optional[str], type[_VllmWrapper]]: + total_gpus = _require_vllm_host( + workload_label="Dual-pool demo", + minimum_gpus=2, + ) + model_id = cli_args.model + if not model_id: + _skip("Pass --model to run vLLM dual-pool demo.") + _assert_vllm_runtime_ready() + + normalized_mode = mode.lower() + if normalized_mode in {"dual", "dual_pool", "optimized"}: + normalized_mode = "dual" + else: + normalized_mode = "shared" + + prefill_ids = _parse_device_list(cli_args.prefill_gpus, "0", total_gpus) + decode_default = "1" if total_gpus > 1 else "0" + decode_ids = _parse_device_list(cli_args.decode_gpus, decode_default, total_gpus) + if not prefill_ids: + prefill_ids = [0] + if not decode_ids: + decode_ids = [1] if total_gpus > 1 else [0] + if normalized_mode == "dual": + if not prefill_ids: + _skip("Dual mode needs at least one prefill GPU.") + if not decode_ids: + _skip("Dual mode needs at least one decode GPU.") + if not (set(prefill_ids) - set(decode_ids)) or not ( + set(decode_ids) - set(prefill_ids) + ): + _skip( + "Dual mode needs at least one GPU dedicated to prefill and one to decode. " + "Adjust VLLM_PREFILL_GPUS/VLLM_DECODE_GPUS." + ) + + handles = _build_handles( + normalized_mode, + prefill_ids, + decode_ids, + gpu_numa=topology_snapshot.gpu_numa, + ) + wrapper_cls = ( + _VllmV1Wrapper + if getattr(cli_args, "use_v1_core_loop", False) + else _VllmWrapper + ) + return ( + normalized_mode, + handles, + model_id, + getattr(cli_args, "attention_backend", None), + wrapper_cls, + ) + + +def create_vllm_routing_session( + mode: str, + *, + topology_snapshot: TopologySnapshot, + cli_args: Optional[argparse.Namespace] = None, + warmup_runs: int = 0, +) -> VllmEngineSession: + """Construct routing engines once so request processing can be timed separately.""" + args = cli_args or _CLI_ARGS + normalized_mode, handles, model_id, attention_backend, wrapper_cls = ( + _routing_session_layout( + mode=mode, + topology_snapshot=topology_snapshot, + cli_args=args, + ) + ) + return VllmEngineSession( + workload_kind="dynamic_router", + mode=normalized_mode, + handles=handles, + model_id=model_id, + attention_backend=attention_backend, + wrapper_cls=wrapper_cls, + warmup_runs=warmup_runs, + ) + + +def create_dual_pool_vllm_session( + mode: str, + *, + topology_snapshot: TopologySnapshot, + cli_args: Optional[argparse.Namespace] = None, + warmup_runs: int = 0, +) -> VllmEngineSession: + """Construct shared or dual-pool engines once for steady-state replay.""" + args = cli_args or _CLI_ARGS + normalized_mode, handles, model_id, attention_backend, wrapper_cls = ( + _dual_pool_session_layout( + mode=mode, + topology_snapshot=topology_snapshot, + cli_args=args, + ) + ) + return VllmEngineSession( + workload_kind="dual_pool", + mode=normalized_mode, + handles=handles, + model_id=model_id, + attention_backend=attention_backend, + wrapper_cls=wrapper_cls, + warmup_runs=warmup_runs, + ) + + +def _manage_engine_session(session_factory): + """Close call-owned sessions on success or failure without masking failures.""" + + def decorate(run_fn): + @wraps(run_fn) + def managed(mode, *args, **kwargs): + session = kwargs.get("engine_session") + owns_session = session is None + if owns_session: + topology_snapshot = kwargs.get("topology_snapshot") + if topology_snapshot is None: + raise TypeError("topology_snapshot must be passed by keyword") + session = session_factory( + mode, + topology_snapshot=topology_snapshot, + cli_args=kwargs.get("cli_args"), + warmup_runs=0, + ) + kwargs["engine_session"] = session + try: + summary = run_fn(mode, *args, **kwargs) + except BaseException as exc: + session.abort_run(exc, retain_inactive_failure=owns_session) + if owns_session: + session.close(preserve_primary_error=True) + raise + if owns_session: + session.close() + if isinstance(summary, dict): + summary.update(session.lifecycle_metrics()) + return summary + + return managed + + return decorate + + def _collect_verification_output_token_ids( engines: Dict[str, _VllmWrapper], request_ids: List[str] ) -> List[int]: @@ -643,6 +1125,7 @@ def _collect_verification_output_token_ids( return framed +@_manage_engine_session(create_vllm_routing_session) def run_vllm_routing_with_topology( mode: str, *, @@ -651,18 +1134,29 @@ def run_vllm_routing_with_topology( max_tokens: Optional[int] = None, cli_args: Optional[argparse.Namespace] = None, prompt_token_ids: torch.Tensor, + engine_session: Optional[VllmEngineSession] = None, ) -> Dict[str, float]: """Run a small vLLM-backed routing demo with a precomputed topology snapshot.""" - if not torch.cuda.is_available(): - _skip("CUDA is required for vLLM routing demo.") - if torch.cuda.device_count() < 2: - _skip("vLLM routing demo requires at least 2 GPUs.") - args = cli_args or _CLI_ARGS - model_id = args.model - if not model_id: - _skip("Pass --model to run vLLM demo.") - _assert_vllm_runtime_ready() + normalized_mode, handles, model_id, attention_backend, wrapper_cls = ( + _routing_session_layout( + mode=mode, + topology_snapshot=topology_snapshot, + cli_args=args, + ) + ) + if engine_session is None: + raise RuntimeError("managed vLLM routing call did not receive an engine session") + session = engine_session + session.validate_layout( + workload_kind="dynamic_router", + mode=normalized_mode, + handles=handles, + model_id=model_id, + attention_backend=attention_backend, + ) + run_phase, request_prefix = session.begin_run() + engines = session.engines prompt_lengths = routing_prompt_lengths(args, req_count=req_count) req_count_val = len(prompt_lengths) @@ -674,14 +1168,6 @@ def run_vllm_routing_with_topology( topo = topology_snapshot gpu_numa = topo.gpu_numa - decode_ids = _parse_device_list(args.decode_gpus, "0,1", torch.cuda.device_count()) - if not decode_ids: - decode_ids = list(range(min(2, torch.cuda.device_count()))) - engines = { - f"gpu{idx}": _VllmWrapper(f"gpu{idx}", idx, model_id, attention_backend=getattr(args, "attention_backend", None)) - for idx in decode_ids - } - # Router selection router = Router() if mode == "optimized" else None if router: @@ -702,7 +1188,7 @@ def run_vllm_routing_with_topology( # Submit all requests up front for i in range(req_count_val): - rid = f"req-{i}" + rid = f"{request_prefix}req-{i}" request_ids.append(rid) req = Request( req_id=rid, @@ -750,6 +1236,8 @@ def run_vllm_routing_with_topology( summary[VERIFICATION_OUTPUT_KEY] = _collect_verification_output_token_ids( engines, request_ids ) + session.finish_run(run_phase) + summary.update(session.lifecycle_metrics()) return summary @@ -774,6 +1262,7 @@ def run_vllm_routing( ) +@_manage_engine_session(create_dual_pool_vllm_session) def run_dual_pool_vllm_with_topology( mode: str, *, @@ -787,28 +1276,20 @@ def run_dual_pool_vllm_with_topology( prefill_ctx_thresh: Optional[int] = None, cli_args: Optional[argparse.Namespace] = None, prompt_token_ids: torch.Tensor, + engine_session: Optional[VllmEngineSession] = None, ) -> Dict[str, float]: """ Dual-pool vLLM experiment: compare shared-pool vs disaggregated prefill/decode. """ - if not torch.cuda.is_available(): - _skip("CUDA is required for vLLM dual-pool demo.") - - total_gpus = torch.cuda.device_count() - if total_gpus < 2: - _skip("Dual-pool demo requires at least 2 GPUs.") - args = cli_args or _CLI_ARGS - model_id = args.model - if not model_id: - _skip("Pass --model to run vLLM dual-pool demo.") - _assert_vllm_runtime_ready() - - normalized_mode = mode.lower() - if normalized_mode in {"dual", "dual_pool", "optimized"}: - normalized_mode = "dual" - else: - normalized_mode = "shared" + normalized_mode, handles, model_id, attention_backend, wrapper_cls = ( + _dual_pool_session_layout( + mode=mode, + topology_snapshot=topology_snapshot, + cli_args=args, + ) + ) + total_gpus = torch.cuda.device_count() long_prompt_tokens = args.long_prompt_tokens if long_prompt_tokens is None else long_prompt_tokens short_prompt_tokens = args.short_prompt_tokens if short_prompt_tokens is None else short_prompt_tokens @@ -839,26 +1320,22 @@ def run_dual_pool_vllm_with_topology( if not decode_ids: decode_ids = [1] if total_gpus > 1 else [0] - if normalized_mode == "dual": - if not set(prefill_ids): - _skip("Dual mode needs at least one prefill GPU.") - if not set(decode_ids): - _skip("Dual mode needs at least one decode GPU.") - if not (set(prefill_ids) - set(decode_ids)) or not (set(decode_ids) - set(prefill_ids)): - _skip("Dual mode needs at least one GPU dedicated to prefill and one to decode. Adjust VLLM_PREFILL_GPUS/VLLM_DECODE_GPUS.") - - topo = topology_snapshot - handles = _build_handles(normalized_mode, prefill_ids, decode_ids, gpu_numa=topo.gpu_numa) prefill_handles = [h for h in handles if h.is_prefill] decode_handles = [h for h in handles if h.is_decode] if not prefill_handles or not decode_handles: _skip("No usable GPUs after parsing pool assignments.") - - wrapper_cls = _VllmV1Wrapper if getattr(args, "use_v1_core_loop", False) else _VllmWrapper - engines = { - h.gpu_id: wrapper_cls(h.gpu_id, h.device_index, model_id, attention_backend=getattr(args, "attention_backend", None)) - for h in handles - } + if engine_session is None: + raise RuntimeError("managed dual-pool call did not receive an engine session") + session = engine_session + session.validate_layout( + workload_kind="dual_pool", + mode=normalized_mode, + handles=handles, + model_id=model_id, + attention_backend=attention_backend, + ) + run_phase, request_prefix = session.begin_run() + engines = session.engines router = Router() for h in handles: @@ -875,7 +1352,7 @@ def run_dual_pool_vllm_with_topology( def _enqueue(n: int, prompt_tokens: int, hint: str) -> None: nonlocal next_id for _ in range(n): - rid = f"req-{next_id}" + rid = f"{request_prefix}req-{next_id}" next_id += 1 workload.append( ( @@ -1004,6 +1481,8 @@ def _enqueue(n: int, prompt_tokens: int, hint: str) -> None: summary[VERIFICATION_OUTPUT_KEY] = _collect_verification_output_token_ids( engines, list(req_roles) ) + session.finish_run(run_phase) + summary.update(session.lifecycle_metrics()) return summary diff --git a/code/tests/test_serving_followthrough_engine_lifecycle.py b/code/tests/test_serving_followthrough_engine_lifecycle.py new file mode 100644 index 000000000..797db544d --- /dev/null +++ b/code/tests/test_serving_followthrough_engine_lifecycle.py @@ -0,0 +1,272 @@ +"""Focused contracts for reusable serving engines and direct 1P1D routing.""" + +from __future__ import annotations + +import importlib +import json +from types import SimpleNamespace + +import pytest + +from labs.cache_aware_disagg_inference.cache_aware_disagg_multigpu_common import ( + DecodeAffinityMode, + DistributedRequestPlan, + _affinity_opportunity_count, + _direct_1p1d_barriers_avoided_per_request, + _use_direct_1p1d_sync_fast_path, +) +from labs.dynamic_router import vllm_runner + + +class _FakeReusableWrapper: + created: list[_FakeReusableWrapper] = [] + + def __init__( + self, + gpu_id: str, + device_index: int, + model_id: str, + *, + attention_backend: str | None, + ) -> None: + self.gpu_id = gpu_id + self.device_index = device_index + self.model_id = model_id + self.attention_backend = attention_backend + self.reset_calls = 0 + self.close_calls = 0 + self.force_close_calls = 0 + self.created.append(self) + + def reset_request_state(self) -> None: + self.reset_calls += 1 + + def close(self, *, force: bool = False) -> None: + self.close_calls += 1 + self.force_close_calls += int(force) + + +def test_engine_session_reuses_engines_and_reports_all_lifecycle_phases( + capsys: pytest.CaptureFixture[str], +) -> None: + _FakeReusableWrapper.created.clear() + handles = [ + vllm_runner._GPUHandle("gpu0", 0, True, False, 0), + vllm_runner._GPUHandle("gpu1", 1, False, True, 1), + ] + session = vllm_runner.VllmEngineSession( + workload_kind="dual_pool", + mode="dual", + handles=handles, + model_id="/models/local", + attention_backend="TRITON_ATTN", + wrapper_cls=_FakeReusableWrapper, + warmup_runs=1, + ) + + warmup_phase, warmup_prefix = session.begin_run() + session.finish_run(warmup_phase) + steady_phase, steady_prefix = session.begin_run() + session.finish_run(steady_phase) + + assert warmup_phase == "warmup" + assert steady_phase == "steady_state" + assert warmup_prefix != steady_prefix + assert len(_FakeReusableWrapper.created) == 2 + assert [wrapper.reset_calls for wrapper in _FakeReusableWrapper.created] == [2, 2] + metrics = session.lifecycle_metrics() + assert metrics["lifecycle.warmup_runs"] == 1.0 + assert metrics["lifecycle.steady_state_runs"] == 1.0 + assert metrics["lifecycle.engine_reuse_count"] == 1.0 + assert "lifecycle.engine_teardown_ms" not in metrics + + session.close() + lifecycle = json.loads(capsys.readouterr().err) + assert lifecycle["event"] == "vllm_engine_lifecycle" + assert lifecycle["disposition"] == "completed" + assert lifecycle["request_state_reset_per_run"] is True + assert len(lifecycle["warmup_request_processing_ms"]) == 1 + assert len(lifecycle["steady_state_request_processing_ms"]) == 1 + assert lifecycle["engine_teardown_ms"] >= 0.0 + assert lifecycle["end_to_end_ms"] >= lifecycle["engine_startup_ms"] + assert lifecycle["shutdown_errors"] == [] + assert [wrapper.close_calls for wrapper in _FakeReusableWrapper.created] == [1, 1] + + +def test_partial_engine_startup_closes_created_engine_and_keeps_primary_error( + capsys: pytest.CaptureFixture[str], +) -> None: + class FailSecondWrapper(_FakeReusableWrapper): + def __init__(self, gpu_id: str, *args, **kwargs) -> None: + if gpu_id == "gpu1": + raise ValueError("second engine startup failed") + super().__init__(gpu_id, *args, **kwargs) + + _FakeReusableWrapper.created.clear() + handles = [ + vllm_runner._GPUHandle("gpu0", 0, True, False, 0), + vllm_runner._GPUHandle("gpu1", 1, False, True, 1), + ] + with pytest.raises(ValueError, match="second engine startup failed"): + vllm_runner.VllmEngineSession( + workload_kind="dual_pool", + mode="dual", + handles=handles, + model_id="/models/local", + attention_backend="TRITON_ATTN", + wrapper_cls=FailSecondWrapper, + warmup_runs=1, + ) + + assert len(_FakeReusableWrapper.created) == 1 + assert _FakeReusableWrapper.created[0].force_close_calls == 1 + lifecycle = json.loads(capsys.readouterr().err) + assert lifecycle["disposition"] == "startup_failed" + assert lifecycle["primary_failure"] == { + "type": "ValueError", + "message": "second engine startup failed", + } + assert lifecycle["shutdown_errors"] == [] + + +def test_owned_session_failure_keeps_primary_error_and_emits_teardown_disposition( + capsys: pytest.CaptureFixture[str], +) -> None: + class FailingCloseWrapper(_FakeReusableWrapper): + def close(self, *, force: bool = False) -> None: + super().close(force=force) + raise RuntimeError("shutdown failed") + + created_sessions: list[vllm_runner.VllmEngineSession] = [] + + def factory(mode: str, **kwargs) -> vllm_runner.VllmEngineSession: + del kwargs + session = vllm_runner.VllmEngineSession( + workload_kind="test", + mode=mode, + handles=[vllm_runner._GPUHandle("gpu0", 0, True, True, 0)], + model_id="/models/local", + attention_backend=None, + wrapper_cls=FailingCloseWrapper, + warmup_runs=0, + ) + created_sessions.append(session) + return session + + @vllm_runner._manage_engine_session(factory) + def failing_run(mode: str, *, topology_snapshot, engine_session=None) -> None: + del mode, topology_snapshot + engine_session.begin_run() + raise ValueError("request processing failed") + + with pytest.raises(ValueError, match="request processing failed") as exc_info: + failing_run("test", topology_snapshot=object()) + + assert len(created_sessions) == 1 + assert created_sessions[0]._closed is True + assert any("shutdown failed" in note for note in (exc_info.value.__notes__ or [])) + lifecycle = json.loads(capsys.readouterr().err) + assert lifecycle["disposition"] == "failed_run_with_teardown_errors" + assert lifecycle["primary_failure"] == { + "type": "ValueError", + "message": "request processing failed", + } + assert lifecycle["failed_runs"][0]["phase"] == "steady_state" + assert lifecycle["shutdown_errors"] == ["gpu0: shutdown failed"] + + +def test_wrapper_reset_requires_an_idle_engine_and_clears_completed_outputs() -> None: + wrapper = vllm_runner._VllmWrapper.__new__(vllm_runner._VllmWrapper) + wrapper.gpu_id = "gpu0" + wrapper.engine = SimpleNamespace(get_num_unfinished_requests=lambda: 0) + wrapper._inflight = {} + wrapper._completed_output_token_ids = {"old-request": (1, 2)} + wrapper._closed = False + + wrapper.reset_request_state() + assert wrapper._completed_output_token_ids == {} + + wrapper.engine = SimpleNamespace(get_num_unfinished_requests=lambda: 1) + with pytest.raises(RuntimeError, match="unfinished requests"): + wrapper.reset_request_state() + + +def test_wrapper_force_close_aborts_vllm_request_list_before_shutdown() -> None: + aborted: list[list[str]] = [] + shutdown_calls: list[bool] = [] + wrapper = vllm_runner._VllmWrapper.__new__(vllm_runner._VllmWrapper) + wrapper.gpu_id = "gpu0" + wrapper.engine = SimpleNamespace( + get_num_unfinished_requests=lambda: 2, + abort_request=lambda request_ids: aborted.append(list(request_ids)), + engine_core=SimpleNamespace(shutdown=lambda: shutdown_calls.append(True)), + ) + wrapper._inflight = {"request-a": object(), "request-b": object()} + wrapper._completed_output_token_ids = {} + wrapper._closed = False + + wrapper.close(force=True) + + assert aborted == [["request-a", "request-b"]] + assert shutdown_calls == [True] + assert wrapper._inflight == {} + assert wrapper._closed is True + + +def test_one_decode_rank_exposes_no_affinity_opportunity_and_exact_fast_path() -> None: + plans = [ + DistributedRequestPlan(0, 0, 0, warm_chunks=2, total_chunks=4), + DistributedRequestPlan(0, 1, 1, warm_chunks=0, total_chunks=4), + ] + + assert _affinity_opportunity_count( + plans, + prefill_ranks=1, + decode_ranks=1, + ) == 0 + assert _affinity_opportunity_count( + plans, + prefill_ranks=1, + decode_ranks=2, + ) > 0 + assert _use_direct_1p1d_sync_fast_path( + affinity_mode=DecodeAffinityMode.STICKY, + world_size=2, + prefill_ranks=1, + decode_ranks=1, + ) + assert not _use_direct_1p1d_sync_fast_path( + affinity_mode=DecodeAffinityMode.ROUND_ROBIN, + world_size=2, + prefill_ranks=1, + decode_ranks=1, + ) + assert not _use_direct_1p1d_sync_fast_path( + affinity_mode=DecodeAffinityMode.STICKY, + world_size=3, + prefill_ranks=1, + decode_ranks=2, + ) + assert _direct_1p1d_barriers_avoided_per_request(plans) == 7.0 + + +@pytest.mark.parametrize( + "module_name", + [ + "labs.dynamic_router.baseline_dynamic_router_vllm", + "labs.dynamic_router.optimized_dynamic_router_vllm", + "labs.dynamic_router.baseline_dual_pool_vllm", + "labs.dynamic_router.optimized_dual_pool_vllm", + ], +) +def test_vllm_benchmarks_separate_five_warmups_from_three_wall_clock_runs( + module_name: str, +) -> None: + benchmark = importlib.import_module(module_name).get_benchmark() + config = benchmark.get_config() + + assert config is not None + assert config.warmup == vllm_runner.WARMUP_ITERATIONS == 5 + assert config.iterations == vllm_runner.STEADY_STATE_ITERATIONS == 3 + assert config.adaptive_iterations is False + assert config.timing_method == "wall_clock" From 09207e5a683a98a6feabed981c80be6a2ad6073f Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 02:39:05 -0700 Subject: [PATCH 06/19] perf: amortize fixed-weight pipeline fill and drain across repeats --- code/ch04/optimized_pipeline_parallel_1f1b.py | 78 +++++++-- code/ch04/pipeline_parallel_common.py | 18 +- ...ne-1f1b-contiguous-performance-intake.yaml | 127 ++++++++++++++ .../test_pipeline_1f1b_contiguous_repeats.py | 156 ++++++++++++++++++ 4 files changed, 363 insertions(+), 16 deletions(-) create mode 100644 code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml create mode 100644 code/tests/test_pipeline_1f1b_contiguous_repeats.py diff --git a/code/ch04/optimized_pipeline_parallel_1f1b.py b/code/ch04/optimized_pipeline_parallel_1f1b.py index 2dc60dee8..9dff9a74c 100644 --- a/code/ch04/optimized_pipeline_parallel_1f1b.py +++ b/code/ch04/optimized_pipeline_parallel_1f1b.py @@ -10,6 +10,7 @@ import argparse import os import time +from collections.abc import Callable, Sequence from typing import Optional import torch @@ -135,6 +136,54 @@ def _run_rank_stage_inplace( return _run_stage_inplace(stages[rank], x) +def _run_contiguous_1f1b_iterations( + *, + rank: int, + world_size: int, + micro_batches_per_iteration: int, + iteration_count: int, + get_rank0_microbatch: Callable[[int], torch.Tensor], + recv_forward_buffers: Sequence[torch.Tensor], + recv_backward_buffers: Sequence[torch.Tensor], + forward_step: Callable[[torch.Tensor], torch.Tensor], + backward_step: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + activation_slots: list[Optional[tuple[int, torch.Tensor]]], + capture: Optional[PipelineIterationCapture] = None, +) -> None: + """Amortize boundaries across fixed-state iterations with no intervening update. + + This is valid for this benchmark's fixed weights and repeated input batch. A + training loop that updates parameters between iterations must drain before + each update instead of using this helper across that boundary. + """ + + if iteration_count <= 0: + raise ValueError("iteration_count must be positive") + if micro_batches_per_iteration <= 0: + raise ValueError("micro_batches_per_iteration must be positive") + if micro_batches_per_iteration < world_size: + raise ValueError("Each logical iteration needs at least one microbatch per stage") + scheduled_micro_batches = micro_batches_per_iteration * iteration_count + + def get_repeated_rank0_microbatch(index: int) -> torch.Tensor: + return get_rank0_microbatch(index % micro_batches_per_iteration) + + run_1f1b_iteration( + rank=rank, + world_size=world_size, + num_micro_batches=scheduled_micro_batches, + get_rank0_microbatch=get_repeated_rank0_microbatch, + # Only references are repeated. The same fixed tensor slots are safe to + # receive into after their prior microbatch has completed its stage. + recv_forward_buffers=list(recv_forward_buffers) * iteration_count, + recv_backward_buffers=list(recv_backward_buffers) * iteration_count, + forward_step=forward_step, + backward_step=backward_step, + activation_slots=activation_slots, + capture=capture, + ) + + def _run_worker( iters: int, warmup: int, @@ -209,11 +258,15 @@ def _get_rank0_microbatch(micro_idx: int) -> torch.Tensor: warmup_steps, 1 ) - def _run_iteration(capture: Optional[PipelineIterationCapture] = None) -> None: - run_1f1b_iteration( + def _run_contiguous_iterations( + iteration_count: int, + capture: Optional[PipelineIterationCapture] = None, + ) -> None: + _run_contiguous_1f1b_iterations( rank=rank, world_size=world_size, - num_micro_batches=num_micro_batches, + micro_batches_per_iteration=num_micro_batches, + iteration_count=iteration_count, get_rank0_microbatch=_get_rank0_microbatch, recv_forward_buffers=recv_micro_batches, recv_backward_buffers=recv_grads, @@ -226,22 +279,21 @@ def _run_iteration(capture: Optional[PipelineIterationCapture] = None) -> None: result_requested = pipeline_child_result_requested() captured_iteration: Optional[PipelineIterationCapture] = None with torch.inference_mode(): - for _ in range(max(warmup, 0)): - _run_iteration() + warmup_iterations = max(warmup, 0) + if warmup_iterations: + _run_contiguous_iterations(warmup_iterations) torch.cuda.synchronize(device) with nvtx_range(PROFILE_NVTX_RANGE, enable=True): start = time.perf_counter() measured_iterations = max(iters, 1) - for iteration in range(max(iters, 1)): - capture = ( - PipelineIterationCapture.create(num_micro_batches) - if result_requested and iteration == measured_iterations - 1 - else None + if result_requested: + captured_iteration = PipelineIterationCapture.create( + num_micro_batches, + first_microbatch_index=(measured_iterations - 1) + * num_micro_batches, ) - _run_iteration(capture) - if capture is not None: - captured_iteration = capture + _run_contiguous_iterations(measured_iterations, captured_iteration) torch.cuda.synchronize(device) elapsed = time.perf_counter() - start diff --git a/code/ch04/pipeline_parallel_common.py b/code/ch04/pipeline_parallel_common.py index 678a46533..da236546e 100644 --- a/code/ch04/pipeline_parallel_common.py +++ b/code/ch04/pipeline_parallel_common.py @@ -58,20 +58,29 @@ @dataclass class PipelineIterationCapture: - """References to one measured iteration, indexed by original microbatch.""" + """References to a contiguous microbatch window from a measured schedule.""" forward_inputs: list[torch.Tensor | None] backward_inputs: list[torch.Tensor | None] backward_outputs: list[torch.Tensor | None] + first_microbatch_index: int = 0 @classmethod - def create(cls, num_micro_batches: int) -> PipelineIterationCapture: + def create( + cls, + num_micro_batches: int, + *, + first_microbatch_index: int = 0, + ) -> PipelineIterationCapture: if num_micro_batches <= 0: raise ValueError("num_micro_batches must be positive") + if first_microbatch_index < 0: + raise ValueError("first_microbatch_index must be non-negative") return cls( forward_inputs=[None] * num_micro_batches, backward_inputs=[None] * num_micro_batches, backward_outputs=[None] * num_micro_batches, + first_microbatch_index=first_microbatch_index, ) def concatenate(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -114,7 +123,10 @@ def _record( tensor: torch.Tensor, ) -> None: if capture is not None: - getattr(capture, collection)[microbatch_index] = tensor + capture_index = microbatch_index - capture.first_microbatch_index + values = getattr(capture, collection) + if 0 <= capture_index < len(values): + values[capture_index] = tensor def _exchange_neighbor( diff --git a/code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml b/code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml new file mode 100644 index 000000000..25ef82c0f --- /dev/null +++ b/code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml @@ -0,0 +1,127 @@ +objective: + primary_kpi: rank0_time_per_iter_ms + secondary_kpis: [p2p_kernel_instances, exact_full_output] + +benchmark: + layer: component + grade: publication + variable_under_test: scheduler_path + +workload: + type: inference + model: "toy BF16 pipeline: 8 forward and 8 backward bias-free Linear+ReLU layers" + prompt_or_seq_len: + p50: 2048 + p95: 2048 + max: 2048 + batch_shape: + p50: [32, 2048, 4096] + p95: [32, 2048, 4096] + max: [32, 2048, 4096] + constraints: + - 2 pipeline stages on 2 NVIDIA B200 GPUs + - 8 microbatches and 3 measured iterations after 5 warmup iterations + - preserve every forward step, backward step, and P2P payload transfer + - preserve the full [32, 2048, 4096] output from each rank + - require exact full-output equality for candidate qualification + - fixed weights and fixed seeded inputs with no optimizer update between logical iterations + +slos: + latency_ms: + p50: null + p95: null + p99: null + availability: null + cost_budget: "none; bounded reuse of the existing 2-B200 validation host" + +environment: + hardware: + gpu: ["NVIDIA B200", 2, "unknown in retained receipt"] + cpu: "not recorded in retained receipt" + ram_gb: null + storage: "not material to timed worker interval" + interconnect: ["GPU P2P; exact NVLink/NVSwitch topology not recorded"] + topology: "two visible GPUs, CUDA_VISIBLE_DEVICES=0,1" + software: + os: "linux" + kernel: "not recorded" + driver: "580.173.02" + cuda: "13.0" + pytorch: "2.9.1+cu130" + triton: "3.5.1" + frameworks: [] + nccl: "2.27.7" + container_orchestrator: "not recorded" + data_pipeline: ["fixed seeded device input reused for each measured iteration"] + +current_baseline: + time_to_train_hours: null + mfu_pct: null + scaling_efficiency_pct: null + tokens_per_second: null + ttft_ms: null + p50_ms: 15.9415460075 + p99_ms: null + jitter_cv_pct: null + gpu_util_percent: null + mem_bw_percent: null + nic_gbps: null + goodput_percent: null + cost_per_token_usd: null + cost_per_request_usd: null + notes: + - "Retained four-block ABBA median: baseline 15.9415460075 ms, optimized 15.7822739955 ms, 1.0100918291x. This was near parity and predates this candidate." + - "The retained ABBA receipt reports exact full-output equality, but its declared tolerance was [0.1, 1.0]; the new candidate still requires a fresh exact comparison." + +hot_path_model: + path: "ch04.optimized_pipeline_parallel_1f1b:_run_worker -> _run_contiguous_1f1b_iterations -> run_1f1b_iteration" + phase: steady_state + invocation_frequency: per_rank + invocations_per_primary_unit: 1 + estimated_costs: + bytes_read: null + bytes_written: null + host_device_bytes: 0 + allocations_or_materializations: "fixed 8 forward/8 backward receive tensor slots per applicable rank; repeated Python references only" + kernel_launches: "retained measured range: 768 compute + 54 P2P kernels across both ranks" + synchronizations_or_locks: "one schedule fill and drain plus blocking waits required by each 1F1B dependency" + api_process_or_rpc_crossings: "two torchrun ranks" + storage_or_network_operations: "48 logical 64 MiB tensor transfers across the link for 3 measured iterations" + operation_cost_source: "retained Wave72 Nsys cuda_gpu_kern_sum.csv and nvtx_gpu_proj_sum.csv" + expected_dominant_cost: "GEMM: 157.749088 ms/1024 kernels; P2P: 54.699815 ms/144 kernels; ReLU: 22.735644 ms/1024 kernels over setup, warmup, and measured work" + best_case_primary_kpi_improvement_pct: null + +evidence: + baseline_artifact: "/Users/admin/.codex/artifacts/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-abba-results/receipt.json" + profile_artifacts: + - "/Users/admin/.codex/artifacts/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-nsys-baseline-results/exports-v2" + - "/Users/admin/.codex/artifacts/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-nsys-optimized-results/exports-v2" + allocation_or_memory_profile: "none" + hardware_counter_artifacts: [] + static_review_findings: + - "The retained measured optimized range has 822 GPU operations: 384 GEMMs, 384 ReLUs, and 54 P2P kernels across two ranks." + - "Three separate 8-microbatch schedules require 54 P2P kernels; one continuous 24-microbatch schedule is expected to require 50 while preserving 48 logical tensor transfers." + - "Across five warmup plus three measured iterations, the trace has 144 P2P kernels; batching warmup and measured phases separately predicts 132." + +candidate_optimizations: + - type: api_shape + change: "run repeated warmup and measured iterations as two continuous 1F1B schedules" + mechanism: "amortize repeated fill/drain boundaries without changing compute or P2P payload work" + expected_measured_p2p_kernel_instances: 50 + expected_total_trace_p2p_kernel_instances: 132 + status: "implemented; CPU/Gloo correctness verified; B200 timing and profile unmeasured" + +known_issues_or_hypotheses: + - "The predicted four-kernel measured reduction is only 0.49% of retained measured GPU operations, so a wall-time gain may be below run variance." + - "Boundary P2P kernels have a long tail in the retained trace, but a fresh candidate trace is required to determine whether continuity reduces that tail." + - "Candidate profile capture is absent; no performance gain is claimed." + - "Continuity is valid for this fixed-state toy worker only; a training loop must drain before an optimizer update and cannot reuse this helper across update boundaries." + +risk_guardrails: + correctness_tests: + - "tests/test_pipeline_1f1b_contiguous_repeats.py" + - "tests/test_pipeline_1f1b_schedule.py" + - "tests/test_pipeline_stage_execution.py" + - "fresh B200 ABBA with exact full-output comparison" + change_windows: + - "run only through the root-owned serial B200 queue" diff --git a/code/tests/test_pipeline_1f1b_contiguous_repeats.py b/code/tests/test_pipeline_1f1b_contiguous_repeats.py new file mode 100644 index 000000000..931124b8c --- /dev/null +++ b/code/tests/test_pipeline_1f1b_contiguous_repeats.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist + +from ch04.optimized_pipeline_parallel_1f1b import _run_contiguous_1f1b_iterations +from ch04.pipeline_parallel_common import ( + PipelineIterationCapture, + verify_and_concatenate_pipeline_capture, +) + +_WORLD_SIZE = 2 +_MICROBATCHES = 4 +_ITERATIONS = 3 +_SHAPE = (_MICROBATCHES, 1, 2) + + +def _continuous_gloo_worker(rank: int, rendezvous: str, result_dir: str) -> None: + os.environ["GLOO_SOCKET_IFNAME"] = "lo0" if sys.platform == "darwin" else "lo" + dist.init_process_group( + "gloo", + init_method=rendezvous, + rank=rank, + world_size=_WORLD_SIZE, + ) + try: + rank0_input = ( + torch.arange(torch.tensor(_SHAPE).prod().item(), dtype=torch.float32) + .reshape(_SHAPE) + .to(torch.bfloat16) + ) + forward_calls = 0 + backward_calls = 0 + + def forward_step(value: torch.Tensor) -> torch.Tensor: + nonlocal forward_calls + forward_calls += 1 + return value + float(rank + 1) + + def backward_step( + _activation: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + nonlocal backward_calls + backward_calls += 1 + return value + float((rank + 1) * 10) + + recv_forward = ( + [torch.empty((1, 1, 2), dtype=torch.bfloat16) for _ in range(_MICROBATCHES)] + if rank > 0 + else [] + ) + recv_backward = ( + [torch.empty((1, 1, 2), dtype=torch.bfloat16) for _ in range(_MICROBATCHES)] + if rank < _WORLD_SIZE - 1 + else [] + ) + scheduled_microbatches = _ITERATIONS * _MICROBATCHES + capture = PipelineIterationCapture.create( + _MICROBATCHES, + first_microbatch_index=scheduled_microbatches - _MICROBATCHES, + ) + + _run_contiguous_1f1b_iterations( + rank=rank, + world_size=_WORLD_SIZE, + micro_batches_per_iteration=_MICROBATCHES, + iteration_count=_ITERATIONS, + get_rank0_microbatch=lambda index: rank0_input[index : index + 1], + recv_forward_buffers=recv_forward, + recv_backward_buffers=recv_backward, + forward_step=forward_step, + backward_step=backward_step, + activation_slots=[None] * max(_WORLD_SIZE - rank - 1, 1), + capture=capture, + ) + + forward_stages = [ + lambda value, increment=float(stage + 1): value + increment + for stage in range(_WORLD_SIZE) + ] + backward_stages = [ + lambda value, increment=float((stage + 1) * 10): value + increment + for stage in range(_WORLD_SIZE) + ] + verify_input, verify_output = verify_and_concatenate_pipeline_capture( + rank=rank, + capture=capture, + forward_stages=forward_stages, + backward_stages=backward_stages, + tolerance=(0.0, 0.0), + ) + torch.save( + { + "forward_calls": forward_calls, + "backward_calls": backward_calls, + "verify_input": verify_input, + "verify_output": verify_output, + "unique_forward_buffers": len({tensor.data_ptr() for tensor in recv_forward}), + "unique_backward_buffers": len({tensor.data_ptr() for tensor in recv_backward}), + }, + Path(result_dir) / f"rank-{rank}.pt", + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available(), reason="torch.distributed is unavailable") +def test_contiguous_repeats_preserve_all_work_and_last_iteration_output( + tmp_path: Path, +) -> None: + rendezvous_path = tmp_path / "gloo-rendezvous" + torch.multiprocessing.spawn( + _continuous_gloo_worker, + args=(f"file://{rendezvous_path}", str(tmp_path)), + nprocs=_WORLD_SIZE, + join=True, + daemon=False, + ) + + expected_calls = _ITERATIONS * _MICROBATCHES + rank0 = torch.load(tmp_path / "rank-0.pt", weights_only=True) + rank1 = torch.load(tmp_path / "rank-1.pt", weights_only=True) + for result in (rank0, rank1): + assert result["forward_calls"] == expected_calls + assert result["backward_calls"] == expected_calls + assert result["verify_input"].shape == _SHAPE + assert result["verify_output"].shape == _SHAPE + + original = ( + torch.arange(torch.tensor(_SHAPE).prod().item(), dtype=torch.float32) + .reshape(_SHAPE) + .to(torch.bfloat16) + ) + torch.testing.assert_close(rank0["verify_input"], original, rtol=0.0, atol=0.0) + torch.testing.assert_close(rank0["verify_output"], original + 33.0, rtol=0.0, atol=0.0) + torch.testing.assert_close(rank1["verify_input"], original + 1.0, rtol=0.0, atol=0.0) + torch.testing.assert_close(rank1["verify_output"], original + 23.0, rtol=0.0, atol=0.0) + + assert rank0["unique_backward_buffers"] == _MICROBATCHES + assert rank1["unique_forward_buffers"] == _MICROBATCHES + + +def test_capture_window_rejects_negative_start_and_requires_complete_window() -> None: + with pytest.raises(ValueError, match="must be non-negative"): + PipelineIterationCapture.create(1, first_microbatch_index=-1) + + capture = PipelineIterationCapture.create(1, first_microbatch_index=4) + with pytest.raises(RuntimeError, match="capture is incomplete"): + capture.concatenate() From aea503d9a88bd0326d93d27d5f480507febc076c Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 02:40:30 -0700 Subject: [PATCH 07/19] feat: expose matched TE precision batch sizes for throughput sweeps --- code/ch13/baseline_precisionfp8_te.py | 31 ++++++- code/ch13/optimized_precisionfp8_te.py | 31 ++++++- code/ch13/te_runtime_common.py | 25 ++++++ ...ecisionfp8_te_batch_sweep_followthrough.py | 85 +++++++++++++++++++ 4 files changed, 170 insertions(+), 2 deletions(-) create mode 100644 code/tests/test_precisionfp8_te_batch_sweep_followthrough.py diff --git a/code/ch13/baseline_precisionfp8_te.py b/code/ch13/baseline_precisionfp8_te.py index 03955f255..b1c33ae4b 100644 --- a/code/ch13/baseline_precisionfp8_te.py +++ b/code/ch13/baseline_precisionfp8_te.py @@ -13,8 +13,10 @@ import torch.nn as nn from ch13.te_runtime_common import ( + TE_PRECISION_DEFAULT_BATCH_SIZE, ensure_te_runtime_initialized, get_te_precision_output_tolerances, + parse_te_precision_batch_size, ) from core.benchmark.verification_mixin import VerificationPayloadMixin from core.harness.benchmark_harness import ( @@ -73,8 +75,9 @@ def __init__(self): self.targets: Optional[torch.Tensor] = None self.optimizer: Optional[torch.optim.Optimizer] = None self.criterion: Optional[nn.Module] = None - self.batch_size = 256 + self.batch_size = TE_PRECISION_DEFAULT_BATCH_SIZE self.hidden_dim = 4096 + self._target_override_error: str | None = None tokens = self.batch_size * self.hidden_dim self._workload = WorkloadMetadata( requests_per_iteration=1.0, @@ -92,7 +95,33 @@ def __init__(self): self._verify_input: Optional[torch.Tensor] = None self._verify_target: Optional[torch.Tensor] = None + def _set_batch_size(self, batch_size: int) -> None: + self.batch_size = batch_size + tokens = self.batch_size * self.hidden_dim + self._workload = WorkloadMetadata( + requests_per_iteration=1.0, + tokens_per_iteration=float(tokens), + ) + self.register_workload_metadata( + requests_per_iteration=1.0, + tokens_per_iteration=float(tokens), + ) + + def apply_target_overrides(self, argv: list[str]) -> None: + """Apply ``aisp bench --target-extra-arg`` batch configuration.""" + try: + batch_size = parse_te_precision_batch_size(argv, default=self.batch_size) + except ValueError as exc: + # The harness logs and suppresses override-hook exceptions. Retain + # the error so setup still rejects an invalid requested workload. + self._target_override_error = str(exc) + raise + self._target_override_error = None + self._set_batch_size(batch_size) + def setup(self) -> None: + if self._target_override_error is not None: + raise ValueError(f"Invalid target override: {self._target_override_error}") _load_te_linear() model = TEFP16MLP(hidden_dim=self.hidden_dim).to(self.device).train().half() self.model = model diff --git a/code/ch13/optimized_precisionfp8_te.py b/code/ch13/optimized_precisionfp8_te.py index d85150308..46a7ed894 100644 --- a/code/ch13/optimized_precisionfp8_te.py +++ b/code/ch13/optimized_precisionfp8_te.py @@ -9,8 +9,10 @@ from torch.optim import Optimizer from ch13.te_runtime_common import ( + TE_PRECISION_DEFAULT_BATCH_SIZE, ensure_te_runtime_initialized, get_te_precision_output_tolerances, + parse_te_precision_batch_size, ) from core.benchmark.verification_mixin import VerificationPayloadMixin from core.harness.benchmark_harness import ( @@ -83,8 +85,9 @@ def __init__(self): self.optimizer: Optional[Optimizer] = None self.criterion: Optional[nn.Module] = None self.fp8_recipe: Optional[object] = None - self.batch_size = 256 + self.batch_size = TE_PRECISION_DEFAULT_BATCH_SIZE self.hidden_dim = 4096 + self._target_override_error: str | None = None self.compute_dtype = torch.float16 self.input_pool: List[torch.Tensor] = [] self.target_pool: List[torch.Tensor] = [] @@ -105,7 +108,33 @@ def __init__(self): tokens_per_iteration=float(tokens), ) + def _set_batch_size(self, batch_size: int) -> None: + self.batch_size = batch_size + tokens = self.batch_size * self.hidden_dim + self._workload = WorkloadMetadata( + requests_per_iteration=1.0, + tokens_per_iteration=float(tokens), + ) + self.register_workload_metadata( + requests_per_iteration=1.0, + tokens_per_iteration=float(tokens), + ) + + def apply_target_overrides(self, argv: list[str]) -> None: + """Apply ``aisp bench --target-extra-arg`` batch configuration.""" + try: + batch_size = parse_te_precision_batch_size(argv, default=self.batch_size) + except ValueError as exc: + # The harness logs and suppresses override-hook exceptions. Retain + # the error so setup still rejects an invalid requested workload. + self._target_override_error = str(exc) + raise + self._target_override_error = None + self._set_batch_size(batch_size) + def setup(self) -> None: + if self._target_override_error is not None: + raise ValueError(f"Invalid target override: {self._target_override_error}") _, _, te_recipe_module = _load_transformer_engine() # Transformer Engine 2.x defaults to HYBRID format (E4M3 forward, # E5M2 backward); keep a long amax history and the standard max policy. diff --git a/code/ch13/te_runtime_common.py b/code/ch13/te_runtime_common.py index 11610f823..d61bd5f3a 100644 --- a/code/ch13/te_runtime_common.py +++ b/code/ch13/te_runtime_common.py @@ -2,7 +2,9 @@ from __future__ import annotations +import argparse import ctypes +from collections.abc import Iterable from functools import lru_cache from pathlib import Path @@ -25,6 +27,29 @@ "parameter.fc2.bias": (0.001, 0.00005), } +TE_PRECISION_DEFAULT_BATCH_SIZE = 256 + + +def parse_te_precision_batch_size( + argv: Iterable[str], + *, + default: int = TE_PRECISION_DEFAULT_BATCH_SIZE, +) -> int: + """Parse the shared batch override used by both TE precision arms.""" + parser = argparse.ArgumentParser( + add_help=False, + allow_abbrev=False, + exit_on_error=False, + ) + parser.add_argument("--batch-size", type=int, default=default) + try: + args, _ = parser.parse_known_args(list(argv)) + except (argparse.ArgumentError, SystemExit) as exc: + raise ValueError("--batch-size must be a positive integer") from exc + if args.batch_size <= 0: + raise ValueError("--batch-size must be a positive integer") + return int(args.batch_size) + def get_te_precision_output_tolerances() -> dict[str, tuple[float, float]]: """Return the calibrated full-output policy for the TE2.18 precision pair.""" diff --git a/code/tests/test_precisionfp8_te_batch_sweep_followthrough.py b/code/tests/test_precisionfp8_te_batch_sweep_followthrough.py new file mode 100644 index 000000000..149b70dd0 --- /dev/null +++ b/code/tests/test_precisionfp8_te_batch_sweep_followthrough.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import pytest + +from ch13.baseline_precisionfp8_te import BaselineTEFP8Benchmark +from ch13.optimized_precisionfp8_te import OptimizedTEFP8Benchmark +from ch13.te_runtime_common import ( + TE_PRECISION_DEFAULT_BATCH_SIZE, +) +from core.harness.benchmark_harness import BenchmarkConfig, BenchmarkHarness + +BENCHMARK_TYPES = (BaselineTEFP8Benchmark, OptimizedTEFP8Benchmark) +FROZEN_TOLERANCES = { + "prediction": (0.4, 1.0), + "parameter.fc1.weight": (0.001, 0.00075), + "parameter.fc1.bias": (0.001, 0.00005), + "parameter.fc2.weight": (0.001, 0.00075), + "parameter.fc2.bias": (0.001, 0.00005), +} + + +@pytest.mark.parametrize("benchmark_type", BENCHMARK_TYPES) +def test_te_precision_batch_override_keeps_256_default_and_frozen_contract( + benchmark_type, +) -> None: + benchmark = benchmark_type() + + assert benchmark.batch_size == TE_PRECISION_DEFAULT_BATCH_SIZE == 256 + assert benchmark._workload.tokens_per_iteration == 256 * 4096 + assert benchmark.get_workload_metadata().tokens_per_iteration == 256 * 4096 + assert benchmark.get_output_tolerances() == FROZEN_TOLERANCES + config = benchmark.get_config() + assert (config.iterations, config.warmup) == (50, 10) + + +@pytest.mark.parametrize("batch_size", (256, 1024, 4096)) +def test_te_precision_batch_override_is_symmetric_and_updates_workload_metadata( + batch_size: int, +) -> None: + benchmarks = [benchmark_type() for benchmark_type in BENCHMARK_TYPES] + + for benchmark in benchmarks: + benchmark.apply_target_overrides(["--batch-size", str(batch_size)]) + assert benchmark.batch_size == batch_size + expected_tokens = float(batch_size * benchmark.hidden_dim) + assert benchmark._workload.tokens_per_iteration == expected_tokens + assert benchmark.get_workload_metadata().tokens_per_iteration == expected_tokens + + assert benchmarks[0].signature_equivalence_group == benchmarks[1].signature_equivalence_group + assert benchmarks[0].signature_equivalence_ignore_fields == ("precision_flags",) + assert benchmarks[1].signature_equivalence_ignore_fields == ("precision_flags",) + + +@pytest.mark.parametrize("benchmark_type", BENCHMARK_TYPES) +def test_harness_routes_target_extra_batch_argument_to_te_precision_pair( + benchmark_type, +) -> None: + config = BenchmarkConfig( + target_label="ch13:precisionfp8_te", + target_extra_args={ + "ch13:precisionfp8_te": ["--batch-size", "1024"], + }, + ) + benchmark = benchmark_type() + + BenchmarkHarness(config=config)._apply_target_overrides(benchmark, config) + + assert benchmark.batch_size == 1024 + assert benchmark.get_workload_metadata().tokens_per_iteration == 1024 * 4096 + + +@pytest.mark.parametrize("benchmark_type", BENCHMARK_TYPES) +@pytest.mark.parametrize("value", ("0", "-1", "not-an-integer")) +def test_te_precision_batch_override_rejects_invalid_values_without_fallback( + benchmark_type, + value: str, +) -> None: + benchmark = benchmark_type() + + with pytest.raises(ValueError, match="--batch-size must be a positive integer"): + benchmark.apply_target_overrides(["--batch-size", value]) + + assert benchmark.batch_size == 256 + with pytest.raises(ValueError, match="Invalid target override"): + benchmark.setup() From d71c4caf88ece6d1c73cd57eb77c993e7d1edcee Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 03:23:28 -0700 Subject: [PATCH 08/19] fix: apply serving profile arguments and close reused engines --- .../dynamic_router/baseline_dual_pool_vllm.py | 45 +++- .../baseline_dynamic_router_vllm.py | 43 +++- .../optimized_dual_pool_vllm.py | 45 +++- .../optimized_dynamic_router_vllm.py | 43 +++- code/labs/dynamic_router/vllm_runner.py | 213 ++++++++++++++--- ...er_vllm_profile_overrides_followthrough.py | 217 ++++++++++++++++++ 6 files changed, 558 insertions(+), 48 deletions(-) create mode 100644 code/tests/test_dynamic_router_vllm_profile_overrides_followthrough.py diff --git a/code/labs/dynamic_router/baseline_dual_pool_vllm.py b/code/labs/dynamic_router/baseline_dual_pool_vllm.py index b4205b2d5..4a13c8873 100644 --- a/code/labs/dynamic_router/baseline_dual_pool_vllm.py +++ b/code/labs/dynamic_router/baseline_dual_pool_vllm.py @@ -29,6 +29,7 @@ class BaselineDualPoolVllmBenchmark(VerificationPayloadMixin, BaseBenchmark): multi_gpu_required = True _is_deterministic = True input_jitter_bounds = {"prompt_token_ids": (0, 2)} + profile_require_teardown = True def __init__(self) -> None: super().__init__() @@ -38,31 +39,53 @@ def __init__(self) -> None: self._metric_output_buffer: Optional[torch.Tensor] = None self._mode_input: Optional[torch.Tensor] = None self._prompt_token_ids: Optional[torch.Tensor] = None - self._prompt_lengths = vllm_runner.dual_pool_prompt_lengths( - vllm_runner._CLI_ARGS - ) + self._cli_args = vllm_runner._CLI_ARGS + self._target_override_error: Optional[str] = None + self._prompt_lengths: list[int] = [] self._topology = None self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False + self._configure_cli_args(self._cli_args) + + def _configure_cli_args(self, cli_args) -> None: + self._cli_args = cli_args + self._prompt_lengths = vllm_runner.dual_pool_prompt_lengths(cli_args) request_count = len(self._prompt_lengths) self.register_workload_metadata( requests_per_iteration=float(request_count), tokens_per_iteration=float( sum(self._prompt_lengths) - + request_count * max(1, vllm_runner._CLI_ARGS.max_tokens) + + request_count * max(1, cli_args.max_tokens) ), ) + def apply_target_overrides(self, argv: list[str]) -> None: + """Apply one exact ``--target-extra-arg`` vector before setup.""" + if self._engine_session is not None: + raise RuntimeError("vLLM target overrides must be applied before setup()") + try: + cli_args = vllm_runner.parse_vllm_target_overrides(argv) + except ValueError as exc: + self._target_override_error = str(exc) + raise + self._target_override_error = None + self._configure_cli_args(cli_args) + def setup(self) -> None: + if self._target_override_error is not None: + raise ValueError(f"Invalid target override: {self._target_override_error}") self._mode_input = scalar_int_buffer(self, "_mode_input", 0) self._prompt_token_ids = vllm_runner.build_prompt_token_ids( self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + vllm_runner.emit_vllm_profile_runtime_receipt( + type(self).__name__, self._cli_args + ) self._engine_session = vllm_runner.create_dual_pool_vllm_session( "shared", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, warmup_runs=vllm_runner.WARMUP_ITERATIONS, ) @@ -76,7 +99,7 @@ def benchmark_fn(self) -> None: self._summary = run_dual_pool_vllm_with_topology( "shared", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, prompt_token_ids=self._prompt_token_ids, engine_session=self._engine_session, ) @@ -104,6 +127,14 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + receipt_error: Optional[Exception] = None + if self._summary_ready: + try: + vllm_runner.emit_vllm_profile_output_receipt( + type(self).__name__, self._cli_args, self._summary + ) + except Exception as exc: + receipt_error = exc if self._engine_session is not None: self._engine_session.close() self._engine_session = None @@ -115,6 +146,8 @@ def teardown(self) -> None: self._topology = None self._summary_ready = False super().teardown() + if receipt_error is not None: + raise receipt_error def get_config(self) -> Optional[BenchmarkConfig]: return BenchmarkConfig( diff --git a/code/labs/dynamic_router/baseline_dynamic_router_vllm.py b/code/labs/dynamic_router/baseline_dynamic_router_vllm.py index 45a02b252..6ca4fb92b 100644 --- a/code/labs/dynamic_router/baseline_dynamic_router_vllm.py +++ b/code/labs/dynamic_router/baseline_dynamic_router_vllm.py @@ -29,6 +29,7 @@ class BaselineDynamicRouterVllmBenchmark(VerificationPayloadMixin, BaseBenchmark multi_gpu_required = True _is_deterministic = True input_jitter_bounds = {"prompt_token_ids": (0, 2)} + profile_require_teardown = True def __init__(self) -> None: super().__init__() @@ -38,28 +39,52 @@ def __init__(self) -> None: self._metric_output_buffer: Optional[torch.Tensor] = None self._mode_input: Optional[torch.Tensor] = None self._prompt_token_ids: Optional[torch.Tensor] = None - self._prompt_lengths = vllm_runner.routing_prompt_lengths(vllm_runner._CLI_ARGS) + self._cli_args = vllm_runner._CLI_ARGS + self._target_override_error: Optional[str] = None + self._prompt_lengths: list[int] = [] self._topology = None self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False + self._configure_cli_args(self._cli_args) + + def _configure_cli_args(self, cli_args) -> None: + self._cli_args = cli_args + self._prompt_lengths = vllm_runner.routing_prompt_lengths(cli_args) request_count = len(self._prompt_lengths) self.register_workload_metadata( requests_per_iteration=float(request_count), tokens_per_iteration=float( - sum(self._prompt_lengths) + request_count * vllm_runner._CLI_ARGS.max_tokens + sum(self._prompt_lengths) + request_count * cli_args.max_tokens ), ) + def apply_target_overrides(self, argv: list[str]) -> None: + """Apply one exact ``--target-extra-arg`` vector before setup.""" + if self._engine_session is not None: + raise RuntimeError("vLLM target overrides must be applied before setup()") + try: + cli_args = vllm_runner.parse_vllm_target_overrides(argv) + except ValueError as exc: + self._target_override_error = str(exc) + raise + self._target_override_error = None + self._configure_cli_args(cli_args) + def setup(self) -> None: + if self._target_override_error is not None: + raise ValueError(f"Invalid target override: {self._target_override_error}") self._mode_input = scalar_int_buffer(self, "_mode_input", 0) self._prompt_token_ids = vllm_runner.build_prompt_token_ids( self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + vllm_runner.emit_vllm_profile_runtime_receipt( + type(self).__name__, self._cli_args + ) self._engine_session = vllm_runner.create_vllm_routing_session( "baseline", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, warmup_runs=vllm_runner.WARMUP_ITERATIONS, ) @@ -73,7 +98,7 @@ def benchmark_fn(self) -> None: self._summary = run_vllm_routing_with_topology( "baseline", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, prompt_token_ids=self._prompt_token_ids, engine_session=self._engine_session, ) @@ -101,6 +126,14 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + receipt_error: Optional[Exception] = None + if self._summary_ready: + try: + vllm_runner.emit_vllm_profile_output_receipt( + type(self).__name__, self._cli_args, self._summary + ) + except Exception as exc: + receipt_error = exc if self._engine_session is not None: self._engine_session.close() self._engine_session = None @@ -112,6 +145,8 @@ def teardown(self) -> None: self._topology = None self._summary_ready = False super().teardown() + if receipt_error is not None: + raise receipt_error def get_config(self) -> Optional[BenchmarkConfig]: return BenchmarkConfig( diff --git a/code/labs/dynamic_router/optimized_dual_pool_vllm.py b/code/labs/dynamic_router/optimized_dual_pool_vllm.py index 1dc4e0e67..b2ae79b4e 100644 --- a/code/labs/dynamic_router/optimized_dual_pool_vllm.py +++ b/code/labs/dynamic_router/optimized_dual_pool_vllm.py @@ -29,6 +29,7 @@ class OptimizedDualPoolVllmBenchmark(VerificationPayloadMixin, BaseBenchmark): multi_gpu_required = True _is_deterministic = True input_jitter_bounds = {"prompt_token_ids": (0, 2)} + profile_require_teardown = True def __init__(self) -> None: super().__init__() @@ -38,31 +39,53 @@ def __init__(self) -> None: self._metric_output_buffer: Optional[torch.Tensor] = None self._mode_input: Optional[torch.Tensor] = None self._prompt_token_ids: Optional[torch.Tensor] = None - self._prompt_lengths = vllm_runner.dual_pool_prompt_lengths( - vllm_runner._CLI_ARGS - ) + self._cli_args = vllm_runner._CLI_ARGS + self._target_override_error: Optional[str] = None + self._prompt_lengths: list[int] = [] self._topology = None self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False + self._configure_cli_args(self._cli_args) + + def _configure_cli_args(self, cli_args) -> None: + self._cli_args = cli_args + self._prompt_lengths = vllm_runner.dual_pool_prompt_lengths(cli_args) request_count = len(self._prompt_lengths) self.register_workload_metadata( requests_per_iteration=float(request_count), tokens_per_iteration=float( sum(self._prompt_lengths) - + request_count * max(1, vllm_runner._CLI_ARGS.max_tokens) + + request_count * max(1, cli_args.max_tokens) ), ) + def apply_target_overrides(self, argv: list[str]) -> None: + """Apply one exact ``--target-extra-arg`` vector before setup.""" + if self._engine_session is not None: + raise RuntimeError("vLLM target overrides must be applied before setup()") + try: + cli_args = vllm_runner.parse_vllm_target_overrides(argv) + except ValueError as exc: + self._target_override_error = str(exc) + raise + self._target_override_error = None + self._configure_cli_args(cli_args) + def setup(self) -> None: + if self._target_override_error is not None: + raise ValueError(f"Invalid target override: {self._target_override_error}") self._mode_input = scalar_int_buffer(self, "_mode_input", 1) self._prompt_token_ids = vllm_runner.build_prompt_token_ids( self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + vllm_runner.emit_vllm_profile_runtime_receipt( + type(self).__name__, self._cli_args + ) self._engine_session = vllm_runner.create_dual_pool_vllm_session( "dual", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, warmup_runs=vllm_runner.WARMUP_ITERATIONS, ) @@ -76,7 +99,7 @@ def benchmark_fn(self) -> None: self._summary = run_dual_pool_vllm_with_topology( "dual", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, prompt_token_ids=self._prompt_token_ids, engine_session=self._engine_session, ) @@ -104,6 +127,14 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + receipt_error: Optional[Exception] = None + if self._summary_ready: + try: + vllm_runner.emit_vllm_profile_output_receipt( + type(self).__name__, self._cli_args, self._summary + ) + except Exception as exc: + receipt_error = exc if self._engine_session is not None: self._engine_session.close() self._engine_session = None @@ -115,6 +146,8 @@ def teardown(self) -> None: self._topology = None self._summary_ready = False super().teardown() + if receipt_error is not None: + raise receipt_error def get_config(self) -> Optional[BenchmarkConfig]: return BenchmarkConfig( diff --git a/code/labs/dynamic_router/optimized_dynamic_router_vllm.py b/code/labs/dynamic_router/optimized_dynamic_router_vllm.py index 887f6f7e2..4e5adf7b6 100644 --- a/code/labs/dynamic_router/optimized_dynamic_router_vllm.py +++ b/code/labs/dynamic_router/optimized_dynamic_router_vllm.py @@ -29,6 +29,7 @@ class OptimizedDynamicRouterVllmBenchmark(VerificationPayloadMixin, BaseBenchmar multi_gpu_required = True _is_deterministic = True input_jitter_bounds = {"prompt_token_ids": (0, 2)} + profile_require_teardown = True def __init__(self) -> None: super().__init__() @@ -38,28 +39,52 @@ def __init__(self) -> None: self._metric_output_buffer: Optional[torch.Tensor] = None self._mode_input: Optional[torch.Tensor] = None self._prompt_token_ids: Optional[torch.Tensor] = None - self._prompt_lengths = vllm_runner.routing_prompt_lengths(vllm_runner._CLI_ARGS) + self._cli_args = vllm_runner._CLI_ARGS + self._target_override_error: Optional[str] = None + self._prompt_lengths: list[int] = [] self._topology = None self._engine_session: Optional[vllm_runner.VllmEngineSession] = None self._summary_ready = False + self._configure_cli_args(self._cli_args) + + def _configure_cli_args(self, cli_args) -> None: + self._cli_args = cli_args + self._prompt_lengths = vllm_runner.routing_prompt_lengths(cli_args) request_count = len(self._prompt_lengths) self.register_workload_metadata( requests_per_iteration=float(request_count), tokens_per_iteration=float( - sum(self._prompt_lengths) + request_count * vllm_runner._CLI_ARGS.max_tokens + sum(self._prompt_lengths) + request_count * cli_args.max_tokens ), ) + def apply_target_overrides(self, argv: list[str]) -> None: + """Apply one exact ``--target-extra-arg`` vector before setup.""" + if self._engine_session is not None: + raise RuntimeError("vLLM target overrides must be applied before setup()") + try: + cli_args = vllm_runner.parse_vllm_target_overrides(argv) + except ValueError as exc: + self._target_override_error = str(exc) + raise + self._target_override_error = None + self._configure_cli_args(cli_args) + def setup(self) -> None: + if self._target_override_error is not None: + raise ValueError(f"Invalid target override: {self._target_override_error}") self._mode_input = scalar_int_buffer(self, "_mode_input", 2) self._prompt_token_ids = vllm_runner.build_prompt_token_ids( self._prompt_lengths ) self._topology = detect_topology(max_gpus=torch.cuda.device_count()) + vllm_runner.emit_vllm_profile_runtime_receipt( + type(self).__name__, self._cli_args + ) self._engine_session = vllm_runner.create_vllm_routing_session( "optimized", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, warmup_runs=vllm_runner.WARMUP_ITERATIONS, ) @@ -73,7 +98,7 @@ def benchmark_fn(self) -> None: self._summary = run_vllm_routing_with_topology( "optimized", topology_snapshot=self._topology, - cli_args=vllm_runner._CLI_ARGS, + cli_args=self._cli_args, prompt_token_ids=self._prompt_token_ids, engine_session=self._engine_session, ) @@ -101,6 +126,14 @@ def capture_verification_payload(self) -> None: ) def teardown(self) -> None: + receipt_error: Optional[Exception] = None + if self._summary_ready: + try: + vllm_runner.emit_vllm_profile_output_receipt( + type(self).__name__, self._cli_args, self._summary + ) + except Exception as exc: + receipt_error = exc if self._engine_session is not None: self._engine_session.close() self._engine_session = None @@ -112,6 +145,8 @@ def teardown(self) -> None: self._topology = None self._summary_ready = False super().teardown() + if receipt_error is not None: + raise receipt_error def get_config(self) -> Optional[BenchmarkConfig]: return BenchmarkConfig( diff --git a/code/labs/dynamic_router/vllm_runner.py b/code/labs/dynamic_router/vllm_runner.py index 21d8271eb..e7691e15d 100644 --- a/code/labs/dynamic_router/vllm_runner.py +++ b/code/labs/dynamic_router/vllm_runner.py @@ -8,15 +8,18 @@ from __future__ import annotations import argparse +import hashlib import importlib import importlib.metadata import io import json +import os import sys import time from contextlib import redirect_stdout from dataclasses import dataclass from functools import wraps +from pathlib import Path from typing import Dict, List, Optional, Sequence, Set, Tuple import torch @@ -52,6 +55,100 @@ def _skip(reason: str) -> None: _EXPECTED_FLASHINFER_DIST_VERSION = _SERVING_STACK_PINS.flashinfer_version WARMUP_ITERATIONS = 5 STEADY_STATE_ITERATIONS = 3 +VLLM_PROFILE_RECEIPT_DIR_ENV = "AISP_VLLM_PROFILE_RECEIPT_DIR" +VLLM_PROFILE_RUNTIME_SCHEMA = "aisp.dynamic-router-vllm-profile-runtime.v1" +VLLM_PROFILE_OUTPUT_SCHEMA = "aisp.dynamic-router-vllm-profile-output.v1" +VLLM_PROFILE_LIFECYCLE_SCHEMA = "aisp.dynamic-router-vllm-profile-lifecycle.v1" + + +def _profile_receipt_dir() -> Optional[Path]: + raw = os.environ.get(VLLM_PROFILE_RECEIPT_DIR_ENV, "").strip() + if not raw: + return None + path = Path(raw).resolve() + if not path.is_dir(): + raise RuntimeError( + f"{VLLM_PROFILE_RECEIPT_DIR_ENV} must name an existing directory: {path}" + ) + return path + + +def _write_profile_receipt(filename: str, payload: Dict[str, object]) -> None: + directory = _profile_receipt_dir() + if directory is None: + return + destination = directory / filename + temporary = directory / f".{filename}.{os.getpid()}.tmp" + temporary.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + os.replace(temporary, destination) + + +def emit_vllm_profile_runtime_receipt( + benchmark_name: str, + cli_args: argparse.Namespace, +) -> None: + """Retain provenance from the process that creates the profiled engines.""" + if _profile_receipt_dir() is None: + return + from core.benchmark.run_manifest import capture_runtime_provenance + + runtime = capture_runtime_provenance() + devices = [] + for index in range(torch.cuda.device_count()): + properties = torch.cuda.get_device_properties(index) + devices.append( + { + "logical_index": index, + "name": properties.name, + "uuid": str(getattr(properties, "uuid", "")) or None, + "compute_capability": f"{properties.major}.{properties.minor}", + } + ) + _write_profile_receipt( + "runtime-provenance.json", + { + "schema": VLLM_PROFILE_RUNTIME_SCHEMA, + "benchmark": benchmark_name, + "cli_args": vars(cli_args), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "vllm_batch_invariant": os.environ.get("VLLM_BATCH_INVARIANT"), + "devices": devices, + "runtime_provenance": runtime.model_dump(mode="json"), + }, + ) + + +def emit_vllm_profile_output_receipt( + benchmark_name: str, + cli_args: argparse.Namespace, + summary: Dict[str, object], +) -> None: + """Retain the full profiled-call model output after the measured range closes.""" + if _profile_receipt_dir() is None: + return + output = summary.get(VERIFICATION_OUTPUT_KEY) + if not isinstance(output, (list, tuple)) or not output: + raise RuntimeError("profile output receipt requires framed generated token ids") + token_ids = [int(value) for value in output] + encoded = json.dumps(token_ids, separators=(",", ":")).encode("utf-8") + scalar_metrics = { + key: value + for key, value in summary.items() + if key != VERIFICATION_OUTPUT_KEY and isinstance(value, (bool, int, float, str)) + } + _write_profile_receipt( + "profile-output.json", + { + "schema": VLLM_PROFILE_OUTPUT_SCHEMA, + "benchmark": benchmark_name, + "cli_args": vars(cli_args), + "scalar_metrics": scalar_metrics, + "framed_token_ids": token_ids, + "framed_token_ids_sha256": hashlib.sha256(encoded).hexdigest(), + }, + ) def _is_vllm_abi_mismatch_error(exc: BaseException) -> bool: @@ -127,8 +224,8 @@ def _assert_vllm_runtime_ready() -> None: _skip(_format_vllm_import_error(exc)) -def _parse_cli_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(add_help=False) +def _build_cli_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(add_help=False, exit_on_error=False) parser.add_argument("--model", type=str, help="Local HF model path/id for vLLM.") parser.add_argument( "--attention-backend", type=str, default=None, @@ -149,7 +246,60 @@ def _parse_cli_args() -> argparse.Namespace: action="store_true", help="Drive vLLM V1 EngineCore directly with the optimized polling loop (Inproc only).", ) - return parser.parse_known_args()[0] + return parser + + +def _validate_cli_args(args: argparse.Namespace) -> argparse.Namespace: + positive_fields = ( + "req_count", + "max_tokens", + "long_prompt_tokens", + "short_prompt_tokens", + "prefill_ctx_thresh", + ) + for field in positive_fields: + value = getattr(args, field) + if value <= 0: + raise ValueError(f"--{field.replace('_', '-')} must be positive") + request_mix_fields = ("prefill_burst", "decode_requests", "continue_requests") + for field in request_mix_fields: + if getattr(args, field) < 0: + raise ValueError(f"--{field.replace('_', '-')} must be non-negative") + if not any(getattr(args, field) > 0 for field in request_mix_fields): + raise ValueError("dual-pool request mix must contain at least one request") + for field in ("prefill_gpus", "decode_gpus"): + raw = getattr(args, field) + if raw is None: + continue + parts = [part.strip() for part in raw.split(",")] + if not parts or any(not part.isdigit() for part in parts): + raise ValueError( + f"--{field.replace('_', '-')} must be a comma-separated list of non-negative GPU ids" + ) + if len(set(parts)) != len(parts): + raise ValueError(f"--{field.replace('_', '-')} must not contain duplicate GPU ids") + return args + + +def _parse_cli_args( + argv: Optional[Sequence[str]] = None, + *, + reject_unknown: bool = False, +) -> argparse.Namespace: + parser = _build_cli_parser() + try: + args, unknown = parser.parse_known_args(argv) + except argparse.ArgumentError as exc: + raise ValueError(str(exc)) from exc + if reject_unknown and unknown: + raise ValueError(f"Unrecognized vLLM target arguments: {unknown}") + return _validate_cli_args(args) + + +def parse_vllm_target_overrides(argv: Sequence[str]) -> argparse.Namespace: + """Parse one harness target override vector without mutating module globals.""" + + return _parse_cli_args(list(argv), reject_unknown=True) _CLI_ARGS = _parse_cli_args() @@ -765,35 +915,42 @@ def lifecycle_metrics(self) -> Dict[str, float]: def _emit_lifecycle(self, disposition: str, errors: Sequence[str]) -> None: failure = self._primary_failure - print( - json.dumps( - { - "event": "vllm_engine_lifecycle", - "disposition": disposition, - "workload_kind": self.workload_kind, - "mode": self.mode, - "setup_engine_startup_ms": self.engine_startup_ms, - "engine_startup_ms": self.engine_startup_ms, - "warmup_request_processing_ms": self._phase_durations_ms["warmup"], - "steady_state_request_processing_ms": self._phase_durations_ms[ - "steady_state" - ], - "failed_runs": self._failed_runs, - "engine_teardown_ms": self._teardown_ms, - "end_to_end_ms": self._end_to_end_ms, - "request_state_reset_per_run": True, - "primary_failure": ( - {"type": type(failure).__name__, "message": str(failure)} - if failure is not None - else None - ), - "shutdown_errors": list(errors), - }, - sort_keys=True, + payload = { + "schema": VLLM_PROFILE_LIFECYCLE_SCHEMA, + "event": "vllm_engine_lifecycle", + "disposition": disposition, + "workload_kind": self.workload_kind, + "mode": self.mode, + "setup_engine_startup_ms": self.engine_startup_ms, + "engine_startup_ms": self.engine_startup_ms, + "warmup_request_processing_ms": self._phase_durations_ms["warmup"], + "steady_state_request_processing_ms": self._phase_durations_ms[ + "steady_state" + ], + "failed_runs": self._failed_runs, + "engine_teardown_ms": self._teardown_ms, + "end_to_end_ms": self._end_to_end_ms, + "request_state_reset_per_run": True, + "primary_failure": ( + {"type": type(failure).__name__, "message": str(failure)} + if failure is not None + else None ), + "shutdown_errors": list(errors), + } + print( + json.dumps(payload, sort_keys=True), file=sys.stderr, flush=True, ) + try: + _write_profile_receipt("lifecycle.json", payload) + except Exception as exc: + print( + f"[profile_warning] Failed to retain vLLM lifecycle receipt: {exc}", + file=sys.stderr, + flush=True, + ) def close(self, *, preserve_primary_error: bool = False) -> List[str]: if self._closed: diff --git a/code/tests/test_dynamic_router_vllm_profile_overrides_followthrough.py b/code/tests/test_dynamic_router_vllm_profile_overrides_followthrough.py new file mode 100644 index 000000000..0220c20b9 --- /dev/null +++ b/code/tests/test_dynamic_router_vllm_profile_overrides_followthrough.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import json +import os +import sys +from dataclasses import replace +from pathlib import Path + +import pytest + +from core.harness import run_benchmarks +from core.profiling.profiler_wrapper import render_nsys_python_profile_wrapper +from labs.dynamic_router import vllm_runner +from labs.dynamic_router.baseline_dual_pool_vllm import ( + BaselineDualPoolVllmBenchmark, +) +from labs.dynamic_router.baseline_dynamic_router_vllm import ( + BaselineDynamicRouterVllmBenchmark, +) +from labs.dynamic_router.optimized_dual_pool_vllm import ( + OptimizedDualPoolVllmBenchmark, +) +from labs.dynamic_router.optimized_dynamic_router_vllm import ( + OptimizedDynamicRouterVllmBenchmark, +) + +TARGET_LABEL = "labs/dynamic_router:dual_pool_vllm" +TARGET_ARGV = [ + "--model", + "/models/gpt-oss-20b", + "--prefill-gpus", + "0", + "--decode-gpus", + "1", + "--attention-backend", + "TRITON_ATTN", + "--max-tokens", + "3", + "--long-prompt-tokens", + "32", + "--short-prompt-tokens", + "4", + "--prefill-burst", + "2", + "--decode-requests", + "3", + "--continue-requests", + "4", +] + + +@pytest.mark.parametrize( + "benchmark_type,expected_requests,expected_tokens", + [ + (BaselineDualPoolVllmBenchmark, 9, 119), + (OptimizedDualPoolVllmBenchmark, 9, 119), + ], +) +def test_dual_pool_target_overrides_are_instance_local_and_refresh_workload( + benchmark_type, expected_requests: int, expected_tokens: int +) -> None: + global_args_before = vars(vllm_runner._CLI_ARGS).copy() + benchmark = benchmark_type() + + benchmark.apply_target_overrides(TARGET_ARGV) + + assert vars(vllm_runner._CLI_ARGS) == global_args_before + assert benchmark._cli_args.model == "/models/gpt-oss-20b" + assert benchmark._cli_args.prefill_gpus == "0" + assert benchmark._cli_args.decode_gpus == "1" + assert benchmark._cli_args.attention_backend == "TRITON_ATTN" + assert len(benchmark._prompt_lengths) == expected_requests + metadata = benchmark.get_workload_metadata() + assert metadata is not None + assert metadata.requests_per_iteration == expected_requests + assert metadata.tokens_per_iteration == expected_tokens + assert benchmark.profile_require_teardown is True + + +@pytest.mark.parametrize( + "benchmark_type", + [BaselineDynamicRouterVllmBenchmark, OptimizedDynamicRouterVllmBenchmark], +) +def test_router_target_overrides_share_parser_without_global_state( + benchmark_type, +) -> None: + benchmark = benchmark_type() + argv = [ + "--model", + "/models/gpt-oss-20b", + "--decode-gpus", + "0,1", + "--attention-backend", + "TRITON_ATTN", + "--req-count", + "7", + "--max-tokens", + "3", + ] + + benchmark.apply_target_overrides(argv) + + assert benchmark._cli_args.req_count == 7 + assert benchmark._prompt_lengths == [64] * 7 + metadata = benchmark.get_workload_metadata() + assert metadata is not None + assert metadata.requests_per_iteration == 7 + assert metadata.tokens_per_iteration == 7 * (64 + 3) + assert benchmark.profile_require_teardown is True + + +@pytest.mark.parametrize( + "argv,error", + [ + (["--unknown-vllm-option", "1"], "Unrecognized vLLM target arguments"), + (["--max-tokens", "0"], "--max-tokens must be positive"), + (["--decode-gpus", "1,1"], "must not contain duplicate GPU ids"), + (["--prefill-gpus", "gpu0"], "comma-separated list"), + ], +) +def test_invalid_target_override_is_retained_for_fail_closed_setup( + argv: list[str], error: str +) -> None: + benchmark = BaselineDualPoolVllmBenchmark() + + with pytest.raises(ValueError, match=error): + benchmark.apply_target_overrides(argv) + with pytest.raises(ValueError, match="Invalid target override"): + benchmark.setup() + + +def test_real_nsys_wrapper_plan_embeds_target_override_and_teardown_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + benchmark = BaselineDualPoolVllmBenchmark() + config = replace( + benchmark.get_config(), + target_label=TARGET_LABEL, + target_extra_args={TARGET_LABEL: list(TARGET_ARGV)}, + profile_env_overrides={"VLLM_BATCH_INVARIANT": "1"}, + validity_profile="portable", + lock_gpu_clocks=True, + ) + target_argv = run_benchmarks._resolve_target_override_argv(config) + assert target_argv == TARGET_ARGV + benchmark_path = ( + Path(__file__).resolve().parents[1] + / "labs/dynamic_router/baseline_dual_pool_vllm.py" + ) + source = render_nsys_python_profile_wrapper( + benchmark_path=benchmark_path, + nvtx_includes=["compute_kernel:profile"], + target_label=TARGET_LABEL, + target_override_argv=target_argv, + validity_profile=config.validity_profile, + lock_gpu_clocks_flag=False, + gpu_sm_clock_mhz=None, + gpu_mem_clock_mhz=None, + ) + compile(source, "", "exec") + assert "_apply_overrides(list(_target_override_argv))" in source + assert "profile_require_teardown" in source + + repo_root = Path(run_benchmarks.__file__).resolve().parents[2] + monkeypatch.delenv("PYTHONNOUSERSITE", raising=False) + with run_benchmarks._temporary_python_profile_launch( + source, + chapter_dir=benchmark_path.parent, + repo_root=repo_root, + config=config, + benchmark=benchmark, + ) as (wrapper_path, command, env, use_torchrun): + assert wrapper_path.is_file() + assert command == [sys.executable, str(wrapper_path)] + assert use_torchrun is False + assert env["VLLM_BATCH_INVARIANT"] == "1" + assert str(repo_root) in env["PYTHONPATH"].split(os.pathsep) + + +def test_profile_output_and_lifecycle_receipts_retain_full_evidence( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv(vllm_runner.VLLM_PROFILE_RECEIPT_DIR_ENV, str(tmp_path)) + args = vllm_runner.parse_vllm_target_overrides(TARGET_ARGV) + summary = { + "mode": "shared", + "requests": 2, + "completed": 2, + "ttft_ms_p95": 1.25, + "_verification_output_token_ids": [2, 11, 12, 1, 13], + } + + vllm_runner.emit_vllm_profile_output_receipt("sentinel", args, summary) + output = json.loads( + (tmp_path / "profile-output.json").read_text(encoding="utf-8") + ) + assert output["schema"] == vllm_runner.VLLM_PROFILE_OUTPUT_SCHEMA + assert output["framed_token_ids"] == [2, 11, 12, 1, 13] + assert output["scalar_metrics"]["ttft_ms_p95"] == 1.25 + assert len(output["framed_token_ids_sha256"]) == 64 + + session = vllm_runner.VllmEngineSession.__new__(vllm_runner.VllmEngineSession) + session._primary_failure = None + session.workload_kind = "dual_pool" + session.mode = "shared" + session.engine_startup_ms = 1.0 + session._phase_durations_ms = {"warmup": [2.0], "steady_state": [3.0]} + session._failed_runs = [] + session._teardown_ms = 4.0 + session._end_to_end_ms = 10.0 + session._emit_lifecycle("completed", []) + lifecycle = json.loads( + (tmp_path / "lifecycle.json").read_text(encoding="utf-8") + ) + assert lifecycle["schema"] == vllm_runner.VLLM_PROFILE_LIFECYCLE_SCHEMA + assert lifecycle["disposition"] == "completed" + assert lifecycle["shutdown_errors"] == [] From d36ef3dda4c1cdf62500b3c053a5719b6b5dfb0f Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 03:30:03 -0700 Subject: [PATCH 09/19] fix: honor explicit Nsight Systems warmup and capture counts --- code/core/harness/run_benchmarks.py | 11 +++++ code/core/profiling/profiler_wrapper.py | 56 +++++++++++++++++++++---- code/tests/test_profiler_wrapper.py | 48 +++++++++++++++++++++ 3 files changed, 106 insertions(+), 9 deletions(-) diff --git a/code/core/harness/run_benchmarks.py b/code/core/harness/run_benchmarks.py index 64855e82c..6d2e1b36e 100755 --- a/code/core/harness/run_benchmarks.py +++ b/code/core/harness/run_benchmarks.py @@ -96,6 +96,7 @@ ) from core.profiling.metrics_extractor import inspect_ncu_app_range_report from core.profiling.profiler_wrapper import ( + _resolve_wrapper_loop_budget, render_ncu_python_profile_wrapper, render_nsys_python_profile_wrapper, render_torch_python_profile_wrapper, @@ -3203,6 +3204,14 @@ def profile_python_benchmark( gpu_mem_clock_mhz = validity_view.gpu_mem_clock_mhz if validity_view else None target_label = getattr(bench_config, "target_label", None) if bench_config else None target_override_argv = _resolve_target_override_argv(bench_config) + if bench_config is not None: + profiling_warmup, profiling_iterations = _resolve_wrapper_loop_budget( + bench_config, + default_warmup=1, + default_iterations=1, + ) + else: + profiling_warmup, profiling_iterations = (1, 1) # chapter_dir points to e.g. /ch10 or /labs/; use global # repository root for package imports like `labs.*` and `core.*`. repo_root = Path(__file__).resolve().parents[2] @@ -3254,6 +3263,8 @@ def profile_python_benchmark( lock_gpu_clocks_flag=lock_gpu_clocks_flag, gpu_sm_clock_mhz=gpu_sm_clock_mhz, gpu_mem_clock_mhz=gpu_mem_clock_mhz, + profiling_warmup=profiling_warmup, + profiling_iterations=profiling_iterations, ) _wrapper_path, target_command, env, use_torchrun = stack.enter_context( _temporary_python_profile_launch( diff --git a/code/core/profiling/profiler_wrapper.py b/code/core/profiling/profiler_wrapper.py index 00f1e38c9..58b7e5740 100644 --- a/code/core/profiling/profiler_wrapper.py +++ b/code/core/profiling/profiler_wrapper.py @@ -5,8 +5,8 @@ from __future__ import annotations -from contextlib import contextmanager import tempfile +from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING, Any, Iterator, Optional @@ -17,17 +17,46 @@ BenchmarkConfig = Any # type: ignore[assignment,misc] -def _resolve_wrapper_loop_budget(config: BenchmarkConfig) -> tuple[int, int]: +def _resolve_wrapper_loop_budget( + config: BenchmarkConfig, + *, + default_warmup: Optional[int] = None, + default_iterations: Optional[int] = None, +) -> tuple[int, int]: """Resolve warmup and profiled iteration counts for wrapper-based captures.""" profiling_warmup = getattr(config, "profiling_warmup", None) if profiling_warmup is None: - profiling_warmup = getattr(config, "warmup", 0) + profiling_warmup = ( + getattr(config, "warmup", 0) + if default_warmup is None + else default_warmup + ) profiling_iterations = getattr(config, "profiling_iterations", None) if profiling_iterations is None: - profiling_iterations = min(getattr(config, "iterations", 1), 10) + profiling_iterations = ( + min(getattr(config, "iterations", 1), 10) + if default_iterations is None + else default_iterations + ) + + return _validate_wrapper_loop_budget(profiling_warmup, profiling_iterations) - return max(int(profiling_warmup), 0), max(int(profiling_iterations), 1) + +def _validate_wrapper_loop_budget( + profiling_warmup: object, profiling_iterations: object +) -> tuple[int, int]: + if isinstance(profiling_warmup, bool) or not isinstance(profiling_warmup, int): + raise ValueError("profiling_warmup must be a non-negative integer") + if profiling_warmup < 0: + raise ValueError("profiling_warmup must be a non-negative integer") + if isinstance(profiling_iterations, bool) or not isinstance( + profiling_iterations, int + ): + raise ValueError("profiling_iterations must be a positive integer") + if profiling_iterations < 1: + raise ValueError("profiling_iterations must be a positive integer") + return profiling_warmup, profiling_iterations @contextmanager @@ -59,8 +88,13 @@ def render_nsys_python_profile_wrapper( lock_gpu_clocks_flag: bool, gpu_sm_clock_mhz: Optional[int], gpu_mem_clock_mhz: Optional[int], + profiling_warmup: int = 1, + profiling_iterations: int = 1, ) -> str: """Render the nsys-specific Python benchmark wrapper.""" + profiling_warmup, profiling_iterations = _validate_wrapper_loop_budget( + profiling_warmup, profiling_iterations + ) return f""" from pathlib import Path @@ -99,6 +133,8 @@ def _run_profile() -> None: lock_gpu_clocks={lock_gpu_clocks_flag!r}, gpu_sm_clock_mhz={gpu_sm_clock_mhz!r}, gpu_mem_clock_mhz={gpu_mem_clock_mhz!r}, + profiling_warmup={profiling_warmup!r}, + profiling_iterations={profiling_iterations!r}, ) benchmark._config = ReadOnlyBenchmarkConfigView.from_config(_profiling_config) lock_ctx = ( @@ -118,10 +154,11 @@ def _run_profile() -> None: print(f"[profile_warning] Failed to ramp GPU clocks before nsys capture: {{exc}}", file=sys.stderr) benchmark.setup() - # Warmup (keep short; profiling is not a timing run) - benchmark.benchmark_fn() + # Warmup stays outside the measured NVTX range. + for _ in range({profiling_warmup}): + benchmark.benchmark_fn() - # Profile exactly one execution. Nsight Systems defers CUDA activity + # Profile the configured steady-state executions. Nsight Systems defers CUDA activity # buffer flushing until cudaProfilerStop(); calling it explicitly is # required before the successful hard exit below. import torch @@ -142,7 +179,8 @@ def _run_profile() -> None: _profiler_started = True with nvtx_range("compute_kernel:profile", enable=True): - benchmark.benchmark_fn() + for _ in range({profiling_iterations}): + benchmark.benchmark_fn() if torch.cuda.is_available(): torch.cuda.synchronize() except BaseException as exc: diff --git a/code/tests/test_profiler_wrapper.py b/code/tests/test_profiler_wrapper.py index 02614a720..5255056ca 100644 --- a/code/tests/test_profiler_wrapper.py +++ b/code/tests/test_profiler_wrapper.py @@ -6,6 +6,8 @@ import sys from pathlib import Path +import pytest + from core.harness.benchmark_harness import BenchmarkConfig from core.profiling.profiler_wrapper import ( _resolve_wrapper_loop_budget, @@ -31,6 +33,31 @@ def test_wrapper_loop_budget_honors_profiling_specific_overrides() -> None: assert _resolve_wrapper_loop_budget(config) == (0, 1) +def test_nsys_loop_budget_preserves_one_plus_one_when_unset() -> None: + config = BenchmarkConfig(iterations=20, warmup=5) + assert _resolve_wrapper_loop_budget( + config, default_warmup=1, default_iterations=1 + ) == (1, 1) + + +@pytest.mark.parametrize( + "field,value,error", + [ + ("profiling_warmup", -1, "non-negative integer"), + ("profiling_warmup", True, "non-negative integer"), + ("profiling_iterations", 0, "positive integer"), + ("profiling_iterations", 1.5, "positive integer"), + ], +) +def test_wrapper_loop_budget_rejects_invalid_explicit_counts( + field: str, value: object, error: str +) -> None: + config = BenchmarkConfig() + setattr(config, field, value) + with pytest.raises(ValueError, match=error): + _resolve_wrapper_loop_budget(config) + + def test_temporary_python_profile_wrapper_cleans_up_file() -> None: wrapper_path: Path | None = None @@ -64,11 +91,32 @@ def test_render_nsys_wrapper_contains_expected_config() -> None: assert "_apply_overrides(list(_target_override_argv))" in wrapper assert "target_extra_args={_target_label: list(_target_override_argv)}" in wrapper assert 'with nvtx_range("compute_kernel:profile", enable=True):' in wrapper + assert "for _ in range(1):" in wrapper assert 'if getattr(benchmark, "profile_require_teardown", False):' in wrapper assert "_os._exit(0)" in wrapper assert "raise SystemExit(0)" not in wrapper +def test_render_nsys_wrapper_honors_explicit_loop_budget() -> None: + wrapper = render_nsys_python_profile_wrapper( + benchmark_path=Path("/tmp/example.py"), + nvtx_includes=["compute_kernel:profile/"], + target_label=None, + target_override_argv=None, + validity_profile="portable", + lock_gpu_clocks_flag=False, + gpu_sm_clock_mhz=None, + gpu_mem_clock_mhz=None, + profiling_warmup=5, + profiling_iterations=3, + ) + + assert "profiling_warmup=5" in wrapper + assert "profiling_iterations=3" in wrapper + assert "for _ in range(5):" in wrapper + assert "for _ in range(3):" in wrapper + + def test_render_ncu_wrapper_contains_expected_config() -> None: wrapper = render_ncu_python_profile_wrapper( benchmark_path=Path("/tmp/example.py"), From 2aae639f0cb10dc7ca01a1e7bbb7992f33a732fd Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 03:30:03 -0700 Subject: [PATCH 10/19] fix: expose the declared Ozaki comparison budget on its reference --- code/labs/ozaki_scheme/accuracy_policy.py | 20 ++++++++- .../ozaki_scheme/baseline_ozaki_scheme.py | 6 ++- code/tests/test_ozaki_scheme_lab.py | 43 +++++++++++++++++++ 3 files changed, 67 insertions(+), 2 deletions(-) diff --git a/code/labs/ozaki_scheme/accuracy_policy.py b/code/labs/ozaki_scheme/accuracy_policy.py index ae17445d6..90d27085b 100644 --- a/code/labs/ozaki_scheme/accuracy_policy.py +++ b/code/labs/ozaki_scheme/accuracy_policy.py @@ -5,7 +5,6 @@ import os from pathlib import Path - POLICY_ID = "ozaki-fp64-emulation-arithmetic-ceilings-v1" REFERENCE_ID = "native-fp64-full-plus-cpu-long-double-edge-v1" DEFAULT_POLICY_PATH = Path(__file__).with_name("accuracy_policy.json") @@ -122,3 +121,22 @@ def configured_accuracy(variant: str) -> tuple[list[str], tuple[float, float]]: return (["--relative-l2-limit", str(relative_arg), "--normalized-max-abs-limit", str(normalized_arg)], (limits["checksum_rtol"], limits["checksum_atol"])) + + +def configured_reference_tolerance() -> tuple[float, float]: + """Expose the declared candidate envelope on the shared native reference. + + The ordinary pair runner uses the baseline's secondary checksum tolerance. + Each candidate still gates its complete array against its own, potentially + stricter, native-FP64 error limits before emitting an accepted checksum. + """ + path = os.environ.get("AISP_OZAKI_ACCURACY_POLICY") + if not path: + return (0.0, 0.0) + policy = load_accuracy_policy(Path(path)) + variants = policy if policy["schema_version"] == 1 else policy["variants"] + limits = [_limits_from_item(variants[name]) for name in QUALIFICATION_VARIANTS if name in variants] + return ( + max((item["checksum_rtol"] for item in limits), default=0.0), + max((item["checksum_atol"] for item in limits), default=0.0), + ) diff --git a/code/labs/ozaki_scheme/baseline_ozaki_scheme.py b/code/labs/ozaki_scheme/baseline_ozaki_scheme.py index 28e6b0eae..8c081f028 100644 --- a/code/labs/ozaki_scheme/baseline_ozaki_scheme.py +++ b/code/labs/ozaki_scheme/baseline_ozaki_scheme.py @@ -8,12 +8,14 @@ from core.benchmark.cuda_binary_benchmark import CudaBinaryBenchmark from core.benchmark.verification import simple_signature from core.harness.benchmark_harness import BaseBenchmark +from labs.ozaki_scheme.accuracy_policy import configured_reference_tolerance class BaselineOzakiSchemeBenchmark(CudaBinaryBenchmark): """Native FP64 accuracy anchor for the Ozaki scheme lab.""" def __init__(self) -> None: + self._checksum_tolerance = configured_reference_tolerance() self._shape = (4096, 4096, 4096) self._run_args = [ "--m", str(self._shape[0]), @@ -68,7 +70,9 @@ def get_custom_metrics(self) -> Optional[dict]: return None def get_output_tolerance(self) -> tuple[float, float]: - return (0.0, 0.0) + # The shared reference advertises the already-declared pair envelope; + # candidate binaries retain their independent full-array gates. + return self._checksum_tolerance def get_benchmark() -> BaseBenchmark: diff --git a/code/tests/test_ozaki_scheme_lab.py b/code/tests/test_ozaki_scheme_lab.py index 68226dcc6..9ce5b08e6 100644 --- a/code/tests/test_ozaki_scheme_lab.py +++ b/code/tests/test_ozaki_scheme_lab.py @@ -1,5 +1,9 @@ from __future__ import annotations +import json + +import pytest + from labs.ozaki_scheme.lab_utils import ( format_result_row, parse_float_csv, @@ -9,6 +13,45 @@ ) +def test_ozaki_reference_exposes_declared_secondary_pair_budget(monkeypatch: pytest.MonkeyPatch) -> None: + from labs.ozaki_scheme.accuracy_policy import ( + DEFAULT_POLICY_PATH, + configured_accuracy, + ) + from labs.ozaki_scheme.baseline_ozaki_scheme import BaselineOzakiSchemeBenchmark + from labs.ozaki_scheme.optimized_ozaki_scheme_dynamic import ( + OptimizedOzakiSchemeDynamicBenchmark, + ) + from labs.ozaki_scheme.optimized_ozaki_scheme_fixed import ( + OptimizedOzakiSchemeFixedBenchmark, + ) + + monkeypatch.setenv("AISP_OZAKI_ACCURACY_POLICY", str(DEFAULT_POLICY_PATH)) + policy = json.loads(DEFAULT_POLICY_PATH.read_text()) + reference = BaselineOzakiSchemeBenchmark() + expected = (0.0, max(item["checksum_atol"] for item in policy["variants"].values())) + assert reference.get_output_tolerance() == expected + assert expected[1] > 0.0 + for variant, benchmark_type in ( + ("dynamic", OptimizedOzakiSchemeDynamicBenchmark), + ("fixed", OptimizedOzakiSchemeFixedBenchmark), + ): + candidate = benchmark_type() + native_gate_args, tolerance = configured_accuracy(variant) + assert candidate.get_output_tolerance() == tolerance + assert tolerance[1] <= reference.get_output_tolerance()[1] + for index in range(0, len(native_gate_args), 2): + flag, value = native_gate_args[index:index + 2] + assert candidate._run_args[candidate._run_args.index(flag) + 1] == value + + +def test_ozaki_reference_stays_exact_without_an_accuracy_policy(monkeypatch: pytest.MonkeyPatch) -> None: + from labs.ozaki_scheme.baseline_ozaki_scheme import BaselineOzakiSchemeBenchmark + + monkeypatch.delenv("AISP_OZAKI_ACCURACY_POLICY", raising=False) + assert BaselineOzakiSchemeBenchmark().get_output_tolerance() == (0.0, 0.0) + + def test_ozaki_lab_parse_metrics_captures_strategy_and_checksum() -> None: stdout = """ VARIANT: ozaki_dynamic From 06a58a79fc058f8a1e351e59d9188b604ebd14c9 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 03:33:18 -0700 Subject: [PATCH 11/19] fix: reject empty execution audits and preserve qualification scope --- code/core/harness/execution_audit.py | 64 +++++++-- .../kv_cache_compression/qualify_accuracy.py | 35 ++++- code/labs/ozaki_scheme/qualify_accuracy.py | 35 ++++- ...est_accuracy_requirements_followthrough.py | 64 ++++++++- code/tests/test_anti_cheat_edge_cases.py | 11 +- code/tests/test_execution_audit_guards.py | 131 +++++++++++++++++- 6 files changed, 308 insertions(+), 32 deletions(-) diff --git a/code/core/harness/execution_audit.py b/code/core/harness/execution_audit.py index 5ab8077b4..046f170f3 100644 --- a/code/core/harness/execution_audit.py +++ b/code/core/harness/execution_audit.py @@ -106,6 +106,7 @@ class OperationPlacementResult: expected_device: str operations_seen: int + expected_device_operations_seen: int operator_counts: tuple[tuple[str, int], ...] operation_evidence: tuple[OperationEvidence, ...] operation_evidence_truncated: bool @@ -115,7 +116,28 @@ class OperationPlacementResult: @property def passed(self) -> bool: - return self.violations_seen == 0 + return self.violations_seen == 0 and self.expected_device_operations_seen > 0 + + @property + def execution_observed(self) -> bool: + """Whether at least one operation touched a tensor on the expected device.""" + + return self.expected_device_operations_seen > 0 + + @property + def failure_reasons(self) -> tuple[str, ...]: + reasons: list[str] = [] + if not self.execution_observed: + if self.operations_seen == 0: + reasons.append("no dispatcher-visible tensor operations were observed") + else: + reasons.append( + "no dispatcher-visible tensor operation touched the expected device " + f"{self.expected_device}" + ) + if self.violations_seen: + reasons.append(f"{self.violations_seen} operation placement violation(s)") + return tuple(reasons) def to_dict(self) -> dict[str, Any]: return { @@ -123,6 +145,9 @@ def to_dict(self) -> dict[str, Any]: "scope": PLACEMENT_SCOPE, "expected_device": self.expected_device, "operations_seen": self.operations_seen, + "expected_device_operations_seen": self.expected_device_operations_seen, + "execution_observed": self.execution_observed, + "failure_reasons": list(self.failure_reasons), "operator_counts": dict(self.operator_counts), "operation_evidence": [item.to_dict() for item in self.operation_evidence], "operation_evidence_truncated": self.operation_evidence_truncated, @@ -169,6 +194,7 @@ def __init__( allowance.operations ) self._operations_seen = 0 + self._expected_device_operations_seen = 0 self._operator_counts: Counter[str] = Counter() self._operation_evidence: list[OperationEvidence] = [] self._violations_seen = 0 @@ -226,6 +252,8 @@ def __torch_dispatch__( *self._collect_tensor_evidence(actual_kwargs, "kwargs", operator), *self._collect_tensor_evidence(result, "output", operator), ] + if any(item.matches_expected_device for item in observed): + self._expected_device_operations_seen += 1 mismatched_paths = tuple( item.path for item in observed @@ -261,6 +289,7 @@ def result(self) -> OperationPlacementResult: return OperationPlacementResult( expected_device=str(self.expected_device), operations_seen=self._operations_seen, + expected_device_operations_seen=self._expected_device_operations_seen, operator_counts=tuple(sorted(self._operator_counts.items())), operation_evidence=tuple(self._operation_evidence), operation_evidence_truncated=self._operations_seen > len(self._operation_evidence), @@ -412,7 +441,9 @@ def raise_for_failure(self) -> None: if self.passed: return diagnostics: list[str] = [] - if not self.placement.passed: + if not self.placement.execution_observed: + diagnostics.append(self.placement.failure_reasons[0]) + if self.placement.violation_evidence: first = self.placement.violation_evidence[0] diagnostics.append( f"{self.placement.violations_seen} operation placement violation(s); " @@ -557,6 +588,7 @@ def _audit_fresh_benchmark(args: argparse.Namespace) -> tuple[ExecutionAuditResu original_argv = sys.argv benchmark = None + primary_error: Exception | None = None teardown_error: Exception | None = None try: sys.argv = [str(benchmark_path), *args.target_arg] @@ -599,6 +631,9 @@ def _audit_fresh_benchmark(args: argparse.Namespace) -> tuple[ExecutionAuditResu "normal_timing_lifecycle_modified": False, } return result, metadata + except Exception as error: + primary_error = error + raise finally: if benchmark is not None and callable(getattr(benchmark, "teardown", None)): try: @@ -607,10 +642,14 @@ def _audit_fresh_benchmark(args: argparse.Namespace) -> tuple[ExecutionAuditResu teardown_error = error sys.argv = original_argv if teardown_error is not None: - raise RuntimeError( + detail = ( f"fresh benchmark teardown failed: {type(teardown_error).__name__}: " f"{teardown_error}" - ) from teardown_error + ) + if primary_error is not None: + primary_error.add_note(detail) + else: + raise RuntimeError(detail) from teardown_error def main(argv: Sequence[str] | None = None) -> int: @@ -620,15 +659,16 @@ def main(argv: Sequence[str] | None = None) -> int: try: result, metadata = _audit_fresh_benchmark(args) except Exception as error: + payload = { + "schema": "aisp.execution-audit.v1", + "passed": False, + "error": f"{type(error).__name__}: {error}", + } + error_notes = [str(note) for note in getattr(error, "__notes__", ())] + if error_notes: + payload["error_notes"] = error_notes print( - json.dumps( - { - "schema": "aisp.execution-audit.v1", - "passed": False, - "error": f"{type(error).__name__}: {error}", - }, - sort_keys=True, - ) + json.dumps(payload, sort_keys=True) ) return 1 payload = result.to_dict() diff --git a/code/labs/kv_cache_compression/qualify_accuracy.py b/code/labs/kv_cache_compression/qualify_accuracy.py index 2849a5547..68d989d57 100644 --- a/code/labs/kv_cache_compression/qualify_accuracy.py +++ b/code/labs/kv_cache_compression/qualify_accuracy.py @@ -20,6 +20,18 @@ load_accuracy_policy, ) +PROVENANCE_FIELDS = ( + "git_commit", + "torch", + "cuda", + "transformer_engine", + "gpu", + "compute_capability", +) +RECEIPT_BINDING = ( + "receipt_consistency_only; execution/source identity requires companion receipts" +) + def required_cases(policy: dict) -> set[tuple[str, str, int]]: cases: set[tuple[str, str, int]] = set() @@ -38,6 +50,7 @@ def assess_receipts(policy: dict, receipts: list[dict]) -> dict: failures: list[str] = [] receipt_results: list[dict] = [] expected_provenance: tuple | None = None + provenance_consistent = True variants = policy["variants"] expected_workload = {key: value for key, value in WORKLOAD.items() if key != "storage_dtype"} @@ -60,15 +73,15 @@ def assess_receipts(policy: dict, receipts: list[dict]) -> dict: reasons.append("reference identity mismatch") if receipt.get("workload") != expected_workload: reasons.append("workload mismatch") - provenance = tuple(receipt.get(name) for name in ( - "git_commit", "torch", "cuda", "transformer_engine", "gpu", "compute_capability" - )) + provenance = tuple(receipt.get(name) for name in PROVENANCE_FIELDS) if any(value in (None, "", []) for value in provenance): reasons.append("hardware/software provenance is incomplete") + provenance_consistent = False elif expected_provenance is None: expected_provenance = provenance elif provenance != expected_provenance: reasons.append("hardware/software provenance differs across receipts") + provenance_consistent = False metrics = receipt.get("metrics") if key[0] in variants and isinstance(metrics, dict): limits = _limits_from_item(variants[key[0]]) @@ -77,8 +90,13 @@ def assess_receipts(policy: dict, receipts: list[dict]) -> dict: name = f"{tensor}.{metric_name}" value = metrics.get(name) limit = getattr(limits, metric_name) - if not isinstance(value, (int, float)) or not math.isfinite(float(value)): - reasons.append(f"{name} is missing or non-finite") + if ( + isinstance(value, bool) + or not isinstance(value, int | float) + or not math.isfinite(float(value)) + or float(value) < 0 + ): + reasons.append(f"{name} is missing, boolean, negative, or non-finite") elif float(value) > limit: reasons.append(f"{name}={float(value):.8g} exceeds {limit:.8g}") else: @@ -91,6 +109,11 @@ def assess_receipts(policy: dict, receipts: list[dict]) -> dict: for key in sorted(required - seen): failures.append(f"missing required receipt: {key}") + declared_provenance = ( + dict(zip(PROVENANCE_FIELDS, expected_provenance, strict=True)) + if expected_provenance is not None and provenance_consistent + else None + ) return { "schema_version": 1, "policy_id": policy["policy_id"], @@ -99,6 +122,8 @@ def assess_receipts(policy: dict, receipts: list[dict]) -> dict: "passing_case_count": sum(item["passed"] for item in receipt_results), "failures": failures, "receipts": receipt_results, + "declared_provenance": declared_provenance, + "binding": RECEIPT_BINDING, "claim_boundary": policy["qualification"]["claim_boundary"], } diff --git a/code/labs/ozaki_scheme/qualify_accuracy.py b/code/labs/ozaki_scheme/qualify_accuracy.py index 4b1396cf8..f9f8c505b 100644 --- a/code/labs/ozaki_scheme/qualify_accuracy.py +++ b/code/labs/ozaki_scheme/qualify_accuracy.py @@ -20,6 +20,16 @@ ) from labs.ozaki_scheme.lab_utils import parse_metrics +PROVENANCE_FIELDS = ( + "gpu_name", + "compute_capability", + "cuda_runtime_version", + "cublas_version", +) +RECEIPT_BINDING = ( + "receipt_consistency_only; execution/source identity requires companion receipts" +) + def _required_cases(policy: dict) -> dict[tuple[str, str, int], dict]: cases = {} @@ -49,6 +59,7 @@ def assess_logs(policy: dict, log_texts: list[str]) -> dict: failures: list[str] = [] results: list[dict] = [] expected_provenance: tuple | None = None + provenance_consistent = True for index, text in enumerate(log_texts): metrics = parse_metrics(text) @@ -62,15 +73,15 @@ def assess_logs(policy: dict, log_texts: list[str]) -> dict: reasons.append("log is not retained measurement-only evidence") if metrics.get("emulation_used") != 1 or int(metrics.get("retained_bits", -1)) < 0: reasons.append("cuBLAS did not report active fixed-point emulation") - provenance = tuple(metrics.get(name) for name in ( - "gpu_name", "compute_capability", "cuda_runtime_version", "cublas_version" - )) + provenance = tuple(metrics.get(name) for name in PROVENANCE_FIELDS) if any(value in (None, "") for value in provenance): reasons.append("GPU/CUDA/cuBLAS provenance is incomplete") + provenance_consistent = False elif expected_provenance is None: expected_provenance = provenance elif provenance != expected_provenance: reasons.append("GPU/CUDA/cuBLAS provenance differs across logs") + provenance_consistent = False variant = str(metrics.get("variant", "")).removeprefix("ozaki_") if variant in policy["variants"]: limits = _limits_from_item(policy["variants"][variant]) @@ -79,8 +90,15 @@ def assess_logs(policy: dict, log_texts: list[str]) -> dict: ("normalized_max_abs_error", "normalized_max_abs"), ): value = metrics.get(metric_name) - if not isinstance(value, (int, float)) or not math.isfinite(float(value)): - reasons.append(f"{metric_name} is missing or non-finite") + if ( + isinstance(value, bool) + or not isinstance(value, int | float) + or not math.isfinite(float(value)) + or float(value) < 0 + ): + reasons.append( + f"{metric_name} is missing, boolean, negative, or non-finite" + ) elif float(value) > limits[limit_name]: reasons.append(f"{metric_name}={float(value):.8g} exceeds {limits[limit_name]:.8g}") if variant == "dynamic" and ( @@ -103,6 +121,11 @@ def assess_logs(policy: dict, log_texts: list[str]) -> dict: for key in sorted(set(required) - seen): failures.append(f"missing required log: {key}") + declared_provenance = ( + dict(zip(PROVENANCE_FIELDS, expected_provenance, strict=True)) + if expected_provenance is not None and provenance_consistent + else None + ) return { "schema_version": 1, "policy_id": policy["policy_id"], @@ -111,6 +134,8 @@ def assess_logs(policy: dict, log_texts: list[str]) -> dict: "passing_case_count": sum(item["passed"] for item in results), "failures": failures, "logs": results, + "declared_provenance": declared_provenance, + "binding": RECEIPT_BINDING, "claim_boundary": policy["qualification"]["claim_boundary"], } diff --git a/code/tests/test_accuracy_requirements_followthrough.py b/code/tests/test_accuracy_requirements_followthrough.py index 15f743ae1..7edda17e3 100644 --- a/code/tests/test_accuracy_requirements_followthrough.py +++ b/code/tests/test_accuracy_requirements_followthrough.py @@ -134,7 +134,19 @@ def test_kv_qualification_requires_all_holdouts_and_retains_failure() -> None: policy = load_accuracy_policy(DEFAULT_POLICY_PATH) receipts = _kv_receipts(policy) - assert assess_receipts(policy, receipts)["status"] == "qualified_arithmetic_gate" + passing = assess_receipts(policy, receipts) + assert passing["status"] == "qualified_arithmetic_gate" + assert passing["declared_provenance"] == { + "git_commit": "0123456789abcdef", + "torch": "2.9.1", + "cuda": "13.0", + "transformer_engine": "2.18.0", + "gpu": "NVIDIA B200", + "compute_capability": [10, 0], + } + assert passing["binding"] == ( + "receipt_consistency_only; execution/source identity requires companion receipts" + ) damaged = deepcopy(receipts) damaged[-1]["metrics"]["cache_v.normalized_max_abs"] = 0.5 @@ -143,6 +155,26 @@ def test_kv_qualification_requires_all_holdouts_and_retains_failure() -> None: assert any("cache_v.normalized_max_abs" in failure for failure in result["failures"]) +@pytest.mark.parametrize("invalid_error", [False, -0.01]) +def test_kv_qualification_rejects_boolean_and_negative_error_metrics( + invalid_error: bool | float, +) -> None: + from labs.kv_cache_compression.accuracy import DEFAULT_POLICY_PATH, load_accuracy_policy + from labs.kv_cache_compression.qualify_accuracy import assess_receipts + + policy = load_accuracy_policy(DEFAULT_POLICY_PATH) + receipts = _kv_receipts(policy) + receipts[0]["metrics"]["cache_k.relative_l2"] = invalid_error + + result = assess_receipts(policy, receipts) + + assert result["status"] == "failed_arithmetic_gate" + assert any( + "cache_k.relative_l2 is missing, boolean, negative, or non-finite" in failure + for failure in result["failures"] + ) + + def test_ozaki_checked_in_policy_rejects_widening(tmp_path: Path) -> None: from labs.ozaki_scheme.accuracy_policy import DEFAULT_POLICY_PATH, load_accuracy_policy @@ -252,9 +284,37 @@ def test_ozaki_qualification_requires_independent_edges_and_rejects_corruption() policy = load_accuracy_policy(DEFAULT_POLICY_PATH) cases = list(_required_cases(policy).values()) logs = [_ozaki_log(case) for case in cases] - assert assess_logs(policy, logs)["status"] == "qualified_arithmetic_gate" + passing = assess_logs(policy, logs) + assert passing["status"] == "qualified_arithmetic_gate" + assert passing["declared_provenance"] == { + "gpu_name": "NVIDIA B200", + "compute_capability": "10.0", + "cuda_runtime_version": 13000, + "cublas_version": 130000, + } + assert passing["binding"] == ( + "receipt_consistency_only; execution/source identity requires companion receipts" + ) logs[-1] = _ozaki_log(cases[-1], relative_l2=0.1) result = assess_logs(policy, logs) assert result["status"] == "failed_arithmetic_gate" assert any("relative_l2_error" in failure for failure in result["failures"]) + + +def test_ozaki_qualification_rejects_negative_error_metric() -> None: + from labs.ozaki_scheme.accuracy_policy import DEFAULT_POLICY_PATH, load_accuracy_policy + from labs.ozaki_scheme.qualify_accuracy import _required_cases, assess_logs + + policy = load_accuracy_policy(DEFAULT_POLICY_PATH) + cases = list(_required_cases(policy).values()) + logs = [_ozaki_log(case) for case in cases] + logs[0] = _ozaki_log(cases[0], relative_l2=-0.01) + + result = assess_logs(policy, logs) + + assert result["status"] == "failed_arithmetic_gate" + assert any( + "relative_l2_error is missing, boolean, negative, or non-finite" in failure + for failure in result["failures"] + ) diff --git a/code/tests/test_anti_cheat_edge_cases.py b/code/tests/test_anti_cheat_edge_cases.py index 2313b0d90..46d81ccb3 100644 --- a/code/tests/test_anti_cheat_edge_cases.py +++ b/code/tests/test_anti_cheat_edge_cases.py @@ -326,16 +326,19 @@ def test_cpu_spillover_single_op_on_cpu(self): assert {tensor.numel for tensor in violation.tensors} == {12} def test_cpu_spillover_data_dependent_branch(self): - """Placement evidence follows the branch that actually executes.""" + """Placement evidence rejects both no execution and a CPU-only branch.""" cpu_value = torch.arange(8, dtype=torch.float32) def execute_branch(use_cpu: bool) -> None: if use_cpu: torch.relu(cpu_value) - clean = audit_callable_once(lambda: execute_branch(False), expected_device="cuda") - assert clean.passed - assert clean.placement.operations_seen == 0 + no_execution = audit_callable_once(lambda: execute_branch(False), expected_device="cuda") + assert not no_execution.passed + assert no_execution.placement.operations_seen == 0 + assert no_execution.placement.failure_reasons == ( + "no dispatcher-visible tensor operations were observed", + ) violation = audit_callable_once(lambda: execute_branch(True), expected_device="cuda") assert not violation.passed diff --git a/code/tests/test_execution_audit_guards.py b/code/tests/test_execution_audit_guards.py index bd584b979..0cfea1d4a 100644 --- a/code/tests/test_execution_audit_guards.py +++ b/code/tests/test_execution_audit_guards.py @@ -29,6 +29,9 @@ def test_cpu_operation_audit_retains_positive_and_negative_extent_evidence() -> clean = audit_callable_once(lambda: value.square(), expected_device="cpu") assert clean.passed assert clean.placement.operations_seen == 1 + assert clean.placement.expected_device_operations_seen == 1 + assert clean.placement.execution_observed + assert clean.placement.failure_reasons == () clean_record = clean.placement.operation_evidence[0] assert clean_record.operator == "aten.pow.Tensor_Scalar" assert clean_record.mismatched_paths == () @@ -38,6 +41,8 @@ def test_cpu_operation_audit_retains_positive_and_negative_extent_evidence() -> violation = audit_callable_once(lambda: value.square(), expected_device="cuda") assert not violation.passed + assert violation.placement.expected_device_operations_seen == 0 + assert not violation.placement.execution_observed assert violation.placement.violations_seen == 1 violation_record = violation.placement.violation_evidence[0] assert violation_record.operator == "aten.pow.Tensor_Scalar" @@ -58,14 +63,19 @@ def test_host_tensor_allowance_requires_exact_identity_and_operator_scope() -> N operations=("aten._local_scalar_dense.default",), ) - clean = audit_callable_once( + allowed_host_only = audit_callable_once( declared_scalar.item, expected_device="cuda", allowed_host_tensors=(allowance,), ) - assert clean.passed - assert clean.placement.operations_seen == 1 - tensor_evidence = clean.placement.operation_evidence[0].tensors + assert not allowed_host_only.passed + assert allowed_host_only.placement.operations_seen == 1 + assert allowed_host_only.placement.violations_seen == 0 + assert allowed_host_only.placement.expected_device_operations_seen == 0 + assert allowed_host_only.placement.failure_reasons == ( + "no dispatcher-visible tensor operation touched the expected device cuda", + ) + tensor_evidence = allowed_host_only.placement.operation_evidence[0].tensors assert len(tensor_evidence) == 1 assert tensor_evidence[0].allowed_host_tensor @@ -86,6 +96,20 @@ def test_host_tensor_allowance_requires_exact_identity_and_operator_scope() -> N assert wrong_operator.placement.violation_evidence[0].operator == "aten.neg.default" +def test_noop_audit_fails_closed_without_scanning_bounded_evidence() -> None: + result = audit_callable_once(lambda: None, expected_device="cuda:0", evidence_limit=1) + + assert not result.passed + assert result.placement.operations_seen == 0 + assert result.placement.expected_device_operations_seen == 0 + assert not result.placement.execution_observed + assert result.placement.failure_reasons == ( + "no dispatcher-visible tensor operations were observed", + ) + with pytest.raises(RuntimeError, match="no dispatcher-visible tensor operations"): + result.raise_for_failure() + + def test_declared_destination_write_coverage_accepts_full_write_and_rejects_partial_write() -> None: source = torch.arange(8, dtype=torch.float32) full_destination = torch.empty_like(source) @@ -135,6 +159,7 @@ def test_destination_write_coverage_refuses_unprovable_tensor_contracts( def _write_cli_benchmark(path: Path, *, mode: str) -> None: write_statement = { + "noop": "pass", "full": "torch.mul(self.input, 2, out=self.output)", "partial": "self.output[:3].copy_(self.input[:3] * 2)", "reassigned": ( @@ -218,6 +243,102 @@ def test_fresh_benchmark_cli_runs_real_out_of_timing_audit( assert payload["passed"] is (mode == "full") +def test_fresh_benchmark_cli_rejects_noop_without_declared_destination(tmp_path: Path) -> None: + benchmark_path = tmp_path / "noop_cpu_benchmark.py" + _write_cli_benchmark(benchmark_path, mode="noop") + + completed = subprocess.run( + [ + sys.executable, + "-m", + "core.harness.execution_audit", + str(benchmark_path), + "--expected-device", + "cuda:0", + ], + cwd=Path(__file__).resolve().parents[1], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + + assert completed.returncode == 2, completed.stderr + payload = json.loads(completed.stdout) + assert payload["passed"] is False + assert payload["placement"]["operations_seen"] == 0 + assert payload["placement"]["expected_device_operations_seen"] == 0 + assert payload["placement"]["execution_observed"] is False + assert payload["placement"]["failure_reasons"] == [ + "no dispatcher-visible tensor operations were observed" + ] + + +@pytest.mark.parametrize("primary_failure", [False, True]) +def test_fresh_benchmark_cli_preserves_primary_failure_when_teardown_also_fails( + tmp_path: Path, + primary_failure: bool, +) -> None: + benchmark_path = tmp_path / "failing_lifecycle_cpu_benchmark.py" + primary_statement = ( + 'raise ValueError("primary execution failure")' if primary_failure else "pass" + ) + benchmark_path.write_text( + textwrap.dedent( + f""" + import torch + + class FailingLifecycleCpuBenchmark: + def setup(self): + self.input = torch.arange(6, dtype=torch.float32) + + def benchmark_fn(self): + self.input.square() + {primary_statement} + + def teardown(self): + raise RuntimeError("secondary teardown failure") + + def get_benchmark(): + return FailingLifecycleCpuBenchmark() + """ + ).strip() + + "\n", + encoding="utf-8", + ) + + completed = subprocess.run( + [ + sys.executable, + "-m", + "core.harness.execution_audit", + str(benchmark_path), + "--expected-device", + "cpu", + ], + cwd=Path(__file__).resolve().parents[1], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + + assert completed.returncode == 1, completed.stderr + payload = json.loads(completed.stdout) + assert payload["passed"] is False + if primary_failure: + assert payload["error"] == "ValueError: primary execution failure" + assert payload["error_notes"] == [ + "fresh benchmark teardown failed: RuntimeError: secondary teardown failure" + ] + else: + assert payload["error"] == ( + "RuntimeError: fresh benchmark teardown failed: RuntimeError: " + "secondary teardown failure" + ) + assert "error_notes" not in payload + + @requires_cuda def test_cuda_operation_audit_detects_real_cpu_spillover() -> None: device = torch.device("cuda", torch.cuda.current_device()) @@ -227,6 +348,7 @@ def test_cuda_operation_audit_detects_real_cpu_spillover() -> None: clean = audit_callable_once(lambda: cuda_value.square(), expected_device=device) assert clean.passed assert clean.placement.operations_seen == 1 + assert clean.placement.expected_device_operations_seen == 1 assert {item.device for item in clean.placement.operation_evidence[0].tensors} == {str(device)} def spill_to_cpu() -> None: @@ -236,6 +358,7 @@ def spill_to_cpu() -> None: violation = audit_callable_once(spill_to_cpu, expected_device=device) assert not violation.passed assert violation.placement.operations_seen == 2 + assert violation.placement.expected_device_operations_seen == 1 assert violation.placement.violations_seen == 1 record = violation.placement.violation_evidence[0] assert record.operator == "aten.pow.Tensor_Scalar" From ad47f66401fc150df9435150eb8e356195e59faa Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 03:33:18 -0700 Subject: [PATCH 12/19] docs: explain matched batch controls and explicit execution audits --- code/ch13/README.md | 18 ++++++++++++++++++ code/core/scripts/refresh_readmes.py | 25 ++++++++++++++++++++++++- code/docs/api-reference.md | 18 ++++++++++++++++++ 3 files changed, 60 insertions(+), 1 deletion(-) diff --git a/code/ch13/README.md b/code/ch13/README.md index 4b485d56c..a4da0e732 100644 --- a/code/ch13/README.md +++ b/code/ch13/README.md @@ -39,6 +39,21 @@ The TE 2.18 pair now verifies its captured prediction and every post-step parame The frozen map passed all holdouts and independently rejected zeroed and localized corrupted copies of all five outputs. These bounds apply to this TE 2.18 workload and establish numerical verification only; they do not establish a speedup. +## Matched Batch Controls +`precisionfp8_te` keeps batch size 256 as its default workload. Omitting a target override preserves that default for both the eager FP16 baseline and eager Transformer Engine FP8 candidate: + +```bash +python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu +``` + +To test an optional larger matched control, pass one pair-wide override through the harness. This sends batch size 1024 to both arms and updates their workload metadata consistently: + +```bash +python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' +``` + +Compare results only when both arms report the same requested batch size and workload signature. The calibrated output policy above was established at batch 256. Batch 1024 keeps that policy but needs a fresh correctness and B200 timing run; it does not inherit the batch-256 evidence or establish a performance gain by itself. + ## Profiler Evidence Use deep-dive runs when you want to see whether the gain came from framework overhead reduction, memory behavior, or the lower-precision path itself: @@ -61,6 +76,7 @@ python -m ch13.compare python -m cli.aisp bench list-targets --chapter ch13 python -m cli.aisp bench run --targets ch13 --profile minimal python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu +python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' ``` ## Learning Goals @@ -97,10 +113,12 @@ python -m cli.aisp bench run --targets ch13 --profile minimal ## Validation Checklist - `python -m ch13.compare --examples training_standard` shows optimized training runs producing higher goodput with identical metrics. - `python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal` confirms Transformer Engine calibration plus NVFP8 execution with max error tolerances enforced. +- `python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024'` exercises the optional matched batch control; accept its timing only after both arms report batch 1024 and pass fresh output verification. - `python -m ch13.memory_profiling --dump` and the optimized variant demonstrate allocator fragmentation dropping after applying the recommended knobs, with memory reduction treated as the primary benchmark outcome. ## Notes - `custom_allocator.py` contains a standalone torch allocator shim that can be re-used in other chapters when debugging fragmentation. - `compiled_autograd.py` doubles as a tutorial on partial graph capture; the README here references it directly. +- `precisionfp8_te` defaults to batch 256. `--batch-size 1024` is an explicit pair-wide workload override, not a new default or a qualified speed claim. - `torchao_quantization_compiled`, `kv_cache_naive_flash_blockwise`, `precisionfp8`, `precisionfp8_rowwise`, and `precisionfp8_rowwise_gw_hp` remain informational variants. - `kv_cache_naive` and `memory_profiling` are memory-goal benchmarks; they are expected to reduce memory pressure even when the timed path is not faster. diff --git a/code/core/scripts/refresh_readmes.py b/code/core/scripts/refresh_readmes.py index 33db0119e..c2859b496 100644 --- a/code/core/scripts/refresh_readmes.py +++ b/code/core/scripts/refresh_readmes.py @@ -2352,6 +2352,25 @@ def lab_entry( The frozen map passed all holdouts and independently rejected zeroed and localized corrupted copies of all five outputs. These bounds apply to this TE 2.18 workload and establish numerical verification only; they do not establish a speedup.""" ), ), + MarkdownSection( + "Matched Batch Controls", + dedent( + """\ + `precisionfp8_te` keeps batch size 256 as its default workload. Omitting a target override preserves that default for both the eager FP16 baseline and eager Transformer Engine FP8 candidate: + + ```bash + python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu + ``` + + To test an optional larger matched control, pass one pair-wide override through the harness. This sends batch size 1024 to both arms and updates their workload metadata consistently: + + ```bash + python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' + ``` + + Compare results only when both arms report the same requested batch size and workload signature. The calibrated output policy above was established at batch 256. Batch 1024 keeps that policy but needs a fresh correctness and B200 timing run; it does not inherit the batch-256 evidence or establish a performance gain by itself.""" + ), + ), MarkdownSection( "Profiler Evidence", dedent( @@ -2381,6 +2400,7 @@ def lab_entry( python -m cli.aisp bench list-targets --chapter ch13 python -m cli.aisp bench run --targets ch13 --profile minimal python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu + python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' ```""" ), ), @@ -2406,11 +2426,13 @@ def lab_entry( validation=[ "`python -m ch13.compare --examples training_standard` shows optimized training runs producing higher goodput with identical metrics.", "`python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal` confirms Transformer Engine calibration plus NVFP8 execution with max error tolerances enforced.", + "`python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024'` exercises the optional matched batch control; accept its timing only after both arms report batch 1024 and pass fresh output verification.", "`python -m ch13.memory_profiling --dump` and the optimized variant demonstrate allocator fragmentation dropping after applying the recommended knobs, with memory reduction treated as the primary benchmark outcome.", ], notes=[ "`custom_allocator.py` contains a standalone torch allocator shim that can be re-used in other chapters when debugging fragmentation.", "`compiled_autograd.py` doubles as a tutorial on partial graph capture; the README here references it directly.", + "`precisionfp8_te` defaults to batch 256. `--batch-size 1024` is an explicit pair-wide workload override, not a new default or a qualified speed claim.", "`torchao_quantization_compiled`, `kv_cache_naive_flash_blockwise`, `precisionfp8`, `precisionfp8_rowwise`, and `precisionfp8_rowwise_gw_hp` remain informational variants.", "`kv_cache_naive` and `memory_profiling` are memory-goal benchmarks; they are expected to reduce memory pressure even when the timed path is not faster.", ], @@ -3701,7 +3723,7 @@ def lab_entry( ], notes=[ "The dual-pool policies produce different batch shapes. On the pinned vLLM 0.16 stack with GPT-OSS-20B, the default backend produced different greedy tokens for identical prompts, including across repeated optimized runs. The explicit batch-invariant Triton configuration above matched all 1,734 output elements on 2×B200. Apply the same backend and environment to both arms; the option does not change the default backend for other workloads. Other models and stacks still require their own correctness check.", - "Harness latency includes constructing both model engines on every benchmark invocation. Treat it as startup plus request processing, rather than steady-state routing throughput. Prefix caching is disabled so every request processes its full declared prompt.", + "The vLLM benchmarks construct each model engine once in `setup()`, execute the harness-required five full-workload warmups, and reuse the idle engines for exactly three steady-state iterations. Every invocation clears completed-request bookkeeping, uses a fresh request-id generation, and still verifies every generated token. Custom metrics report engine startup, warmup request processing, and steady-state request processing separately. Teardown emits a `vllm_engine_lifecycle` JSON record with teardown and end-to-end wall time, so moving engine construction outside the steady-state timer cannot be presented as an end-to-end speedup. Prefix caching remains disabled so every request processes its full declared prompt.", "The harness prepares live CPU prompt IDs during setup. The topology-aware runner requires this input; standalone entrypoints create default prompts before calling it. Conversion to the Python token lists required by vLLM remains part of request admission, and GPU-resident prompt inputs fail explicitly before conversion.", "`driver.py` accepts knobs such as `--prefill-gpus`, `--decode-gpus`, and `--migration-budget` to stress different regimes.", "vLLM integration now takes flags (`--model`, `--prefill-gpus`, `--decode-gpus`, etc.) plus locally available tokenizer/model weights.", @@ -3758,6 +3780,7 @@ def lab_entry( ], notes=[ "With two GPUs, the distributed target has one prefill and one decode rank. Both placement policies select the same decode rank, so this topology cannot demonstrate a reduction in migrations between decode ranks. September 7, 2026 repeated ABBA measurements on two B200s passed every full 2,048-element output comparison but found no speedup: median 14.331648 ms baseline and 14.499441 ms optimized (0.988428x; eight observations per arm, four fresh seeds). Standard deviations were 0.470046 and 0.517233 ms. These portable, unlocked observations did not reproduce the earlier single live-input run's 1.64306x. Both-arm Nsight traces are retained. Use at least two decode ranks to investigate migration benefits; this topology cannot establish that mechanism.", + "Multi-GPU results now expose `cache_aware.decode_rank_count`, `cache_aware.affinity_opportunity_count`, and `cache_aware.affinity_placement_distinguishable`. A direct 1P1D run is classified as a comparison surface. The optimized 1P1D path removes redundant global barriers between blocking point-to-point transfers while preserving the final per-request drain; `cache_aware.direct_1p1d_sync_fast_path` and `cache_aware.global_barriers_avoided_per_request` identify that separate synchronization mechanism. Its performance effect remains unmeasured until a repeated B200 rerun with full-output checks and retained traces.", "This lab is intentionally a logical reproduction of the scheduler/caching story, not a full serving engine.", "Treat single-GPU `cache_aware_disagg` as a locality-comparison benchmark with a local comparison contract. The stable value on one GPU is the cache hit rate, KV transfer volume, and worker affinity improvement; the timed delta is recorded, but it is not a trustworthy headline speed gate on this host.", "Judge the single-GPU target by cache hit rate, KV transfer volume, and worker affinity before raw wall-clock speedup.", diff --git a/code/docs/api-reference.md b/code/docs/api-reference.md index cd2aa55c5..46ab6995d 100644 --- a/code/docs/api-reference.md +++ b/code/docs/api-reference.md @@ -357,6 +357,24 @@ engine.benchmark.speed_test() # Quick GEMM/attention test **Note:** `aisp benchmark ...` commands are diagnostic microbenchmarks (`hw_*` tools) and do not use the harness. +**Explicit execution audit:** from `code/`, run one fresh setup and callback outside +benchmark timing, then emit a JSON placement and destination-write receipt: + +```bash +python -m core.harness.execution_audit ch05/optimized_vectorization.py \ + --expected-device cuda:0 --destination _output_buffer +``` + +Repeat `--destination ATTRIBUTE` for additional preallocated contiguous floating-point +or complex outputs. The audit poisons those exact tensors and rejects incomplete writes +or replaced destination identities. Placement checks cover PyTorch dispatcher-visible +operations on the current thread and require at least one operation touching the +expected device. No-op and host-only callbacks cannot pass a CUDA execution audit. +These checks do not inspect arbitrary extension internals, +other processes, or general uninitialized-memory provenance. An intentional host tensor +can be allowed for one exact operator with `--allow-host-tensor ATTRIBUTE=aten.operator.overload`. +This standalone audit produces correctness evidence, not performance measurements. + --- ### 9. AI Domain From 15301517a938ac76df8c30c4737423df6773c56d Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 03:50:11 -0700 Subject: [PATCH 13/19] docs: clarify profile loop budgets and Ozaki comparison envelope --- code/docs/api-reference.md | 7 +++++++ code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md | 5 +++++ 2 files changed, 12 insertions(+) diff --git a/code/docs/api-reference.md b/code/docs/api-reference.md index 46ab6995d..4276920af 100644 --- a/code/docs/api-reference.md +++ b/code/docs/api-reference.md @@ -174,6 +174,13 @@ Profiling with Nsight Systems, Nsight Compute, and torch.profiler. **NSYS timeout hardening (CLI + MCP + harness):** all NSYS entrypoints now route through `NsightAutomation.profile_nsys`, default to the safer `preset='light'`, support `wait_mode` (`primary`/`all`), and use a graceful timeout finalization window (`finalize_grace_seconds`) before hard termination. `profile_nsys` also supports `sanitize_python_startup=true` to prefix a safe `sitecustomize` shim for profiler subprocesses. +The Python benchmark Nsys wrapper honors explicit `BenchmarkConfig.profiling_warmup` +and `profiling_iterations`: warmups run before capture, and measured calls run +inside the profile range. If unset, this wrapper retains one warmup and one +measured call. Counts must be integers, with nonnegative warmups and positive +measured iterations. Reused serving engines can therefore warm up completely +before a steady-state mechanism capture. + **Python API:** ```python engine.profile.flame_graph() # Flame graph data diff --git a/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md b/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md index 5579ae3a0..fab97c501 100644 --- a/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md +++ b/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md @@ -29,6 +29,11 @@ The checksum tolerances are secondary harness bounds. They follow from the full- relative-L2 ceiling, Cauchy-Schwarz, and the declared uniform input bound: `atol = relative_l2 * m * n * k * input_scale^2`; `rtol` is zero. Candidate results cannot pass on the checksum alone because the executable first gates the complete array. +The ordinary pair runner reads its secondary envelope from the native-FP64 baseline. +That baseline exposes the largest declared candidate checksum envelope so either +variant can be compared. Each executable still enforces its own full-array limits; +the fixed variant retains its stricter requirements. Without an explicit policy, +the baseline keeps an exact-zero comparison envelope. These limits were fixed before new candidate execution. Prior measurement-only logs remain diagnostics and were not used to widen a bound. `load_accuracy_policy()` rejects From f10a5b7553825909b66b011f5c60dbb9809713b9 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 03:54:02 -0700 Subject: [PATCH 14/19] fix: derive KV pairwise tolerance from the independent reference budgets --- .../ACCURACY_REQUIREMENTS.md | 27 +- code/labs/kv_cache_compression/accuracy.py | 238 +++++++++++++++--- .../kv_cache_compression/accuracy_policy.json | 12 +- .../kv_cache_compression/baseline_kv_cache.py | 85 ++++++- .../test_kv_pairwise_reference_envelope.py | 158 ++++++++++++ 5 files changed, 467 insertions(+), 53 deletions(-) create mode 100644 code/tests/test_kv_pairwise_reference_envelope.py diff --git a/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md b/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md index a47b8bb6c..cc9efb537 100644 --- a/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md +++ b/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md @@ -9,10 +9,10 @@ storage-alias, relative-L2, and maximum-error checks all remain mandatory. The source ceilings are: -| Variant | Full-cache relative L2 | Maximum error / maximum reference | Pairwise rtol | Pairwise atol | -| --- | ---: | ---: | ---: | ---: | -| Delayed-scaling FP8 E4M3 | `0.0625` (`2^-4`) | `0.0625` (`2^-4`) | `0.25` | `0.0625` | -| NVFP4 E2M1 | `0.25` (`2^-2`) | `0.25` (`2^-2`) | `0.25` | `0.0625` | +| Variant | Full-cache relative L2 | Maximum error / maximum reference | +| --- | ---: | ---: | +| Delayed-scaling FP8 E4M3 | `0.0625` (`2^-4`) | `0.0625` (`2^-4`) | +| NVFP4 E2M1 | `0.25` (`2^-2`) | `0.25` (`2^-2`) | E4M3 stores three fraction bits, making half the spacing within a normal binade `2^-4`; E2M1 stores one fraction bit, making the analogous quantity `2^-2`. @@ -20,7 +20,24 @@ Those representation-scale quantities define the independent full-cache engineer ceilings. NVFP4 also uses a per-16-element E4M3 block scale and a global FP32 scale, as described in the [Transformer Engine NVFP4 documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/features/low_precision_training/nvfp4/nvfp4.html). The pairwise allowance is secondary: each arm must first pass its own full-cache -reference check. +reference check. It is a shared-reference envelope rather than a raw +`torch.allclose(rtol=0.25, atol=0.0625)` call. For the common BF16 reference `R`, +the FP8 cache `B`, the NVFP4 cache `O`, and `M = max(abs(R))`, the two unchanged +maximum-error requirements and the triangle inequality give: + +```text +max(abs(B - O)) + <= max(abs(B - R)) + max(abs(O - R)) + <= (0.0625 + 0.25) * M +``` + +The ordinary pair therefore transports the complete raw K/V arrays and uses an +exact-keyed output policy with `rtol=0` and +`atol=(0.0625 + 0.25) * max(abs(R))`. Both arms must derive an identical policy +from their independent reference pass or the harness rejects the comparison. +The coefficient `0.3125` is fixed by the two declared format ceilings; it is not +calibrated or fitted to an observed pairwise difference. A zero reference yields +zero absolute tolerance and admits only exact zero outputs. These limits were fixed from the declared formats before the new candidate runs. Prior calibration errors remain diagnostics and were not used to widen a bound. diff --git a/code/labs/kv_cache_compression/accuracy.py b/code/labs/kv_cache_compression/accuracy.py index 85f762cc3..3d14de3c3 100644 --- a/code/labs/kv_cache_compression/accuracy.py +++ b/code/labs/kv_cache_compression/accuracy.py @@ -21,6 +21,7 @@ REFERENCE_ID = "pytorch-unquantized-bf16-full-cache-v1" POLICY_ID = "kv-cache-projection-format-ceilings-v1" +PAIRWISE_ENVELOPE_METHOD = "shared-reference-max-triangle-envelope-v1" DEFAULT_POLICY_PATH = Path(__file__).with_name("accuracy_policy.json") WORKLOAD = { "batch_size": 8, @@ -44,10 +45,21 @@ class AccuracyLimits: relative_l2: float normalized_max_abs: float - pairwise_rtol: float - pairwise_atol: float + # Retained only for schema-v1 exact-zero fixtures. Schema-v2 policies use + # PairwiseEnvelope because raw torch.allclose rtol/atol have different units. + pairwise_rtol: float = 0.0 + pairwise_atol: float = 0.0 def __post_init__(self): + for name in ( + "relative_l2", + "normalized_max_abs", + "pairwise_rtol", + "pairwise_atol", + ): + value = getattr(self, name) + if isinstance(value, bool): + raise ValueError(f"{name} must be numeric, not boolean") for name in ("relative_l2", "normalized_max_abs", "pairwise_rtol"): value = getattr(self, name) if not math.isfinite(value) or not 0 <= value < 1: @@ -56,39 +68,113 @@ def __post_init__(self): raise ValueError("pairwise_atol must be finite and nonnegative") +@dataclass(frozen=True) +class PairwiseEnvelope: + """Reference-normalized full-output envelope shared by both benchmark arms.""" + + normalized_max_abs: float + output_rtol: float = 0.0 + method: str = PAIRWISE_ENVELOPE_METHOD + reference_id: str = REFERENCE_ID + + def __post_init__(self) -> None: + for name in ("normalized_max_abs", "output_rtol"): + value = getattr(self, name) + if isinstance(value, bool): + raise ValueError(f"pairwise envelope {name} must be numeric, not boolean") + if not math.isfinite(value) or not 0 <= value < 1: + raise ValueError( + f"pairwise envelope {name} must be finite and in [0, 1)" + ) + if self.output_rtol != 0: + raise ValueError("pairwise envelope output_rtol must be zero") + if self.method != PAIRWISE_ENVELOPE_METHOD: + raise ValueError(f"pairwise envelope method must be {PAIRWISE_ENVELOPE_METHOD}") + if self.reference_id != REFERENCE_ID: + raise ValueError(f"pairwise envelope reference_id must be {REFERENCE_ID}") + + +@dataclass(frozen=True) +class CacheAccuracyEvidence: + metrics: dict[str, float] + reference_max_abs: float + + # These ceilings are set from the quantized operand representations, before # candidate execution. E4M3 has three stored fraction bits, so half an ULP at # a normal binade is 2^-4. E2M1 has one stored fraction bit, so the analogous # bound is 2^-2. The full-cache aggregate and global-maximum requirements use -# those format-scale bounds. The pairwise check compares FP8 and NVFP4 caches -# only after each arm passes its independent reference check, so it uses the -# coarser E2M1 ceiling plus one E4M3-scale absolute allowance near zero. +# those format-scale bounds. The pairwise envelope is derived separately from +# the sum of both selected normalized-maximum limits and the shared reference +# magnitude. It is never fitted to a candidate output. ENGINEERING_CEILINGS = { "fp8": AccuracyLimits( relative_l2=2.0**-4, normalized_max_abs=2.0**-4, - pairwise_rtol=2.0**-2, - pairwise_atol=2.0**-4, ), "nvfp4": AccuracyLimits( relative_l2=2.0**-2, normalized_max_abs=2.0**-2, - pairwise_rtol=2.0**-2, - pairwise_atol=2.0**-4, ), } def _limits_from_item(item: dict) -> AccuracyLimits: + def policy_float(name: str, *, default: float | None = None) -> float: + if name not in item: + if default is not None: + return default + raise KeyError(name) + value = item[name] + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError(f"KV accuracy policy {name} must be numeric, not boolean") + return float(value) + try: - return AccuracyLimits(**{ - name: float(item[name]) - for name in ("relative_l2", "normalized_max_abs", "pairwise_rtol", "pairwise_atol") - }) + return AccuracyLimits( + relative_l2=policy_float("relative_l2"), + normalized_max_abs=policy_float("normalized_max_abs"), + pairwise_rtol=policy_float("pairwise_rtol", default=0.0), + pairwise_atol=policy_float("pairwise_atol", default=0.0), + ) except KeyError as exc: raise ValueError(f"KV accuracy policy missing {exc.args[0]}") from exc +def _pairwise_envelope_from_policy(policy: dict) -> PairwiseEnvelope: + if policy.get("schema_version") == 1: + return PairwiseEnvelope(normalized_max_abs=0.0) + item = policy.get("pairwise_envelope") + if not isinstance(item, dict): + raise ValueError("KV accuracy policy requires a pairwise_envelope object") + if item.get("method") != PAIRWISE_ENVELOPE_METHOD: + raise ValueError( + f"KV pairwise envelope requires method={PAIRWISE_ENVELOPE_METHOD}" + ) + if item.get("reference_id") != REFERENCE_ID: + raise ValueError(f"KV pairwise envelope requires reference_id={REFERENCE_ID}") + if item.get("variants") != QUALIFICATION_VARIANTS: + raise ValueError("KV pairwise envelope variants do not match the source contract") + if item.get("coefficient") != "sum_variant_normalized_max_abs": + raise ValueError( + "KV pairwise envelope coefficient must be sum_variant_normalized_max_abs" + ) + output_rtol = item.get("output_rtol") + if isinstance(output_rtol, bool) or not isinstance(output_rtol, int | float): + raise ValueError("KV pairwise envelope output_rtol must be numeric") + variants = policy.get("variants", {}) + coefficient = sum( + _limits_from_item(variants[variant]).normalized_max_abs + for variant in QUALIFICATION_VARIANTS + ) + return PairwiseEnvelope( + normalized_max_abs=coefficient, + output_rtol=float(output_rtol), + method=str(item["method"]), + reference_id=str(item["reference_id"]), + ) + + def load_accuracy_policy(path: Path) -> dict: """Load and validate the reviewed policy contract, without running a candidate.""" policy = json.loads(path.read_text()) @@ -100,10 +186,18 @@ def load_accuracy_policy(path: Path) -> dict: if variant not in policy: continue limits = _limits_from_item(policy[variant]) - if any(getattr(limits, name) != 0 for name in ( - "relative_l2", "normalized_max_abs", "pairwise_rtol", "pairwise_atol" - )): - raise ValueError("schema_version=1 is permitted only for exact-zero test policies") + if any( + getattr(limits, name) != 0 + for name in ( + "relative_l2", + "normalized_max_abs", + "pairwise_rtol", + "pairwise_atol", + ) + ): + raise ValueError( + "schema_version=1 is permitted only for exact-zero test policies" + ) return policy if schema_version != 2: raise ValueError("KV accuracy policy requires schema_version=2") @@ -114,9 +208,13 @@ def load_accuracy_policy(path: Path) -> dict: if policy.get("workload") != WORKLOAD: raise ValueError("KV accuracy policy workload does not match the benchmark contract") qualification = policy.get("qualification", {}) - if (qualification.get("variants") != QUALIFICATION_VARIANTS or - qualification.get("required_receipts") != QUALIFICATION_RECEIPTS): - raise ValueError("KV accuracy policy qualification matrix does not match the source contract") + if ( + qualification.get("variants") != QUALIFICATION_VARIANTS + or qualification.get("required_receipts") != QUALIFICATION_RECEIPTS + ): + raise ValueError( + "KV accuracy policy qualification matrix does not match the source contract" + ) variants = policy.get("variants") if not isinstance(variants, dict): raise ValueError("KV accuracy policy requires a variants object") @@ -124,16 +222,34 @@ def load_accuracy_policy(path: Path) -> dict: if variant not in variants: raise ValueError(f"KV accuracy policy missing variant {variant}") limits = _limits_from_item(variants[variant]) - for name in ("relative_l2", "normalized_max_abs", "pairwise_rtol", "pairwise_atol"): + legacy_pairwise = { + name + for name in ("pairwise_rtol", "pairwise_atol") + if name in variants[variant] + } + if legacy_pairwise: + raise ValueError( + f"{variant} uses obsolete raw allclose fields: {sorted(legacy_pairwise)}" + ) + for name in ("relative_l2", "normalized_max_abs"): if getattr(limits, name) > getattr(ceiling, name): raise ValueError( f"{variant}.{name} exceeds the source-defined engineering ceiling " f"{getattr(ceiling, name):.8g}" ) + envelope = _pairwise_envelope_from_policy(policy) + source_envelope_ceiling = sum( + limits.normalized_max_abs for limits in ENGINEERING_CEILINGS.values() + ) + if envelope.normalized_max_abs > source_envelope_ceiling: + raise ValueError( + "pairwise envelope exceeds the source-derived normalized-maximum ceiling " + f"{source_envelope_ceiling:.8g}" + ) return policy -def load_accuracy_limits(variant: str) -> AccuracyLimits: +def load_accuracy_contract(variant: str) -> tuple[AccuracyLimits, PairwiseEnvelope]: path = os.environ.get("AISP_KV_CACHE_ACCURACY_POLICY") if not path: raise RuntimeError( @@ -147,7 +263,11 @@ def load_accuracy_limits(variant: str) -> AccuracyLimits: item = policy[variant] if policy["schema_version"] == 1 else policy["variants"][variant] except KeyError as exc: raise ValueError(f"KV accuracy policy missing variant {variant}") from exc - return _limits_from_item(item) + return _limits_from_item(item), _pairwise_envelope_from_policy(policy) + + +def load_accuracy_limits(variant: str) -> AccuracyLimits: + return load_accuracy_contract(variant)[0] def reference_cache(model, groups, cache: KVCache) -> KVCache: @@ -178,9 +298,10 @@ def reference_cache(model, groups, cache: KVCache) -> KVCache: return reference -def cache_accuracy(actual: KVCache, expected: KVCache) -> dict[str, float]: +def cache_accuracy_evidence(actual: KVCache, expected: KVCache) -> CacheAccuracyEvidence: """Measure full K/V tensors without checksum cancellation or a giant FP64 copy.""" result = {} + reference_max_abs = 0.0 for name in ("cache_k", "cache_v"): got, ref = getattr(actual, name), getattr(expected, name) if got.shape != ref.shape or got.dtype != ref.dtype or not got.numel(): @@ -190,7 +311,8 @@ def cache_accuracy(actual: KVCache, expected: KVCache) -> dict[str, float]: error_squared = reference_squared = max_error = max_reference = 0.0 flat_got, flat_ref = got.reshape(-1), ref.reshape(-1) for start in range(0, got.numel(), 1 << 20): - g, r = flat_got[start:start + (1 << 20)].double(), flat_ref[start:start + (1 << 20)].double() + g = flat_got[start : start + (1 << 20)].double() + r = flat_ref[start : start + (1 << 20)].double() if not torch.isfinite(g).all() or not torch.isfinite(r).all(): raise AssertionError(f"{name}: non-finite output/reference") error = g - r @@ -198,18 +320,58 @@ def cache_accuracy(actual: KVCache, expected: KVCache) -> dict[str, float]: reference_squared += float(torch.sum(r * r)) max_error = max(max_error, float(error.abs().max())) max_reference = max(max_reference, float(r.abs().max())) - result[f"{name}.relative_l2"] = (math.sqrt(error_squared / reference_squared) - if reference_squared else (0.0 if error_squared == 0 else math.inf)) - result[f"{name}.normalized_max_abs"] = (max_error / max_reference - if max_reference else (0.0 if max_error == 0 else math.inf)) - return result - - -def assert_cache_accuracy(actual: KVCache, expected: KVCache, limits: AccuracyLimits) -> dict[str, float]: - metrics = cache_accuracy(actual, expected) - failures = [f"{name}={value:.8g} > {getattr(limits, name.split('.')[-1]):.8g}" - for name, value in metrics.items() - if not math.isfinite(value) or value > getattr(limits, name.split('.')[-1])] + reference_max_abs = max(reference_max_abs, max_reference) + result[f"{name}.relative_l2"] = ( + math.sqrt(error_squared / reference_squared) + if reference_squared + else (0.0 if error_squared == 0 else math.inf) + ) + result[f"{name}.normalized_max_abs"] = ( + max_error / max_reference + if max_reference + else (0.0 if max_error == 0 else math.inf) + ) + return CacheAccuracyEvidence(metrics=result, reference_max_abs=reference_max_abs) + + +def cache_accuracy(actual: KVCache, expected: KVCache) -> dict[str, float]: + return cache_accuracy_evidence(actual, expected).metrics + + +def _assert_accuracy_evidence( + evidence: CacheAccuracyEvidence, + limits: AccuracyLimits, +) -> CacheAccuracyEvidence: + failures = [ + f"{name}={value:.8g} > {getattr(limits, name.split('.')[-1]):.8g}" + for name, value in evidence.metrics.items() + if not math.isfinite(value) or value > getattr(limits, name.split(".")[-1]) + ] if failures: raise AssertionError("KV cache accuracy failed: " + "; ".join(failures)) - return metrics + return evidence + + +def assert_cache_accuracy_evidence( + actual: KVCache, + expected: KVCache, + limits: AccuracyLimits, +) -> CacheAccuracyEvidence: + return _assert_accuracy_evidence(cache_accuracy_evidence(actual, expected), limits) + + +def assert_cache_accuracy( + actual: KVCache, + expected: KVCache, + limits: AccuracyLimits, +) -> dict[str, float]: + return assert_cache_accuracy_evidence(actual, expected, limits).metrics + + +def pairwise_absolute_tolerance( + envelope: PairwiseEnvelope, + reference_max_abs: float, +) -> float: + if not math.isfinite(reference_max_abs) or reference_max_abs < 0: + raise ValueError("pairwise reference_max_abs must be finite and nonnegative") + return envelope.normalized_max_abs * reference_max_abs diff --git a/code/labs/kv_cache_compression/accuracy_policy.json b/code/labs/kv_cache_compression/accuracy_policy.json index 004a162c7..2c10fb234 100644 --- a/code/labs/kv_cache_compression/accuracy_policy.json +++ b/code/labs/kv_cache_compression/accuracy_policy.json @@ -19,19 +19,23 @@ "operand_format": "E4M3 forward operands with delayed per-tensor scaling", "relative_l2": 0.0625, "normalized_max_abs": 0.0625, - "pairwise_rtol": 0.25, - "pairwise_atol": 0.0625, "rationale": "The independent full-cache aggregate and global-maximum ceilings are one E4M3 half-ULP at a normal binade (2^-4)." }, "nvfp4": { "operand_format": "E2M1 values with per-16-element E4M3 block scale and FP32 tensor scale", "relative_l2": 0.25, "normalized_max_abs": 0.25, - "pairwise_rtol": 0.25, - "pairwise_atol": 0.0625, "rationale": "The independent full-cache aggregate and global-maximum ceilings are one E2M1 half-ULP at a normal binade (2^-2); finite and overflow checks remain mandatory." } }, + "pairwise_envelope": { + "method": "shared-reference-max-triangle-envelope-v1", + "reference_id": "pytorch-unquantized-bf16-full-cache-v1", + "variants": ["fp8", "nvfp4"], + "coefficient": "sum_variant_normalized_max_abs", + "output_rtol": 0.0, + "description": "After both independent full-cache gates pass, compare the raw FP8 and NVFP4 caches with atol=max_abs(shared BF16 reference)*(0.0625+0.25)." + }, "qualification": { "required_receipts": [ {"cohort": "nominal", "seeds": [2026]}, diff --git a/code/labs/kv_cache_compression/baseline_kv_cache.py b/code/labs/kv_cache_compression/baseline_kv_cache.py index 1ba2b87ad..ace20e30e 100644 --- a/code/labs/kv_cache_compression/baseline_kv_cache.py +++ b/code/labs/kv_cache_compression/baseline_kv_cache.py @@ -13,7 +13,12 @@ from core.env import apply_env_defaults from core.harness.benchmark_harness import BaseBenchmark, BenchmarkConfig from labs.kv_cache_compression.accuracy import ( - assert_cache_accuracy, cache_accuracy, load_accuracy_limits, reference_cache, + PairwiseEnvelope, + assert_cache_accuracy_evidence, + cache_accuracy_evidence, + load_accuracy_contract, + pairwise_absolute_tolerance, + reference_cache, ) from labs.kv_cache_compression.kv_cache_common import ( KVCache, @@ -94,6 +99,9 @@ def __init__(self) -> None: self._verify_output_buffer: Optional[torch.Tensor] = None self._accuracy_variant = "fp8" self._accuracy_limits = None + self._pairwise_envelope: Optional[PairwiseEnvelope] = None + self._pairwise_reference_max_abs: Optional[float] = None + self._pairwise_absolute_tolerance: Optional[float] = None self._accuracy_metrics: dict[str, float] = {} def _resolve_device(self) -> torch.device: @@ -106,8 +114,13 @@ def _setup_with_recipe(self, recipe, *, require_accuracy_policy: bool = True) -> if not TE_AVAILABLE or recipe is None: raise RuntimeError(f"SKIPPED: Transformer Engine not available: {TE_IMPORT_ERROR}") - self._accuracy_limits = (load_accuracy_limits(self._accuracy_variant) - if require_accuracy_policy else None) + if require_accuracy_policy: + self._accuracy_limits, self._pairwise_envelope = load_accuracy_contract( + self._accuracy_variant + ) + else: + self._accuracy_limits = None + self._pairwise_envelope = None self.device = self._resolve_device() # Preserve common BF16 weights for an independent unquantized reference. # TE autocast still selects FP8/NVFP4 GEMMs during the benchmark. @@ -211,8 +224,15 @@ def measure_accuracy(self) -> dict[str, float]: raise RuntimeError("benchmark_fn() must run before accuracy measurement") reference = reference_cache(self.model, self._prefill_groups + self._decode_groups, self.cache) if self._accuracy_limits is None: - return cache_accuracy(self.cache, reference) - return assert_cache_accuracy(self.cache, reference, self._accuracy_limits) + evidence = cache_accuracy_evidence(self.cache, reference) + else: + evidence = assert_cache_accuracy_evidence( + self.cache, + reference, + self._accuracy_limits, + ) + self._pairwise_reference_max_abs = evidence.reference_max_abs + return evidence.metrics def _build_verification_output(self) -> torch.Tensor: if self.cache is None or not self._cache_output_ready: @@ -227,6 +247,24 @@ def capture_verification_payload(self) -> None: self.output = self._build_verification_output() if self._batch_size_tensor is None or self._seq_meta_tensor is None: raise RuntimeError("setup() must initialize verification metadata tensors") + if self._pairwise_reference_max_abs is None: + raise RuntimeError("independent full-cache reference gate must run before pair mapping") + if self._accuracy_limits is None: + raise RuntimeError("KV accuracy limits must be available before pair mapping") + if ( + self._accuracy_limits.pairwise_rtol != 0 + or self._accuracy_limits.pairwise_atol != 0 + ): + raise RuntimeError("obsolete raw allclose pairwise fields must remain zero") + envelope = self._pairwise_envelope + if envelope is None: + if self._accuracy_limits.relative_l2 != 0 or self._accuracy_limits.normalized_max_abs != 0: + raise RuntimeError("pairwise reference envelope is missing") + envelope = PairwiseEnvelope(normalized_max_abs=0.0) + self._pairwise_absolute_tolerance = pairwise_absolute_tolerance( + envelope, + self._pairwise_reference_max_abs, + ) self._set_verification_payload( inputs={ "batch_size": self._batch_size_tensor, @@ -246,7 +284,13 @@ def capture_verification_payload(self) -> None: "bf16": self.tensor_dtype == torch.bfloat16, "tf32": torch.backends.cuda.matmul.allow_tf32, }, - output_tolerance=(self._accuracy_limits.pairwise_rtol, self._accuracy_limits.pairwise_atol), + output_tolerance=(0.0, 0.0), + output_tolerances={ + "output": ( + envelope.output_rtol, + self._pairwise_absolute_tolerance, + ) + }, ) def teardown(self) -> None: @@ -261,6 +305,9 @@ def teardown(self) -> None: self._seq_meta_tensor = None self._verify_output_buffer = None self._cache_output_ready = False + self._pairwise_envelope = None + self._pairwise_reference_max_abs = None + self._pairwise_absolute_tolerance = None self._accuracy_metrics = {} torch.cuda.empty_cache() @@ -289,11 +336,37 @@ def get_custom_metrics(self) -> Optional[dict]: tensors = (self.cache.cache_k, self.cache.cache_v) storage_bytes = sum(t.numel() * t.element_size() for t in tensors) bf16_bytes = sum(t.numel() * 2 for t in tensors) + pairwise_metrics = {} + pairwise_envelope = getattr(self, "_pairwise_envelope", None) + pairwise_reference_max_abs = getattr( + self, "_pairwise_reference_max_abs", None + ) + pairwise_absolute_tolerance = getattr( + self, "_pairwise_absolute_tolerance", None + ) + if ( + pairwise_envelope is not None + and pairwise_reference_max_abs is not None + and pairwise_absolute_tolerance is not None + ): + pairwise_metrics = { + "kv_cache.accuracy.pairwise_reference_max_abs": ( + pairwise_reference_max_abs + ), + "kv_cache.accuracy.pairwise_normalized_max_abs": ( + pairwise_envelope.normalized_max_abs + ), + "kv_cache.accuracy.pairwise_absolute_tolerance": ( + pairwise_absolute_tolerance + ), + "kv_cache.accuracy.pairwise_output_rtol": pairwise_envelope.output_rtol, + } return { "kv_cache.storage_bytes": float(storage_bytes), "kv_cache.storage_bits_per_element": float(8 * storage_bytes / sum(t.numel() for t in tensors)), "kv_cache.compression_ratio": float(bf16_bytes / storage_bytes), **{f"kv_cache.accuracy.{k}": v for k, v in self._accuracy_metrics.items()}, + **pairwise_metrics, "kv_cache.batch_size": float(self.batch_size), "kv_cache.seq_len": float(total_tokens), "kv_cache.hidden_dim": float(self.hidden_dim), diff --git a/code/tests/test_kv_pairwise_reference_envelope.py b/code/tests/test_kv_pairwise_reference_envelope.py new file mode 100644 index 000000000..ab0b893cc --- /dev/null +++ b/code/tests/test_kv_pairwise_reference_envelope.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import json +from copy import deepcopy + +import pytest +import torch + +from core.benchmark.verification import resolve_output_tolerances +from core.benchmark.verify_runner import VerifyRunner +from labs.kv_cache_compression.accuracy import ( + DEFAULT_POLICY_PATH, + ENGINEERING_CEILINGS, + PairwiseEnvelope, + assert_cache_accuracy_evidence, + cache_accuracy_evidence, + load_accuracy_contract, + load_accuracy_policy, + pairwise_absolute_tolerance, +) +from labs.kv_cache_compression.kv_cache_common import KVCache + + +def _cache(k: list[float], v: list[float]) -> KVCache: + return KVCache( + torch.tensor(k, dtype=torch.float64), + torch.tensor(v, dtype=torch.float64), + ) + + +def _raw(cache: KVCache) -> torch.Tensor: + return torch.cat((cache.cache_k.reshape(-1), cache.cache_v.reshape(-1))) + + +def test_policy_derives_pairwise_envelope_from_unchanged_independent_limits( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + policy = load_accuracy_policy(DEFAULT_POLICY_PATH) + assert "pairwise_rtol" not in policy["variants"]["fp8"] + assert "pairwise_atol" not in policy["variants"]["nvfp4"] + + monkeypatch.setenv("AISP_KV_CACHE_ACCURACY_POLICY", str(DEFAULT_POLICY_PATH)) + fp8_limits, fp8_envelope = load_accuracy_contract("fp8") + nvfp4_limits, nvfp4_envelope = load_accuracy_contract("nvfp4") + assert fp8_limits.normalized_max_abs == 2**-4 + assert nvfp4_limits.normalized_max_abs == 2**-2 + assert fp8_envelope == nvfp4_envelope + assert fp8_envelope.normalized_max_abs == 2**-4 + 2**-2 + assert fp8_envelope.output_rtol == 0 + + tighter = deepcopy(policy) + tighter["variants"]["fp8"].update(relative_l2=2**-5, normalized_max_abs=2**-5) + tighter["variants"]["nvfp4"].update(relative_l2=2**-3, normalized_max_abs=2**-3) + tighter_path = tmp_path / "tighter.json" + tighter_path.write_text(json.dumps(tighter)) + monkeypatch.setenv("AISP_KV_CACHE_ACCURACY_POLICY", str(tighter_path)) + _, tighter_envelope = load_accuracy_contract("nvfp4") + assert tighter_envelope.normalized_max_abs == 2**-5 + 2**-3 + + +@pytest.mark.parametrize( + "mutation,match", + [ + (("fp8", "pairwise_rtol", 0.25), "obsolete raw allclose fields"), + (("nvfp4", "normalized_max_abs", False), "must be numeric, not boolean"), + ], +) +def test_policy_rejects_obsolete_or_non_numeric_pairwise_inputs( + tmp_path, mutation: tuple[str, str, object], match: str +) -> None: + policy = json.loads(DEFAULT_POLICY_PATH.read_text()) + variant, name, value = mutation + policy["variants"][variant][name] = value + path = tmp_path / "invalid.json" + path.write_text(json.dumps(policy)) + with pytest.raises(ValueError, match=match): + load_accuracy_policy(path) + + +def test_reference_triangle_envelope_handles_cancellation_and_rejects_corruption() -> None: + reference = _cache([16.0, 0.0], [8.0, 0.0]) + fp8 = _cache([16.0, 1.0], [8.0, 0.5]) + nvfp4 = _cache([16.0, -4.0], [8.0, -2.0]) + + fp8_evidence = assert_cache_accuracy_evidence( + fp8, reference, ENGINEERING_CEILINGS["fp8"] + ) + nvfp4_evidence = assert_cache_accuracy_evidence( + nvfp4, reference, ENGINEERING_CEILINGS["nvfp4"] + ) + assert fp8_evidence.reference_max_abs == nvfp4_evidence.reference_max_abs == 16 + assert set(fp8_evidence.metrics) == { + "cache_k.relative_l2", + "cache_k.normalized_max_abs", + "cache_v.relative_l2", + "cache_v.normalized_max_abs", + } + + raw_fp8, raw_nvfp4 = _raw(fp8), _raw(nvfp4) + assert raw_fp8.numel() == fp8.cache_k.numel() + fp8.cache_v.numel() + assert raw_nvfp4.numel() == nvfp4.cache_k.numel() + nvfp4.cache_v.numel() + assert not torch.allclose(raw_fp8, raw_nvfp4, rtol=0.25, atol=0.0625) + + envelope = PairwiseEnvelope(normalized_max_abs=2**-4 + 2**-2) + absolute = pairwise_absolute_tolerance(envelope, fp8_evidence.reference_max_abs) + tolerance_map = {"output": (0.0, absolute)} + assert absolute == 5.0 + assert VerifyRunner().compare_perf_outputs(raw_fp8, raw_nvfp4, tolerance_map).passed + + corrupted = _cache([16.0, -4.125], [8.0, -2.0]) + with pytest.raises(AssertionError, match="KV cache accuracy failed"): + assert_cache_accuracy_evidence( + corrupted, reference, ENGINEERING_CEILINGS["nvfp4"] + ) + corrupted_raw = _raw(corrupted) + assert not VerifyRunner().compare_perf_outputs( + raw_fp8, corrupted_raw, tolerance_map + ).passed + + +def test_zero_reference_is_exact_only_and_invalid_scale_fails_closed() -> None: + reference = _cache([0.0, 0.0], [0.0, 0.0]) + exact = _cache([0.0, 0.0], [0.0, 0.0]) + evidence = cache_accuracy_evidence(exact, reference) + assert evidence.reference_max_abs == 0 + assert all(value == 0 for value in evidence.metrics.values()) + + envelope = PairwiseEnvelope(normalized_max_abs=2**-4 + 2**-2) + assert pairwise_absolute_tolerance(envelope, 0) == 0 + assert VerifyRunner().compare_perf_outputs( + _raw(exact), _raw(exact), {"output": (0.0, 0.0)} + ).passed + + corrupt = _cache([0.0, 0.0], [0.0, 1.0]) + with pytest.raises(AssertionError, match="KV cache accuracy failed"): + assert_cache_accuracy_evidence( + corrupt, reference, ENGINEERING_CEILINGS["nvfp4"] + ) + for bad in (-1.0, float("nan"), float("inf")): + with pytest.raises(ValueError, match="finite and nonnegative"): + pairwise_absolute_tolerance(envelope, bad) + + +def test_pairwise_tolerance_maps_must_match_exactly() -> None: + expected = {"output": (0.0, 5.0)} + assert resolve_output_tolerances( + expected, + dict(expected), + baseline_output_names={"output"}, + optimized_output_names={"output"}, + ) == expected + with pytest.raises(ValueError, match="maps must be identical"): + resolve_output_tolerances( + expected, + {"output": (0.0, 4.999)}, + baseline_output_names={"output"}, + optimized_output_names={"output"}, + ) From 8a7540714c067b8ee0893614d4a503f5262eb0c6 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 04:02:33 -0700 Subject: [PATCH 15/19] docs: record B200 follow-through gains, accuracy gates and profiler limits --- code/ch13/README.md | 22 +- code/core/scripts/refresh_readmes.py | 22 +- .../2026-09-08-b200-followthrough-results.md | 281 ++++++++++++++++++ 3 files changed, 315 insertions(+), 10 deletions(-) create mode 100644 docs/reviews/2026-09-08-b200-followthrough-results.md diff --git a/code/ch13/README.md b/code/ch13/README.md index a4da0e732..0a9abd3a4 100644 --- a/code/ch13/README.md +++ b/code/ch13/README.md @@ -27,7 +27,15 @@ Representative validated results from `artifacts/runs/20260303_163946__bench__pr This chapter is one of the easiest places to fool yourself with framework overhead. That is why the benchmark contract and side-by-side baseline/optimized structure matter here more than almost anywhere else. -The prior `precisionfp8_te` number compared eager FP16 with CUDA-graph-replayed FP8, so it is retired. The pair now runs both sides eagerly to isolate Transformer Engine FP8; publish a new speed result only after a fresh B200 correctness and timing run. +The prior `precisionfp8_te` number compared eager FP16 with CUDA-graph-replayed FP8, so it is retired. The pair now runs both sides eagerly to isolate Transformer Engine FP8. A fresh direct-B200 sweep with Transformer Engine 2.18 measured the full training update after five setup and ten warmup updates: + +| Matched batch | Eager FP16 median | Eager TE FP8 median | FP16 / FP8 | Disposition | +| ---: | ---: | ---: | ---: | --- | +| 256 | `0.4786 ms` | `0.6662 ms` | `0.7183x` | no measured speedup | +| 1,024 | `0.5483 ms` | `0.6644 ms` | `0.8252x` | no measured speedup | +| 4,096 | `1.2303 ms` | `0.9616 ms` | `1.2795x` | candidate speedup for this workload | + +The 24 observations cover two seeds, two repeats, both arms, and all three batches. Every observation compared the actual prediction and all 67,121,152 post-step parameters and passed the frozen output policy. Batch 256 remains the default, so the batch-4,096 result establishes a workload-specific crossover rather than a general default-workload speedup. Whole-call cost is retained separately from the CUDA-event update timing. The TE 2.18 pair now verifies its captured prediction and every post-step parameter with a calibrated per-output policy. Previously, all outputs inherited the global `(rtol=0.5, atol=5.0)` threshold. B200 calibration of the unchanged batch-256, hidden-4096 training step selected these stricter budgets across seed 44 and fixed holdouts 45, 1044, and 1045: @@ -46,13 +54,14 @@ The frozen map passed all holdouts and independently rejected zeroed and localiz python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu ``` -To test an optional larger matched control, pass one pair-wide override through the harness. This sends batch size 1024 to both arms and updates their workload metadata consistently: +To test optional larger matched controls, pass one pair-wide override through the harness. This sends the requested batch to both arms and updates their workload metadata consistently: ```bash python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' +python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 4096' ``` -Compare results only when both arms report the same requested batch size and workload signature. The calibrated output policy above was established at batch 256. Batch 1024 keeps that policy but needs a fresh correctness and B200 timing run; it does not inherit the batch-256 evidence or establish a performance gain by itself. +Compare results only when both arms report the same requested batch size and workload signature. Fresh B200 checks pass at 256, 1,024, and 4,096 with the frozen policy. The first two batches are slower under FP8; the observed 1.2795x result applies only to the matched batch-4,096 workload and does not change the default. ## Profiler Evidence Use deep-dive runs when you want to see whether the gain came from framework overhead reduction, memory behavior, or the lower-precision path itself: @@ -68,6 +77,8 @@ Those targets cover three different PyTorch optimization stories: - `autograd_standard`: framework/compile overhead - `precisionfp8_te`: lower-precision execution with real library support +Matched batch-4,096 Nsys captures pass for both `precisionfp8_te` arms. The FP8 trace shows shorter main GEMMs alongside quantization and scale-update work. Each trace contains one profiled update after setup and one profiler warmup, so trace durations are diagnostic; the repeated sweep above remains the timing authority. + The torchao FP8 recipe demos (`precisionfp8`, `precisionfp8_rowwise`, `precisionfp8_rowwise_gw_hp`) remain useful implementation references, but they are treated as informational examples rather than canonical speed-claim surfaces. ## Repro Commands @@ -77,6 +88,7 @@ python -m cli.aisp bench list-targets --chapter ch13 python -m cli.aisp bench run --targets ch13 --profile minimal python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' +python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 4096' ``` ## Learning Goals @@ -113,12 +125,12 @@ python -m cli.aisp bench run --targets ch13 --profile minimal ## Validation Checklist - `python -m ch13.compare --examples training_standard` shows optimized training runs producing higher goodput with identical metrics. - `python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal` confirms Transformer Engine calibration plus NVFP8 execution with max error tolerances enforced. -- `python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024'` exercises the optional matched batch control; accept its timing only after both arms report batch 1024 and pass fresh output verification. +- Matched B200 `precisionfp8_te` checks at batches 256, 1,024, and 4,096 pass the full-output policy; only batch 4,096 shows a measured candidate speedup (`1.2795x`) for this exact workload. - `python -m ch13.memory_profiling --dump` and the optimized variant demonstrate allocator fragmentation dropping after applying the recommended knobs, with memory reduction treated as the primary benchmark outcome. ## Notes - `custom_allocator.py` contains a standalone torch allocator shim that can be re-used in other chapters when debugging fragmentation. - `compiled_autograd.py` doubles as a tutorial on partial graph capture; the README here references it directly. -- `precisionfp8_te` defaults to batch 256. `--batch-size 1024` is an explicit pair-wide workload override, not a new default or a qualified speed claim. +- `precisionfp8_te` defaults to batch 256. Larger `--batch-size` values are explicit pair-wide workload overrides; the B200 crossover appeared at batch 4,096 and does not change the default. - `torchao_quantization_compiled`, `kv_cache_naive_flash_blockwise`, `precisionfp8`, `precisionfp8_rowwise`, and `precisionfp8_rowwise_gw_hp` remain informational variants. - `kv_cache_naive` and `memory_profiling` are memory-goal benchmarks; they are expected to reduce memory pressure even when the timed path is not faster. diff --git a/code/core/scripts/refresh_readmes.py b/code/core/scripts/refresh_readmes.py index c2859b496..afa35c7dc 100644 --- a/code/core/scripts/refresh_readmes.py +++ b/code/core/scripts/refresh_readmes.py @@ -2339,7 +2339,15 @@ def lab_entry( This chapter is one of the easiest places to fool yourself with framework overhead. That is why the benchmark contract and side-by-side baseline/optimized structure matter here more than almost anywhere else. - The prior `precisionfp8_te` number compared eager FP16 with CUDA-graph-replayed FP8, so it is retired. The pair now runs both sides eagerly to isolate Transformer Engine FP8; publish a new speed result only after a fresh B200 correctness and timing run. + The prior `precisionfp8_te` number compared eager FP16 with CUDA-graph-replayed FP8, so it is retired. The pair now runs both sides eagerly to isolate Transformer Engine FP8. A fresh direct-B200 sweep with Transformer Engine 2.18 measured the full training update after five setup and ten warmup updates: + + | Matched batch | Eager FP16 median | Eager TE FP8 median | FP16 / FP8 | Disposition | + | ---: | ---: | ---: | ---: | --- | + | 256 | `0.4786 ms` | `0.6662 ms` | `0.7183x` | no measured speedup | + | 1,024 | `0.5483 ms` | `0.6644 ms` | `0.8252x` | no measured speedup | + | 4,096 | `1.2303 ms` | `0.9616 ms` | `1.2795x` | candidate speedup for this workload | + + The 24 observations cover two seeds, two repeats, both arms, and all three batches. Every observation compared the actual prediction and all 67,121,152 post-step parameters and passed the frozen output policy. Batch 256 remains the default, so the batch-4,096 result establishes a workload-specific crossover rather than a general default-workload speedup. Whole-call cost is retained separately from the CUDA-event update timing. The TE 2.18 pair now verifies its captured prediction and every post-step parameter with a calibrated per-output policy. Previously, all outputs inherited the global `(rtol=0.5, atol=5.0)` threshold. B200 calibration of the unchanged batch-256, hidden-4096 training step selected these stricter budgets across seed 44 and fixed holdouts 45, 1044, and 1045: @@ -2362,13 +2370,14 @@ def lab_entry( python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu ``` - To test an optional larger matched control, pass one pair-wide override through the harness. This sends batch size 1024 to both arms and updates their workload metadata consistently: + To test optional larger matched controls, pass one pair-wide override through the harness. This sends the requested batch to both arms and updates their workload metadata consistently: ```bash python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' + python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 4096' ``` - Compare results only when both arms report the same requested batch size and workload signature. The calibrated output policy above was established at batch 256. Batch 1024 keeps that policy but needs a fresh correctness and B200 timing run; it does not inherit the batch-256 evidence or establish a performance gain by itself.""" + Compare results only when both arms report the same requested batch size and workload signature. Fresh B200 checks pass at 256, 1,024, and 4,096 with the frozen policy. The first two batches are slower under FP8; the observed 1.2795x result applies only to the matched batch-4,096 workload and does not change the default.""" ), ), MarkdownSection( @@ -2388,6 +2397,8 @@ def lab_entry( - `autograd_standard`: framework/compile overhead - `precisionfp8_te`: lower-precision execution with real library support + Matched batch-4,096 Nsys captures pass for both `precisionfp8_te` arms. The FP8 trace shows shorter main GEMMs alongside quantization and scale-update work. Each trace contains one profiled update after setup and one profiler warmup, so trace durations are diagnostic; the repeated sweep above remains the timing authority. + The torchao FP8 recipe demos (`precisionfp8`, `precisionfp8_rowwise`, `precisionfp8_rowwise_gw_hp`) remain useful implementation references, but they are treated as informational examples rather than canonical speed-claim surfaces.""" ), ), @@ -2401,6 +2412,7 @@ def lab_entry( python -m cli.aisp bench run --targets ch13 --profile minimal python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024' + python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile deep_dive --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 4096' ```""" ), ), @@ -2426,13 +2438,13 @@ def lab_entry( validation=[ "`python -m ch13.compare --examples training_standard` shows optimized training runs producing higher goodput with identical metrics.", "`python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal` confirms Transformer Engine calibration plus NVFP8 execution with max error tolerances enforced.", - "`python -m cli.aisp bench run --targets ch13:precisionfp8_te --profile minimal --single-gpu --target-extra-arg 'ch13:precisionfp8_te=--batch-size 1024'` exercises the optional matched batch control; accept its timing only after both arms report batch 1024 and pass fresh output verification.", + "Matched B200 `precisionfp8_te` checks at batches 256, 1,024, and 4,096 pass the full-output policy; only batch 4,096 shows a measured candidate speedup (`1.2795x`) for this exact workload.", "`python -m ch13.memory_profiling --dump` and the optimized variant demonstrate allocator fragmentation dropping after applying the recommended knobs, with memory reduction treated as the primary benchmark outcome.", ], notes=[ "`custom_allocator.py` contains a standalone torch allocator shim that can be re-used in other chapters when debugging fragmentation.", "`compiled_autograd.py` doubles as a tutorial on partial graph capture; the README here references it directly.", - "`precisionfp8_te` defaults to batch 256. `--batch-size 1024` is an explicit pair-wide workload override, not a new default or a qualified speed claim.", + "`precisionfp8_te` defaults to batch 256. Larger `--batch-size` values are explicit pair-wide workload overrides; the B200 crossover appeared at batch 4,096 and does not change the default.", "`torchao_quantization_compiled`, `kv_cache_naive_flash_blockwise`, `precisionfp8`, `precisionfp8_rowwise`, and `precisionfp8_rowwise_gw_hp` remain informational variants.", "`kv_cache_naive` and `memory_profiling` are memory-goal benchmarks; they are expected to reduce memory pressure even when the timed path is not faster.", ], diff --git a/docs/reviews/2026-09-08-b200-followthrough-results.md b/docs/reviews/2026-09-08-b200-followthrough-results.md new file mode 100644 index 000000000..3f4ef921b --- /dev/null +++ b/docs/reviews/2026-09-08-b200-followthrough-results.md @@ -0,0 +1,281 @@ +# B200 follow-through: serving, arithmetic requirements, and profiler recovery + +Status: scoped target validation complete, with no-win outcomes and an unresolved collective-replay limitation. This report extends the [September 7 results](2026-09-07-b200-review-results.md); it does not replace their retained failures or claim that every example is faster. + +## Source and execution scope + +DDP and the initial arithmetic matrix use `ce4681dd63740ca10a66ea3627d83a6e2d9e7a1b`; +the subsequent pipeline, cache, and FP8 checkpoint is +`aea503d9a88bd0326d93d27d5f480507febc076c`. +The final Ozaki and vLLM reruns use clean checkpoint +`2aae639f0cb10dc7ca01a1e7bbb7992f33a732fd`; their receipts bind the exact source +files and runtime used. The fresh KV envelope rerun uses clean checkpoint +`f10a5b7553825909b66b011f5c60dbb9809713b9`. Results from these checkpoints are +kept distinct below. +Execution uses one or two B200 GPUs directly, with the repository harness and an +owned, serialized process supervisor. The host reports virtualization, so these +are explicitly portable development results. They are not canonical bare-metal +qualification. CUDA, PyTorch, and the driver remain unchanged. + +Receipts use the run names below under `artifacts/parallel_runs/followthrough_20260908`. +Full-output verification receipts, failed attempts, traces, and supervision receipts are retained +privately; large tensors and profiler binaries are not committed to this public repository. +The final materialized copy contains 645 files totaling 2,797,167,819 bytes. +Remote and local SHA-256 inventories are byte-identical, with inventory digest +`0ff1e4dbeb322d02360422ed4e61ff30e91ad3f2aa45851abd3967b7cb1c5e13`. + +## Changes and measured disposition + +| Area | Change | Current disposition | +| --- | --- | --- | +| DDP | Store sampled losses on the device and read the retained history once after training | Exact outputs pass on one and two B200s; no training-throughput win established | +| Pipeline | Amortize fill/drain boundaries across repeated fixed-weight iterations | 16 B200 runs pass exact outputs; 1.024x iteration / 1.003x process medians with mixed seed results; paired traces pass their exact mechanism contracts | +| Cache-aware inference | Remove redundant barriers for the exact 1P1D topology; report zero placement opportunities | 16 timing runs and both traces pass; observed median ratio 1.574x from reduced synchronization | +| vLLM routing and dual pools | Reuse engines; separate startup, five warmups, three wall-clock measurements, and teardown | Final dual-pool run passes exact tokens at 1.472x steady-state; dynamic routing passes exact tokens at 0.885x and is a no-win | +| TE FP8 | Expose the same explicit batch-size override in FP16 and FP8 while preserving batch 256 by default | 24 full-output observations show a workload-specific crossover: 0.718x at batch 256, 0.825x at 1024, and 1.279x at 4096 | +| Arithmetic requirements | Freeze independent numerical ceilings and require nominal, holdout, edge, and shared-reference checks | Fresh KV and Ozaki runs pass correctness; both speed goals fail without a claimed win | +| Execution hardening | Add explicit operation-placement and declared-destination write-coverage audits | Final 27 B200 tests and a real CUDA audit CLI pass, with no skips | +| Nsight Compute | Isolated newer tool versions and exactly five metrics for minimal captures | Paired 2026.2.1 selected-GEMM captures pass; application-range and coordinated collective-kernel replay time out at 180 seconds | + +### DDP + +The one-GPU comparison uses batch 16 and 100 training steps. The two-GPU +comparison uses batch 32 per rank and all 32 steps available from its loader. +Each topology has two fresh seeds and an ABBA order, yielding four observations +per arm. Full inputs and final outputs match exactly, and worker runtime receipts +match the workers that produced those outputs. + +| Scope | Control median | Candidate median | Control / candidate | +| --- | ---: | ---: | ---: | +| One-GPU training iteration | 46.9307 ms | 46.7522 ms | 1.0038x | +| One-GPU process wall time | 13,985.3 ms | 14,012.4 ms | 0.9981x | +| Two-GPU training iteration | 84.6293 ms | 84.7961 ms | 0.9980x | +| Two-GPU process wall time | 17,533.9 ms | 17,211.8 ms | 1.0187x | + +These are observed ratios, not a qualified speedup. The one-GPU control includes +a large cold-start outlier; the worker-loop results are near parity. Deferring +logging does not make its cost disappear from process wall time. The early +two-GPU drivers incorrectly expected 100 and then 64 steps; both were rejected +when actual execution reported 32, retained, and replaced by a fresh 32-step plan. +Receipts: `ddp-progress-abba/receipt.json` (completed one-GPU rows; later driver +failure retained) and `ddp-progress-abba-w2-32steps/receipt.json`. + +### Pipeline and cache-aware inference + +The pipeline comparison completed four seeds in ABBA order: 16 observations, +complete two-rank runtime receipts, and bitwise-equal full outputs. Median +iteration times were 15.5584 ms versus 15.1887 ms, while process wall times were +8,514.49 ms versus 8,488.87 ms. Per-seed iteration ratios ranged from 0.967x +to 1.071x. The improvement is too small and inconsistent to claim a robust win. +The change applies to this synthetic fixed-weight repeated workload; it cannot +combine iterations separated by real optimizer updates. + +The baseline Nsys capture initially failed during table export because its +SQLite cache was older than the report. A separate recovery directory preserves +the original failure and uses the supported `--force-export=true` option. The +recovered baseline and fresh optimized traces both pass their exact operation-count +contracts: each contains 1,024 GEMMs, 1,024 ReLUs, and 128 NCCL sends plus 128 +receives. The optimized trace reduces P2P kernels from 256 to 132 and total GPU +operations in the two measured ranges from 864 to 818. These counts explain the +mechanism; the separate ABBA timing remains the performance evidence. Receipts: +`pipeline-contiguous-abba/receipt.json`, +`pipeline-contiguous-nsys-baseline/recovery-force-export-v1/recovery-receipt.json`, +and `pipeline-contiguous-nsys-optimized/receipt.json`. + +Cache-aware inference completed 16 timing runs across four seeds and two matched +Nsys captures, with complete outputs, worker/runtime PID parity, and per-rank +application-clock checks. Baseline and optimized medians were 14.0537 ms and +8.9308 ms. All four block ratios favored the candidate, ranging from 1.441x +to 1.668x. Baseline/optimized standard deviations were 1.1614/0.2731 ms. + +The traces show float32 NCCL all-reduce launches dropping from 1,328 to 176 and +`cudaStreamSynchronize` calls from 1,424 to 272, while send/receive kernel counts +remain 660 in both arms. The reduction of 1,152 matches eight avoided barriers +per request across eight requests, nine warmup/measured invocations, and two +ranks. Both arms report the same cache hit rate and transfer volume. There is +only one decode rank, so neither arm has an alternative cache placement. +Receipt: `cache-aware-locked-abba/receipt.json` reports +`FULL_OUTPUT_RUNTIME_CLOCK_ABBA_NSYS_PASS` with 18 rows. + +### FP8 training + +Fresh matched batch-256 Nsys captures show the optimizer's four tensor-update +kernels taking about 106–107 microseconds in each arm. The FP8 capture adds +approximately 76 microseconds of quantization and scale-update kernels while +its GEMMs become shorter. A subsequent sweep used the full 67,121,152-parameter +model, compared the actual prediction and every post-step parameter, and passed the +frozen output policy for all 24 observations: two seeds, two repeats, two arms, +and three matched batches. + +| Batch | Eager FP16 median | Eager TE FP8 median | FP16 / FP8 | Disposition | +| ---: | ---: | ---: | ---: | --- | +| 256 | 0.4786 ms | 0.6662 ms | 0.7183x | No measured speedup | +| 1,024 | 0.5483 ms | 0.6644 ms | 0.8252x | No measured speedup | +| 4,096 | 1.2303 ms | 0.9616 ms | 1.2795x | Candidate speedup for this workload | + +The primary timing is one complete CUDA-event training update after five setup +and ten warmup updates; whole-call time is retained separately. Batch 256 remains +the default, so the 4,096 result establishes a workload-specific crossover rather +than a general default-workload speedup. Matched batch-4,096 Nsys captures also +pass for both arms. They show shorter main FP8 GEMMs together with quantization +and scale-update work; their single profiled update is diagnostic and is not used +as benchmark timing. Receipts: `te218-matched-batch-sweep-v2-receipt.json` and +`te218-batch4096-profiles-receipt.json`. + +The first private sweep driver rejected an actual baseline signature because +the driver omitted the mixin's output shape and dtype from its expected fields. +The failed receipt is retained. The driver was corrected to require those fields; +the benchmark and all numerical tolerances remain unchanged. + +### Reused serving engines + +Three ordinary-harness runs of `labs/dynamic_router:dual_pool_vllm` +passed exact token-output checks and runtime parity. Each arm starts two engines +once, runs five warmups, then measures three complete 102-request workloads. +Prefix caching is disabled and request state resets between invocations. + +| Run | Shared-pool wall time | Dual-pool wall time | Observed ratio | +| --- | ---: | ---: | ---: | +| 1 | 1,870.687 ms | 1,280.290 ms | 1.461x | +| 2 | 1,872.504 ms | 1,273.115 ms | 1.471x | +| Final source | 1,866.333 ms | 1,268.244 ms | 1.472x | + +The second run reports 54.47 versus 80.12 requests/second. Startup is reported +separately at 27.21–27.64 seconds, and final lifecycle receipts retain 1.91–2.64 +seconds of teardown. Including startup, warmups, measured requests, bookkeeping, +and teardown, shared-pool sessions took 47.32 and 47.11 seconds; dual-pool +sessions took 44.89 and 44.70 seconds. The steady-state ratio must not be applied +to those total session durations. All four lifecycle receipts for the first two +runs report completed disposition, no failed requests, and no shutdown errors. + +The matched final-source Nsys pair passes full token equality (1,734 integers), +runtime parity, and the five-warmup/one-steady lifecycle contract. Baseline puts +all 19,222 captured kernels on device 0. The dual-pool arm splits 15,550 kernels +onto device 0 and 12,178 onto device 1; kernel-busy interval union falls from +1.812668 to 1.218329 seconds, and 99.0% of device 1 busy time overlaps device 0. +Summed GPU duration still rises 5.3% and launches rise 44.3%, while MoE BMM time +is within 1.8% and attention within 0.5%. This supports concurrent placement, +not removed arithmetic. Profiled durations are diagnostic; the repeated ordinary +runs above remain the timing evidence. + +The reused dynamic-routing pair also passes exact token checks and lifecycle +separation, but is slower: 209.280 ms for the static control versus 236.455 ms +for dynamic routing, or 0.885x. That run is retained as `failed_no_speedup`, not +as a routing win. Receipts: `vllm-dual-pool-reuse-{1,2}`, their adjacent +`-lifecycle-final.json` extracts, `vllm-dual-pool-final-2aae`, +`vllm-dual-pool-nsys-2aae/{baseline,optimized}/receipt.json`, and +`vllm-dynamic-routing-reuse-2aae`. + +### Numerical acceptance + +The policies were fixed before the qualification runs and reject widened JSON +limits. The retained 10-case matrices passed one nominal case, two holdouts, and +two edge cases for each variant on the recorded B200 stack. The later Ozaki +ordinary-harness run also passes full correctness for both variants, while its +speed goal correctly fails: dynamic reaches 0.545x and fixed reaches 0.724x +relative to native FP64. These are no-speedup results. + +| Variant | Relative-L2 ceiling | Maximum error / maximum reference ceiling | +| --- | ---: | ---: | +| KV FP8 E4M3 | 0.0625 | 0.0625 | +| KV NVFP4 E2M1 | 0.25 | 0.25 | +| Ozaki dynamic, max 16 bits, offset -56 | 0.03125 | 0.0625 | +| Ozaki fixed 12 bits | 0.000244140625 | 0.00048828125 | + +KV checks every stored cache element against an unquantized BF16 PyTorch +projection. Ozaki compares complete production-size arrays with native FP64, +and the small rectangular edges additionally use a CPU long-double reference. +The x86 host ran that independent reference successfully. The limits are explicit +engineering requirements, not arbitrary-matrix error theorems or application-quality +guarantees. See the [KV requirements](../../code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md) +and [Ozaki requirements](../../code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md). +Receipts: `accuracy-matrix/kv-qualification.json`, +`accuracy-matrix/ozaki-qualification.json`, and `ozaki-frozen-policy-2aae`. +The updated assessors also accept all 20 retained cases while preserving their +original declared provenance. This is a reassessment of those receipts, not a +claim that their kernels reran at the later source checkpoint. + +The earlier ordinary KV pair failed its secondary raw `allclose` criterion even +though both independent gates passed. Near cancellation, that local relative +criterion is inconsistent with the declared reference-normalized limits. The +replacement follows the triangle inequality: the full raw pairwise difference +must stay within `(0.0625 + 0.25) * max(abs(reference))`, with exactly matching +per-output maps in both arms. Both independent L2 and maximum-error gates remain +mandatory and unchanged. The original failed receipt remains failed. Ozaki's +ordinary runner similarly needed the baseline to expose the already-declared +secondary comparison envelope; its full-array requirements were unchanged. + +The fresh KV ordinary-harness run passes input, full-output, pairwise-reference, +and runtime checks. Its shared reference has maximum magnitude 17.75; the frozen +0.3125 coefficient gives an absolute envelope of 5.546875, and the observed raw +maximum difference is 2.8125. FP8 reaches maximum relative-L2 0.040992 and +normalized-maximum error 0.044118; NVFP4 reaches 0.146155 and 0.159467, all under +the unchanged ceilings above. Baseline and optimized times are 557.582 and +544.132 ms, a 1.0247x ratio below the required 1.05x threshold. The receipt +`kv-cache-reference-envelope-final` therefore retains `failed_no_speedup` while +its correctness gates pass. + +### Practical hardening + +The new standalone `python -m core.harness.execution_audit` tool uses a fresh +benchmark instance outside normal timing. It observes dispatcher-visible tensor +operations and can poison explicitly declared floating-point or complex output +buffers to check that their complete logical extent was written. It records +operation/device evidence and rejects replaced destination identities. + +The placement audit covers the current Python thread's PyTorch dispatcher. +Destination coverage applies to the exact declared contiguous buffers. Neither +claims to inspect arbitrary extension internals, other processes, or all +uninitialized-memory provenance. Nine stale declarations now execute real checks, +reducing the explicit missing-protection declarations from 42 to 33; those counts +are declarations, not distinct confirmed bugs. +The audit also rejects no-op and allowed-host-only callbacks that provide no +evidence of execution on the requested device, and preserves a primary failure +when teardown itself fails. + +### Nsight recovery + +Nsight Compute 2026.1.1 and 2026.2.1 were unpacked into private tool directories +after checking NVIDIA package sizes and SHA-256 digests. No system package hooks, +driver update, or shared CUDA/PyTorch change was used. + +The old `minimal` automation path also selected NVIDIA's `basic` sections, adding +many counters beyond the requested list. The shared automation and harness now +request exactly the five minimal metrics; explicit `basic` still selects NVIDIA +sections. CLI/MCP descriptions and regression coverage agree with this behavior. + +On 2026.2.1, a selected GEMM on device 0 completed with both pipeline ranks and +exactly five requested metrics in each arm. The baseline and optimized kernels +reported 164,160 and 162,848 ns, 91.62% and 92.04% SM throughput, and 15.41% and +15.52% DRAM throughput. Both receipts include the two executed rank PIDs and their +runtime provenance. This is one selected kernel's mechanism evidence, not a +complete-workload timing comparison. The earlier 2026.1.1 parser omitted the L2 +counter prefix; a separate retained-artifact validation corrected that parser +error without replacing the original receipt. + +Range capture rejected `cuThreadExchangeStreamCaptureMode`; on 2026.2.1 both an +application-range replay and a shared-memory NCCL attempt without NVTX timed out +after 180 seconds. Each timed-out capture was preserved and fully drained. +NVIDIA documents mandatory concurrent-kernel +coordination and replay constraints in its [Nsight Compute CLI guide](https://docs.nvidia.com/nsight-compute/NsightComputeCli/index.html#mandatory-concurrent-kernels). +The selected-kernel path is recovered; full-range replay remains unresolved. + +## Checks and remaining limitations + +- Final affected CPU source integration: 913 passed, 103 explicit capability or missing-declaration skips, and 30 warnings across 17 modules. +- Profiler command, CLI, MCP-document, and harness contracts: 132 passed, 1 skip. +- KV reference-envelope compatibility: 57 passed, 1 capability skip. +- Earlier B200 execution guards and deferred-progress behavior: 23 passed, no skips. +- Late-source B200 audits: 27 passed with no skips, and the real vectorization CLI passes at `f10a5b755`. +- Retained B200 arithmetic matrix: 20 of 20 cases pass; fresh KV and Ozaki ordinary-harness correctness passes with no qualified speedup. +- Syntax passes for all 48 changed Python files. Full-file Ruff has the same 75 diagnostics as the base commit and no new diagnostics; focused changed-implementation lint passes. +- All target runs are finished and their owned processes drained; artifact hash verification passes. + +Full collective NCU replay still needs a working tool/runtime combination; the +bounded selected-kernel captures and matched Nsys traces provide the usable +profiling paths on this stack. DDP and pipeline need a repeatable measured gain +before stronger performance claims. Small-batch FP8, dynamic routing, KV, and +Ozaki remain no-win examples for the measured workloads. The arithmetic budgets +cover these lab outputs; downstream model-quality acceptance remains a separate +requirement. The remaining 33 protection declarations require individual review, +not treatment as 33 confirmed bugs. From 73ab9ef8517376b04afa933eb61c9e30b08c8e25 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 04:36:13 -0700 Subject: [PATCH 16/19] ci: leave enough time for complete CPU validation on hosted runners --- .github/workflows/benchmark-validation.yml | 4 +++- docs/reviews/2026-09-08-b200-followthrough-results.md | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/.github/workflows/benchmark-validation.yml b/.github/workflows/benchmark-validation.yml index 6c5ec3bdc..e3d528340 100644 --- a/.github/workflows/benchmark-validation.yml +++ b/.github/workflows/benchmark-validation.yml @@ -55,7 +55,9 @@ defaults: jobs: validate: runs-on: ubuntu-latest - timeout-minutes: 30 + # Full CPU validation can approach 30 minutes on slower hosted runners. + # Keep the complete suite and leave time for the final audits and artifacts. + timeout-minutes: 35 steps: - name: Checkout repository diff --git a/docs/reviews/2026-09-08-b200-followthrough-results.md b/docs/reviews/2026-09-08-b200-followthrough-results.md index 3f4ef921b..54cf30ff9 100644 --- a/docs/reviews/2026-09-08-b200-followthrough-results.md +++ b/docs/reviews/2026-09-08-b200-followthrough-results.md @@ -271,6 +271,12 @@ The selected-kernel path is recovered; full-range replay remains unresolved. - Syntax passes for all 48 changed Python files. Full-file Ruff has the same 75 diagnostics as the base commit and no new diagnostics; focused changed-implementation lint passes. - All target runs are finished and their owned processes drained; artifact hash verification passes. +The first hosted CPU validation attempt reached 98% before its 30-minute job +budget expired. Its cancelled result is retained. The workflow now allows 35 +minutes, preserving the entire test suite and all final audits. A prior run had +also exhausted 30 minutes after its full suite and linter passed; the added +margin addresses observed hosted-runner variability rather than skipping work. + Full collective NCU replay still needs a working tool/runtime combination; the bounded selected-kernel captures and matched Nsys traces provide the usable profiling paths on this stack. DDP and pipeline need a repeatable measured gain From 7db240c6689f8e3d93c83fcd33d9f51e7cfbdfa6 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 04:51:01 -0700 Subject: [PATCH 17/19] test: preserve profiler range checks for contiguous pipeline schedules --- .../tests/test_ch04_direct_profiler_ranges.py | 84 ++++++++++++++++--- .../2026-09-08-b200-followthrough-results.md | 11 ++- 2 files changed, 81 insertions(+), 14 deletions(-) diff --git a/code/tests/test_ch04_direct_profiler_ranges.py b/code/tests/test_ch04_direct_profiler_ranges.py index d31b931dd..b72ef51bf 100644 --- a/code/tests/test_ch04_direct_profiler_ranges.py +++ b/code/tests/test_ch04_direct_profiler_ranges.py @@ -168,19 +168,77 @@ def test_target_range_encloses_only_the_post_warmup_timed_loop( for keyword in range_call.keywords ) - warmup_loops = [ - node - for node in ast.walk(worker) - if isinstance(node, ast.For) and ast.unparse(node.iter) == "range(max(warmup, 0))" - ] - timed_loops = [ - node - for node in ast.walk(profile_range) - if isinstance(node, ast.For) and ast.unparse(node.iter) == "range(max(iters, 1))" - ] - assert len(warmup_loops) == 1 - assert len(timed_loops) == 1 - assert warmup_loops[0].end_lineno < profile_range.lineno < timed_loops[0].lineno + if module_name == "ch04.optimized_pipeline_parallel_1f1b": + iteration_assignments = [ + (target.id, node) + for node in ast.walk(worker) + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) + and target.id in {"warmup_iterations", "measured_iterations"} + ] + assert len(iteration_assignments) == 2 + assignments = dict(iteration_assignments) + assert set(assignments) == {"warmup_iterations", "measured_iterations"} + assert ast.unparse(assignments["warmup_iterations"].value) == "max(warmup, 0)" + assert ast.unparse(assignments["measured_iterations"].value) == "max(iters, 1)" + + iteration_calls = [ + node + for node in ast.walk(worker) + if isinstance(node, ast.Call) + and _call_name(node) == "_run_contiguous_iterations" + ] + assert len(iteration_calls) == 2 + warmup_calls = [ + node + for node in iteration_calls + if node.args and ast.unparse(node.args[0]) == "warmup_iterations" + ] + measured_calls = [ + node + for node in iteration_calls + if node.args and ast.unparse(node.args[0]) == "measured_iterations" + ] + assert len(warmup_calls) == 1 + assert len(measured_calls) == 1 + + warmup_guards = [ + node + for node in ast.walk(worker) + if isinstance(node, ast.If) + and isinstance(node.test, ast.Name) + and node.test.id == "warmup_iterations" + and warmup_calls[0] in ast.walk(node) + ] + assert len(warmup_guards) == 1 + profile_node_ids = {id(node) for node in ast.walk(profile_range)} + assert id(warmup_calls[0]) not in profile_node_ids + assert id(measured_calls[0]) in profile_node_ids + assert ( + assignments["warmup_iterations"].lineno + < warmup_guards[0].lineno + < warmup_calls[0].lineno + < profile_range.lineno + < assignments["measured_iterations"].lineno + < measured_calls[0].lineno + ) + else: + warmup_loops = [ + node + for node in ast.walk(worker) + if isinstance(node, ast.For) + and ast.unparse(node.iter) == "range(max(warmup, 0))" + ] + timed_loops = [ + node + for node in ast.walk(profile_range) + if isinstance(node, ast.For) + and ast.unparse(node.iter) == "range(max(iters, 1))" + ] + assert len(warmup_loops) == 1 + assert len(timed_loops) == 1 + assert warmup_loops[0].end_lineno < profile_range.lineno < timed_loops[0].lineno range_calls = { _call_name(node) diff --git a/docs/reviews/2026-09-08-b200-followthrough-results.md b/docs/reviews/2026-09-08-b200-followthrough-results.md index 54cf30ff9..0985971ac 100644 --- a/docs/reviews/2026-09-08-b200-followthrough-results.md +++ b/docs/reviews/2026-09-08-b200-followthrough-results.md @@ -268,7 +268,7 @@ The selected-kernel path is recovered; full-range replay remains unresolved. - Earlier B200 execution guards and deferred-progress behavior: 23 passed, no skips. - Late-source B200 audits: 27 passed with no skips, and the real vectorization CLI passes at `f10a5b755`. - Retained B200 arithmetic matrix: 20 of 20 cases pass; fresh KV and Ozaki ordinary-harness correctness passes with no qualified speedup. -- Syntax passes for all 48 changed Python files. Full-file Ruff has the same 75 diagnostics as the base commit and no new diagnostics; focused changed-implementation lint passes. +- Syntax passes for all 49 changed Python files. Full-file Ruff has the same 75 diagnostics as the base commit and no new diagnostics; focused changed-implementation lint passes. - All target runs are finished and their owned processes drained; artifact hash verification passes. The first hosted CPU validation attempt reached 98% before its 30-minute job @@ -276,6 +276,15 @@ budget expired. Its cancelled result is retained. The workflow now allows 35 minutes, preserving the entire test suite and all final audits. A prior run had also exhausted 30 minutes after its full suite and linter passed; the added margin addresses observed hosted-runner variability rather than skipping work. +The cancelled log also contained a failing profiler-range contract test. A focused +reproducer identified its obsolete assumption that the optimized 1F1B worker must +use per-iteration `for` loops. The test now checks the contiguous schedule's exact +warmup and measured counts and their positions outside and inside the NVTX range, +respectively. Timer and synchronization checks remain in place; the benchmark +implementation is unchanged from the B200 profile captures. +The corrected profiler-range module passes all 42 tests, the broader pipeline +regression set passes 48 tests with six hardware skips, and all 24 workflow +configuration tests pass. Full collective NCU replay still needs a working tool/runtime combination; the bounded selected-kernel captures and matched Nsys traces provide the usable From e582b906abfe0823934acf40248e047cedd7c439 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 05:07:52 -0700 Subject: [PATCH 18/19] Use portable paths in validation documentation --- ...6-09-08-pipeline-1f1b-contiguous-performance-intake.yaml | 6 +++--- code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md | 2 +- code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml b/code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml index 25ef82c0f..abcb2e71c 100644 --- a/code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml +++ b/code/docs/reviews/2026-09-08-pipeline-1f1b-contiguous-performance-intake.yaml @@ -92,10 +92,10 @@ hot_path_model: best_case_primary_kpi_improvement_pct: null evidence: - baseline_artifact: "/Users/admin/.codex/artifacts/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-abba-results/receipt.json" + baseline_artifact: "/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-abba-results/receipt.json" profile_artifacts: - - "/Users/admin/.codex/artifacts/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-nsys-baseline-results/exports-v2" - - "/Users/admin/.codex/artifacts/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-nsys-optimized-results/exports-v2" + - "/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-nsys-baseline-results/exports-v2" + - "/ai-perf-remaining-20260906/wave72/wave72-validation-20260908T055519Z/pipeline-nsys-optimized-results/exports-v2" allocation_or_memory_profile: "none" hardware_counter_artifacts: [] static_review_findings: diff --git a/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md b/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md index cc9efb537..96c8edb5c 100644 --- a/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md +++ b/code/labs/kv_cache_compression/ACCURACY_REQUIREMENTS.md @@ -56,7 +56,7 @@ reason in its summary. Run these commands serially on the requested B200 from `code/`: ```bash -accuracy_out=/tmp/ai-perf-followthrough-20260908-private/accuracy +accuracy_out="${TMPDIR:-/tmp}/ai-perf-accuracy" mkdir -p "$accuracy_out" for variant in fp8 nvfp4; do diff --git a/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md b/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md index fab97c501..09518195f 100644 --- a/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md +++ b/code/labs/ozaki_scheme/ACCURACY_REQUIREMENTS.md @@ -50,7 +50,7 @@ From `code/labs/ozaki_scheme/` on the requested B200, build once: ```bash make ARCH=sm_100 all -accuracy_out=/tmp/ai-perf-followthrough-20260908-private/accuracy +accuracy_out="${TMPDIR:-/tmp}/ai-perf-accuracy" mkdir -p "$accuracy_out" ``` @@ -109,8 +109,8 @@ Qualify the complete set from `code/`: ```bash python -m labs.ozaki_scheme.qualify_accuracy \ --policy labs/ozaki_scheme/accuracy_policy.json \ - --output /tmp/ai-perf-followthrough-20260908-private/accuracy/ozaki-qualification.json \ - /tmp/ai-perf-followthrough-20260908-private/accuracy/ozaki-*.log + --output "$accuracy_out/ozaki-qualification.json" \ + "$accuracy_out"/ozaki-*.log ``` Only after the summary says `qualified_arithmetic_gate`, run the ordinary pair with From 4b1cfb801da70ca58aa02b5960894adb0bd20bbd Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Tue, 8 Sep 2026 05:46:27 -0700 Subject: [PATCH 19/19] Refresh FP8 documentation regression after measured sweep --- .../test_wave2_medium_ch06_ch09_ch13_regressions.py | 11 ++++++++++- .../reviews/2026-09-08-b200-followthrough-results.md | 12 +++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/code/tests/test_wave2_medium_ch06_ch09_ch13_regressions.py b/code/tests/test_wave2_medium_ch06_ch09_ch13_regressions.py index 7efdf1bff..8928d88dc 100644 --- a/code/tests/test_wave2_medium_ch06_ch09_ch13_regressions.py +++ b/code/tests/test_wave2_medium_ch06_ch09_ch13_regressions.py @@ -150,7 +150,16 @@ def test_transformer_engine_precision_pair_keeps_optimized_path_eager() -> None: readme = (CODE_ROOT / "ch13" / "README.md").read_text(encoding="utf-8") assert "compared eager FP16 with CUDA-graph-replayed FP8" in readme - assert "publish a new speed result only after a fresh B200" in readme + assert "`precisionfp8_te` keeps batch size 256 as its default workload" in readme + assert ( + "| 256 | `0.4786 ms` | `0.6662 ms` | `0.7183x` | no measured speedup |" + in readme + ) + assert ( + "batch-4,096 result establishes a workload-specific crossover rather than " + "a general default-workload speedup" + in readme + ) def test_transformer_engine_eager_benchmark_runs_one_fp8_training_step() -> None: diff --git a/docs/reviews/2026-09-08-b200-followthrough-results.md b/docs/reviews/2026-09-08-b200-followthrough-results.md index 0985971ac..bb6c41e59 100644 --- a/docs/reviews/2026-09-08-b200-followthrough-results.md +++ b/docs/reviews/2026-09-08-b200-followthrough-results.md @@ -268,7 +268,7 @@ The selected-kernel path is recovered; full-range replay remains unresolved. - Earlier B200 execution guards and deferred-progress behavior: 23 passed, no skips. - Late-source B200 audits: 27 passed with no skips, and the real vectorization CLI passes at `f10a5b755`. - Retained B200 arithmetic matrix: 20 of 20 cases pass; fresh KV and Ozaki ordinary-harness correctness passes with no qualified speedup. -- Syntax passes for all 49 changed Python files. Full-file Ruff has the same 75 diagnostics as the base commit and no new diagnostics; focused changed-implementation lint passes. +- Syntax passes for all 50 changed Python files. Changed-file Ruff `F` checks have the same 75 diagnostics as the base commit and no new `F` diagnostics; the repository-wide CI correctness rules pass. - All target runs are finished and their owned processes drained; artifact hash verification passes. The first hosted CPU validation attempt reached 98% before its 30-minute job @@ -286,6 +286,16 @@ The corrected profiler-range module passes all 42 tests, the broader pipeline regression set passes 48 tests with six hardware skips, and all 24 workflow configuration tests pass. +The next completed hosted CPU suite reported 5,528 passed, 527 skipped, and one +failing documentation assertion. That assertion still expected the README to ask +for a fresh B200 FP8 run, although the matched sweep is now completed and reported. +The corrected test retains eager-execution and default-expectation guards while +checking the documented batch-256 no-speedup result and workload-specific +batch-4,096 crossover. Its focused TE regression set passes 23 tests with four +hardware skips. The remaining hosted post-suite checks also pass locally: five +shell entrypoints, zero silent-fallback findings, and 936 benchmark files with +zero contract errors or warnings. The failed hosted receipt remains retained. + Full collective NCU replay still needs a working tool/runtime combination; the bounded selected-kernel captures and matched Nsys traces provide the usable profiling paths on this stack. DDP and pipeline need a repeatable measured gain