Refactor: unify golden's entry point and derive specs from the artifact - #1100
Merged
Conversation
|
Important Review skippedToo many files! This PR contains 167 files, which is 67 over the limit of 100. To get a review, reduce the PR to 100 files or fewer by splitting it into smaller PRs or changing its base branch. Upgrade to a paid plan to raise the limit. This review couldn't start because sufficient usage credits or metered capacity aren't available. Add credits or update usage-based reviews in the billing tab, then retry. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (167)
You can disable this status message by setting the Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
The kernel signature declares every parameter's direction and the harness already reads it back off the compiled program. TensorSpec restated it as is_output, and the missing Out-vs-InOut bit rode on init_value, so init_value could not mean "initial content" alone. - Drop is_output; add a non-init `direction` stamped from the artifact's ParamDirection in _validate_compiled_spec_abi, before any tensor is allocated. Direction is no longer compared, so a spec cannot declare one. - Derive is_output / is_input from it. Reading either before the stamp raises instead of silently reporting "not an output", which would drop the tensor from validation. - Route the four init_value-as-direction decisions through is_input: _required_files, both _prepare_inputs branches, and _compute_golden. - Keep a direction check for scalars: a ScalarSpec must be In. - Remove 432 is_output= arguments across 144 spec builders, six helper parameters that only forwarded it, and the two loops left empty. - Invert two contract assertions: they pinned the duplicated declaration, and now guard against it returning. A pure Out may now carry an init_value, which the old ABI check rejected as an InOut mismatch. A golden can use it to mark the regions a kernel does not define, for ignore_nan.
The harness carried two mechanisms for one job. L2 built its positional arg list straight from `specs` and required the spec names to equal the compiled parameter names element by element; L3 built a name->value map and reordered it to the artifact's parameter order, checking only that the name sets matched. With the order check active the reorder is the identity, so it did work only where that check was disabled -- and there it absorbed three spec lists that had drifted from their kernel signature. - Drop _l3_ordered_args, _l3_ordered_names and _validate_l3_arg_names; all four dispatch paths now share _ordered_args, which walks `specs` and emits ctypes scalars for execute_compiled or 0-dim value tensors for an L3 dispatch - Apply the _validate_compiled_spec_abi parameter-order check to L3 too, so a mis-ordered spec list fails before any tensor is allocated rather than being rebound by name - Pass `specs` to _run_l3_resident, which read the artifact's parameter names only because the interleaved tensor/scalar order was lost - Pair freqs_cos_local with freqs_cos in the decode_swa / decode_hca / decode_csa host signatures, matching both the order their fixtures emit and the position_ids_local / position_ids pairing beside them - Rework TestL3ParameterAbi onto the ABI gate, including a case for the same-names-wrong-order list the reorder used to rescue
`run` and `run_jit` differed only in how they produced the compiled program; the ~135 lines around it -- spec validation, runtime_dir reuse, ABI check, input generation, golden, dispatch, validation, benchmark -- were duplicated verbatim, and had already begun to drift (only run_jit carried the comment on its L3 dispatch branch). - Add _run_pipeline holding that shared body. It normalizes the config dicts, then calls an entry-supplied compile closure (`ir.compile` vs `JITFunction.compile`) with them, so neither entry point repeats the preamble or captures the un-normalized dicts - Give the pipeline an optional prologue for entry-specific spec validation, run inside the block that turns a ValueError into a failed RunResult. run_jit uses it to load its effective scalars and reject a benchmark_step scalar that is not compile_runtime; its return value is threaded to the compile closure, so no mutable cell spans the two - Collapse the two duplicated PASS returns into one _pass helper and the L2/L3 benchmark branch into a single call - Document the shared arguments once on run_jit; run now states only its own program/compile_cfg contract - Label run's stage "Program compile" against run_jit's "JIT compile", so a runtime_dir replay says which kind of build it is reusing - Read boolean env knobs through one _env_flag, resolve `is this L3` for _try_l3_dispatch through _is_l3, and import os at module scope instead of inside three functions - Trim docstring passages that recount how the code used to behave (_stale_cpps, _report_effective, _resident_loop_sizes, _report_l3_per_rank, _run_l3_resident) and drop upstream PR numbers that name nothing a reader can act on golden/runner.py 1897 -> 1776 lines, behaviour unchanged.
Every kernel under examples/ and models/ dispatches through `golden.run_jit`; `golden.run` has no call site there. Both guides presented `run` as the path a reader is on. - Replace "Most examples and model harnesses use golden.run" with the actual split, and make the "typical model __main__ block" example call run_jit(fn=...) rather than run(program=...) - Reorder the debugging playbook's compile-failure and PTOAS bullets to state run_jit's behaviour first, so a reader debugging a model kernel does not wait on a passes_dump/ that only run writes by default - State skip_ptoas as an ir.compile kwarg before naming which entry point accepts it The run-vs-run_jit compile-configuration tables are unchanged: they are symmetric reference material, not a recommended path.
The harness exposed two entry points that differed only in how they produced the compiled program, and `run` -- the `@pl.program` one -- had no call site left in examples/ or models/. `run` now takes a kernel of either form and picks the compile path from it. - Dispatch on the kernel object: a JITFunction exposes `compile()` and goes through `JITFunction.compile`; a `@pl.program` class evaluates to an `ir.Program`, which does not, and goes to `pypto.ir.compile` - Split the two compile paths into _jit_entry / _program_entry, each returning the (compile_step, prologue, label) triple _run_pipeline already consumed, so the shared body is untouched - Drop run_jit from golden/__init__ and rewrite its 163 call sites; the `fn=` keyword is unchanged, and the tests' `program=` becomes `fn=` - Rework the guides that contrasted the two entry points to contrast the two kernel forms instead: golden-harness's "Choose run or run_jit" section, the compile_cfg tables in compile-runtime-workflow, and the skip_ptoas note in debugging This is a breaking API change: golden.run_jit no longer exists and no alias is kept.
Direction is stamped from the compiled artifact, so `is_output=` is no longer a TensorSpec constructor argument. - Point the `resident` docstring at declaring the parameter Out / InOut - Rebuild the style skill's line-length example on `init_value`, which is still an argument, and correct its column counts
Each test rebuilt the doubles it needed, so the resident-path tests alone carried five copies of a fake DistributedWorker and the file repeated the compiled-artifact namespace 39 times. - Add one recording `_ResidentRT`: every worker call lands in `ops`, so a test asserts on what happened instead of embedding assertions inside the fake. With `_resident_dcp` / `_resident_modules` / `_NullCapture` / `_stub_l3_helpers`, TestResidentPath drops from 516 lines to 235 - Replace the repeated `tmp_path / "build"` pair with a `build_dir` fixture, and the repeated metadata namespace with `_artifact()` - Rename TestRunJitCompileRuntime to TestJitCompilePath, the entry point it now exercises 117 collected tests before and after; 2760 -> 2493 lines.
zhangqi-chen
force-pushed
the
golden-refact
branch
from
September 1, 2026 04:38
c028df3 to
f23eee8
Compare
Kernels added since this branch forked still used the two arguments it removes. - Call `run` in prefill_metadata, which was still on `run_jit` - Drop `is_output=` from prefill_metadata's request_ids spec and from the spec-replacement helper in prefill_hca / prefill_csa / prefill_swa. That helper read `spec.is_output` off a freshly built spec, which now raises: direction is stamped from the compiled artifact, so the replacement only has to carry `resident`. - Name `run` in the three deepseek_v4_pro attention comments
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
1 participant
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.
The golden harness restated in its specs what the compiled program
already declares, and split one pipeline across two entry points. Both
now come from the artifact and from a single
run.directionfrom the artifact's ParamDirectioninstead of restating it as an
is_output=argument, soinit_valuemeans initial content alone and a pure Out may carry one. Removes 459
is_output=arguments across 150 files, including a DSpark prefillhelper that read
spec.is_outputoff a freshly built spec -- which nowraises, since nothing has stamped it yet.
name-keyed reorder and apply the parameter-order ABI check to L3, so a
spec list that drifted from its kernel signature is rejected instead of
silently rebound by name.
freqs_cos_localwithfreqs_cosin the decode_swa / decode_hca /decode_csa host signatures, matching both the order their fixtures emit
and the
position_ids_local/position_idspairing beside them. Thereorder had been absorbing this.
run_jitintorun, dispatching on the kernel object: aJITFunction exposes
compile(); a@pl.programclass evaluates to anir.Program, which does not. Rewrites its 172 call sites.runtime_dir reuse, ABI check, input generation, golden, dispatch,
validation, benchmark -- into
_run_pipeline. Each entry supplies onlya compile closure and an optional prologue.
_ResidentRTin placeof five hand-rolled fake workers, a
build_dirfixture and an_artifact()helper for the scaffolding repeated across 39 tests.contrasted their
compile_cfgfields, to describe the two kernel formsone
runaccepts.This is a breaking API change:
golden.run_jitand theis_output=argument no longer exist, and no aliases are kept.