From 8325d30282234a50abaf12747abc0bb0453a6bb3 Mon Sep 17 00:00:00 2001
From: Jan Bielak <jbielak@nvidia.com>
Date: Wed, 12 Aug 2026 20:37:03 +0000
Subject: [PATCH] Use wide instructions in SR to reduce issue-bound bottleneck

Signed-off-by: Jan Bielak <jbielak@nvidia.com>
---
 transformer_engine/common/util/curanddx.hpp |  9 ++++++--
 transformer_engine/common/util/ptx.cuh      | 23 ++++++++++++++-------
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/transformer_engine/common/util/curanddx.hpp b/transformer_engine/common/util/curanddx.hpp
index 6dd0b57177..d05546a381 100644
--- a/transformer_engine/common/util/curanddx.hpp
+++ b/transformer_engine/common/util/curanddx.hpp
@@ -7,6 +7,8 @@
 #ifndef TRANSFORMER_ENGINE_COMMON_UTIL_CURANDDX_HPP_
 #define TRANSFORMER_ENGINE_COMMON_UTIL_CURANDDX_HPP_
 
+#include <cstdint>
+
 namespace transformer_engine {
 namespace curanddx {
 namespace detail {
@@ -18,8 +20,11 @@ inline constexpr unsigned int philox4x32_m4x32_1 = 0xCD9E8D57U;
 
 __forceinline__ __device__ unsigned int mulhilo32(unsigned int a, unsigned int b,
                                                   unsigned int* hip) {
-  *hip = __umulhi(a, b);
-  return a * b;
+  // Returns uint64_t(a) * b in two uint32 halves
+  uint64_t product;
+  asm("mul.wide.u32 %0, %1, %2;" : "=l"(product) : "r"(a), "r"(b));
+  *hip = static_cast<unsigned int>(product >> 32);
+  return static_cast<unsigned int>(product);
 }
 
 __forceinline__ __device__ uint4 single_round(uint4 ctr, uint2 key) {
diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh
index 8dbb8b3044..2cbf1f019d 100644
--- a/transformer_engine/common/util/ptx.cuh
+++ b/transformer_engine/common/util/ptx.cuh
@@ -951,6 +951,8 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding(
     } else if constexpr (std::is_same<SCALING_COEFFICIENT_TYPE, float>::value) {
       asm volatile(
           "{\n"
+          ".reg.b64 scaling_coeff_2x; \n\t"
+          "mov.b64 scaling_coeff_2x, {%3, %3}; \n\t"
           ".reg.b16 v0_bf16, v1_bf16, v2_bf16, v3_bf16, v4_bf16, v5_bf16, v6_bf16, v7_bf16; \n\t"
           "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16}, %1; \n\t"
           "mov.b64 {v4_bf16, v5_bf16, v6_bf16, v7_bf16}, %2; \n\t"
@@ -965,14 +967,19 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding(
           "cvt.f32.bf16 v6, v6_bf16; \n\t"
           "cvt.f32.bf16 v7, v7_bf16; \n\t"
 
-          "mul.f32 v0, v0, %3; \n\t"
-          "mul.f32 v1, v1, %3; \n\t"
-          "mul.f32 v2, v2, %3; \n\t"
-          "mul.f32 v3, v3, %3; \n\t"
-          "mul.f32 v4, v4, %3; \n\t"
-          "mul.f32 v5, v5, %3; \n\t"
-          "mul.f32 v6, v6, %3; \n\t"
-          "mul.f32 v7, v7, %3; \n\t"
+          ".reg.b64 v01, v23, v45, v67; \n\t"
+          "mov.b64 v01, {v0, v1}; \n\t"
+          "mov.b64 v23, {v2, v3}; \n\t"
+          "mov.b64 v45, {v4, v5}; \n\t"
+          "mov.b64 v67, {v6, v7}; \n\t"
+          "mul.f32x2 v01, v01, scaling_coeff_2x; \n\t"
+          "mul.f32x2 v23, v23, scaling_coeff_2x; \n\t"
+          "mul.f32x2 v45, v45, scaling_coeff_2x; \n\t"
+          "mul.f32x2 v67, v67, scaling_coeff_2x; \n\t"
+          "mov.b64 {v0, v1}, v01; \n\t"
+          "mov.b64 {v2, v3}, v23; \n\t"
+          "mov.b64 {v4, v5}, v45; \n\t"
+          "mov.b64 {v6, v7}, v67; \n\t"
           ".reg.b16 b03, b47; \n\t"
           // Elements reordered to match e2m1x4 packing order (v3,v2,v1,v0)
           "cvt.rs.satfinite.e2m1x4.f32 b03, {v3, v2, v1, v0}, %4; \n\t"