[BUG] Grouped MXFP8 activation kernels return all-zero output when tensor widths differ (CUDA 12.x builds)

Read transformer_engine/common/cast/core/grouped_tma.cuh and inspect the linkage of g_tensor_maps and update_tma_descriptors; the issue links PR #3704, so check its status before starting. Reproduce with the stated *GroupedFusedCastMXFP8* test filter on a CUDA 12.8 Blackwell build. Done means grouped activation outputs and scales match the CPU reference for varying-width shapes.

Describe the bug

When Transformer Engine is built with CUDA 12.x, the grouped MXFP8 activation functions (nvte_group_gelu,
nvte_group_dgelu, and the other grouped activations) return an output that is entirely zero, with no CUDA
error, when the tensors in the group have different widths (VARYING_LAST_DIM or VARYING_BOTH_DIMS). Plain grouped
quantization (nvte_group_quantize) gives correct results for the same inputs.

Cause. For groups with different widths, two kernels work together:

1. update_tma_descriptors writes one TMA descriptor plus rows/cols/offset per tensor into a GPU table,
   g_tensor_maps.
2. group_quantize_mxfp8_kernel reads that table and does the quantization.

In transformer_engine/common/cast/core/grouped_tma.cuh the table and the writer kernel have different linkage:

// Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs.
static __device__ TensorMapStorage g_tensor_maps;          // one private copy per .cu file

template <typename IType, typename OType>
__global__ void __launch_bounds__(THREADS_PER_WARP)
    update_tma_descriptors(...)                              // externally linked template

