[faresobeid]

Qwen3.5-MoE (qwen3_5_moe_text) training with ep > 1 and activation checkpointing crashes in backward with ValueError: Pointer argument cannot be accessed from Triton (cpu tensor?) in torchtitan's generate_permute_indices.

Root cause

• FLA memoizes cu_seqlens metadata (prepare_chunk_indices, …) with @tensor_cache, which matches arguments by tensor identity.
• All decoder layers get the same cu_seqlens object. The first linear-attention layer misses that cache in the forward. In its AC recompute (it runs last in backward) it hits the cache, so the recompute runs fewer ops (fewer aten::_to_copy, among others).
• torch's selective checkpointing replays saved outputs by (op, invocation index). The AC policy MUST_SAVEs CUDA→CPU copies, which includes the EP dispatcher's .cpu() of the split sizes. Because of the shifted index, the .to(torch.int32) inside generate_permute_indices gets that saved CPU tensor back. m_sizes/write_offsets end up on CPU, and the Triton fill-indices kernel rejects them.
• Debug run: m_sizes and write_offsets were on cpu only in the 4th recompute (layer 0), while their input total_tokens_per_expert was on cuda:0.
• Without EP there is no saved _to_copy after the mismatch, so the shift was silent.

Fix

Qwen3_5GatedDeltaNet takes a per-layer cu_seqlens.clone(). Every layer then starts with a cache miss, in the forward and in the recompute alike, so both passes run the same op sequence. FLA still caches within a layer.

Checks (1 node, 2 GPUs per run, released)

SFT, 3 steps, fake data, default AC. The tiny HF-format Qwen3.5-MoE has 4 layers (3 linear + 1 full attention), 16 experts, top-4, hidden 512, random HF init:

run	losses steps 1–3	grad norms
main, ep=2	crash: Pointer argument cannot be accessed from Triton (cpu tensor?)
main, ep=1	12.3580 12.9846 12.6149	65.573 56.601 60.373
this PR, ep=1	12.3582 12.9862 12.6147	65.576 56.594 60.386
this PR, ep=2	12.3582 12.9858 12.6145	65.568 56.599 60.369

model.debug 4-layer random-init runs with ep=2, all of which complete 3 steps with this PR: Qwen/Qwen3.5-35B-A3B (crashes on main), Qwen/Qwen3-30B-A3B, zai-org/GLM-4.5-Air. (random_init leaves zero weights, so these only check the code path.)

ruff check / ruff format --check are clean. pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train: 441 passed, 1 failed. The failure is test_cpu_adam.py::test_native_cpu_adamw_matches_fused_torch_and_preserves_gradients, which also fails on origin/main here.

Same pattern in other models

qwen3_8_flash_next/gated_delta_net.py and nemotron_h/mamba.py also pass the shared cu_seqlens into FLA's identity cache inside checkpointed layers. Both get the same one-line clone.

• Nemotron-3-Nano-30B-A3B, 4-layer debug.random_init, ep=2, 3 steps: main and this PR both finish with the same losses (11.7835 ×3). Each Nemotron-H block holds only one mixer (mamba, MoE, or attention), so on main the shifted replay inside a mamba block never reaches a saved MoE .cpu(). The clone keeps the forward and the recompute consistent there anyway.
• Qwen3.8-Flash-Next: not run on GPU. The shared venv's prime_kernels has no indexed_attention kernel, so main and this PR both fail to build the model. The fix is by analogy with Qwen3.5: the GatedDeltaNet and the MoE share one decoder layer.

Relation to #3864 (closed)

#3864 fixed the same FLA-cache-vs-AC-recompute mismatch for cp > 1. It builds the FLA CP context once in the model, before the layer loop. It touches the same forward in both GatedDeltaNet files and removes the per-layer build_cp_context. This PR covers the default cp = 1 path, where FLA derives its metadata from cu_seqlens inside each layer. The two fixes are independent. If #3864 is reopened, the conflict is a small context conflict at the top of forward. Under CP the clone only feeds the cu_seqlens.to(...) passed to build_cp_context, so it is harmless there.

[faresobeid]

Closing: superseded by #3833, which sets fla.utils.FLA_DISABLE_TENSOR_CACHE=True globally (also covers Nemotron-H) and passes cu_seqlens_cpu — fixes the same activation-checkpoint replay mismatch.