File size: 1,357 Bytes
7218492 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | """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())
@contextmanager
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
|