• Because of static, every .cu file that includes the header gets its own copy of the table. Since #2983 split
  the grouped activations into separate files, 17 files do this (cast/cast_grouped*.cu and
  activation/*_grouped*.cu).
• With CUDA 12.x, nvcc gives template kernel launchers (host stubs) external linkage, so the 17 identical
  instantiations are merged at link time. The built libtransformer_engine.so has one launcher for
  update_tma_descriptors<__nv_bfloat16, __nv_fp8_e4m3> (nm -C), so every launch runs one file's kernel and
  writes that file's table.
• The activation kernels are compiled only in their own files (e.g. gelu_grouped.cu) and read their own copy,
  which nothing ever wrote. They see rows == 0, return early, and leave the output untouched.

Why it depends on the CUDA version. From CUDA 13.0, nvcc's -static-global-template-stub defaults to true,
which gives template kernel launchers internal linkage in whole-program mode; before 13.0 it defaults to false
(NVIDIA blog, which
describes this exact failure mode). TE builds without -rdc and doesn't set the flag. Measured on the same commit:

Build	Launcher copies in .so	*GroupedFusedCastMXFP8*
CUDA 12.8, unchanged	1	240 failed, 1,055 passed
CUDA 12.8 + -static-global-template-stub=true	17	0 failed, 1,295 passed
CUDA 12.8 + proposed fix (below)	17	0 failed, 1,295 passed
CUDA 13.0, unchanged	17	not run (driver on that machine only supports 12.8)
CUDA 13.0 + -static-global-template-stub=false	1	not run

This is probably why CI doesn't show it. In a 12.x build, which file's launcher survives depends on link order, which
is presumably why plain cast happens to work. A change elsewhere (adding or reordering .cu files, a different
linker) could move the problem to nvte_group_quantize, which PyTorch and JAX use.

Steps/Code to reproduce bug

Build the common library and the C++ tests with CUDA 12.8 for a Blackwell GPU, then run:

tests/cpp/build/operator/test_operator --gtest_filter='*GroupedFusedCastMXFP8*'

Note: on current main, tests/cpp doesn't compile with CUDA < 13.4 (separate problem, see Additional context).
I built the tests with a one-line local workaround in tests/cpp/test_common.h (if → if constexpr in
TypeInfo::Helper::getType); the library itself was unmodified.

Result at 33f3b356 (RTX 5090, CUDA 12.8): 240 failed, 1,055 passed, 1,044 skipped. The failures are exactly
{CAST_ACT, CAST_DACT} × {GeLU, SiLU, ReLU, QGeLU, SReLU} × {VARYING_LAST_DIM (9 cases), VARYING_BOTH_DIMS (15 cases)}, across input types (FP32/FP16/BF16), output types (E4M3/E5M2) and scaling directions.
SAME_BOTH_DIMS, VARYING_FIRST_DIM and CAST_ONLY all pass. Example:

CAST_ACTXGeLU_ROWWISE_VARYING_LAST_DIM_N_3_SHAPE_256X896_bfloat16_float8e4m3
Error in rowwise_scales: Mismatch at (0): 0 vs 118
33 mismatche(s) which is more than tolerable mismatch limit of 32.

The test reports only one bad scale because its compare helpers stop counting after the first mismatch. With
temporary instrumentation in the test (earlier run, 5759fa0f), all 229,376 output values and every scale are 0.
The failure is deterministic (3/3 runs; running the op 2–3 times in a row doesn't change it).
compute-sanitizer memcheck, initcheck and synccheck report no errors.

Expected behavior

Grouped activation + quantize should produce the same result as the CPU reference for every shape mode and every
supported CUDA version, as nvte_group_quantize already does.

Proposed fix

Give the updater the same internal linkage as the table, so each file uses its own updater with its own table:

 template <typename IType, typename OType>
-__global__ void __launch_bounds__(THREADS_PER_WARP)
+static __global__ void __launch_bounds__(THREADS_PER_WARP)
     update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input,

On CUDA 12.8 this gives 17 launcher copies and 0 failures in *GroupedFusedCastMXFP8* and in the wider
*Grouped*:*MXFP8*:*mxfp8*:*Group* sweep. It works on every supported CUDA version, and on CUDA 13 it matches what
the compiler already does by default. An alternative is building with -static-global-template-stub=true when
CUDA >= 12.8, which would also cover any other template kernels with the same pattern, but the flag doesn't exist
before 12.8 (TE supports 12.1+).

I'm happy to open a PR with this.

Environment overview

• Environment location: Cloud (RunPod), Docker container
• Method of install: from source. Common library only, configured directly with CMake:
  cmake -GNinja -S transformer_engine/common -B build -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES=120 -DCUDNN_FRONTEND_INCLUDE_DIR=<nvidia-cudnn-frontend>/include -DNVTE_WITH_NCCL_EP=OFF -DNVTE_WITH_CUTEDSL=OFF;
  C++ tests with cmake -GNinja -S tests/cpp -B tests/cpp/build -DCMAKE_CUDA_ARCHITECTURES=120 -DTE_LIB_PATH=build

Environment details

• OS: Ubuntu 24.04.5, gcc 13.3.0
• Transformer Engine: main at 33f3b356
• CUDA: 12.8 (nvcc 12.8, driver 570.195.03); CUDA 13.0 nvcc for the compile-only rows
• cuDNN: 9.27 (pip, nvidia-cudnn-cu12 / -cu13), cudnn-frontend 1.31.0
• PyTorch/Python: not used (C++ tests only)

Device details

• GPU: NVIDIA GeForce RTX 5090 (compute capability 12.0)

Additional context

• The GPU model shouldn't matter: any Blackwell GPU (these tests skip below compute capability 10.0) running a
  CUDA 12.x build of TE should fail the same way. I haven't run it on B200.
• #3630 (@kainzhong) is a different problem with the same g_tensor_maps: concurrent launches on different streams
  share one table. This issue is the opposite: the writer and reader use different copies. If #3630 is fixed by moving
  the descriptors out of the static __device__ global (e.g. into per-call storage), that would fix this too. The
  one-line change above is independent and could go in first.
• group_dequantize_mxfp8.cuh has a similar structure (per-file __device__ tables plus an externally linked updater
  template), but only one file calls it, so it isn't affected today.
• Separate problem found while reproducing: since #3325, tests/cpp fails to compile with CUDA < 13.4
  (static assertion failed with "tuple index must be in range" in test_common.h). TypeInfo::types only includes
  fp8ue5m3 when CUDA_VERSION >= 13040, but DType::kFloat8UE5M3 always exists, and the non-constexpr if in
  Helper::getType() instantiates tuple_element<11>. I can file this separately.