[sraman-rgb]

Accept NVFP4BlockScaling as the EP dispatch_fwd_quant_recipe and combine_bwd_quant_recipe. Tokens travel as packed NVFP4 rows (E2M1 codes, E4M3 block scales and a power-of-two FP32 token scale) through the unscaled dispatch / combine-backward path.

Under NVFP4 autocast in graph mode, the receiver builds the expert GEMM's NVFP4 GroupedTensor straight from the rows in one pass: the rowwise operand reuses the sender's codes and the columnwise RHT operand is quantized with a bound-derived per-expert scale, so no BF16 round trip or amax pass is needed. The combine backward uses stochastic rounding per the recipe. Otherwise the rows are decoded to BF16.

Add nvte_ep_nvfp4_pack / _unpack / _direct_receive to the common EP API, the tex bindings, single-GPU tests and multi-GPU EP harness tests.

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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 PR appears safe to merge; the latest change fixes the capture test's pass selection.

Summary

Adds NVFP4 token transport for EP dispatch and combine backward.

• Expert-parallel dispatch can send tokens as packed NVFP4 rows.
• Combine backward can return NVFP4 gradient rows for expert GEMMs.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[BF16 tokens or gradients] --> B[Pack NVFP4 rows]
  B --> C[EP transport]
  C --> D{Direct receive enabled?}
  D -->|Yes| E[NVFP4 GroupedTensor]
  D -->|No| F[Decode to BF16]

Loading

_{Reviews (5) · Last reviewed commit: "[PyTorch] EP test: keep the alignment-pa..." · Reviewed by Greptile}

[phu0ngng]

Two questions here

1. I think you will need to sync the power-of-two FP32 token scale before quantization within the DP group already. Why do you need to send it one more time here?
2. When we don't need to send the extra FP32 tensor scale, we should be able to use the current existing infrastructure to dispatch token and block scale with this packing.