[Common] Use wide instructions in SR to reduce issue-bound bottleneck

[janekb04]

Description

Optimizes the performance of stochastic rounding kernel by using wider instructions.

A simple benchmark shows a 4-10% throughput improvement on a B200 for the NVFP4 specialized 1D quantization kernel in SR mode. The SR kernel is issue-bound, so using fewer, wider instructions, improves performance.

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:

• In the Philox RNG, mul.wide.u32 is used explicitly, as I noticed, while working on [Common] Experimental CuTeDSL NVFP4 backend #3223 that the multiply and add would not get automatically fused into an FMA.
• In mul_cvt_bf16_to_fp4_8x_stochastic_rounding, use mul.f32x2 to halve the number of mul instructions.

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

Raw results

version                    M     RN GB/s     SR GB/s
baseline                1024        41.1        39.0
baseline                2048       173.6       156.7
baseline                3072       412.5       363.8
baseline                4096       738.7       587.0
baseline                5120      1059.6       852.3
baseline                6144      1315.6      1048.0
baseline                7168      1557.5      1202.3
baseline                8192      1787.7      1391.3
baseline                9216      2087.7      1543.9
baseline               10240      2234.6      1668.4
baseline               11264      2419.8      1807.6
baseline               12288      2531.7      1948.7
baseline               13312      2722.5      1990.5
baseline               14336      2807.6      2076.6
baseline               15360      2951.8      2174.3
baseline               16384      3044.5      2220.1
wide                    1024        45.8        40.8
wide                    2048       182.2       169.5
wide                    3072       420.6       361.3
wide                    4096       748.3       635.2
wide                    5120      1079.3       897.8
wide                    6144      1370.9      1114.4
wide                    7168      1610.5      1307.9
wide                    8192      1828.8      1499.3
wide                    9216      2107.0      1689.5
wide                   10240      2264.2      1830.9
wide                   11264      2478.6      1971.9
wide                   12288      2639.5      2081.5
wide                   13312      2699.4      2154.4
wide                   14336      2877.9      2257.1
wide                   15360      2986.0      2339.4
wide                   16384      3075.0      2386.7

[janekb04]

@Oleg-Goncharov @ksivaman, I think you might be interested in reviewing.

[greptile-apps]

Greptile Summary

The PR optimizes stochastic-rounding GPU arithmetic without changing its public interfaces.

• Replaces separate Philox low/high multiplication operations with one mul.wide.u32.
• Replaces eight scalar FP32 scaling operations with four packed mul.f32x2 operations in the NVFP4 stochastic-rounding path.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/util/curanddx.hpp	Uses a single wide PTX multiplication to produce the two 32-bit Philox product halves; identified compilation routes remain compatible.
transformer_engine/common/util/ptx_arch_spec.cuh	Packs float pairs for wide multiplication while preserving the existing lane order and architecture gating.

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into wide_instructio..." | Re-trigger Greptile}

[janekb04]

/te-ci

[Oleg-Goncharov]

LGTM. I would just check whether static_cast<uint64_t>(a) * b generates the same SASS as mul.wide.u32.

[Oleg-Goncharov]

/te-ci

[janekb04]

/te-ci