Hybrid quantization follow-ups

Tracking follow-ups from PR #2817. Detailed implementation notes are already in code comments / xfails.

Recipe / validation

• Validate per-GEMM scaling-mode compatibility for hybrid qfactories.
  • See HybridQuantizer._get_compatible_recipe() in transformer_engine/pytorch/tensor/hybrid_tensor.py.
• Support delayed-scaling requests inside HybridQuantizer.
  • See HybridQuantizer.__init__() in `transformer_engine/pytorch/tensor/hybrid_tensor.py
• Support producing both directional representations at the fused normalization boundary (transformer_engine/pytorch/ops/basic/layer_norm.py)

Performance / kernels

• Fused kernels: one read -> two writes.
• Introduce a policy for when to materialize a tensor for columnwise_source="rowwise_dequantized".
• Support module-/role-specific CustomRecipe quantization alignment. The current recipe-global quantization_alignment must use the maximum requirement across all factory outputs, which can overpad lower-alignment MXFP8/FP8 layers in mixed-format models. Derive alignment from canonical cached module quantizers and propagate it to the corresponding Fp8Padding without speculative qfactory calls.
• Support HybridQuantizer with standard Userbuffers.
• Add native columnwise-only per-tensor FP8 quantization for Hopper, covering both CurrentScaling and DelayedScaling. The legacy per-tensor cast-transpose kernel currently requires both rowwise and columnwise output buffers.

TP/SP

• Add native HybridQuantizer dispatch to gather_along_first_dim.
• Preserve existing SP amax-reduction semantics and cover them in distributed tests.

FSDP2

• Optimize hybrid FSDP2 communication buffers.
  • See HybridQuantizedTensor.fsdp_pre_all_gather() in transformer_engine/pytorch/tensor/hybrid_tensor.py.
• Fix HybridFloat8BlockScaling FSDP2 xfail.
  • See _HYBRID_FLOAT8_BLOCK_FSDP2_XFAIL_REASON in tests/pytorch/distributed/fsdp2_tests/conftest.py.
• Add NVFP4 hybrid sub-storage FSDP2 hooks.
  • See TestHybridFsdpPreAllGatherProtocol.test_nvfp4_sub_storage_raises_on_pre_all_gather() in tests/pytorch/test_hybrid_quantization.py.
  • Non-hybrid FSDP alignment: reuse fsdp extract buffers and fsdp_assign_gather for the current base tensor class (non-hybrid)
• Support Hybrid sub-storages that fall back to a high-precision tensor for an unquantizable local shard (TransformerEngine/transformer_engine/pytorch/tensor/hybrid_tensor.py).

GEMM / quantization

• Support HybridQuantizer / IdentityQuantizer as GEMM output quantizers.
  • See _reject_unsupported_output_quantizer() in transformer_engine/pytorch/cpp_extensions/gemm.py.

GroupedLinear / grouped storage

• Support IdentityQuantizer and HybridQuantizer with GroupedLinear(single_grouped_weight=True).

Distributed optimizer / Megatron

• Support per-block hybrid sub-quantizers in quantize_master_weights.
• Hybrid distributed optimizer: partial-master support with columnwise_source="original". The current one-payload distributed-optimizer path cannot preserve both independently quantized Hybrid directions when each rank owns only a partial master weight. A safety guard rejects this configuration unless full-master data or an FSDP Hybrid shard containing both directions is provided. Unblock by adding a two-payload sharding/all-gather contract for Hybrid rowwise and columnwise storage, including their scale/amax metadata, then relax _validate_hybrid_partial_master_policy.
• Hybrid distributed optimizer with columnwise_source="rowwise_dequantized": reconstruct the column only after the rowwise update/all-gather.
• Complete Megatron-LM quantized_model_init + --fp{4,8}-param-gather + dist opt.
  • See PR #2817 integration notes.
• Complete Megatron-FSDP + --fp{4,8}-param-gather.
  • See PR #2817 integration notes.
• Complete Torch FSDP2 + --fp{4,8}-param-gather.
  • See PR #2817 integration notes and tests/pytorch/distributed/fsdp2_tests/.

Activation recompute

• Investigate vanilla torch.utils.checkpoint(use_reentrant=False) with TE weight-workspace cache.
  • See xfails in TestHybridActivationRecompute in tests/pytorch/test_hybrid_quantization.py.
  • te.checkpoint path is already covered and works.

Validation

• Convergence validation of base non-hybrid recipes.