diff --git a/.claude/skills/fmt-coding-style/SKILL.md b/.claude/skills/fmt-coding-style/SKILL.md index 436c3dec2..e72e3a582 100644 --- a/.claude/skills/fmt-coding-style/SKILL.md +++ b/.claude/skills/fmt-coding-style/SKILL.md @@ -60,12 +60,12 @@ the long line — when a statement cannot be split into ops (a single 120 columns instead of trailing one keyword onto a continuation line. ```python -# before — 96 + 15 columns split across two lines +# before — 92 + 23 columns split across two lines next_hidden_spec = TensorSpec("next_pre_hc_hidden", [N_RANKS, T, HC_MULT, D], torch.float32, - is_output=True) + init_value=torch.zeros) -# after — 112 columns, one line -next_hidden_spec = TensorSpec("next_pre_hc_hidden", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True) +# after — 116 columns, one line +next_hidden_spec = TensorSpec("next_pre_hc_hidden", [N_RANKS, T, HC_MULT, D], torch.float32, init_value=torch.zeros) ``` Past ~120, go back to Rule 1: split the op, or group the arguments as below. @@ -163,7 +163,7 @@ Structures that stay one-item-per-line: - the **kernel signature** — one annotated parameter per line, already the convention; -- the **harness call**, `run_jit(...)` / `run(...)` — leave it exactly as it is. +- the **harness call**, `run(...)` — leave it exactly as it is. One option per line is what makes it easy to toggle `compile_cfg` / `runtime_cfg` / `rtol` / `compare_fn` entries while testing; - list literals whose elements are themselves full-width expressions, such as diff --git a/.claude/skills/test-with-golden/SKILL.md b/.claude/skills/test-with-golden/SKILL.md index ebf6c1b30..53bb5ca9b 100644 --- a/.claude/skills/test-with-golden/SKILL.md +++ b/.claude/skills/test-with-golden/SKILL.md @@ -15,7 +15,7 @@ snapshot-layout, invalidation, and CLI-wiring reference. whose specs, inputs, golden logic, and intended numerics remain unchanged. Do not use it for a precision investigation that needs fixture, seed, dtype, rounding, quantization, or reference changes. -2. Inspect the target's `__main__` and its `run` or `run_jit` call. Check +2. Inspect the target's `__main__` and its `run` call. Check whether both `--save-data` and `--golden-data` are already exposed. 3. If either flag is missing, add only the minimal argparse option and keyword forwarding shown in the canonical guide. Do not change the default run: diff --git a/docs/debug-and-tune/debugging.md b/docs/debug-and-tune/debugging.md index e30a39117..717987e0d 100644 --- a/docs/debug-and-tune/debugging.md +++ b/docs/debug-and-tune/debugging.md @@ -7,7 +7,7 @@ does), [performance-tuning.md](performance-tuning.md) (perf), and [precision-tuning.md](precision-tuning.md) (numerical fidelity — cast modes, dtype alignment, the `error_distribution` sweep). -The harness exposes most of these as both a `run` / `run_jit` kwarg and a +The harness exposes most of these as both a `run` kwarg and a CLI flag; a typical model `__main__` wires them up like: ```python @@ -18,7 +18,7 @@ parser.add_argument( ) parser.add_argument("--enable-dep-gen", action="store_true") ... -result = run_jit( +result = run( fn=indexer_test, specs=build_tensor_specs(...), golden_fn=golden_indexer, @@ -43,18 +43,17 @@ location. Most compile failures are fixed at the cited site without any further tooling; read the message before reaching for the heavier mechanisms below. -- **Compile failure** — scripts using `golden.run` call `ir.compile` - directly, whose default `dump_passes=True` writes per-pass IR under - `build_output/<...>/passes_dump/`. `golden.run_jit` constructs a - `RunConfig`, whose default is `dump_passes=False`; pass - `compile_cfg=dict(dump_passes=True)` when you need the same files. Diff the - last clean pass against the first failing one to see which pass rejected - the IR. `report/` holds scheduling diagnostics. -- **PTOAS failure** — the error quotes the `.pto` op. With `golden.run`, +- **Compile failure** — a `@pl.jit` kernel compiles through a `RunConfig`, + whose default is `dump_passes=False`; pass `compile_cfg=dict(dump_passes=True)` + to write per-pass IR under `build_output/<...>/passes_dump/`. A `@pl.program` + kernel goes through `ir.compile`, whose default `dump_passes=True` writes + those files already. Diff the last clean pass against the first failing one to + see which pass rejected the IR. `report/` holds scheduling diagnostics. +- **PTOAS failure** — the error quotes the `.pto` op. `compile_cfg=dict(skip_ptoas=True)` keeps the raw `.pto` MLIR and stops before the C++ wrapper, isolating whether the regression is in PyPTO's - IR-to-MLIR path or in PTOAS. `skip_ptoas` is not a `RunConfig` field and - therefore is not accepted by `golden.run_jit`. + IR-to-MLIR path or in PTOAS. It is an `ir.compile` kwarg, not a `RunConfig` + field, so it works on a `@pl.program` kernel only. - **Runtime crash** — rerun on the matching simulator (`-p a2a3sim` / `a5sim`); it gives more diagnostic output than the device backend and reproduces most lowering bugs. @@ -183,7 +182,7 @@ pl.dump_tag(h_tile_i8) # capture this one tensor under partial dump ``` ```python -run_jit(..., runtime_cfg=dict(platform=..., enable_dump_args=1)) +run(..., runtime_cfg=dict(platform=..., enable_dump_args=1)) ``` `pl.dump_tag` works on plain function args and on internal @@ -202,7 +201,7 @@ python -m simpler_setup.tools.dump_viewer /dfx_outputs/args_dump`. This section is the dump *mechanism*. For the end-to-end @@ -240,7 +239,7 @@ round-trip, which drops the read-dep and lets the downstream task race). | Symptom | Tool | Kwarg / flag | |---------|------|--------------| -| Compile / PTOAS error | `passes_dump/`; `skip_ptoas` for `golden.run` only | `compile_cfg=dict(dump_passes=True)`; with `run`, `compile_cfg=dict(skip_ptoas=True)` | +| Compile / PTOAS error | `passes_dump/`; `skip_ptoas` on a `@pl.program` kernel only | `compile_cfg=dict(dump_passes=True)`; `compile_cfg=dict(skip_ptoas=True)` | | Need to reproduce on the same inputs | golden-data replay (§2) | `golden_data=` / `--golden-data` | | Iterating on generated `.cpp` / `.pto` | runtime-dir reuse (§3) | `runtime_dir=` / `--runtime-dir` | | Run hangs / deadlocks (§4) | device log | `runtime_cfg["log_level"]="v0"` + `ASCEND_PROCESS_LOG_PATH` | diff --git a/docs/debug-and-tune/performance-tuning.md b/docs/debug-and-tune/performance-tuning.md index 50dfdaec6..5b33d1767 100644 --- a/docs/debug-and-tune/performance-tuning.md +++ b/docs/debug-and-tune/performance-tuning.md @@ -15,7 +15,7 @@ compute core. ## Measuring — the benchmark loop (`PYPTO_BENCH`) Tuning needs a number before and after. Set `PYPTO_BENCH=1` and every -`run` / `run_jit` call in the process times the kernel on device after its +`run` call in the process times the kernel on device after its correctness dispatch — no `--benchmark` flag, no edit to the model file: ```bash @@ -366,7 +366,7 @@ python models/deepseek_v4_flash_mtp/decode_sparse_attn.py -p a2a3 -d 0 --enable- ``` Not every kernel exposes `--enable-pmu`; a kernel that does not can still be -captured by passing `runtime_cfg={"enable_pmu": 2}` to its `run` / `run_jit` +captured by passing `runtime_cfg={"enable_pmu": 2}` to its `run` call (the harness bundles it into the runtime's DFX options). For a per-kernel intra-core swimlane, use diff --git a/docs/debug-and-tune/precision-tuning.md b/docs/debug-and-tune/precision-tuning.md index af2dc09f9..de3b92222 100644 --- a/docs/debug-and-tune/precision-tuning.md +++ b/docs/debug-and-tune/precision-tuning.md @@ -173,7 +173,7 @@ comparator from the golden harness to see the precision distribution instead: ```python from golden import error_distribution -run_jit( +run( ..., compare_fn={ "x_next": error_distribution(), # measure, never fails diff --git a/docs/debug-and-tune/ring-heap-and-scope-stats.md b/docs/debug-and-tune/ring-heap-and-scope-stats.md index 73ceb91f5..1a8f4f105 100644 --- a/docs/debug-and-tune/ring-heap-and-scope-stats.md +++ b/docs/debug-and-tune/ring-heap-and-scope-stats.md @@ -154,7 +154,7 @@ prefill — size the rings explicitly. Each knob takes either a scalar this ring at the compile default". Precedence is just two tiers now: `RunConfig` field > compile default. -**This works on the L3 (distributed) path only.** golden's `run_jit` forwards +**This works on the L3 (distributed) path only.** golden's `run` forwards whatever `runtime_cfg` keys are `RunConfig` fields to the per-rank dispatch, which builds a `CallConfig` and transcribes the ring sizes into `runtime_env`: @@ -175,7 +175,7 @@ Live examples: [`models/deepseek_v4_pro/prefill_fwd.py`](../../models/deepseek_v An L2 (single-chip) entry cannot size its rings today, whichever way it is dispatched: -- Through golden's `run_jit`, the single-chip path calls +- Through golden's `run`, the single-chip path calls `execute_compiled(work_dir, args, ...)`, which has no ring parameters — a `runtime_cfg["ring_heap"]` raises `TypeError`. - Through `CompiledProgram.__call__(*args, config=RunConfig(ring_heap=...))`, diff --git a/docs/get-started/first-kernel.md b/docs/get-started/first-kernel.md index fed1b6aaf..784b6e91f 100644 --- a/docs/get-started/first-kernel.md +++ b/docs/get-started/first-kernel.md @@ -13,7 +13,7 @@ The file contains four parts: 1. a module-level `@pl.jit` kernel; 2. `TensorSpec` and `ScalarSpec` inputs for the Golden Harness; 3. a PyTorch golden function; -4. a CLI that calls `golden.run_jit` and exits non-zero on failure. +4. a CLI that calls `golden.run` and exits non-zero on failure. ## Run on the A2/A3 simulator diff --git a/docs/get-started/installation.md b/docs/get-started/installation.md index 97bdfb160..bffe0e6b6 100644 --- a/docs/get-started/installation.md +++ b/docs/get-started/installation.md @@ -147,7 +147,7 @@ re-cloned rather than reused in place. ## Verify ```bash -python -c "import pypto, torch; from golden import run, run_jit; print(torch.__version__)" +python -c "import pypto, torch; from golden import run; print(torch.__version__)" python examples/beginner/hello_world.py -p a2a3sim ``` diff --git a/docs/models/deepseek_v4_pro/index.md b/docs/models/deepseek_v4_pro/index.md index 5d0f8ca28..d96605ebf 100644 --- a/docs/models/deepseek_v4_pro/index.md +++ b/docs/models/deepseek_v4_pro/index.md @@ -196,7 +196,7 @@ python models/deepseek_v4_pro/synthetic_token_loop.py --variant flash \ The full prefill and decode programs carry runtime `num_tokens` and `moe_epoch_base` scalars in their compiled ABI. Their `ScalarSpec`s use -`compile_runtime=True`, so `run_jit` passes `pl.RUNTIME` during +`compile_runtime=True`, so `run` passes `pl.RUNTIME` during signature-driven compilation instead of folding the initial values into generated task arguments. `num_tokens` follows the real prompt/decode row count, while callers advance the epoch scalar by diff --git a/docs/pypto-coding/distributed-programming.md b/docs/pypto-coding/distributed-programming.md index 5d229c126..ec55620e4 100644 --- a/docs/pypto-coding/distributed-programming.md +++ b/docs/pypto-coding/distributed-programming.md @@ -342,7 +342,7 @@ The harness compiles an L3 program when `compile_cfg` carries a ```python from pypto.ir.distributed_compiled_program import DistributedConfig -result = run_jit( +result = run( fn=l3_moe, # the @pl.jit.host driver specs=build_tensor_specs(), golden_fn=golden_moe, diff --git a/docs/pypto-coding/pypto-coding-style.md b/docs/pypto-coding/pypto-coding-style.md index 3b8e85414..1857c4b3a 100644 --- a/docs/pypto-coding/pypto-coding-style.md +++ b/docs/pypto-coding/pypto-coding-style.md @@ -91,13 +91,14 @@ Every tensor the golden test compares **must** declare an explicit direction on the **orchestration entry** — the `@pl.jit` entry, its `@pl.jit.host` driver, or the `@pl.function(type=Opaque)` / `Orchestration` method. A plain `pl.Tensor` is treated as `In`: the runtime skips its device→host copy-back, so the tensor -**reads back as all-zeros on the host** and golden silently fails. The direction -is decided by the tensor's `TensorSpec`: - -| `TensorSpec` | meaning | annotation | -|--------------|---------|------------| -| `is_output=True`, no `init_value` | pure output (write-only) | `pl.Out[pl.Tensor[...]]` | -| `is_output=True` **and** `init_value` | inout — read-modify-write (e.g. a paged KV cache the kernel reads history from and appends to; recurrent state) | `pl.InOut[pl.Tensor[...]]` | +**reads back as all-zeros on the host** and golden silently fails. The +annotation is the only place direction is declared — the harness reads it back +off the compiled artifact, so a `TensorSpec` never restates it: + +| annotation | meaning | `TensorSpec` | +|------------|---------|--------------| +| `pl.Out[pl.Tensor[...]]` | pure output (write-only); validated | no `init_value` needed — the host buffer is not uploaded | +| `pl.InOut[pl.Tensor[...]]` | inout — read-modify-write (e.g. a paged KV cache the kernel reads history from and appends to; recurrent state); validated | `init_value` is the uploaded initial state | Annotate the **entry only**. `@pl.jit.inline` sub-kernels keep bare `pl.Tensor`: they are spliced at the call site before SSA conversion, so a parameter is diff --git a/docs/run-and-validate/compile-runtime-workflow.md b/docs/run-and-validate/compile-runtime-workflow.md index a86aaf402..4e5820306 100644 --- a/docs/run-and-validate/compile-runtime-workflow.md +++ b/docs/run-and-validate/compile-runtime-workflow.md @@ -1,11 +1,11 @@ # Compile and Runtime Workflow What usually happens when you run `python .py -p `. -Most examples and model harnesses use `golden.run` for `@pl.program` kernels -or `golden.run_jit` for module-level `@pl.jit` kernels. A few specialized -smoke, artifact-regeneration, and external-runtime drivers call PyPTO's -compile/runtime APIs directly; their `__main__` blocks are the authority for -what a command actually validates. +Examples and model harnesses go through `golden.run`, which takes a kernel of +either form — a module-level `@pl.jit` function or a `@pl.program` class — and +picks its compile path from it. A few specialized smoke, artifact-regeneration, and +external-runtime drivers call PyPTO's compile/runtime APIs directly; their +`__main__` blocks are the authority for what a command actually validates. ## CLI shape @@ -19,8 +19,8 @@ parser.add_argument("--enable-chip-swimlane", action="store_true") args = parser.parse_args() result = run( - program=build_qwen3_decode_program(...), # @pl.program class - specs=build_tensor_specs(...), # ordered TensorSpec / ScalarSpec list + fn=qwen3_decode, # module-level @pl.jit function + specs=build_tensor_specs(...), # TensorSpec / ScalarSpec, in fn's param order golden_fn=golden_qwen3_decode, # PyTorch reference compile_cfg=dict(dump_passes=True), runtime_cfg=dict(platform=args.platform, device_id=args.device, @@ -29,11 +29,10 @@ result = run( ) ``` -A kernel written as a module-level `@pl.jit` function calls **`run_jit`** -instead, passing `fn=` in place of `program=`. Both entry points -share tensor specs, golden computation, runtime dispatch, and validation, but -their `compile_cfg` fields and defaults differ; see -[Compile configuration](#compile-configuration). +A kernel built as a `@pl.program` class passes that program as `fn=` +instead. Both forms share tensor specs, golden computation, runtime dispatch, +and validation; only the compile step and the accepted `compile_cfg` fields +differ, see [Compile configuration](#compile-configuration). | Flag | Purpose | |------|---------| @@ -83,10 +82,10 @@ each phase, so the console log is the authoritative trace of what ran: ### 1. Compile (pypto) -Driven by the **pypto** repo. `run` calls -`pypto.ir.compile(program, backend_type=..., **compile_cfg)` directly. -`run_jit` builds a `pypto.runtime.RunConfig` from `compile_cfg` and calls -`fn.compile(..., config=...)`, which specializes the JIT function before +Driven by the **pypto** repo. For a `@pl.program` kernel, `run` calls +`pypto.ir.compile(program, backend_type=..., **compile_cfg)` directly. For a +`@pl.jit` kernel it builds a `pypto.runtime.RunConfig` from `compile_cfg` and +calls `fn.compile(..., config=...)`, which specializes the JIT function before entering the same IR compiler. Both paths run a **pass pipeline** followed by a **codegen pipeline** and normally write `build_output/_/`. @@ -96,9 +95,9 @@ a **codegen pipeline** and normally write `PassManager.get_strategy(strategy).run_passes(program, ...)` runs an ordered sequence of passes that progressively rewrites the IR. The exact pass list changes often — consult the pypto repo for the current pipeline, -and look at `passes_dump/` when `dump_passes` is enabled. Direct `run` -compilation inherits `ir.compile`'s enabled default; `run_jit` inherits -`RunConfig`'s disabled default. +and look at `passes_dump/` when `dump_passes` is enabled. The direct +`@pl.program` path inherits `ir.compile`'s enabled default; the `@pl.jit` path +inherits `RunConfig`'s disabled default. The end state, regardless of which passes ran, is the same: @@ -157,7 +156,8 @@ build_output/_/ #### Compile configuration -For **`run`**, `compile_cfg` is forwarded to `ir.compile`. Common fields are: +For a **`@pl.program`** kernel, `compile_cfg` is forwarded to `ir.compile`. +Common fields are: | `compile_cfg` field | Purpose | |---|---| @@ -173,8 +173,8 @@ The harness derives `backend_type` and `platform` from `runtime_cfg["platform"]` unless the direct compile configuration already supplies them. -For **`run_jit`**, `compile_cfg` must instead contain fields accepted by -`pypto.runtime.RunConfig`. The JIT layer maps its compile-side fields into +For a **`@pl.jit`** kernel, `compile_cfg` must instead contain fields accepted +by `pypto.runtime.RunConfig`. The JIT layer maps its compile-side fields into `ir.compile`: | `compile_cfg` field | Compiler mapping | @@ -186,9 +186,10 @@ For **`run_jit`**, `compile_cfg` must instead contain fields accepted by | `distributed_config`, `analyze_auto_scopes_for_deps`, `memory_planner` | Forwarded when set. | `output_dir`, `profiling`, `skip_ptoas`, and `verification_level` are not -`RunConfig` field names and therefore cannot be copied unchanged from a -`run` call into `run_jit`. Unknown fields raise while constructing -`RunConfig`; unknown direct-compiler fields raise in `ir.compile`. +`RunConfig` field names, so a `compile_cfg` written for a `@pl.program` kernel +cannot be copied unchanged onto a `@pl.jit` one. Unknown fields raise while +constructing `RunConfig`; unknown direct-compiler fields raise in +`ir.compile`. To stop after compile without touching the device, see `compile_only` under [Skipping phases](#skipping-phases). @@ -198,7 +199,11 @@ To stop after compile without touching the device, see `compile_only` under Each entry of `specs` is a `TensorSpec` (named tensor, shape, dtype, direction) or a `ScalarSpec` (named scalar, dtype, value); see `golden/spec.py`. The list is ordered to match the parameter order of the -top opaque function. For each entry, allocate a torch tensor: +top opaque function — single-chip and distributed alike. The harness compares +the spec names against the compiled artifact's parameters element by element +and fails with `compiled parameter ABI mismatch (parameter order ...)` before +allocating anything; it never rebinds a mis-ordered list by name. For each +entry, allocate a torch tensor: - Pure inputs and inout initial values are filled via `spec.create_tensor()` (`init_value=None` creates zeros; random data requires an explicit factory @@ -358,7 +363,7 @@ false; an uncaught compile/runtime exception is already a nonzero failure. ## Skipping phases -`run` / `run_jit` knobs that short-circuit the pipeline: +`run` knobs that short-circuit the pipeline: | Knob | Effect | |------|--------| diff --git a/docs/run-and-validate/golden-harness.md b/docs/run-and-validate/golden-harness.md index a2713f325..1c8e57ed0 100644 --- a/docs/run-and-validate/golden-harness.md +++ b/docs/run-and-validate/golden-harness.md @@ -5,8 +5,8 @@ correctness-validation path. Its public entry points are exported from [`golden/__init__.py`](../../golden/__init__.py): - `TensorSpec` and `ScalarSpec` describe ordered kernel arguments; -- `run` drives a `@pl.program` program; -- `run_jit` drives a module-level `@pl.jit` function; +- `run` drives a kernel of either form — a module-level `@pl.jit` function or + a `@pl.program` program; - validation helpers provide output-specific comparison policies. ## Describe the arguments @@ -21,7 +21,6 @@ TensorSpec( shape, dtype, init_value=None, - is_output=False, ) ``` @@ -38,9 +37,13 @@ TensorSpec( Random input is therefore explicit: use `init_value=torch.randn` or another random factory. `init_value=None` does not generate random values. -Set `is_output=True` for every tensor that validation must compare. An output -with a non-`None` initializer is an input/output state tensor: its initial -value is supplied to the runtime and its final value is validated. +A spec does not declare a direction. The harness reads each parameter's +`In` / `Out` / `InOut` from the compiled artifact and stamps it onto the spec +before any tensor is allocated, so the kernel signature is the single source of +truth: every `pl.Out` / `pl.InOut` parameter is validated, and a `pl.InOut` +tensor's `init_value` is uploaded as its initial state. A pure `pl.Out` +parameter's host buffer is never uploaded, so an `init_value` there reaches only +the golden reference, not the device. ### ScalarSpec @@ -49,9 +52,9 @@ represents a scalar kernel argument. The harness stores it as a zero-dimensional PyTorch tensor and converts it to the runtime ABI form during dispatch. The name and position must still match the kernel signature. -`run_jit` normally specializes scalar values into the artifact. Set +A `@pl.jit` kernel normally specializes scalar values into the artifact. Set `compile_runtime=True` when dispatches must supply different values to the same -artifact. If any scalar is marked, `run_jit` compiles from the JIT function's +artifact. If any scalar is marked, `run` compiles from the JIT function's fully annotated signature, passes marked scalars as `pl.RUNTIME`, and keeps unmarked scalars specialized to their `value`. Every tensor parameter therefore needs a complete `pl.Tensor[[shape...], dtype]` annotation on this path. @@ -62,7 +65,7 @@ warmup launches, receives `value + i * benchmark_step`. Stepped scalars require resident specs: L2 benchmarks and the non-resident L3 benchmark both reject them, because those paths repeat one argument list per launch instead of providing the persistent-window contract. A -stepped scalar compiled through `run_jit` must also use `compile_runtime=True`; +stepped scalar on a `@pl.jit` kernel must also use `compile_runtime=True`; otherwise the compiler is allowed to fold the initial value into the artifact. Stepped scalars cannot be combined with `runtime_cfg={"enable_chip_swimlane": True}` (nor its pre-rename spelling `enable_l2_swimlane`): that mode may @@ -94,12 +97,12 @@ def golden_hello_world(values): The harness gives the golden function cloned inputs and separate zero-filled pure outputs, so runtime writes cannot mutate the already computed reference. -## Choose run or run_jit +## Pass either kernel form -Use `run_jit` for a module-level `@pl.jit` function: +`run` takes the kernel itself. A module-level `@pl.jit` function: ```python -result = run_jit( +result = run( fn=hello_world, specs=build_specs(), golden_fn=golden_hello_world, @@ -112,11 +115,11 @@ result = run_jit( ) ``` -Use `run` for a built `@pl.program`: +A built `@pl.program`: ```python result = run( - program=build_program(), + fn=build_program(), specs=build_specs(), golden_fn=golden_fn, runtime_cfg={ @@ -126,9 +129,9 @@ result = run( ) ``` -Both entry points perform the same input, golden, runtime, and validation -stages after their respective compile path. The detailed sequence and -configuration groups are documented in +`run` picks the compile path from the kernel it is handed, then performs the +same input, golden, runtime, and validation stages either way. The detailed +sequence and configuration groups are documented in [Compile and Runtime Workflow](compile-runtime-workflow.md). ## Validation @@ -154,7 +157,7 @@ a correctness check. ## Handle RunResult -`run` and `run_jit` return: +`run` returns: ```python RunResult( diff --git a/docs/run-and-validate/index.md b/docs/run-and-validate/index.md index 8acbc74e8..9bc392dec 100644 --- a/docs/run-and-validate/index.md +++ b/docs/run-and-validate/index.md @@ -25,7 +25,7 @@ Use the following references: - [Compile and Runtime Workflow](compile-runtime-workflow.md) explains the generated artifacts and each stage in detail. - [Golden Harness](golden-harness.md) introduces `TensorSpec`, `ScalarSpec`, - `run`, `run_jit`, validation, and `RunResult`. + `run`, validation, and `RunResult`. - [Save and Replay Golden Data](save-and-replay.md) freezes inputs and expected outputs for repeated performance-oriented runs. diff --git a/docs/run-and-validate/save-and-replay.md b/docs/run-and-validate/save-and-replay.md index 47b2d23b9..5f54af2ea 100644 --- a/docs/run-and-validate/save-and-replay.md +++ b/docs/run-and-validate/save-and-replay.md @@ -29,10 +29,10 @@ failure is the explicit goal. ## Harness API -`run` and `run_jit` share two keyword arguments: +`run` takes two keyword arguments for this: ```python -result = run_jit( +result = run( fn=kernel, specs=specs, golden_fn=golden_fn, @@ -66,7 +66,7 @@ parser.add_argument( ) ``` -Then forward `args.save_data` and `args.golden_data` to `run` or `run_jit`. +Then forward `args.save_data` and `args.golden_data` to `run`. Do not assume every existing entry point already exposes both flags; check its `--help` and call site. diff --git a/examples/advanced/allreduce.py b/examples/advanced/allreduce.py index 0bef81610..4d00400a5 100644 --- a/examples/advanced/allreduce.py +++ b/examples/advanced/allreduce.py @@ -109,7 +109,7 @@ def init_inputs(): return [ TensorSpec("inputs", [N_RANKS, 1, SIZE], torch.float32, init_value=init_inputs), - TensorSpec("outputs", [N_RANKS, 1, SIZE], torch.float32, is_output=True), + TensorSpec("outputs", [N_RANKS, 1, SIZE], torch.float32), ] @@ -122,7 +122,7 @@ def golden_allreduce(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser() @@ -136,7 +136,7 @@ def golden_allreduce(tensors): device_ids = [int(d) for d in args.device.split(",")] assert len(device_ids) == N_RANKS, f"need exactly {N_RANKS} devices, got {device_ids}" - result = run_jit( + result = run( fn=l3_allreduce, specs=build_tensor_specs(), golden_fn=golden_allreduce, diff --git a/examples/advanced/gemm_eltwise.py b/examples/advanced/gemm_eltwise.py index 1b4cab9b8..8e6418728 100644 --- a/examples/advanced/gemm_eltwise.py +++ b/examples/advanced/gemm_eltwise.py @@ -58,7 +58,7 @@ def build_tensor_specs( TensorSpec("attn_out", [batch, hidden], torch.bfloat16, init_value=torch.randn), TensorSpec("hidden_states", [batch, hidden], torch.bfloat16, init_value=torch.randn), TensorSpec("wo", [hidden, hidden], torch.bfloat16, init_value=torch.randn), - TensorSpec("resid", [batch, hidden], torch.float32, is_output=True), + TensorSpec("resid", [batch, hidden], torch.float32), ] @@ -71,7 +71,7 @@ def golden_gemm_eltwise(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -80,7 +80,7 @@ def golden_gemm_eltwise(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=gemm_eltwise, specs=build_tensor_specs(), golden_fn=golden_gemm_eltwise, diff --git a/examples/advanced/multi_proj.py b/examples/advanced/multi_proj.py index dab5377d7..d584835e9 100644 --- a/examples/advanced/multi_proj.py +++ b/examples/advanced/multi_proj.py @@ -75,9 +75,9 @@ def init_w(): TensorSpec("wq", [HIDDEN, HIDDEN], torch.bfloat16, init_value=init_w), TensorSpec("wk", [HIDDEN, HIDDEN], torch.bfloat16, init_value=init_w), TensorSpec("wv", [HIDDEN, HIDDEN], torch.bfloat16, init_value=init_w), - TensorSpec("q_out", [BATCH, HIDDEN], torch.float32, is_output=True), - TensorSpec("k_out", [BATCH, HIDDEN], torch.float32, is_output=True), - TensorSpec("v_out", [BATCH, HIDDEN], torch.float32, is_output=True), + TensorSpec("q_out", [BATCH, HIDDEN], torch.float32), + TensorSpec("k_out", [BATCH, HIDDEN], torch.float32), + TensorSpec("v_out", [BATCH, HIDDEN], torch.float32), ] @@ -91,7 +91,7 @@ def golden_qkv_proj(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -100,7 +100,7 @@ def golden_qkv_proj(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=qkv_proj, specs=build_tensor_specs(), golden_fn=golden_qkv_proj, diff --git a/examples/advanced/topk.py b/examples/advanced/topk.py index c72b4fb0b..f16cf5477 100644 --- a/examples/advanced/topk.py +++ b/examples/advanced/topk.py @@ -52,8 +52,8 @@ def build_tensor_specs(): return [ TensorSpec("scores", [ROWS, N], torch.float32, init_value=torch.randn), - TensorSpec("topk_vals", [ROWS, K], torch.float32, is_output=True), - TensorSpec("topk_idx", [ROWS, K], torch.int32, is_output=True), + TensorSpec("topk_vals", [ROWS, K], torch.float32), + TensorSpec("topk_idx", [ROWS, K], torch.int32), ] @@ -68,7 +68,7 @@ def golden_topk(tensors): if __name__ == "__main__": import argparse - from golden import run_jit, topk_pair_compare + from golden import run, topk_pair_compare parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -77,7 +77,7 @@ def golden_topk(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=topk, specs=build_tensor_specs(), golden_fn=golden_topk, diff --git a/examples/beginner/hello_world.py b/examples/beginner/hello_world.py index bc6d7fd9b..1968faf33 100644 --- a/examples/beginner/hello_world.py +++ b/examples/beginner/hello_world.py @@ -47,7 +47,7 @@ def build_specs( return [ TensorSpec("x", [rows, cols], torch.float32, init_value=torch.randn), ScalarSpec("a", torch.float32, a), - TensorSpec("y", [rows, cols], torch.float32, is_output=True), + TensorSpec("y", [rows, cols], torch.float32), ] @@ -57,7 +57,7 @@ def golden_hello_world(values): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -66,7 +66,7 @@ def golden_hello_world(values): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=hello_world, specs=build_specs(), golden_fn=golden_hello_world, diff --git a/examples/beginner/matmul.py b/examples/beginner/matmul.py index 6e2f4bc44..0e4123250 100644 --- a/examples/beginner/matmul.py +++ b/examples/beginner/matmul.py @@ -48,7 +48,7 @@ def build_tensor_specs( return [ TensorSpec("a", [m, k], torch.float32, init_value=torch.randn), TensorSpec("b", [k, n], torch.float32, init_value=torch.randn), - TensorSpec("c", [m, n], torch.float32, is_output=True), + TensorSpec("c", [m, n], torch.float32), ] @@ -58,7 +58,7 @@ def golden_matmul(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -67,7 +67,7 @@ def golden_matmul(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=matmul, specs=build_tensor_specs(), golden_fn=golden_matmul, diff --git a/examples/intermediate/gemm.py b/examples/intermediate/gemm.py index 2abb580f6..d897f3c9b 100644 --- a/examples/intermediate/gemm.py +++ b/examples/intermediate/gemm.py @@ -56,7 +56,7 @@ def build_tensor_specs( return [ TensorSpec("a", [m, k], torch.float32, init_value=torch.randn), TensorSpec("b", [k, n], torch.float32, init_value=torch.randn), - TensorSpec("c", [m, n], torch.float32, is_output=True), + TensorSpec("c", [m, n], torch.float32), ] @@ -66,7 +66,7 @@ def golden_gemm(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -75,7 +75,7 @@ def golden_gemm(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=gemm, specs=build_tensor_specs(), golden_fn=golden_gemm, diff --git a/examples/intermediate/layer_norm.py b/examples/intermediate/layer_norm.py index 462e04fb8..3dcd6152e 100644 --- a/examples/intermediate/layer_norm.py +++ b/examples/intermediate/layer_norm.py @@ -67,7 +67,7 @@ def build_tensor_specs( TensorSpec("x", [rows, hidden], torch.float32, init_value=torch.randn), TensorSpec("gamma", [1, hidden], torch.float32, init_value=torch.randn), TensorSpec("beta", [1, hidden], torch.float32, init_value=torch.randn), - TensorSpec("y", [rows, hidden], torch.float32, is_output=True), + TensorSpec("y", [rows, hidden], torch.float32), ] @@ -84,7 +84,7 @@ def golden_layer_norm(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -93,7 +93,7 @@ def golden_layer_norm(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=layer_norm, specs=build_tensor_specs(), golden_fn=golden_layer_norm, diff --git a/examples/intermediate/rms_norm.py b/examples/intermediate/rms_norm.py index d3b1e73c1..716a94d4a 100644 --- a/examples/intermediate/rms_norm.py +++ b/examples/intermediate/rms_norm.py @@ -68,7 +68,7 @@ def build_tensor_specs( return [ TensorSpec("x", [rows, hidden], torch.float32, init_value=torch.randn), TensorSpec("gamma", [1, hidden], torch.float32, init_value=torch.randn), - TensorSpec("y", [rows, hidden], torch.float32, is_output=True), + TensorSpec("y", [rows, hidden], torch.float32), ] @@ -83,7 +83,7 @@ def golden_rms_norm(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -92,7 +92,7 @@ def golden_rms_norm(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=rms_norm, specs=build_tensor_specs(), golden_fn=golden_rms_norm, diff --git a/examples/intermediate/rope.py b/examples/intermediate/rope.py index ef8aa609f..ddd1cb05c 100644 --- a/examples/intermediate/rope.py +++ b/examples/intermediate/rope.py @@ -69,7 +69,7 @@ def build_tensor_specs( TensorSpec("x", [total_rows, head_dim], torch.float32, init_value=torch.randn), TensorSpec("cos", [1, head_dim], torch.float32, init_value=torch.randn), TensorSpec("sin", [1, head_dim], torch.float32, init_value=torch.randn), - TensorSpec("y", [total_rows, head_dim], torch.float32, is_output=True), + TensorSpec("y", [total_rows, head_dim], torch.float32), ] @@ -92,7 +92,7 @@ def golden_rope(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -101,7 +101,7 @@ def golden_rope(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=rope, specs=build_tensor_specs(), golden_fn=golden_rope, diff --git a/examples/intermediate/softmax.py b/examples/intermediate/softmax.py index 8f9a4e8f3..4d5e71e82 100644 --- a/examples/intermediate/softmax.py +++ b/examples/intermediate/softmax.py @@ -47,7 +47,7 @@ def build_tensor_specs( return [ TensorSpec("x", [rows, cols], torch.float32, init_value=torch.randn), - TensorSpec("y", [rows, cols], torch.float32, is_output=True), + TensorSpec("y", [rows, cols], torch.float32), ] @@ -59,7 +59,7 @@ def golden_softmax(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -68,7 +68,7 @@ def golden_softmax(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=softmax, specs=build_tensor_specs(), golden_fn=golden_softmax, diff --git a/golden/__init__.py b/golden/__init__.py index b947655ef..9fdb90e16 100644 --- a/golden/__init__.py +++ b/golden/__init__.py @@ -42,7 +42,7 @@ def _configure_golden_threads() -> int: _GOLDEN_NUM_THREADS = _configure_golden_threads() -from .runner import RunResult, run, run_jit +from .runner import RunResult, run from .spec import ScalarSpec, TensorSpec from .validation import ( error_distribution, @@ -66,5 +66,4 @@ def _configure_golden_threads() -> int: "topk_pair_compare", "RunResult", "run", - "run_jit", ] diff --git a/golden/runner.py b/golden/runner.py index c539572cb..b65c76208 100644 --- a/golden/runner.py +++ b/golden/runner.py @@ -9,9 +9,10 @@ """Compile PyPTO programs, run them on device, and validate against goldens. -Public entry points: :func:`run` and :func:`run_jit`. +Public entry points: :func:`run` and :func:`run`. """ +import os import statistics import time from collections.abc import Callable, Sequence @@ -64,18 +65,15 @@ def _required_files(spec: TensorSpec | ScalarSpec) -> list[tuple[str, str]]: :attr:`ScalarSpec.value` tensor). - :class:`TensorSpec` pure input: ``in/{name}.pt``. - :class:`TensorSpec` pure output: ``out/{name}.pt``. - - :class:`TensorSpec` inout (``is_output`` + ``init_value``): - both ``in/{name}.pt`` and ``out/{name}.pt``. + - :class:`TensorSpec` inout: both ``in/{name}.pt`` and ``out/{name}.pt``. """ if isinstance(spec, ScalarSpec): return [("in", f"{spec.name}.pt")] files: list[tuple[str, str]] = [] - if not spec.is_output: + if spec.is_input: files.append(("in", f"{spec.name}.pt")) - else: + if spec.is_output: files.append(("out", f"{spec.name}.pt")) - if spec.init_value is not None: - files.append(("in", f"{spec.name}.pt")) return files @@ -194,10 +192,8 @@ def _stale_cpps(work_dir: Path) -> list[Path]: - any existing sibling ``.so``/``.o`` is older than the cpp itself (cpp was edited after its last build). - Both cases require a rebuild; reporting them uniformly through this - helper keeps the runner's log message honest (previously a missing - binary would log ``no cpp edits ... reusing cached binaries`` even - though ``compile_and_assemble`` would silently rebuild it). + Both cases require a rebuild, so the caller's log line must report them + together. """ stale: list[Path] = [] # Single-chip / L2 builds keep kernels/ + orchestration/ at the root; an L3 @@ -340,7 +336,7 @@ def _prepare_inputs( input_snapshot = { spec.name: tensors[spec.name].clone() for spec in tensor_specs - if not spec.is_output or spec.init_value is not None + if spec.is_input } if save_data: in_dir = work_dir / "data" / "in" @@ -360,11 +356,13 @@ def _prepare_inputs( raise ValueError(f"golden_data is missing files: {missing}") print(f"[RUN] cache hit: {data_dir / 'in'}", flush=True) - # Load inputs + inout initial values from {dir}/in/; pure outputs stay zero-init. - input_names = [s.name for s in tensor_specs if not s.is_output or s.init_value is not None] + # Load inputs + inout initial values from {dir}/in/. A pure output carries + # no input data, so its host buffer -- the read-back destination -- stays + # zero-init rather than re-running the spec's init_value. + input_names = [s.name for s in tensor_specs if s.is_input] tensors = _load_tensors(data_dir, "in", input_names) for spec in tensor_specs: - if spec.is_output and spec.init_value is None: + if not spec.is_input: tensors[spec.name] = torch.zeros(spec.shape, dtype=spec.dtype) scalar_specs_eff = _effective_scalar_specs(scalar_specs, data_dir) @@ -372,6 +370,40 @@ def _prepare_inputs( return tensors, scalar_specs_eff, {} +def _ordered_args( + specs: list[TensorSpec | ScalarSpec], + tensors: dict[str, torch.Tensor], + scalar_specs_eff: dict[str, ScalarSpec], + *, + ctypes_scalars: bool, + benchmark_dispatch_index: int | None = None, +) -> list[Any]: + """Positional dispatch args in spec order. + + Spec order *is* the compiled parameter order: + :func:`_validate_compiled_spec_abi` rejects any artifact whose parameter + names differ from the spec names element by element, so no name-keyed + reordering is needed here. + + ``execute_compiled`` takes ctypes scalars; an L3 dispatch takes the 0-dim + value tensor. *benchmark_dispatch_index* advances a stepped scalar to its + value for that physical benchmark dispatch. + """ + args: list[Any] = [] + for spec in specs: + if isinstance(spec, TensorSpec): + args.append(tensors[spec.name]) + continue + scalar = scalar_specs_eff[spec.name] + if ctypes_scalars: + args.append(scalar.to_ctypes()) + elif benchmark_dispatch_index is None: + args.append(scalar.value) + else: + args.append(scalar.value_for_benchmark_dispatch(benchmark_dispatch_index)) + return args + + def _execute_via_runner( work_dir: Path, specs: list[TensorSpec | ScalarSpec], @@ -379,13 +411,10 @@ def _execute_via_runner( scalar_specs_eff: dict[str, ScalarSpec], runtime_cfg: dict[str, Any], ) -> None: - """Reorder args to orchestration param order and dispatch via ``execute_compiled``.""" + """Dispatch via ``execute_compiled`` in orchestration param order.""" from pypto.runtime import execute_compiled - ordered: list[Any] = [ - tensors[s.name] if isinstance(s, TensorSpec) else scalar_specs_eff[s.name].to_ctypes() - for s in specs - ] + ordered = _ordered_args(specs, tensors, scalar_specs_eff, ctypes_scalars=True) execute_compiled(work_dir, ordered, **_execute_compiled_kwargs(runtime_cfg)) @@ -412,17 +441,20 @@ def _is_l3(compiled: Any) -> bool: _BENCH_WARMUP_DEFAULT = 5 +def _env_flag(name: str) -> bool: + """True when env var *name* holds anything but empty / ``0`` / ``false``.""" + return os.environ.get(name, "").strip() not in ("", "0", "false", "False") + + def _bench_enabled() -> bool: """True when ``PYPTO_BENCH`` is set truthy. Benchmarking is entirely env-driven so no model file needs a ``--benchmark`` - flag and ``run_jit`` needs no extra parameters: daily CI's a2a3 job sets - ``PYPTO_BENCH=1`` and every ``run_jit`` call then times the kernel over + flag and ``run`` needs no extra parameters: daily CI's a2a3 job sets + ``PYPTO_BENCH=1`` and every ``run`` call then times the kernel over :func:`_bench_loop_sizes` rounds (warmup discarded). """ - import os - - return os.environ.get("PYPTO_BENCH", "").strip() not in ("", "0", "false", "False") + return _env_flag("PYPTO_BENCH") def _bench_env_int(name: str, default: int, minimum: int) -> int: @@ -431,8 +463,6 @@ def _bench_env_int(name: str, default: int, minimum: int) -> int: A malformed or out-of-range value warns and uses the default rather than raising: a mistyped tuning knob must not fail an otherwise good run. """ - import os - raw = os.environ.get(name, "").strip() if not raw: return default @@ -457,7 +487,7 @@ def _bench_loop_sizes() -> tuple[int, int]: 100-round default is ~0.1 s of device time for a decode step but minutes for a long prefill or a multi-card L3 run, and while iterating on a kernel a handful of rounds is usually enough. Both are read per run (not cached), so - a sweep can vary them between :func:`run_jit` calls in one process. + a sweep can vary them between :func:`run` calls in one process. Daily CI sets neither, so its numbers stay comparable across runs. Warmup is allowed to be 0; rounds must be at least 1. @@ -472,10 +502,8 @@ def _resident_loop_sizes() -> tuple[int, int]: """:func:`_bench_loop_sizes` with ``warmup`` forced to at least 1. The resident L3 path spends its first warmup launch on the validation - dispatch, so unlike ``benchmark()``'s own loops it cannot honour - ``warmup=0``: that would emit ``rounds + 1`` dispatches per rank against a - declared ``rounds + 0``, which no longer segments evenly and drops the whole - run into the flatten fallback. + dispatch, so ``warmup=0`` would emit ``rounds + 1`` dispatches per rank + against a declared ``rounds + 0`` and stop segmenting evenly. """ rounds, warmup = _bench_loop_sizes() return rounds, max(warmup, 1) @@ -489,9 +517,7 @@ def _bench_raw_enabled() -> bool: suspicious — start-up drift, a bimodal rank, one card lagging — and the individual samples are needed to see the shape. """ - import os - - return os.environ.get("PYPTO_BENCH_RAW", "").strip() not in ("", "0", "false", "False") + return _env_flag("PYPTO_BENCH_RAW") def _benchmark_unavailable(error: RuntimeError) -> bool: @@ -513,8 +539,8 @@ def _run_benchmark( L2 single-chip only: delegates to :func:`pypto.runtime.benchmark`, which opens one :class:`~pypto.runtime.ChipWorker`, registers *compiled* once, and reads each launch's on-NPU span tree from the runtime's ``[STRACE]`` - markers (simpler PR #1177). Args are reordered to the - orchestration parameter order exactly as :func:`_execute_via_runner` does. + markers. Args are built in spec order by :func:`_ordered_args`, exactly as + :func:`_execute_via_runner` does. Returns the :class:`~pypto.runtime.BenchmarkStats`, or ``None`` when the runtime emits no markers (built without ``SIMPLER_PROFILING``). """ @@ -527,10 +553,7 @@ def _run_benchmark( from pypto.runtime import benchmark - ordered: list[Any] = [ - tensors[s.name] if isinstance(s, TensorSpec) else scalar_specs_eff[s.name].to_ctypes() - for s in specs - ] + ordered = _ordered_args(specs, tensors, scalar_specs_eff, ctypes_scalars=True) platform = runtime_cfg.get("platform") device_id = runtime_cfg.get("device_id") stats = None @@ -580,12 +603,9 @@ def _report_effective(stats: Any) -> None: every rank's per-dispatch samples into the same window. The Effective window is the framework's post-graph-build execution window - (``orch``∪``sched``, the old device-log "Total"), surfaced directly by - ``BenchmarkStats.per_round("effective")`` — L2: each launch's window; L3: - per-round max across ranks. This replaces the old hand-rolled span math, - which also hardcoded the pre-#1210 ``run_prepared`` span names; - ``per_round`` resolves the names from the installed runtime. The aggregate is - over the measured rounds (warmup excluded). + (``orch``∪``sched``), read from ``BenchmarkStats.per_round("effective")`` + so the span names come from the installed runtime. The aggregate covers the + measured rounds; warmup is excluded. """ if stats.all_zero_device: print( @@ -718,13 +738,11 @@ def _report_l3_per_rank(stats: Any) -> None: surface the cross-card imbalance the headline (per-round max across ranks) hides. No-op for L2 and the flatten fallback (``per_rank`` returns ``{}``). - Because a rank entry **sums** that card's dispatches (a card runs them - serially), one nested ``slot`` line per dispatch follows each rank line, - read from ``per_dispatch`` and labelled with the orchestration function - ``dispatch_tasks()`` names. Those appear only when some rank dispatches more - than once per round (see :func:`_per_dispatch_effective`); then every rank's - dispatches are listed, so single-dispatch ranks show one slot line restating - their rank line and the block stays a complete table. + A rank entry **sums** that card's dispatches, so each rank line is followed + by one nested ``slot`` line per dispatch, labelled with the orchestration + ``dispatch_tasks()`` name. Slot lines appear only when some rank dispatches + more than once per round (see :func:`_per_dispatch_effective`), and then for + every rank, so the block stays a complete table. All lines use an ``eff_us`` token, so the Daily-CI collector's ``effective_us`` match never selects them. @@ -744,37 +762,6 @@ def _report_l3_per_rank(stats: Any) -> None: _print_eff_summary(label, samples, indent=7) -def _l3_ordered_args( - compiled: Any, - specs: list[TensorSpec | ScalarSpec], - tensors: dict[str, torch.Tensor], - scalar_specs_eff: dict[str, ScalarSpec], - *, - benchmark_dispatch_index: int | None = None, -) -> list[Any]: - """Positional dispatch args for an L3 program, in orchestration param order. - - Builds a name→value map from *specs* (tensors as host tensors, scalars as - their Python value) then reorders it to the compiled program's parameter - order, stripping SSA suffixes ``orig__ssa_vN`` -> ``orig`` (the same mapping - :func:`_try_l3_dispatch` uses). - """ - arg_map: dict[str, Any] = {} - for s in specs: - if isinstance(s, TensorSpec): - arg_map[s.name] = tensors[s.name] - else: - scalar = scalar_specs_eff[s.name] - arg_map[s.name] = ( - scalar.value - if benchmark_dispatch_index is None - else scalar.value_for_benchmark_dispatch(benchmark_dispatch_index) - ) - ordered_names = _l3_ordered_names(compiled) - _validate_l3_arg_names(ordered_names, [s.name for s in specs]) - return [arg_map[name] for name in ordered_names] - - def _run_benchmark_l3( compiled: Any, specs: list[TensorSpec | ScalarSpec], @@ -813,7 +800,7 @@ def _run_benchmark_l3( # L3 dispatch reads IO through the fork-inherited shared mapping; validation # (after this) then reads the device-written outputs back from these buffers. _share_in_place(tensors) - ordered = _l3_ordered_args(compiled, specs, tensors, scalar_specs_eff) + ordered = _ordered_args(specs, tensors, scalar_specs_eff, ctypes_scalars=False) stats = None with _Stage("benchmark"): try: @@ -851,14 +838,10 @@ def _try_l3_dispatch( top-level ``kernel_config.py``); the compiled object is callable directly with ``pypto.runtime.RunConfig``. """ - try: - from pypto.ir.distributed_compiled_program import DistributedCompiledProgram - except ImportError: - return False - if not isinstance(compiled, DistributedCompiledProgram): + if not _is_l3(compiled): return False - ordered = _l3_ordered_args(compiled, specs, tensors, scalar_specs_eff) + ordered = _ordered_args(specs, tensors, scalar_specs_eff, ctypes_scalars=False) run_config = _l3_run_config(runtime_cfg) compiled(*ordered, config=run_config) return True @@ -887,52 +870,38 @@ def _strip_ssa_suffix(name: str) -> str: return base if marker and version.isdigit() else name -def _l3_ordered_names(compiled: Any) -> list[str]: - """Parameter names in orchestration order (SSA suffix ``orig__ssa_vN`` -> ``orig``).""" - param_infos, _, _ = compiled._get_metadata() - names = [_strip_ssa_suffix(p.name) for p in param_infos] - if len(set(names)) != len(names): - raise ValueError("compiled L3 parameters collide after stripping SSA suffixes") - return names - - -def _validate_l3_arg_names(compiled_names: list[str], provided_names: list[str]) -> None: - """Require an exact source-spec ↔ compiled-artifact parameter ABI match.""" - duplicate_specs = sorted( - {name for name in provided_names if provided_names.count(name) > 1} - ) - compiled_set = set(compiled_names) - provided_set = set(provided_names) - missing_specs = sorted(compiled_set - provided_set) - stale_specs = sorted(provided_set - compiled_set) - if not duplicate_specs and not missing_specs and not stale_specs: - return +def _direction_names() -> dict[Any, str]: + """``ParamDirection`` -> :attr:`TensorSpec.direction` string, built per call + so a patched ``ParamDirection`` is never shadowed by a cached map.""" + from pypto.ir import ParamDirection - details = [] - if duplicate_specs: - details.append(f"duplicate specs={duplicate_specs}") - if missing_specs: - details.append(f"compiled parameters without specs={missing_specs}") - if stale_specs: - details.append(f"specs absent from compiled artifact={stale_specs}") - raise ValueError( - "L3 parameter ABI mismatch (" + "; ".join(details) + "); recompile the artifact" - ) + return { + ParamDirection.In: "in", + ParamDirection.Out: "out", + ParamDirection.InOut: "inout", + } def _validate_compiled_spec_abi( compiled: Any, specs: list[TensorSpec | ScalarSpec], ) -> None: - """Validate the complete spec ABI of any live compiled artifact. + """Validate the spec ABI of a live compiled artifact and stamp directions. Signature-driven JIT compilation does not consume tensor sample arguments, so a successful compile alone cannot prove that the caller's specs still match the annotated program. Compare the normalized parameter name, kind, - shape, dtype, and direction before either compile-only success or replay. - A compiled ``-1`` dimension is dynamic and therefore accepts the matching - concrete spec dimension. Lightweight test doubles without metadata are - ignored; real L2 and L3 compiled programs both expose ``_get_metadata``. + shape, and dtype before either compile-only success or replay. A compiled + ``-1`` dimension is dynamic and therefore accepts the matching concrete spec + dimension. + + Direction is not compared but **copied**: the kernel signature owns it, so + each :class:`TensorSpec` takes its :attr:`~TensorSpec.direction` from the + artifact here, before any tensor is allocated. Every later + ``spec.is_output`` / ``spec.is_input`` read resolves against it. + + Lightweight test doubles without metadata are ignored; real L2 and L3 + compiled programs both expose ``_get_metadata``. """ metadata_getter = getattr(compiled, "_get_metadata", None) if not callable(metadata_getter): @@ -941,9 +910,9 @@ def _validate_compiled_spec_abi( if not isinstance(metadata, tuple) or len(metadata) != 3: return - from pypto.ir import ParamDirection from pypto.ir.compiled_program import _to_torch_dtype + directions = _direction_names() param_infos, _, _ = metadata compiled_names = [_strip_ssa_suffix(info.name) for info in param_infos] if len(set(compiled_names)) != len(compiled_names): @@ -965,7 +934,7 @@ def _validate_compiled_spec_abi( + "; ".join(details) + "); recompile the artifact" ) - if not _is_l3(compiled) and compiled_names != provided_names: + if compiled_names != provided_names: raise ValueError( "compiled parameter ABI mismatch (parameter order " f"spec={provided_names} artifact={compiled_names}); recompile the artifact" @@ -974,9 +943,6 @@ def _validate_compiled_spec_abi( specs_by_name = {spec.name: spec for spec in specs} mismatches: list[str] = [] - def _direction_name(direction: Any) -> str: - return getattr(direction, "name", repr(direction)) - for name, info in zip(compiled_names, param_infos, strict=True): spec = specs_by_name[name] artifact_shape = None if info.shape is None else tuple(info.shape) @@ -986,19 +952,15 @@ def _direction_name(direction: Any) -> str: artifact_dtype = None if isinstance(spec, ScalarSpec): - expected_direction = ParamDirection.In if artifact_shape is not None: mismatches.append( f"{name}: expected scalar, artifact is tensor shape={artifact_shape}" ) + if directions.get(info.direction) != "in": + mismatches.append( + f"{name}: scalar direction must be In, artifact={info.direction!r}" + ) else: - if not spec.is_output: - expected_direction = ParamDirection.In - elif spec.init_value is None: - expected_direction = ParamDirection.Out - else: - expected_direction = ParamDirection.InOut - expected_shape = tuple(spec.shape) if artifact_shape is None: mismatches.append(f"{name}: expected tensor shape={expected_shape}, artifact is scalar") @@ -1016,11 +978,10 @@ def _direction_name(direction: Any) -> str: mismatches.append( f"{name}: dtype spec={spec.dtype} artifact={artifact_dtype}" ) - if info.direction != expected_direction: - mismatches.append( - f"{name}: direction spec={_direction_name(expected_direction)} " - f"artifact={_direction_name(info.direction)}" - ) + if isinstance(spec, TensorSpec): + spec.direction = directions.get(info.direction) + if spec.direction is None: + mismatches.append(f"{name}: unknown artifact direction {info.direction!r}") if mismatches: raise ValueError( @@ -1035,10 +996,9 @@ def _l3_pure_out_names(compiled: Any) -> set[str]: from pypto.ir import ParamDirection param_infos, _, _ = compiled._get_metadata() - normalized_names = _l3_ordered_names(compiled) return { - name - for name, p in zip(normalized_names, param_infos, strict=True) + _strip_ssa_suffix(p.name) + for p in param_infos if p.direction == ParamDirection.Out } @@ -1124,7 +1084,7 @@ def _readback_resident_outputs( def _run_l3_resident( compiled: Any, - tensor_specs: list[TensorSpec], + specs: list[TensorSpec | ScalarSpec], tensors: dict[str, torch.Tensor], scalar_specs_eff: dict[str, ScalarSpec], runtime_cfg: dict[str, Any], @@ -1140,26 +1100,22 @@ def _run_l3_resident( that can build worker-resident :class:`~pypto.runtime.DeviceTensor` buffers. Each resident input / ``InOut`` spec is uploaded once via ``rt.alloc_tensor(init=...)`` and reused across the validation dispatch and - every benchmark round. A pure ``Out`` resident is allocated uninitialized, - because its host tensor is only an output destination and uploading its - zero-filled placeholder would be wasted work. Resident outputs are read back - once before golden validation via :func:`_readback_resident_outputs`. - - When *benchmark_enabled* is true (or defaults to - :func:`_bench_enabled` via ``PYPTO_BENCH``), the resident weights are reused - for :func:`_bench_loop_sizes` timed rounds. This cannot go through - :func:`pypto.runtime.benchmark` — that owns its own ``prepare()``, and a - resident buffer allocated on our worker is invisible to a second, separately - forked one — so it mirrors ``benchmark``'s L3 path by hand: raise the runtime - log level to ``v9`` and set up the fd-level ``[STRACE]`` capture *around* - ``prepare()`` (the forked chip workers inherit fd 2 at fork time), then parse - the captured markers into a :class:`BenchmarkStats` with real per-round L3 - device / effective timing (max across ranks) — not just host wall. - - Validation runs on the first dispatch (a correctness gate that propagates an - ``AssertionError``); the benchmark rounds that follow are never a correctness - gate (a failure there is logged, not raised). Returns a :class:`BenchmarkStats` - or ``None``. + every benchmark round; a pure ``Out`` resident is allocated uninitialized. + Resident outputs are read back once before golden validation via + :func:`_readback_resident_outputs`. + + When *benchmark_enabled* is true (default: :func:`_bench_enabled`), the + resident weights are reused for :func:`_bench_loop_sizes` timed rounds. + :func:`pypto.runtime.benchmark` cannot serve this — it owns its own + ``prepare()``, and a buffer allocated on our worker is invisible to a + second, separately forked one — so the capture is mirrored here by hand: + raise the runtime log level to ``v9`` and wrap ``prepare()`` in the + fd-level ``[STRACE]`` capture (the forked chip workers inherit fd 2 at fork + time), then parse the markers into a :class:`BenchmarkStats`. + + Validation runs on the first dispatch and propagates its ``AssertionError``; + a failure in the benchmark rounds that follow is logged, not raised. Returns + a :class:`BenchmarkStats` or ``None``. """ try: from pypto.ir.distributed_compiled_program import DistributedCompiledProgram @@ -1177,14 +1133,11 @@ def _run_l3_resident( # Per-call IO + resident upload sources must be shared memory before prepare(). _share_in_place(tensors) - ordered_names = _l3_ordered_names(compiled) - _validate_l3_arg_names( - ordered_names, - [*tensors.keys(), *scalar_specs_eff.keys()], - ) + ordered_names = [spec.name for spec in specs] pure_out_names = _l3_pure_out_names(compiled) run_config = _l3_run_config(runtime_cfg) - resident_specs = [s for s in tensor_specs if s.is_resident] + tensor_specs = [spec for spec in specs if isinstance(spec, TensorSpec)] + resident_specs = [spec for spec in tensor_specs if spec.is_resident] bench = _bench_enabled() if benchmark_enabled is None else benchmark_enabled def _dispatch_resident( @@ -1378,7 +1331,7 @@ def _maybe_reload_l3( Returns ``None`` for a single-chip / L2 build (which keeps using ``execute_compiled``). An L3 build is identified by the - ``distributed_meta.json`` sidecar written at compile time (pypto #1689); + ``distributed_meta.json`` sidecar written at compile time; :meth:`DistributedCompiledProgram.from_dir` rebuilds its metadata without re-running the pypto compile, so the existing :func:`_try_l3_dispatch` path can dispatch it. The run's ``platform`` and ``distributed_config`` override @@ -1419,8 +1372,9 @@ def _compute_golden( """Produce golden output tensors for validation. With *data_dir* set, load from ``{data_dir}/out/``. Otherwise call - *golden_fn* on a scratch dict (inputs cloned from *input_snapshot*, - outputs zero-init) and, when *save_data* is True, persist results into + *golden_fn* on a scratch dict (input tensors cloned from *input_snapshot*, + pure outputs from their own ``init_value``) and, when + *save_data* is True, persist results into ``{work_dir}/data/out/``. """ with _Stage("compute golden"): @@ -1433,10 +1387,10 @@ def _compute_golden( for spec in specs: if isinstance(spec, ScalarSpec): scratch[spec.name] = scalar_specs_eff[spec.name].to_python() - elif spec.is_output and spec.init_value is None: - scratch[spec.name] = torch.zeros(spec.shape, dtype=spec.dtype) - else: + elif spec.is_input: scratch[spec.name] = input_snapshot[spec.name].clone() + else: + scratch[spec.name] = spec.create_tensor() golden_fn(scratch) golden_outputs = {spec.name: scratch[spec.name] for spec in tensor_specs if spec.is_output} if save_data: @@ -1476,40 +1430,47 @@ def _validate( ) -def run( - program: Any, +def _run_pipeline( specs: list[TensorSpec | ScalarSpec], - golden_fn: Callable | None = None, - golden_data: str | None = None, - compile_cfg: dict[str, Any] | None = None, - runtime_cfg: dict[str, Any] | None = None, - rtol: float = 1e-5, - atol: float = 1e-5, - compare_fn: dict[str, Callable] | None = None, - compile_only: bool = False, - runtime_dir: str | None = None, - save_data: bool = False, + compile_step: Callable[[dict[str, Any], dict[str, Any], Any], Any], + compile_label: str, + prologue: Callable[[list[ScalarSpec], Path | None], Any] | None, + golden_fn: Callable | None, + golden_data: str | None, + compile_cfg: dict[str, Any] | None, + runtime_cfg: dict[str, Any] | None, + rtol: float, + atol: float, + compare_fn: dict[str, Callable] | None, + compile_only: bool, + runtime_dir: str | None, + save_data: bool, ) -> RunResult: - """Compile *program*, run on device, and validate against golden. + """Shared body of :func:`run` and :func:`run`. + + *prologue* runs entry-specific spec validation, may raise ``ValueError``, + and returns whatever state its *compile_step* needs. *compile_step* then + returns the ``CompiledProgram`` for a fresh compile, given the normalized + configs and that state. Everything around the two is identical for both + entry points. Args: - program: ``@pl.program`` class or ``ir.Program``. - specs: :class:`TensorSpec` / :class:`ScalarSpec` list in orchestration - parameter order. + specs: :class:`TensorSpec` / :class:`ScalarSpec` list in the compiled + program's parameter order. A mismatched order is rejected by + :func:`_validate_compiled_spec_abi`, never reordered. golden_fn: ``golden_fn(values)`` that fills outputs in-place; *values* maps spec name to tensor clone or Python scalar. Ignored when *golden_data* is set; if neither is given, validation is skipped. golden_data: Directory with ``in/{name}.pt`` and ``out/{name}.pt``; loads inputs and expected outputs (read-only). Takes precedence over *golden_fn*. - compile_cfg: Kwargs forwarded to :func:`pypto.ir.compile`. Unknown - keys raise there. + compile_cfg: Entry-specific compile kwargs; consumed by *compile_step* + and by the ``runtime_dir`` L3 reload. runtime_cfg: Kwargs forwarded to :func:`pypto.runtime.execute_compiled` (``platform``, ``device_id``, ``enable_chip_swimlane``, ...). Unknown keys raise there, except - the harness-only key ``log_level``, which is consumed up-front - to configure the PyPTO runtime logger via - :func:`pypto.runtime.log_config.configure_log`. + the harness-only key ``log_level``, consumed up-front by + :func:`_consume_runtime_harness_keys`. rtol, atol: Golden comparison tolerances. compare_fn: Per-output-name overrides for ``torch.allclose``; see :func:`golden.validation.validate_golden`. @@ -1518,22 +1479,17 @@ def run( compile and invalidates cached ``.so``/``.bin`` so cpp edits rebuild; *compile_cfg* is ignored, *compile_only* is rejected, and ``PYPTO_BENCH`` is skipped because replay is correctness-only. - save_data: When True, persist generated inputs to - ``{work_dir}/data/in/`` and golden outputs to - ``{work_dir}/data/out/`` for later replay via *golden_data*. - Defaults to False, skipping the on-disk ``.pt`` snapshot; - validation still runs against the in-memory golden. Enable it - when you need to replay the exact inputs/outputs later. + save_data: Persist generated inputs to ``{work_dir}/data/in/`` and + golden outputs to ``{work_dir}/data/out/`` for later replay via + *golden_data*. Off by default; validation still runs against the + in-memory golden. Returns: :class:`RunResult`. """ - from pypto import ir - compile_cfg = compile_cfg or {} runtime_cfg = dict(runtime_cfg or {}) # copy: we pop harness-only keys compare_fn = compare_fn or {} - _consume_runtime_harness_keys(runtime_cfg) if compile_only and runtime_dir is not None: @@ -1552,17 +1508,19 @@ def _fail(error: str) -> RunResult: execution_time=time.time() - start, work_dir=work_dir, ) + compile_state: Any = None try: _validate_unique_spec_names(specs) _validate_stepped_swimlane(scalar_specs, runtime_cfg) + if prologue is not None: + compile_state = prologue(scalar_specs, data_dir) except ValueError as e: return _fail(str(e)) - # Compile (or pick runtime_dir) - compiled: Any = None + compiled: Any = None # the CompiledProgram, when we compiled it this call if runtime_dir is not None: try: - work_dir = _setup_runtime_dir(runtime_dir, compile_label="compile") + work_dir = _setup_runtime_dir(runtime_dir, compile_label=compile_label) except ValueError as e: return _fail(str(e)) # An L3 build has no live compiled object here (compile was skipped); @@ -1571,23 +1529,13 @@ def _fail(error: str) -> RunResult: compiled = _maybe_reload_l3(work_dir, runtime_cfg, compile_cfg) else: with _Stage("compile"): - compile_kwargs = dict(compile_cfg) - platform = runtime_cfg.get("platform") - if platform is not None: - compile_kwargs.setdefault("backend_type", _backend_for_platform(platform)) - # L3 distributed programs bake the platform into compiled.platform - # at compile time (the runtime config's platform is ignored when - # assembling chip callables). Without this, compiled.platform falls - # back to the backend's default sim platform, so a `-p a2a3` run - # silently compiles incore kernels for a2a3sim (g++-15) instead of - # the real device (ccec). - compile_kwargs.setdefault("platform", platform) - compiled = ir.compile(program, **compile_kwargs) + compiled = compile_step(compile_cfg, runtime_cfg, compile_state) work_dir = Path(compiled.output_dir) - # A live L3 object is available after both a fresh compile and a runtime-dir - # reload. Validate its complete metadata ABI before allocating any inputs or - # allowing compile-only to report success. + # Neither a signature-driven compile (which trusts annotations over tensor + # samples) nor a runtime-dir replay (which trusts a persisted artifact) can + # prove the specs still describe the program. Reject stale ones before + # allocating any input or letting compile-only report success. try: _validate_compiled_spec_abi(compiled, specs) except ValueError as e: @@ -1597,7 +1545,6 @@ def _fail(error: str) -> RunResult: print(f"[RUN] PASS ({total:.2f}s)", flush=True) return RunResult(passed=True, execution_time=total, work_dir=work_dir) - # Generate Inputs try: with _Stage("generate inputs"): tensors, scalar_specs_eff, input_snapshot = _prepare_inputs( @@ -1607,7 +1554,6 @@ def _fail(error: str) -> RunResult: except ValueError as e: return _fail(str(e)) - # Compute Golden golden_outputs: dict[str, torch.Tensor] | None = None if golden_fn is not None or golden_data is not None: golden_outputs = _compute_golden( @@ -1623,27 +1569,36 @@ def _fail(error: str) -> RunResult: ) benchmark_enabled = False - # Resident-weight path: keep resident specs device-resident across - # the validation dispatch and any benchmark rounds via the L3 prepare() - # worker (validation + benchmark are handled inside; return early). + def _pass(bench: Any) -> RunResult: + total = time.time() - start + skip_note = ( + ", validation skipped: no golden_fn or golden_data" + if golden_outputs is None else "" + ) + print(f"[RUN] PASS ({total:.2f}s{skip_note})", flush=True) + return RunResult( + passed=True, execution_time=total, work_dir=work_dir, bench=bench, + ) + + # Resident-weight path: keep resident specs device-resident across the + # validation dispatch and any benchmark rounds via the L3 prepare() worker + # (validation + benchmark are handled inside; return early). if any(s.is_resident for s in tensor_specs): with _Stage("runtime"): try: bench = _run_l3_resident( - compiled, tensor_specs, tensors, scalar_specs_eff, + compiled, specs, tensors, scalar_specs_eff, runtime_cfg, golden_outputs, rtol, atol, compare_fn, benchmark_enabled=benchmark_enabled, ) except (AssertionError, ValueError) as e: return _fail(str(e)) - validation_skipped = golden_outputs is None - total = time.time() - start - skip_note = ", validation skipped: no golden_fn or golden_data" if validation_skipped else "" - print(f"[RUN] PASS ({total:.2f}s{skip_note})", flush=True) - return RunResult(passed=True, execution_time=total, work_dir=work_dir, bench=bench) + return _pass(bench) - # Runtime with _Stage("runtime"): + # An L3 ``DistributedCompiledProgram`` (a @pl.jit.host kernel compiled + # with distributed_config) dispatches per-rank via _try_l3_dispatch; + # everything else runs through the single-chip runner. if compiled is None or not _try_l3_dispatch( compiled, specs, tensors, scalar_specs_eff, runtime_cfg, ): @@ -1654,13 +1609,8 @@ def _fail(error: str) -> RunResult: if golden_outputs is not None: try: _validate( - tensor_specs, - tensors, - golden_outputs, - rtol, - atol, - compare_fn, - scalar_specs_eff, + tensor_specs, tensors, golden_outputs, + rtol, atol, compare_fn, scalar_specs_eff, ) except AssertionError as e: return _fail(str(e)) @@ -1672,24 +1622,113 @@ def _fail(error: str) -> RunResult: bench = None if benchmark_enabled: rounds, warmup = _bench_loop_sizes() - if _is_l3(compiled): - bench = _run_benchmark_l3( - compiled, specs, tensors, scalar_specs_eff, runtime_cfg, - rounds, warmup, - ) - else: - bench = _run_benchmark( - compiled, specs, tensors, scalar_specs_eff, runtime_cfg, - rounds, warmup, + run_bench = _run_benchmark_l3 if _is_l3(compiled) else _run_benchmark + bench = run_bench( + compiled, specs, tensors, scalar_specs_eff, runtime_cfg, rounds, warmup, + ) + return _pass(bench) + + +def _program_entry( + fn: Any, specs: list[TensorSpec | ScalarSpec] +) -> tuple[Callable[[dict[str, Any], dict[str, Any], Any], Any], None, str]: + """``(compile_step, prologue, label)`` for a ``@pl.program`` class / ``ir.Program``.""" + del specs # the program path derives everything from the compiled artifact + + def _compile( + compile_cfg: dict[str, Any], runtime_cfg: dict[str, Any], _state: Any + ) -> Any: + from pypto import ir + + compile_kwargs = dict(compile_cfg) + platform = runtime_cfg.get("platform") + if platform is not None: + compile_kwargs.setdefault("backend_type", _backend_for_platform(platform)) + # L3 distributed programs bake the platform into compiled.platform at + # compile time (the runtime config's platform is ignored when + # assembling chip callables). Without this, compiled.platform falls + # back to the backend's default sim platform, so a `-p a2a3` run + # silently compiles incore kernels for a2a3sim (g++-15) instead of + # the real device (ccec). + compile_kwargs.setdefault("platform", platform) + return ir.compile(fn, **compile_kwargs) + + return _compile, None, "Program compile" + + +def _jit_entry( + fn: Any, specs: list[TensorSpec | ScalarSpec] +) -> tuple[ + Callable[[dict[str, Any], dict[str, Any], Any], Any], + Callable[[list[ScalarSpec], Path | None], Any], + str, +]: + """``(compile_step, prologue, label)`` for a ``@pl.jit`` callable. + + The prologue resolves the scalar values the specialization key needs and + hands them to the compile step. + """ + + def _prologue( + scalar_specs: list[ScalarSpec], data_dir: Path | None + ) -> dict[str, ScalarSpec]: + compile_scalars = _effective_scalar_specs(scalar_specs, data_dir) + # A stepped scalar must survive specialization as a runtime parameter; + # a literal-specialized one would bake dispatch 0's value into the code. + stepped = sorted( + spec.name for spec in scalar_specs + if spec.has_benchmark_step and not spec.compile_runtime + ) + if stepped: + raise ValueError( + "ScalarSpec benchmark_step requires compile_runtime=True; " + f"stepped scalars: {stepped}" ) + return compile_scalars + + def _compile( + compile_cfg: dict[str, Any], + runtime_cfg: dict[str, Any], + compile_scalars: dict[str, ScalarSpec], + ) -> Any: + from pypto.runtime import RunConfig + + cfg = dict(compile_cfg) + platform = runtime_cfg.get("platform") + if platform is not None: + cfg["platform"] = platform + scalar_specs = [s for s in specs if isinstance(s, ScalarSpec)] + if any(spec.compile_runtime for spec in scalar_specs): + import pypto.language as pl + + # pl.RUNTIME is accepted only by annotation-driven signature + # compilation: omit every tensor sample and provide all scalar + # parameters by name. Unmarked scalars retain their literal + # specialization semantics. + scalar_compile_args = { + spec.name: ( + pl.RUNTIME + if spec.compile_runtime + else compile_scalars[spec.name].to_python() + ) + for spec in scalar_specs + } + return fn.compile(config=RunConfig(**cfg), **scalar_compile_args) + # Dummy args carry shape/dtype and scalar values into the specialization + # key; real tensors of the same shape hit the same JIT cache entry at + # dispatch. + dummy_args = [ + compile_scalars[spec.name].to_python() + if isinstance(spec, ScalarSpec) + else torch.empty(spec.shape, dtype=spec.dtype) + for spec in specs + ] + return fn.compile(*dummy_args, config=RunConfig(**cfg)) - total = time.time() - start - skip_note = ", validation skipped: no golden_fn or golden_data" if golden_outputs is None else "" - print(f"[RUN] PASS ({total:.2f}s{skip_note})", flush=True) - return RunResult(passed=True, execution_time=total, work_dir=work_dir, bench=bench) + return _compile, _prologue, "JIT compile" -def run_jit( +def run( fn: Any, specs: list[TensorSpec | ScalarSpec], golden_fn: Callable | None = None, @@ -1703,23 +1742,29 @@ def run_jit( runtime_dir: str | None = None, save_data: bool = False, ) -> RunResult: - """JIT-flavoured :func:`run`: compile via ``@pl.jit``, then same harness. + """Compile *fn*, run it on device, and validate against golden. + + Accepts either kernel form. A ``@pl.jit`` callable exposes ``compile`` and + is specialized through ``JITFunction.compile``; a ``@pl.program`` class or + an ``ir.Program`` goes straight to :func:`pypto.ir.compile`. The two differ + only in the compile step and in which *compile_cfg* keys they accept. Args: - fn: ``@pl.jit`` decorated callable. - specs: :class:`TensorSpec` / :class:`ScalarSpec` list in the JIT - function's parameter order. + fn: ``@pl.jit`` callable, ``@pl.program`` class, or ``ir.Program``. + specs: :class:`TensorSpec` / :class:`ScalarSpec` list in *fn*'s + parameter order. A mismatched order is rejected, never reordered. golden_fn: ``golden_fn(values)`` that fills outputs in-place; *values* maps spec name to tensor clone or Python scalar. Ignored when *golden_data* is set; if neither is given, validation is skipped. golden_data: Directory with ``in/{name}.pt`` and ``out/{name}.pt``; loads inputs and expected outputs (read-only). Takes precedence over *golden_fn*. - compile_cfg: Compile-side ``RunConfig`` fields (``dump_passes`` / - ``distributed_config`` / ``compile_profiling`` / ...) carried into - ``JITFunction.compile``; ``platform`` is supplied separately - (typically via *runtime_cfg*). Unknown keys raise when the - ``RunConfig`` is built. + compile_cfg: For a ``@pl.jit`` kernel, compile-side ``RunConfig`` + fields (``dump_passes`` / ``distributed_config`` / + ``compile_profiling`` / ...) carried into ``JITFunction.compile``; + ``platform`` is supplied separately, typically via *runtime_cfg*. + For a ``@pl.program`` kernel, kwargs forwarded to + :func:`pypto.ir.compile`. Unknown keys raise either way. runtime_cfg: Kwargs forwarded to :func:`pypto.runtime.execute_compiled` (``platform``, ``device_id``, ``enable_chip_swimlane``, ...). Unknown keys raise there, except @@ -1744,202 +1789,12 @@ def run_jit( Returns: :class:`RunResult`. """ - compile_cfg = compile_cfg or {} - runtime_cfg = dict(runtime_cfg or {}) # copy: we pop harness-only keys - compare_fn = compare_fn or {} - - _consume_runtime_harness_keys(runtime_cfg) - - if compile_only and runtime_dir is not None: - return RunResult(passed=False, error="runtime_dir is incompatible with compile_only") - - data_dir = Path(golden_data) if golden_data is not None else None - tensor_specs = [s for s in specs if isinstance(s, TensorSpec)] - scalar_specs = [s for s in specs if isinstance(s, ScalarSpec)] - - start = time.time() - work_dir: Path | None = None - - def _fail(error: str) -> RunResult: - return RunResult( - passed=False, error=error, - execution_time=time.time() - start, work_dir=work_dir, - ) - - try: - _validate_unique_spec_names(specs) - compile_scalar_specs_eff = _effective_scalar_specs(scalar_specs, data_dir) - _validate_stepped_swimlane(scalar_specs, runtime_cfg) - except ValueError as e: - return _fail(str(e)) - - invalid_stepped_scalars = sorted( - spec.name - for spec in scalar_specs - if spec.has_benchmark_step and not spec.compile_runtime + # A JITFunction exposes compile(); a @pl.program class evaluates to an + # ir.Program, which does not. + entry = _jit_entry if callable(getattr(fn, "compile", None)) else _program_entry + compile_step, prologue, compile_label = entry(fn, specs) + return _run_pipeline( + specs, compile_step, compile_label, prologue, + golden_fn, golden_data, compile_cfg, runtime_cfg, + rtol, atol, compare_fn, compile_only, runtime_dir, save_data, ) - if invalid_stepped_scalars: - return _fail( - "run_jit ScalarSpec benchmark_step requires compile_runtime=True; " - f"stepped scalars: {invalid_stepped_scalars}" - ) - - # Compile - compiled: Any = None # the CompiledProgram, when we compiled it this call - if runtime_dir is not None: - try: - work_dir = _setup_runtime_dir(runtime_dir, compile_label="JIT compile") - except ValueError as e: - return _fail(str(e)) - # An L3 build has no live compiled object here (JIT compile was skipped); - # reconstruct it from the build dir so the L3 dispatch path below runs - # instead of falling through to the single-chip execute_compiled. - compiled = _maybe_reload_l3(work_dir, runtime_cfg, compile_cfg) - else: - with _Stage("compile"): - from pypto.runtime import RunConfig - - cfg = dict(compile_cfg) - platform = runtime_cfg.get("platform") - if platform is not None: - cfg["platform"] = platform - # Public compile-only entry: same specialize → cache → ir.compile - # pipeline as __call__, minus on-device dispatch. Returns a - # DistributedCompiledProgram for an L3 host orchestrator. - if any(spec.compile_runtime for spec in scalar_specs): - import pypto.language as pl - - # pl.RUNTIME is accepted only by annotation-driven signature - # compilation: omit every tensor sample and provide all scalar - # parameters by name. Unmarked scalars retain their literal - # specialization semantics. - scalar_compile_args = { - spec.name: ( - pl.RUNTIME - if spec.compile_runtime - else compile_scalar_specs_eff[spec.name].to_python() - ) - for spec in scalar_specs - } - compiled = fn.compile( - config=RunConfig(**cfg), - **scalar_compile_args, - ) - else: - # Dummy args carry shape/dtype and scalar values into the - # specialization key; real tensors of the same shape hit the - # same JIT cache entry at dispatch. - dummy_args = [ - compile_scalar_specs_eff[spec.name].to_python() - if isinstance(spec, ScalarSpec) - else torch.empty(spec.shape, dtype=spec.dtype) - for spec in specs - ] - compiled = fn.compile(*dummy_args, config=RunConfig(**cfg)) - work_dir = Path(compiled.output_dir) - - # Signature-driven compilation trusts the function annotations rather than - # tensor samples, and a runtime-dir replay trusts a persisted artifact. In - # both cases, reject stale specs before allocating any inputs or dispatching. - try: - _validate_compiled_spec_abi(compiled, specs) - except ValueError as e: - return _fail(str(e)) - if compile_only: - total = time.time() - start - print(f"[RUN] PASS ({total:.2f}s)", flush=True) - return RunResult(passed=True, execution_time=total, work_dir=work_dir) - - # Generate Inputs - try: - with _Stage("generate inputs"): - tensors, scalar_specs_eff, input_snapshot = _prepare_inputs( - specs, tensor_specs, scalar_specs, data_dir, work_dir, save_data, - need_snapshot=golden_fn is not None, - ) - except ValueError as e: - return _fail(str(e)) - - # Compute Golden - golden_outputs: dict[str, torch.Tensor] | None = None - if golden_fn is not None or golden_data is not None: - golden_outputs = _compute_golden( - specs, tensor_specs, scalar_specs_eff, input_snapshot, - work_dir, data_dir, golden_fn, save_data, - ) - - benchmark_enabled = _bench_enabled() - if benchmark_enabled and runtime_dir is not None: - print( - "[RUN] benchmark skipped: runtime_dir replay is correctness-only", - flush=True, - ) - benchmark_enabled = False - - # Resident-weight path: keep resident specs device-resident across - # the validation dispatch and any benchmark rounds via the L3 prepare() - # worker (validation + benchmark are handled inside; return early). - if any(s.is_resident for s in tensor_specs): - with _Stage("runtime"): - try: - bench = _run_l3_resident( - compiled, tensor_specs, tensors, scalar_specs_eff, - runtime_cfg, golden_outputs, rtol, atol, compare_fn, - benchmark_enabled=benchmark_enabled, - ) - except (AssertionError, ValueError) as e: - return _fail(str(e)) - validation_skipped = golden_outputs is None - total = time.time() - start - skip_note = ", validation skipped: no golden_fn or golden_data" if validation_skipped else "" - print(f"[RUN] PASS ({total:.2f}s{skip_note})", flush=True) - return RunResult(passed=True, execution_time=total, work_dir=work_dir, bench=bench) - - # Runtime - with _Stage("runtime"): - # An L3 ``DistributedCompiledProgram`` (a @pl.jit.host kernel compiled - # with distributed_config) dispatches per-rank via _try_l3_dispatch; - # everything else runs through the single-chip runner. - if compiled is None or not _try_l3_dispatch( - compiled, specs, tensors, scalar_specs_eff, runtime_cfg, - ): - _execute_via_runner(work_dir, specs, tensors, scalar_specs_eff, runtime_cfg) - - # Validate the dedicated correctness dispatch before benchmark launches - # mutate output or inout tensors in place. - if golden_outputs is not None: - try: - _validate( - tensor_specs, - tensors, - golden_outputs, - rtol, - atol, - compare_fn, - scalar_specs_eff, - ) - except AssertionError as e: - return _fail(str(e)) - - # Benchmark (L2 via _run_benchmark, non-resident L3 via _run_benchmark_l3). - # Runs only after the correctness dispatch has been validated. A runtime-dir - # replay is correctness-only even when an L3 object was reconstructed from - # metadata. Entirely env-gated via PYPTO_BENCH=1 (daily CI). - bench = None - if benchmark_enabled: - rounds, warmup = _bench_loop_sizes() - if _is_l3(compiled): - bench = _run_benchmark_l3( - compiled, specs, tensors, scalar_specs_eff, runtime_cfg, - rounds, warmup, - ) - else: - bench = _run_benchmark( - compiled, specs, tensors, scalar_specs_eff, runtime_cfg, - rounds, warmup, - ) - - total = time.time() - start - skip_note = ", validation skipped: no golden_fn or golden_data" if golden_outputs is None else "" - print(f"[RUN] PASS ({total:.2f}s{skip_note})", flush=True) - return RunResult(passed=True, execution_time=total, work_dir=work_dir, bench=bench) diff --git a/golden/spec.py b/golden/spec.py index d35960106..ab6adf7fa 100644 --- a/golden/spec.py +++ b/golden/spec.py @@ -45,15 +45,18 @@ class TensorSpec: one of the supported ``torch`` factory functions (``torch.randn``, ``torch.rand``, ``torch.zeros``, ``torch.ones``) that will be called with ``(shape, dtype=dtype)``. - is_output: If ``True``, the tensor is an output to be validated against the - golden reference. + + The value is the tensor's initial host content and nothing else. + The runtime does not upload a pure ``Out`` parameter's host buffer + (see ``docs/pypto-coding/pypto-coding-style.md``), so an + ``init_value`` there reaches only the golden reference. resident: Keep this tensor device-resident (``child_memory``): the harness uploads inputs once and reuses them across the validation dispatch and every benchmark round, skipping the per-dispatch host→device upload and device→host readback. Only supported for L3 distributed programs. A pure ``Out`` resident is allocated uninitialized; its zero-filled - host placeholder is not uploaded. Combine with ``is_output=True`` - for either a pure output or a read-write **resident state buffer** + host placeholder is not uploaded. Declare the parameter ``Out`` or + ``InOut`` for either a pure output or a read-write **resident state buffer** (e.g. a KV cache): an ``InOut`` state is uploaded once from ``init_value``, both kinds are updated on-device, and they are read back **once** at the end (via ``copy_stacked_from`` / ``copy_from``) @@ -74,11 +77,16 @@ class TensorSpec: each rank's slice must reside on the card that consumes it. ``True`` is rejected as ambiguous — pass an int worker id instead. + direction: ``"in"`` / ``"out"`` / ``"inout"``, stamped by the harness + from the compiled artifact's ``ParamDirection`` before any tensor is + allocated. Not an init argument: the kernel signature owns the + direction, and reading :attr:`is_output` / :attr:`is_input` before + the stamp raises. Example: >>> import torch >>> TensorSpec("query", [32, 128], torch.bfloat16, init_value=torch.randn) - >>> TensorSpec("out", [32, 128], torch.float32, is_output=True) + >>> TensorSpec("out", [32, 128], torch.float32) # direction comes from the artifact >>> TensorSpec("wq_a", [2, 4096, 1536], torch.bfloat16, init_value=torch.randn, resident="stacked") >>> TensorSpec("bias", [128], torch.float32, init_value=torch.randn, resident=1) # whole on card 1 """ @@ -87,11 +95,11 @@ class TensorSpec: shape: list[int] dtype: torch.dtype init_value: int | float | torch.Tensor | Callable | None = field(default=None) - is_output: bool = False resident: int | str | bool | None = None + direction: str | None = field(default=None, init=False, repr=False, compare=False) def __post_init__(self) -> None: - # Validate the ``resident`` mode. ``resident`` + ``is_output`` is allowed: + # Validate the ``resident`` mode. A resident output is allowed: # a read-write resident state buffer (e.g. a KV cache) uploaded once, # updated in place on-device across dispatches, and read back once at the # end for validation. @@ -115,6 +123,28 @@ def __post_init__(self) -> None: f"got {r!r}." ) + def _stamped_direction(self) -> str: + """The artifact-stamped direction; reading it before the stamp is a bug.""" + if self.direction is None: + raise RuntimeError( + f"TensorSpec {self.name!r}: direction not stamped -- the compiled " + f"artifact's parameter metadata has not been inspected yet" + ) + return self.direction + + @property + def is_output(self) -> bool: + """True if the compiled parameter is written by the kernel (``Out`` / ``InOut``).""" + return self._stamped_direction() in ("out", "inout") + + @property + def is_input(self) -> bool: + """True if the host tensor's initial content is input data (``In`` / ``InOut``). + + Overlaps :attr:`is_output` on ``InOut``, mirroring ``ParamDirection``. + """ + return self._stamped_direction() in ("in", "inout") + @property def is_resident(self) -> bool: """True if this spec is device-resident in any mode (int worker id or ``"stacked"``).""" @@ -219,7 +249,7 @@ class ScalarSpec: *dtype* (used directly). After ``__post_init__`` runs, ``value`` is **always** a 0-dim ``torch.Tensor`` carrying the dtype-precise representation, so cache I/O is just ``torch.save`` / ``torch.load``. - compile_runtime: Whether :func:`golden.run_jit` must leave this scalar + compile_runtime: Whether :func:`golden.run` must leave this scalar unspecialized in the compiled artifact. Marked scalars are passed to ``JITFunction.compile`` as ``pl.RUNTIME`` and remain real runtime ABI parameters. This flag has no effect on the direct diff --git a/models/deepseek_v4_flash_dspark/decode_compressor_ratio128.py b/models/deepseek_v4_flash_dspark/decode_compressor_ratio128.py index 773d61b75..d05014fc9 100644 --- a/models/deepseek_v4_flash_dspark/decode_compressor_ratio128.py +++ b/models/deepseek_v4_flash_dspark/decode_compressor_ratio128.py @@ -565,8 +565,8 @@ def init_cmp_slot_mapping(): ) return [ TensorSpec("x", [batch * S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [batch * S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [batch * S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [batch, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -574,7 +574,7 @@ def init_cmp_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("cos", [batch, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_cos), TensorSpec("sin", [batch, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), - TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache, is_output=True), + TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache), TensorSpec("position_ids", [batch * S], torch.int32, init_value=lambda: init_position_ids().reshape(-1)), TensorSpec("cmp_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_cmp_slot_mapping().reshape(-1)), TensorSpec("state_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_state_slot_mapping().reshape(-1)), @@ -583,7 +583,7 @@ def init_cmp_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -602,7 +602,7 @@ def init_cmp_slot_mapping(): if args.batch < 4 or args.batch > B or args.batch % 4 != 0: parser.error(f"--batch must be a multiple of 4 in [4, {B}], got {args.batch}") - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(args.start_pos, batch=args.batch), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_flash_dspark/decode_compressor_ratio4.py b/models/deepseek_v4_flash_dspark/decode_compressor_ratio4.py index 66fdedea6..913e92075 100644 --- a/models/deepseek_v4_flash_dspark/decode_compressor_ratio4.py +++ b/models/deepseek_v4_flash_dspark/decode_compressor_ratio4.py @@ -529,8 +529,8 @@ def init_cmp_slot_mapping(): return [ TensorSpec("x", [batch * S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [batch * S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [state_block_num, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [batch * S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [state_block_num, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [batch, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -538,7 +538,7 @@ def init_cmp_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("cos", [batch * S, ROPE_HEAD_DIM], torch.float32, init_value=init_cos), TensorSpec("sin", [batch * S, ROPE_HEAD_DIM], torch.float32, init_value=init_sin), - TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache, is_output=True), + TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache), TensorSpec("position_ids", [batch * S], torch.int32, init_value=lambda: init_position_ids().reshape(-1)), TensorSpec("cmp_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_cmp_slot_mapping().reshape(-1)), TensorSpec("state_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_state_slot_mapping().reshape(-1)), @@ -547,7 +547,7 @@ def init_cmp_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -575,7 +575,7 @@ def init_cmp_slot_mapping(): parser.error(f"--start-pos must contain integers, got {args.start_pos!r}") start_pos = start_values[0] if len(start_values) == 1 else start_values - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(start_pos, batch=args.batch), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_flash_dspark/decode_cp_token_allgather.py b/models/deepseek_v4_flash_dspark/decode_cp_token_allgather.py index 2dc2c38bf..94108e76e 100644 --- a/models/deepseek_v4_flash_dspark/decode_cp_token_allgather.py +++ b/models/deepseek_v4_flash_dspark/decode_cp_token_allgather.py @@ -247,7 +247,7 @@ def init_hidden_local(): return [ TensorSpec("hidden_local", [FIXTURE_ROUNDS, TP_SIZE, local_t, D], torch.bfloat16, init_value=init_hidden_local), - TensorSpec("group_out", [FIXTURE_ROUNDS, TP_SIZE, group_t, D], torch.bfloat16, is_output=True), + TensorSpec("group_out", [FIXTURE_ROUNDS, TP_SIZE, group_t, D], torch.bfloat16), ] @@ -262,7 +262,7 @@ def golden_decode_cp_token_allgather(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser(description="Standalone context-parallel decode token-row all-gather test.") @@ -286,7 +286,7 @@ def golden_decode_cp_token_allgather(tensors): if not 1 <= args.local_t <= DECODE_LOCAL_CAP: parser.error(f"--local-t must be in [1, {DECODE_LOCAL_CAP}], got {args.local_t}") - result = run_jit( + result = run( fn=l3_decode_cp_token_allgather_fixture, specs=build_tensor_specs(args.local_t), golden_fn=golden_decode_cp_token_allgather, diff --git a/models/deepseek_v4_flash_dspark/decode_csa.py b/models/deepseek_v4_flash_dspark/decode_csa.py index 07bdf15bb..935873ee6 100644 --- a/models/deepseek_v4_flash_dspark/decode_csa.py +++ b/models/deepseek_v4_flash_dspark/decode_csa.py @@ -634,8 +634,8 @@ def l3_decode_csa( gamma_cq: pl.Tensor[[TP_SIZE, Q_LORA], pl.BF16], gamma_ckv: pl.Tensor[[TP_SIZE, HEAD_DIM], pl.BF16], freqs_cos_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], - freqs_sin_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_cos: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], + freqs_sin_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], cmp_freqs_cos: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], cmp_freqs_sin: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], @@ -1575,7 +1575,7 @@ def init_wo_b(): TensorSpec("cmp_norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_cmp_norm_w), TensorSpec( "compress_state", [main_state_block_num, MAIN_STATE_BLOCK_SIZE, MAIN_STATE_DIM], - torch.float32, init_value=init_compress_state, is_output=True, + torch.float32, init_value=init_compress_state, ), TensorSpec("compress_state_block_table", [batch, MAIN_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("idx_wq_b", [Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_wq_b_i8), @@ -1588,22 +1588,22 @@ def init_wo_b(): TensorSpec("inner_norm_w", [IDX_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), TensorSpec( "inner_compress_state", [inner_state_block_num, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], - torch.float32, init_value=init_inner_compress_state, is_output=True, + torch.float32, init_value=init_inner_compress_state, ), TensorSpec("inner_compress_state_block_table", [batch, INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), TensorSpec( "cmp_kv", [cmp_block_num, BLOCK_SIZE, 1, HEAD_DIM], - torch.bfloat16, init_value=init_cmp_kv, is_output=True, + torch.bfloat16, init_value=init_cmp_kv, ), TensorSpec("cmp_block_table", [batch, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec( "idx_kv_cache", [idx_cache_block_num, BLOCK_SIZE, 1, IDX_HEAD_DIM], - torch.int8, init_value=lambda: shared_idx_kv_cache_i8.clone(), is_output=True, + torch.int8, init_value=lambda: shared_idx_kv_cache_i8.clone(), ), TensorSpec( "idx_kv_scale", [idx_cache_block_num, BLOCK_SIZE, 1, 1], - torch.float32, init_value=lambda: shared_idx_kv_scale.clone(), is_output=True, + torch.float32, init_value=lambda: shared_idx_kv_scale.clone(), ), TensorSpec("idx_block_table", [batch, IDX_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("ori_slot_mapping", [tokens], torch.int64, init_value=init_ori_slot_mapping), @@ -1619,7 +1619,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [tokens, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [tokens, HC_MULT, D], torch.float32), ] @@ -1687,7 +1687,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): if spec.name == "x_out": specs.append(TensorSpec( "x_out", [TP_SIZE, local_t, HC_MULT, D], - torch.float32, is_output=True, + torch.float32, )) continue @@ -1734,7 +1734,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): local_name = f"{spec.name}_local" if spec.name in dual_names else spec.name distributed_spec = TensorSpec( local_name, list(rank_value.shape), spec.dtype, - init_value=rank_value, is_output=spec.is_output, + init_value=rank_value, ) if spec.name in resident_names: distributed_spec.resident = "stacked" @@ -1957,7 +1957,7 @@ def mapped(pool, local_mapping, **kwargs): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser() @@ -2011,7 +2011,7 @@ def mapped(pool, local_mapping, **kwargs): local_t = batch * S if TP_SIZE == 1: - result = run_jit( + result = run( fn=decode_csa_tp1_test, specs=build_tensor_specs(start_pos=start_pos, batch=batch), golden_fn=golden_decode_csa_tp1, @@ -2042,7 +2042,7 @@ def mapped(pool, local_mapping, **kwargs): dump_passes=args.dump_passes, distributed_config=DistributedConfig(device_ids=device_ids, num_sub_workers=0), ) - result = run_jit( + result = run( fn=l3_decode_csa, specs=build_distributed_tensor_specs(local_t, start_pos=start_pos), golden_fn=golden_decode_csa, diff --git a/models/deepseek_v4_flash_dspark/decode_fwd.py b/models/deepseek_v4_flash_dspark/decode_fwd.py index 57ec0e5d2..eea035c62 100644 --- a/models/deepseek_v4_flash_dspark/decode_fwd.py +++ b/models/deepseek_v4_flash_dspark/decode_fwd.py @@ -1388,11 +1388,10 @@ def l3_decode_fwd( "hca_kv_seq_lens": "kv_seq_lens", } -def _copy_spec(name, source, *, is_output=None): +def _copy_spec(name, source): from golden import TensorSpec - output = source.is_output if is_output is None else is_output - copied = TensorSpec(name, list(source.shape), source.dtype, init_value=source.init_value, is_output=output) + copied = TensorSpec(name, list(source.shape), source.dtype, init_value=source.init_value) copied.resident = source.resident return copied @@ -1445,7 +1444,7 @@ def init_value(): packed[:, ordinal * extent : (ordinal + 1) * extent].fill_(ordinal + 1) return packed - spec = TensorSpec(name, shape, source.dtype, init_value=init_value, is_output=True) + spec = TensorSpec(name, shape, source.dtype, init_value=init_value) spec.resident = "stacked" return spec @@ -1501,7 +1500,7 @@ def init_value(source=source): spec = TensorSpec( source.name, [N_RANKS, *source.shape[1:]], source.dtype, - init_value=init_value, is_output=source.is_output, + init_value=init_value, ) spec.resident = source.resident specs[spec.name] = spec @@ -1567,30 +1566,30 @@ def init_embed_weight(): "logit_row_indices", [N_RANKS, MAX_LOGIT_ROWS], torch.int32, init_value=lambda: build_active_logit_row_indices_host(local_t), ), - "hidden_workspace": TensorSpec("hidden_workspace", [N_RANKS, local_t, D], torch.bfloat16, is_output=True), + "hidden_workspace": TensorSpec("hidden_workspace", [N_RANKS, local_t, D], torch.bfloat16), "x_ping": TensorSpec( "x_ping", [N_RANKS, local_t, HC_MULT, D], torch.float32, - init_value=zero_active, is_output=True, + init_value=zero_active, ), "x_pong": TensorSpec( "x_pong", [N_RANKS, local_t, HC_MULT, D], torch.float32, - init_value=zero_active, is_output=True, + init_value=zero_active, ), "x_attn_active": TensorSpec( "x_attn_active", [N_RANKS, local_t, HC_MULT, D], torch.float32, - init_value=zero_active, is_output=True, + init_value=zero_active, ), "x_moe_next": TensorSpec( "x_moe_next", [N_RANKS, MOE_TOKENS, HC_MULT, D], torch.float32, - init_value=lambda: torch.zeros(N_RANKS, MOE_TOKENS, HC_MULT, D, dtype=torch.float32), is_output=True, + init_value=lambda: torch.zeros(N_RANKS, MOE_TOKENS, HC_MULT, D, dtype=torch.float32), ), "pre_hc_hidden_out": TensorSpec( - "pre_hc_hidden_out", [N_RANKS, local_t, HC_MULT, D], torch.float32, is_output=True, + "pre_hc_hidden_out", [N_RANKS, local_t, HC_MULT, D], torch.float32, ), - "x_out": TensorSpec("x_out", [N_RANKS, local_t, D], torch.bfloat16, is_output=True), - "logits": TensorSpec("logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, is_output=True), + "x_out": TensorSpec("x_out", [N_RANKS, local_t, D], torch.bfloat16), + "logits": TensorSpec("logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32), "sampled_ids": TensorSpec( - "sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, is_output=True, + "sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, ), } @@ -1651,7 +1650,7 @@ def _parse_start_pos(raw): def main(): import argparse - from golden import run_jit + from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser(description="DeepSeek-V4 D-Spark decode-forward integration") @@ -1704,7 +1703,7 @@ def main(): runtime_case = None if weight_bank_size == MAIN_LAYER_COUNT else args.runtime_case specs = build_tensor_specs(start_pos=start_pos, weight_bank_size=weight_bank_size, runtime_case=runtime_case) - result = run_jit( + result = run( fn=l3_decode_fwd, specs=specs, save_data=args.save_data, diff --git a/models/deepseek_v4_flash_dspark/decode_hca.py b/models/deepseek_v4_flash_dspark/decode_hca.py index 81dda847c..0d0148463 100644 --- a/models/deepseek_v4_flash_dspark/decode_hca.py +++ b/models/deepseek_v4_flash_dspark/decode_hca.py @@ -503,8 +503,8 @@ def l3_decode_hca( gamma_cq: pl.Tensor[[TP_SIZE, Q_LORA], pl.BF16], gamma_ckv: pl.Tensor[[TP_SIZE, HEAD_DIM], pl.BF16], freqs_cos_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], - freqs_sin_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_cos: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], + freqs_sin_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], cmp_freqs_cos: pl.Tensor[[TP_SIZE, KV_B_DYN, ROPE_HEAD_DIM // 2], pl.FP32], cmp_freqs_sin: pl.Tensor[[TP_SIZE, KV_B_DYN, ROPE_HEAD_DIM // 2], pl.FP32], @@ -1320,10 +1320,10 @@ def init_wo_b(): TensorSpec("cmp_wgate", [MAIN_OUT_DIM, D], torch.bfloat16, init_value=init_cmp_wgate), TensorSpec("cmp_ape", [COMPRESS_RATIO, MAIN_OUT_DIM], torch.float32, init_value=init_cmp_ape), TensorSpec("cmp_norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_cmp_norm_w), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [batch, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), - TensorSpec("cmp_kv", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), + TensorSpec("cmp_kv", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("cmp_block_table", list(cmp_block_table.shape), torch.int32, init_value=init_cmp_block_table), TensorSpec("ori_slot_mapping", [tokens], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("window_swa_indices", [tokens, WIN], torch.int32, init_value=init_window_swa_indices), @@ -1336,7 +1336,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [tokens, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [tokens, HC_MULT, D], torch.float32), ] @@ -1386,7 +1386,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): for spec in build_tensor_specs(start_pos=start_pos, batch=group_batch): if spec.name == "x_out": specs.append(TensorSpec( - "x_out", [TP_SIZE, local_t, HC_MULT, D], torch.float32, is_output=True, + "x_out", [TP_SIZE, local_t, HC_MULT, D], torch.float32, )) continue @@ -1431,7 +1431,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): local_name = f"{spec.name}_local" if spec.name in dual_names else spec.name distributed_spec = TensorSpec( local_name, list(rank_value.shape), spec.dtype, - init_value=rank_value, is_output=spec.is_output, + init_value=rank_value, ) if spec.name in resident_names: distributed_spec.resident = "stacked" @@ -1449,7 +1449,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): if __name__ == "__main__": import argparse - from golden import mapped_pool_ratio_allclose, ratio_reldiff, run_jit + from golden import mapped_pool_ratio_allclose, ratio_reldiff, run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser() @@ -1505,7 +1505,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): for local_t in token_counts: if TP_SIZE == 1: - result = run_jit( + result = run( fn=decode_hca_tp1_test, specs=build_tensor_specs(start_pos=args.start_pos, batch=local_t // S), golden_fn=golden_decode_hca_tp1, @@ -1545,7 +1545,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): full_x_out_max_diff = 2 if args.platform == "a2a3sim" else 1 mapping_shape = (TP_SIZE, local_t) full_mapping_shape = (TP_SIZE, TP_SIZE * local_t) - result = run_jit( + result = run( fn=l3_decode_hca, specs=build_distributed_tensor_specs(local_t, start_pos=args.start_pos), golden_fn=golden_decode_hca, diff --git a/models/deepseek_v4_flash_dspark/decode_indexer.py b/models/deepseek_v4_flash_dspark/decode_indexer.py index 9d659f738..9bdff6583 100644 --- a/models/deepseek_v4_flash_dspark/decode_indexer.py +++ b/models/deepseek_v4_flash_dspark/decode_indexer.py @@ -1259,11 +1259,11 @@ def init_idx_slot_mapping(): TensorSpec("inner_wgate", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wgate), TensorSpec("inner_ape", [COMPRESS_RATIO, INNER_OUT_DIM], torch.float32, init_value=init_inner_ape), TensorSpec("inner_norm_w", [INNER_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), - TensorSpec("idx_kv_cache", [idx_physical_blocks, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_kv_i8, is_output=True), - TensorSpec("idx_kv_scale", [idx_physical_blocks, BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: idx_kv_sc, is_output=True), + TensorSpec("idx_kv_cache", [idx_physical_blocks, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_kv_i8), + TensorSpec("idx_kv_scale", [idx_physical_blocks, BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: idx_kv_sc), TensorSpec("idx_block_table", [batch, IDX_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), - TensorSpec("topk_scores", [tokens, IDX_TOPK], torch.float32, is_output=True), - TensorSpec("topk_idxs", [tokens, IDX_TOPK], torch.int32, is_output=True), + TensorSpec("topk_scores", [tokens, IDX_TOPK], torch.float32), + TensorSpec("topk_idxs", [tokens, IDX_TOPK], torch.int32), TensorSpec("position_ids", [batch * S], torch.int32, init_value=lambda: init_position_ids().reshape(-1)), TensorSpec("idx_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_idx_slot_mapping().reshape(-1)), TensorSpec("inner_state_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_inner_state_slot_mapping().reshape(-1)), @@ -1273,7 +1273,7 @@ def init_idx_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit, topk_pair_compare + from golden import ratio_allclose, run, topk_pair_compare parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1303,7 +1303,7 @@ def init_idx_slot_mapping(): ) start_pos = start_values[0] if len(start_values) == 1 else start_values - result = run_jit( + result = run( fn=indexer_test, specs=build_tensor_specs(start_pos, batch=args.batch), golden_fn=golden_indexer, diff --git a/models/deepseek_v4_flash_dspark/decode_indexer_compressor.py b/models/deepseek_v4_flash_dspark/decode_indexer_compressor.py index 48b6788b5..a1aaf44b5 100644 --- a/models/deepseek_v4_flash_dspark/decode_indexer_compressor.py +++ b/models/deepseek_v4_flash_dspark/decode_indexer_compressor.py @@ -602,8 +602,8 @@ def init_idx_slot_mapping(): return [ TensorSpec("x", [batch * S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [batch * S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [state_block_num, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [batch * S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [state_block_num, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [batch, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -612,8 +612,8 @@ def init_idx_slot_mapping(): TensorSpec("cos", [batch * S, ROPE_HEAD_DIM], torch.float32, init_value=init_cos), TensorSpec("sin", [batch * S, ROPE_HEAD_DIM], torch.float32, init_value=init_sin), TensorSpec("hadamard", [HEAD_DIM, HEAD_DIM], torch.bfloat16, init_value=init_hadamard), - TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, is_output=True), - TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, is_output=True), + TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache), + TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale), TensorSpec("position_ids", [batch * S], torch.int32, init_value=lambda: init_position_ids().reshape(-1)), TensorSpec("idx_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_idx_slot_mapping().reshape(-1)), TensorSpec("inner_state_slot_mapping", [batch * S], torch.int64, init_value=lambda: init_inner_state_slot_mapping().reshape(-1)), @@ -622,7 +622,7 @@ def init_idx_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -650,7 +650,7 @@ def init_idx_slot_mapping(): parser.error(f"--start-pos must contain integers, got {args.start_pos!r}") start_pos = start_values[0] if len(start_values) == 1 else start_values - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(start_pos, batch=args.batch), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_flash_dspark/decode_layer.py b/models/deepseek_v4_flash_dspark/decode_layer.py index 2a8e24d92..bd3dfcf73 100644 --- a/models/deepseek_v4_flash_dspark/decode_layer.py +++ b/models/deepseek_v4_flash_dspark/decode_layer.py @@ -1845,7 +1845,6 @@ def init_value(): [N_RANKS, *spec.shape[1:]], spec.dtype, init_value=init_value, - is_output=spec.is_output, ) expanded.resident = spec.resident return expanded @@ -1891,19 +1890,16 @@ def build_swa_layer_specs(start_pos=None, layer_id=0): "x_attn_active", [N_RANKS, local_t, HC_MULT, D], torch.float32, - is_output=True, ), TensorSpec( "x_moe_next", [N_RANKS, MOE_TOKENS, HC_MULT, D], torch.float32, - is_output=True, ), TensorSpec( "x_next", [N_RANKS, local_t, HC_MULT, D], torch.float32, - is_output=True, ), ScalarSpec("layer_id", torch.int32, layer_id), ScalarSpec("local_t", torch.int32, local_t), @@ -1965,7 +1961,6 @@ def init_value(): [N_RANKS, *spec.shape[1:]], spec.dtype, init_value=init_value, - is_output=spec.is_output, ) expanded.resident = spec.resident return expanded @@ -2011,19 +2006,16 @@ def build_hca_layer_specs(start_pos=None, layer_id=3): "x_attn_active", [N_RANKS, local_t, HC_MULT, D], torch.float32, - is_output=True, ), TensorSpec( "x_moe_next", [N_RANKS, MOE_TOKENS, HC_MULT, D], torch.float32, - is_output=True, ), TensorSpec( "x_next", [N_RANKS, local_t, HC_MULT, D], torch.float32, - is_output=True, ), ScalarSpec("layer_id", torch.int32, layer_id), ScalarSpec("local_t", torch.int32, local_t), @@ -2085,7 +2077,6 @@ def init_value(): [N_RANKS, *spec.shape[1:]], spec.dtype, init_value=init_value, - is_output=spec.is_output, ) expanded.resident = spec.resident return expanded @@ -2131,19 +2122,16 @@ def build_csa_layer_specs(start_pos=None, layer_id=2): "x_attn_active", [N_RANKS, local_t, HC_MULT, D], torch.float32, - is_output=True, ), TensorSpec( "x_moe_next", [N_RANKS, MOE_TOKENS, HC_MULT, D], torch.float32, - is_output=True, ), TensorSpec( "x_next", [N_RANKS, local_t, HC_MULT, D], torch.float32, - is_output=True, ), ScalarSpec("layer_id", torch.int32, layer_id), ScalarSpec("local_t", torch.int32, local_t), @@ -2207,7 +2195,7 @@ def main(): from golden import ( mapped_pool_ratio_allclose, ratio_reldiff, - run_jit, + run, ) from pypto.ir.distributed_compiled_program import DistributedConfig @@ -2417,7 +2405,7 @@ def main(): "x_next": ratio_reldiff(diff_thd=0.01, pct_thd=0.05), } - result = run_jit( + result = run( fn=layer_fn, specs=specs, golden_fn=golden_fn, diff --git a/models/deepseek_v4_flash_dspark/decode_metadata.py b/models/deepseek_v4_flash_dspark/decode_metadata.py index 4163429f5..fd9dbb556 100644 --- a/models/deepseek_v4_flash_dspark/decode_metadata.py +++ b/models/deepseek_v4_flash_dspark/decode_metadata.py @@ -466,14 +466,14 @@ def build_tensor_specs(): ("csa_state_slot_mapping", [T], torch.int64), ("csa_inner_state_slot_mapping", [T], torch.int64), ): - specs.append(TensorSpec(name, shape, dtype, is_output=True)) + specs.append(TensorSpec(name, shape, dtype)) return specs if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( @@ -486,7 +486,7 @@ def build_tensor_specs(): parser.add_argument("--compile-only", action="store_true") args = parser.parse_args() - result = run_jit( + result = run( fn=decode_metadata, specs=build_tensor_specs(), golden_fn=golden_decode_metadata, diff --git a/models/deepseek_v4_flash_dspark/decode_o_proj.py b/models/deepseek_v4_flash_dspark/decode_o_proj.py index 8c62ac289..8d4cbdefb 100644 --- a/models/deepseek_v4_flash_dspark/decode_o_proj.py +++ b/models/deepseek_v4_flash_dspark/decode_o_proj.py @@ -430,7 +430,7 @@ def init_attention_grouped(): TensorSpec("attention_grouped", attention_grouped_shape, torch.bfloat16, init_value=init_attention_grouped), TensorSpec( "attention_local_groups", attention_local_shape, torch.bfloat16, - init_value=FIXTURE_OUTPUT_SENTINEL, is_output=True, + init_value=FIXTURE_OUTPUT_SENTINEL, ), ScalarSpec("local_t", torch.int32, local_t), ] @@ -721,7 +721,7 @@ def golden_decode_o_proj_tp1(o_packed_heads, wo_a, wo_b, wo_b_scale, tokens): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser() @@ -751,7 +751,7 @@ def golden_decode_o_proj_tp1(o_packed_heads, wo_a, wo_b, wo_b_scale, tokens): if not 1 <= a2a_local_t <= LOCAL_T: parser.error(f"--local-t must be in [1, {LOCAL_T}], got {a2a_local_t}") - result = run_jit( + result = run( fn=l3_o_group_a2a, specs=build_o_group_a2a_specs(a2a_local_t), golden_fn=golden_o_group_a2a, diff --git a/models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py b/models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py index 19ed9284d..5a32c4930 100644 --- a/models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py +++ b/models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py @@ -735,13 +735,13 @@ def init_sin(): TensorSpec("attn_sink", [H], torch.float32, init_value=init_attn_sink), TensorSpec("freqs_cos", [tokens, ROPE_DIM], torch.bfloat16, init_value=init_cos), TensorSpec("freqs_sin", [tokens, ROPE_DIM], torch.bfloat16, init_value=init_sin), - TensorSpec("o_packed_heads", [O_GROUPS, T_PAD, O_GROUP_IN], torch.bfloat16, is_output=True), + TensorSpec("o_packed_heads", [O_GROUPS, T_PAD, O_GROUP_IN], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -784,7 +784,7 @@ def init_sin(): print(f"compress_ratio={COMPRESS_RATIO} -> TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}", flush=True) - result = run_jit( + result = run( fn=sparse_attn_csa_test, specs=build_tensor_specs( args.causal_regression_fixture, diff --git a/models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py b/models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py index da7ff34f6..29e7dee17 100644 --- a/models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py +++ b/models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py @@ -855,13 +855,13 @@ def init_sin(): TensorSpec("attn_sink", [H], torch.float32, init_value=init_attn_sink), TensorSpec("freqs_cos", [tokens, ROPE_DIM], torch.bfloat16, init_value=init_cos), TensorSpec("freqs_sin", [tokens, ROPE_DIM], torch.bfloat16, init_value=init_sin), - TensorSpec("o_packed_heads", [O_GROUPS, T_PAD, O_GROUP_IN], torch.bfloat16, is_output=True), + TensorSpec("o_packed_heads", [O_GROUPS, T_PAD, O_GROUP_IN], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -900,7 +900,7 @@ def init_sin(): ) print(f"compress_ratio={COMPRESS_RATIO} {workload}", flush=True) - result = run_jit( + result = run( fn=sparse_attn_hca_test, specs=build_tensor_specs( args.causal_regression_fixture, diff --git a/models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py b/models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py index 895fba120..827018618 100644 --- a/models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py +++ b/models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py @@ -563,13 +563,13 @@ def init_sin(): TensorSpec("attn_sink", [H], torch.float32, init_value=init_attn_sink), TensorSpec("freqs_cos", [tokens, ROPE_DIM], torch.bfloat16, init_value=init_cos), TensorSpec("freqs_sin", [tokens, ROPE_DIM], torch.bfloat16, init_value=init_sin), - TensorSpec("o_packed_heads", [O_GROUPS, T_PAD * HEADS_PER_GROUP, HEAD_DIM], torch.bfloat16, is_output=True), + TensorSpec("o_packed_heads", [O_GROUPS, T_PAD * HEADS_PER_GROUP, HEAD_DIM], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -595,7 +595,7 @@ def init_sin(): print(f"TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}", flush=True) - result = run_jit( + result = run( fn=sparse_attn_swa_test, specs=build_tensor_specs( args.causal_regression_fixture, diff --git a/models/deepseek_v4_flash_dspark/decode_swa.py b/models/deepseek_v4_flash_dspark/decode_swa.py index 2b63478b1..71abd54ad 100644 --- a/models/deepseek_v4_flash_dspark/decode_swa.py +++ b/models/deepseek_v4_flash_dspark/decode_swa.py @@ -450,8 +450,8 @@ def l3_decode_swa( gamma_cq: pl.Tensor[[TP_SIZE, Q_LORA], pl.BF16], gamma_ckv: pl.Tensor[[TP_SIZE, HEAD_DIM], pl.BF16], freqs_cos_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], - freqs_sin_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_cos: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], + freqs_sin_local: pl.Tensor[[TP_SIZE, T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[TP_SIZE, KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], kv_cache: pl.InOut[pl.Tensor[[TP_SIZE, ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], swa_slot_mapping: pl.Tensor[[TP_SIZE, KV_T_DYN], pl.INT64], @@ -920,7 +920,7 @@ def init_wo_b(): TensorSpec("gamma_ckv", [HEAD_DIM], torch.bfloat16, init_value=init_gamma_ckv), TensorSpec("freqs_cos", [tokens, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [tokens, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), TensorSpec("swa_slot_mapping", [tokens], torch.int64, init_value=init_swa_slot_mapping), TensorSpec("swa_indices", [tokens, WIN], torch.int32, init_value=init_swa_indices), TensorSpec("swa_lens", [tokens], torch.int32, init_value=init_swa_lens), @@ -929,7 +929,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [tokens, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [tokens, HC_MULT, D], torch.float32), ] @@ -971,7 +971,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): for spec in build_tensor_specs(start_pos=start_pos, batch=group_batch): if spec.name == "x_out": specs.append(TensorSpec( - "x_out", [TP_SIZE, local_t, HC_MULT, D], torch.float32, is_output=True, + "x_out", [TP_SIZE, local_t, HC_MULT, D], torch.float32, )) continue @@ -1013,7 +1013,7 @@ def build_distributed_tensor_specs(local_t, start_pos=None): local_name = f"{spec.name}_local" if spec.name in dual_names else spec.name distributed_spec = TensorSpec( local_name, list(rank_value.shape), spec.dtype, - init_value=rank_value, is_output=spec.is_output, + init_value=rank_value, ) if spec.name in resident_names: distributed_spec.resident = "stacked" @@ -1096,7 +1096,7 @@ def golden_decode_swa(tensors): if __name__ == "__main__": import argparse - from golden import mapped_pool_ratio_allclose, ratio_reldiff, run_jit + from golden import mapped_pool_ratio_allclose, ratio_reldiff, run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser() @@ -1148,7 +1148,7 @@ def golden_decode_swa(tensors): for local_t in token_counts: if TP_SIZE == 1: - result = run_jit( + result = run( fn=decode_swa_tp1_test, specs=build_tensor_specs(start_pos=args.start_pos, batch=local_t // S), golden_fn=golden_decode_swa_tp1, @@ -1177,7 +1177,7 @@ def golden_decode_swa(tensors): }, ) else: - result = run_jit( + result = run( fn=l3_decode_swa, specs=build_distributed_tensor_specs(local_t, start_pos=args.start_pos), golden_fn=golden_decode_swa, diff --git a/models/deepseek_v4_flash_dspark/dspark_attention.py b/models/deepseek_v4_flash_dspark/dspark_attention.py index 22bbebcc5..ee367fc92 100644 --- a/models/deepseek_v4_flash_dspark/dspark_attention.py +++ b/models/deepseek_v4_flash_dspark/dspark_attention.py @@ -481,7 +481,6 @@ def init_wkv(): [KV_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, - is_output=True, ), TensorSpec("slot_mapping", [T], torch.int64, init_value=init_slot_mapping), TensorSpec("swa_indices", [B, INDEX_WIDTH], torch.int32, init_value=init_swa_indices), @@ -491,14 +490,13 @@ def init_wkv(): "o_packed_heads", [O_GROUPS, LOCAL_T_PAD * HEADS_PER_GROUP, HEAD_DIM], torch.bfloat16, - is_output=True, ), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="DeepSeek-V4 DSpark drafter attention validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -508,7 +506,7 @@ def init_wkv(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=dspark_attention_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_dspark_attention, diff --git a/models/deepseek_v4_flash_dspark/dspark_context_kv.py b/models/deepseek_v4_flash_dspark/dspark_context_kv.py index 0a1b9982e..9ff59bd2d 100644 --- a/models/deepseek_v4_flash_dspark/dspark_context_kv.py +++ b/models/deepseek_v4_flash_dspark/dspark_context_kv.py @@ -263,14 +263,13 @@ def init_kv_cache(): [KV_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, - is_output=True, ), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="DeepSeek-V4 DSpark drafter context-KV validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -286,7 +285,7 @@ def init_kv_cache(): for mode in (modes if args.mode == "all" else [args.mode]): batch, seq = modes[mode] print(f"--- dspark_context_kv_test {mode}: T={batch * seq} ---") - result = run_jit( + result = run( fn=dspark_context_kv_test, specs=build_tensor_specs(batch, seq), golden_fn=golden_dspark_context_kv, diff --git a/models/deepseek_v4_flash_dspark/dspark_drafter.py b/models/deepseek_v4_flash_dspark/dspark_drafter.py index 33724d53e..286673151 100644 --- a/models/deepseek_v4_flash_dspark/dspark_drafter.py +++ b/models/deepseek_v4_flash_dspark/dspark_drafter.py @@ -1172,8 +1172,8 @@ def init_wkv(): weight[:, layer * D + diagonal, diagonal] = 1 return weight - def ranked(name, shape, dtype, init_value=0, *, output=False, resident=False): - spec = TensorSpec(name, [N_RANKS, *shape], dtype, init_value=init_value, is_output=output) + def ranked(name, shape, dtype, init_value=0, *, resident=False): + spec = TensorSpec(name, [N_RANKS, *shape], dtype, init_value=init_value) if resident: spec.resident = "stacked" return spec @@ -1201,13 +1201,11 @@ def init_kv_caches(): "initial_hidden", [N_RANKS, DSPARK_MAX_BATCH * DSPARK_QUERY_PAD, HC_MULT, D], torch.float32, - is_output=True, ), TensorSpec( "intermediate_hidden", [N_RANKS, DSPARK_DRAFT_LAYERS, DSPARK_MAX_BATCH * DSPARK_QUERY_PAD, HC_MULT, D], torch.float32, - is_output=True, ), ranked("main_proj_weight", [D, MAIN_IN], torch.bfloat16, init_value=init_main_proj_weight, resident=True), ranked("main_norm_weight", [D], torch.bfloat16, init_value=1, resident=True), @@ -1267,7 +1265,6 @@ def init_kv_caches(): [DSPARK_DRAFT_LAYERS, ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_caches, - output=True, ), ranked("attn_sink", [DSPARK_DRAFT_LAYERS * H], torch.float32, resident=True), ranked("wo_a", [DSPARK_DRAFT_LAYERS * LOCAL_O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, resident=True), @@ -1310,7 +1307,6 @@ def init_kv_caches(): "head_hidden", [N_RANKS, batch, DSPARK_QUERY_WIDTH, D], torch.bfloat16, - is_output=True, ), ] ) @@ -1486,7 +1482,7 @@ def hc_post_zero_output(residual, post, combine): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description="Validate the multi-rank DeepSeek V4 DSpark drafter.") parser.add_argument("--batch", type=int, choices=DSPARK_SUPPORTED_BATCHES, default=4) @@ -1502,7 +1498,7 @@ def hc_post_zero_output(residual, post, combine): assert args.tp == TP_SIZE assert args.ep == N_RANKS assert len(device_ids) >= N_RANKS - result = run_jit( + result = run( fn=l3_dspark_drafter, specs=build_tensor_specs(args.batch), golden_fn=golden_dspark_drafter, diff --git a/models/deepseek_v4_flash_dspark/dspark_markov.py b/models/deepseek_v4_flash_dspark/dspark_markov.py index 0b33ff071..2e0f48eac 100644 --- a/models/deepseek_v4_flash_dspark/dspark_markov.py +++ b/models/deepseek_v4_flash_dspark/dspark_markov.py @@ -876,13 +876,11 @@ def init_logit_row_indices(): "draft_token_ids", with_world(batch, DSPARK_QUERY_WIDTH), torch.int32, - is_output=True, ), TensorSpec( "confidence_probs", with_world(batch, DSPARK_QUERY_WIDTH), torch.float32, - is_output=True, ), ]) return specs @@ -959,7 +957,7 @@ def golden_distributed_markov(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description="Validate the DeepSeek V4 DSpark Markov sampler.") parser.add_argument("--batch", type=int, choices=DSPARK_SUPPORTED_BATCHES, default=4) @@ -990,7 +988,7 @@ def golden_distributed_markov(tensors): else: runtime_cfg["device_id"] = int(args.device) - result = run_jit( + result = run( fn=fn, specs=build_tensor_specs(args.batch, distributed=args.distributed), golden_fn=golden_fn, diff --git a/models/deepseek_v4_flash_dspark/dspark_prefill.py b/models/deepseek_v4_flash_dspark/dspark_prefill.py index c0dd26af6..b3b2cc129 100644 --- a/models/deepseek_v4_flash_dspark/dspark_prefill.py +++ b/models/deepseek_v4_flash_dspark/dspark_prefill.py @@ -26,7 +26,7 @@ if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description="Validate DeepSeek V4 DSpark prompt prefill and drafting.") parser.add_argument("--batch", type=int, choices=DSPARK_SUPPORTED_BATCHES, default=4) @@ -42,7 +42,7 @@ assert args.tp == TP_SIZE assert args.ep == N_RANKS assert len(device_ids) >= N_RANKS - result = run_jit( + result = run( fn=l3_dspark_drafter, specs=build_tensor_specs(args.batch, mode="prefill"), golden_fn=golden_dspark_drafter, diff --git a/models/deepseek_v4_flash_dspark/dspark_proj.py b/models/deepseek_v4_flash_dspark/dspark_proj.py index 41f640b33..2a1d3e26f 100644 --- a/models/deepseek_v4_flash_dspark/dspark_proj.py +++ b/models/deepseek_v4_flash_dspark/dspark_proj.py @@ -102,13 +102,13 @@ def init_main_proj_w(): TensorSpec("main_hidden", [t, MAIN_HIDDEN_DIM], torch.bfloat16, init_value=init_main_hidden), TensorSpec("main_proj_w", [D, MAIN_HIDDEN_DIM], torch.bfloat16, init_value=init_main_proj_w), TensorSpec("main_norm_w", [D], torch.bfloat16, init_value=lambda: torch.randn(D) * 0.1 + 1.0), - TensorSpec("main_x", [t, D], torch.bfloat16, is_output=True), + TensorSpec("main_x", [t, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="DeepSeek-V4 DSpark main-hidden projection validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -125,7 +125,7 @@ def init_main_proj_w(): for mode in (modes if args.mode == "all" else [args.mode]): batch, seq = modes[mode] print(f"--- dspark_proj_test {mode}: T={batch * seq} ---") - result = run_jit( + result = run( fn=dspark_proj_test, specs=build_tensor_specs(batch, seq), golden_fn=golden_dspark_proj, diff --git a/models/deepseek_v4_flash_dspark/expert_routed.py b/models/deepseek_v4_flash_dspark/expert_routed.py index 28bae8fa6..1c8a03fea 100644 --- a/models/deepseek_v4_flash_dspark/expert_routed.py +++ b/models/deepseek_v4_flash_dspark/expert_routed.py @@ -448,13 +448,13 @@ def init_recv_weights(): TensorSpec("routed_w3_scale", [N_LOCAL_EXPERTS, MOE_INTER], torch.float32, init_value=lambda: w3_s), TensorSpec("routed_w2", [N_LOCAL_EXPERTS, D, MOE_INTER], torch.int8, init_value=lambda: w2_i8), TensorSpec("routed_w2_scale", [N_LOCAL_EXPERTS, D], torch.float32, init_value=lambda: w2_s), - TensorSpec("recv_y", [N_LOCAL_EXPERTS, RECV_MAX, D], torch.bfloat16, is_output=True), + TensorSpec("recv_y", [N_LOCAL_EXPERTS, RECV_MAX, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -464,7 +464,7 @@ def init_recv_weights(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=expert_routed_test, specs=build_tensor_specs(), golden_fn=golden_expert_routed, diff --git a/models/deepseek_v4_flash_dspark/expert_shared.py b/models/deepseek_v4_flash_dspark/expert_shared.py index 79925075c..92a23b7d3 100644 --- a/models/deepseek_v4_flash_dspark/expert_shared.py +++ b/models/deepseek_v4_flash_dspark/expert_shared.py @@ -388,13 +388,13 @@ def build_tensor_specs(): TensorSpec("shared_w3_scale", [MOE_INTER], torch.float32, init_value=lambda: sw3_s), TensorSpec("shared_w2", [D, MOE_INTER], torch.int8, init_value=lambda: sw2_i8), TensorSpec("shared_w2_scale", [D], torch.float32, init_value=lambda: sw2_s), - TensorSpec("sh", [T, D], torch.bfloat16, is_output=True), + TensorSpec("sh", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -404,7 +404,7 @@ def build_tensor_specs(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=expert_shared_test, specs=build_tensor_specs(), golden_fn=golden_expert_shared, diff --git a/models/deepseek_v4_flash_dspark/gate.py b/models/deepseek_v4_flash_dspark/gate.py index ee2ca92ca..127f7dc7e 100644 --- a/models/deepseek_v4_flash_dspark/gate.py +++ b/models/deepseek_v4_flash_dspark/gate.py @@ -396,10 +396,10 @@ def init_input_ids(): ScalarSpec("num_tokens", torch.int32, num_tokens), TensorSpec("tid2eid", [VOCAB, TOPK], torch.int32, init_value=init_tid2eid), TensorSpec("input_ids", [T], torch.int64, init_value=init_input_ids), - TensorSpec("x_norm_i8", [T, D], torch.int8, is_output=True), - TensorSpec("x_norm_scale", [T, 1], torch.float32, is_output=True), - TensorSpec("indices", [T, TOPK], torch.int32, is_output=True), - TensorSpec("weights", [T, TOPK], torch.float32, is_output=True), + TensorSpec("x_norm_i8", [T, D], torch.int8), + TensorSpec("x_norm_scale", [T, 1], torch.float32), + TensorSpec("indices", [T, TOPK], torch.int32), + TensorSpec("weights", [T, TOPK], torch.float32), ] @@ -411,7 +411,7 @@ def gate_active_rows(num_tokens): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit, topk_pair_compare + from golden import ratio_allclose, run, topk_pair_compare parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -423,7 +423,7 @@ def gate_active_rows(num_tokens): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=gate_test, specs=build_tensor_specs(layer_id=args.layer_id, num_tokens=args.num_tokens), golden_fn=golden_gate_core, diff --git a/models/deepseek_v4_flash_dspark/hc_head.py b/models/deepseek_v4_flash_dspark/hc_head.py index a3516874a..2892f1390 100644 --- a/models/deepseek_v4_flash_dspark/hc_head.py +++ b/models/deepseek_v4_flash_dspark/hc_head.py @@ -332,14 +332,14 @@ def init_hc_head_fn(): torch.float32, init_value=lambda: torch.tensor([5.9166, -3.6223, -2.9324, -3.3124]), ), - TensorSpec("y", [token_count, D], torch.bfloat16, is_output=True), + TensorSpec("y", [token_count, D], torch.bfloat16), ] if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument( @@ -356,7 +356,7 @@ def init_hc_head_fn(): args = parser.parse_args() torch.manual_seed(args.seed) - result = run_jit( + result = run( fn=hc_head_test, specs=build_tensor_specs(args.token_count), golden_fn=golden_hc_head, diff --git a/models/deepseek_v4_flash_dspark/hc_post.py b/models/deepseek_v4_flash_dspark/hc_post.py index 768a699a5..ccf793bcb 100644 --- a/models/deepseek_v4_flash_dspark/hc_post.py +++ b/models/deepseek_v4_flash_dspark/hc_post.py @@ -190,13 +190,13 @@ def init_comb(): TensorSpec("residual", [T, HC_MULT, D], torch.float32, init_value=init_residual), TensorSpec("post", [T, HC_MULT], torch.float32, init_value=init_post), TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32, init_value=init_comb), - TensorSpec("y", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("y", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run MODES = { "decode": (DECODE_BATCH // TP, DECODE_SEQ), @@ -219,7 +219,7 @@ def init_comb(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- hc_post {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=hc_post_test, specs=build_tensor_specs(B, S), golden_fn=golden_hc_post, diff --git a/models/deepseek_v4_flash_dspark/hc_pre.py b/models/deepseek_v4_flash_dspark/hc_pre.py index b16059612..885bee255 100644 --- a/models/deepseek_v4_flash_dspark/hc_pre.py +++ b/models/deepseek_v4_flash_dspark/hc_pre.py @@ -409,15 +409,15 @@ def init_hc_base(): TensorSpec("hc_fn", [MIX_HC, HC_DIM], torch.float32, init_value=init_hc_fn), TensorSpec("hc_scale", [3], torch.float32, init_value=init_hc_scale), TensorSpec("hc_base", [MIX_HC], torch.float32, init_value=init_hc_base), - TensorSpec("x_mixed", [T, D], torch.bfloat16, is_output=True), - TensorSpec("post", [T, HC_MULT], torch.float32, is_output=True), - TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32, is_output=True), + TensorSpec("x_mixed", [T, D], torch.bfloat16), + TensorSpec("post", [T, HC_MULT], torch.float32), + TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH // TP, DECODE_SEQ), @@ -440,7 +440,7 @@ def init_hc_base(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- hc_pre {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=hc_pre_test, specs=build_tensor_specs(B, S), golden_fn=golden_hc_pre, diff --git a/models/deepseek_v4_flash_dspark/lm_head.py b/models/deepseek_v4_flash_dspark/lm_head.py index d87a0be8a..9dda6106a 100644 --- a/models/deepseek_v4_flash_dspark/lm_head.py +++ b/models/deepseek_v4_flash_dspark/lm_head.py @@ -510,14 +510,14 @@ def init_logit_row_indices(): "lm_head_weight", [WORLD_SIZE, VOCAB_PER_TP, D], torch.bfloat16, init_value=init_lm_head_weight, resident="stacked", ), - TensorSpec("logits", [WORLD_SIZE, MAX_LOGIT_ROWS, VOCAB], torch.float32, is_output=True), + TensorSpec("logits", [WORLD_SIZE, MAX_LOGIT_ROWS, VOCAB], torch.float32), TensorSpec("logit_row_indices", [WORLD_SIZE, MAX_LOGIT_ROWS], torch.int32, init_value=init_logit_row_indices), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -544,7 +544,7 @@ def init_logit_row_indices(): specs = build_tensor_specs(args.num_tokens) golden_fn = golden_lm_head - result = run_jit( + result = run( fn=fn, specs=specs, golden_fn=golden_fn, diff --git a/models/deepseek_v4_flash_dspark/lookup_embedding.py b/models/deepseek_v4_flash_dspark/lookup_embedding.py index 34f5b951f..94e040765 100644 --- a/models/deepseek_v4_flash_dspark/lookup_embedding.py +++ b/models/deepseek_v4_flash_dspark/lookup_embedding.py @@ -94,14 +94,14 @@ def init_embed_weight(): return [ TensorSpec("input_ids", [token_count], torch.int64, init_value=init_input_ids), TensorSpec("embed_weight", [vocab_size, D], torch.bfloat16, init_value=init_embed_weight), - TensorSpec("hidden_states", [token_count, D], torch.bfloat16, is_output=True), - TensorSpec("x_hc", [token_count, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("hidden_states", [token_count, D], torch.bfloat16), + TensorSpec("x_hc", [token_count, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run MODES = {"decode": DECODE_TOKENS, "prefill": PREFILL_TOKENS} TEST_VOCAB_SIZE = 256 @@ -116,7 +116,7 @@ def init_embed_weight(): for mode_name in modes_to_run: token_count = MODES[mode_name] print(f"--- lookup_embedding_test {mode_name}: T={token_count} ---") - result = run_jit( + result = run( fn=lookup_embedding_test, specs=build_tensor_specs(token_count, TEST_VOCAB_SIZE), golden_fn=golden_lookup_embedding_test, diff --git a/models/deepseek_v4_flash_dspark/markov_head.py b/models/deepseek_v4_flash_dspark/markov_head.py index f0bf95b9a..dc18d37ca 100644 --- a/models/deepseek_v4_flash_dspark/markov_head.py +++ b/models/deepseek_v4_flash_dspark/markov_head.py @@ -124,14 +124,14 @@ def init_markov_w2(): TensorSpec("token_ids", [token_count], torch.int64, init_value=init_token_ids), TensorSpec("markov_w1", [vocab_size, MARKOV_RANK], torch.bfloat16, init_value=init_markov_w1), TensorSpec("markov_w2", [vocab_size, MARKOV_RANK], torch.bfloat16, init_value=init_markov_w2), - TensorSpec("logits_bias", [token_count, vocab_size], torch.float32, is_output=True), - TensorSpec("markov_embed", [token_count, MARKOV_RANK], torch.bfloat16, is_output=True), + TensorSpec("logits_bias", [token_count, vocab_size], torch.float32), + TensorSpec("markov_embed", [token_count, MARKOV_RANK], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run TEST_VOCAB_SIZE = 4096 @@ -144,7 +144,7 @@ def init_markov_w2(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=markov_head_test, specs=build_tensor_specs(args.token_count, args.vocab_size), golden_fn=golden_markov_head, diff --git a/models/deepseek_v4_flash_dspark/moe.py b/models/deepseek_v4_flash_dspark/moe.py index 2966875d0..009f14dfb 100644 --- a/models/deepseek_v4_flash_dspark/moe.py +++ b/models/deepseek_v4_flash_dspark/moe.py @@ -1202,7 +1202,7 @@ def init_input_ids(): TensorSpec("shared_w3_scale", [N_RANKS, MOE_INTER], torch.float32, init_value=lambda: sw3_s), TensorSpec("shared_w2", [N_RANKS, D, MOE_INTER], torch.int8, init_value=lambda: sw2_i8), TensorSpec("shared_w2_scale", [N_RANKS, D], torch.float32, init_value=lambda: sw2_s), - TensorSpec("x_next", x_hc_shape, torch.float32, is_output=True), + TensorSpec("x_next", x_hc_shape, torch.float32), ScalarSpec("layer_id", torch.int32, layer_id), ScalarSpec("num_tokens", torch.int32, num_tokens), ] @@ -1214,7 +1214,7 @@ def init_input_ids(): # H2D/D2H. Covers the routed/shared expert weights and their scales, the gate, # the HC-FFN constants, the RMSNorm gamma, and the static tid2eid route table — # but NOT the per-step activation (x_hc), per-step input_ids, or the output. - # All resident names are inputs (is_output=False), so the flag is always valid. + # All resident names are pure inputs, so the flag is always valid. RESIDENT_WEIGHT_NAMES = frozenset([ "hc_ffn_fn", "hc_ffn_scale", "hc_ffn_base", "norm_w", "gate_w", "gate_bias", "tid2eid", @@ -1242,7 +1242,7 @@ def build_rounds_tensor_specs(layer_id=0, num_tokens=T, balanced_routing=False): if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1273,7 +1273,7 @@ def build_rounds_tensor_specs(layer_id=0, num_tokens=T, balanced_routing=False): golden_data = args.golden_data - result = run_jit( + result = run( fn=l3_moe, specs=build_rounds_tensor_specs( layer_id=args.layer_id, diff --git a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py index f46d4d62d..7e20a09ab 100644 --- a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py +++ b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py @@ -688,7 +688,6 @@ def init_state_slot_mapping(): [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, - is_output=True, ), TensorSpec( "compress_state_block_table", @@ -707,7 +706,6 @@ def init_state_slot_mapping(): [CMP_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), TensorSpec("cmp_slot_mapping", [token_count], torch.int64, init_value=init_cmp_slot_mapping), @@ -717,7 +715,7 @@ def init_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser( description="Standalone token-major DeepSeek V4 prefill compressor ratio128 validation." @@ -750,7 +748,7 @@ def init_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_compressor_ratio128_test, specs=build_tensor_specs(args.start_pos, args.token_count), golden_fn=golden_prefill_compressor_ratio128, diff --git a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py index c753e0396..941747b89 100644 --- a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py +++ b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py @@ -813,7 +813,6 @@ def init_state_slot_mapping(): [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, - is_output=True, ), TensorSpec( "compress_state_block_table", @@ -832,7 +831,6 @@ def init_state_slot_mapping(): [CMP_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), TensorSpec("cmp_slot_mapping", [token_count], torch.int64, init_value=init_cmp_slot_mapping), @@ -842,7 +840,7 @@ def init_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser( description="Standalone physical-dynamic DeepSeek V4 prefill compressor ratio4 validation." @@ -879,7 +877,7 @@ def init_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_compressor_ratio4_test, specs=build_tensor_specs(args.start_pos, args.token_count), golden_fn=golden_prefill_compressor_ratio4, diff --git a/models/deepseek_v4_flash_dspark/prefill_cp_token_allgather.py b/models/deepseek_v4_flash_dspark/prefill_cp_token_allgather.py index 64549ffcc..7472d31e7 100644 --- a/models/deepseek_v4_flash_dspark/prefill_cp_token_allgather.py +++ b/models/deepseek_v4_flash_dspark/prefill_cp_token_allgather.py @@ -228,7 +228,7 @@ def init_hidden_local(): return [ TensorSpec("hidden_local", [FIXTURE_ROUNDS, TP_SIZE, local_t, D], torch.bfloat16, init_value=init_hidden_local), - TensorSpec("group_out", [FIXTURE_ROUNDS, TP_SIZE, group_t, D], torch.bfloat16, is_output=True), + TensorSpec("group_out", [FIXTURE_ROUNDS, TP_SIZE, group_t, D], torch.bfloat16), ] @@ -243,7 +243,7 @@ def golden_prefill_cp_token_allgather(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig parser = argparse.ArgumentParser(description="Standalone context-parallel prefill token-row all-gather test.") @@ -267,7 +267,7 @@ def golden_prefill_cp_token_allgather(tensors): if not 1 <= args.local_t <= PREFILL_LOCAL_CAP: parser.error(f"--local-t must be in [1, {PREFILL_LOCAL_CAP}], got {args.local_t}") - result = run_jit( + result = run( fn=l3_prefill_cp_token_allgather_fixture, specs=build_tensor_specs(args.local_t), golden_fn=golden_prefill_cp_token_allgather, diff --git a/models/deepseek_v4_flash_dspark/prefill_csa.py b/models/deepseek_v4_flash_dspark/prefill_csa.py index 540f2bed0..00d22ec32 100644 --- a/models/deepseek_v4_flash_dspark/prefill_csa.py +++ b/models/deepseek_v4_flash_dspark/prefill_csa.py @@ -1071,7 +1071,6 @@ def init_wo_b(): [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, - is_output=True, ), TensorSpec( "compress_state_block_table", @@ -1103,7 +1102,6 @@ def init_wo_b(): [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], torch.float32, init_value=init_inner_compress_state, - is_output=True, ), TensorSpec( "inner_compress_state_block_table", @@ -1116,7 +1114,6 @@ def init_wo_b(): [CSA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, - is_output=True, ), TensorSpec("ori_block_table", [1, SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), TensorSpec("ori_slot_mapping", [token_count], torch.int64, init_value=init_ori_slot_mapping), @@ -1125,7 +1122,6 @@ def init_wo_b(): [CSA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("cmp_block_table", [1, SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec( @@ -1133,14 +1129,12 @@ def init_wo_b(): [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, - is_output=True, ), TensorSpec( "idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, - is_output=True, ), TensorSpec("idx_block_table", [1, IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), @@ -1158,7 +1152,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [token_count, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [token_count, HC_MULT, D], torch.float32), ] @@ -1841,11 +1835,11 @@ def build_cp_tensor_specs( init_value=torch.stack(shards).contiguous(), )) elif spec.name == "x_out": - specs.append(TensorSpec("x_out_full", [tp_size, token_count, HC_MULT, D], spec.dtype, is_output=True)) + specs.append(TensorSpec("x_out_full", [tp_size, token_count, HC_MULT, D], spec.dtype)) else: specs.append(TensorSpec( spec.name, [tp_size, *spec.shape], spec.dtype, - init_value=cp_stack(value, tp_size), is_output=spec.is_output, + init_value=cp_stack(value, tp_size), )) return specs @@ -1996,7 +1990,7 @@ def build_ragged2_cp_tensor_specs(tp_size: int = TP_SIZE): continue replacement_spec = TensorSpec( spec.name, list(value.shape), spec.dtype, init_value=value, - is_output=spec.is_output, resident=spec.resident, + resident=spec.resident, ) specs.append(replacement_spec) return specs @@ -2043,7 +2037,7 @@ def golden_prefill_attention_csa_cp(tensors): if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser( description="Standalone DeepSeek V4 packed prefill CSA correctness test." @@ -2102,7 +2096,7 @@ def golden_prefill_attention_csa_cp(tensors): parser.error(f"--token-count must be a multiple of --tp={TP_SIZE}, got {args.token_count}") if TP_SIZE == 1: - result = run_jit( + result = run( fn=prefill_attention_csa_test, specs=build_tensor_specs(args.start_pos, args.token_count), golden_fn=golden_prefill_attention_csa, @@ -2135,7 +2129,7 @@ def golden_prefill_attention_csa_cp(tensors): if args.case == "ragged2" else build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE) ) - result = run_jit( + result = run( fn=l3_prefill_attention_csa_cp, specs=specs, golden_fn=golden_prefill_attention_csa_cp, diff --git a/models/deepseek_v4_flash_dspark/prefill_fwd.py b/models/deepseek_v4_flash_dspark/prefill_fwd.py index 4aed45a9f..1234f7325 100644 --- a/models/deepseek_v4_flash_dspark/prefill_fwd.py +++ b/models/deepseek_v4_flash_dspark/prefill_fwd.py @@ -15,7 +15,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import run_jit +from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig from moe import ( @@ -1337,7 +1337,7 @@ def init_value(): return torch.stack(layer_values, dim=1) # Mutable caches are fixture outputs. - return TensorSpec(name, packed_shape, spec.dtype, init_value=init_value, is_output=name in RESIDENT_CACHE_OUTPUT_NAMES) + return TensorSpec(name, packed_shape, spec.dtype, init_value=init_value) def _make_o_proj_tp_stacked_spec(name, base_specs): @@ -1374,7 +1374,7 @@ def init_value(): target.copy_(source) return packed - return TensorSpec(name, packed_shape, spec.dtype, init_value=init_value, is_output=False) + return TensorSpec(name, packed_shape, spec.dtype, init_value=init_value) def _make_shared_spec(name, base_specs): @@ -1386,7 +1386,7 @@ def _make_shared_spec(name, base_specs): def init_value(): return _expand_rank_axis(_spec_value(spec, torch), torch) - return TensorSpec(name, [N_RANKS, *spec.shape[1:]], spec.dtype, init_value=init_value, is_output=False) + return TensorSpec(name, [N_RANKS, *spec.shape[1:]], spec.dtype, init_value=init_value) def _align_up(value, alignment): @@ -1536,7 +1536,7 @@ def kind_specs(build_fn, build_ragged_fn): ] tensor_specs = [ - TensorSpec(name, list(src.shape), src.dtype, init_value=src.init_value, is_output=src.is_output) + TensorSpec(name, list(src.shape), src.dtype, init_value=src.init_value) for name, src in attention_specs ] @@ -1572,7 +1572,7 @@ def init_input_ids(spec=spec): else: tensor_specs.append(spec) - tensor_specs.append(TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) + tensor_specs.append(TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32)) tensor_by_name = {spec.name: spec for spec in tensor_specs} missing = [name for name in HOST_TENSOR_ORDER if name not in tensor_by_name] if missing: @@ -1735,7 +1735,7 @@ def init_x_hc(tokens=physical_tokens, active_tokens=num_tokens, dtype=base.dtype group_ids = torch.arange(N_RANKS, dtype=torch.int64) // TP_SIZE return (global_x[group_ids] * 0.05).to(dtype).contiguous() - specs.append(TensorSpec(name, x_hc_shape, base.dtype, init_value=init_x_hc, is_output=True)) + specs.append(TensorSpec(name, x_hc_shape, base.dtype, init_value=init_x_hc)) elif name == "position_ids_local": if fixture_case == "ragged2": specs.append(_make_shared_spec(name, base_specs)) @@ -1786,12 +1786,10 @@ def init_padded_mapping(spec=base_spec, active_tokens=num_tokens): N_RANKS, O_PROJ_SCRATCH_GROUPS, O_PROJ_SCRATCH_RANK, O_PROJ_SCRATCH_INPUT, dtype=torch.bfloat16, ), - is_output=False, ), TensorSpec( "o_proj_wo_b_full", [N_RANKS, O_PROJ_SCRATCH_D, O_PROJ_SCRATCH_COLS], torch.int8, init_value=lambda: torch.zeros(N_RANKS, O_PROJ_SCRATCH_D, O_PROJ_SCRATCH_COLS, dtype=torch.int8), - is_output=False, ), ] for spec in o_proj_scratch_specs: @@ -1800,7 +1798,6 @@ def init_padded_mapping(spec=base_spec, active_tokens=num_tokens): attn_stage = TensorSpec( "attn_stage", [N_RANKS, stage_tokens, HC_MULT, D], torch.float32, init_value=lambda: torch.zeros(N_RANKS, stage_tokens, HC_MULT, D, dtype=torch.float32), - is_output=True, ) attn_stage.resident = "stacked" specs.append(attn_stage) @@ -1808,22 +1805,18 @@ def init_padded_mapping(spec=base_spec, active_tokens=num_tokens): TensorSpec( "x_mixed", [N_RANKS, stage_tokens, D], torch.bfloat16, init_value=lambda: torch.zeros(N_RANKS, stage_tokens, D, dtype=torch.bfloat16), - is_output=False, ), TensorSpec( "post_ffn", [N_RANKS, stage_tokens, HC_MULT], torch.float32, init_value=lambda: torch.zeros(N_RANKS, stage_tokens, HC_MULT, dtype=torch.float32), - is_output=False, ), TensorSpec( "comb_ffn", [N_RANKS, stage_tokens, HC_MULT * HC_MULT], torch.float32, init_value=lambda: torch.zeros(N_RANKS, stage_tokens, HC_MULT * HC_MULT, dtype=torch.float32), - is_output=False, ), TensorSpec( "ffn_out", [N_RANKS, local_tokens, D], torch.bfloat16, init_value=lambda: torch.zeros(N_RANKS, local_tokens, D, dtype=torch.bfloat16), - is_output=False, ), ] for spec in moe_stage_specs: @@ -1870,9 +1863,9 @@ def init_hidden_workspace(): TensorSpec("lm_head_weight", [N_RANKS, VOCAB_PER_TP, D], torch.bfloat16, init_value=init_lm_head_weight), TensorSpec("logit_row_indices", [N_RANKS, MAX_LOGIT_ROWS], torch.int32, init_value=init_logit_row_indices), TensorSpec("hidden_workspace", [N_RANKS, stage_tokens, D], torch.bfloat16, init_value=init_hidden_workspace), - TensorSpec("x_out", [N_RANKS, stage_tokens, D], torch.bfloat16, is_output=True), - TensorSpec("logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, is_output=True), - TensorSpec("sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, is_output=True), + TensorSpec("x_out", [N_RANKS, stage_tokens, D], torch.bfloat16), + TensorSpec("logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32), + TensorSpec("sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32), ] for spec in head_specs: spec.resident = "stacked" @@ -2091,7 +2084,7 @@ def main(): fixture_case=args.case, ) - result = run_jit( + result = run( fn=l3_prefill_fwd, specs=specs, golden_fn=golden_prefill_fwd, diff --git a/models/deepseek_v4_flash_dspark/prefill_hca.py b/models/deepseek_v4_flash_dspark/prefill_hca.py index 184c2f252..df68c9a29 100644 --- a/models/deepseek_v4_flash_dspark/prefill_hca.py +++ b/models/deepseek_v4_flash_dspark/prefill_hca.py @@ -825,7 +825,6 @@ def init_wo_b(): [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, - is_output=True, ), TensorSpec( "compress_state_block_table", @@ -838,7 +837,6 @@ def init_wo_b(): [HCA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, - is_output=True, ), TensorSpec("ori_slot_mapping", [token_count], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("ori_block_table", [1, SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), @@ -847,7 +845,6 @@ def init_wo_b(): [HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("cmp_block_table", [1, SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), @@ -857,7 +854,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [token_count, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [token_count, HC_MULT, D], torch.float32), ] @@ -1429,11 +1426,11 @@ def build_cp_tensor_specs( init_value=torch.stack(shards).contiguous(), )) elif spec.name == "x_out": - specs.append(TensorSpec("x_out_full", [tp_size, token_count, HC_MULT, D], spec.dtype, is_output=True)) + specs.append(TensorSpec("x_out_full", [tp_size, token_count, HC_MULT, D], spec.dtype)) else: specs.append(TensorSpec( spec.name, [tp_size, *spec.shape], spec.dtype, - init_value=cp_stack(value, tp_size), is_output=spec.is_output, + init_value=cp_stack(value, tp_size), )) return specs @@ -1552,7 +1549,7 @@ def build_ragged2_cp_tensor_specs(tp_size: int = TP_SIZE): continue replacement_spec = TensorSpec( spec.name, list(value.shape), spec.dtype, init_value=value, - is_output=spec.is_output, resident=spec.resident, + resident=spec.resident, ) specs.append(replacement_spec) return specs @@ -1599,7 +1596,7 @@ def golden_prefill_attention_hca_cp(tensors): if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser( description="Standalone DeepSeek V4 packed prefill HCA correctness test." @@ -1646,7 +1643,7 @@ def golden_prefill_attention_hca_cp(tensors): parser.error(f"--token-count must be a multiple of --tp={TP_SIZE}, got {args.token_count}") if TP_SIZE == 1: - result = run_jit( + result = run( fn=prefill_attention_hca_test, specs=build_tensor_specs(args.start_pos, args.token_count), golden_fn=golden_prefill_attention_hca, @@ -1675,7 +1672,7 @@ def golden_prefill_attention_hca_cp(tensors): if args.case == "ragged2" else build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE) ) - result = run_jit( + result = run( fn=l3_prefill_attention_hca_cp, specs=specs, golden_fn=golden_prefill_attention_hca_cp, diff --git a/models/deepseek_v4_flash_dspark/prefill_indexer.py b/models/deepseek_v4_flash_dspark/prefill_indexer.py index 02190817e..4d6953f84 100644 --- a/models/deepseek_v4_flash_dspark/prefill_indexer.py +++ b/models/deepseek_v4_flash_dspark/prefill_indexer.py @@ -1216,7 +1216,6 @@ def init_sin(): [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], torch.float32, init_value=init_inner_compress_state, - is_output=True, ), TensorSpec( "inner_compress_state_block_table", @@ -1233,17 +1232,15 @@ def init_sin(): [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, - is_output=True, ), TensorSpec( "idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, - is_output=True, ), TensorSpec("idx_block_table", [1, IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), - TensorSpec("topk_idxs", [token_count, IDX_TOPK], torch.int32, is_output=True), + TensorSpec("topk_idxs", [token_count, IDX_TOPK], torch.int32), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), TensorSpec("local_request_ids", [token_count], torch.int32, init_value=init_local_request_ids), TensorSpec( @@ -1264,7 +1261,7 @@ def init_sin(): if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run from utils import int8_quant_per_row parser = argparse.ArgumentParser( @@ -1402,7 +1399,7 @@ def compare(actual, expected, *, actual_outputs, expected_outputs, inputs, rtol, compare.__name__ = f"mapped_active_rows({mapping_name})" return compare - result = run_jit( + result = run( fn=prefill_indexer_test, specs=build_tensor_specs(args.start_pos, args.token_count), golden_fn=golden_prefill_indexer, diff --git a/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py b/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py index 6d007a926..2b74d0ddb 100644 --- a/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py +++ b/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py @@ -939,7 +939,6 @@ def init_inner_state_slot_mapping(): [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, - is_output=True, ), TensorSpec( "inner_compress_state_block_table", @@ -959,14 +958,12 @@ def init_inner_state_slot_mapping(): [IDX_CACHE_MAX_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, - is_output=True, ), TensorSpec( "idx_kv_scale", [IDX_CACHE_MAX_BLOCKS, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, - is_output=True, ), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), TensorSpec("idx_slot_mapping", [token_count], torch.int64, init_value=init_idx_slot_mapping), @@ -981,7 +978,7 @@ def init_inner_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser( description="Standalone token-major DeepSeek V4 prefill indexer compressor validation." @@ -1014,7 +1011,7 @@ def init_inner_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_indexer_compressor_test, specs=build_tensor_specs(args.start_pos, args.token_count), golden_fn=golden_prefill_indexer_compressor, diff --git a/models/deepseek_v4_flash_dspark/prefill_layer.py b/models/deepseek_v4_flash_dspark/prefill_layer.py index 3d55aadd3..b056c28b2 100644 --- a/models/deepseek_v4_flash_dspark/prefill_layer.py +++ b/models/deepseek_v4_flash_dspark/prefill_layer.py @@ -15,7 +15,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import ScalarSpec, TensorSpec, ratio_allclose, ratio_reldiff, run_jit +from golden import ScalarSpec, TensorSpec, ratio_allclose, ratio_reldiff, run from pypto.ir.distributed_compiled_program import DistributedConfig from moe import ( @@ -790,7 +790,7 @@ def init_ranked(source=source, name=name): shape[0] = N_RANKS specs_by_name[name] = TensorSpec( name, shape, source.dtype, - init_value=init_ranked, is_output=name in _CACHE_STATE_NAMES, + init_value=init_ranked, ) def init_attn_stage(): @@ -799,16 +799,16 @@ def init_attn_stage(): stage_specs = [ TensorSpec( "attn_stage", [N_RANKS, physical_tokens, HC_MULT, D], torch.float32, - init_value=init_attn_stage, is_output=True, + init_value=init_attn_stage, ), - TensorSpec("x_mixed", [N_RANKS, physical_tokens, D], torch.bfloat16, is_output=True), - TensorSpec("post_ffn", [N_RANKS, physical_tokens, HC_MULT], torch.float32, is_output=True), - TensorSpec("comb_ffn", [N_RANKS, physical_tokens, HC_MULT * HC_MULT], torch.float32, is_output=True), - TensorSpec("ffn_out", [N_RANKS, local_tokens, D], torch.bfloat16, is_output=True), + TensorSpec("x_mixed", [N_RANKS, physical_tokens, D], torch.bfloat16), + TensorSpec("post_ffn", [N_RANKS, physical_tokens, HC_MULT], torch.float32), + TensorSpec("comb_ffn", [N_RANKS, physical_tokens, HC_MULT * HC_MULT], torch.float32), + TensorSpec("ffn_out", [N_RANKS, local_tokens, D], torch.bfloat16), ] specs_by_name.update({spec.name: spec for spec in stage_specs}) specs_by_name["x_next"] = TensorSpec( - "x_next", [N_RANKS, physical_tokens, HC_MULT, D], torch.float32, is_output=True + "x_next", [N_RANKS, physical_tokens, HC_MULT, D], torch.float32 ) for name in _RESIDENT_WEIGHT_NAMES: @@ -1053,7 +1053,7 @@ def main(): import torch torch.manual_seed(args.seed) - result = run_jit( + result = run( fn=l3_prefill_layer, specs=build_tensor_specs(args.layer_id, start_pos=args.start_pos, token_count=args.token_count), golden_fn=golden_prefill_layer, diff --git a/models/deepseek_v4_flash_dspark/prefill_metadata.py b/models/deepseek_v4_flash_dspark/prefill_metadata.py index 30f5e33e7..cdeb503ea 100644 --- a/models/deepseek_v4_flash_dspark/prefill_metadata.py +++ b/models/deepseek_v4_flash_dspark/prefill_metadata.py @@ -66,7 +66,7 @@ def build_tensor_specs(): return [ TensorSpec("query_start_loc", list(query_start_loc.shape), torch.int32, init_value=query_start_loc), ScalarSpec("local_base", torch.int32, local_base), - TensorSpec("request_ids", list(request_ids.shape), torch.int32, is_output=True), + TensorSpec("request_ids", list(request_ids.shape), torch.int32), ] @@ -95,7 +95,7 @@ def golden_prefill_metadata(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("-p", "--platform", default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -103,7 +103,7 @@ def golden_prefill_metadata(tensors): parser.add_argument("--compile-only", action="store_true") args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_metadata_test, specs=build_tensor_specs(), golden_fn=golden_prefill_metadata, diff --git a/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py b/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py index dbf89f8a6..0265e22e8 100644 --- a/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py +++ b/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py @@ -1529,13 +1529,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("attn_out", [token_count, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [token_count, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument( @@ -1556,7 +1556,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_sparse_attn_test, specs=build_tensor_specs(args.compress_ratio, args.tokens, args.ori_block_num, args.cmp_block_num), golden_fn=golden_prefill_sparse_attn, diff --git a/models/deepseek_v4_flash_dspark/prefill_swa.py b/models/deepseek_v4_flash_dspark/prefill_swa.py index 9a8246f19..2d3aa7411 100644 --- a/models/deepseek_v4_flash_dspark/prefill_swa.py +++ b/models/deepseek_v4_flash_dspark/prefill_swa.py @@ -834,7 +834,7 @@ def init_wo_b(): TensorSpec("freqs_cos", [token_count, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [token_count, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), TensorSpec("kv_cache", [BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, - init_value=init_kv_cache, is_output=True), + init_value=init_kv_cache), TensorSpec("block_table", [1, BLOCK_TABLE_BLOCKS], torch.int32, init_value=init_block_table), TensorSpec("ori_slot_mapping", [token_count], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), @@ -843,7 +843,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [token_count, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [token_count, HC_MULT, D], torch.float32), ] @@ -904,11 +904,11 @@ def build_cp_tensor_specs( init_value=torch.stack(shards).contiguous(), )) elif spec.name == "x_out": - specs.append(TensorSpec("x_out_full", [tp_size, token_count, HC_MULT, D], spec.dtype, is_output=True)) + specs.append(TensorSpec("x_out_full", [tp_size, token_count, HC_MULT, D], spec.dtype)) else: specs.append(TensorSpec( spec.name, [tp_size, *spec.shape], spec.dtype, - init_value=cp_stack(value, tp_size), is_output=spec.is_output, + init_value=cp_stack(value, tp_size), )) return specs @@ -971,7 +971,7 @@ def build_ragged2_cp_tensor_specs(tp_size: int = TP_SIZE): continue replacement_spec = TensorSpec( spec.name, list(value.shape), spec.dtype, init_value=value, - is_output=spec.is_output, resident=spec.resident, + resident=spec.resident, ) specs.append(replacement_spec) return specs @@ -1018,7 +1018,7 @@ def golden_prefill_attention_swa_cp(tensors): if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser( description="Standalone DeepSeek V4 packed prefill SWA correctness test." @@ -1062,7 +1062,7 @@ def golden_prefill_attention_swa_cp(tensors): parser.error(f"--token-count must be a multiple of --tp={TP_SIZE}, got {args.token_count}") if TP_SIZE == 1: - result = run_jit( + result = run( fn=prefill_attention_swa_test, specs=build_tensor_specs(args.start_pos, args.token_count), golden_fn=golden_prefill_attention_swa, @@ -1089,7 +1089,7 @@ def golden_prefill_attention_swa_cp(tensors): if args.case == "ragged2" else build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE) ) - result = run_jit( + result = run( fn=l3_prefill_attention_swa_cp, specs=specs, golden_fn=golden_prefill_attention_swa_cp, diff --git a/models/deepseek_v4_flash_dspark/qkv_proj_rope.py b/models/deepseek_v4_flash_dspark/qkv_proj_rope.py index 5fd8bcfcb..69bbe433f 100644 --- a/models/deepseek_v4_flash_dspark/qkv_proj_rope.py +++ b/models/deepseek_v4_flash_dspark/qkv_proj_rope.py @@ -963,10 +963,10 @@ def build_split_tensor_specs(): ), base["gamma_cq"], base["gamma_ckv"], - TensorSpec("q", [SPLIT_T_LOCAL, H, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("qr", [SPLIT_T_LOCAL, Q_LORA], torch.int8, is_output=True), - TensorSpec("qr_scale", [SPLIT_T_LOCAL, 1], torch.float32, is_output=True), - TensorSpec("kv", [SPLIT_T_FULL, HEAD_DIM], torch.bfloat16, is_output=True), + TensorSpec("q", [SPLIT_T_LOCAL, H, HEAD_DIM], torch.bfloat16), + TensorSpec("qr", [SPLIT_T_LOCAL, Q_LORA], torch.int8), + TensorSpec("qr_scale", [SPLIT_T_LOCAL, 1], torch.float32), + TensorSpec("kv", [SPLIT_T_FULL, HEAD_DIM], torch.bfloat16), ] @@ -1089,16 +1089,16 @@ def init_gamma_ckv(): TensorSpec("rope_sin", [T, ROPE_DIM], torch.bfloat16, init_value=init_sin), TensorSpec("gamma_cq", [Q_LORA], torch.bfloat16, init_value=init_gamma_cq), TensorSpec("gamma_ckv", [HEAD_DIM], torch.bfloat16, init_value=init_gamma_ckv), - TensorSpec("q", [T, H, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("kv", [T, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("qr", [T, Q_LORA], torch.int8, is_output=True), - TensorSpec("qr_scale", [T, 1], torch.float32, is_output=True), + TensorSpec("q", [T, H, HEAD_DIM], torch.bfloat16), + TensorSpec("kv", [T, HEAD_DIM], torch.bfloat16), + TensorSpec("qr", [T, Q_LORA], torch.int8), + TensorSpec("qr_scale", [T, 1], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH // TP, DECODE_SEQ), @@ -1140,7 +1140,7 @@ def init_gamma_ckv(): B, S = MODES[mode_name] fn, specs, golden = qkv_proj_rope_test, build_tensor_specs(B, S), golden_qkv_proj_rope print(f"--- qkv_proj_rope {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=fn, specs=specs, golden_fn=golden, diff --git a/models/deepseek_v4_flash_dspark/rmsnorm.py b/models/deepseek_v4_flash_dspark/rmsnorm.py index 5cdf6e21c..14f017998 100644 --- a/models/deepseek_v4_flash_dspark/rmsnorm.py +++ b/models/deepseek_v4_flash_dspark/rmsnorm.py @@ -178,13 +178,13 @@ def init_norm_w(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), TensorSpec("norm_w", [D], torch.bfloat16, init_value=init_norm_w), - TensorSpec("x_normed", [T, D], torch.bfloat16, is_output=True), + TensorSpec("x_normed", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH // TP, DECODE_SEQ), @@ -208,7 +208,7 @@ def init_norm_w(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- rms_norm_test {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=rms_norm_test, specs=build_tensor_specs(B, S), golden_fn=golden_rms_norm_test, diff --git a/models/deepseek_v4_flash_mtp/decode_compressor_ratio128.py b/models/deepseek_v4_flash_mtp/decode_compressor_ratio128.py index 030a8c68b..c074314a4 100644 --- a/models/deepseek_v4_flash_mtp/decode_compressor_ratio128.py +++ b/models/deepseek_v4_flash_mtp/decode_compressor_ratio128.py @@ -547,8 +547,8 @@ def init_cmp_slot_mapping(): ) return [ TensorSpec("x", [B, S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [B, S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [B, S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -556,7 +556,7 @@ def init_cmp_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("cos", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_cos), TensorSpec("sin", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), - TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache, is_output=True), + TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("cmp_slot_mapping", [B, S], torch.int64, init_value=init_cmp_slot_mapping), TensorSpec("state_slot_mapping", [B, S], torch.int64, init_value=init_state_slot_mapping), @@ -565,7 +565,7 @@ def init_cmp_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -578,7 +578,7 @@ def init_cmp_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_flash_mtp/decode_compressor_ratio4.py b/models/deepseek_v4_flash_mtp/decode_compressor_ratio4.py index bf82e69c4..81f6ec8aa 100644 --- a/models/deepseek_v4_flash_mtp/decode_compressor_ratio4.py +++ b/models/deepseek_v4_flash_mtp/decode_compressor_ratio4.py @@ -529,8 +529,8 @@ def init_cmp_slot_mapping(): return [ TensorSpec("x", [B, S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [B, S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [B, S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -538,7 +538,7 @@ def init_cmp_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("cos", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_cos), TensorSpec("sin", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), - TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache, is_output=True), + TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("cmp_slot_mapping", [B, S], torch.int64, init_value=init_cmp_slot_mapping), TensorSpec("state_slot_mapping", [B, S], torch.int64, init_value=init_state_slot_mapping), @@ -547,7 +547,7 @@ def init_cmp_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -562,7 +562,7 @@ def init_cmp_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_flash_mtp/decode_csa.py b/models/deepseek_v4_flash_mtp/decode_csa.py index a65f679c8..210c73ad4 100644 --- a/models/deepseek_v4_flash_mtp/decode_csa.py +++ b/models/deepseek_v4_flash_mtp/decode_csa.py @@ -900,7 +900,7 @@ def init_wo_b(): TensorSpec("inner_norm_w", [IDX_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), TensorSpec("inner_compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], torch.float32, init_value=init_inner_compress_state), TensorSpec("inner_compress_state_block_table", [B, INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), TensorSpec("cmp_kv", [CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("cmp_block_table", [B, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: shared_idx_kv_cache_i8.clone()), @@ -919,13 +919,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -941,7 +941,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=attention_csa_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_attention_csa, diff --git a/models/deepseek_v4_flash_mtp/decode_fwd.py b/models/deepseek_v4_flash_mtp/decode_fwd.py index 2172f82f4..6a2c0e18c 100644 --- a/models/deepseek_v4_flash_mtp/decode_fwd.py +++ b/models/deepseek_v4_flash_mtp/decode_fwd.py @@ -14,7 +14,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import run_jit +from golden import run from hc_head import hc_head from lm_head import ( GROUP_LOGIT_ROWS, @@ -1100,13 +1100,12 @@ def init_value(): init_value=init_value, # Caches the kernel writes in place (kv_cache) are read back for # validation; every other stacked tensor is a plain input. - is_output=name in RESIDENT_CACHE_OUTPUT_NAMES, ) def _make_shared_spec(name, base_spec, out_name=None): from golden import TensorSpec - return TensorSpec(out_name or name, list(base_spec.shape), base_spec.dtype, init_value=base_spec.init_value if out_name is None else None, is_output=out_name is not None) + return TensorSpec(out_name or name, list(base_spec.shape), base_spec.dtype, init_value=base_spec.init_value if out_name is None else None) def _make_hc_head_spec(name): @@ -1623,7 +1622,6 @@ def init_block_table(): cache_shapes[name], spec.dtype, init_value=lambda shape=cache_shapes[name], dtype=spec.dtype: torch.zeros(shape, dtype=dtype), - is_output=spec.is_output, ) if name in cache_shapes else spec ) @@ -1636,7 +1634,6 @@ def init_block_table(): [N_RANKS, *spec.shape], spec.dtype, init_value=_ranked_init(spec, replicated=name in replicated_attention), - is_output=name == "kv_cache", ) for name, spec in attention_specs ] @@ -1664,7 +1661,7 @@ def init_input_ids(): specs.append(moe_tensor_specs[spec.name]) specs.extend([ - TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32), ScalarSpec("layer_id", torch.int32, layer_id), ]) return specs @@ -1805,7 +1802,7 @@ def init_logit_row_indices(): if spec.name in RESIDENT_WEIGHT_NAMES or spec.name in CACHE_POOL_NAMES: spec.resident = "stacked" - specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) + specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32)) specs.append(TensorSpec( "lm_head_weight", [N_RANKS, VOCAB_PER_TP, D], torch.bfloat16, init_value=init_lm_head_weight, resident="stacked", @@ -1836,13 +1833,12 @@ def init_logit_row_indices(): torch.int32, init_value=lambda: torch.zeros(N_RANKS, MAX_LOGIT_ROWS, dtype=torch.int32), )) - specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16, is_output=True)) - specs.append(TensorSpec("logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, is_output=True)) + specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16)) + specs.append(TensorSpec("logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32)) specs.append(TensorSpec( "sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, - is_output=True, )) specs.append(TensorSpec( "num_tokens_per_owner", @@ -1907,7 +1903,7 @@ def main(): inner_state_block_num=args.inner_state_block_num, ) - result = run_jit( + result = run( fn=l3_decode_fwd, specs=specs, golden_fn=None, diff --git a/models/deepseek_v4_flash_mtp/decode_fwd_mtp.py b/models/deepseek_v4_flash_mtp/decode_fwd_mtp.py index 3d934a0f8..076dcc084 100644 --- a/models/deepseek_v4_flash_mtp/decode_fwd_mtp.py +++ b/models/deepseek_v4_flash_mtp/decode_fwd_mtp.py @@ -1064,8 +1064,6 @@ def build_tensor_specs( inner_state_block_num=inner_state_block_num, ) ) - for name in ("input_ids", "position_ids", "kv_seq_lens"): - specs[name] = replace(specs[name], is_output=True) def init_tail_token_ids(): tokens = torch.arange(B, dtype=torch.int64) + 10 @@ -1105,14 +1103,12 @@ def init_state_meta(): [N_RANKS, B], torch.int64, init_value=init_tail_token_ids, - is_output=True, ), "mtp_tail_positions": TensorSpec( "mtp_tail_positions", [N_RANKS, B], torch.int32, init_value=init_tail_positions, - is_output=True, ), "mtp_tail_slot_ids": TensorSpec( "mtp_tail_slot_ids", @@ -1131,7 +1127,6 @@ def init_state_meta(): [N_RANKS, B, STATE_TOKEN_WIDTH], torch.int64, init_value=init_state_tokens, - is_output=True, resident="stacked", ), "mtp_state_meta": TensorSpec( @@ -1139,33 +1134,28 @@ def init_state_meta(): [N_RANKS, B, STATE_META_WIDTH], torch.int32, init_value=init_state_meta, - is_output=True, resident="stacked", ), "mtp_input_ids": replace( mtp_specs["input_ids"], name="mtp_input_ids", init_value=None, - is_output=True, ), "mtp_position_ids": replace( mtp_specs["position_ids"], name="mtp_position_ids", init_value=None, - is_output=True, ), "mtp_accepted_counts": TensorSpec( "mtp_accepted_counts", [N_RANKS, B], torch.int32, - is_output=True, ), "mtp_tail_pre_hc_pool": TensorSpec( "mtp_tail_pre_hc_pool", [N_RANKS, B, HC_MULT, D], torch.float32, init_value=lambda: torch.randn(N_RANKS, B, HC_MULT, D), - is_output=True, resident="stacked", ), "mtp_hidden_out": replace(mtp_specs["hidden_out"], name="mtp_hidden_out"), @@ -1209,17 +1199,6 @@ def init_state_meta(): continue specs[f"mtp_{name}"] = replace(spec, name=f"mtp_{name}") - sampling_names = { - "sampling_temperatures", - "sampling_top_ks", - "sampling_seeds", - "sampling_positions", - } - for name in sampling_names: - specs[name] = replace(specs[name], is_output=True) - mtp_name = f"mtp_{name}" - specs[mtp_name] = replace(specs[mtp_name], is_output=True) - param_names = l3_decode_fwd_mtp._param_names() missing = set(param_names) - specs.keys() extra = specs.keys() - set(param_names) @@ -1231,7 +1210,7 @@ def init_state_meta(): def main(): - from golden import run_jit + from golden import run parser = argparse.ArgumentParser( description="DeepSeek-V4 fused main-decode, verification, and MTP decode driver." @@ -1293,7 +1272,7 @@ def main(): device_ids = [int(device) for device in args.device.split(",")] assert len(device_ids) >= N_RANKS, f"need at least {N_RANKS} devices, got {device_ids}" - result = run_jit( + result = run( fn=l3_decode_fwd_mtp, specs=build_tensor_specs( start_pos=args.start_pos, diff --git a/models/deepseek_v4_flash_mtp/decode_hca.py b/models/deepseek_v4_flash_mtp/decode_hca.py index 7a68ee52b..9dcdba681 100644 --- a/models/deepseek_v4_flash_mtp/decode_hca.py +++ b/models/deepseek_v4_flash_mtp/decode_hca.py @@ -658,7 +658,7 @@ def init_wo_b(): TensorSpec("cmp_norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_cmp_norm_w), TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), TensorSpec("cmp_kv", [CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("cmp_block_table", [B, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), @@ -672,13 +672,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -693,7 +693,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=attention_hca_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_attention_hca, diff --git a/models/deepseek_v4_flash_mtp/decode_indexer.py b/models/deepseek_v4_flash_mtp/decode_indexer.py index e196fb72d..b5769a3ac 100644 --- a/models/deepseek_v4_flash_mtp/decode_indexer.py +++ b/models/deepseek_v4_flash_mtp/decode_indexer.py @@ -685,12 +685,12 @@ def init_idx_slot_mapping(): TensorSpec("inner_wgate", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wgate), TensorSpec("inner_ape", [COMPRESS_RATIO, INNER_OUT_DIM], torch.float32, init_value=init_inner_ape), TensorSpec("inner_norm_w", [INNER_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), - TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_kv_i8, is_output=True), - TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: idx_kv_sc, is_output=True), + TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_kv_i8), + TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: idx_kv_sc), TensorSpec("idx_block_table", [B, IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), # Outputs are fixed to SCORE_LEN; positions past cache_len are -inf for score and -1 for topk_idxs. - TensorSpec("score", [B, S, SCORE_LEN], torch.float32, is_output=True), - TensorSpec("topk_idxs", [B, S, SCORE_LEN], torch.int32, is_output=True), + TensorSpec("score", [B, S, SCORE_LEN], torch.float32), + TensorSpec("topk_idxs", [B, S, SCORE_LEN], torch.int32), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("idx_slot_mapping", [B, S], torch.int64, init_value=init_idx_slot_mapping), TensorSpec("inner_state_slot_mapping", [B, S], torch.int64, init_value=init_inner_state_slot_mapping), @@ -702,7 +702,7 @@ def init_idx_slot_mapping(): if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit, topk_pair_compare + from golden import ratio_allclose, run, topk_pair_compare parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -752,7 +752,7 @@ def score_valid_compare(actual, expected, *, actual_outputs, expected_outputs, i ) score_valid_compare.__name__ = "score_valid_region_compare" - result = run_jit( + result = run( fn=indexer_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_indexer, diff --git a/models/deepseek_v4_flash_mtp/decode_indexer_compressor.py b/models/deepseek_v4_flash_mtp/decode_indexer_compressor.py index 2e73688f1..ea390d243 100644 --- a/models/deepseek_v4_flash_mtp/decode_indexer_compressor.py +++ b/models/deepseek_v4_flash_mtp/decode_indexer_compressor.py @@ -589,8 +589,8 @@ def init_idx_slot_mapping(): return [ TensorSpec("x", [B, S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [B, S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [B, S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -599,8 +599,8 @@ def init_idx_slot_mapping(): TensorSpec("cos", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_cos), TensorSpec("sin", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), TensorSpec("hadamard", [HEAD_DIM, HEAD_DIM], torch.bfloat16, init_value=init_hadamard), - TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, is_output=True), - TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, is_output=True), + TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache), + TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("idx_slot_mapping", [B, S], torch.int64, init_value=init_idx_slot_mapping), TensorSpec("inner_state_slot_mapping", [B, S], torch.int64, init_value=init_inner_state_slot_mapping), @@ -609,7 +609,7 @@ def init_idx_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -623,7 +623,7 @@ def init_idx_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_flash_mtp/decode_layer.py b/models/deepseek_v4_flash_mtp/decode_layer.py index afc66beb5..5867c7883 100644 --- a/models/deepseek_v4_flash_mtp/decode_layer.py +++ b/models/deepseek_v4_flash_mtp/decode_layer.py @@ -813,7 +813,6 @@ def init_block_table(): [N_RANKS, *spec.shape], spec.dtype, init_value=_ranked_init(spec, replicated=name in replicated_attention), - is_output=name == "kv_cache", ) for name, spec in attention_specs ] @@ -847,7 +846,7 @@ def init_input_ids(): # KV/state cache or per-step metadata); the MoE set adds the FFN/gate/expert # weights and the static tid2eid route table. The KV/state caches # (RESIDENT_CACHE_NAMES) are kept resident too — the written kv_cache is also - # is_output=True (line above) and read back once at the end for validation; + # an InOut (line above) and read back once at the end for validation; # the others are read-only inputs. NOT resident: the per-step slot mappings / # block tables / ids / position_ids / kv_seq_lens, the input activation # (x_hc), and the output (x_next), which change per token. @@ -868,7 +867,7 @@ def init_input_ids(): spec.resident = "stacked" specs.extend([ - TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32), ScalarSpec("layer_id", torch.int32, layer_id), ]) return specs @@ -876,7 +875,7 @@ def init_input_ids(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -900,7 +899,7 @@ def init_input_ids(): host_fn = l3_decode_layer golden_fn = golden_decode_layer_auto - result = run_jit( + result = run( fn=host_fn, specs=build_tensor_specs( start_pos=args.start_pos, diff --git a/models/deepseek_v4_flash_mtp/decode_mtp.py b/models/deepseek_v4_flash_mtp/decode_mtp.py index 44df01288..2d35bc4f0 100644 --- a/models/deepseek_v4_flash_mtp/decode_mtp.py +++ b/models/deepseek_v4_flash_mtp/decode_mtp.py @@ -484,12 +484,12 @@ def init(): return init -def _ranked_spec(name, spec, *, replicated=False, is_output=False): +def _ranked_spec(name, spec, *, replicated=False): from golden import TensorSpec return TensorSpec( name, [N_RANKS, *spec.shape], spec.dtype, - init_value=_ranked_init(spec, replicated=replicated), is_output=is_output, + init_value=_ranked_init(spec, replicated=replicated), ) @@ -760,7 +760,7 @@ def init_kv_cache(): cache_spec = TensorSpec( name, [N_RANKS, ori_block_num, BLOCK_SIZE, 1, HEAD_DIM], cache_dtype, - init_value=init_kv_cache, is_output=swa_specs[name].is_output, + init_value=init_kv_cache, ) specs.append(cache_spec) elif name in swa_metadata_specs: @@ -770,7 +770,7 @@ def init_kv_cache(): else: ranked_spec = _ranked_spec( name, swa_specs[name], - replicated=name in replicated_attention, is_output=swa_specs[name].is_output, + replicated=name in replicated_attention, ) specs.append(ranked_spec) @@ -832,24 +832,23 @@ def init_kv_cache(): ) hidden_out_spec = TensorSpec( "hidden_out", [N_RANKS, T, D], torch.bfloat16, - is_output=True, resident="stacked", + resident="stacked", ) specs.append(hidden_out_spec) next_hidden_spec = TensorSpec( "next_pre_hc_hidden", [N_RANKS, T, HC_MULT, D], torch.float32, - is_output=True, resident="stacked", + resident="stacked", ) specs.append(next_hidden_spec) logits_spec = TensorSpec( "logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, - is_output=True, resident="stacked", + resident="stacked", ) specs.append(logits_spec) sampled_ids_spec = TensorSpec( "sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, - is_output=True, ) specs.append(sampled_ids_spec) row_indices_spec = TensorSpec( @@ -945,7 +944,7 @@ def golden_decode_mtp(tensors): def main(): - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser(description="DeepSeek-V4 MTP decode layer driver.") parser.add_argument( @@ -983,7 +982,7 @@ def main(): device_ids = [int(d) for d in args.device.split(",")] assert len(device_ids) >= N_RANKS, f"need at least {N_RANKS} devices, got {device_ids}" - result = run_jit( + result = run( fn=l3_decode_mtp, specs=build_tensor_specs( start_pos=args.start_pos, diff --git a/models/deepseek_v4_flash_mtp/decode_sparse_attn_csa.py b/models/deepseek_v4_flash_mtp/decode_sparse_attn_csa.py index 757cae04e..9f70a2f1a 100644 --- a/models/deepseek_v4_flash_mtp/decode_sparse_attn_csa.py +++ b/models/deepseek_v4_flash_mtp/decode_sparse_attn_csa.py @@ -777,13 +777,13 @@ def init_wo_b_scale(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -806,7 +806,7 @@ def init_wo_b_scale(): print(f"compress_ratio={COMPRESS_RATIO} -> TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}", flush=True) - result = run_jit( + result = run( fn=sparse_attn_test, specs=build_tensor_specs( args.causal_regression_fixture, diff --git a/models/deepseek_v4_flash_mtp/decode_sparse_attn_hca.py b/models/deepseek_v4_flash_mtp/decode_sparse_attn_hca.py index c6b23e3cb..b0f093551 100644 --- a/models/deepseek_v4_flash_mtp/decode_sparse_attn_hca.py +++ b/models/deepseek_v4_flash_mtp/decode_sparse_attn_hca.py @@ -755,13 +755,13 @@ def init_wo_b_scale(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -785,7 +785,7 @@ def init_wo_b_scale(): print(f"compress_ratio={COMPRESS_RATIO} -> TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}", flush=True) - result = run_jit( + result = run( fn=sparse_attn_test, specs=build_tensor_specs( args.causal_regression_fixture, diff --git a/models/deepseek_v4_flash_mtp/decode_sparse_attn_swa.py b/models/deepseek_v4_flash_mtp/decode_sparse_attn_swa.py index 53b0a8db9..0165255df 100644 --- a/models/deepseek_v4_flash_mtp/decode_sparse_attn_swa.py +++ b/models/deepseek_v4_flash_mtp/decode_sparse_attn_swa.py @@ -633,13 +633,13 @@ def init_wo_b_scale(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -659,7 +659,7 @@ def init_wo_b_scale(): print(f"TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}", flush=True) - result = run_jit( + result = run( fn=sparse_attn_test, specs=build_tensor_specs( args.causal_regression_fixture, diff --git a/models/deepseek_v4_flash_mtp/decode_swa.py b/models/deepseek_v4_flash_mtp/decode_swa.py index 583aad61d..666cfbc7f 100644 --- a/models/deepseek_v4_flash_mtp/decode_swa.py +++ b/models/deepseek_v4_flash_mtp/decode_swa.py @@ -462,7 +462,7 @@ def init_wo_b(): TensorSpec("gamma_ckv", [HEAD_DIM], torch.bfloat16, init_value=init_gamma_ckv), TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), TensorSpec("swa_slot_mapping", [T], torch.int64, init_value=init_swa_slot_mapping), TensorSpec("swa_indices", [T, WIN], torch.int32, init_value=init_swa_indices), TensorSpec("swa_lens", [T], torch.int32, init_value=init_swa_lens), @@ -471,13 +471,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -492,7 +492,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=attention_swa_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_attention_swa, diff --git a/models/deepseek_v4_flash_mtp/expert_routed.py b/models/deepseek_v4_flash_mtp/expert_routed.py index 26f94a5b2..06ffb4522 100644 --- a/models/deepseek_v4_flash_mtp/expert_routed.py +++ b/models/deepseek_v4_flash_mtp/expert_routed.py @@ -448,13 +448,13 @@ def init_recv_weights(): TensorSpec("routed_w3_scale", [N_LOCAL_EXPERTS, MOE_INTER], torch.float32, init_value=lambda: w3_s), TensorSpec("routed_w2", [N_LOCAL_EXPERTS, D, MOE_INTER], torch.int8, init_value=lambda: w2_i8), TensorSpec("routed_w2_scale", [N_LOCAL_EXPERTS, D], torch.float32, init_value=lambda: w2_s), - TensorSpec("recv_y", [N_LOCAL_EXPERTS, RECV_MAX, D], torch.bfloat16, is_output=True), + TensorSpec("recv_y", [N_LOCAL_EXPERTS, RECV_MAX, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -464,7 +464,7 @@ def init_recv_weights(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=expert_routed_test, specs=build_tensor_specs(), golden_fn=golden_expert_routed, diff --git a/models/deepseek_v4_flash_mtp/expert_shared.py b/models/deepseek_v4_flash_mtp/expert_shared.py index 79925075c..92a23b7d3 100644 --- a/models/deepseek_v4_flash_mtp/expert_shared.py +++ b/models/deepseek_v4_flash_mtp/expert_shared.py @@ -388,13 +388,13 @@ def build_tensor_specs(): TensorSpec("shared_w3_scale", [MOE_INTER], torch.float32, init_value=lambda: sw3_s), TensorSpec("shared_w2", [D, MOE_INTER], torch.int8, init_value=lambda: sw2_i8), TensorSpec("shared_w2_scale", [D], torch.float32, init_value=lambda: sw2_s), - TensorSpec("sh", [T, D], torch.bfloat16, is_output=True), + TensorSpec("sh", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -404,7 +404,7 @@ def build_tensor_specs(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=expert_shared_test, specs=build_tensor_specs(), golden_fn=golden_expert_shared, diff --git a/models/deepseek_v4_flash_mtp/gate.py b/models/deepseek_v4_flash_mtp/gate.py index beaac91c6..4a21e0ecc 100644 --- a/models/deepseek_v4_flash_mtp/gate.py +++ b/models/deepseek_v4_flash_mtp/gate.py @@ -394,10 +394,10 @@ def init_input_ids(): ScalarSpec("num_tokens", torch.int32, num_tokens), TensorSpec("tid2eid", [VOCAB, TOPK], torch.int32, init_value=init_tid2eid), TensorSpec("input_ids", [T], torch.int64, init_value=init_input_ids), - TensorSpec("x_norm_i8", [T, D], torch.int8, is_output=True), - TensorSpec("x_norm_scale", [T, 1], torch.float32, is_output=True), - TensorSpec("indices", [T, TOPK], torch.int32, is_output=True), - TensorSpec("weights", [T, TOPK], torch.float32, is_output=True), + TensorSpec("x_norm_i8", [T, D], torch.int8), + TensorSpec("x_norm_scale", [T, 1], torch.float32), + TensorSpec("indices", [T, TOPK], torch.int32), + TensorSpec("weights", [T, TOPK], torch.float32), ] @@ -409,7 +409,7 @@ def gate_active_rows(num_tokens): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit, topk_pair_compare + from golden import ratio_allclose, run, topk_pair_compare parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -421,7 +421,7 @@ def gate_active_rows(num_tokens): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=gate_test, specs=build_tensor_specs(layer_id=args.layer_id, num_tokens=args.num_tokens), golden_fn=golden_gate_core, diff --git a/models/deepseek_v4_flash_mtp/hc_head.py b/models/deepseek_v4_flash_mtp/hc_head.py index f23666318..7a1774fea 100644 --- a/models/deepseek_v4_flash_mtp/hc_head.py +++ b/models/deepseek_v4_flash_mtp/hc_head.py @@ -205,14 +205,14 @@ def init_hc_head_fn(): init_value=lambda: torch.tensor([0.076099])), TensorSpec("hc_head_base", [HC_MULT], torch.float32, init_value=lambda: torch.tensor([5.9166, -3.6223, -2.9324, -3.3124])), - TensorSpec("y", [T, D], torch.bfloat16, is_output=True), + TensorSpec("y", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -226,7 +226,7 @@ def init_hc_head_fn(): args = parser.parse_args() torch.manual_seed(args.seed) - result = run_jit( + result = run( fn=hc_head_test, specs=build_tensor_specs(), golden_fn=golden_hc_head, diff --git a/models/deepseek_v4_flash_mtp/hc_post.py b/models/deepseek_v4_flash_mtp/hc_post.py index f658829fc..f1e3cf93c 100644 --- a/models/deepseek_v4_flash_mtp/hc_post.py +++ b/models/deepseek_v4_flash_mtp/hc_post.py @@ -188,13 +188,13 @@ def init_comb(): TensorSpec("residual", [T, HC_MULT, D], torch.float32, init_value=init_residual), TensorSpec("post", [T, HC_MULT], torch.float32, init_value=init_post), TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32, init_value=init_comb), - TensorSpec("y", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("y", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -217,7 +217,7 @@ def init_comb(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- hc_post {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=hc_post_test, specs=build_tensor_specs(B, S), golden_fn=golden_hc_post, diff --git a/models/deepseek_v4_flash_mtp/hc_pre.py b/models/deepseek_v4_flash_mtp/hc_pre.py index 8933ad6e2..8786c8e2b 100644 --- a/models/deepseek_v4_flash_mtp/hc_pre.py +++ b/models/deepseek_v4_flash_mtp/hc_pre.py @@ -379,15 +379,15 @@ def init_hc_base(): TensorSpec("hc_fn", [MIX_HC, HC_DIM], torch.float32, init_value=init_hc_fn), TensorSpec("hc_scale", [3], torch.float32, init_value=init_hc_scale), TensorSpec("hc_base", [MIX_HC], torch.float32, init_value=init_hc_base), - TensorSpec("x_mixed", [T, D], torch.bfloat16, is_output=True), - TensorSpec("post", [T, HC_MULT], torch.float32, is_output=True), - TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32, is_output=True), + TensorSpec("x_mixed", [T, D], torch.bfloat16), + TensorSpec("post", [T, HC_MULT], torch.float32), + TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -410,7 +410,7 @@ def init_hc_base(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- hc_pre {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=hc_pre_test, specs=build_tensor_specs(B, S), golden_fn=golden_hc_pre, diff --git a/models/deepseek_v4_flash_mtp/lm_head.py b/models/deepseek_v4_flash_mtp/lm_head.py index 73c3dff91..73fd5adae 100644 --- a/models/deepseek_v4_flash_mtp/lm_head.py +++ b/models/deepseek_v4_flash_mtp/lm_head.py @@ -549,13 +549,11 @@ def init_logit_row_indices(): "logits", [WORLD_SIZE, MAX_LOGIT_ROWS, VOCAB], torch.float32, - is_output=True, ), TensorSpec( "sampled_ids", [WORLD_SIZE, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, - is_output=True, ), TensorSpec( "logit_row_indices", @@ -606,7 +604,7 @@ def compare_sampled_ids(actual, _expected, *, actual_outputs, **_): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -642,7 +640,7 @@ def compare_sampled_ids(actual, _expected, *, actual_outputs, **_): "sampled_ids": compare_sampled_ids, } - result = run_jit( + result = run( fn=fn, specs=specs, golden_fn=golden_fn, diff --git a/models/deepseek_v4_flash_mtp/lookup_embedding.py b/models/deepseek_v4_flash_mtp/lookup_embedding.py index 8cb89c5bf..0cf65ec17 100644 --- a/models/deepseek_v4_flash_mtp/lookup_embedding.py +++ b/models/deepseek_v4_flash_mtp/lookup_embedding.py @@ -79,13 +79,13 @@ def init_embed_weight(): return [ TensorSpec("input_ids", [token_count], torch.int64, init_value=init_input_ids), TensorSpec("embed_weight", [vocab_size, D], torch.bfloat16, init_value=init_embed_weight), - TensorSpec("hidden_states", [token_count, D], torch.bfloat16, is_output=True), + TensorSpec("hidden_states", [token_count, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run MODES = {"decode": DECODE_TOKENS, "prefill": PREFILL_TOKENS} TEST_VOCAB_SIZE = 256 @@ -100,7 +100,7 @@ def init_embed_weight(): for mode_name in modes_to_run: token_count = MODES[mode_name] print(f"--- lookup_embedding_test {mode_name}: T={token_count} ---") - result = run_jit( + result = run( fn=lookup_embedding_test, specs=build_tensor_specs(token_count, TEST_VOCAB_SIZE), golden_fn=golden_lookup_embedding_test, diff --git a/models/deepseek_v4_flash_mtp/moe.py b/models/deepseek_v4_flash_mtp/moe.py index 2a6b4682f..57ab27fd3 100644 --- a/models/deepseek_v4_flash_mtp/moe.py +++ b/models/deepseek_v4_flash_mtp/moe.py @@ -938,7 +938,7 @@ def init_input_ids(): TensorSpec("shared_w3_scale", [N_RANKS, MOE_INTER], torch.float32, init_value=lambda: sw3_s), TensorSpec("shared_w2", [N_RANKS, D, MOE_INTER], torch.int8, init_value=lambda: sw2_i8), TensorSpec("shared_w2_scale", [N_RANKS, D], torch.float32, init_value=lambda: sw2_s), - TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32), ScalarSpec("layer_id", torch.int32, layer_id), ScalarSpec("num_tokens", torch.int32, num_tokens), ] @@ -950,7 +950,7 @@ def init_input_ids(): # H2D/D2H. Covers the routed/shared expert weights and their scales, the gate, # the HC-FFN constants, the RMSNorm gamma, and the static tid2eid route table — # but NOT the per-step activation (x_hc), per-step input_ids, or the output. - # All resident names are inputs (is_output=False), so the flag is always valid. + # All resident names are pure inputs, so the flag is always valid. RESIDENT_WEIGHT_NAMES = frozenset([ "hc_ffn_fn", "hc_ffn_scale", "hc_ffn_base", "norm_w", "gate_w", "gate_bias", "tid2eid", @@ -969,7 +969,7 @@ def init_input_ids(): if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1000,7 +1000,7 @@ def init_input_ids(): golden_data = args.golden_data - result = run_jit( + result = run( fn=l3_moe, specs=build_tensor_specs( layer_id=args.layer_id, diff --git a/models/deepseek_v4_flash_mtp/mtp_projection.py b/models/deepseek_v4_flash_mtp/mtp_projection.py index 177d77383..0d08a03c7 100644 --- a/models/deepseek_v4_flash_mtp/mtp_projection.py +++ b/models/deepseek_v4_flash_mtp/mtp_projection.py @@ -329,13 +329,13 @@ def init_h_proj_w_scale(): TensorSpec("h_proj_w", [D, D], torch.int8, init_value=init_h_proj_w), TensorSpec("h_proj_w_scale", [D], torch.float32, init_value=init_h_proj_w_scale), TensorSpec("h_proj_smooth", [D], torch.float32, init_value=lambda: torch.ones(D)), - TensorSpec("hidden_states_out", [t, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("hidden_states_out", [t, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument( @@ -357,7 +357,7 @@ def init_h_proj_w_scale(): } for mode in (modes if args.mode == "all" else [args.mode]): batch, seq = modes[mode] - result = run_jit( + result = run( fn=mtp_projection_test, specs=build_tensor_specs(batch, seq), golden_fn=golden_mtp_projection, diff --git a/models/deepseek_v4_flash_mtp/prefill_compressor_ratio128.py b/models/deepseek_v4_flash_mtp/prefill_compressor_ratio128.py index 9b8898a14..1917d6269 100644 --- a/models/deepseek_v4_flash_mtp/prefill_compressor_ratio128.py +++ b/models/deepseek_v4_flash_mtp/prefill_compressor_ratio128.py @@ -424,7 +424,7 @@ def init_state_slot_mapping(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), - TensorSpec("compress_state", [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("compress_state", [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [HCA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -432,7 +432,7 @@ def init_state_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), - TensorSpec("cmp_kv", [HCA_CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, is_output=True), + TensorSpec("cmp_kv", [HCA_CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, num_tokens), TensorSpec("cmp_slot_mapping", [T], torch.int64, init_value=init_cmp_slot_mapping), @@ -442,7 +442,7 @@ def init_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill compressor ratio128 validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -467,7 +467,7 @@ def init_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_compressor_ratio128_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_prefill_compressor_ratio128, diff --git a/models/deepseek_v4_flash_mtp/prefill_compressor_ratio4.py b/models/deepseek_v4_flash_mtp/prefill_compressor_ratio4.py index 04eabc8af..212a6532b 100644 --- a/models/deepseek_v4_flash_mtp/prefill_compressor_ratio4.py +++ b/models/deepseek_v4_flash_mtp/prefill_compressor_ratio4.py @@ -550,7 +550,7 @@ def init_state_slot_mapping(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), - TensorSpec("compress_state", [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("compress_state", [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [CSA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -558,7 +558,7 @@ def init_state_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), - TensorSpec("cmp_kv", [PREFILL_CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, is_output=True), + TensorSpec("cmp_kv", [PREFILL_CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, T), TensorSpec("cmp_slot_mapping", [T], torch.int64, init_value=init_cmp_slot_mapping), @@ -568,7 +568,7 @@ def init_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill compressor ratio4 validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -581,7 +581,7 @@ def init_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_compressor_ratio4_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_prefill_compressor_ratio4, diff --git a/models/deepseek_v4_flash_mtp/prefill_cp_csa_draft.py b/models/deepseek_v4_flash_mtp/prefill_cp_csa_draft.py index 61b4ad299..ddab2a4b9 100644 --- a/models/deepseek_v4_flash_mtp/prefill_cp_csa_draft.py +++ b/models/deepseek_v4_flash_mtp/prefill_cp_csa_draft.py @@ -896,15 +896,6 @@ def build_tensor_specs(cp_size: int = CP_SIZE): list(value.shape), value.dtype, init_value=value, - is_output=name - in { - "compress_state", - "inner_compress_state", - "kv_cache", - "cmp_kv", - "idx_kv_cache", - "idx_kv_scale", - }, ) ) for name, value in raw.items(): @@ -1020,7 +1011,7 @@ def build_tensor_specs(cp_size: int = CP_SIZE): ) specs.append( TensorSpec( - "x_out", list(x_hc.shape), torch.float32, is_output=True + "x_out", list(x_hc.shape), torch.float32 ) ) golden_prefill_cp_csa._ctx = { @@ -3288,14 +3279,14 @@ def golden_prefill_cp_csa(tensors): parser.add_argument("--dump-passes", action="store_true") parser.add_argument("--enable-chip-swimlane", action="store_true") args = parser.parse_args() - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run if args.cp != CP_SIZE: raise SystemExit(f"--cp={args.cp} does not match import-time CP_SIZE={CP_SIZE}") device_ids = [int(device) for device in args.device.split(",")] if len(device_ids) < args.cp: raise SystemExit(f"CP{args.cp} requires {args.cp} devices, got {device_ids}") - result = run_jit( + result = run( fn=prefill_cp_csa_test, specs=build_tensor_specs(args.cp), golden_fn=golden_prefill_cp_csa, diff --git a/models/deepseek_v4_flash_mtp/prefill_cp_fwd_draft.py b/models/deepseek_v4_flash_mtp/prefill_cp_fwd_draft.py index 0bff70d33..939f2db50 100644 --- a/models/deepseek_v4_flash_mtp/prefill_cp_fwd_draft.py +++ b/models/deepseek_v4_flash_mtp/prefill_cp_fwd_draft.py @@ -187,7 +187,7 @@ def _parse_static_bool(name: str, default: bool = False) -> bool: STATE_RECORDS_PER_WINDOW, STATE_WINDOW_ROWS, ) -from golden import TensorSpec, run_jit +from golden import TensorSpec, run # Phase 3 final tail: HC head + final RMSNorm (inlined in the FWD child) and # the LM head (host-launched per rank). All are accepted leaf math; only the # composition is added here. @@ -2362,7 +2362,7 @@ def _stack_spec(spec: TensorSpec, num_layers: int) -> TensorSpec: shape[0] = num_layers * shape[0] return TensorSpec( spec.name, shape, spec.dtype, - init_value=spec.init_value, is_output=spec.is_output, + init_value=spec.init_value, resident=spec.resident, ) @@ -2379,7 +2379,7 @@ def _make_stacked_swa_attn_spec(name: str, base_spec: TensorSpec, shape = [num_layers * base_spec.shape[0]] + list(base_spec.shape[1:]) return TensorSpec( name, shape, base_spec.dtype, init_value=stacked, - is_output=base_spec.is_output, resident=base_spec.resident, + resident=base_spec.resident, ) @@ -2397,7 +2397,7 @@ def _make_stacked_moe_spec(name: str, base_spec: TensorSpec, ) return TensorSpec( name, shape, base_spec.dtype, init_value=stacked, - is_output=base_spec.is_output, resident=base_spec.resident, + resident=base_spec.resident, ) @@ -2422,10 +2422,10 @@ def _build_input_ids_spec(cp_size: int, active_lengths_spec, prefix_seed: int): def _rename_spec(base_spec: TensorSpec, new_name: str) -> TensorSpec: """Return a copy of ``base_spec`` under a new name, preserving shape, - dtype, init_value, is_output, and resident.""" + dtype, init_value, and resident.""" return TensorSpec( new_name, list(base_spec.shape), base_spec.dtype, - init_value=base_spec.init_value, is_output=base_spec.is_output, + init_value=base_spec.init_value, resident=base_spec.resident, ) @@ -2463,7 +2463,7 @@ def _stack_type_spec( ) return TensorSpec( name, shape, base_spec.dtype, init_value=stacked, - is_output=base_spec.is_output, resident=base_spec.resident, + resident=base_spec.resident, ) @@ -2529,7 +2529,7 @@ def build_tensor_specs(cp_size: int = CP_SIZE): specs_by_name["kv_cache"] = TensorSpec( "kv_cache", [cp_size, FWD_NUM_LAYERS * ORI_MAX_BLOCKS, BLOCK_ROWS, 1, HEAD_DIM], - torch.bfloat16, init_value=cache_fwd, is_output=True, + torch.bfloat16, init_value=cache_fwd, ) # cmp_kv: compressed KV cache, FWD_NUM_LAYERS per-layer pools (every @@ -2543,7 +2543,7 @@ def build_tensor_specs(cp_size: int = CP_SIZE): specs_by_name["cmp_kv"] = TensorSpec( "cmp_kv", [cp_size, FWD_NUM_LAYERS * PREFILL_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], - torch.bfloat16, init_value=cmp_kv_fwd, is_output=True, + torch.bfloat16, init_value=cmp_kv_fwd, ) # --- HCA type-specific (layers 3 and 5) -------------------------------- @@ -2747,15 +2747,14 @@ def _init_logit_row_indices(): specs_by_name["pre_hc_hidden_out"] = TensorSpec( "pre_hc_hidden_out", [cp_size, LOCAL_PARTS, MAX_SEGMENT_TILES, T, HC_MULT, D], - torch.float32, is_output=True, + torch.float32, ) specs_by_name["hidden_out"] = TensorSpec( "hidden_out", [cp_size, LOCAL_ROWS, D], torch.bfloat16, - is_output=True, ) specs_by_name["logits"] = TensorSpec( "logits", [cp_size, LM_HEAD_MAX_LOGIT_ROWS, LM_HEAD_VOCAB], - torch.float32, is_output=True, + torch.float32, ) # Verify the host ABI matches. @@ -2773,7 +2772,7 @@ def _init_logit_row_indices(): # Harness-only Phase-3 tail sanity comparators (§8.9). # # The multi-layer FWD has no mathematical golden. Under --check-outputs the -# harness supplies a no-op golden_fn (so every is_output spec gets a +# harness supplies a no-op golden_fn (so every output spec gets a # zero-filled expected tensor) and a per-output compare_fn: persistent outputs # pass through _accept_output, and the three Phase-3 tail outputs # (pre_hc_hidden_out, hidden_out, logits) are inspected for finite + nonzero @@ -2964,6 +2963,17 @@ def _check_logits(actual, _expected, *, inputs, **_kwargs): return True, "logits sanity OK" +# Every pl.Out / pl.InOut parameter of l3_cp_prefill_fwd. Spelled out because +# compare_fn is built before compilation, and a spec learns its direction only +# once the harness stamps it from the compiled artifact. +_OUTPUT_NAMES = ( + "kv_cache", "cmp_kv", "idx_kv_cache", "idx_kv_scale", + "hca_compress_state", "csa_compress_state", "csa_inner_compress_state", + "stage_token", "completion_anchor", + "pre_hc_hidden_out", "hidden_out", "logits", +) + + def _build_outputs_compare_fn(specs): """Build a compare_fn dict: persistent outputs pass through _accept_output; the three Phase-3 tail outputs use dedicated sanity @@ -2972,7 +2982,7 @@ def _build_outputs_compare_fn(specs): for spec in specs: if not isinstance(spec, TensorSpec): continue - if not spec.is_output: + if spec.name not in _OUTPUT_NAMES: continue if spec.name == "pre_hc_hidden_out": compare_fn[spec.name] = _check_pre_hc_hidden_out @@ -3090,7 +3100,7 @@ def _build_outputs_compare_fn(specs): specs, ctx = build_tensor_specs(cp_size=args.cp) if args.check_outputs: - # No-op golden fills every is_output spec with zeros; the compare_fn + # No-op golden fills every output spec with zeros; the compare_fn # ignores expected for the tail outputs and pass-throughs for the # persistent caches. golden_fn = lambda _scratch: None @@ -3099,7 +3109,7 @@ def _build_outputs_compare_fn(specs): golden_fn = None compare_fn = None - result = run_jit( + result = run( fn=l3_prefill_cp_fwd, specs=specs, golden_fn=golden_fn, diff --git a/models/deepseek_v4_flash_mtp/prefill_cp_hca_draft.py b/models/deepseek_v4_flash_mtp/prefill_cp_hca_draft.py index 5d1ae3045..bb5c90cf3 100644 --- a/models/deepseek_v4_flash_mtp/prefill_cp_hca_draft.py +++ b/models/deepseek_v4_flash_mtp/prefill_cp_hca_draft.py @@ -1668,7 +1668,6 @@ def build_tensor_specs(cp_size: int = CP_SIZE): list(state.shape), state.dtype, init_value=state, - is_output=True, ), TensorSpec( "compress_state_block_table", @@ -1681,14 +1680,12 @@ def build_tensor_specs(cp_size: int = CP_SIZE): list(kv_cache.shape), kv_cache.dtype, init_value=kv_cache, - is_output=True, ), TensorSpec( "cmp_kv", list(cmp_cache.shape), cmp_cache.dtype, init_value=cmp_cache, - is_output=True, ), TensorSpec( "cmp_block_table", @@ -1740,7 +1737,6 @@ def build_tensor_specs(cp_size: int = CP_SIZE): "x_out", list(x_hc.shape), torch.float32, - is_output=True, ) ) @@ -2129,12 +2125,12 @@ def golden_prefill_cp_hca(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true") args = parser.parse_args() - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run device_ids = [int(device) for device in args.device.split(",")] if len(device_ids) < args.cp: raise SystemExit(f"CP{args.cp} requires {args.cp} devices, got {device_ids}") - result = run_jit( + result = run( fn=prefill_cp_hca_test, specs=build_tensor_specs(args.cp), golden_fn=golden_prefill_cp_hca, diff --git a/models/deepseek_v4_flash_mtp/prefill_cp_layer_draft.py b/models/deepseek_v4_flash_mtp/prefill_cp_layer_draft.py index 856dc95c2..0cfb35d28 100644 --- a/models/deepseek_v4_flash_mtp/prefill_cp_layer_draft.py +++ b/models/deepseek_v4_flash_mtp/prefill_cp_layer_draft.py @@ -114,7 +114,7 @@ golden_prefill_cp_csa, prefill_cp_csa_core, ) -from golden import ScalarSpec, TensorSpec, ratio_allclose, ratio_reldiff, run_jit +from golden import ScalarSpec, TensorSpec, ratio_allclose, ratio_reldiff, run # --------------------------------------------------------------------------- # Static CP/EP contract @@ -1982,7 +1982,7 @@ def _build_swa_specs(layer_id: int, cp_size: int): x_next_spec = TensorSpec( "x_next", [cp_size, LOCAL_PARTS, MAX_SEGMENT_TILES, T, HC_MULT, D], - torch.float32, is_output=True, + torch.float32, ) layer_id_spec = ScalarSpec("layer_id", torch.int32, layer_id) @@ -2018,7 +2018,7 @@ def _build_hca_specs(layer_id: int, cp_size: int): x_next_spec = TensorSpec( "x_next", [cp_size, LOCAL_PARTS, MAX_SEGMENT_TILES, T, HC_MULT, D], - torch.float32, is_output=True, + torch.float32, ) layer_id_spec = ScalarSpec("layer_id", torch.int32, layer_id) @@ -2080,7 +2080,7 @@ def _build_csa_specs(layer_id: int, cp_size: int): x_next_spec = TensorSpec( "x_next", [cp_size, LOCAL_PARTS, MAX_SEGMENT_TILES, T, HC_MULT, D], - torch.float32, is_output=True, + torch.float32, ) layer_id_spec = ScalarSpec("layer_id", torch.int32, layer_id) @@ -2173,14 +2173,14 @@ def golden_prefill_layer_cp(tensors): ) if layer_id == SWA_LAYER_ID: - # SWA golden context is installed by __main__ before run_jit. + # SWA golden context is installed by __main__ before run. swa_tensors = dict(tensors) swa_tensors["x_out"] = x_attn golden_prefill_cp_swa(swa_tensors) x_attn = swa_tensors["x_out"] # SWA golden mutates tensors["kv_cache"] in place; that stands. elif layer_id == HCA_LAYER_ID: - # HCA golden context is installed by __main__ before run_jit. + # HCA golden context is installed by __main__ before run. hca_tensors = dict(tensors) hca_tensors["x_out"] = x_attn golden_prefill_cp_hca(hca_tensors) @@ -2292,7 +2292,7 @@ def golden_prefill_layer_cp(tensors): else: raise RuntimeError(f"unsupported layer_id={args.layer_id}") - result = run_jit( + result = run( fn=host_fn, specs=specs, golden_fn=None if args.no_golden else golden_prefill_layer_cp, diff --git a/models/deepseek_v4_flash_mtp/prefill_cp_swa_draft.py b/models/deepseek_v4_flash_mtp/prefill_cp_swa_draft.py index e9bc26536..17056419d 100644 --- a/models/deepseek_v4_flash_mtp/prefill_cp_swa_draft.py +++ b/models/deepseek_v4_flash_mtp/prefill_cp_swa_draft.py @@ -805,7 +805,7 @@ def build_tensor_specs(cp_size: int = CP_SIZE): "freqs_cos", "freqs_sin", ): specs.append(TensorSpec(name, list(base[name].shape), base[name].dtype, init_value=base[name])) - specs.append(TensorSpec("kv_cache", list(cache.shape), torch.bfloat16, init_value=cache, is_output=True)) + specs.append(TensorSpec("kv_cache", list(cache.shape), torch.bfloat16, init_value=cache)) for name in tail_names: specs.append(TensorSpec(name, list(base[name].shape), base[name].dtype, init_value=base[name])) segment_starts = meta["segment_starts"] @@ -816,14 +816,14 @@ def build_tensor_specs(cp_size: int = CP_SIZE): ): value = meta[name] specs.append(TensorSpec(name, list(value.shape), value.dtype, init_value=value)) - # Spec order must match the kernel signature: run_jit binds its dummy compile + # Spec order must match the kernel signature: run binds its dummy compile # args positionally, so owner_rank_table sits between reverse_index and the # final_win_* triple exactly as prefill_cp_swa_test declares them. specs.append(TensorSpec("reverse_index", list(meta["reverse_index"].shape), meta["reverse_index"].dtype, init_value=meta["reverse_index"])) specs.append(TensorSpec("owner_rank_table", list(owner_rank.shape), owner_rank.dtype, init_value=owner_rank)) for name in ("final_win_seg_src", "final_win_row_src", "final_slot_mapping"): specs.append(TensorSpec(name, list(meta[name].shape), meta[name].dtype, init_value=meta[name])) - specs.append(TensorSpec("x_out", list(x.shape), torch.float32, is_output=True)) + specs.append(TensorSpec("x_out", list(x.shape), torch.float32)) return specs, ctx @@ -968,14 +968,14 @@ def golden_prefill_cp_swa(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run device_ids = [int(device) for device in args.device.split(",")] if len(device_ids) < args.cp: raise SystemExit(f"CP{args.cp} requires {args.cp} devices, got {device_ids}") specs, ctx = build_tensor_specs(args.cp) golden_prefill_cp_swa._ctx = ctx - result = run_jit( + result = run( fn=prefill_cp_swa_test, specs=specs, golden_fn=golden_prefill_cp_swa, diff --git a/models/deepseek_v4_flash_mtp/prefill_cp_zigzag.py b/models/deepseek_v4_flash_mtp/prefill_cp_zigzag.py index 42c9fc99e..7f327d843 100644 --- a/models/deepseek_v4_flash_mtp/prefill_cp_zigzag.py +++ b/models/deepseek_v4_flash_mtp/prefill_cp_zigzag.py @@ -380,14 +380,14 @@ def build_tensor_specs(): TensorSpec("owner_rank_table", [NUM_SEGMENTS], torch.int32, init_value=owner_rank), TensorSpec("final_win_seg_src", [TAIL_ROWS], torch.int32, init_value=final_segment.to(torch.int32)), TensorSpec("final_win_row_src", [TAIL_ROWS], torch.int32, init_value=final_row.to(torch.int32)), - TensorSpec("logical_tails_out", [CP_SIZE, CP_TAIL_WINDOW_ROWS, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("decode_raw_window_out", [CP_SIZE, TAIL_ROWS, HEAD_DIM], torch.bfloat16, is_output=True), + TensorSpec("logical_tails_out", [CP_SIZE, CP_TAIL_WINDOW_ROWS, HEAD_DIM], torch.bfloat16), + TensorSpec("decode_raw_window_out", [CP_SIZE, TAIL_ROWS, HEAD_DIM], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description="Standalone context-parallel zigzag exchange test.") parser.add_argument("-p", "--platform", default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -402,7 +402,7 @@ def build_tensor_specs(): if len(device_ids) < args.cp: raise SystemExit(f"CP{args.cp} requires {args.cp} devices, got {device_ids}") - result = run_jit( + result = run( fn=prefill_cp_zigzag_kv_tail_exchange_test, specs=build_tensor_specs(), golden_fn=golden_prefill_cp_zigzag_kv_tail_exchange, diff --git a/models/deepseek_v4_flash_mtp/prefill_csa.py b/models/deepseek_v4_flash_mtp/prefill_csa.py index 2eaf1af10..1380cf6b6 100644 --- a/models/deepseek_v4_flash_mtp/prefill_csa.py +++ b/models/deepseek_v4_flash_mtp/prefill_csa.py @@ -900,7 +900,7 @@ def init_wo_b(): ), TensorSpec("inner_compress_state_block_table", [INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), TensorSpec("kv_cache", [CSA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, - init_value=init_kv_cache, is_output=True), + init_value=init_kv_cache), TensorSpec("ori_block_table", [SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), # Compressor / indexer caches are written in-place but not validated here @@ -911,7 +911,6 @@ def init_wo_b(): [CSA_CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("cmp_block_table", [SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec( @@ -919,14 +918,12 @@ def init_wo_b(): [PREFILL_IDX_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, - is_output=True, ), TensorSpec( "idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, - is_output=True, ), TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), @@ -938,7 +935,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ScalarSpec("num_tokens", torch.int32, num_tokens), ] @@ -956,7 +953,7 @@ def _quant_w_per_output_channel_local(w): if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser(description="Standalone DeepSeek V4 packed prefill CSA correctness test.") parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -980,7 +977,7 @@ def _quant_w_per_output_channel_local(w): # elements), but keep the 0.5% fraction bar identical to full prefill. x_out_diff_thd, x_out_max_diff = (8e-3, 2) if args.start_pos else (5e-3, 1) - result = run_jit( + result = run( fn=prefill_attention_csa_test, specs=build_tensor_specs( args.start_pos, diff --git a/models/deepseek_v4_flash_mtp/prefill_fwd.py b/models/deepseek_v4_flash_mtp/prefill_fwd.py index 3d684cce4..15c567727 100644 --- a/models/deepseek_v4_flash_mtp/prefill_fwd.py +++ b/models/deepseek_v4_flash_mtp/prefill_fwd.py @@ -27,7 +27,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import run_jit +from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig # prefill_fwd is self-contained: it imports kernels, constants, and per-kind @@ -1053,7 +1053,6 @@ def init_value(): # cache). return TensorSpec( name, packed_shape, spec.dtype, init_value=init_value, - is_output=name in RESIDENT_CACHE_OUTPUT_NAMES, ) @@ -1125,7 +1124,7 @@ def init_value(): # Any remaining shared metadata: smoke zeros. return torch.zeros(list(spec.shape), dtype=spec.dtype) - return TensorSpec(name, list(spec.shape), spec.dtype, init_value=init_value, is_output=False) + return TensorSpec(name, list(spec.shape), spec.dtype, init_value=init_value) def _make_hc_head_spec(name): @@ -1402,7 +1401,6 @@ def kind_specs(build_fn): src.dtype, init_value=(_ranked_x_hc_init(src, N_RANKS, active_tokens, torch) if name == "x_hc" else _ranked_init(src, N_RANKS, torch)), - is_output=src.is_output, ) for name, src in attention_specs ] @@ -1435,7 +1433,7 @@ def init_input_ids(spec=spec): else: tensor_specs.append(spec) - tensor_specs.append(TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) + tensor_specs.append(TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32)) tensor_by_name = {spec.name: spec for spec in tensor_specs} missing = [name for name in HOST_TENSOR_ORDER if name not in tensor_by_name] if missing: @@ -1561,7 +1559,7 @@ def init_joined(parts=per_tile): def init_x_hc(shape=x_hc_shape, dtype=base.dtype): return (torch.randn(shape) * 0.05).to(dtype) - specs.append(TensorSpec(name, x_hc_shape, base.dtype, init_value=init_x_hc, is_output=False)) + specs.append(TensorSpec(name, x_hc_shape, base.dtype, init_value=init_x_hc)) elif name in TILED_NAMES: specs.append(make_tiled_shared_spec(name)) elif name in SHARED_NAMES: @@ -1577,12 +1575,12 @@ def init_x_hc(shape=x_hc_shape, dtype=base.dtype): # (child_memory): each shard uploaded once to its card and reused across # dispatches, skipping per-dispatch H2D/D2H. RESIDENT_WEIGHT_NAMES are static # weights; RESIDENT_CACHE_NAMES are the KV/state caches (the written kv_cache - # is also is_output=True and read back at the end via RESIDENT_CACHE_OUTPUT_NAMES). + # is also an InOut, read back at the end via RESIDENT_CACHE_OUTPUT_NAMES). for spec in specs: if spec.name in RESIDENT_WEIGHT_NAMES or spec.name in RESIDENT_CACHE_NAMES: spec.resident = "stacked" - specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) + specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32)) specs.append(TensorSpec( "lm_head_weight", [N_RANKS, VOCAB_PER_TP, D], @@ -1590,12 +1588,11 @@ def init_x_hc(shape=x_hc_shape, dtype=base.dtype): init_value=init_lm_head_weight, resident="stacked", )) - specs.append(TensorSpec("hidden_out", [N_RANKS, num_tiles * T, D], torch.bfloat16, is_output=True)) + specs.append(TensorSpec("hidden_out", [N_RANKS, num_tiles * T, D], torch.bfloat16)) specs.append(TensorSpec( "logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, - is_output=True, )) specs.append(TensorSpec( "num_tokens_per_owner", @@ -1666,7 +1663,7 @@ def main(): num_tiles=args.num_tiles, ) - result = run_jit( + result = run( fn=l3_prefill_fwd, specs=specs, golden_fn=None, diff --git a/models/deepseek_v4_flash_mtp/prefill_hca.py b/models/deepseek_v4_flash_mtp/prefill_hca.py index 9c4b5ace3..85c56b1e5 100644 --- a/models/deepseek_v4_flash_mtp/prefill_hca.py +++ b/models/deepseek_v4_flash_mtp/prefill_hca.py @@ -659,7 +659,6 @@ def init_wo_b(): [HCA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, - is_output=True, ), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("ori_block_table", [SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), @@ -668,7 +667,6 @@ def init_wo_b(): [HCA_CMP_BLOCK_NUM, CMP_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("cmp_block_table", [SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), @@ -678,14 +676,14 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ScalarSpec("num_tokens", torch.int32, num_tokens), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser(description="Standalone DeepSeek V4 packed prefill HCA correctness test.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -702,7 +700,7 @@ def init_wo_b(): args = parser.parse_args() compare_tokens = args.num_tokens - result = run_jit( + result = run( fn=prefill_attention_hca_test, specs=build_tensor_specs( args.start_pos, diff --git a/models/deepseek_v4_flash_mtp/prefill_indexer.py b/models/deepseek_v4_flash_mtp/prefill_indexer.py index 4524875f3..0540b0c1f 100644 --- a/models/deepseek_v4_flash_mtp/prefill_indexer.py +++ b/models/deepseek_v4_flash_mtp/prefill_indexer.py @@ -970,11 +970,11 @@ def init_sin(): TensorSpec("inner_wgate", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wgate), TensorSpec("inner_ape", [COMPRESS_RATIO, INNER_OUT_DIM], torch.float32, init_value=init_inner_ape), TensorSpec("inner_norm_w", [INNER_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), - TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, is_output=True), - TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, is_output=True), + TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache), + TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale), TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), - TensorSpec("score", [T, INDEXER_SCORE_CAP], torch.float32, is_output=True), - TensorSpec("topk_idxs", [T, INDEXER_TOPK_CAP], torch.int32, is_output=True), + TensorSpec("score", [T, INDEXER_SCORE_CAP], torch.float32), + TensorSpec("topk_idxs", [T, INDEXER_TOPK_CAP], torch.int32), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, num_tokens), TensorSpec("idx_slot_mapping", [T], torch.int64, init_value=init_idx_slot_mapping), @@ -985,7 +985,7 @@ def init_sin(): if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit, topk_pair_compare + from golden import ratio_allclose, run, topk_pair_compare parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill indexer validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1030,7 +1030,7 @@ def topk_idxs_compare(actual, expected, *, actual_outputs, expected_outputs, inp ) topk_idxs_compare.__name__ = "topk_pair_compare" - result = run_jit( + result = run( fn=prefill_indexer_test, specs=build_tensor_specs(args.start_pos, args.num_tokens), golden_fn=golden_prefill_indexer, diff --git a/models/deepseek_v4_flash_mtp/prefill_indexer_compressor.py b/models/deepseek_v4_flash_mtp/prefill_indexer_compressor.py index f349a849b..393afcaf5 100644 --- a/models/deepseek_v4_flash_mtp/prefill_indexer_compressor.py +++ b/models/deepseek_v4_flash_mtp/prefill_indexer_compressor.py @@ -743,8 +743,8 @@ def init_inner_state_slot_mapping(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [MAX_CMP_WRITES, HEAD_DIM], torch.int8, is_output=True), - TensorSpec("compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [MAX_CMP_WRITES, HEAD_DIM], torch.int8), + TensorSpec("compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("inner_compress_state_block_table", [INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -753,8 +753,8 @@ def init_inner_state_slot_mapping(): TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), TensorSpec("hadamard", [HEAD_DIM, HEAD_DIM], torch.bfloat16, init_value=init_hadamard), - TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, is_output=True), - TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, is_output=True), + TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache), + TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, IDX_STORAGE_BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale), TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, T), @@ -765,7 +765,7 @@ def init_inner_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill indexer compressor validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -783,7 +783,7 @@ def init_inner_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_indexer_compressor_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_prefill_indexer_compressor, diff --git a/models/deepseek_v4_flash_mtp/prefill_layer.py b/models/deepseek_v4_flash_mtp/prefill_layer.py index 31ba8ee11..0a884d9b0 100644 --- a/models/deepseek_v4_flash_mtp/prefill_layer.py +++ b/models/deepseek_v4_flash_mtp/prefill_layer.py @@ -846,7 +846,6 @@ def init(): [N_RANKS, *src.shape], src.dtype, init_value=make_init(), - is_output=packed_name in _HISTORY_CACHE_NAMES, ) ) @@ -866,7 +865,7 @@ def init_tid2eid(spec=spec): else: tensor_specs.append(spec) - tensor_specs.append(TensorSpec("x_next", [N_RANKS, total_tokens, HC_MULT, D], torch.float32, is_output=True)) + tensor_specs.append(TensorSpec("x_next", [N_RANKS, total_tokens, HC_MULT, D], torch.float32)) # Keep static weight parameters device-resident (child_memory), sharded per # rank. Cache/state/table tensors remain host tensors for output validation. @@ -1004,7 +1003,7 @@ def tile_buffer(packed_per_rank, rank, base, _valid, feature_shape, dtype): if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1024,7 +1023,7 @@ def tile_buffer(packed_per_rank, rank, base, _valid, feature_shape, dtype): device_ids = [int(d) for d in args.device.split(",")] assert len(device_ids) >= N_RANKS, f"need at least {N_RANKS} devices, got {device_ids}" - result = run_jit( + result = run( fn=l3_prefill_layer, specs=build_tensor_specs(layer_id=args.layer_id), golden_fn=golden_prefill_layer, diff --git a/models/deepseek_v4_flash_mtp/prefill_mtp.py b/models/deepseek_v4_flash_mtp/prefill_mtp.py index c1189e422..032618284 100644 --- a/models/deepseek_v4_flash_mtp/prefill_mtp.py +++ b/models/deepseek_v4_flash_mtp/prefill_mtp.py @@ -14,7 +14,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import ratio_allclose, ratio_reldiff, run_jit +from golden import ratio_allclose, ratio_reldiff, run from pypto.ir.distributed_compiled_program import DistributedConfig import config @@ -371,10 +371,10 @@ def l3_mtp_prefill_fwd( ) -def _ranked(spec, torch, is_output=False): +def _ranked(spec, torch): from golden import TensorSpec - return TensorSpec(spec.name, list(spec.shape), spec.dtype, init_value=spec.init_value, is_output=is_output) + return TensorSpec(spec.name, list(spec.shape), spec.dtype, init_value=spec.init_value) def _projection_specs(): @@ -542,7 +542,6 @@ def init_kv_cache(): cache_spec = TensorSpec( name, [N_RANKS, ori_block_num, BLOCK_SIZE, 1, HEAD_DIM], cache_dtype, init_value=init_kv_cache, - is_output=True, resident="stacked", ) specs.append(cache_spec) @@ -587,19 +586,16 @@ def init_ori_slot_mapping(): specs.append(lm_head_spec) hidden_out_spec = TensorSpec( "hidden_out", [N_RANKS, T, D], torch.bfloat16, - is_output=True, resident="stacked", ) specs.append(hidden_out_spec) pre_hc_hidden_spec = TensorSpec( "pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32, - is_output=True, resident="stacked", ) specs.append(pre_hc_hidden_spec) logits_spec = TensorSpec( "logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, - is_output=True, resident="stacked", ) specs.append(logits_spec) @@ -762,7 +758,7 @@ def main(): device_ids = [int(d) for d in args.device.split(",")] assert len(device_ids) >= N_RANKS, f"need at least {N_RANKS} devices, got {device_ids}" - result = run_jit( + result = run( fn=l3_mtp_prefill_fwd, specs=build_tensor_specs( start_pos=args.start_pos, diff --git a/models/deepseek_v4_flash_mtp/prefill_sparse_attn.py b/models/deepseek_v4_flash_mtp/prefill_sparse_attn.py index beedaf3cf..9333b68bd 100644 --- a/models/deepseek_v4_flash_mtp/prefill_sparse_attn.py +++ b/models/deepseek_v4_flash_mtp/prefill_sparse_attn.py @@ -728,13 +728,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -752,7 +752,7 @@ def init_wo_b(): args = parser.parse_args() torch.manual_seed(args.seed) - result = run_jit( + result = run( fn=prefill_sparse_attn_test, specs=build_tensor_specs(args.compress_ratio, args.num_tokens, args.ori_block_num, args.cmp_block_num), golden_fn=golden_prefill_sparse_attn, diff --git a/models/deepseek_v4_flash_mtp/prefill_swa.py b/models/deepseek_v4_flash_mtp/prefill_swa.py index 7ea440405..6bf103df7 100644 --- a/models/deepseek_v4_flash_mtp/prefill_swa.py +++ b/models/deepseek_v4_flash_mtp/prefill_swa.py @@ -482,7 +482,7 @@ def init_wo_b(): TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), TensorSpec("kv_cache", [BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, - init_value=init_kv_cache, is_output=True), + init_value=init_kv_cache), TensorSpec("block_table", [BLOCK_NUM], torch.int32, init_value=init_block_table), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), @@ -490,14 +490,14 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ScalarSpec("num_tokens", torch.int32, num_tokens), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, ratio_reldiff, run_jit + from golden import ratio_allclose, ratio_reldiff, run parser = argparse.ArgumentParser(description="Standalone DeepSeek V4 packed prefill SWA correctness test.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -514,7 +514,7 @@ def init_wo_b(): args = parser.parse_args() compare_tokens = args.num_tokens - result = run_jit( + result = run( fn=prefill_attention_swa_test, specs=build_tensor_specs( args.start_pos, diff --git a/models/deepseek_v4_flash_mtp/qkv_proj_rope.py b/models/deepseek_v4_flash_mtp/qkv_proj_rope.py index 6e1978859..138b8d145 100644 --- a/models/deepseek_v4_flash_mtp/qkv_proj_rope.py +++ b/models/deepseek_v4_flash_mtp/qkv_proj_rope.py @@ -537,16 +537,16 @@ def init_gamma_ckv(): TensorSpec("rope_sin", [T, ROPE_DIM], torch.bfloat16, init_value=init_sin), TensorSpec("gamma_cq", [Q_LORA], torch.bfloat16, init_value=init_gamma_cq), TensorSpec("gamma_ckv", [HEAD_DIM], torch.bfloat16, init_value=init_gamma_ckv), - TensorSpec("q", [T, H, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("kv", [T, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("qr", [T, Q_LORA], torch.int8, is_output=True), - TensorSpec("qr_scale", [T, 1], torch.float32, is_output=True), + TensorSpec("q", [T, H, HEAD_DIM], torch.bfloat16), + TensorSpec("kv", [T, HEAD_DIM], torch.bfloat16), + TensorSpec("qr", [T, Q_LORA], torch.int8), + TensorSpec("qr_scale", [T, 1], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -573,7 +573,7 @@ def init_gamma_ckv(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- qkv_proj_rope {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=qkv_proj_rope_test, specs=build_tensor_specs(B, S), golden_fn=golden_qkv_proj_rope, diff --git a/models/deepseek_v4_flash_mtp/rmsnorm.py b/models/deepseek_v4_flash_mtp/rmsnorm.py index 794e96734..ab93f11c6 100644 --- a/models/deepseek_v4_flash_mtp/rmsnorm.py +++ b/models/deepseek_v4_flash_mtp/rmsnorm.py @@ -104,13 +104,13 @@ def init_norm_w(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), TensorSpec("norm_w", [D], torch.bfloat16, init_value=init_norm_w), - TensorSpec("x_normed", [T, D], torch.bfloat16, is_output=True), + TensorSpec("x_normed", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -134,7 +134,7 @@ def init_norm_w(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- rms_norm_test {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=rms_norm_test, specs=build_tensor_specs(B, S), golden_fn=golden_rms_norm_test, diff --git a/models/deepseek_v4_flash_mtp/sample.py b/models/deepseek_v4_flash_mtp/sample.py index 7153db77d..458872bb5 100644 --- a/models/deepseek_v4_flash_mtp/sample.py +++ b/models/deepseek_v4_flash_mtp/sample.py @@ -607,7 +607,7 @@ def repeat_rows(values, dtype): torch.int32, init_value=lambda: repeat_rows([0, 1, 17, 1024, 4, 55, 4096, 32767], torch.int32), ), - TensorSpec("sampled_ids", [SAMPLE_ROWS, SAMPLED_IDS_PAD], torch.int32, is_output=True), + TensorSpec("sampled_ids", [SAMPLE_ROWS, SAMPLED_IDS_PAD], torch.int32), ] @@ -643,7 +643,7 @@ def golden_sample(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument( @@ -660,7 +660,7 @@ def golden_sample(tensors): if args.temperature is not None and args.temperature < 0.0: parser.error(f"--temperature must be non-negative, got {args.temperature}") - result = run_jit( + result = run( fn=sample_test, specs=build_tensor_specs(args.temperature, args.top_k), golden_fn=golden_sample, diff --git a/models/deepseek_v4_pro/decode_attention_csa.py b/models/deepseek_v4_pro/decode_attention_csa.py index cfe87b2be..8f7000647 100644 --- a/models/deepseek_v4_pro/decode_attention_csa.py +++ b/models/deepseek_v4_pro/decode_attention_csa.py @@ -843,7 +843,7 @@ def init_wo_b(): TensorSpec("cmp_wgate", [MAIN_OUT_DIM, D], torch.bfloat16, init_value=init_cmp_wgate), TensorSpec("cmp_ape", [COMPRESS_RATIO, MAIN_OUT_DIM], torch.float32, init_value=init_cmp_ape), TensorSpec("cmp_norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_cmp_norm_w), - TensorSpec("compress_state", [MAIN_STATE_BLOCK_NUM, MAIN_STATE_BLOCK_SIZE, MAIN_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("compress_state", [MAIN_STATE_BLOCK_NUM, MAIN_STATE_BLOCK_SIZE, MAIN_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, MAIN_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("idx_wq_b", [Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_wq_b_i8), TensorSpec("idx_wq_b_scale", [IDX_N_HEADS * IDX_HEAD_DIM], torch.float32, init_value=lambda: idx_wq_b_scale), @@ -853,13 +853,13 @@ def init_wo_b(): TensorSpec("inner_wgate", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wgate), TensorSpec("inner_ape", [COMPRESS_RATIO, INNER_OUT_DIM], torch.float32, init_value=init_inner_ape), TensorSpec("inner_norm_w", [IDX_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), - TensorSpec("inner_compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], torch.float32, init_value=init_inner_compress_state, is_output=True), + TensorSpec("inner_compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], torch.float32, init_value=init_inner_compress_state), TensorSpec("inner_compress_state_block_table", [B, INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), - TensorSpec("cmp_kv", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), + TensorSpec("cmp_kv", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("cmp_block_table", [B, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), - TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: shared_idx_kv_cache_i8.clone(), is_output=True), - TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: shared_idx_kv_scale.clone(), is_output=True), + TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: shared_idx_kv_cache_i8.clone()), + TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: shared_idx_kv_scale.clone()), TensorSpec("idx_block_table", [B, IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("window_swa_indices", [T, WIN], torch.int32, init_value=init_window_swa_indices), @@ -874,13 +874,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -896,7 +896,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=attention_csa_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_attention_csa, diff --git a/models/deepseek_v4_pro/decode_attention_hca.py b/models/deepseek_v4_pro/decode_attention_hca.py index 57471d815..a7fbdc635 100644 --- a/models/deepseek_v4_pro/decode_attention_hca.py +++ b/models/deepseek_v4_pro/decode_attention_hca.py @@ -627,10 +627,10 @@ def init_wo_b(): TensorSpec("cmp_wgate", [MAIN_OUT_DIM, D], torch.bfloat16, init_value=init_cmp_wgate), TensorSpec("cmp_ape", [COMPRESS_RATIO, MAIN_OUT_DIM], torch.float32, init_value=init_cmp_ape), TensorSpec("cmp_norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_cmp_norm_w), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), - TensorSpec("cmp_kv", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), + TensorSpec("cmp_kv", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("cmp_block_table", [B, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("window_swa_indices", [T, WIN], torch.int32, init_value=init_window_swa_indices), @@ -643,13 +643,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -664,7 +664,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=attention_hca_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_attention_hca, diff --git a/models/deepseek_v4_pro/decode_attention_swa.py b/models/deepseek_v4_pro/decode_attention_swa.py index bf7b47407..ffa987744 100644 --- a/models/deepseek_v4_pro/decode_attention_swa.py +++ b/models/deepseek_v4_pro/decode_attention_swa.py @@ -471,7 +471,7 @@ def init_wo_b(): TensorSpec("gamma_ckv", [HEAD_DIM], torch.bfloat16, init_value=init_gamma_ckv), TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), - TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), + TensorSpec("kv_cache", [ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache), TensorSpec("swa_slot_mapping", [T], torch.int64, init_value=init_swa_slot_mapping), TensorSpec("swa_indices", [T, WIN], torch.int32, init_value=init_swa_indices), TensorSpec("swa_lens", [T], torch.int32, init_value=init_swa_lens), @@ -480,13 +480,13 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -508,7 +508,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=attention_swa_test, specs=build_tensor_specs(args.start_pos, args.unmapped_visible_page_fixture), golden_fn=golden_attention_swa, diff --git a/models/deepseek_v4_pro/decode_compressor_ratio128.py b/models/deepseek_v4_pro/decode_compressor_ratio128.py index 28aa6621d..94ee1c361 100644 --- a/models/deepseek_v4_pro/decode_compressor_ratio128.py +++ b/models/deepseek_v4_pro/decode_compressor_ratio128.py @@ -487,8 +487,8 @@ def init_cmp_slot_mapping(): ) return [ TensorSpec("x", [B, S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [B, S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [B, S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -496,7 +496,7 @@ def init_cmp_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("cos", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_cos), TensorSpec("sin", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), - TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache, is_output=True), + TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("cmp_slot_mapping", [B, S], torch.int64, init_value=init_cmp_slot_mapping), TensorSpec("state_slot_mapping", [B, S], torch.int64, init_value=init_state_slot_mapping), @@ -505,7 +505,7 @@ def init_cmp_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -518,7 +518,7 @@ def init_cmp_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_pro/decode_compressor_ratio4.py b/models/deepseek_v4_pro/decode_compressor_ratio4.py index d8544827b..1e21a401c 100644 --- a/models/deepseek_v4_pro/decode_compressor_ratio4.py +++ b/models/deepseek_v4_pro/decode_compressor_ratio4.py @@ -497,8 +497,8 @@ def init_cmp_slot_mapping(): return [ TensorSpec("x", [B, S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [B, S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [B, S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -506,7 +506,7 @@ def init_cmp_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("cos", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_cos), TensorSpec("sin", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), - TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache, is_output=True), + TensorSpec("cmp_kv_cache", [CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv_cache), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("cmp_slot_mapping", [B, S], torch.int64, init_value=init_cmp_slot_mapping), TensorSpec("state_slot_mapping", [B, S], torch.int64, init_value=init_state_slot_mapping), @@ -515,7 +515,7 @@ def init_cmp_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -530,7 +530,7 @@ def init_cmp_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_pro/decode_fwd.py b/models/deepseek_v4_pro/decode_fwd.py index 742927c9d..451709285 100644 --- a/models/deepseek_v4_pro/decode_fwd.py +++ b/models/deepseek_v4_pro/decode_fwd.py @@ -15,7 +15,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import run_jit +from golden import run from hc_head import hc_head from input_pack import VOCAB_DYN as EMBED_VOCAB_DYN, pack_x_hc from lm_head import ( @@ -1050,13 +1050,12 @@ def init_value(): init_value=init_value, # Caches the kernel writes in place (kv_cache) are read back for # validation; every other stacked tensor is a plain input. - is_output=name in RESIDENT_CACHE_OUTPUT_NAMES, ) def _make_shared_spec(name, base_spec, out_name=None): from golden import TensorSpec - return TensorSpec(out_name or name, list(base_spec.shape), base_spec.dtype, init_value=base_spec.init_value if out_name is None else None, is_output=out_name is not None) + return TensorSpec(out_name or name, list(base_spec.shape), base_spec.dtype, init_value=base_spec.init_value if out_name is None else None) def _make_hc_head_spec(name): @@ -1515,7 +1514,6 @@ def init_block_table(): [N_RANKS, *spec.shape], spec.dtype, init_value=_ranked_init(spec, replicated=name in replicated_attention), - is_output=name == "kv_cache", ) for name, spec in attention_specs ] @@ -1543,7 +1541,7 @@ def init_input_ids(): specs.append(moe_tensor_specs[spec.name]) specs.extend([ - TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32), ScalarSpec("layer_id", torch.int32, layer_id), ]) return specs @@ -1603,7 +1601,7 @@ def init_logit_row_indices(): # (child_memory): each shard uploaded once to its card and reused across # dispatches, skipping per-dispatch H2D/D2H. RESIDENT_WEIGHT_NAMES are static # weights; CACHE_POOL_NAMES are the KV/state caches (the written kv_cache is - # also is_output=True and read back at the end via RESIDENT_CACHE_OUTPUT_NAMES). + # also an InOut, read back at the end via RESIDENT_CACHE_OUTPUT_NAMES). for spec in specs: if spec.name in RESIDENT_WEIGHT_NAMES or spec.name in CACHE_POOL_NAMES: spec.resident = "stacked" @@ -1621,19 +1619,17 @@ def init_logit_row_indices(): torch.int32, init_value=init_logit_row_indices, )) - specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) - specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16, is_output=True)) + specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32)) + specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16)) specs.append(TensorSpec( "logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, - is_output=True, )) specs.append(TensorSpec( "sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, - is_output=True, )) specs.append(ScalarSpec("num_tokens", torch.int32, num_tokens, compile_runtime=True)) specs.append(ScalarSpec("moe_epoch_base", torch.int32, 0, compile_runtime=True, benchmark_step=LAST_MOE_EPOCH)) @@ -1681,7 +1677,7 @@ def main(): count = apply_real_weights(specs, args.weights, ep=N_RANKS, tp=LM_HEAD_TP_SIZE) print(f"[RUN] real weights: {count} tensors from {args.weights}", flush=True) - result = run_jit( + result = run( fn=l3_decode_fwd, specs=specs, golden_fn=None, diff --git a/models/deepseek_v4_pro/decode_indexer.py b/models/deepseek_v4_pro/decode_indexer.py index 495b648e3..74cb3e9e1 100644 --- a/models/deepseek_v4_pro/decode_indexer.py +++ b/models/deepseek_v4_pro/decode_indexer.py @@ -1007,18 +1007,18 @@ def init_idx_slot_mapping(): TensorSpec("sin", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), TensorSpec("hadamard", [IDX_HEAD_DIM, IDX_HEAD_DIM], torch.bfloat16, init_value=init_hadamard), TensorSpec("inner_kv", [B, S, INNER_HEAD_DIM], torch.float32), - TensorSpec("inner_compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], torch.float32, init_value=init_inner_compress_state, is_output=True), + TensorSpec("inner_compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], torch.float32, init_value=init_inner_compress_state), TensorSpec("inner_compress_state_block_table", [B, INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), TensorSpec("inner_wkv", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wkv), TensorSpec("inner_wgate", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wgate), TensorSpec("inner_ape", [COMPRESS_RATIO, INNER_OUT_DIM], torch.float32, init_value=init_inner_ape), TensorSpec("inner_norm_w", [INNER_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), - TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_kv_i8, is_output=True), - TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: idx_kv_sc, is_output=True), + TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=lambda: idx_kv_i8), + TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=lambda: idx_kv_sc), TensorSpec("idx_block_table", [B, IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), # Output tails use -inf scores and -1 indices. - TensorSpec("score", [B, S, SCORE_LEN], torch.float32, is_output=True), - TensorSpec("topk_idxs", [B, S, SCORE_LEN], torch.int32, is_output=True), + TensorSpec("score", [B, S, SCORE_LEN], torch.float32), + TensorSpec("topk_idxs", [B, S, SCORE_LEN], torch.int32), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("idx_slot_mapping", [B, S], torch.int64, init_value=init_idx_slot_mapping), TensorSpec("inner_state_slot_mapping", [B, S], torch.int64, init_value=init_inner_state_slot_mapping), @@ -1033,7 +1033,7 @@ def init_idx_slot_mapping(): mapped_idx_cache_ratio_allclose, mapped_inner_state_ratio_allclose, ) - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1048,7 +1048,7 @@ def init_idx_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=indexer_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_indexer, diff --git a/models/deepseek_v4_pro/decode_indexer_compressor.py b/models/deepseek_v4_pro/decode_indexer_compressor.py index c4e8a2900..b99da5e63 100644 --- a/models/deepseek_v4_pro/decode_indexer_compressor.py +++ b/models/deepseek_v4_pro/decode_indexer_compressor.py @@ -602,8 +602,8 @@ def init_idx_slot_mapping(): return [ TensorSpec("x", [B, S, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [B, S, HEAD_DIM], torch.float32, is_output=True), - TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [B, S, HEAD_DIM], torch.float32), + TensorSpec("compress_state", [COMPRESS_STATE_BLOCK_NUM, COMPRESS_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [B, COMPRESS_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -612,8 +612,8 @@ def init_idx_slot_mapping(): TensorSpec("cos", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_cos), TensorSpec("sin", [B, ROPE_HEAD_DIM // 2], torch.float32, init_value=init_sin), TensorSpec("hadamard", [HEAD_DIM, HEAD_DIM], torch.bfloat16, init_value=init_hadamard), - TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, is_output=True), - TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, is_output=True), + TensorSpec("idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache), + TensorSpec("idx_kv_scale", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale), TensorSpec("position_ids", [B, S], torch.int32, init_value=init_position_ids), TensorSpec("idx_slot_mapping", [B, S], torch.int64, init_value=init_idx_slot_mapping), TensorSpec("inner_state_slot_mapping", [B, S], torch.int64, init_value=init_inner_state_slot_mapping), @@ -622,7 +622,7 @@ def init_idx_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -636,7 +636,7 @@ def init_idx_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=compressor_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_compressor, diff --git a/models/deepseek_v4_pro/decode_layer.py b/models/deepseek_v4_pro/decode_layer.py index 445a866fb..2ebcb6e6b 100644 --- a/models/deepseek_v4_pro/decode_layer.py +++ b/models/deepseek_v4_pro/decode_layer.py @@ -891,7 +891,6 @@ def init_block_table(): TensorSpec( name, [N_RANKS, *spec.shape], spec.dtype, init_value=_ranked_init(spec, replicated=name in replicated_attention), - is_output=name in mutable_cache_names, ) for name, spec in attention_specs ] @@ -944,7 +943,7 @@ def init_input_ids(): spec.resident = "stacked" specs.extend([ - TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32), ScalarSpec("layer_id", torch.int32, layer_id), ScalarSpec("moe_epoch", torch.int32, 1, compile_runtime=True, benchmark_step=1), ]) @@ -954,7 +953,7 @@ def init_input_ids(): if __name__ == "__main__": import argparse import torch - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1095,7 +1094,7 @@ def init_input_ids(): count = apply_real_layer_weights(specs, args.weights, layer_id=args.layer_id, ep=N_RANKS) print(f"[RUN] real weights: layer {args.layer_id}, {count} tensors from {args.weights}", flush=True) - result = run_jit( + result = run( fn=host_fn, specs=specs, golden_fn=golden_fn, diff --git a/models/deepseek_v4_pro/decode_mtp.py b/models/deepseek_v4_pro/decode_mtp.py index cd79d7af4..9dceca5ee 100644 --- a/models/deepseek_v4_pro/decode_mtp.py +++ b/models/deepseek_v4_pro/decode_mtp.py @@ -392,7 +392,7 @@ def init(): return init -def _ranked_spec(name, spec, *, replicated=False, is_output=False): +def _ranked_spec(name, spec, *, replicated=False): from golden import TensorSpec return TensorSpec( @@ -400,7 +400,6 @@ def _ranked_spec(name, spec, *, replicated=False, is_output=False): [N_RANKS, *spec.shape], spec.dtype, init_value=_ranked_init(spec, replicated=replicated), - is_output=is_output, ) @@ -621,7 +620,6 @@ def build_tensor_specs(start_pos=DECODE_START_POS, num_tokens=T): name, swa_specs[name], replicated=name in replicated_attention, - is_output=swa_specs[name].is_output, ) ) @@ -640,8 +638,8 @@ def build_tensor_specs(start_pos=DECODE_START_POS, num_tokens=T): if spec.name in resident_names: spec.resident = "stacked" - specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16, is_output=True)) - specs.append(TensorSpec("next_pre_hc_hidden", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) + specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16)) + specs.append(TensorSpec("next_pre_hc_hidden", [N_RANKS, T, HC_MULT, D], torch.float32)) specs.append(ScalarSpec("num_tokens", torch.int32, num_tokens)) return specs @@ -716,7 +714,7 @@ def golden_mtp_decode_layer(tensors): def main(): import torch - from golden import mapped_pool_ratio_reldiff, run_jit + from golden import mapped_pool_ratio_reldiff, run parser = argparse.ArgumentParser(description="DeepSeek-V4 MTP decode layer driver.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -740,7 +738,7 @@ def main(): device_ids = [int(d) for d in args.device.split(",")] assert len(device_ids) >= N_RANKS, f"need at least {N_RANKS} devices, got {device_ids}" - result = run_jit( + result = run( fn=l3_mtp_decode_layer, specs=build_tensor_specs(start_pos=args.start_pos, num_tokens=args.num_tokens), golden_fn=golden_mtp_decode_layer, diff --git a/models/deepseek_v4_pro/decode_sparse_attn.py b/models/deepseek_v4_pro/decode_sparse_attn.py index 35ef6e10a..f1fe66104 100644 --- a/models/deepseek_v4_pro/decode_sparse_attn.py +++ b/models/deepseek_v4_pro/decode_sparse_attn.py @@ -787,13 +787,13 @@ def init_wo_b_scale(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -821,7 +821,7 @@ def init_wo_b_scale(): summary = f"compress_ratio={compress_ratio} -> TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}" print(summary, flush=True) - result = run_jit( + result = run( fn=sparse_attn_test, specs=build_tensor_specs( compress_ratio, diff --git a/models/deepseek_v4_pro/decode_sparse_attn_hca.py b/models/deepseek_v4_pro/decode_sparse_attn_hca.py index ad420b94e..5f04626e7 100644 --- a/models/deepseek_v4_pro/decode_sparse_attn_hca.py +++ b/models/deepseek_v4_pro/decode_sparse_attn_hca.py @@ -707,13 +707,13 @@ def init_wo_b_scale(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -741,7 +741,7 @@ def init_wo_b_scale(): summary = f"compress_ratio={compress_ratio} -> TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}" print(summary, flush=True) - result = run_jit( + result = run( fn=sparse_attn_test, specs=build_tensor_specs( compress_ratio, diff --git a/models/deepseek_v4_pro/decode_sparse_attn_swa.py b/models/deepseek_v4_pro/decode_sparse_attn_swa.py index 1ce292a0f..509fa9a98 100644 --- a/models/deepseek_v4_pro/decode_sparse_attn_swa.py +++ b/models/deepseek_v4_pro/decode_sparse_attn_swa.py @@ -610,13 +610,13 @@ def init_wo_b_scale(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=init_wo_b), TensorSpec("wo_b_scale", [D], torch.float32, init_value=init_wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -643,7 +643,7 @@ def init_wo_b_scale(): print(f"TOPK={TOPK} SPARSE_BLOCKS={SPARSE_BLOCKS} PADDED_TOPK={PADDED_TOPK}", flush=True) - result = run_jit( + result = run( fn=sparse_attn_test, specs=build_tensor_specs( args.causal_regression_fixture, diff --git a/models/deepseek_v4_pro/expert_routed.py b/models/deepseek_v4_pro/expert_routed.py index 727edbe3d..d3f6b997d 100644 --- a/models/deepseek_v4_pro/expert_routed.py +++ b/models/deepseek_v4_pro/expert_routed.py @@ -475,7 +475,7 @@ def init_recv_weights(): TensorSpec("routed_w3_scale", [N_LOCAL_EXPERTS, MOE_INTER], torch.float32, init_value=lambda: w3_s), TensorSpec("routed_w2", [N_LOCAL_EXPERTS, D, MOE_INTER], torch.int8, init_value=lambda: w2_i8), TensorSpec("routed_w2_scale", [N_LOCAL_EXPERTS, D], torch.float32, init_value=lambda: w2_s), - TensorSpec("recv_y", [N_LOCAL_EXPERTS, RECV_MAX, D], torch.bfloat16, is_output=True), + TensorSpec("recv_y", [N_LOCAL_EXPERTS, RECV_MAX, D], torch.bfloat16), ] @@ -554,7 +554,7 @@ def compare(actual, expected, **kwargs): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -564,7 +564,7 @@ def compare(actual, expected, **kwargs): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=expert_routed_test, specs=build_tensor_specs(), golden_fn=golden_expert_routed, diff --git a/models/deepseek_v4_pro/expert_shared.py b/models/deepseek_v4_pro/expert_shared.py index 41e592b58..bd3907dd4 100644 --- a/models/deepseek_v4_pro/expert_shared.py +++ b/models/deepseek_v4_pro/expert_shared.py @@ -278,13 +278,13 @@ def build_tensor_specs(): TensorSpec("shared_w3_scale", [MOE_INTER], torch.float32, init_value=lambda: sw3_s), TensorSpec("shared_w2", [D, MOE_INTER], torch.int8, init_value=lambda: sw2_i8), TensorSpec("shared_w2_scale", [D], torch.float32, init_value=lambda: sw2_s), - TensorSpec("sh", [T, D], torch.bfloat16, is_output=True), + TensorSpec("sh", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -293,7 +293,7 @@ def build_tensor_specs(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=expert_shared_test, specs=build_tensor_specs(), golden_fn=golden_expert_shared, diff --git a/models/deepseek_v4_pro/gate.py b/models/deepseek_v4_pro/gate.py index ffaa4eab6..1cb56fadd 100644 --- a/models/deepseek_v4_pro/gate.py +++ b/models/deepseek_v4_pro/gate.py @@ -661,10 +661,10 @@ def init_input_ids(): ScalarSpec("num_tokens", torch.int32, num_tokens), TensorSpec("tid2eid", [VOCAB, TOPK], torch.int32, init_value=init_tid2eid), TensorSpec("input_ids", [T], torch.int64, init_value=init_input_ids), - TensorSpec("x_norm_i8", [T, D], torch.int8, is_output=True), - TensorSpec("x_norm_scale", [T, 1], torch.float32, is_output=True), - TensorSpec("indices", [T, TOPK], torch.int32, is_output=True), - TensorSpec("weights", [T, TOPK], torch.float32, is_output=True), + TensorSpec("x_norm_i8", [T, D], torch.int8), + TensorSpec("x_norm_scale", [T, 1], torch.float32), + TensorSpec("indices", [T, TOPK], torch.int32), + TensorSpec("weights", [T, TOPK], torch.float32), ] @@ -685,7 +685,7 @@ def gate_x_norm_scale_compare(num_tokens): if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -698,7 +698,7 @@ def gate_x_norm_scale_compare(num_tokens): args = parser.parse_args() torch.manual_seed(args.seed) - result = run_jit( + result = run( fn=gate_test, specs=build_tensor_specs(layer_id=args.layer_id, num_tokens=args.num_tokens), golden_fn=golden_gate_core, diff --git a/models/deepseek_v4_pro/hc_head.py b/models/deepseek_v4_pro/hc_head.py index da275ec06..033c6a6f4 100644 --- a/models/deepseek_v4_pro/hc_head.py +++ b/models/deepseek_v4_pro/hc_head.py @@ -250,14 +250,14 @@ def init_hc_head_fn(): "hc_head_base", [HC_MULT], torch.float32, init_value=lambda: torch.tensor([5.9166, -3.6223, -2.9324, -3.3124]), ), - TensorSpec("y", [T, D], torch.bfloat16, is_output=True), + TensorSpec("y", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -268,7 +268,7 @@ def init_hc_head_fn(): args = parser.parse_args() torch.manual_seed(args.seed) - result = run_jit( + result = run( fn=hc_head_test, specs=build_tensor_specs(), golden_fn=golden_hc_head, diff --git a/models/deepseek_v4_pro/hc_post.py b/models/deepseek_v4_pro/hc_post.py index 4f366050f..3a145116e 100644 --- a/models/deepseek_v4_pro/hc_post.py +++ b/models/deepseek_v4_pro/hc_post.py @@ -204,13 +204,13 @@ def init_comb(): TensorSpec("residual", [T, HC_MULT, D], torch.float32, init_value=init_residual), TensorSpec("post", [T, HC_MULT], torch.float32, init_value=init_post), TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32, init_value=init_comb), - TensorSpec("y", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("y", [T, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -232,7 +232,7 @@ def init_comb(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- hc_post {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=hc_post_test, specs=build_tensor_specs(B, S), golden_fn=golden_hc_post, diff --git a/models/deepseek_v4_pro/hc_pre.py b/models/deepseek_v4_pro/hc_pre.py index 622353397..726a2b235 100644 --- a/models/deepseek_v4_pro/hc_pre.py +++ b/models/deepseek_v4_pro/hc_pre.py @@ -833,15 +833,15 @@ def init_hc_base(): TensorSpec("hc_fn", [MIX_HC, HC_DIM], torch.float32, init_value=init_hc_fn), TensorSpec("hc_scale", [3], torch.float32, init_value=init_hc_scale), TensorSpec("hc_base", [MIX_HC], torch.float32, init_value=init_hc_base), - TensorSpec("x_mixed", [T, D], torch.bfloat16, is_output=True), - TensorSpec("post", [T, HC_MULT], torch.float32, is_output=True), - TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32, is_output=True), + TensorSpec("x_mixed", [T, D], torch.bfloat16), + TensorSpec("post", [T, HC_MULT], torch.float32), + TensorSpec("comb", [T, HC_MULT * HC_MULT], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -881,7 +881,7 @@ def init_hc_base(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- hc_pre {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=hc_pre_test, specs=build_tensor_specs(B, S), golden_fn=golden_hc_pre, diff --git a/models/deepseek_v4_pro/input_pack.py b/models/deepseek_v4_pro/input_pack.py index 99e74e555..f0889a667 100644 --- a/models/deepseek_v4_pro/input_pack.py +++ b/models/deepseek_v4_pro/input_pack.py @@ -102,14 +102,13 @@ def init_input_ids(): "x_hc", [token_count, HC_MULT, D], torch.float32, - is_output=True, ), ] if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run modes = {"decode": DECODE_TOKENS, "prefill": PREFILL_TOKENS} test_vocab_size = 256 @@ -131,7 +130,7 @@ def init_input_ids(): for mode_name in modes_to_run: token_count = modes[mode_name] print(f"--- pack_x_hc_test {mode_name}: T={token_count} ---") - result = run_jit( + result = run( fn=pack_x_hc_test, specs=build_tensor_specs(token_count, test_vocab_size), golden_fn=golden_pack_x_hc, diff --git a/models/deepseek_v4_pro/lm_head.py b/models/deepseek_v4_pro/lm_head.py index 65ad41bbd..5cc2c88fd 100644 --- a/models/deepseek_v4_pro/lm_head.py +++ b/models/deepseek_v4_pro/lm_head.py @@ -564,13 +564,11 @@ def init_logit_row_indices(): "logits", [DP_SIZE, MAX_LOGIT_ROWS, VOCAB], torch.float32, - is_output=True, ), TensorSpec( "sampled_ids", [DP_SIZE, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, - is_output=True, ), TensorSpec( "logit_row_indices", @@ -621,7 +619,7 @@ def compare_sampled_ids(actual, _expected, *, actual_outputs, **_): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -657,7 +655,7 @@ def compare_sampled_ids(actual, _expected, *, actual_outputs, **_): "sampled_ids": compare_sampled_ids, } - result = run_jit( + result = run( fn=fn, specs=specs, golden_fn=golden_fn, diff --git a/models/deepseek_v4_pro/moe.py b/models/deepseek_v4_pro/moe.py index c4cf73587..a58065acc 100644 --- a/models/deepseek_v4_pro/moe.py +++ b/models/deepseek_v4_pro/moe.py @@ -995,7 +995,7 @@ def init_input_ids(): TensorSpec("shared_w3_scale", [N_RANKS, MOE_INTER], torch.float32, init_value=lambda: sw3_s), TensorSpec("shared_w2", [N_RANKS, D, MOE_INTER], torch.int8, init_value=lambda: sw2_i8), TensorSpec("shared_w2_scale", [N_RANKS, D], torch.float32, init_value=lambda: sw2_s), - TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32), ScalarSpec("layer_id", torch.int32, layer_id), ScalarSpec("num_tokens", torch.int32, num_tokens), ScalarSpec("moe_epoch", torch.int32, 1, compile_runtime=True, benchmark_step=1), @@ -1008,7 +1008,7 @@ def init_input_ids(): # H2D/D2H. Covers the routed/shared expert weights and their scales, the gate, # the HC-FFN constants, the RMSNorm gamma, and the static tid2eid route table — # but NOT the per-step activation (x_hc), per-step input_ids, or the output. - # All resident names are inputs (is_output=False), so the flag is always valid. + # All resident names are pure inputs, so the flag is always valid. RESIDENT_WEIGHT_NAMES = frozenset([ "hc_ffn_fn", "hc_ffn_scale", "hc_ffn_base", "norm_w", "gate_w", "gate_bias", "tid2eid", @@ -1095,7 +1095,7 @@ def compare(actual, expected, **kwargs): import argparse import torch - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1143,7 +1143,7 @@ def compare(actual, expected, **kwargs): ), } - result = run_jit( + result = run( fn=l3_moe, specs=build_tensor_specs( layer_id=args.layer_id, diff --git a/models/deepseek_v4_pro/mtp_projection.py b/models/deepseek_v4_pro/mtp_projection.py index 2c802f1c8..181f60f40 100644 --- a/models/deepseek_v4_pro/mtp_projection.py +++ b/models/deepseek_v4_pro/mtp_projection.py @@ -312,14 +312,14 @@ def init_h_proj_w_scale(): TensorSpec("h_proj_w", [D, D], torch.int8, init_value=init_h_proj_w), TensorSpec("h_proj_w_scale", [D], torch.float32, init_value=init_h_proj_w_scale), TensorSpec("h_proj_smooth", [D], torch.float32, init_value=lambda: torch.ones(D)), - TensorSpec("hidden_states_out", [t, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("hidden_states_out", [t, HC_MULT, D], torch.float32), ] if __name__ == "__main__": import argparse import torch - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -337,7 +337,7 @@ def init_h_proj_w_scale(): } for mode in (modes if args.mode == "all" else [args.mode]): batch, seq = modes[mode] - result = run_jit( + result = run( fn=mtp_projection_test, specs=build_tensor_specs(batch, seq), golden_fn=golden_mtp_projection, diff --git a/models/deepseek_v4_pro/prefill_attention_csa.py b/models/deepseek_v4_pro/prefill_attention_csa.py index d2cc98c2c..163ea0b54 100644 --- a/models/deepseek_v4_pro/prefill_attention_csa.py +++ b/models/deepseek_v4_pro/prefill_attention_csa.py @@ -131,7 +131,7 @@ # ring; measured on a5, every ring at 512 MiB is already enough for all three prefill # attention variants, so 1 GiB is one doubling of headroom over the measured need. # All four rings, not just ring 2: ring 2 alone (what prefill_fwd.py sets) does not -# clear it. Applied through run_jit's runtime_cfg, which reaches the device only on +# clear it. Applied through run's runtime_cfg, which reaches the device only on # the ChipWorker route -- see golden/runner.py::_execute_via_runner. PREFILL_ATTN_RING_HEAP = (1024 * 1024 * 1024,) * 4 @@ -879,7 +879,6 @@ def init_wo_b(): [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, - is_output=True, ), TensorSpec("compress_state_block_table", [CSA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("hadamard_idx", [IDX_HEAD_DIM, IDX_HEAD_DIM], torch.bfloat16, init_value=init_hadamard_idx), @@ -895,11 +894,10 @@ def init_wo_b(): [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], torch.float32, init_value=init_inner_compress_state, - is_output=True, ), TensorSpec("inner_compress_state_block_table", [INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), TensorSpec("kv_cache", [CSA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, - init_value=init_kv_cache, is_output=True), + init_value=init_kv_cache), TensorSpec("ori_block_table", [SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), TensorSpec( @@ -907,7 +905,6 @@ def init_wo_b(): [CSA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("cmp_block_table", [SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec( @@ -915,14 +912,12 @@ def init_wo_b(): [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, - is_output=True, ), TensorSpec( "idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, - is_output=True, ), TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), @@ -934,7 +929,7 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ScalarSpec("num_tokens", torch.int32, num_tokens), ] @@ -952,7 +947,7 @@ def _quant_w_per_output_channel_local(w): if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser(description="Standalone DeepSeek V4 packed prefill CSA correctness test.") parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -978,7 +973,7 @@ def _quant_w_per_output_channel_local(w): from pypto.runtime import RunConfig - result = run_jit( + result = run( fn=prefill_attention_csa_test, specs=build_tensor_specs( args.start_pos, diff --git a/models/deepseek_v4_pro/prefill_attention_hca.py b/models/deepseek_v4_pro/prefill_attention_hca.py index ccad226e8..0c0f25e41 100644 --- a/models/deepseek_v4_pro/prefill_attention_hca.py +++ b/models/deepseek_v4_pro/prefill_attention_hca.py @@ -114,7 +114,7 @@ # ring; measured on a5, every ring at 512 MiB is already enough for all three prefill # attention variants, so 1 GiB is one doubling of headroom over the measured need. # All four rings, not just ring 2: ring 2 alone (what prefill_fwd.py sets) does not -# clear it. Applied through run_jit's runtime_cfg, which reaches the device only on +# clear it. Applied through run's runtime_cfg, which reaches the device only on # the ChipWorker route -- see golden/runner.py::_execute_via_runner. PREFILL_ATTN_RING_HEAP = (1024 * 1024 * 1024,) * 4 @@ -662,7 +662,6 @@ def init_wo_b(): [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, - is_output=True, ), TensorSpec("compress_state_block_table", [HCA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec( @@ -670,7 +669,6 @@ def init_wo_b(): [HCA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, - is_output=True, ), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("ori_block_table", [SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), @@ -679,7 +677,6 @@ def init_wo_b(): [HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, - is_output=True, ), TensorSpec("cmp_block_table", [SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), @@ -689,14 +686,14 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ScalarSpec("num_tokens", torch.int32, num_tokens), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser(description="Standalone DeepSeek V4 packed prefill HCA correctness test.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -715,7 +712,7 @@ def init_wo_b(): from pypto.runtime import RunConfig - result = run_jit( + result = run( fn=prefill_attention_hca_test, specs=build_tensor_specs( args.start_pos, diff --git a/models/deepseek_v4_pro/prefill_attention_swa.py b/models/deepseek_v4_pro/prefill_attention_swa.py index 5ff9b5fe7..1ebda9249 100644 --- a/models/deepseek_v4_pro/prefill_attention_swa.py +++ b/models/deepseek_v4_pro/prefill_attention_swa.py @@ -115,7 +115,7 @@ # ring; measured on a5, every ring at 512 MiB is already enough for all three prefill # attention variants, so 1 GiB is one doubling of headroom over the measured need. # All four rings, not just ring 2: ring 2 alone (what prefill_fwd.py sets) does not -# clear it. Applied through run_jit's runtime_cfg, which reaches the device only on +# clear it. Applied through run's runtime_cfg, which reaches the device only on # the ChipWorker route -- see golden/runner.py::_execute_via_runner. PREFILL_ATTN_RING_HEAP = (1024 * 1024 * 1024,) * 4 @@ -658,7 +658,7 @@ def init_wo_b(): TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), TensorSpec("kv_cache", [BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, - init_value=init_kv_cache, is_output=True), + init_value=init_kv_cache), TensorSpec("block_table", [BLOCK_NUM], torch.int32, init_value=init_block_table), TensorSpec("ori_slot_mapping", [T], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), @@ -666,14 +666,14 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("x_out", [T, HC_MULT, D], torch.float32, is_output=True), + TensorSpec("x_out", [T, HC_MULT, D], torch.float32), ScalarSpec("num_tokens", torch.int32, num_tokens), ] if __name__ == "__main__": import argparse - from golden import ratio_reldiff, run_jit + from golden import ratio_reldiff, run parser = argparse.ArgumentParser(description="Standalone DeepSeek V4 packed prefill SWA correctness test.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -692,7 +692,7 @@ def init_wo_b(): from pypto.runtime import RunConfig - result = run_jit( + result = run( fn=prefill_attention_swa_test, specs=build_tensor_specs( args.start_pos, diff --git a/models/deepseek_v4_pro/prefill_compressor_ratio128.py b/models/deepseek_v4_pro/prefill_compressor_ratio128.py index 200a4ecfb..2d10ae54a 100644 --- a/models/deepseek_v4_pro/prefill_compressor_ratio128.py +++ b/models/deepseek_v4_pro/prefill_compressor_ratio128.py @@ -443,7 +443,7 @@ def init_state_slot_mapping(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), - TensorSpec("compress_state", [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("compress_state", [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [HCA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -451,7 +451,7 @@ def init_state_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), - TensorSpec("cmp_kv", [HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, is_output=True), + TensorSpec("cmp_kv", [HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, num_tokens), TensorSpec("cmp_slot_mapping", [T], torch.int64, init_value=init_cmp_slot_mapping), @@ -461,7 +461,7 @@ def init_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill compressor ratio128 validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -486,7 +486,7 @@ def init_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_compressor_ratio128_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_prefill_compressor_ratio128, diff --git a/models/deepseek_v4_pro/prefill_compressor_ratio4.py b/models/deepseek_v4_pro/prefill_compressor_ratio4.py index 5c282521b..6bfc21ab0 100644 --- a/models/deepseek_v4_pro/prefill_compressor_ratio4.py +++ b/models/deepseek_v4_pro/prefill_compressor_ratio4.py @@ -569,7 +569,7 @@ def init_state_slot_mapping(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), - TensorSpec("compress_state", [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("compress_state", [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("compress_state_block_table", [CSA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -577,7 +577,7 @@ def init_state_slot_mapping(): TensorSpec("norm_w", [HEAD_DIM], torch.bfloat16, init_value=init_norm_w), TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), - TensorSpec("cmp_kv", [PREFILL_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv, is_output=True), + TensorSpec("cmp_kv", [PREFILL_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, T), TensorSpec("cmp_slot_mapping", [T], torch.int64, init_value=init_cmp_slot_mapping), @@ -587,7 +587,7 @@ def init_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill compressor ratio4 validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -605,7 +605,7 @@ def init_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_compressor_ratio4_test, specs=build_tensor_specs(args.start_pos), golden_fn=golden_prefill_compressor_ratio4, diff --git a/models/deepseek_v4_pro/prefill_fwd.py b/models/deepseek_v4_pro/prefill_fwd.py index dca5ef923..a453e6d21 100644 --- a/models/deepseek_v4_pro/prefill_fwd.py +++ b/models/deepseek_v4_pro/prefill_fwd.py @@ -29,7 +29,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import run_jit +from golden import run from pypto.ir.distributed_compiled_program import DistributedConfig from input_pack import VOCAB_DYN as EMBED_VOCAB_DYN, pack_x_hc @@ -1138,7 +1138,6 @@ def init_value(): # cache). return TensorSpec( name, packed_shape, spec.dtype, init_value=init_value, - is_output=name in RESIDENT_CACHE_OUTPUT_NAMES, ) @@ -1179,7 +1178,7 @@ def init_value(): # Any remaining shared metadata: smoke zeros. return torch.zeros(list(spec.shape), dtype=spec.dtype) - return TensorSpec(name, list(spec.shape), spec.dtype, init_value=init_value, is_output=False) + return TensorSpec(name, list(spec.shape), spec.dtype, init_value=init_value) def _make_hc_head_spec(name): @@ -1462,7 +1461,6 @@ def kind_specs(build_fn): src.dtype, init_value=(_ranked_x_hc_init(src, N_RANKS, active_tokens, torch) if name == "x_hc" else _ranked_init(src, N_RANKS, torch)), - is_output=src.is_output, ) for name, src in attention_specs ] @@ -1495,7 +1493,7 @@ def init_input_ids(spec=spec): else: tensor_specs.append(spec) - tensor_specs.append(TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) + tensor_specs.append(TensorSpec("x_next", [N_RANKS, T, HC_MULT, D], torch.float32)) tensor_by_name = {spec.name: spec for spec in tensor_specs} missing = [name for name in HOST_TENSOR_ORDER if name not in tensor_by_name] if missing: @@ -1580,13 +1578,13 @@ def init_logit_row_indices(): # (child_memory): each shard uploaded once to its card and reused across # dispatches, skipping per-dispatch H2D/D2H. RESIDENT_WEIGHT_NAMES are static # weights; RESIDENT_CACHE_NAMES are the KV/state caches (the written kv_cache - # is also is_output=True and read back at the end via RESIDENT_CACHE_OUTPUT_NAMES). + # is also an InOut, read back at the end via RESIDENT_CACHE_OUTPUT_NAMES). for spec in specs: if spec.name in RESIDENT_WEIGHT_NAMES or spec.name in RESIDENT_CACHE_NAMES: spec.resident = "stacked" - specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) - specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16, is_output=True)) + specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32)) + specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16)) specs.append(TensorSpec( "lm_head_weight", [N_RANKS, VOCAB_PER_TP, D], @@ -1598,13 +1596,11 @@ def init_logit_row_indices(): "logits", [N_RANKS, MAX_LOGIT_ROWS, LM_HEAD_VOCAB], torch.float32, - is_output=True, )) specs.append(TensorSpec( "sampled_ids", [N_RANKS, MAX_LOGIT_ROWS, SAMPLED_IDS_PAD], torch.int32, - is_output=True, )) specs.append(TensorSpec( "logit_row_indices", @@ -1726,7 +1722,7 @@ def main(): golden_fn = golden_prefill_fwd compare_fn = build_validate_compare_fn(args.num_tokens) - result = run_jit( + result = run( fn=l3_prefill_fwd, specs=specs, golden_fn=golden_fn, diff --git a/models/deepseek_v4_pro/prefill_indexer.py b/models/deepseek_v4_pro/prefill_indexer.py index c53668ae8..bf80c7370 100644 --- a/models/deepseek_v4_pro/prefill_indexer.py +++ b/models/deepseek_v4_pro/prefill_indexer.py @@ -1043,17 +1043,17 @@ def init_sin(): TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), TensorSpec("hadamard", [IDX_HEAD_DIM, IDX_HEAD_DIM], torch.bfloat16, init_value=init_hadamard), - TensorSpec("inner_compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], torch.float32, init_value=init_inner_compress_state, is_output=True), + TensorSpec("inner_compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], torch.float32, init_value=init_inner_compress_state), TensorSpec("inner_compress_state_block_table", [INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), TensorSpec("inner_wkv", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wkv), TensorSpec("inner_wgate", [INNER_OUT_DIM, D], torch.bfloat16, init_value=init_inner_wgate), TensorSpec("inner_ape", [COMPRESS_RATIO, INNER_OUT_DIM], torch.float32, init_value=init_inner_ape), TensorSpec("inner_norm_w", [INNER_HEAD_DIM], torch.bfloat16, init_value=init_inner_norm_w), - TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, is_output=True), - TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, is_output=True), + TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], torch.int8, init_value=init_idx_kv_cache), + TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale), TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), - TensorSpec("score", [T, INDEXER_SCORE_CAP], torch.float32, is_output=True), - TensorSpec("topk_idxs", [T, INDEXER_SCORE_CAP], torch.int32, is_output=True), + TensorSpec("score", [T, INDEXER_SCORE_CAP], torch.float32), + TensorSpec("topk_idxs", [T, INDEXER_SCORE_CAP], torch.int32), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, num_tokens), TensorSpec("idx_slot_mapping", [T], torch.int64, init_value=init_idx_slot_mapping), @@ -1063,7 +1063,7 @@ def init_sin(): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill indexer validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1083,7 +1083,7 @@ def init_sin(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_indexer_test, specs=build_tensor_specs(args.start_pos, args.num_tokens), golden_fn=golden_prefill_indexer, diff --git a/models/deepseek_v4_pro/prefill_indexer_compressor.py b/models/deepseek_v4_pro/prefill_indexer_compressor.py index ae237493f..e45d9933c 100644 --- a/models/deepseek_v4_pro/prefill_indexer_compressor.py +++ b/models/deepseek_v4_pro/prefill_indexer_compressor.py @@ -1079,8 +1079,8 @@ def init_inner_state_slot_mapping(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), - TensorSpec("kv", [MAX_CMP_WRITES, HEAD_DIM], torch.int8, is_output=True), - TensorSpec("compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state, is_output=True), + TensorSpec("kv", [MAX_CMP_WRITES, HEAD_DIM], torch.int8), + TensorSpec("compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], torch.float32, init_value=init_compress_state), TensorSpec("inner_compress_state_block_table", [INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table), TensorSpec("wkv", [OUT_DIM, D], torch.bfloat16, init_value=init_wkv), TensorSpec("wgate", [OUT_DIM, D], torch.bfloat16, init_value=init_wgate), @@ -1089,8 +1089,8 @@ def init_inner_state_slot_mapping(): TensorSpec("freqs_cos", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_cos), TensorSpec("freqs_sin", [MAX_SEQ_LEN, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), TensorSpec("hadamard", [HEAD_DIM, HEAD_DIM], torch.bfloat16, init_value=init_hadamard), - TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache, is_output=True), - TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale, is_output=True), + TensorSpec("idx_kv_cache", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.int8, init_value=init_idx_kv_cache), + TensorSpec("idx_kv_scale", [PREFILL_IDX_BLOCK_NUM, BLOCK_SIZE, 1, 1], torch.float32, init_value=init_idx_kv_scale), TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("position_ids", [T], torch.int32, init_value=init_position_ids), ScalarSpec("num_tokens", torch.int32, num_tokens), @@ -1101,7 +1101,7 @@ def init_inner_state_slot_mapping(): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser(description="Standalone token-major DeepSeek V4 prefill indexer compressor validation.") parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1121,7 +1121,7 @@ def init_inner_state_slot_mapping(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_indexer_compressor_test, specs=build_tensor_specs(args.start_pos, args.num_tokens), golden_fn=golden_prefill_indexer_compressor, diff --git a/models/deepseek_v4_pro/prefill_layer.py b/models/deepseek_v4_pro/prefill_layer.py index e3dbc70da..c04cfdda7 100644 --- a/models/deepseek_v4_pro/prefill_layer.py +++ b/models/deepseek_v4_pro/prefill_layer.py @@ -1203,8 +1203,7 @@ def init(): dim0 = batch * src.shape[0] tensor_specs.append(TensorSpec(packed_name, [N_RANKS, dim0, *src.shape[1:]], - src.dtype, init_value=make_init(), - is_output=is_global_pool)) + src.dtype, init_value=make_init())) # Batch metadata. tensor_specs.append(TensorSpec("seq_lens", [N_RANKS, batch], torch.int32, init_value=replicate(seq_lens_t))) @@ -1232,7 +1231,7 @@ def init_tid2eid(spec=spec): # InOut, not Out: the kernel writes only the packed chunk rows, and the host zeros # must reach the device so valid_ratio_reldiff can check the pad rows are untouched. tensor_specs.append(TensorSpec("x_next", [N_RANKS, total_tokens, HC_MULT, D], torch.float32, - init_value=torch.zeros, is_output=True)) + init_value=torch.zeros)) # Keep static weight parameters device-resident (child_memory), sharded per # rank. Dynamic cache/state/table tensors must stay as host tensors because @@ -1610,7 +1609,7 @@ def compare(actual, expected, **kwargs): import argparse import torch - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1712,7 +1711,7 @@ def compare(actual, expected, **kwargs): count = apply_real_layer_weights(specs, args.weights, layer_id=args.layer_id, ep=N_RANKS) print(f"[RUN] real weights: layer {args.layer_id}, {count} tensors from {args.weights}", flush=True) - result = run_jit( + result = run( fn=l3_prefill_layer, specs=specs, golden_fn=golden_prefill_layer, diff --git a/models/deepseek_v4_pro/prefill_mtp.py b/models/deepseek_v4_pro/prefill_mtp.py index 79740b482..c7aef0b12 100644 --- a/models/deepseek_v4_pro/prefill_mtp.py +++ b/models/deepseek_v4_pro/prefill_mtp.py @@ -14,7 +14,7 @@ import pypto.language as pl import pypto.language.distributed as pld -from golden import mapped_pool_ratio_reldiff, ratio_reldiff, run_jit +from golden import mapped_pool_ratio_reldiff, ratio_reldiff, run from pypto.ir.distributed_compiled_program import DistributedConfig import config @@ -326,12 +326,10 @@ def l3_mtp_prefill_fwd( ) -def _ranked(spec, torch, is_output=None): +def _ranked(spec, torch): from golden import TensorSpec - if is_output is None: - is_output = spec.is_output - return TensorSpec(spec.name, list(spec.shape), spec.dtype, init_value=spec.init_value, is_output=is_output) + return TensorSpec(spec.name, list(spec.shape), spec.dtype, init_value=spec.init_value) def _projection_specs(): @@ -458,8 +456,8 @@ def build_tensor_specs(start_pos=0, num_tokens=T): else: specs.append(_ranked(base[name], torch)) - specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16, is_output=True)) - specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32, is_output=True)) + specs.append(TensorSpec("hidden_out", [N_RANKS, T, D], torch.bfloat16)) + specs.append(TensorSpec("pre_hc_hidden_out", [N_RANKS, T, HC_MULT, D], torch.float32)) specs.append(ScalarSpec("num_tokens", torch.int32, num_tokens)) return specs @@ -584,7 +582,7 @@ def main(): if len(device_ids) < N_RANKS: raise ValueError(f"need at least {N_RANKS} devices, got {device_ids}") - result = run_jit( + result = run( fn=l3_mtp_prefill_fwd, specs=build_tensor_specs(start_pos=args.start_pos, num_tokens=args.num_tokens), golden_fn=golden_mtp_prefill_fwd, diff --git a/models/deepseek_v4_pro/prefill_sparse_attn.py b/models/deepseek_v4_pro/prefill_sparse_attn.py index 68e7f1633..9dd4117d1 100644 --- a/models/deepseek_v4_pro/prefill_sparse_attn.py +++ b/models/deepseek_v4_pro/prefill_sparse_attn.py @@ -724,12 +724,12 @@ def init_wo_b(): TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), TensorSpec("wo_b_scale", [D], torch.float32, init_value=lambda: wo_b_scale), - TensorSpec("attn_out", [T, D], torch.bfloat16, is_output=True), + TensorSpec("attn_out", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) @@ -742,7 +742,7 @@ def init_wo_b(): parser.add_argument("--dump-passes", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_sparse_attn_test, specs=build_tensor_specs(args.compress_ratio), golden_fn=golden_prefill_sparse_attn, diff --git a/models/deepseek_v4_pro/qkv_proj_rope.py b/models/deepseek_v4_pro/qkv_proj_rope.py index 4a20e41b3..c538f7aba 100644 --- a/models/deepseek_v4_pro/qkv_proj_rope.py +++ b/models/deepseek_v4_pro/qkv_proj_rope.py @@ -910,16 +910,16 @@ def init_gamma_ckv(): TensorSpec("rope_sin", [T, ROPE_DIM], torch.bfloat16, init_value=init_sin), TensorSpec("gamma_cq", [Q_LORA], torch.bfloat16, init_value=init_gamma_cq), TensorSpec("gamma_ckv", [HEAD_DIM], torch.bfloat16, init_value=init_gamma_ckv), - TensorSpec("q", [T, H, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("kv", [T, HEAD_DIM], torch.bfloat16, is_output=True), - TensorSpec("qr", [T, Q_LORA], torch.int8, is_output=True), - TensorSpec("qr_scale", [T, 1], torch.float32, is_output=True), + TensorSpec("q", [T, H, HEAD_DIM], torch.bfloat16), + TensorSpec("kv", [T, HEAD_DIM], torch.bfloat16), + TensorSpec("qr", [T, Q_LORA], torch.int8), + TensorSpec("qr_scale", [T, 1], torch.float32), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -949,7 +949,7 @@ def init_gamma_ckv(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- qkv_proj_rope {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=qkv_proj_rope_test, specs=build_tensor_specs(B, S), golden_fn=golden_qkv_proj_rope, diff --git a/models/deepseek_v4_pro/rmsnorm.py b/models/deepseek_v4_pro/rmsnorm.py index d7e637e6e..8fd0af5c4 100644 --- a/models/deepseek_v4_pro/rmsnorm.py +++ b/models/deepseek_v4_pro/rmsnorm.py @@ -109,13 +109,13 @@ def init_norm_w(): return [ TensorSpec("x", [T, D], torch.bfloat16, init_value=init_x), TensorSpec("norm_w", [D], torch.bfloat16, init_value=init_norm_w), - TensorSpec("x_normed", [T, D], torch.bfloat16, is_output=True), + TensorSpec("x_normed", [T, D], torch.bfloat16), ] if __name__ == "__main__": import argparse - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run MODES = { "decode": (DECODE_BATCH, DECODE_SEQ), @@ -139,7 +139,7 @@ def init_norm_w(): for mode_name in modes_to_run: B, S = MODES[mode_name] print(f"--- rms_norm_test {mode_name}: B={B}, S={S} ---") - result = run_jit( + result = run( fn=rms_norm_test, specs=build_tensor_specs(B, S), golden_fn=golden_rms_norm_test, diff --git a/models/deepseek_v4_pro/synthetic_token_loop.py b/models/deepseek_v4_pro/synthetic_token_loop.py index 22cdcae1b..91856a1d3 100644 --- a/models/deepseek_v4_pro/synthetic_token_loop.py +++ b/models/deepseek_v4_pro/synthetic_token_loop.py @@ -333,7 +333,7 @@ def _build_resident_hosts(prefill, prefill_compiled, decode_compiled, weight_spe ) if tensor.dtype != expected_dtype: # Fixture inits may build in a wider dtype (e.g. fp32) than the - # kernel ABI declares; run_jit casts on materialization too. + # kernel ABI declares; run casts on materialization too. tensor = tensor.to(expected_dtype) else: tensor = _empty_host_tensor(name, info, prefill.MODEL_CONFIG.vocab_size) diff --git a/models/deepseek_v4_pro/utils.py b/models/deepseek_v4_pro/utils.py index ef9529095..4f0526cbd 100644 --- a/models/deepseek_v4_pro/utils.py +++ b/models/deepseek_v4_pro/utils.py @@ -46,7 +46,7 @@ FWD-layer id, and the CSA/HCA-compact stacks are packed by per-kind order (ascending layer id of that kind). All cache slices handed to the leaf goldens are torch views of the stacked tensors, so their in-place slot updates land in -the ``is_output`` cache tensors that validation reads back. +the output cache tensors that validation reads back. Real DeepSeek-V4-Flash checkpoint loader ---------------------------------------- @@ -900,7 +900,7 @@ def _moe_views(tensors, layer, x_hc, x_next, num_tokens): def golden_prefill_fwd(tensors): - """Fill every ``is_output`` tensor of prefill_fwd's spec list in place.""" + """Fill every output tensor of prefill_fwd's spec list in place.""" import torch num_tokens = int(tensors["num_tokens"]) diff --git a/models/qwen3_14b/decode_fwd.py b/models/qwen3_14b/decode_fwd.py index 5570036dc..bf5ddb41c 100644 --- a/models/qwen3_14b/decode_fwd.py +++ b/models/qwen3_14b/decode_fwd.py @@ -1662,7 +1662,7 @@ def _backend_type(platform: str) -> BackendType: # the golden with `golden_decode_layer` (a torch reference mirroring the kernel's # math AND its bf16 cast points), runs decode_fwd_layers with _CHUNK_NLAYERS == 1 # (a single fused decode layer, hidden -> hidden, no LM head) on device through the -# `golden/` harness (golden.run_jit), and validates the device output against the +# `golden/` harness (golden.run), and validates the device output against the # golden — no pre-generated data files needed. # # Fixture scales are chosen so the (unnormalized) residual-stream output stays @@ -1853,7 +1853,7 @@ def _build_specs(inputs: dict) -> list: TensorSpec(name, list(inputs[name].shape), inputs[name].dtype, init_value=inputs[name]) for name in INPUT_NAMES ] - specs.append(TensorSpec("out", [BATCH_PAD, HIDDEN], torch.bfloat16, is_output=True)) + specs.append(TensorSpec("out", [BATCH_PAD, HIDDEN], torch.bfloat16)) return specs @@ -2025,7 +2025,7 @@ def _build_specs(inputs: dict) -> list: # ── Default single-layer unit test: RANDOM inputs, on-the-fly torch golden, # on-device run + compare, all through the golden/ harness. ── if not args.validate_fwd: - from golden import ratio_allclose, run_jit + from golden import ratio_allclose, run inputs = random_inputs( full_seq=args.max_seq, @@ -2042,7 +2042,7 @@ def _build_specs(inputs: dict) -> list: # rtol/atol=3e-3 (one bf16 ULP at value 1 is 2**-8 ≈ 0.0039 > 3e-3), so allow # up to 2% outliers — the codebase's ratio_allclose convention. Remaining # mismatches are 1-2 ULP bf16 quantization, not errors. - result = run_jit( + result = run( fn=decode_fwd_layers, specs=specs, golden_fn=golden_decode_layer, diff --git a/models/qwen3_14b/decode_tq_draft.py b/models/qwen3_14b/decode_tq_draft.py index c548fa3c7..322dec644 100644 --- a/models/qwen3_14b/decode_tq_draft.py +++ b/models/qwen3_14b/decode_tq_draft.py @@ -833,7 +833,7 @@ def init_lm_head_weight(): TensorSpec("final_norm_weight", [1, hidden_size], torch.float32, init_value=init_final_norm_weight), TensorSpec("lm_head_weight", [vocab, hidden_size], torch.bfloat16, init_value=init_lm_head_weight), # Outputs. - TensorSpec("out", [batch, vocab], torch.float32, is_output=True), + TensorSpec("out", [batch, vocab], torch.float32), ] @@ -1253,7 +1253,7 @@ def golden_decode_fwd_fp(tensors): if str(repo_root) not in sys.path: sys.path.insert(0, str(repo_root)) - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -1280,7 +1280,7 @@ def golden_decode_fwd_fp(tensors): f"decode_fwd_tq currently supports max_seq <= {MAX_SEQ}" ) - result = run_jit( + result = run( fn=decode_fwd_tq, specs=build_tensor_specs( batch=args.batch, diff --git a/models/qwen3_14b/greedy_sample.py b/models/qwen3_14b/greedy_sample.py index c814f3f86..bb0224d21 100644 --- a/models/qwen3_14b/greedy_sample.py +++ b/models/qwen3_14b/greedy_sample.py @@ -132,7 +132,7 @@ def init_logits(): return [ TensorSpec("logits", [BATCH_PAD, VOCAB], torch.float32, init_value=init_logits), - TensorSpec("sampled_ids", [BATCH_PAD, SAMPLED_IDS_PAD], torch.int32, is_output=True), + TensorSpec("sampled_ids", [BATCH_PAD, SAMPLED_IDS_PAD], torch.int32), ] @@ -148,7 +148,7 @@ def golden_greedy_sample(tensors): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument("-p", "--platform", type=str, default="a2a3", @@ -157,7 +157,7 @@ def golden_greedy_sample(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=greedy_sample_fwd, specs=build_tensor_specs(), golden_fn=golden_greedy_sample, diff --git a/models/qwen3_14b/prefill_fwd.py b/models/qwen3_14b/prefill_fwd.py index c4094d6a2..88fa75ecb 100644 --- a/models/qwen3_14b/prefill_fwd.py +++ b/models/qwen3_14b/prefill_fwd.py @@ -1941,7 +1941,7 @@ def init_embed_weight(): init_value=init_lm_head_weight), TensorSpec("embed_weight", [vocab, hidden_size], torch.bfloat16, init_value=init_embed_weight), - TensorSpec("out", [batch, vocab], torch.float32, is_output=True), + TensorSpec("out", [batch, vocab], torch.float32), ] @@ -2209,7 +2209,7 @@ def tiled_lm_head(lhs, rhs_t, k_chunk, vocab_chunk): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument( @@ -2258,7 +2258,7 @@ def tiled_lm_head(lhs, rhs_t, k_chunk, vocab_chunk): help="skip host golden computation and output validation") args = parser.parse_args() - result = run_jit( + result = run( fn=prefill_fwd, specs=build_tensor_specs( batch=args.batch, diff --git a/models/qwen3_14b/prefill_tq_draft.py b/models/qwen3_14b/prefill_tq_draft.py index 6afd96e90..098eee4f6 100644 --- a/models/qwen3_14b/prefill_tq_draft.py +++ b/models/qwen3_14b/prefill_tq_draft.py @@ -973,7 +973,7 @@ def init_lm_head_weight(): init_value=init_final_norm_weight), TensorSpec("lm_head_weight", [vocab, hidden_size], torch.bfloat16, init_value=init_lm_head_weight), - TensorSpec("out", [batch, vocab], torch.float32, is_output=True), + TensorSpec("out", [batch, vocab], torch.float32), ] @@ -1319,7 +1319,7 @@ def tq_dequant(indices, scales, rot_matrix_f): if __name__ == "__main__": import argparse - from golden import run_jit + from golden import run parser = argparse.ArgumentParser() parser.add_argument( @@ -1347,7 +1347,7 @@ def tq_dequant(indices, scales, rot_matrix_f): torch.manual_seed(args.seed) - result = run_jit( + result = run( fn=prefill_fwd_tq, specs=build_tensor_specs( batch=args.batch, num_layers=args.num_layers, diff --git a/models/qwen3_14b/test_paged_attention_cce.py b/models/qwen3_14b/test_paged_attention_cce.py index 88fa6820e..ffe46da31 100644 --- a/models/qwen3_14b/test_paged_attention_cce.py +++ b/models/qwen3_14b/test_paged_attention_cce.py @@ -14,7 +14,7 @@ import torch -from golden import TensorSpec, run_jit +from golden import TensorSpec, run from paged_attention_cce import ( BATCH_PAD, BLOCK_SIZE, @@ -72,7 +72,6 @@ def build_specs( "out", [batch, NUM_HEADS, HEAD_DIM], torch.bfloat16, - is_output=True, ), ] @@ -160,7 +159,7 @@ def main() -> None: else: golden_fn = golden_attention if args.check else None - result = run_jit( + result = run( fn=fn, specs=specs, golden_fn=golden_fn, diff --git a/models/qwen3_14b/test_paged_attention_pypto.py b/models/qwen3_14b/test_paged_attention_pypto.py index a2bf7f205..a161735d2 100644 --- a/models/qwen3_14b/test_paged_attention_pypto.py +++ b/models/qwen3_14b/test_paged_attention_pypto.py @@ -36,7 +36,7 @@ import pypto.language as pl import torch -from golden import ScalarSpec, TensorSpec, run_jit +from golden import ScalarSpec, TensorSpec, run from paged_attention_pypto import ( BATCH, BLOCK_SIZE, @@ -568,14 +568,12 @@ def initializer(name: str): cache_shape, torch.bfloat16, init_value=initializer("key_cache"), - is_output=True, ), TensorSpec( "value_cache", cache_shape, torch.bfloat16, init_value=initializer("value_cache"), - is_output=True, ), TensorSpec( "block_table", @@ -652,7 +650,6 @@ def initializer(name: str): "out", [case.batch, NUM_HEADS, HEAD_DIM], torch.bfloat16, - is_output=True, ), ] @@ -831,7 +828,7 @@ def _run_case(case: DynamicCase, args: argparse.Namespace) -> dict[str, object]: _validate_case(case) compile_only = args.compile_only or args.platform.endswith("sim") fixture = None if compile_only else make_fixture(case) - result = run_jit( + result = run( fn=paged_attention_pypto_dynamic, specs=build_specs(case, fixture), golden_fn=None if fixture is None else lambda values: golden_attention(values, case), diff --git a/models/qwen3_14b/topk_select.py b/models/qwen3_14b/topk_select.py index bd1a7a983..65c2f3aa4 100644 --- a/models/qwen3_14b/topk_select.py +++ b/models/qwen3_14b/topk_select.py @@ -317,8 +317,8 @@ def init_logits(): torch.int32, init_value=lambda: torch.tensor([2 if selection_k == TOPK else 1, selection_k], dtype=torch.int32), ), - TensorSpec("topk_values", [BATCH_PAD, TOPK], torch.float32, is_output=True), - TensorSpec("topk_indices", [BATCH_PAD, TOPK], torch.int32, is_output=True), + TensorSpec("topk_values", [BATCH_PAD, TOPK], torch.float32), + TensorSpec("topk_indices", [BATCH_PAD, TOPK], torch.int32), ] @@ -345,7 +345,7 @@ def golden_topk_select(tensors): if __name__ == "__main__": import argparse - from golden import run_jit, topk_pair_compare + from golden import run, topk_pair_compare parser = argparse.ArgumentParser() parser.add_argument( @@ -356,7 +356,7 @@ def golden_topk_select(tensors): parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) args = parser.parse_args() - result = run_jit( + result = run( fn=topk_select_fwd, specs=build_tensor_specs(args.selection_k), golden_fn=golden_topk_select, diff --git a/tests/contract/test_deepseek_v4_pro_moe_protocol.py b/tests/contract/test_deepseek_v4_pro_moe_protocol.py index 8c8534059..83fff3241 100644 --- a/tests/contract/test_deepseek_v4_pro_moe_protocol.py +++ b/tests/contract/test_deepseek_v4_pro_moe_protocol.py @@ -106,7 +106,9 @@ def test_decode_layer_cache_specs_match_static_inout_abi(): assert inout_names == ast.literal_eval(cache_assignment.value) assert out_names == {"x_next"} - assert "is_output=name in mutable_cache_names" in ast.unparse(build_specs) + # Direction is stamped from the compiled artifact, so the spec builder must + # not re-declare it. + assert "is_output" not in ast.unparse(build_specs) def test_prefill_mtp_ranked_specs_preserve_inout_direction(): @@ -123,7 +125,7 @@ def test_prefill_mtp_ranked_specs_preserve_inout_direction(): assert inout_names == {"kv_cache"} assert out_names == {"hidden_out", "pre_hc_hidden_out"} - assert "is_output = spec.is_output" in ast.unparse(_function(tree, "_ranked")) + assert "is_output" not in ast.unparse(_function(tree, "_ranked")) main = _function(tree, "main") cache_comparators = [ diff --git a/tests/golden/test_runner.py b/tests/golden/test_runner.py index 9b985120a..a7e93d062 100644 --- a/tests/golden/test_runner.py +++ b/tests/golden/test_runner.py @@ -22,13 +22,12 @@ import pytest import torch -from golden import ScalarSpec, TensorSpec, run, run_jit +from golden import ScalarSpec, TensorSpec, run from golden.runner import ( RunResult, _backend_for_platform, _bench_loop_sizes, _format_stale_paths, - _l3_ordered_args, _maybe_reload_l3, _prepare_inputs, _report_effective, @@ -43,6 +42,7 @@ _setup_runtime_dir, _share_in_place, _stale_cpps, + _validate_compiled_spec_abi, ) @@ -92,15 +92,44 @@ def _l3_abi_environment(): yield +def _make_build_dir(tmp_path): + """A created ``build_output``-style directory for a compiled-artifact double.""" + build = tmp_path / "build" + build.mkdir(exist_ok=True) + return build + + +def _artifact(output_dir, *infos): + """A compiled-artifact double exposing *infos* as its parameter metadata.""" + return types.SimpleNamespace( + output_dir=output_dir, + _get_metadata=lambda: (list(infos), None, None), + ) + + +@pytest.fixture +def build_dir(tmp_path): + """`_make_build_dir` as a fixture, for tests that need no other tmp files.""" + return _make_build_dir(tmp_path) + + +def _stamped(specs, directions): + """Apply the direction stamp ``_validate_compiled_spec_abi`` normally copies + from the compiled artifact, for doubles that expose no parameter metadata.""" + for spec in specs: + if isinstance(spec, TensorSpec): + spec.direction = directions[spec.name] + return specs + + @pytest.fixture def three_kinds_specs(): """TensorSpec trio covering pure input / pure output / inout.""" - return [ + return _stamped([ TensorSpec("x", [4], torch.float32, init_value=torch.randn), # pure input - TensorSpec("y", [4], torch.float32, is_output=True), # pure output - TensorSpec("state", [4], torch.float32, init_value=torch.zeros, # inout - is_output=True), - ] + TensorSpec("y", [4], torch.float32), # pure output + TensorSpec("state", [4], torch.float32, init_value=torch.zeros), # inout + ], {"x": "in", "y": "out", "state": "inout"}) @pytest.fixture @@ -155,10 +184,8 @@ def fake_execute(work_dir, tensors, **kwargs): class TestGoldenDataCacheHit: """``golden_data`` points at a complete cache: skip generate + compute.""" - def test_hit_skips_generate_and_golden_fn(self, populated_cache, three_kinds_specs, tmp_path): + def test_hit_skips_generate_and_golden_fn(self, populated_cache, three_kinds_specs, build_dir): """With cache hit: create_tensor and golden_fn must not run; validate passes.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() # Simulate a correct kernel: it writes the cached golden values back into # the y and state tensors so validate_golden passes. @@ -172,10 +199,10 @@ def golden_fn_should_not_run(tensors): def _no_create_tensor(self): pytest.fail(f"TensorSpec.create_tensor must not run for {self.name}") - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, write_outputs) + compile_p, exec_p = _patch_compile_and_execute(build_dir, write_outputs) with compile_p, exec_p, patch.object(TensorSpec, "create_tensor", _no_create_tensor): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn_should_not_run, golden_data=str(populated_cache), @@ -183,24 +210,22 @@ def _no_create_tensor(self): assert r.passed, f"unexpected failure: {r.error}" # Read-only: no data/ written under compiled.output_dir. - assert not (compiled_dir / "data").exists() + assert not (build_dir / "data").exists() def test_hit_without_golden_fn_still_validates( - self, populated_cache, three_kinds_specs, tmp_path, + self, populated_cache, three_kinds_specs, build_dir, ): """golden_fn=None + golden_data set → validation still runs via loaded out/.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() # Same setup as the previous test but no golden_fn. y_golden = torch.tensor([2.0, 3.0, 4.0, 5.0]) state_out = torch.tensor([11.0, 22.0, 33.0, 44.0]) write_outputs = [None, y_golden, state_out] - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, write_outputs) + compile_p, exec_p = _patch_compile_and_execute(build_dir, write_outputs) with compile_p, exec_p: r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=None, golden_data=str(populated_cache), @@ -209,20 +234,18 @@ def test_hit_without_golden_fn_still_validates( assert r.passed, f"unexpected failure: {r.error}" def test_hit_with_mismatched_device_output_fails( - self, populated_cache, three_kinds_specs, tmp_path, + self, populated_cache, three_kinds_specs, build_dir, ): """If device writes values that differ from cached golden → validation fails.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() bad_y = torch.full((4,), 99.0) bad_state = torch.full((4,), -1.0) write_outputs = [None, bad_y, bad_state] - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, write_outputs) + compile_p, exec_p = _patch_compile_and_execute(build_dir, write_outputs) with compile_p, exec_p: r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=None, golden_data=str(populated_cache), @@ -232,12 +255,10 @@ def test_hit_with_mismatched_device_output_fails( assert "does not match golden" in (r.error or "") def test_hit_loads_inout_initial_value_from_in( - self, populated_cache, three_kinds_specs, tmp_path, + self, populated_cache, three_kinds_specs, build_dir, ): """Verify that the tensor handed to execute_compiled for the inout "state" is the value from in/state.pt, not a freshly created one.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() observed: dict[str, torch.Tensor] = {} @@ -249,11 +270,11 @@ def capture_execute(work_dir, tensors, **kwargs): tensors[1][:] = torch.tensor([2.0, 3.0, 4.0, 5.0]) # y_golden tensors[2][:] = torch.tensor([11.0, 22.0, 33.0, 44.0]) # state_out - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) with patch("pypto.ir.compile", return_value=fake), \ patch("pypto.runtime.execute_compiled", side_effect=capture_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=None, golden_data=str(populated_cache), @@ -271,12 +292,11 @@ class TestGoldenDataCacheMiss: def test_empty_dir_lists_all_missing(self, three_kinds_specs, tmp_path): empty = tmp_path / "empty_cache" empty.mkdir() - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) compile_p, exec_p = _patch_compile_and_execute(compiled_dir) with compile_p, exec_p: r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=lambda t: None, golden_data=str(empty), @@ -295,12 +315,11 @@ def test_partial_cache_still_fails(self, three_kinds_specs, tmp_path): "y": torch.zeros(4), "state": torch.zeros(4), }) - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) compile_p, exec_p = _patch_compile_and_execute(compiled_dir) with compile_p, exec_p: r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=None, golden_data=str(partial), @@ -317,10 +336,8 @@ class TestGoldenFnPath: ``golden_fn``, and persists ``data/in/`` + ``data/out/`` under the compiled output directory.""" - def test_golden_fn_called_and_matches(self, three_kinds_specs, tmp_path): + def test_golden_fn_called_and_matches(self, three_kinds_specs, build_dir): """``golden_fn`` runs, writes expected outputs, and validation passes.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() # golden_fn is called with a {name: tensor} dict — mutate y/state in place. def golden_fn(tensors): @@ -334,11 +351,11 @@ def fake_execute(work_dir, tensors, **_kwargs): tensors[1][:] = tensors[0] + 1 tensors[2][:] = tensors[2] + 100 - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) with patch("pypto.ir.compile", return_value=fake), \ patch("pypto.runtime.execute_compiled", side_effect=fake_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn, save_data=True, @@ -346,19 +363,17 @@ def fake_execute(work_dir, tensors, **_kwargs): assert r.passed, f"unexpected failure: {r.error}" # Persistence: data/in/ and data/out/ written under compiled.output_dir. - assert (compiled_dir / "data" / "in" / "x.pt").is_file() - assert (compiled_dir / "data" / "in" / "state.pt").is_file() - assert (compiled_dir / "data" / "out" / "y.pt").is_file() - assert (compiled_dir / "data" / "out" / "state.pt").is_file() + assert (build_dir / "data" / "in" / "x.pt").is_file() + assert (build_dir / "data" / "in" / "state.pt").is_file() + assert (build_dir / "data" / "out" / "y.pt").is_file() + assert (build_dir / "data" / "out" / "state.pt").is_file() def test_golden_fn_sees_cloned_inputs_not_live_tensors( - self, three_kinds_specs, tmp_path, + self, three_kinds_specs, build_dir, ): """``golden_fn`` receives a *clone* of inputs, not the live tensors handed to ``execute_compiled`` — so device writes don't corrupt the golden computation.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() captured = {} @@ -374,11 +389,11 @@ def fake_execute(work_dir, tensors, **_kwargs): tensors[1][:] = tensors[0] + 1 tensors[2][:] = tensors[2] + 100 - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) with patch("pypto.ir.compile", return_value=fake), \ patch("pypto.runtime.execute_compiled", side_effect=fake_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn, ) @@ -387,10 +402,8 @@ def fake_execute(work_dir, tensors, **_kwargs): # The golden_fn copy must not share storage with the device tensor. assert captured["x_ptr"] != device_x_ptrs["x"] - def test_golden_fn_mismatch_fails(self, three_kinds_specs, tmp_path): + def test_golden_fn_mismatch_fails(self, three_kinds_specs, build_dir): """Device output diverges from golden_fn output → FAIL.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() def golden_fn(tensors): tensors["y"][:] = tensors["x"] + 1 @@ -400,11 +413,11 @@ def bad_execute(work_dir, tensors, **_kwargs): tensors[1][:] = tensors[0] - 99 # wrong tensors[2][:] = tensors[2] + 100 - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) with patch("pypto.ir.compile", return_value=fake), \ patch("pypto.runtime.execute_compiled", side_effect=bad_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn, ) @@ -417,12 +430,10 @@ class TestSaveData: """``save_data=False`` skips the ``data/`` snapshot but still validates.""" def test_save_data_false_skips_persist_but_validates( - self, three_kinds_specs, tmp_path, + self, three_kinds_specs, build_dir, ): """With save_data=False: validation runs against the in-memory golden, but no data/in/ or data/out/ files are written.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() def golden_fn(tensors): tensors["y"][:] = tensors["x"] + 1 @@ -432,11 +443,11 @@ def fake_execute(_work_dir, tensors, **_kwargs): tensors[1][:] = tensors[0] + 1 tensors[2][:] = tensors[2] + 100 - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) with patch("pypto.ir.compile", return_value=fake), \ patch("pypto.runtime.execute_compiled", side_effect=fake_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn, save_data=False, @@ -444,26 +455,24 @@ def fake_execute(_work_dir, tensors, **_kwargs): assert r.passed, f"unexpected failure: {r.error}" # Nothing persisted under the compiled output directory. - assert not (compiled_dir / "data").exists() + assert not (build_dir / "data").exists() class TestNoValidation: """Neither ``golden_fn`` nor ``golden_data`` — validation is skipped.""" def test_skip_validation_passes_even_on_nonsense_outputs( - self, three_kinds_specs, tmp_path, + self, three_kinds_specs, build_dir, ): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() def fake_execute(work_dir, tensors, **_kwargs): tensors[1][:] = torch.full_like(tensors[1], 9999.0) - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) with patch("pypto.ir.compile", return_value=fake), \ patch("pypto.runtime.execute_compiled", side_effect=fake_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=None, golden_data=None, @@ -472,19 +481,17 @@ def fake_execute(work_dir, tensors, **_kwargs): assert r.passed # Inputs are still persisted (classic path), outputs are NOT computed/saved. - assert (compiled_dir / "data" / "in" / "x.pt").is_file() - assert not (compiled_dir / "data" / "out").exists() + assert (build_dir / "data" / "in" / "x.pt").is_file() + assert not (build_dir / "data" / "out").exists() class TestCompileOnly: """``compile_only=True`` short-circuits after compile.""" def test_compile_only_skips_runtime_and_validation( - self, three_kinds_specs, tmp_path, + self, three_kinds_specs, build_dir, ): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) def exec_must_not_run(*_args, **_kwargs): pytest.fail("execute_compiled must not run when compile_only=True") @@ -495,7 +502,7 @@ def golden_fn_must_not_run(_tensors): with patch("pypto.ir.compile", return_value=fake), \ patch("pypto.runtime.execute_compiled", side_effect=exec_must_not_run): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, compile_only=True, golden_fn=golden_fn_must_not_run, @@ -504,7 +511,7 @@ def golden_fn_must_not_run(_tensors): assert r.passed assert r.error is None # compile_only path must not persist anything under data/. - assert not (compiled_dir / "data").exists() + assert not (build_dir / "data").exists() def test_duplicate_specs_fail_before_compile(self): specs = [ @@ -512,23 +519,14 @@ def test_duplicate_specs_fail_before_compile(self): ScalarSpec("duplicate", torch.int32, 0), ] with patch("pypto.ir.compile") as compile_fn: - result = run(program=object(), specs=specs, compile_only=True) + result = run(fn=object(), specs=specs, compile_only=True) assert not result.passed assert "duplicate spec names" in result.error compile_fn.assert_not_called() - def test_compile_only_validates_exact_l3_abi_before_success(self, tmp_path): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compiled = types.SimpleNamespace( - output_dir=compiled_dir, - _get_metadata=lambda: ( - [_l3_info("x__ssa_v0", shape=[5], dtype=torch.float32)], - None, - None, - ), - ) + def test_compile_only_validates_exact_l3_abi_before_success(self, build_dir): + compiled = _artifact(build_dir, _l3_info("x__ssa_v0", shape=[5], dtype=torch.float32)) specs = [TensorSpec("x", [4], torch.float32)] with ( @@ -537,7 +535,7 @@ def test_compile_only_validates_exact_l3_abi_before_success(self, tmp_path): patch.object(TensorSpec, "create_tensor") as create_tensor, patch("pypto.runtime.execute_compiled") as execute, ): - result = run(program=object(), specs=specs, compile_only=True) + result = run(fn=object(), specs=specs, compile_only=True) assert not result.passed assert "shape" in (result.error or "") @@ -546,11 +544,9 @@ def test_compile_only_validates_exact_l3_abi_before_success(self, tmp_path): execute.assert_not_called() -class TestRunJitCompileRuntime: - def test_marked_scalar_uses_signature_mode_and_runtime_marker(self, tmp_path): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compiled = _FakeCompiled(compiled_dir) +class TestJitCompilePath: + def test_marked_scalar_uses_signature_mode_and_runtime_marker(self, build_dir): + compiled = _FakeCompiled(build_dir) fn = types.SimpleNamespace(compile=MagicMock(return_value=compiled)) runtime_marker = object() @@ -577,7 +573,7 @@ def __init__(self, **kwargs): create=True, ), ): - result = run_jit( + result = run( fn, specs, compile_cfg={"dump_passes": False}, @@ -597,7 +593,7 @@ def __init__(self, **kwargs): def test_stepped_scalar_requires_runtime_compilation(self): fn = types.SimpleNamespace(compile=MagicMock()) - result = run_jit( + result = run( fn, [ScalarSpec("epoch", torch.int32, 0, benchmark_step=1)], compile_only=True, @@ -611,7 +607,7 @@ def test_stepped_scalar_rejects_multi_pass_swimlane_before_compile(self): # Both the current key and its pre-rename spelling arm the guard. for key in ("enable_chip_swimlane", "enable_l2_swimlane"): fn = types.SimpleNamespace(compile=MagicMock()) - result = run_jit( + result = run( fn, [ ScalarSpec( @@ -635,7 +631,7 @@ def test_duplicate_specs_fail_before_compile(self): ScalarSpec("duplicate", torch.int32, 0, compile_runtime=True), ] - result = run_jit(fn, specs, compile_only=True) + result = run(fn, specs, compile_only=True) assert not result.passed assert "duplicate spec names" in result.error @@ -649,21 +645,12 @@ def test_duplicate_specs_fail_before_compile(self): ["epoch__ssa_v0", "epoch__ssa_v1"], ], ) - def test_compile_only_validates_l3_parameter_abi(self, tmp_path, compiled_names): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compiled = types.SimpleNamespace( - output_dir=compiled_dir, - _get_metadata=lambda: ( - [_l3_info(name) for name in compiled_names], - None, - None, - ), - ) + def test_compile_only_validates_l3_parameter_abi(self, build_dir, compiled_names): + compiled = _artifact(build_dir, *[_l3_info(name) for name in compiled_names]) fn = types.SimpleNamespace(compile=MagicMock(return_value=compiled)) with _l3_abi_environment(): - result = run_jit( + result = run( fn, [ScalarSpec("epoch", torch.int32, 0, compile_runtime=True)], compile_only=True, @@ -674,20 +661,12 @@ def test_compile_only_validates_l3_parameter_abi(self, tmp_path, compiled_names) @pytest.mark.parametrize(("artifact_shape", "passed"), [([4], True), ([5], False)]) def test_compile_only_validates_l2_annotation_abi( - self, tmp_path, artifact_shape, passed + self, build_dir, artifact_shape, passed ): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compiled = types.SimpleNamespace( - output_dir=compiled_dir, - _get_metadata=lambda: ( - [ - _l3_info("x__ssa_v0", shape=artifact_shape, dtype=torch.float32), - _l3_info("epoch__ssa_v0"), - ], - None, - None, - ), + compiled = _artifact( + build_dir, + _l3_info("x__ssa_v0", shape=artifact_shape, dtype=torch.float32), + _l3_info("epoch__ssa_v0"), ) fn = types.SimpleNamespace(compile=MagicMock(return_value=compiled)) specs = [ @@ -701,7 +680,7 @@ def test_compile_only_validates_l2_annotation_abi( patch.object(TensorSpec, "create_tensor") as create_tensor, patch("pypto.runtime.execute_compiled") as execute, ): - result = run_jit(fn, specs, compile_only=True) + result = run(fn, specs, compile_only=True) assert result.passed is passed if not passed: @@ -710,24 +689,16 @@ def test_compile_only_validates_l2_annotation_abi( create_tensor.assert_not_called() execute.assert_not_called() - def test_compile_only_rejects_l2_parameter_order_mismatch(self, tmp_path): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compiled = types.SimpleNamespace( - output_dir=compiled_dir, - _get_metadata=lambda: ( - [ - _l3_info("b__ssa_v0", shape=[4], dtype=torch.float32), - _l3_info("a__ssa_v0", shape=[4], dtype=torch.float32), - ], - None, - None, - ), + def test_compile_only_rejects_l2_parameter_order_mismatch(self, build_dir): + compiled = _artifact( + build_dir, + _l3_info("b__ssa_v0", shape=[4], dtype=torch.float32), + _l3_info("a__ssa_v0", shape=[4], dtype=torch.float32), ) fn = types.SimpleNamespace(compile=MagicMock(return_value=compiled)) with _l3_abi_environment(), patch("golden.runner._is_l3", return_value=False): - result = run_jit( + result = run( fn, [ TensorSpec("a", [4], torch.float32), @@ -745,7 +716,6 @@ def test_compile_only_rejects_l2_parameter_order_mismatch(self, tmp_path): ("x", [5], torch.float32, _FakeParamDirection.In, "shape"), ("x", [4], torch.int32, _FakeParamDirection.In, "dtype"), ("x", None, torch.float32, _FakeParamDirection.In, "expected tensor"), - ("x", [4], torch.float32, _FakeParamDirection.Out, "direction"), ("epoch", [1], torch.int32, _FakeParamDirection.In, "expected scalar"), ("epoch", None, torch.int64, _FakeParamDirection.In, "dtype"), ("epoch", None, torch.int32, _FakeParamDirection.InOut, "direction"), @@ -753,7 +723,7 @@ def test_compile_only_rejects_l2_parameter_order_mismatch(self, tmp_path): ) def test_compile_only_validates_exact_l3_parameter_abi( self, - tmp_path, + build_dir, target, shape, dtype, @@ -761,8 +731,6 @@ def test_compile_only_validates_exact_l3_parameter_abi( error, ): """Signature compilation must not hide a stale tensor/scalar spec.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() infos = { "x": _l3_info("x__ssa_v0", shape=[4], dtype=torch.float32), "epoch": _l3_info("epoch__ssa_v0"), @@ -774,7 +742,7 @@ def test_compile_only_validates_exact_l3_parameter_abi( direction=direction, ) compiled = types.SimpleNamespace( - output_dir=compiled_dir, + output_dir=build_dir, _get_metadata=lambda: ([infos["x"], infos["epoch"]], None, None), ) fn = types.SimpleNamespace(compile=MagicMock(return_value=compiled)) @@ -788,7 +756,7 @@ def test_compile_only_validates_exact_l3_parameter_abi( patch.object(TensorSpec, "create_tensor") as create_tensor, patch("pypto.runtime.execute_compiled") as execute, ): - result = run_jit(fn, specs, compile_only=True) + result = run(fn, specs, compile_only=True) assert not result.passed assert error in (result.error or "") @@ -796,30 +764,22 @@ def test_compile_only_validates_exact_l3_parameter_abi( create_tensor.assert_not_called() execute.assert_not_called() - def test_compile_only_accepts_dynamic_l3_tensor_dimension(self, tmp_path): - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compiled = types.SimpleNamespace( - output_dir=compiled_dir, - _get_metadata=lambda: ( - [ - _l3_info("x__ssa_v0", shape=[-1], dtype=torch.float32), - _l3_info( - "state__ssa_v0", - shape=[4], - dtype=torch.float32, - direction=_FakeParamDirection.InOut, - ), - _l3_info("epoch__ssa_v0"), - ], - None, - None, - ), + def test_compile_only_accepts_dynamic_l3_tensor_dimension(self, build_dir): + compiled = _artifact( + build_dir, + _l3_info("x__ssa_v0", shape=[-1], dtype=torch.float32), + _l3_info( + "state__ssa_v0", + shape=[4], + dtype=torch.float32, + direction=_FakeParamDirection.InOut, + ), + _l3_info("epoch__ssa_v0"), ) fn = types.SimpleNamespace(compile=MagicMock(return_value=compiled)) with _l3_abi_environment(): - result = run_jit( + result = run( fn, [ TensorSpec("x", [4], torch.float32), @@ -828,7 +788,6 @@ def test_compile_only_accepts_dynamic_l3_tensor_dimension(self, tmp_path): [4], torch.float32, init_value=torch.zeros, - is_output=True, ), ScalarSpec("epoch", torch.int32, 0, compile_runtime=True), ], @@ -846,8 +805,7 @@ def test_signature_compile_uses_cached_static_scalar(self, tmp_path): "epoch": torch.tensor(86, dtype=torch.int32), }, ) - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) fn = types.SimpleNamespace( compile=MagicMock(return_value=_FakeCompiled(compiled_dir)) ) @@ -860,7 +818,7 @@ def test_signature_compile_uses_cached_static_scalar(self, tmp_path): with patch.object( sys.modules["pypto.language"], "RUNTIME", runtime_marker, create=True ): - result = run_jit( + result = run( fn, specs, golden_data=str(cache), @@ -876,13 +834,12 @@ def test_legacy_compile_uses_cached_static_scalar(self, tmp_path): _save_tensors( cache / "in", {"num_tokens": torch.tensor(9, dtype=torch.int32)} ) - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) fn = types.SimpleNamespace( compile=MagicMock(return_value=_FakeCompiled(compiled_dir)) ) - result = run_jit( + result = run( fn, [ScalarSpec("num_tokens", torch.int32, 4)], golden_data=str(cache), @@ -906,7 +863,7 @@ def test_bad_cached_scalar_fails_before_compile(self, tmp_path, cached): _save_tensors(cache / "in", {"num_tokens": cached}) fn = types.SimpleNamespace(compile=MagicMock()) - result = run_jit( + result = run( fn, [ScalarSpec("num_tokens", torch.int32, 4)], golden_data=str(cache), @@ -938,7 +895,7 @@ def fake_execute(work_dir, tensors, **_kwargs): with patch("pypto.ir.compile", side_effect=compile_must_not_run), \ patch("pypto.runtime.execute_compiled", side_effect=fake_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, runtime_dir=str(prebuilt), ) @@ -964,7 +921,7 @@ def fake_execute(_work_dir, tensors, **_kwargs): with patch("pypto.ir.compile", side_effect=lambda *a, **kw: pytest.fail("compile must not run")), \ patch("pypto.runtime.execute_compiled", side_effect=fake_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn, runtime_dir=str(prebuilt), @@ -995,7 +952,7 @@ def test_runtime_dir_l3_routes_to_l3_dispatch(self, three_kinds_specs, tmp_path) patch("golden.runner._maybe_reload_l3", return_value=fake_l3) as reload, patch("golden.runner._try_l3_dispatch", return_value=True) as l3, ): - r = run(program=object(), specs=three_kinds_specs, runtime_dir=str(prebuilt)) + r = run(fn=object(), specs=three_kinds_specs, runtime_dir=str(prebuilt)) assert r.passed, f"unexpected failure: {r.error}" reload.assert_called_once() @@ -1005,14 +962,7 @@ def test_runtime_dir_l3_routes_to_l3_dispatch(self, three_kinds_specs, tmp_path) def test_runtime_dir_l3_abi_mismatch_fails_before_input_or_runtime(self, tmp_path): prebuilt = tmp_path / "prebuilt" prebuilt.mkdir() - compiled = types.SimpleNamespace( - output_dir=prebuilt, - _get_metadata=lambda: ( - [_l3_info("x__ssa_v0", shape=[5], dtype=torch.float32)], - None, - None, - ), - ) + compiled = _artifact(prebuilt, _l3_info("x__ssa_v0", shape=[5], dtype=torch.float32)) specs = [TensorSpec("x", [4], torch.float32)] with ( @@ -1023,7 +973,7 @@ def test_runtime_dir_l3_abi_mismatch_fails_before_input_or_runtime(self, tmp_pat patch("golden.runner._try_l3_dispatch") as l3_dispatch, patch("pypto.runtime.execute_compiled") as execute, ): - result = run(program=object(), specs=specs, runtime_dir=str(prebuilt)) + result = run(fn=object(), specs=specs, runtime_dir=str(prebuilt)) assert not result.passed assert "shape" in (result.error or "") @@ -1032,17 +982,10 @@ def test_runtime_dir_l3_abi_mismatch_fails_before_input_or_runtime(self, tmp_pat l3_dispatch.assert_not_called() execute.assert_not_called() - def test_run_jit_runtime_dir_l3_abi_mismatch_fails_before_input(self, tmp_path): + def test_runtime_dir_l3_abi_mismatch_fails_before_input(self, tmp_path): prebuilt = tmp_path / "prebuilt" prebuilt.mkdir() - compiled = types.SimpleNamespace( - output_dir=prebuilt, - _get_metadata=lambda: ( - [_l3_info("x__ssa_v0", shape=[4], dtype=torch.int32)], - None, - None, - ), - ) + compiled = _artifact(prebuilt, _l3_info("x__ssa_v0", shape=[4], dtype=torch.int32)) fn = types.SimpleNamespace(compile=MagicMock()) with ( @@ -1052,7 +995,7 @@ def test_run_jit_runtime_dir_l3_abi_mismatch_fails_before_input(self, tmp_path): patch("golden.runner._try_l3_dispatch") as l3_dispatch, patch("pypto.runtime.execute_compiled") as execute, ): - result = run_jit( + result = run( fn, [TensorSpec("x", [4], torch.float32)], runtime_dir=str(prebuilt), @@ -1073,14 +1016,7 @@ def test_runtime_dir_l3_skips_requested_benchmark( ): prebuilt = tmp_path / "prebuilt" prebuilt.mkdir() - compiled = types.SimpleNamespace( - output_dir=prebuilt, - _get_metadata=lambda: ( - [_l3_info("x__ssa_v0", shape=[1], dtype=torch.float32)], - None, - None, - ), - ) + compiled = _artifact(prebuilt, _l3_info("x__ssa_v0", shape=[1], dtype=torch.float32)) specs = [TensorSpec("x", [1], torch.float32)] monkeypatch.setenv("PYPTO_BENCH", "1") @@ -1091,7 +1027,7 @@ def test_runtime_dir_l3_skips_requested_benchmark( patch("golden.runner._run_benchmark_l3") as benchmark, patch("pypto.runtime.execute_compiled") as execute, ): - result = run(program=object(), specs=specs, runtime_dir=str(prebuilt)) + result = run(fn=object(), specs=specs, runtime_dir=str(prebuilt)) assert result.passed, result.error assert result.bench is None @@ -1111,7 +1047,7 @@ def exec_must_not_run(*_args, **_kwargs): with patch("pypto.ir.compile", side_effect=compile_must_not_run), \ patch("pypto.runtime.execute_compiled", side_effect=exec_must_not_run): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, runtime_dir=str(missing), ) @@ -1134,7 +1070,7 @@ def exec_must_not_run(*_args, **_kwargs): with patch("pypto.ir.compile", side_effect=compile_must_not_run), \ patch("pypto.runtime.execute_compiled", side_effect=exec_must_not_run): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, compile_only=True, runtime_dir=str(prebuilt), @@ -1165,7 +1101,7 @@ def golden_fn_should_not_run(_tensors): with patch("pypto.ir.compile", side_effect=lambda *a, **kw: pytest.fail("compile must not run")), \ patch("pypto.runtime.execute_compiled", side_effect=fake_execute): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn_should_not_run, golden_data=str(populated_cache), @@ -1216,23 +1152,21 @@ def test_fail_without_error(self): @pytest.fixture def mixed_specs(): """Mix of TensorSpec input + ScalarSpec + TensorSpec output.""" - return [ + return _stamped([ TensorSpec("x", [4], torch.float32, init_value=torch.randn), ScalarSpec("alpha", torch.float32, 2.5), - TensorSpec("y", [4], torch.float32, is_output=True), - ] + TensorSpec("y", [4], torch.float32), + ], {"x": "in", "y": "out"}) class TestScalarMixedSpecs: """Mixed TensorSpec + ScalarSpec exercises the scalar path through run().""" - def test_scalar_passed_as_ctypes_to_execute(self, mixed_specs, tmp_path): + def test_scalar_passed_as_ctypes_to_execute(self, mixed_specs, build_dir): """run() forwards args in the user-declared spec order: for ``[Tensor x, Scalar alpha, Tensor y]`` the args list passed to execute_compiled is ``[x, alpha, y]`` (scalars are encoded via ctypes but stay in their declared position).""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() observed: dict[str, object] = {} @@ -1246,23 +1180,21 @@ def fake_execute(work_dir, args, **_kwargs): def golden_fn(scratch): scratch["y"][:] = scratch["x"] + scratch["alpha"] - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, fake_execute=fake_execute) + compile_p, exec_p = _patch_compile_and_execute(build_dir, fake_execute=fake_execute) with compile_p, exec_p: - r = run(program=object(), specs=mixed_specs, golden_fn=golden_fn) + r = run(fn=object(), specs=mixed_specs, golden_fn=golden_fn) assert r.passed, f"unexpected failure: {r.error}" - assert r.work_dir == compiled_dir + assert r.work_dir == build_dir # Spec order: x (input tensor), alpha (scalar), y (output tensor) assert isinstance(observed["arg0"], torch.Tensor) assert isinstance(observed["arg1"], ctypes.c_float) assert isinstance(observed["arg2"], torch.Tensor) assert observed["arg1"].value == pytest.approx(2.5) - def test_scalar_persisted_to_pt(self, mixed_specs, tmp_path): + def test_scalar_persisted_to_pt(self, mixed_specs, build_dir): """After a successful run, work_dir/data/in/{name}.pt must exist with the spec's value as a 0-dim tensor of the spec's dtype.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() def fake_execute(_work_dir, args, **_kwargs): # Spec order: [x (in tensor), alpha (scalar), y (out tensor)] @@ -1271,12 +1203,12 @@ def fake_execute(_work_dir, args, **_kwargs): def golden_fn(scratch): scratch["y"][:] = scratch["x"] + scratch["alpha"] - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, fake_execute=fake_execute) + compile_p, exec_p = _patch_compile_and_execute(build_dir, fake_execute=fake_execute) with compile_p, exec_p: - r = run(program=object(), specs=mixed_specs, golden_fn=golden_fn, save_data=True) + r = run(fn=object(), specs=mixed_specs, golden_fn=golden_fn, save_data=True) assert r.passed, f"unexpected failure: {r.error}" - scalar_path = compiled_dir / "data" / "in" / "alpha.pt" + scalar_path = build_dir / "data" / "in" / "alpha.pt" assert scalar_path.is_file() loaded = torch.load(scalar_path, weights_only=True) assert loaded.ndim == 0 @@ -1329,8 +1261,7 @@ def test_duplicate_spec_names_rejected_before_dict_conversion(self, tmp_path): def test_scalar_pt_loaded_from_cache(self, mixed_specs, tmp_path): """When golden_data has {name}.pt, the cached value (not the spec value) must be used for ctypes encoding.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) cache = tmp_path / "cache" # Pre-populate the cache: x, y, alpha.pt — alpha=10.0 (different from spec's 2.5) x = torch.tensor([1.0, 2.0, 3.0, 4.0]) @@ -1350,7 +1281,7 @@ def fake_execute(_work_dir, args, **_kwargs): compile_p, exec_p = _patch_compile_and_execute(compiled_dir, fake_execute=fake_execute) with compile_p, exec_p: r = run( - program=object(), + fn=object(), specs=mixed_specs, golden_data=str(cache), ) @@ -1362,8 +1293,7 @@ def fake_execute(_work_dir, args, **_kwargs): def test_custom_comparator_receives_cached_scalar(self, mixed_specs, tmp_path): """Validation exposes the replayed scalar, not the spec default.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) cache = tmp_path / "cache" x = torch.tensor([1.0, 2.0, 3.0, 4.0]) y_golden = torch.tensor([11.0, 12.0, 13.0, 14.0]) @@ -1388,7 +1318,7 @@ def compare(actual, expected, *, inputs, **_kwargs): ) with compile_p, exec_p: result = run( - program=object(), + fn=object(), specs=mixed_specs, golden_data=str(cache), compare_fn={"y": compare}, @@ -1401,8 +1331,7 @@ def compare(actual, expected, *, inputs, **_kwargs): def test_missing_scalar_pt_in_cache_fails(self, mixed_specs, tmp_path): """golden_data with a ScalarSpec must include {name}.pt — missing it should produce a ``missing files`` error.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) cache = tmp_path / "cache" x = torch.tensor([1.0, 2.0, 3.0, 4.0]) y_golden = torch.tensor([3.5, 4.5, 5.5, 6.5]) @@ -1412,15 +1341,14 @@ def test_missing_scalar_pt_in_cache_fails(self, mixed_specs, tmp_path): compile_p, exec_p = _patch_compile_and_execute(compiled_dir) with compile_p, exec_p: - r = run(program=object(), specs=mixed_specs, golden_data=str(cache)) + r = run(fn=object(), specs=mixed_specs, golden_data=str(cache)) assert not r.passed assert "alpha.pt" in (r.error or "") def test_scalar_pt_non_zero_dim_fails(self, mixed_specs, tmp_path): """A non-0-dim tensor in {name}.pt must fail via RunResult, not raise.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) cache = tmp_path / "cache" x = torch.tensor([1.0, 2.0, 3.0, 4.0]) y_golden = torch.tensor([3.5, 4.5, 5.5, 6.5]) @@ -1431,15 +1359,14 @@ def test_scalar_pt_non_zero_dim_fails(self, mixed_specs, tmp_path): compile_p, exec_p = _patch_compile_and_execute(compiled_dir) with compile_p, exec_p: - r = run(program=object(), specs=mixed_specs, golden_data=str(cache)) + r = run(fn=object(), specs=mixed_specs, golden_data=str(cache)) assert not r.passed assert "0-dim" in (r.error or "") def test_scalar_pt_dtype_mismatch_fails(self, mixed_specs, tmp_path): """If {name}.pt has a different dtype than the spec, fail loudly.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() + compiled_dir = _make_build_dir(tmp_path) cache = tmp_path / "cache" x = torch.tensor([1.0, 2.0, 3.0, 4.0]) y_golden = torch.tensor([3.5, 4.5, 5.5, 6.5]) @@ -1450,15 +1377,13 @@ def test_scalar_pt_dtype_mismatch_fails(self, mixed_specs, tmp_path): compile_p, exec_p = _patch_compile_and_execute(compiled_dir) with compile_p, exec_p: - r = run(program=object(), specs=mixed_specs, golden_data=str(cache)) + r = run(fn=object(), specs=mixed_specs, golden_data=str(cache)) assert not r.passed assert "dtype mismatch" in (r.error or "") - def test_golden_fn_receives_scalar_python_value(self, mixed_specs, tmp_path): + def test_golden_fn_receives_scalar_python_value(self, mixed_specs, build_dir): """golden_fn(scratch) must see the scalar as a python float keyed by name.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() captured: dict[str, object] = {} @@ -1471,9 +1396,9 @@ def fake_execute(_work_dir, args, **_kwargs): # Spec order: [x (in tensor), alpha (scalar), y (out tensor)] args[2][:] = args[0] + args[1].value - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, fake_execute=fake_execute) + compile_p, exec_p = _patch_compile_and_execute(build_dir, fake_execute=fake_execute) with compile_p, exec_p: - r = run(program=object(), specs=mixed_specs, golden_fn=golden_fn) + r = run(fn=object(), specs=mixed_specs, golden_fn=golden_fn) assert r.passed, f"unexpected failure: {r.error}" assert captured["alpha"] == pytest.approx(2.5) @@ -1483,10 +1408,8 @@ def fake_execute(_work_dir, args, **_kwargs): class TestStageOrder: """compute_golden runs before runtime — fail-fast on golden_fn errors.""" - def test_compute_golden_runs_before_runtime(self, three_kinds_specs, tmp_path): + def test_compute_golden_runs_before_runtime(self, three_kinds_specs, build_dir): """golden_fn is invoked before execute_compiled.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() order: list[str] = [] @@ -1502,11 +1425,11 @@ def fake_execute(_work_dir, tensors, **_kwargs): tensors[2][:] = tensors[2] + 100 compile_p, exec_p = _patch_compile_and_execute( - compiled_dir, fake_execute=fake_execute, + build_dir, fake_execute=fake_execute, ) with compile_p, exec_p: r = run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=golden_fn, ) @@ -1514,10 +1437,8 @@ def fake_execute(_work_dir, tensors, **_kwargs): assert r.passed, f"unexpected failure: {r.error}" assert order == ["golden", "runtime"] - def test_golden_fn_error_short_circuits_runtime(self, three_kinds_specs, tmp_path): + def test_golden_fn_error_short_circuits_runtime(self, three_kinds_specs, build_dir): """A typo / shape bug in golden_fn surfaces before execute_compiled runs.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() def bad_golden(_tensors): raise RuntimeError("typo in golden_fn") @@ -1526,11 +1447,11 @@ def exec_must_not_run(*_args, **_kwargs): pytest.fail("execute_compiled ran despite golden_fn error") compile_p, exec_p = _patch_compile_and_execute( - compiled_dir, fake_execute=exec_must_not_run, + build_dir, fake_execute=exec_must_not_run, ) with compile_p, exec_p, pytest.raises(RuntimeError, match="typo in golden_fn"): run( - program=object(), + fn=object(), specs=three_kinds_specs, golden_fn=bad_golden, ) @@ -1539,11 +1460,9 @@ def exec_must_not_run(*_args, **_kwargs): class TestConfigForwarding: """compile_cfg / runtime_cfg pass-through to pypto entry points.""" - def test_compile_cfg_forwarded_to_ir_compile(self, three_kinds_specs, tmp_path): + def test_compile_cfg_forwarded_to_ir_compile(self, three_kinds_specs, build_dir): """Keys in compile_cfg reach ir.compile as kwargs.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - fake = _FakeCompiled(compiled_dir) + fake = _FakeCompiled(build_dir) captured: dict = {} @@ -1554,7 +1473,7 @@ def fake_compile(_program, **kwargs): with patch("pypto.ir.compile", side_effect=fake_compile), \ patch("pypto.runtime.execute_compiled"): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, compile_cfg=dict(dump_passes=False, profiling=True), ) @@ -1564,21 +1483,19 @@ def fake_compile(_program, **kwargs): assert captured["profiling"] is True def test_runtime_cfg_forwarded_to_execute_compiled( - self, three_kinds_specs, tmp_path, + self, three_kinds_specs, build_dir, ): """Non-DFX keys in runtime_cfg reach execute_compiled as kwargs.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() captured: dict = {} def fake_execute(_work_dir, _tensors, **kwargs): captured.update(kwargs) - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, fake_execute=fake_execute) + compile_p, exec_p = _patch_compile_and_execute(build_dir, fake_execute=fake_execute) with compile_p, exec_p: r = run( - program=object(), + fn=object(), specs=three_kinds_specs, runtime_cfg=dict( platform="a2a3sim", @@ -1592,10 +1509,8 @@ def fake_execute(_work_dir, _tensors, **kwargs): assert captured["device_id"] == 3 assert captured["pto_isa_commit"] == "deadbeef" - def test_dump_args_forwarded_as_dfx_option(self, three_kinds_specs, tmp_path): + def test_dump_args_forwarded_as_dfx_option(self, three_kinds_specs, build_dir): """enable_dump_args is bundled into the execute_compiled DFX options.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() captured: dict = {} dfx = object() @@ -1606,14 +1521,14 @@ def test_dump_args_forwarded_as_dfx_option(self, three_kinds_specs, tmp_path): def fake_execute(_work_dir, _tensors, **kwargs): captured.update(kwargs) - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, fake_execute=fake_execute) + compile_p, exec_p = _patch_compile_and_execute(build_dir, fake_execute=fake_execute) with ( compile_p, exec_p, patch.dict(sys.modules, {"pypto.runtime.runner": runner_mod}), ): r = run( - program=object(), + fn=object(), specs=three_kinds_specs, runtime_cfg=dict(enable_dump_args=2), ) @@ -1769,49 +1684,43 @@ def test_missing_dir_raises(self, tmp_path): class TestLogLevelConsumption: """`runtime_cfg['log_level']` is consumed as a harness-only key.""" - def test_log_level_invokes_configure_log(self, three_kinds_specs, tmp_path): + def test_log_level_invokes_configure_log(self, three_kinds_specs, build_dir): """runtime_cfg.log_level → configure_log(level) is called.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compile_p, exec_p = _patch_compile_and_execute(compiled_dir) + compile_p, exec_p = _patch_compile_and_execute(build_dir) with compile_p, exec_p, \ patch("pypto.runtime.log_config.configure_log") as mock_cfg: run( - program=object(), + fn=object(), specs=three_kinds_specs, runtime_cfg=dict(platform="a2a3sim", device_id=0, log_level="debug"), ) mock_cfg.assert_called_once_with("debug") def test_log_level_not_forwarded_to_execute_compiled( - self, three_kinds_specs, tmp_path, + self, three_kinds_specs, build_dir, ): """log_level is popped — execute_compiled does NOT receive it.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() captured: dict = {} def fake_execute(_w, _t, **kw): captured.update(kw) - compile_p, exec_p = _patch_compile_and_execute(compiled_dir, fake_execute=fake_execute) + compile_p, exec_p = _patch_compile_and_execute(build_dir, fake_execute=fake_execute) with compile_p, exec_p, patch("pypto.runtime.log_config.configure_log"): run( - program=object(), + fn=object(), specs=three_kinds_specs, runtime_cfg=dict(platform="a2a3sim", device_id=0, log_level="debug"), ) assert "log_level" not in captured - def test_no_log_level_skips_configure_log(self, three_kinds_specs, tmp_path): + def test_no_log_level_skips_configure_log(self, three_kinds_specs, build_dir): """No log_level key → configure_log not called.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - compile_p, exec_p = _patch_compile_and_execute(compiled_dir) + compile_p, exec_p = _patch_compile_and_execute(build_dir) with compile_p, exec_p, \ patch("pypto.runtime.log_config.configure_log") as mock_cfg: run( - program=object(), + fn=object(), specs=three_kinds_specs, runtime_cfg=dict(platform="a2a3sim", device_id=0), ) @@ -1893,7 +1802,7 @@ def _benchmark(compiled, args, **kwargs): fake_runtime.benchmark = _benchmark compiled = object() tensors = {"x": torch.zeros(1)} - monkeypatch.setattr("golden.runner._l3_ordered_args", lambda *_a: ["ORDERED"]) + monkeypatch.setattr("golden.runner._ordered_args", lambda *_a, **_k: ["ORDERED"]) monkeypatch.setattr("golden.runner._l3_run_config", lambda _cfg: "RUNCFG") with patch.dict(sys.modules, {"pypto.runtime": fake_runtime}): @@ -1927,7 +1836,7 @@ def test_nonresident_validation_precedes_benchmark_mutation( ): """Benchmark reuse must not overwrite the dedicated correctness result.""" compiled = _FakeCompiled(tmp_path) - specs = [TensorSpec("y", [1], torch.float32, is_output=True)] + specs = _stamped([TensorSpec("y", [1], torch.float32)], {"y": "out"}) monkeypatch.setenv("PYPTO_BENCH", "1") def _dispatch(_compiled, _specs, tensors, _scalars, _runtime_cfg): @@ -1946,10 +1855,10 @@ def _benchmark(_compiled, _specs, tensors, *_args): ): if use_jit: fn = types.SimpleNamespace(compile=MagicMock(return_value=compiled)) - result = run_jit(fn, specs, golden_fn=lambda values: values["y"].zero_()) + result = run(fn, specs, golden_fn=lambda values: values["y"].zero_()) else: result = run( - program=object(), specs=specs, + fn=object(), specs=specs, golden_fn=lambda values: values["y"].zero_(), ) @@ -1967,7 +1876,7 @@ def _benchmark(*_args, **_kwargs): with patch.dict(sys.modules, {"pypto.runtime": fake_runtime}): with pytest.raises(RuntimeError, match="DEVICE DISPATCH FAILED"): if l3: - monkeypatch.setattr("golden.runner._l3_ordered_args", lambda *_a: []) + monkeypatch.setattr("golden.runner._ordered_args", lambda *_a, **_k: []) monkeypatch.setattr("golden.runner._l3_run_config", lambda _cfg: "RUNCFG") _run_benchmark_l3(object(), [], {}, {}, {}, rounds=1, warmup=1) else: @@ -1980,7 +1889,7 @@ def _benchmark(*_args, **_kwargs): fake_runtime = types.ModuleType("pypto.runtime") fake_runtime.benchmark = _benchmark - monkeypatch.setattr("golden.runner._l3_ordered_args", lambda *_a: []) + monkeypatch.setattr("golden.runner._ordered_args", lambda *_a, **_k: []) monkeypatch.setattr("golden.runner._l3_run_config", lambda _cfg: "RUNCFG") with patch.dict(sys.modules, {"pypto.runtime": fake_runtime}): assert _run_benchmark_l3(object(), [], {}, {}, {}, rounds=1, warmup=1) is None @@ -2018,48 +1927,182 @@ def test_l2_benchmark_rejects_stepped_scalar(): class TestL3ParameterAbi: + """L3 goes through the same spec-order ABI gate as L2: nothing rebinds a + spec list that drifted from the kernel signature.""" + @staticmethod - def _compiled(*names): - return types.SimpleNamespace( - _get_metadata=lambda: ( - [types.SimpleNamespace(name=name) for name in names], - None, - None, - ) - ) + def _compiled(*infos): + return types.SimpleNamespace(_get_metadata=lambda: (list(infos), None, None)) + + @staticmethod + def _x_info(name="x__ssa_v0"): + return _l3_info(name, shape=[1], dtype=torch.float32) + + @staticmethod + def _epoch_info(name="epoch__ssa_v0"): + return _l3_info(name, dtype=torch.int32) def test_extra_spec_rejected_for_stale_artifact(self): - compiled = self._compiled("x__ssa_v0") + compiled = self._compiled(self._x_info()) specs = [ TensorSpec("x", [1], torch.float32), ScalarSpec("moe_epoch_base", torch.int32, 0), ] - with pytest.raises(ValueError, match="moe_epoch_base.*recompile"): - _l3_ordered_args( - compiled, - specs, - {"x": torch.zeros(1)}, - {"moe_epoch_base": specs[1]}, - ) + with _l3_abi_environment(): + with pytest.raises(ValueError, match="moe_epoch_base.*recompile"): + _validate_compiled_spec_abi(compiled, specs) def test_compiled_parameter_without_spec_rejected(self): - compiled = self._compiled("x__ssa_v0", "moe_epoch_base__ssa_v0") - specs = [TensorSpec("x", [1], torch.float32)] - with pytest.raises(ValueError, match="moe_epoch_base.*recompile"): - _l3_ordered_args(compiled, specs, {"x": torch.zeros(1)}, {}) - - def test_exact_abi_reorders_and_strips_terminal_ssa_suffix(self): - compiled = self._compiled("epoch__ssa_v3", "x__ssa_v0") - epoch = ScalarSpec("epoch", torch.int32, 7) - x = torch.zeros(1) - ordered = _l3_ordered_args( - compiled, - [TensorSpec("x", [1], torch.float32), epoch], - {"x": x}, - {"epoch": epoch}, + compiled = self._compiled( + self._x_info(), self._epoch_info("moe_epoch_base__ssa_v0") ) - assert ordered[0] is epoch.value - assert ordered[1] is x + specs = [TensorSpec("x", [1], torch.float32)] + with _l3_abi_environment(): + with pytest.raises(ValueError, match="moe_epoch_base.*recompile"): + _validate_compiled_spec_abi(compiled, specs) + + def test_spec_order_mismatch_rejected(self): + """Same names, wrong order — previously rescued by the name-keyed reorder.""" + compiled = self._compiled(self._epoch_info("epoch__ssa_v3"), self._x_info()) + specs = [ + TensorSpec("x", [1], torch.float32), + ScalarSpec("epoch", torch.int32, 7), + ] + with _l3_abi_environment(): + with pytest.raises(ValueError, match="parameter order"): + _validate_compiled_spec_abi(compiled, specs) + + def test_matching_order_accepted_and_terminal_ssa_suffix_stripped(self): + compiled = self._compiled(self._x_info(), self._epoch_info("epoch__ssa_v3")) + specs = [ + TensorSpec("x", [1], torch.float32), + ScalarSpec("epoch", torch.int32, 7), + ] + with _l3_abi_environment(): + _validate_compiled_spec_abi(compiled, specs) + assert specs[0].direction == "in" + + +class _ResidentRT: + """Recording stand-in for the prepared ``DistributedWorker``. + + Every worker call lands in :attr:`ops` as an ``(op, payload)`` pair, so a + test asserts on what happened instead of embedding assertions in the fake. + *on_dispatch* is called with the 1-based dispatch index and may raise; + *readback* is the tensor ``copy_stacked_from`` writes into its destination. + """ + + def __init__(self, *, stacked_handle=None, on_dispatch=None, readback=None): + self.ops: list[tuple[str, object]] = [] + self._stacked_handle = stacked_handle + self._on_dispatch = on_dispatch + self._readback = readback + + def payloads(self, op: str) -> list: + return [payload for name, payload in self.ops if name == op] + + def kinds(self) -> list[str]: + return [name for name, _ in self.ops] + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + def alloc_stacked_tensor(self, host, worker_ids=None): + self.ops.append(("alloc_stacked", (tuple(host.shape), worker_ids, host))) + if self._stacked_handle is not None: + return self._stacked_handle + return types.SimpleNamespace(full_shape=tuple(host.shape)) + + def alloc_tensor(self, shape, dtype, *, init=None, worker_id=0): + self.ops.append(("alloc", (tuple(shape), dtype, init, worker_id))) + return types.SimpleNamespace(worker_id=worker_id) + + def free_stacked_tensor(self, handle): + self.ops.append(("free_stacked", handle)) + + def free_tensor(self, handle, *, worker_id=0): + self.ops.append(("free", (handle, worker_id))) + + def copy_stacked_from(self, handle, host): + self.ops.append(("readback", handle)) + if self._readback is not None: + host.copy_(self._readback) + + def __call__(self, *args, config=None): + self.ops.append(("dispatch", (args, config))) + if self._on_dispatch is not None: + self._on_dispatch(len(self.payloads("dispatch"))) + + +def _resident_dcp(rt, prepare_log=None): + """A ``DistributedCompiledProgram`` stand-in whose ``prepare()`` yields *rt*.""" + + class _DCP: + def prepare(self, *args, **kwargs): + if prepare_log is not None: + prepare_log.append((args, kwargs)) + return rt + + return _DCP + + +def _resident_modules(dcp_cls, *, bench_capture=None, parse_stats=None, runtime=None): + """``sys.modules`` entries `_run_l3_resident` imports, as fakes. + + Passing *bench_capture* / *parse_stats* adds the ``pypto.runtime.bench`` and + ``pypto.runtime.log_config`` stubs the benchmark branch needs; *runtime* + installs a ``pypto.runtime`` carrying ``StackedDeviceTensor``. + """ + dcp_mod = types.ModuleType("pypto.ir.distributed_compiled_program") + dcp_mod.DistributedCompiledProgram = dcp_cls + mods = {"pypto.ir.distributed_compiled_program": dcp_mod} + if bench_capture is not None: + bench = types.ModuleType("pypto.runtime.bench") + bench._STRACE_LOG_LEVEL = "v9" + bench._capture_fd_stderr = bench_capture + bench._parse_stats_from_strace = parse_stats + log = types.ModuleType("pypto.runtime.log_config") + log.configure_log = lambda _level: None + log.current_level = lambda: "v0" + mods["pypto.runtime.bench"] = bench + mods["pypto.runtime.log_config"] = log + if runtime is not None: + mods["pypto.runtime"] = runtime + return mods + + +class _NullCapture: + """Stand-in for ``_capture_fd_stderr``: touches the log path, diverts nothing.""" + + def __init__(self, path): + self.path = path + + def __enter__(self): + self.path.touch() + return None + + def __exit__(self, *_a): + return False + + +def _stub_l3_helpers(monkeypatch, pure_out=frozenset()): + """Bypass the real metadata / RunConfig helpers around `_run_l3_resident`.""" + import golden.runner as R + + monkeypatch.setattr(R, "_l3_pure_out_names", lambda _c: set(pure_out)) + monkeypatch.setattr(R, "_l3_run_config", lambda _cfg: "RUNCFG") + + +_RESIDENT_KWARGS = dict( + runtime_cfg={"platform": "a2a3"}, + golden_outputs=None, + rtol=1e-5, + atol=1e-5, + compare_fn={}, +) class TestResidentPath: @@ -2070,57 +2113,37 @@ def _resident_specs(self): TensorSpec("x", [4], torch.float32, init_value=torch.randn), # per-call input TensorSpec("w", [4], torch.float32, init_value=torch.ones, # whole-tensor resident resident=0), - TensorSpec("y", [4], torch.float32, is_output=True), # output + TensorSpec("y", [4], torch.float32), # output ] - def test_resident_routes_to_l3_resident_not_single_chip(self, tmp_path): + def test_resident_routes_to_l3_resident_not_single_chip(self, build_dir): """A resident spec dispatches via _run_l3_resident; execute_compiled never runs.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - fake = _FakeCompiled(compiled_dir) with ( - patch("pypto.ir.compile", return_value=fake), + patch("pypto.ir.compile", return_value=_FakeCompiled(build_dir)), patch( "pypto.runtime.execute_compiled", side_effect=lambda *a, **k: pytest.fail("single-chip path must not run for resident"), ), patch("golden.runner._run_l3_resident", return_value=None) as l3res, ): - r = run(program=object(), specs=self._resident_specs()) + r = run(fn=object(), specs=self._resident_specs()) assert r.passed, f"unexpected failure: {r.error}" l3res.assert_called_once() - def test_runtime_dir_resident_disables_embedded_benchmark( - self, - tmp_path, - monkeypatch, - ): + def test_runtime_dir_resident_disables_embedded_benchmark(self, tmp_path, monkeypatch): prebuilt = tmp_path / "prebuilt" prebuilt.mkdir() specs = self._resident_specs() - directions = { - "x": _FakeParamDirection.In, - "w": _FakeParamDirection.In, - "y": _FakeParamDirection.Out, - } - compiled = types.SimpleNamespace( - output_dir=prebuilt, - _get_metadata=lambda: ( - [ - _l3_info( - f"{spec.name}__ssa_v0", - shape=spec.shape, - dtype=spec.dtype, - direction=directions[spec.name], - ) - for spec in specs - ], - None, - None, - ), - ) + directions = {"x": "In", "w": "In", "y": "Out"} + compiled = _artifact(prebuilt, *[ + _l3_info( + f"{spec.name}__ssa_v0", shape=spec.shape, dtype=spec.dtype, + direction=getattr(_FakeParamDirection, directions[spec.name]), + ) + for spec in specs + ]) monkeypatch.setenv("PYPTO_BENCH", "1") with ( @@ -2129,463 +2152,183 @@ def test_runtime_dir_resident_disables_embedded_benchmark( patch("golden.runner._run_l3_resident", return_value=None) as l3res, patch("pypto.runtime.execute_compiled") as execute, ): - result = run(program=object(), specs=specs, runtime_dir=str(prebuilt)) + result = run(fn=object(), specs=specs, runtime_dir=str(prebuilt)) assert result.passed, result.error assert l3res.call_args.kwargs["benchmark_enabled"] is False execute.assert_not_called() - def test_resident_benchmark_reuses_handle_and_advances_stepped_scalar( - self, monkeypatch - ): + def test_resident_benchmark_reuses_handle_and_advances_stepped_scalar(self, monkeypatch): """The resident L3 benchmark reuses one handle in persistent mode.""" - import golden.runner as R - - calls = {"prepare": None, "events": []} state_handle = object() - initial_state = torch.arange(8, dtype=torch.float32).reshape(2, 4) state_spec = TensorSpec( - "state", [2, 4], torch.float32, init_value=initial_state, - is_output=True, resident="stacked", + "state", [2, 4], torch.float32, + init_value=torch.arange(8, dtype=torch.float32).reshape(2, 4), + resident="stacked", ) state_init = state_spec.create_tensor() - - class _FakeRT: - def __enter__(self): - return self - - def __exit__(self, *_a): - return False - - def alloc_stacked_tensor(self, host, worker_ids=None): - assert worker_ids is None - assert torch.equal(host, state_init) - calls["events"].append(("alloc", state_handle)) - return state_handle - - def free_stacked_tensor(self, handle): - assert handle is state_handle - calls["events"].append(("free", handle)) - - def __call__(self, *args, config=None): - assert config == "RUNCFG" - assert len(args) == 2 - calls["events"].append(("dispatch", (args[0], args[1].item()))) - - class _FakeDCP: - def prepare(self, *args, **kwargs): - calls["prepare"] = (args, kwargs) - return _FakeRT() - - class _Capture: - def __init__(self, path): - self.path = path - - def __enter__(self): - self.path.touch() - return None - - def __exit__(self, *_a): - return False - - fake_dcp = types.ModuleType("pypto.ir.distributed_compiled_program") - fake_dcp.DistributedCompiledProgram = _FakeDCP - fake_bench = types.ModuleType("pypto.runtime.bench") - fake_bench._STRACE_LOG_LEVEL = "v9" - fake_bench._capture_fd_stderr = _Capture - fake_bench._parse_stats_from_strace = lambda *_a, **_k: types.SimpleNamespace( - host_wall_us=[] - ) - fake_log = types.ModuleType("pypto.runtime.log_config") - fake_log.configure_log = lambda _level: None - fake_log.current_level = lambda: "v0" + epoch_spec = ScalarSpec("epoch", torch.int32, 0, benchmark_step=43) + rt = _ResidentRT(stacked_handle=state_handle) + prepared = [] monkeypatch.setenv("PYPTO_BENCH", "1") monkeypatch.setenv("PYPTO_BENCH_ROUNDS", "3") monkeypatch.setenv("PYPTO_BENCH_WARMUP", "2") - monkeypatch.setattr(R, "_l3_ordered_names", lambda _c: ["state", "epoch"]) - monkeypatch.setattr(R, "_l3_pure_out_names", lambda _c: set()) - monkeypatch.setattr(R, "_l3_run_config", lambda _cfg: "RUNCFG") + _stub_l3_helpers(monkeypatch) - with patch.dict( - sys.modules, - { - "pypto.ir.distributed_compiled_program": fake_dcp, - "pypto.runtime.bench": fake_bench, - "pypto.runtime.log_config": fake_log, - }, - ): - result = R._run_l3_resident( - compiled=_FakeDCP(), - tensor_specs=[state_spec], - tensors={"state": state_init}, - scalar_specs_eff={ - "epoch": ScalarSpec( - "epoch", torch.int32, 0, benchmark_step=43 - ) - }, - runtime_cfg={"platform": "a2a3"}, - golden_outputs=None, - rtol=1e-5, - atol=1e-5, - compare_fn={}, + dcp = _resident_dcp(rt, prepare_log=prepared) + mods = _resident_modules( + dcp, + bench_capture=_NullCapture, + parse_stats=lambda *_a, **_k: types.SimpleNamespace(host_wall_us=[]), + ) + with patch.dict(sys.modules, mods): + result = _run_l3_resident( + compiled=dcp(), specs=[state_spec, epoch_spec], + tensors={"state": state_init}, scalar_specs_eff={"epoch": epoch_spec}, + **_RESIDENT_KWARGS, ) assert result is None - assert calls["prepare"] == ( - ("RUNCFG",), - {"persistent": True, "reset_persistent_windows": False}, - ) - assert [kind for kind, _ in calls["events"]] == [ - "alloc", "dispatch", "dispatch", "dispatch", "dispatch", "dispatch", "free", - ] - dispatches = [value for kind, value in calls["events"] if kind == "dispatch"] - assert [epoch for _, epoch in dispatches] == [0, 43, 86, 129, 172] - assert all(handle is state_handle for handle, _ in dispatches) + assert prepared == [(("RUNCFG",), {"persistent": True, "reset_persistent_windows": False})] + assert rt.kinds() == ["alloc_stacked", *["dispatch"] * 5, "free_stacked"] + dispatched = rt.payloads("dispatch") + assert [args[1].item() for args, _cfg in dispatched] == [0, 43, 86, 129, 172] + assert all(args[0] is state_handle for args, _cfg in dispatched) + assert torch.equal(rt.payloads("alloc_stacked")[0][2], state_init) def test_resident_benchmark_propagates_later_dispatch_failure(self, monkeypatch): - import golden.runner as R - - calls = 0 - - class _FakeRT: - def __enter__(self): - return self - - def __exit__(self, *_args): - return False - - def __call__(self, *_args, **_kwargs): - nonlocal calls - calls += 1 - if calls == 2: - raise RuntimeError("persistent dispatch failed") - - class _FakeDCP: - def prepare(self, *_args, **_kwargs): - return _FakeRT() - - class _Capture: - def __init__(self, path): - self.path = path + def _fail_on_second(index): + if index == 2: + raise RuntimeError("persistent dispatch failed") - def __enter__(self): - self.path.touch() - - def __exit__(self, *_args): - return False - - fake_dcp = types.ModuleType("pypto.ir.distributed_compiled_program") - fake_dcp.DistributedCompiledProgram = _FakeDCP - fake_bench = types.ModuleType("pypto.runtime.bench") - fake_bench._STRACE_LOG_LEVEL = "v9" - fake_bench._capture_fd_stderr = _Capture - fake_bench._parse_stats_from_strace = MagicMock() - fake_log = types.ModuleType("pypto.runtime.log_config") - fake_log.configure_log = lambda _level: None - fake_log.current_level = lambda: "v0" + rt = _ResidentRT(on_dispatch=_fail_on_second) + parse_stats = MagicMock() monkeypatch.setenv("PYPTO_BENCH", "1") monkeypatch.setenv("PYPTO_BENCH_ROUNDS", "2") monkeypatch.setenv("PYPTO_BENCH_WARMUP", "1") - monkeypatch.setattr(R, "_l3_ordered_names", lambda _compiled: []) - monkeypatch.setattr(R, "_l3_pure_out_names", lambda _compiled: set()) - monkeypatch.setattr(R, "_l3_run_config", lambda _cfg: "RUNCFG") + _stub_l3_helpers(monkeypatch) + dcp = _resident_dcp(rt) + mods = _resident_modules(dcp, bench_capture=_NullCapture, parse_stats=parse_stats) with ( - patch.dict( - sys.modules, - { - "pypto.ir.distributed_compiled_program": fake_dcp, - "pypto.runtime.bench": fake_bench, - "pypto.runtime.log_config": fake_log, - }, - ), + patch.dict(sys.modules, mods), pytest.raises(RuntimeError, match="persistent dispatch failed"), ): - R._run_l3_resident( - compiled=_FakeDCP(), tensor_specs=[], tensors={}, scalar_specs_eff={}, - runtime_cfg={}, golden_outputs=None, rtol=1e-5, atol=1e-5, - compare_fn={}, + _run_l3_resident( + compiled=dcp(), specs=[], tensors={}, scalar_specs_eff={}, + runtime_cfg={}, golden_outputs=None, rtol=1e-5, atol=1e-5, compare_fn={}, ) - assert calls == 2 - fake_bench._parse_stats_from_strace.assert_not_called() - - @staticmethod - def _fake_dcp_module(): - """A stub ``pypto.ir.distributed_compiled_program`` module exposing a - ``DistributedCompiledProgram`` class, so the isinstance branch in - ``_run_l3_resident`` is exercised deterministically even where the real - (heavy) submodule is not importable.""" - mod = types.ModuleType("pypto.ir.distributed_compiled_program") - class _DCP: # noqa: N801 — mirror the real class name for isinstance - pass - mod.DistributedCompiledProgram = _DCP - return mod - - def test_resident_on_non_l3_fails_cleanly(self, tmp_path): - """A resident spec against a non-L3 compiled program fails via RunResult.""" - compiled_dir = tmp_path / "build" - compiled_dir.mkdir() - fake = _FakeCompiled(compiled_dir) # not a DistributedCompiledProgram - - with ( - patch.dict( - sys.modules, - {"pypto.ir.distributed_compiled_program": self._fake_dcp_module()}, - ), - patch("pypto.ir.compile", return_value=fake), - ): - r = run(program=object(), specs=self._resident_specs()) - - assert not r.passed - assert "only supported for L3" in (r.error or "") + assert len(rt.payloads("dispatch")) == 2 + parse_stats.assert_not_called() def test_run_l3_resident_stacked_uses_alloc_stacked(self, monkeypatch): - """A resident="stacked" spec uploads via alloc_stacked_tensor and frees via free_stacked_tensor.""" - import golden.runner as R - - calls = {"stacked": [], "freed": 0, "dispatched": 0} - - class _FakeRT: - last_run_timing = None - - def __enter__(self): - return self - - def __exit__(self, *_a): - return False - - def alloc_stacked_tensor(self, host, worker_ids=None): - calls["stacked"].append((tuple(host.shape), worker_ids)) - return ("stacked_handle", tuple(host.shape)) - - def alloc_tensor(self, *_a, **_k): - raise AssertionError('resident="stacked" must not use alloc_tensor') - - def free_stacked_tensor(self, _h): - calls["freed"] += 1 - - def free_tensor(self, _h): - raise AssertionError("stacked handle must be freed via free_stacked_tensor") - - def __call__(self, *_args, config=None): - calls["dispatched"] += 1 - - class _FakeDCP: - def prepare(self): - return _FakeRT() - - fake_mod = types.ModuleType("pypto.ir.distributed_compiled_program") - fake_mod.DistributedCompiledProgram = _FakeDCP - - specs = [TensorSpec("w", [2, 4], torch.float32, init_value=torch.ones, resident="stacked")] - tensors = {"w": torch.ones(2, 4)} - # Avoid real pypto.runtime / backend by stubbing the metadata + config helpers. - monkeypatch.setattr(R, "_l3_ordered_names", lambda _c: ["w"]) - monkeypatch.setattr(R, "_l3_pure_out_names", lambda _c: set()) - monkeypatch.setattr(R, "_l3_run_config", lambda _cfg: "RUNCFG") - - with patch.dict(sys.modules, {"pypto.ir.distributed_compiled_program": fake_mod}): - out = R._run_l3_resident( - compiled=_FakeDCP(), - tensor_specs=specs, - tensors=tensors, - scalar_specs_eff={}, - runtime_cfg={"platform": "a2a3"}, - golden_outputs=None, - rtol=1e-5, - atol=1e-5, - compare_fn={}, + """A resident="stacked" spec uploads and frees as a stack, never per tensor.""" + rt = _ResidentRT() + _stub_l3_helpers(monkeypatch) + dcp = _resident_dcp(rt) + + with patch.dict(sys.modules, _resident_modules(dcp)): + out = _run_l3_resident( + compiled=dcp(), + specs=[TensorSpec("w", [2, 4], torch.float32, init_value=torch.ones, + resident="stacked")], + tensors={"w": torch.ones(2, 4)}, scalar_specs_eff={}, + **_RESIDENT_KWARGS, ) assert out is None - assert calls["dispatched"] == 1 - assert calls["stacked"] == [((2, 4), None)] # identity worker_ids - assert calls["freed"] == 1 + assert rt.kinds() == ["alloc_stacked", "dispatch", "free_stacked"] + shape, worker_ids, _host = rt.payloads("alloc_stacked")[0] + assert (shape, worker_ids) == ((2, 4), None) # identity worker_ids def test_run_l3_resident_pure_out_stacked_skips_zero_upload(self, monkeypatch): """A write-only stacked resident uses empty per-rank allocations.""" - import golden.runner as R - - calls = { - "alloc": [], - "stacked": [], - "uploaded": 0, - "freed": 0, - "dispatched": 0, - } + built = [] class _FakeStackedDeviceTensor: def __init__(self, shards, full_shape, worker_ids): - calls["stacked"].append( - (len(shards), tuple(full_shape), tuple(worker_ids)) - ) - - class _FakeRT: - def __enter__(self): - return self - - def __exit__(self, *_a): - return False - - def alloc_stacked_tensor(self, _host, worker_ids=None): - calls["uploaded"] += 1 - raise AssertionError("pure Out must not upload its host placeholder") - - def alloc_tensor(self, shape, dtype, *, init=None, worker_id=0): - calls["alloc"].append((tuple(shape), dtype, init, worker_id)) - return types.SimpleNamespace(worker_id=worker_id) - - def free_stacked_tensor(self, _h): - calls["freed"] += 1 - - def free_tensor(self, _h, *, worker_id=0): - raise AssertionError(f"successful stacked handle must be freed as a stack: {worker_id}") - - def __call__(self, *_args, config=None): - calls["dispatched"] += 1 - - class _FakeDCP: - def prepare(self): - return _FakeRT() - - fake_mod = types.ModuleType("pypto.ir.distributed_compiled_program") - fake_mod.DistributedCompiledProgram = _FakeDCP - fake_runtime = types.ModuleType("pypto.runtime") - fake_runtime.StackedDeviceTensor = _FakeStackedDeviceTensor - - specs = [TensorSpec("y", [2, 4], torch.float32, is_output=True, resident="stacked")] - tensors = {"y": torch.zeros(2, 4)} - monkeypatch.setattr(R, "_l3_ordered_names", lambda _c: ["y"]) - monkeypatch.setattr(R, "_l3_pure_out_names", lambda _c: {"y"}) - monkeypatch.setattr(R, "_l3_run_config", lambda _cfg: "RUNCFG") - - with patch.dict( - sys.modules, - { - "pypto.ir.distributed_compiled_program": fake_mod, - "pypto.runtime": fake_runtime, - }, - ): - R._run_l3_resident( - compiled=_FakeDCP(), - tensor_specs=specs, - tensors=tensors, - scalar_specs_eff={}, - runtime_cfg={"platform": "a2a3"}, - golden_outputs=None, - rtol=1e-5, - atol=1e-5, - compare_fn={}, + built.append((len(shards), tuple(full_shape), tuple(worker_ids))) + + rt = _ResidentRT() + _stub_l3_helpers(monkeypatch, pure_out={"y"}) + dcp = _resident_dcp(rt) + runtime = types.ModuleType("pypto.runtime") + runtime.StackedDeviceTensor = _FakeStackedDeviceTensor + + with patch.dict(sys.modules, _resident_modules(dcp, runtime=runtime)): + _run_l3_resident( + compiled=dcp(), + specs=[TensorSpec("y", [2, 4], torch.float32, resident="stacked")], + tensors={"y": torch.zeros(2, 4)}, scalar_specs_eff={}, + **_RESIDENT_KWARGS, ) - assert calls["uploaded"] == 0 - assert calls["dispatched"] == 1 - assert calls["freed"] == 1 - assert calls["stacked"] == [(2, (2, 4), (0, 1))] - assert calls["alloc"] == [ + assert rt.payloads("alloc_stacked") == [] # no host placeholder upload + assert rt.kinds() == ["alloc", "alloc", "dispatch", "free_stacked"] + assert rt.payloads("alloc") == [ ((4,), torch.float32, None, 0), ((4,), torch.float32, None, 1), ] + assert built == [(2, (2, 4), (0, 1))] def test_run_l3_resident_output_reads_back(self, monkeypatch): - """A resident+is_output spec (state buffer) is read back via copy_stacked_from - before validation, so _validate sees the device's final state, not the stale host.""" + """A resident output is read back before validation, so _validate sees the + device's final state rather than the stale host buffer.""" import golden.runner as R - calls = {"readback": 0, "validated_value": None} - - class _FakeRT: - def __enter__(self): - return self - - def __exit__(self, *_a): - return False - - def alloc_stacked_tensor(self, host, worker_ids=None): - return types.SimpleNamespace(full_shape=tuple(host.shape)) - - def free_stacked_tensor(self, _h): - pass - - def __call__(self, *_args, config=None): - pass - - def copy_stacked_from(self, _handle, host): - calls["readback"] += 1 - host.fill_(7.0) # simulate the device's final in-place-updated state - - class _FakeDCP: - def prepare(self): - return _FakeRT() - - fake_mod = types.ModuleType("pypto.ir.distributed_compiled_program") - fake_mod.DistributedCompiledProgram = _FakeDCP + validated = {} + rt = _ResidentRT(readback=torch.full((2, 4), 7.0)) + _stub_l3_helpers(monkeypatch) + monkeypatch.setattr( + R, "_validate", + lambda _ts, tensors, *_a, **_k: validated.update(kv=tensors["kv"].clone()), + ) + dcp = _resident_dcp(rt) + specs = _stamped( + [TensorSpec("kv", [2, 4], torch.float32, init_value=torch.zeros, + resident="stacked")], + {"kv": "inout"}, + ) - specs = [ - TensorSpec( - "kv", [2, 4], torch.float32, init_value=torch.zeros, - is_output=True, resident="stacked", + with patch.dict(sys.modules, _resident_modules(dcp)): + _run_l3_resident( + compiled=dcp(), specs=specs, tensors={"kv": torch.zeros(2, 4)}, + scalar_specs_eff={}, + runtime_cfg={"platform": "a2a3"}, + golden_outputs={"kv": torch.full((2, 4), 7.0)}, + rtol=1e-5, atol=1e-5, compare_fn={}, ) - ] - tensors = {"kv": torch.zeros(2, 4)} - golden = {"kv": torch.full((2, 4), 7.0)} - monkeypatch.setattr(R, "_l3_ordered_names", lambda _c: ["kv"]) - monkeypatch.setattr(R, "_l3_pure_out_names", lambda _c: set()) - monkeypatch.setattr(R, "_l3_run_config", lambda _cfg: "RUNCFG") + assert len(rt.payloads("readback")) == 1 + assert torch.equal(validated["kv"], torch.full((2, 4), 7.0)) - def _fake_validate( - tensor_specs, - tensors, - golden_outputs, - rtol, - atol, - compare_fn, - scalar_specs_eff=None, - ): - calls["validated_value"] = tensors["kv"].clone() - - monkeypatch.setattr(R, "_validate", _fake_validate) + def test_resident_on_non_l3_fails_cleanly(self, build_dir): + """A resident spec against a non-L3 compiled program fails via RunResult.""" - with patch.dict(sys.modules, {"pypto.ir.distributed_compiled_program": fake_mod}): - R._run_l3_resident( - compiled=_FakeDCP(), - tensor_specs=specs, - tensors=tensors, - scalar_specs_eff={}, - runtime_cfg={"platform": "a2a3"}, - golden_outputs=golden, - rtol=1e-5, - atol=1e-5, - compare_fn={}, - ) + with ( + patch.dict(sys.modules, _resident_modules(_resident_dcp(None))), + patch("pypto.ir.compile", return_value=_FakeCompiled(build_dir)), + ): + r = run(fn=object(), specs=self._resident_specs()) - assert calls["readback"] == 1 - # _validate must have seen the read-back device state (7.0), not the stale 0.0. - assert calls["validated_value"] is not None - assert torch.equal(calls["validated_value"], torch.full((2, 4), 7.0)) + assert not r.passed + assert "only supported for L3" in (r.error or "") def test_run_l3_resident_rejects_non_l3(self): """The helper itself raises ValueError for a non-L3 compiled object.""" - with patch.dict( - sys.modules, - {"pypto.ir.distributed_compiled_program": self._fake_dcp_module()}, + with ( + patch.dict(sys.modules, _resident_modules(_resident_dcp(None))), + pytest.raises(ValueError, match="only supported for L3"), ): - with pytest.raises(ValueError, match="only supported for L3"): - _run_l3_resident( - compiled=object(), - tensor_specs=[TensorSpec("w", [4], torch.float32, resident=0)], - tensors={"w": torch.ones(4)}, - scalar_specs_eff={}, - runtime_cfg={"platform": "a2a3"}, - golden_outputs=None, - rtol=1e-5, - atol=1e-5, - compare_fn={}, - ) + _run_l3_resident( + compiled=object(), + specs=[TensorSpec("w", [4], torch.float32, resident=0)], + tensors={"w": torch.ones(4)}, scalar_specs_eff={}, + **_RESIDENT_KWARGS, + ) class _FakeInv: diff --git a/tests/golden/test_spec.py b/tests/golden/test_spec.py index f020b17e8..9bdd4d14e 100644 --- a/tests/golden/test_spec.py +++ b/tests/golden/test_spec.py @@ -65,12 +65,31 @@ def test_unsupported_init_value_raises(self): with pytest.raises(TypeError, match="Unsupported init_value type"): spec.create_tensor() - def test_is_output_flag(self): - """is_output flag is stored correctly and defaults to False.""" - spec_in = TensorSpec("a", [4], torch.float32) - spec_out = TensorSpec("b", [4], torch.float32, is_output=True) - assert spec_in.is_output is False - assert spec_out.is_output is True + def test_direction_is_not_an_init_argument(self): + """The kernel signature owns direction; a spec cannot declare one.""" + with pytest.raises(TypeError): + TensorSpec("a", [4], torch.float32, direction="out") + with pytest.raises(TypeError): + TensorSpec("a", [4], torch.float32, is_output=True) + + def test_direction_derives_is_output_and_is_input(self): + """Both flags read off the stamped direction, covering all three kinds.""" + cases = {"in": (False, True), "out": (True, False), "inout": (True, True)} + for direction, (is_output, is_input) in cases.items(): + spec = TensorSpec("a", [4], torch.float32) + spec.direction = direction + assert spec.is_output is is_output + assert spec.is_input is is_input + + def test_unstamped_direction_raises(self): + """Reading either flag before the harness stamps the artifact's direction + is a harness bug, not a silent False.""" + spec = TensorSpec("a", [4], torch.float32) + assert spec.direction is None + with pytest.raises(RuntimeError, match="direction not stamped"): + spec.is_output + with pytest.raises(RuntimeError, match="direction not stamped"): + spec.is_input def test_resident_defaults_off(self): """resident defaults to None (off) and is_resident is False.""" @@ -111,13 +130,15 @@ def test_resident_negative_worker_rejected(self): TensorSpec("w", [4], torch.float32, resident=-1) def test_resident_output_accepted(self): - """resident + is_output is a read-write resident state buffer (e.g. KV cache).""" - spec = TensorSpec("kv", [4], torch.float32, is_output=True, resident=0) + """A resident output is a read-write resident state buffer (e.g. KV cache).""" + spec = TensorSpec("kv", [4], torch.float32, resident=0) + spec.direction = "inout" assert spec.resident == 0 and spec.is_resident is True and spec.is_output is True def test_resident_stacked_output_accepted(self): - """resident="stacked" + is_output is a per-rank read-write state buffer.""" - spec = TensorSpec("kv", [2, 4], torch.float32, is_output=True, resident="stacked") + """resident="stacked" output is a per-rank read-write state buffer.""" + spec = TensorSpec("kv", [2, 4], torch.float32, resident="stacked") + spec.direction = "inout" assert spec.resident == "stacked" and spec.is_resident is True and spec.is_output is True def test_tensor_init_ignores_spec_shape(self):