Skip to content

Use FP16 operands for UMA radial FC2 - #2150

Draft
mlazos wants to merge 1 commit into
mainfrom
uma-fp16-radial-fc2
Draft

Use FP16 operands for UMA radial FC2#2150
mlazos wants to merge 1 commit into
mainfrom
uma-fp16-radial-fc2

Conversation

@mlazos

@mlazos mlazos commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary

This adds a default-off fp16_radial_fc2_blocks path for the second unified
UMA radial projection. Forward and frozen first-order backward use FP16 GEMM
operands with FP32 outputs, while layer normalization, activation, and
surrounding values remain FP32.

The implementation keeps nonpersistent FP16 weight caches stable across CUDA
graph replay, refreshes them after state or device changes, preserves state
dictionaries, and rejects non-FP32 base precision and Hessian inference.

Performance

With torch.compile(dynamic=True), the option improved the full endpoint from
17.01415 to 16.82399 ms with fresh caches and from 17.01992 to 16.82917 ms with
warm caches. The paired mean improvement is 0.19046 ms (1.12%).

A separate static-shape comparison improved 17.71949 to 17.55717 ms, a
0.16232 ms (0.92%) reduction. Maximum observed energy, force, and stress
differences were 0.006781, 0.000299, and 0.0000167.

Configuration: H100, TF32, UMA-S-1p2, 1,000 atoms, energy/forces/stress,
umas_fast_gpu, merge_mole=True, external_graph_gen=False, internal graph
v3, internal_graph_skin=0, compile=True, compile_dynamic_shapes=True, and
full CUDA graph replay. Steady-state execution had zero recompiles and zero
dynamic CUDA graph fallbacks.

Testing

PYTHONPATH=src /home/mlazos/.conda/envs/pytorch-3.12/bin/python -m pytest -q -c packages/fairchem-core/pyproject.toml tests/core/models/uma/nn/test_unified_radial.py tests/core/models/uma/uma_fast/test_execution_backends.py tests/core/units/mlip_unit/test_inference_settings.py -m 'not gpu'
pre-commit run --files docs/core/common_tasks/ase_calculator.md src/fairchem/core/models/uma/nn/execution_backends.py src/fairchem/core/models/uma/nn/unified_radial.py src/fairchem/core/units/mlip_unit/api/inference.py tests/core/models/uma/nn/test_unified_radial.py tests/core/models/uma/uma_fast/test_execution_backends.py tests/core/units/mlip_unit/test_inference_settings.py

The focused suite passed 21 tests, and all pre-commit hooks passed. The added
GPU coverage exercises eager and compiled VJPs plus CUDA graph replay.

Add a default-off fast-inference path for the second unified radial
projection. Forward and frozen first-order backward use FP16 GEMM operands
with FP32 outputs while layer normalization, activation, and surrounding
values remain FP32. Nonpersistent weight caches keep stable storage across
CUDA graph replay and refresh after state or device changes. Reject non-FP32
base precision and Hessian inference.

On H100 with UMA-S-1p2 at 1,000 atoms, the isolated option reduced dynamic
full-CUDA-graph latency from 17.017035 ms to 16.826578 ms, a 0.190457 ms
(1.12%) improvement averaged across independent fresh- and warm-cache
processes. The configuration used energy, forces, and stress; TF32;
merge_mole=True; external_graph_gen=False; internal graph v3; skin 0;
compile_dynamic_shapes=True; and full CUDA graph replay. A separate static
measurement improved 17.719488 ms to 17.557168 ms. Maximum observed energy,
force, and stress differences were 0.006781, 0.000299, and 0.0000167.

Test Plan:

```
PYTHONPATH=src /home/mlazos/.conda/envs/pytorch-3.12/bin/python -m pytest -q -c packages/fairchem-core/pyproject.toml tests/core/models/uma/nn/test_unified_radial.py tests/core/models/uma/uma_fast/test_execution_backends.py tests/core/units/mlip_unit/test_inference_settings.py -m 'not gpu'
pre-commit run --files docs/core/common_tasks/ase_calculator.md src/fairchem/core/models/uma/nn/execution_backends.py src/fairchem/core/models/uma/nn/unified_radial.py src/fairchem/core/units/mlip_unit/api/inference.py tests/core/models/uma/nn/test_unified_radial.py tests/core/models/uma/uma_fast/test_execution_backends.py tests/core/units/mlip_unit/test_inference_settings.py
```
@meta-cla meta-cla Bot added the cla signed label Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant