Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
janekb04 wants to merge 2 commits into
base: main
Choose a base branch
from
+22 −10
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8325d30
Use wide instructions in SR to reduce issue-bound bottleneck
janekb04 Aug 12, 2026
0eafb55
Merge branch 'main' into wide_instruction
janekb04 Aug 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions transformer_engine/common/util/curanddx.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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) {
Expand Down
23 changes: 15 additions & 8 deletions transformer_engine/common/util/ptx.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand Down
Loading