Summary

This PR bundles two related fixes in MultiTokenPredictionLayer._checkpointed_forward, both addressed by aligning MTP's recompute plumbing with transformer_block._checkpointed_forward.

Fix 1 — TypeError crash on packed sequences + full recompute

MultiTokenPredictionLayer._checkpointed_forward forwards every kwarg positionally to tensor_parallel.checkpoint. CheckpointFunction.forward calls ctx.save_for_backward(*args), which only accepts tensors and None — so any non-tensor kwarg (today: packed_seq_params: PackedSeqParams; tomorrow potentially inference_params: InferenceContext) triggers:

TypeError: save_for_backward can only save variables, but argument N is of type PackedSeqParams

This breaks THD packed sequences + decoder full activation recompute (--recompute-granularity full) whenever MTP is enabled (--mtp-num-layers >= 1).

TransformerBlock._checkpointed_forward and attention.py:_checkpointed_attention_forward have always avoided this by capturing non-tensor objects via Python closure and only forwarding tensor / None args. MTP's _checkpointed_forward (introduced in !3330 "feat(MoE): support CP and recompute for MTP") didn't follow that convention.

Fix 2 — silent grad-drop / FP8 recompute crash (same as #4766)

This PR also incorporates the fix from #4766. Two distinct problems on the FP8 / MXFP8 + MTP + full recompute path were reported there:

• TE's reentrant te_checkpoint only tracks positional tensor inputs as checkpoint inputs. The pre-existing MTP path passed hidden_states / decoder_input as kwargs to te_checkpoint, so those tensors were not represented in the recompute backward path — gradients for MTP internals were silently detached, causing missing DDP grad-ready hooks and MTP loss divergence vs BF16.
• _proj_and_transformer_layer enters its fp8_autocast inside the function body, so by the time TE's _CheckpointFunction.forward reads FP8GlobalStateManager.is_fp8_enabled() to gate the phase-1 amax-buffer stash (keyed by global_fp8_buffer_pos_fwd_recompute), FP8 is not yet active. Phase 1 skips the stash and phase 2 raises KeyError: global_fp8_buffer_pos_fwd_recompute. (Historically this was masked because TE never ran backward through kwarg tensors — fixing fix 2a above exposes this.)

Changes (mirroring transformer_block._checkpointed_forward)

• _checkpointed_forward now takes named parameters that match _proj_and_transformer_layer's signature.
• A custom_forward closure captures attention_bias, inference_params, packed_seq_params from the enclosing scope.
• Only tensor / None arguments are forwarded positionally to tensor_parallel.checkpoint / te_checkpoint. This applies to both the non-FP8 and FP8 paths, so differentiable activations are correctly tracked in the recompute backward.
• An outer get_fp8_context(self.config) is entered around te_checkpoint only when fp8 + delayed scaling is active (per TE's copy/get_old_fp8_meta_tensors_for_recompute, only delayed scaling actually uses the recompute amax buffer; MXFP8 BlockScaling, Float8CurrentScaling, NVFP4BlockScaling are all noops there). Mirrors transformer_block's outer_quantization_context decision.
• Quantized recompute (fp8 or fp4) now routes through te_checkpoint; non-quantized stays on tensor_parallel.checkpoint. Mirrors transformer_block.py:526. MTP's internal fp4 path is still a noop pending numerical validation; this change only prepares the recompute plumbing.
• The single call site drops the now-unused forward_func first argument.

No semantic change on the existing non-FP8 path where all kwargs are tensor / None — the closure is empty and the call is identical.

Companion PR for dev: #4592

E2E testing:

Note that main does not yet carry the SFT-mock-data plumbing (MockSFTDataset, --sft-mock-dataset-config-json) that lives on dev, so on main we exercise the bug only via the unit test (which still fully hits the affected code path).

Test plan

• Added tests/unit_tests/transformer/test_multi_token_prediction.py::TestMultiTokenPrediction::test_packed_sequences_with_full_recompute
  • Without the fix: fails with the exact TypeError: save_for_backward can only save variables, but argument 10 is of type PackedSeqParams.
  • With the fix: forward + backward complete; all params get main_grad.
• Existing MTP tests still pass on a single-rank smoke run on main: test_constructor_local[1], test_packed_sequences[1-1], test_roll_tensor_with_packed_sequences[1], test_fp8_support[False], test_fp8_support[True] (the last one exercises the FP8 + full-recompute path that [Dev] Fix MTP layer recompute  #4766 reported).
• E2E FP8/MXFP8 + MTP + full recompute closure validation referenced in [Dev] Fix MTP layer recompute  #4766 (Qwen3.5 MXFP8 vs BF16 MTP loss gap closes; see the chart attached on that PR).

🤖 Generated with Claude Code