core.tuning.selection

Choosing one Triton config without measuring anything.

Every rule here is a pure function of the candidate list and the call signature, so each rank computes the same answer. That property, not which config wins, is what makes a run reproducible: the cheapest config is as deterministic as the fastest one, it is just slower.

Module Contents

Functions

reset_warnings	Forget which once-per-policy notices were issued, so a new policy reports again.
_arch_for_device
arch_tag	sm<major><minor> for the current device, or unknown off-GPU.
kernel_name	Underlying kernel function name, unwrapping JIT and decorator layers.
kernel_module	Module the kernel is defined in, used to scope which kernels are pinned.
tuning_key	Shape/dtype signature for one autotuner invocation.
config_signature	Stable text form of one Triton config, for logs and cross-rank compare.
config_data	Serializable launch options; hooks remain on the live Config object.
estimate_config_cost	Estimate shared-memory cost of a config. Lower is cheaper.
filter_configs_by_block_sizes	Match explicit kernel block-size arguments against live candidates.
cheapest	The deterministic fallback: cheapest config by static estimate.
deterministic_choice	Pick one config, preferring a tuned entry, never measuring.
chaos_choice	Pick a different config per rank on purpose, reproducibly.
autotune_configs	Reduce an in-tree kernel’s config list under deterministic mode.
warn_if_mamba_env_ignored	Point out, once, that MAMBA_DETERMINISTIC does not pin Megatron’s kernels.

Data

logger
_untuned_kernels_warned
_mamba_env_warned

API

core.tuning.selection.logger

‘getLogger(…)’

core.tuning.selection._untuned_kernels_warned: set

‘set(…)’

core.tuning.selection._mamba_env_warned

False

core.tuning.selection.reset_warnings() → None

Forget which once-per-policy notices were issued, so a new policy reports again.

core.tuning.selection._arch_for_device(index: int) → str

core.tuning.selection.arch_tag() → str

sm<major><minor> for the current device, or unknown off-GPU.

core.tuning.selection.kernel_name(autotuner) → str

Underlying kernel function name, unwrapping JIT and decorator layers.

core.tuning.selection.kernel_module(autotuner) → str

Module the kernel is defined in, used to scope which kernels are pinned.

core.tuning.selection.tuning_key(autotuner, args, kwargs) → str

Shape/dtype signature for one autotuner invocation.

Built from the autotuner’s own keys plus the dtypes of its tensor arguments, so it is at least as fine as what the tiling depends on. It does not have to match Triton’s key byte for byte: a drift produces a table miss, which falls back to a deterministic choice, never to a timed one.

core.tuning.selection.config_signature(config) → str

Stable text form of one Triton config, for logs and cross-rank compare.

core.tuning.selection.config_data(config) → dict

Serializable launch options; hooks remain on the live Config object.

core.tuning.selection.estimate_config_cost(cfg)

Estimate shared-memory cost of a config. Lower is cheaper.

Returns (block_cost, num_warps) so ties in block cost break deterministically on warp count.

core.tuning.selection.filter_configs_by_block_sizes(configs, block_sizes=())

Match explicit kernel block-size arguments against live candidates.

core.tuning.selection.cheapest(configs)

The deterministic fallback: cheapest config by static estimate.

core.tuning.selection.deterministic_choice(

autotuner,

candidates,

args,

kwargs,

*,

table=None,

on_miss='min_cost',

block_sizes=(),

)

Pick one config, preferring a tuned entry, never measuring.

Order: the tuned table entry for this kernel and shape, then an explicit block-size override, then the cheapest config.

core.tuning.selection.chaos_choice(autotuner, candidates, args, kwargs)

Pick a different config per rank on purpose, reproducibly.

A positive control: every other check is a negative one, and “the runs matched” cannot distinguish a working divergence detector from a blind one.

core.tuning.selection.autotune_configs(configs)

Reduce an in-tree kernel’s config list under deterministic mode.

Used by Megatron’s own Triton kernels at decoration time, where there is no autotuner object to intercept. Cached autotuning (TRITON_CACHE_AUTOTUNING=1) is not sufficient here: the first benchmark is still timed, so the cached winner varies per process, per GPU and per cache directory.

core.tuning.selection.warn_if_mamba_env_ignored(policy) → None

Point out, once, that MAMBA_DETERMINISTIC does not pin Megatron’s kernels.

The variable still controls the external mamba_ssm package, but Megatron’s own Triton kernels follow deterministic_mode and PyTorch’s deterministic flag, so setting only the variable leaves them on timed autotuning.