perf: hardware-adaptive GPU optimization (DGX Spark UMA + A100/H100 HBM)

[thc1006]

Problem

All GPU code runs as vanilla PyTorch with no hardware-specific optimization. Leaves significant performance on the table across all target platforms.

Current State (measured on DGX Spark GB10)

Feature	Available	Utilized
TF32 matmul	Yes (Blackwell)	No (allow_tf32 = False)
BF16	Yes	No (all FP32/FP64)
128GB UMA	Yes	No (unnecessary .cpu().numpy() roundtrips)
Tensor Cores (5th gen)	Yes	No
torch.compile	Yes	No

Target Platforms

Platform	Memory	Key Optimization
DGX Spark GB10	128GB UMA (shared)	Minimize .cpu()/.to(device) — zero-copy
A100 (40/80GB)	Dedicated HBM2e	Maximize batch size, fuse transfers
H100 (80GB)	Dedicated HBM3	FP8 tensor cores, transformer engine

Proposed Changes (by effort/impact)

Quick wins (1-2 lines each)

• torch.backends.cuda.matmul.allow_tf32 = True
• torch.set_float32_matmul_precision('high')
• torch.compile(nqs) for NQS inference

Medium effort (~50 lines)

• torch.amp.autocast('cuda', dtype=torch.bfloat16) around NQS training loop
• Audit and remove unnecessary .cpu().numpy() → .to(device) roundtrips (especially in UMA)

Larger effort (separate PR)

• Hardware-adaptive runtime: detect UMA vs HBM and adjust data movement strategy
• CUDA kernel optimization for diagonal_elements_batch and PT2 scoring
• Numba/CUDA acceleration for get_connections inner loop (currently pure Python)

Related

• Issue [RFC] Adaptive sampling strategy for 40Q+ CAS molecules (Issue #21 Phase 4) #25 (adaptive sampling — also performance-sensitive)
• ADR-005 (PT2 selection — compute_pt2_scores is CPU-bound Python loop)

[thc1006]

Additional scope from PT2 selection (PR #31)

40Q benchmark (N₂ CAS(10,20)) revealed the dominant bottleneck is not GPU underutilization but the PT2 scoring Python loop:

• compute_pt2_scores: 2000 candidates × ~7875 connections/config = 15.75M dict lookups in pure Python (~15-30s/iteration)
• get_connections: called once per candidate, each returns ~7875 connected configs via Slater-Condon rules (also Python/Numba)

New items to add to scope:

1. Vectorize PT2 scoring — batch get_connections + vectorized hash lookup (replace Python dict with torch.searchsorted)
2. Cache connections — ConnectionCache already exists but not used in PT2 path
3. Numba-accelerate PT2 inner loop — the coupling calculation sum(H_ij * c_j) is embarrassingly parallel

These are CPU optimizations, not GPU. But they'd cut 40Q time from ~813s to ~200s (estimated).