[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 NVFP4 stochastic rounding by reducing the instruction count while preserving the existing arithmetic.

• Replaces separate Philox high/low multiplication operations with one wide unsigned multiply.
• Packs eight BF16-derived values into pairs and scales them using four mul.f32x2 instructions.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete correctness, compatibility, or security failures identified.

The new wide PTX operations preserve the prior arithmetic, and the packed floating-point instruction is limited to compatible architecture-specific paths.

Important Files Changed

Filename	Overview
transformer_engine/common/util/curanddx.hpp	Replaces separate high- and low-half multiplication with an equivalent mul.wide.u32 operation supported by the configured CUDA targets.
transformer_engine/common/util/ptx.cuh	Replaces eight scalar multiplications with four packed multiplications under architecture guards compatible with the instruction.

_{Reviews (1): 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]

Is inline PTX necessary here? Could we simply use:
uint64_t product = static_cast<uint64_t>(a) * b;
and let the compiler generate the appropriate instruction?