torch.compile compatibility: progress and remaining work

[vmoens]

Summary

Branch torch-compile-compat makes mujoco-torch compatible with torch.compile. The compiled path produces numerically identical results to eager, and steady-state compiled performance is at parity with eager (~1.0-1.1x on CPU).

Fixes applied

Phase 1: Core compatibility (commit de115ed)

• Solver: Converted _Context/_LSPoint/_LSContext to NamedTuples for while_loop compatibility; fixed tensor aliasing and dtype mismatches.
• Scan: Converted numpy index arrays to torch.LongTensor; cached grouping indices.
• Collision: Pre-resolved collision functions and GeomType enums to avoid numpy key lookups.
• Removed named_scope decorator — caused ~12 unnecessary recompilations per step.
• Conditional _check_output — skips validation during compilation.
• Bulk replace() — uses clone._tensordict._tensordict.update(kwargs) to avoid per-key _set_str recompilations.

Phase 2: Graph break elimination (commit 935d7dd)

• Replaced Model.__getattribute__ override with a targeted names property descriptor. The __getattribute__ override caused graph breaks on EVERY attribute access on Model instances, making torch.compile skip entire frames (_advance, _euler).
• Pre-built scan caches in device_put — scan.flat / scan.body_tree contain Python for loops over groups with varying tensor shapes that cause O(groups) recompilations. Grouping caches are now pre-built in device_put so these functions don't need @torch.compiler.disable at runtime. The scan callbacks are compiled via torch.vmap internally.
• @torch.compiler.disable on _collide_hfield_geoms — heightfield collision contains non-traceable operations.

Current state

Metric	Before	After
Recompilations (per 3 warm-up steps)	130+	~20 (bounded)
Graph breaks (unique skipped frames)	42+	10
Compiled vs eager speedup (CPU, 1000 steps)	0.5x	~1.0-1.1x
Warm-up compile time	~20s	~35s
Numerical accuracy	identical	identical

Benchmark (ant model, 1000 steps, CPU)

MuJoCo (C)                         8.6 ms  (116,038 steps/s)
mujoco-torch eager              24353.6 ms  (41 steps/s)
mujoco-torch compiled           23655.4 ms  (42 steps/s)
MJX (JAX jit)                    2201.0 ms  (454 steps/s)

Remaining work

Compilability

• Reduce scan recompilations (high impact) — scan.flat and scan.body_tree use Python for loops over groups with varying tensor shapes, causing O(groups) recompilations that fragment the compiled graph. Options: (1) pad tensors to uniform shapes for a single batched vmap call, (2) use torch._higher_order_ops.scan for the tree traversal, (3) compile individual scan callbacks separately.
• Make TensorClass constructors traceable (tensordict fix) — TensorClass.__init__ raises RuntimeError during tracing, causing graph breaks in collision and constraint code. Requires upstream fix in tensordict. (pytorch/tensordict@c07686f)

CUDA graph support (mode="reduce-overhead")

• Fix tensor aliasing in _take() — CUDA graph capture fails because _take returns views that alias the input tensor's storage. CUDA graph replay requires non-aliased outputs. Options: (1) clone the sliced tensor, (2) restructure scan.flat/scan.body_tree to avoid aliased returns, (3) wait for PyTorch upstream support.

Performance: smooth.py cleanup

