[arrdel]

Description

Adds step-indexed trajectory activation caching (DiffusionTrajectoryCache) and step-selective activation patching across denoising steps for diffusion LMs on TransformerBridge.diffusion_generate(), wires LLaDAArchitectureAdapter into diffusion_generate(), and adds demos/Diffusion_Trajectory_Patching_Demo.ipynb.

Motivation

Autoregressive generate(..., return_cache=True) recovers generation activations from a single post-hoc run_with_cache(output_tokens) pass because causal masking keeps earlier positions invariant as new tokens are appended. Masked discrete-diffusion LMs (LLaDA, Dream, LLaDA-2, GIDD) use bidirectional self-attention (is_causal=False) and iteratively unmask a canvas $x_0 \to x_1 \to \dots \to x_S$ over $S$ denoising steps:

1. A post-hoc pass on the final unmasked sequence $x_S$ does not recover intermediate denoising activations, which depend on the partially-masked canvas $x_s$ at step $s$.
2. Wrapping diffusion_generate() in add_caching_hooks() overwrites cache[name] on every forward pass, discarding steps $0 \dots S-2$ and retaining only step $S-1$.
3. Caching all hooks at every step on GPU across a 256-step trajectory on an 8B model would exhaust device memory unless filtered by step (cache_steps), position (pos_slice), and offload device (device="cpu").

What this PR adds

• DiffusionTrajectoryCache (transformer_lens/ActivationCache.py):
  • Mapping from 0-based denoising step index s to a per-step ActivationCache, plus step_tokens[s] recording the partially-masked input canvas $x_s$ entering step $s$ and num_steps_run.
  • stack(key, steps=None) stacks cache[s][key] along dim 0 into [n_steps, ...].
  • make_step_patch_hook(hook_name, steps=..., positions=...) builds a forward hook that reads bridge.current_diffusion_step and patches activations from a recorded trajectory at selected denoising steps and token positions.
• TransformerBridge.diffusion_generate(..., return_cache=True) (transformer_lens/model_bridge/transformer_bridge.py):
  • Adds keyword-only return_cache, names_filter, cache_steps (int, iterable of int, callable predicate, or None), pos_slice, device, and remove_batch_dim, plus bridge.current_diffusion_step.
  • Skips tensor detachment and device transfers on steps excluded by cache_steps.
• LLaDAArchitectureAdapter sampler binding (transformer_lens/model_bridge/supported_architectures/llada.py):
  • Binds LLaDA's blockwise low-confidence remasking loop (diffusion_generate) in prepare_model when not already present on the model and declares native_sampler = "diffusion_generate" with DiffusionContractTests coverage.
• Demo notebook (demos/Diffusion_Trajectory_Patching_Demo.ipynb):
  • Exercises trajectory-step caching, cache_steps/pos_slice filtering, and pre- vs. post-commitment activation patching across denoising steps on a compact CPU LLaDAModelLM bridge (runs in ~4s under nbval), alongside measured LLaDA-8B-Instruct and Dream-v0-Instruct-7B ARC-Challenge commitment-order tables ($T_{\text{conv}}$, $T_{\text{commit}}$, unwritten reasoning fraction $U$, and order gap) and an opt-in TL_RUN_RELEASED_LLADA=1 section for GSAI-ML/LLaDA-8B-Instruct.

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

Thank you for building this out @arrdel, the trajectory cache fills a gap for the diffusion adapters and the step-filtered caching is the right approach. A handful of comments below that need addressing before I can merge

[jlarson4]

A forward pre-hook on original_model only fires through __call__. LLaDA2's and Gidd's released generate call self.forward(...) directly, meaning on those adapters the counter never advances, every caching hook returns early, and return_cache=True returns an empty trajectory with no error. Can we drive the counter by something every sampler passes through, with an error when a cache was requested but no step was observed?

[jlarson4]

Slice(pos_slice) already accepts a slice object, so the hand-built instance can go; if the SliceInput alias is what's in the way, a cast or adding slice to the alias is simpler.

[jlarson4]

This flips LLaDA from skipping verify_models to running all four phases. Can you validate this change by completing a verify_models run on a LLaDA checkpoint to confirm the phases pass? If you do not have access to the necessary compute, I can pull down your branch and run some, just let me know.

[jlarson4]

We'd rather not add a demo notebook for this. Most of it rebuilds the tiny LLaDA module the integration test already has, and section 4 is a typed-in results table we cannot reproduce in the repo, so it would be documenting a research claim rather than the feature. Please drop the notebook and its docs/CI registrations (including the test_make_docs.py list). There is one check we'd like to keep, that positions already unmasked before a step are unchanged by a patch at that step. Please move it into the LLaDA integration test file. It would also be good to give the sampler port direct tests of its own (temperature > 0, the random and unsupported remasking paths, a multi-block schedule, and the attention_mask branch), since today it only ever runs with defaults.