core.ssm.ops.determinism

Module Contents

Functions

use_deterministic_mode	Use torch deterministic mode.
set_deterministic_mode	Set torch deterministic mode.
_estimate_config_cost	Estimate shared memory cost of a config. Lower is cheaper.
_filter_configs_by_block_sizes	Filter configs by TRITON_AUTOTUNE_BLOCK_* env vars.
autotune_configs	Select autotune configs for deterministic mode.
alloc_tile_workspace	Allocate buffer for deterministic per-program reductions.
finalize_tile_workspace	Finalize tile workspace.

Data

TRITON_HAS_CACHE_RESULTS
_autotune_warning_issued
_deterministic_override

API

core.ssm.ops.determinism.TRITON_HAS_CACHE_RESULTS

None

core.ssm.ops.determinism._autotune_warning_issued

False

core.ssm.ops.determinism._deterministic_override

None

core.ssm.ops.determinism.use_deterministic_mode()

Use torch deterministic mode.

core.ssm.ops.determinism.set_deterministic_mode(value)

Set torch deterministic mode.

core.ssm.ops.determinism._estimate_config_cost(cfg)

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

Returns a tuple (block_cost, num_warps) so that ties in block cost are broken deterministically by warp count (fewer warps = cheaper).

core.ssm.ops.determinism._filter_configs_by_block_sizes(configs)

Filter configs by TRITON_AUTOTUNE_BLOCK_* env vars.

Scans environment for any variable matching TRITON_AUTOTUNE_BLOCK_* (e.g. TRITON_AUTOTUNE_BLOCK_SIZE_M, TRITON_AUTOTUNE_BLOCK_SIZE_H, TRITON_AUTOTUNE_BLOCK_T, TRITON_AUTOTUNE_BLOCK_C, TRITON_AUTOTUNE_BLOCK_SIZE) and maps them to the corresponding kernel kwarg (BLOCK_SIZE_M, BLOCK_SIZE_H, BLOCK_T, BLOCK_C, BLOCK_SIZE).

core.ssm.ops.determinism.autotune_configs(configs)

Select autotune configs for deterministic mode.

Uses cached autotuning (TRITON_CACHE_AUTOTUNING=1) if Triton >= 3.4.0, otherwise auto-selects the cheapest config by block size * stages.

core.ssm.ops.determinism.alloc_tile_workspace(

base_shape,

tile_dim,

dtype,

device,

deterministic,

*,

zero_init=True,

)

Allocate buffer for deterministic per-program reductions.

core.ssm.ops.determinism.finalize_tile_workspace(tensor, deterministic)

Finalize tile workspace.