mstar.engine.fp32_precision

Per-node float32 math precision, applied by the engine around a node’s forward.

apply_torch_config sets float32_matmul_precision('high') (TF32) for the whole process. A submodule that needs full fp32 for its own work declares it (NodeSubmodule.float32_matmul_precision / cudnn_allow_tf32) instead of flipping the process-wide flags itself, and the engine sets them around that node’s forward launch and its warmup and capture only.

TF32 is a kernel choice made on the host when a kernel is launched, compiled or captured, so the scope has to cover the launch, the first (compiling) call and the capture; a replayed graph keeps what it was captured with. Dynamo cannot trace a flag change inside a compiled region, which is why this wraps the compiled callable rather than living in the model’s forward.

The flags are process-wide, not per thread. While a node’s forward is being launched on the GPU thread, other threads of the worker (the main thread’s postprocess of the previous step, the plan thread’s pre-plan) see the node’s setting too. The engine only changes them on the thread that runs forwards (and the one that captures, before serving), so its save and restore never interleave with each other. The process default is set once at import; the executor threads’ initializers apply only the per-thread dynamo config (apply_dynamo_config), so a thread starting mid-step cannot reset a node’s setting. A model that flips the flags itself still can.

Functions

fp32_precision_scope(precision)	Apply precision for the duration of the block, then restore what was there before, also when the block raises.
node_fp32_precision(node_name, submodule)	The submodule's declared settings, checked when the engine loads it.

Classes

Fp32Precision(matmul, cudnn_allow_tf32)

A node's float32 settings; None leaves the process default in place.

class mstar.engine.fp32_precision.Fp32Precision(matmul, cudnn_allow_tf32)

Bases: NamedTuple

A node’s float32 settings; None leaves the process default in place.

Parameters:

• matmul (str | None)

• cudnn_allow_tf32 (bool | None)

cudnn_allow_tf32: bool | None

Alias for field number 1

matmul: str | None

Alias for field number 0

mstar.engine.fp32_precision.fp32_precision_scope(precision)

Apply precision for the duration of the block, then restore what was there before, also when the block raises.

Parameters:

precision (Fp32Precision | None)

mstar.engine.fp32_precision.node_fp32_precision(node_name, submodule)

The submodule’s declared settings, checked when the engine loads it. None when it declares neither.

Parameters:

node_name (str)

Return type:

Fp32Precision | None