TE-FL Upgrade: Synchronization with TE Release V2.14

[lxd-cumt]

Merge upstream release_v2.14 into main

Integrates NVIDIA TransformerEngine upstream release v2.14 (304 commits, v2.9.0 → v2.14.0) into the TransformerEngine-FL fork via tree replacement strategy, preserving the
custom plugin system while incorporating upstream enhancements.

Upstream Enhancements

Quantization & Precision

• MXFP8 grouped GEMM with persistent quantization kernels and tensor-scaled FP8 support
• NVFP4 grouped quantization with Hadamard transform for MoE workloads
• QuantizedTensor support in FusedAdam optimizer for MXFP8/Float8 block scaling

Architecture Support

• Blackwell (sm120) fused attention support with cuDNN 9.18.1+
• Deterministic training on Blackwell with cuDNN ≥9.18.1
• Grouped GEMM cuBLAS bindings with bias support and tensor swizzling

Distributed Training

• FSDP2 support with DTensor-aware optimizer states and allgather optimizations
• Collective GEMM with FP8/MXFP8 for JAX
• GroupedTensor torch ops for DDP and distributed optimizer

Operators

• Fused RMSNorm dLN with add-through via cuDNN
• MoE grouped MLP ops with split dBias and router kernel JAX bindings
• Configurable philox rounds for stochastic rounding

FlagOS Features

Plugin System Preservation

• Synced plugin OP API signatures with upstream csrc changes (fused_attn_fwd/bwd parameters, attention backend dispatch)
• Patched new upstream CUDA hardcoding to te_device_type() for multi-backend compatibility
• Fixed stale references to renamed upstream symbols (e.g., CPUOffloadEnabled → is_cpu_offload_enabled())

Verification

• Build & import validation passed
• Unit & integration tests passed
• FlagScale end to end training test, summary as follows:

Qwen3-32B, 16 layers, 20 iters, 1node x 8 gpus

Config	Status	Avg Throughput (tokens/s/gpu)	Note
vendor-flash	PASS	124.33
vendor-fused	PASS	121.95
vendor-unfused	PASS	108.35
flagos-flash	PASS	94.20
flagos-fused	FAIL	—	No fused attention backend supports for flagos backend
flagos-unfused	PASS	65.01
reference-flash	PASS	93.14
reference-fused	FAIL	—	No fused attention backend support for reference backend
reference-unfused	PASS	66.89

DeepSeek-V3 16BA3B, 18 layers with 1 mtp layer, 20 iters, 1node x 8gpus, there is no flash-attn or fused-attn support for multi-latent attention

Config	Status	Avg Throughput (tokens/s/gpu)	Note
vendor-unfused	PASS	47.00
flagos-unfused	PASS	18.57
reference-unfused	PASS	20.68

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 17 committers have signed the CLA.

✅ BrianPei
❌ timmoon10
❌ phu0ngng
❌ ptrendx
❌ ksivaman
❌ zhongbozhu
❌ vthumbe1503
❌ jomitchellnv
❌ vcherepanov-nv
❌ CarlosGomes98
❌ KshitijLakhani
❌ lixianduo
❌ cspades
❌ Oleg-Goncharov
❌ jberchtold-nvidia
❌ sudhakarsingh27
❌ pstjohn

---

lixianduo seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}