[cael-ling]

Description

Add a single-launch grouped row-scaled NVFP4 cast to the common layer, replacing the per-expert quantize loop used by MoE activations, together with a CUDA-graph-safe variant for dynamic-routing experts.

The grouped cast runs in two single-launch kernels: a grouped amax kernel over the packed (sum_M, K) buffer that produces per-expert rowwise [sum_M] and columnwise [num_tensors*K] amax vectors (both directions fused into one pass), followed by a grouped cast that consumes them.

Both steps have a graph-safe counterpart that reads routing from the grouped tensor's device first_dims/tensor_offsets instead of host split metadata, so an MoE cast can be captured once and replayed with changed routing. The host-split and graph-safe kernels share the same device body, so their outputs are byte-identical.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• 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

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The changes appear safe to merge, with no outstanding findings.

Summary

Adds grouped row-scaled NVFP4 quantization with device-side routing for CUDA graph replay.

• Grouped casts encode packed expert rows as row-scaled NVFP4.
• Grouped amax calls compute per-row values for row-scaled NVFP4.
• PyTorch grouped quantization routes row-scaled NVFP4 calls through the grouped cast.

The latest change adds the header that declares nvte_group_amax. No new issues were found. All six earlier findings are addressed: the per-expert fallback preserves supported shapes, empty inputs, and quantizer checks; grouped input is made contiguous; uniform expert sizes require 128-row alignment; and the operator test includes the required declaration.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Packed BF16 input] --> B[Grouped amax]
  R[Device routing] --> B
  B --> C[Grouped NVFP4 cast]
  R --> C
  C --> D[Packed FP4 data and scales]

Loading

_{Reviews (10) · Last reviewed commit: "Add missing header include" · Reviewed by Greptile}

[timmoon10]

This file is dedicated to NCCL EP support. The new helper function is not related to EP, and in fact is general enough that it could be used beyond MoEs. transformer_engine/pytorch/tensor/grouped_tensor.py would be a more natural place.

@phu0ngng We should consider renaming this file ncclep.py to avoid confusion.

[cael-ling]

@timmoon10 Thanks, the reviews have been replied separately, please help review once you have a window.

[timmoon10]

Now that we have the graph-safe path, robustness is more important than performance for the split-quantize path. I suggest simplifying the logic:

Suggested change

	const bool has_row_scaled_nvfp4 =
	std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(),
	[](const std::unique_ptr<Quantizer> &quantizer) {
	return static_cast<NVFP4Quantizer *>(quantizer.get())->row_scaled_nvfp4;
	});
	if (has_row_scaled_nvfp4) {
	allocation_method = AllocationMethod::UNFUSED;
	quantization_method = QuantizationMethod::UNFUSED;
	} else {
	allocation_method = AllocationMethod::BULK_NVFP4;
	quantization_method = QuantizationMethod::FUSED_NVFP4;
	}

[timmoon10]

/te-ci

[timmoon10]

LGTM, pending CI

[timmoon10]

/te-ci