NVFP4 performance regression in TE main branch

[cael-ling]

Describe the bug

TransformerEngine PR-2411 introduced optimizations to improve NVFP4 computation performance. Previously, using an image built with this PR delivered significant performance gains during the pre-training of the Moonlight-16B model. This PR was merged into the TE main branch two weeks ago. However, when building an image from the main branch and running the same NVFP4 training configuration, the throughput per GPU dropped dramatically and became highly unstable.

Steps/Code to reproduce bug

To reproduce the issue, perform the following comparative experiments:

Step1: Build two Slurm images:
Refer to the script at link to build two Slurm images, one image is based on PR-2411, and another image is based on TE main (commit id c988548）

Step2: Launch a training job:
Refer to the script at link to launch a training job. It doesn’t necessarily have to be the Moonlight model, other models are also likely to experience similar performance issues.

Expected behavior

The training performance based on these two images are expected to be similar, but in reality, the performance gap is huge (570 TFLOPS VS 240-400 TFLOPS), see log1,log2, which makes NVFP4 training solution unusable.

Environment overview (please complete the following information)

• Environment location: Slurm cluster
• Method of Transformer Engine install: please refer to link
• If method of install is [Docker], provide docker pull & docker run commands used
  please refer to link

[zhongbozhu]

Please double check padding & NVTE_USE_FAST_MATH toggle.

[zhongbozhu]

Here is the expected kernel trace for 128 multiple padding:

[zhongbozhu]

For 64 padding, we lose the row-col padding, so three kernels in quantization instead of two.

Without 64 padding, no group kernel will be performed.

Padding overhead also a big factor without router padding.

[cael-ling]

Please double check padding & NVTE_USE_FAST_MATH toggle.

export NVTE_USE_FAST_MATH=1
--enable-experimental
--moe-router-padding-for-quantization

No changes after enabling these toggles

[cael-ling]

Here is the expected kernel trace for 128 multiple padding:

I can see two kernels: group_amax and row_col_rht_gemm when use the image built from pr-2411. However, in the main-branch build, the timeline shows a long idle period inside groupedlinear, which appears to be caused by cudaMallocAsync, and this is severely impacting performance.

Long idle in the timeline using main branch build image to run the train:

[zhongbozhu]

You are right, this is a bug...

We shouldn't allocate this workspace for the atomic tile scheduler but instead making it a torch tensor that's always live. cc @timmoon10.

We didn't have this problem because we had a clc based scheduler for this kernel which doesn't need a workspace buffer. However clc scheduler assumes knowing problem shape on host, which won't be cuda graph compatible in the future, so we removed it.

[zhongbozhu]

bug fix #2564

[zhongbozhu]

Will close this issue since it's resolved. @cael-ling can you confirm?

[cael-ling]

Will close this issue since it's resolved. @cael-ling can you confirm?

Just tested it, and the issue is fixed now, please close the issue, nice work.