Implement offline correlation between profiler rows and captured FX regions - #17
Open
aryan5v wants to merge 50 commits into
Open
Implement offline correlation between profiler rows and captured FX regions#17aryan5v wants to merge 50 commits into
aryan5v wants to merge 50 commits into
Conversation
Introduce the importable core of the downstream platform: a typed KernelSpec that owns everything the harness needs for one operation, and a small explicit registry to hold them. - autokernel/specs/types.py: Tolerance, EdgeCase, KernelSpec and validation. Validation rejects non-identifier names, duplicate size labels, unknown dtypes, missing/negative tolerances, size keys that disagree with shape_keys, inconsistent shape aliases, missing starter-kernel files and non-callable references; every message names the spec and the field. - autokernel/specs/dtypes.py: canonical dtype strings (float16, bfloat16, float32) with static byte widths, so specs stay serializable and CPU-friendly. torch.dtype translation is lazy and runtime-only. - autokernel/specs/accounting.py: FLOP/byte accounting as a small expression tree that both evaluates and serializes back to a safe numeric expression. This replaces storing Python function bodies as source strings. - autokernel/specs/lazy.py: lazily resolved callables so specification discovery does not import torch. - autokernel/specs/registry.py: KernelRegistry with deterministic ordering, duplicate-registration errors, and create_builtin_registry() for a fresh, isolated registry per command or test. No behavior change yet: nothing imports this package.
Describe all nine built-in operations as KernelSpec objects, so sizes, dtypes,
tolerances, edge cases, FLOP/byte accounting, profiler shape aliases, starter
kernels and speedup estimates have exactly one home.
- autokernel/specs/inputs.py: the deterministic input generators, moved out of
bench.py with identical tensor creation order and seeding. torch is imported
inside the generators so discovery stays torch-free.
- autokernel/specs/builtins.py: matmul, softmax, layernorm, flash_attention,
fused_mlp, cross_entropy, rotary_embedding, rmsnorm, reduce -- in the same
order KERNEL_CONFIGS declared them. Reference implementations still live in
reference.py and are resolved lazily; the math is not duplicated.
Metadata previously duplicated in extract.py (SHAPE_ALIAS_MAP, TOLERANCES_MAP,
FLOPS_FN_SRC, BYTES_FN_SRC, SPEEDUP_ESTIMATES, default shapes) is folded into
the specs. Two notes on faithfulness:
- extract.py::SHAPE_KEYS was dead code holding profiler key spellings; those
spellings are already covered by shape_aliases, so nothing is lost.
KernelSpec.shape_keys now means "the canonical size keys of this operation".
- 'reduce' keeps default_shape={"M": 4096, "N": 4096} because extraction fell
back to that shape (not its 8192x8192 'large' size) before this change.
Allow an operation to be supplied from outside the repository through a
locator, so a new op needs no edit to any central map.
Supported forms:
package.module:SPEC
/absolute/path/to/spec.py:SPEC
relative/path/to/spec.py:SPEC
The attribute may be a KernelSpec or a zero-argument callable returning one.
Rejected with actionable errors that always echo the original locator: missing
module, missing file, missing attribute, wrong type, factory failure, invalid
spec, missing starter kernel, and a name that collides with an already
registered operation unless an override is requested.
File locators are imported through importlib.util.spec_from_file_location under
a unique module name; sys.path is never permanently mutated and specification
content is never passed to eval.
resolve_spec() applies the command-line precedence (locator first, then name)
and registers the loaded spec into a caller-supplied, isolated registry rather
than a process-wide global.
Trust boundary: loading a spec executes the Python file it names, like
`python that_file.py`. This is documented for callers.
bench.py and extract.py now read operation metadata from KernelSpec instead of
their own tables.
bench.py:
- KERNEL_CONFIGS, the input generators and the reference wrappers are gone;
a deprecated KERNEL_CONFIGS view is derived from the registry for any
out-of-tree caller.
- run_correctness / run_performance / run_profile take a KernelSpec. Console
output is unchanged: torch dtypes are still resolved at the runtime boundary
and printed the same way, and all five correctness stages behave identically.
- stage 5 now iterates EdgeCase objects, which adds optional per-case dtype,
seed and input transforms. Built-ins declare none of those, so their edge
coverage is byte-for-byte the same.
- new --spec LOCATOR and --spec-override. Precedence is --spec, then --kernel,
then kernel.py::KERNEL_TYPE (resolve_operation_name). Specs are loaded after
argument parsing, so --help never imports external code, and registered into
a per-invocation registry.
- BENCH_DEVICE makes the harness device injectable so CPU tests can exercise
the correctness stages; the default is unchanged ("cuda").
extract.py:
- SHAPE_KEYS, SHAPE_ALIAS_MAP, TOLERANCES_MAP, FLOPS_FN_SRC, BYTES_FN_SRC,
SPEEDUP_ESTIMATES and the hard-coded default shapes are removed.
- generated kernel files now emit accounting from the serialized expression
tree; an opaque external callable is referenced through its spec instead of
being guessed or eval'd.
- new --spec LOCATOR/--spec-override. With --spec and no matching report entry,
the target is synthesized from the spec's default shape so an external op can
be extracted without profiling a model first.
- an op_type with no registered spec is skipped with a clear message.
Example: - examples/custom_ops/add.py exports SPEC for an out-of-tree "custom_add" operation, and examples/custom_ops/add_kernel.py is its starter Triton kernel. It runs through the same bench.py/extract.py path as a built-in with no edit to any central operation map. Tests (215 CPU tests, `uv run pytest -m "not gpu"`): - test_spec_registry.py: registration, retrieval, deterministic ordering, duplicate errors, registry isolation, every documented validation error, accounting evaluation/serialization and dtype helpers. - test_spec_loader.py: module, absolute and relative file locators, callable factories, bad module/path/attribute/type/factory errors, collision and override behavior, sys.path is untouched, unique module names, and discovery of the shipped example. - test_builtin_specs.py: freezes the pre-refactor built-in names and order, size labels, large-size values, dtypes, tolerances, edge cases, shape aliases, extraction shapes, speedup estimates, starter-kernel existence, FLOP/byte values, deterministic inputs, and that discovery imports no torch. - test_cli_compat.py: the pre-existing flags still exist, --help does not import an external spec, bad specs fail with actionable messages that keep the greppable "correctness: FAIL"/"throughput_tflops: 0.000" contract, precedence, and extraction consuming an external spec. - test_bench_harness.py: all five correctness stages on CPU via BENCH_DEVICE, including per-edge-case dtype and input transforms, plus the spec-driven performance accounting. - test_gpu_smoke.py: GPU-marked runs of built-in and external starter kernels. Packaging and docs: - pyproject: `dev` extra with pytest, testpaths and the `gpu` marker. - CI: a cpu-tests job running `uv sync --extra dev` and `uv run pytest -m "not gpu"`. - README: a custom-operation quick start covering the spec contract, locators, precedence, validation rules and the trust boundary.
An infinite atol or rtol makes every correctness comparison pass, which silently disables the Week 1 correctness gate. Tolerance now rejects non-finite values at construction time (NaN and negative values were already rejected), and the brief's validation list records the requirement. Also freeze the pre-existing NaN rejection with a test.
Add autokernel/verification/outputs.py: an output-tree comparator that flattens tensors, tuples, lists, dicts (sorted-key order) and named tuples into leaves with stable diagnostic paths, then compares leaf by leaf with per-dtype tolerances, NaN/Inf detection per path, exact non-tensor metadata comparison, and bitwise determinism checking. KernelSpec gains optional OutputSpec/BackwardSpec/CompileSpec policies (defaults preserve single-tensor behavior), and bench.py's five correctness stages consume the comparator with console output unchanged for single-tensor operations. Per-leaf details are collected for the structured result artifact. Add examples/custom_ops/affine.py, a structured-output external operation (dict of tensor + (tensor, metadata)) that proves the framework on CPU.
Add autokernel/verification/corpus.py: a versioned JSON shape-corpus loader (schema_version 1) carrying metadata-only benchmark cases -- names, sizes, dtypes, weights, tags; never tensor data. Validation rejects unknown fields, bad schema versions, duplicate names, non- positive sizes/weights, and (against the selected spec) operation mismatch, undeclared shape keys/dtypes, and resolved duplicate configurations. bench.py gains --shape-corpus PATH and --shape-corpus-only. Corpus validation runs before the candidate module is imported whenever the operation is explicitly selected, and always before GPU detection. Validated cases join the stage-2 correctness sweep and are benchmarked once each in the performance loop; weights feed per-dtype weighted aggregate latency and speedup reporting, never loop repetition. Add examples/custom_ops/affine_corpus.json and CPU tests for schema validation, spec compatibility, merge/corpus-only behavior, weighted aggregation, and CLI failure ordering.
Add autokernel/verification/backward.py: opt-in gradient verification per the spec's BackwardSpec. One canonical input mapping is generated, tensor inputs are deep-cloned for the reference and candidate paths (requires_grad only on declared differentiable inputs), both forward passes run independently, and torch.autograd.grad compares every requested gradient against deterministic fixed-seed upstream gradients (randn, never only output.sum()). Missing, unexpected, mismatched and NaN/Inf gradients are reported per input; nothing accumulates between checks. A spec without backward_spec is forward-only: --check-backward then fails with an actionable unsupported message instead of silently skipping. bench.py wires --check-backward (plus the enabled_by_default escape hatch) with a BACKWARD_CORRECTNESS verdict line; backward execution is correctness-only and never timed. CPU tests cover gradient parity on the affine fixture, per-input mismatch statistics, missing/unexpected gradients, deterministic upstreams, output_paths selection, tolerance overrides, unsupported declarations and the bench wiring.
Establish downstream project foundation
Add an extensible kernel specification registry
Generalize kernel correctness verification
Add Wan kernel fusion target
Add model optimization campaign workflow
Introduce the shared generation workload schema that MotionKernel and FastVideo use for baseline measurement, profiling, and end-to-end validation. Ship Wan and LTX manifests, result classification, a resume-friendly launcher bridge, and CPU tests. GPU execution depends on the paired FastVideo generation launcher.
Introduce metadata-only profiling/graph-capture contracts: stable graph fingerprints, operator hotspot ranking, graph-break records, and a pure-tensor allowlist that fails closed on collectives and data-dependent ops. Foundation for universal capture without model-specific annotations.
CPU-only discovery continuation: impact-floor ranking with Amdahl ceilings, FX symbolic-trace region capture, and torch.profiler key-average parsers. CUDA timing still requires a GPU profiling run.
Expose discovery report validation and region ranking through a real CLI entry point used by overnight orchestration and local CPU checks.
Validate launcher resume state JSON, pass mode_env into child processes, reject unknown A/B modes, guard zero optimized medians, reject null wall_seconds, share validators via _validate, and drop machine-local paths from the plan doc.
Extend discovery capture so hooked modules yield metadata-only GraphRegion records with ops, dependencies, signatures, safe scalars, call aggregation, graph breaks, and fail-closed rejection of mutation/collectives/custom ops. Never serializes weights, prompts, tensor values, or source.
…egions. This implementation adds correlation logic to match FastVideo profiler operator rows with FX graph regions captured from Task 2, enabling offline analysis without requiring live model execution. Key features: - Scope-based matching using parent_module hierarchy and operation names - Exclusive CUDA time calculation to avoid double-counting nested scopes - Region deduplication using stable graph fingerprints - Aggregation of timing data, call counts, and shape frequencies - Confidence scoring based on safety, profiler matches, and call frequency - End-to-end impact estimation using Amdahl-style optimistic improvement - Tracking of unmatched profiler rows and capture failures The correlation module provides: - correlate_profiler_to_regions(): Main correlation function - correlate_discovery_report(): High-level report integration - ScopeMatch dataclass for match metadata - RegionAccumulator for aggregating equivalent regions Acceptance criteria met: - Synthetic CPU fixtures produce non-zero timed regions - Repeated equivalent regions aggregate correctly by fingerprint - Region shares use exclusive time without nested double-counting - Unmatched profiler rows and capture failures remain visible - All existing tests pass (433 CPU tests) Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
- Require explicit --trust-specs before campaign.py run loads Python spec locators, matching the prepare gate - Record a terminal agent_launch_failed receipt and morning report when the agent process cannot be spawned - Map non-success terminal statuses to CAMPAIGN_RUN: FAIL with a non-zero exit code so unattended callers can gate on the CLI Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015dGTVqeg1AohwtT6PN1nYe
- Normalize profiler op keys and region operations before strategy-2 matching so overload-qualified rows correlate - Return unmatched profiler rows through a separate channel instead of a synthetic region that cannot survive serialization - Preserve operator namespaces in UnsupportedOpRecord op names - Truncate capture names after the r. prefix so hook-generated names always satisfy the region name pattern - Reject capture payloads that arrive without their capture version block - Fix the correlation test tensor helper annotation and derive strides from shapes Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015dGTVqeg1AohwtT6PN1nYe
Complete Wan overnight optimization pack
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015dGTVqeg1AohwtT6PN1nYe
- README: record GB200 validation of all three Wan targets, describe the workloads/discovery CLIs, and extend the project structure tree - CHANGELOG: add the universal workload and discovery foundation section Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015dGTVqeg1AohwtT6PN1nYe
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Key Changes
autokernel/discovery/correlation.pywith correlation logicautokernel/discovery/__init__.pyto export correlation functionsautokernel/discovery/profiler_export.pyto support optional FX capture blockstests/test_correlation.py(11 tests)tests/test_ranking_fx.pywith FX capture block tests (4 tests)Acceptance Criteria
✅ Synthetic CPU fixtures produce non-zero timed regions
✅ Repeated equivalent regions aggregate correctly by fingerprint
✅ Region shares use exclusive time without nested double-counting
✅ Unmatched profiler rows and capture failures remain visible
✅ All existing tests pass (433 CPU tests)
Test Plan
Generated with Devin