| """Pinned precision policy and backend identity for request-owned state.""" | |
| from contextlib import contextmanager | |
| import torch | |
| ARITHMETIC_VERSION='current_recipe_b200_alignment_dense_attention_output_attention_8192_v8' | |
| def execution_identity(device): | |
| device=torch.device(device) | |
| capability=tuple(torch.cuda.get_device_capability(device)) if device.type=='cuda' else None | |
| backend='native_b200_mma' if capability and capability==(10,0) else 'canonical_b200_mma' | |
| matmul=torch.backends.cuda.matmul | |
| return (ARITHMETIC_VERSION,str(torch.__version__),torch.version.cuda,str(device),capability,backend, | |
| bool(matmul.allow_tf32),bool(matmul.allow_bf16_reduced_precision_reduction), | |
| bool(matmul.allow_fp16_reduced_precision_reduction),torch.get_float32_matmul_precision()) | |
| def release_precision(): | |
| # The public facade holds its process-wide execution lock around this scope. | |
| # Both switches are restored even if evaluation raises. | |
| matmul=torch.backends.cuda.matmul | |
| tf32=matmul.allow_tf32;bf16=matmul.allow_bf16_reduced_precision_reduction | |
| try: | |
| matmul.allow_tf32=False | |
| matmul.allow_bf16_reduced_precision_reduction=False | |
| yield | |
| finally: | |
| matmul.allow_bf16_reduced_precision_reduction=bf16 | |
| matmul.allow_tf32=tf32 | |