Support paged stashing for GroupedLinear activations

[lhb8125]

Description

Enable paged stashing integrations to identify the activation storage saved by the device-initiated PyTorch GroupedLinear autograd function.

• Mark input_to_save immediately before prepare_for_saving / save_for_backward.
• Extend mark_grouped_tensor to support unquantized BF16/FP16 rowwise storage.
• Preserve the existing quantized columnwise data and scale-inverse markers.
• Add unit coverage for rowwise and columnwise storage.

This is a follow-up to #3224. The paired Megatron-LM integration is NVIDIA/Megatron-LM#6828.

Validation

• Black 24.4.2 with the repository pre-commit arguments
• Ruff
• Full TE license check
• Four-rank GB300 Megatron-LM integration test using HybridEP, MXFP8, device-initiated GroupedLinear, and GPU-only paged stash: TestPagedStashingGroupedTensor::test_forward_backward_without_op_fuser
  • all four ranks passed
  • activation metadata captured
  • stash/reload completed without overflow
  • output and input-gradient parity passed