kernel: avoid partial Triton lazy initialization

[nascheme]

blockwise_quantization.py and bsr_triton_ops.py both used a module-level _triton_initialized = False flag set to True before the kernel handles were assigned. Under free-threaded Python, one thread can set the flag and yield during the slow import triton, while another thread observes the flag set but the kernel globals still None.

Replace the bool-flag pattern with a lock-guarded builder that publishes an immutable dataclass of the built handles only once construction completes, so callers either see nothing or the fully-built set.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4553

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures

As of commit 1d28ff8 with merge base d1e0e2d ():

NEW FAILURES - The following jobs have failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run 1xH100 Tests / test (H100, mt-l-x86iamx-22-225-h100, --pre torch torchvision torchaudio mslk --index-url https:/... / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t 9d3effcdb49bc61fa1957f04a74bdbe07ecf34a09387ac5182b55ddfa6e6a90e /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vkuzo]

@claude review

[claude]

Claude finished @vkuzo's task in 1m 30s —— View job

---

Code Review: avoid partial Triton lazy initialization

• Read the diff and both changed files
• Analyze correctness of the new lazy-init pattern
• Check for behavior changes and edge cases

---

Nice fix. The core change is correct: replacing the _triton_initialized = True (set before the kernels are assigned) with a lock-guarded builder that publishes an immutable, fully-constructed dataclass in a single atomic reference assignment. Under free-threaded Python, a reader now sees either None or the complete _BlockwiseFp8Impls/_BsrTritonImpls — never a half-built state. The double-checked locking (fast unlocked read, re-check under the lock) is the right idiom here and is safe because the published object is frozen=True and constructed entirely before assignment.

A few observations:

1. Behavior change in bsr_dense_addmm (worth confirming intent) — bsr_triton_ops.py:276

Previously bsr_dense_addmm called _lazy_init_triton() with no availability check. When Triton was absent, _bsr_strided_addmm_kernel stayed None and the function ran on until None[grid](...) raised a cryptic TypeError. The new code adds:

impls = _get_triton_impls()
if not impls.available:
    raise AssertionError("unsupported without triton")

This is a strict improvement (fail fast, clear message) and matches the pattern in blockwise_quantization.py, but it is a new early raise on a path that previously didn't have one. Just flagging it as intentional — I don't think anything relied on the old late-failure behavior.

2. AssertionError for a runtime capability check — minor

Both files raise AssertionError("unsupported without triton"). Assertions are stripped under python -O, which would let execution fall through to impls.gemm(...) where gemm is None → TypeError. RuntimeError (or NotImplementedError) would be more robust. That said, the blockwise_quantization.py calls already used AssertionError before this PR, so this only introduces the pattern newly in bsr_triton_ops.py. Non-blocking; feel free to leave for consistency.

3. The global _bsr_strided_addmm_kernel declaration was correctly removed

bsr_dense_addmm no longer needs the global statement it had before — good cleanup, and I confirmed there are no remaining references to the old module-level names (_triton_initialized, _bsr_strided_addmm_kernel, _triton_available, _blockwise_fp8_gemm_impl, _lazy_init_triton) anywhere in torchao/ or test/.

4. Optional: dedup the identical _get_triton_impls boilerplate

The _get_triton_impls() double-checked-locking helper is now byte-for-byte identical between the two modules (only the dataclass type differs). Not worth blocking on, but a small shared helper (e.g. a lazy_singleton(builder, lock) in torchao/kernel/) would prevent the two copies from drifting. Minor.

Overall this looks correct and is a solid concurrency fix. Nothing blocking from my side — points 1 and 2 are the only ones I'd ask you to confirm are intentional.
• ft-audit-kernel-lazy-init