`core.tuning.table`#

Tuned Triton configs, stored as data rather than as source.

A tuned table lets a pinned run use the fastest config instead of the cheapest one while staying a pure function of its inputs. Tables live as JSON so adding an architecture is a file drop, not a source edit and a rebuild:

torchrun ... pretrain_gpt.py ... --triton-autotune-record-path /tmp/rec
python -m megatron.core.tuning merge /tmp/rec.rank*.json -o ~/.mcore/tuning/sm103.json
torchrun ... --triton-autotune-mode pinned --triton-autotune-table-path ~/.mcore/tuning

Each file records one architecture plus the provenance needed to notice when it has gone stale::

{"arch": "sm103",
 "triton": "3.6.0",
 "packages": {"mamba_ssm": "2.3.1"},
 "source": "nemotron_3_ultra/96gpu",
 "kernels": {"<kernel>": {"<shape key>": {"kwargs": {...},
                                          "num_warps": 4, "num_stages": 3}}}}

A miss is never an error by default: the caller falls back to a deterministic choice. It never falls back to a timed one.

Module Contents#

Classes#

TunedTable

Tuned configs for one architecture, with membership validation.

Functions#

`_triton_version`
`_package_versions`
`_search_dirs`
`load`	Load the first table for `arch` found on the search path.
`_read_recording`	Load one per-rank recording, rejecting files of any other shape by name.
`count_votes`	Count, per (arch, kernel, key), how many ranks recorded each config.
`merge_votes`	Majority winner per (arch, kernel, key); ties break on the serialized config.
`disagreements`	The (arch, kernel, key) entries for which ranks recorded more than one config.
`merge_records`	Merge per-rank recordings by majority vote, keyed by architecture.
`disagreement_report`	Count, per (arch, kernel, key), how many distinct winners ranks chose.
`write`	Write one architecture’s table, with the provenance to date it.

Data#

`logger`
`_PACKAGED`

API#

core.tuning.table.logger#: ‘getLogger(…)’

core.tuning.table._PACKAGED#: None

core.tuning.table._triton_version() → str#

core.tuning.table._package_versions() → dict#

class core.tuning.table.TunedTable(arch: str, kernels: dict, provenance: dict | None = None)#

Tuned configs for one architecture, with membership validation.

Initialization

__bool__() → bool#

lookup(kernel: str, key: str, candidates)#

Return the tuned config for this kernel and shape, or None.

The stored entry is matched back against the kernel’s own candidate list rather than rebuilt, so its pre_hook survives, and a stale entry that no longer names a real candidate degrades to a miss instead of an invalid launch.

core.tuning.table._search_dirs(extra) → list#

core.tuning.table.load(arch: str, table_path=()) → core.tuning.table.TunedTable#

Load the first table for arch found on the search path.

User directories are searched before the packaged defaults, so a locally recorded table wins without touching the tree. A table recorded against a different Triton warns rather than being dropped: entries are validated against the live candidate list anyway, so the worst case is a miss.

core.tuning.table._read_recording(path) → dict#: Load one per-rank recording, rejecting files of any other shape by name.

core.tuning.table.count_votes(paths) → dict#

Count, per (arch, kernel, key), how many ranks recorded each config.

Raises OSError for unreadable files and ValueError naming the file for anything that is not a per-rank recording.

core.tuning.table.merge_votes(votes: dict) → dict#: Majority winner per (arch, kernel, key); ties break on the serialized config.

core.tuning.table.disagreements(votes: dict) → dict#: The (arch, kernel, key) entries for which ranks recorded more than one config.

core.tuning.table.merge_records(paths) → dict#

Merge per-rank recordings by majority vote, keyed by architecture.

Ranks routinely disagree about the winner; that disagreement is the very variance a table removes, so the merge counts votes rather than letting the last file win. Ties break on the serialized config, so a given set of recordings always produces the same table.

core.tuning.table.disagreement_report(paths) → dict#

Count, per (arch, kernel, key), how many distinct winners ranks chose.

Anything above one is timing variance the table is about to remove, and is worth seeing before trusting a recording.

core.tuning.table.write(arch: str, kernels: dict, path, source: str = '') → None#: Write one architecture’s table, with the provenance to date it.

core.tuning.table#