{"resourceId":"alphafold2-tpu-inference-software-layer-20260928","versions":[{"version":"external-da95b2af97a65ae57e02967c548deea5a7c4f251e27ed968671b63c14c72688f","resource":{"id":"alphafold2-tpu-inference-software-layer-20260928","title":"AlphaFold2 on Cloud TPU v5e: 0.47 s/chip vs 13.1 s T4, but default path uses 1 of 8 chips; pmap 6.5–7.9×","organization":"Stanford Summer Session / LUISS Guido Carli / Al Akhawayn University — arXiv preprint","sector":"Higher education research computing / scientific AI workloads (protein structure inference)","geography":"United States (Stanford TPU access; Colab baselines); transferable methods with cloud-tenancy limits","publishedAt":"September 28, 2026 (arXiv; revised September 29, 2026 v2)","publicationDate":"2026-09-28","eventDate":null,"sourceName":"arXiv 2609.34818","sourceLabel":"Academic preprint (arXiv)","sourceUrl":"https://arxiv.org/abs/2609.34818","evidenceClass":"academic-research","outcomeClass":"mixed","topics":["infrastructure","developers-agents","operating-model"],"finding":"Pazienza and El Bani measure identical AlphaFold2 inference code on Colab CPU, NVIDIA T4 GPU, and an eight-chip Cloud TPU v5e slice. Steady-state single-chip TPU time is 0.47 s/call versus 13.1 s on T4 in the August campaign (~27.8×), but the default path places work on only one of eight billed chips. jax.vmap never exceeds single-query throughput; jax.pmap yields 6.5–7.9× eight-over-one-chip throughput on a matched length grid. Automatic sharding leaves per-chip footprint unchanged (consistent with replication). About three quarters of a traced first-call span is JAX tracing/compilation. September Colab reruns did not reproduce August CPU/GPU baselines (~1.6–2× drift), so hardware ratios are campaign-specific. Paper revised 29 Sep 2026 (v2).","sledRelevance":"Campus research-computing and AI-for-science programs buying accelerators need evidence that software mapping, compilation caching, and multi-device programming—not device choice alone—determine realized scientific throughput and cost.","evidence":"arXiv abs/html 2609.34818 opened via HTTPS (curl UA SLED-research-bot/1.0); steady-state timings, idle-chip finding, pmap/vmap results, compilation share, and cross-session variability taken from opened abstract and paper body (submitted 28 Sep 2026; revised 29 Sep 2026).","architectureImplications":"Treat JAX/XLA compilation caches, explicit pmap/device mapping, and sharding annotations as first-class architecture for scientific JAX workloads; default jit paths may leave multi-chip slices mostly idle while still billing the pod.","governanceImplications":"Publish allocation and cost models that distinguish single-chip useful work from multi-chip billed capacity; require workload-level benchmarks with stated software stack, not brochure FLOPS, before claiming research readiness.","securityPrivacyImplications":"Cloud Colab tenancy and unpinned dependencies complicate auditability of scientific timing studies; retain environment captures and avoid treating unreproducible cloud baselines as procurement constants.","caveats":"Preprint; untrained AlphaFold2 weights for timing; mostly single TPU runs; August baselines partially unreproducible; T4 vs v5e not a campus GPU-cluster comparison; AlphaFold3 release tested lacks TPU backend.","streamIds":["research"],"roles":{"sales":"Sales framing: Research-computing buyers evaluating AI accelerators for scientific workloads need measured end-to-end inference economics under the software path users actually run—not peak FLOPS. This AlphaFold2 study supports a bounded pitch that TPU chips can be much faster than a T4 in steady state (0.47 s vs 13.1 s in one campaign) while default eight-chip allocation and missing sharding leave most chips idle. Do not sell portable 27.8× speedups or guaranteed cost wins: Colab baselines drifted across sessions and pod billing can cancel chip advantage when seven chips sit unused.","engineering":"Engineering framing: Fit is a JAX scientific inference stack where device placement, compilation caching, and explicit parallelism (pmap vs vmap vs automatic sharding) dominate realized throughput. Integrate per-device memory checks, compilation caches across process restarts, and annotated sharding before sizing multi-chip allocations. Prerequisites include pinned JAX/Haiku versions, shape-stable batching strategy, and cost models that bill useful chips. Proof-of-value should measure steady-state and cold-start separately with retained traces.","delivery":"Delivery framing: Stand up a scientific JAX benchmark harness with environment capture, single- vs multi-device runs, and compilation-cache policy before migrating AlphaFold-class jobs to multi-chip cloud or campus slices. Ownership: research computing owns device pools and cost accounting; application owners own sharding annotations and model entry points. Acceptance requires documented chips-in-use, cold/steady ratios, and matched-grid throughput—not a single brochure timing. Risks include idle billed capacity, unreproducible cloud baselines, and overclaiming hardware constants from one campaign."},"roleDetails":{"sales":{"customerProblem":"Scientific AI inference buyers risk paying for multi-chip accelerators while software paths leave most chips idle.","stakeholders":["Research computing / HPC directors","VPR / computational science leads","PI structural-biology groups","Cloud/finops for research","Platform engineering"],"discoveryQuestions":["How many chips does the default job actually use (device memory nonzero)?","Are cold-start compilation costs amortized for the workload mix?","Is cost modeled per busy chip or per billed slice?","Can baselines be reproduced across sessions with pinned environments?"],"valueHypothesis":"Explicit multi-device mapping and compilation caching can convert accelerator CapEx into measured scientific inference throughput when default paths leave chips idle.","potentialEngagement":"Two-week AlphaFold-class benchmark on target accelerators with single-chip, pmap, and cost-per-1k-prediction reports.","unsupportedClaims":["Guaranteed 27.8× over every campus GPU","Automatic sharding always splits AlphaFold2","Stable cloud Colab baselines","AlphaFold3 TPU portability"]},"engineering":{"fit":"JAX/XLA scientific inference on multi-chip TPU/GPU with explicit device mapping and compilation cache.","architectureIntegration":"Wire jit + pmap entry points, per-device HBM telemetry, persistent compilation cache, and pinned dependency lockfiles into the research workflow.","prerequisites":["Pinned JAX/Haiku stack","Shape-stable input policy","Multi-device allocator","Environment/provenance capture","Cost telemetry per chip-hour"],"constraints":["Unannotated models may replicate","Cold-start compilation dominance","Cloud tenancy variability","Preprint/single-campaign ratios"],"security":"Capture and retain run environments; avoid leaking research credentials via notebook sharing; treat cloud host reassignment as an audit risk for timing claims.","proofOfValue":["Steady-state time/chip with chips-in-use","pmap speedup on matched grid","Cold/steady ratio with/without cache","Cost per 1k predictions at busy vs idle allocation"]},"delivery":{"work":"Implement benchmark harness, pin environments, run single/multi-device campaigns, publish cost and chips-in-use dashboards.","dependencies":["Accelerator allocation","JAX scientific app owners","Storage for traces/results","Finops/cost APIs"],"ownership":"Research computing owns pools and billing interpretation; app teams own model parallelism annotations; PIs own scientific validity of structures.","skillsAdoption":"Train users to check device memory across chips and prefer pmap/annotated sharding over assuming automatic split.","governanceCheckpoints":["Environment pin freeze","Benchmark protocol review","Cost-model sign-off","Reproducibility rerun gate"],"acceptanceCriteria":["Documented chips-in-use ≥ N for billed jobs","Published cold and steady timings","Matched-grid multi-device report"],"risks":["Idle billed capacity","Unreproducible baselines","Overfit to one sequence length","Missing sharding annotations"]}},"retrievedAt":"2026-09-29T22:01:24-05:00","enrichedAt":"2026-09-29T22:02:46-05:00","enrichmentBasis":"retrieved source","accessibilityWorkforceImplications":"Multi-device JAX competence becomes a workforce prerequisite; without training, researchers may under-use expensive slices while still exhausting allocations.","procurementImplications":"Require workload benchmarks stating software path, chips-in-use, and cold-start behavior in accelerator RFPs; price busy useful capacity, not peak chip counts alone.","operatingModelImplications":"Operating model must include compilation-cache ops, multi-device programming support, and session-reproducible benchmarking—not only queue access to accelerators.","sourceVerification":{"openedUrl":"https://arxiv.org/abs/2609.34818","referenceExcerpt":"0.47 s per call in steady state on a single chip against 13.1 s on the T4","promptVersion":"sled-research-v4.0","model":null,"basis":"agent-reported inspection"}}}]}