Feature request

We are planning to propose native contextvars.ContextVar.get() and .set() support to TorchDynamo.
Once this lands in TorchDynamo, CompileableContextVar can be replaced with a plain ContextVar that works directly inside torch.compile - no wrapper class, no compile-time special-casing, and thread-safe by default.

Motivation

We are planning to add native ContextVar.get() and .set() support to TorchDynamo making it tracable by default and therefore eliminating the need for maintaining a seperate CompileableContextVar.

How is it useful for Transformers?

• CompileableContextVar can be replaced with a plain ContextVar hence no wrapper class, no compile-time special-casing, thread-safe by default
• @torch.compiler.allow_in_graph on can_record_outputs in modeling_utils.py becomes unnecessary
• Eager hook pre-installation in backbone_utils.py post_init() becomes unnecessary
• CompileConfig.fullgraph = False default could be revisited for models that currently tolerate graph breaks only because of this limitation

The oss-model-graph-break-corpus (hundreds of HF models tested systematically) shows this as one of the top graph-break clusters:, affecting over a hundred model classes.

CompileableContextVar in output_capturing.py exists because TorchDynamo cannot trace ContextVar.get(). It switches between a real ContextVar in eager mode and a plain global during compilation, and is read/written on every @capture_outputs-decorated forward across hundreds of model directories. This will all be simply eliminated.

Beyond Transformers, this also unblocks torch.compile interoperability with observability libraries, structured logging, async-safe inference serving like SGL and vLLM, and other projects that rely on ContextVar for per-request or per-task state.

Your contribution

I am planning to implement the native ContextVar support on the PyTorch/Dynamo side and take care of it end-to-end. The approach follows established Dynamo patterns:

• Phase 1 (in progress): Read-only .get() support - cv.get() inside torch.compile resolves at trace time with a guard that re-checks at cache-validation time. .set() and .reset() graph-break with a clear message.
• Phase 2: .set() and .reset() tracked as side effects - mutations replayed in generated code, Token return values properly modeled. This is the phase that enables replacing CompileableContextVar.
• Phase 3: Context.run() scope boundaries inside compiled code.
• Phase 4: copy_context() and Context mapping protocol support.

I have a Detailed analysis, research and design Google doc here: ContextVars Support in TorchDynamo
You can gp through exact API changes proposed and the implementation plan here.

Once the Dynamo-side change lands and the approach is finalized here, I will open a follow-up PR to replace CompileableContextVar with a plain ContextVar in Transformers.