• Remove redundant clone in factor_m — qld = qld / qld[...] already creates a new tensor; the subsequent .clone() was unnecessary. Also replaced clone-per-iteration with out-of-place scatter(). (PR [smooth] Performance cleanup + reduce-overhead benchmarks #29)
• Eliminate per-iteration clones in solve_m — replaced x = x.clone(); x[idx] = val with x = x.scatter(0, idx, val), matching the pattern already used in solver.py. (PR [smooth] Performance cleanup + reduce-overhead benchmarks #29)
• Precompute tendon structural indices — moved np.nonzero/np.isin from per-step tendon() to one-time device_put(), stored as UnbatchedTensor fields on the Model. (PR [smooth] Performance cleanup + reduce-overhead benchmarks #29)
• Precompute tendon tensors — moved torch.as_tensor()/torch.arange() from per-step tendon() to one-time device_put(), stored as UnbatchedTensor fields. tendon() now does zero numpy computation and zero tensor creation at runtime. (PR [smooth] Performance cleanup + reduce-overhead benchmarks #29)

Compile warmup

• AOT compilation / cache shipping — see Ahead-of-time compilation: ship pre-compiled models to eliminate warmup #27

Testing & benchmarking

• GPU benchmarks — pytest-benchmark suite in benchmarks/ with 5 models x 5 batch sizes, comparing C, loop, vmap, compile, compile+H4, compile+H4+H8, and MJX.

[vmoens]

Compile-time analysis and optimization notes

Current compile-time bottleneck

The main source of slow compile times is torch._inductor.config.coordinate_descent_tuning (H4 optimization). When enabled, the Inductor backend performs exhaustive tile-size searches for every new triton kernel. For a single mujoco_torch.step() call on the ant model, this can mean 20+ minutes of initial autotuning. The cost compounds when:

1. torch._dynamo.reset() is called between tests — forces full recompilation and re-autotuning
2. Different fixed_iterations settings produce different graphs — H4+H8 creates a structurally different computation graph from H4-only, requiring a fresh round of autotuning
3. Large batch sizes (32768) + complex models (humanoid, ant) create bigger kernels with more tuning parameters

The Inductor's disk cache (~/.cache/torch/inductor) mitigates this within a single session (same model, same graph structure), but any graph change triggers fresh autotuning.

Pathological case: ant-32768 with H4+H8

The combination of coordinate_descent_tuning=True + fixed_iterations=True + batch size 32768 on the ant model caused 300+ GB RAM usage and a multi-hour hang. The fixed_iterations flag unrolls the solver loop, creating a much larger computation graph that the coordinate descent tuner struggles with.

mode="reduce-overhead" (CUDA graphs) — still blocked

CUDA graph capture fails with tensor aliasing errors in scan.py:_take(). The _take function returns views/slices of the input tensor, and CUDA graphs require that all output tensors have unique storage. This is a fundamental limitation:

RuntimeError: Expected all output_idx values when replaying cuda graphs
to be unique but found the following indices duplicated ...

Root cause: _take() returns indexed views that alias the input tensor's storage. CUDA graph replay needs non-aliased outputs.

Possible fixes:

• Clone the sliced tensor in _take() (costs memory bandwidth)
• Restructure scan.flat / scan.body_tree to avoid returning aliased tensors
• Wait for PyTorch to support aliased outputs in CUDA graph replay (upstream issue)

smooth.py optimization opportunities (not yet implemented)

Several patterns in smooth.py generate unnecessary tensor copies that inflate both eager and compiled execution time:

• Redundant clone in factor_m — qld = qld / qld[...] already creates a new tensor; the subsequent qld = qld.clone() is unnecessary.

• Per-iteration full-tensor clones in solve_m — each loop iteration clones the entire x vector just to do an indexed add. Could be replaced with scatter_add.

• Numpy calls at trace time in tendon() — np.nonzero and np.isin are called every step on constant data. These structural indices should be precomputed in device_put() and stored on the model.

• Late numpy-to-tensor conversion in tendon() — torch.as_tensor() from numpy arrays on every call. Should be precomputed and cached on the model, similar to existing _device_precomp pattern.

These are tracked in TODO/25-smooth-cleanup.md.

[vmoens]

Progress update (2026-03-03)

mode="reduce-overhead" (CUDA graphs) now working

After fixing three PyTorch bugs in torch._inductor.cudagraph_trees, torch.compile(mode="reduce-overhead") now works end-to-end with partitioned graphs (caused by while_loop/scan higher-order ops).

PyTorch fixes required (all on vmoens/nomerg-sum-prs):

1. check_memory_pool crash — Inductor passed partition-carried state (model inputs) as partition outputs without cloning, so their storage wasn't in the CUDA graph pool. Fix: clone non-pool outputs in cudagraph_trees.py. (PR #176267)
2. Buffer reuse crash — CUDA graph replay overwrote output buffers that were simultaneously inputs to the next replay (iterative d = step(m, d) loop). Fix: clone inputs before deallocation when they alias prior outputs.
3. Cross-partition input invalidation — alignment check in deferred_cudagraphify accessed tensors from a prior partition's CUDA graph that had already been overwritten. Fix: defer alignment checks past the clone boundary.

mujoco-torch fix: Added .clone() in _take() (scan.py) to break aliasing between scan input/output tensors — required for CUDA graph capture.

Updated benchmark results (H200, float64, B=32768)

Model	compile	reduce-overhead	compile (tuned)
humanoid	1,184,616	1,187,356	2,304,722
ant	501,108	507,384	2,009,884
halfcheetah	2,233,241	1,804,789	3,366,696
walker2d	504,352	448,015	757,820
hopper	2,717,048	3,059,814	3,876,643

reduce-overhead is roughly on par with default compile at B=32768. The tuned variant (coordinate-descent tile tuning) remains significantly faster, suggesting the bottleneck at large batch sizes is kernel compute, not launch overhead. reduce-overhead + tuning is a natural next combination to try.

Remaining items updated

• Fix tensor aliasing in _take() — resolved by adding .clone() after indexing in scan.py.
• CUDA graph support — working with the three PyTorch fixes above.
• Reduce compile time — fullgraph=True produces a massive single graph (~40K+ lines of generated code) due to loop unrolling in scan.flat/scan.body_tree. Compilation takes 7-10 min per model. Investigating graph simplification opportunities.
• reduce-overhead + tuned — combine CUDA graphs with Inductor tile-size tuning for maximum throughput.

[vmoens]

Compile graph simplification analysis (2026-03-03)

The compiled graph for a single step() call produces ~40K+ lines of generated Inductor code, and compilation takes 7-10 minutes per model. Here's a breakdown of where the complexity comes from and actionable opportunities to reduce it.

Root cause: unrolled Python loops

Dynamo unrolls every Python for loop into the graph. A single step() contains ~150-300 individual torch.vmap subgraphs that all get inlined:

Source	Call sites per step	Groups per call	Total unrolled vmaps
scan.flat	6 (cdof, act_dot, get_force, _next_act, integrate_pos, ...)	5-17 per call	~30-100
scan.body_tree	4 (kinematics, crb, com_vel, xfrc_accumulate)	10-20 per call	~40-80
smooth.transmission	1	m.nu iterations (17 for humanoid)	~17
factor_m / solve_m	2	O(nv) depth (~20-50)	~40-100
collision groups	1	varies	~5-10
sensor groups	3 (pos/vel/acc)	varies	~10-15
RK4 multiplier	—	—	×3

Guard overhead: ~310 fields × ~100+ accesses

• Model: ~230 fields (150 MuJoCo + 55 torch-impl + 25 scalar/option)
• Data: ~80 fields
• Each d.qpos, m.nv access through TensorClass generates Dynamo guards
• _take() adds isinstance + .shape[0] guards per group per arg

Simplification opportunities (ranked by impact)

1. Batch scan groups with padding (very high impact, medium difficulty)

Replace the Python for loop over groups in scan.flat / scan.body_tree with a single padded+masked vmap:

# Current: O(groups) vmaps, each unrolled
for i, group in enumerate(groups):
    result[i] = torch.vmap(f)(take(args, group))

# Proposed: single vmap over padded batch
padded_args = pad_groups_to_max_size(args, groups)  # precomputed
mask = compute_mask(groups)                          # precomputed
all_results = torch.vmap(f)(padded_args)             # ONE vmap call
results = [all_results[mask[i]] for i in range(n_groups)]

This would:

• Turn ~150-300 vmaps into ~10-15 vmaps (one per call site)
• Reduce generated code by ~5-10×
• Cut compilation from ~8 min to ~1-2 min
• Make reduce-overhead more effective (fewer CUDA graph partitions)
• May also improve runtime (fewer kernel launches, better occupancy)

The padding needs to be done at device_put time so it's free at step time. Groups with different function signatures (different trntype etc.) would still need separate vmaps.

2. Use torch.ops.higher_order.scan for tree traversals (very high impact, high difficulty)

Replace scan.body_tree Python loop with the native scan HOP:

# Current: Python for loop over tree depth, unrolled
for depth in range(max_depth):
    carry = vmap(f)(carry, x[depth])

# Proposed: native scan (Inductor sees structure, not unrolled blob)
carry, ys = torch.ops.higher_order.scan(f, [carry], [xs])

Currently blocked by module-level cache access in forward.py:319-323. The scan function references the Model's flat/body_tree caches which are set up outside the compiled region.

3. Vectorize smooth.transmission (medium impact, low difficulty)

Replace for i in range(m.nu) with grouped vmap calls:

# Current: unrolled loop with per-actuator branching
for i in range(m.nu):
    if trntype[i] == TrnType.JOINT: ...
    elif trntype[i] == TrnType.TENDON: ...

# Proposed: group actuators by trntype, single vmap per group
for trntype_group in actuator_groups:  # precomputed, ~2-3 groups
    torch.vmap(transmission_fn)(group_args)

4. Reduce TensorClass attribute-access guards (medium impact, medium difficulty — tensordict)

Each d.field goes through TensorClass.__getattr__ → TensorDict._get_str. Dynamo guards on the key string, dict structure, etc. Options:

• Specialized Dynamo variable type for TensorClass that elides structural guards
• Extract frequently-used fields into local variables before the hot path
• Make clone(recurse=False) lighter under Dynamo (it copies the full 80-field internal dict)

5. Make UnbatchedTensor.clone() cheaper under Dynamo (medium impact, medium difficulty — tensordict)

Every torch.vmap call clones non-batched inputs (in_dims=None). Model is passed as non-batched to many vmaps. After adding 7 UnbatchedTensor fields (tendon precomputation, PR #29), each vmap clone traces through UnbatchedTensor.clone() which is more expensive under Dynamo than regular tensor clone. Options:

• Make UnbatchedTensor.clone() a no-op when the data is static (const under compile)
• Reduce the number of vmap calls (ties back to opportunity torch.compile compatibility: progress and remaining work #1)

6. Remove unnecessary guards in _take() (low-medium impact, low difficulty)

_take() contains isinstance checks and x.shape[0] empty-tensor guards. Since scan groups are precomputed and types are known:

• Remove isinstance(obj, np.ndarray) path (never hit at runtime under compile)
• Remove x.shape[0] == 0 empty checks if groups are guaranteed non-empty
• Specialize _take for the compiled path

Summary

The single biggest win is batching scan groups (#1). It's the most impactful change and doesn't require PyTorch upstream changes. An agent could work on this by:

1. Understanding the _build_flat_cache / _build_body_tree_cache structure
2. Adding padding/masking logic at cache-build time
3. Replacing the for loop in scan.flat / scan.body_tree with a single padded vmap
4. Verifying numerical correctness against MJX

The remaining opportunities (#2-#6) are complementary and can be pursued independently.

[vmoens]

Single-vmap scan rewrite: results and analysis

What was attempted

The better-compile-again branch rewrites scan.flat and scan.body_tree to use a single torch.vmap call per invocation (flat) or per tree-depth level (body_tree), instead of one vmap call per type-group. The motivation was to reduce the number of vmap subgraphs that Dynamo unrolls (from ~150-300 down to ~10-15), cutting compile time.

The approach:

1. Pad per-joint data to uniform shapes across all joint types (e.g. qpos padded to MAX_QPOS_PER_JNT=7 elements per joint)
2. Replace Python if/elif on joint type with torch.where conditional chains in callbacks
3. Single vmap over all bodies at once, with a n_joints mask for multi-joint bodies

What worked

• Compile time improved significantly — 47-58% faster on 4 of 5 tested models (the scan loop unrolling was indeed the dominant compile-time cost)
• All 345 tests pass with identical numerical results

What didn't work: runtime performance

The single-vmap approach introduced a steady-state execution regression across all models:

Model	Joint types	Compile (main → branch)	Exec (main → branch)	Steps/sec delta
ant	hinge only	7.42s → 3.14s (-58%)	0.99s → 1.03s	-4%
humanoid	free + hinge	8.94s → 4.01s (-55%)	1.16s → 1.98s	-41%
halfcheetah	slide + hinge	6.62s → 3.12s (-53%)	0.51s → 0.65s	-22%
hopper	slide + hinge	6.68s → 7.96s (+19%)	0.37s → 0.46s	-20%
walker2d	slide + hinge	7.06s → 3.72s (-47%)	3.57s → 4.16s	-14%

(CPU, float64, B=4096, 5 steps × 3 repeats)

Root causes of the runtime regression

1. Fixed-size padding overhead — the initial implementation padded all qpos to 7 elements per joint and dof to 6, even for hinge-only models that only need 1. This was 7× wasted memory/compute for qpos and 6× for dof.

2. torch.where eager evaluation — torch.where(cond, a, b) evaluates both a and b unconditionally. The FREE branch computes q[3:7] (quaternion normalize, etc.) even when all joints are HINGE. Combined with padding, this multiplied the overhead.

3. Per-joint validity masking — the single-vmap approach requires is_valid = j < n_jnts guards and torch.where(is_valid, new_val, old_val) at each iteration, adding overhead per joint per body.

Attempted mitigation: model-adaptive padding (#33)

To recover runtime, we made the padding and branching adaptive to the joint types actually present in each model:

• Added joint_types_present, max_qpos_per_jnt, max_dof_per_jnt to Model
• Replaced fixed MAX_QPOS_PER_JNT/MAX_DOF_PER_JNT with model-derived values in all cache builders
• Guarded torch.where branches with static if has_free: / if has_ball: checks (Dynamo constant-folds these)

For ant (hinge-only), this reduced padding from 7×/6× to 1×/1× and eliminated all FREE/BALL branch computations entirely. This recovered roughly half the regression (e.g. ant went from -14% to -4%), but a gap remained across all models:

• Remaining gap: 4-41% slower execution depending on model
• Cause: inherent overhead of the single-vmap architecture itself — the for j in range(max_j) loop, carry propagation, and validity masking add computation that the per-type-group approach on main avoids entirely

Conclusion

The single-vmap scan rewrite is a compile-time win (~50% faster) but a runtime loss (4-41% slower). The adaptive padding optimization closes part of the gap but can't eliminate the structural overhead of the single-vmap approach.

Possible next steps

1. Hybrid approach: keep main's per-type-group scan for runtime, but reduce compile-time cost through other means:

   • Pre-compile each group's vmap callback separately and cache the compiled function
   • Use torch.compile(dynamic=True) to avoid re-tracing for different group sizes
   • Investigate whether torch._higher_order_ops.scan can replace the Python depth loop in body_tree

2. Accept the compile-time cost: the ~7-9s compile time on main may be acceptable given the runtime performance advantage, especially if AOT compilation (Ahead-of-time compilation: ship pre-compiled models to eliminate warmup #27) can eliminate the warm-up cost entirely

3. Kernel-level optimization: investigate whether Inductor is generating suboptimal code for the per-type-group vmaps (e.g. unnecessary kernel launches that could be fused)

Draft PR with the adaptive padding experiment: #33