Skip to content

feat: default --xla_gpu_autotune_level=0 in jax_wrapper (env-respecting) #131

Description

@Jammy2211

Overview

Rollout of autolens_profiling#74 verdict 2: the pathological ~7m30 cold GPU compile is XLA autotuning. Evidence (A100): autotune genuinely off gives a 17× cold-probe reduction (498 s → 29 s), −40 % cold full fit (2081 s → 1253 s), bit-identical fixed-input logL, and steady-eval parity across the measured matrix (worst ~4 % on one cell; the 4800-eval fit was faster end-to-end). Stacked with the default cache (#127/#128), worst-case first-fit UX drops ~70 min → ~30 s.

Plan

  • autoconf/jax_wrapper.py: append --xla_gpu_autotune_level=0 to the assembled XLA_FLAGS unless the user's flags already contain an --xla_gpu_autotune_level setting (any value — respect it). Same append/env-respecting pattern as feat: JAX persistent compilation cache by default + XLA_FLAGS append fix (#127) #128. INFO log line.
  • Extend test_autoconf/test_jax_wrapper.py: flag added when absent; user-preset level respected; composes with the constant_folding append.
  • Ship validation: fresh-process check that the flag reaches XLA_FLAGS and a preset level survives.
  • Downstream PyAutoFit suite against the branch (base-layer change).

Prompt: PyAutoMind draft/feature/autoconf/default_xla_gpu_autotuning_off_in_the.md · branch feature/autotune-off-default

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions