From c11c8c721e6324dc58fb5710a429e98bbfa20bf0 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Fri, 30 Oct 2026 08:03:42 -0700 Subject: [PATCH] perf(cake_gated_act_mxfp8): band-path program for the SM107 forward both entry below 2^20 - SM107 forward both, M * K below 2^20: the plain program with the call band path at 4 CTAs per SM (new record). - From 2^20 up to 2^23 the previous plain program still runs; its record now starts at 2^20 on SM107. Larger sizes and SM100 / SM103 are unchanged. Co-Authored-By: Claude --- .../cake_gated_act_mxfp8_backend/build.json | 23 + .../cake_gated_act_mxfp8_backend/cake_jit.py | 77 +- ..._act_mxfp8_e62d5f6dda3083f31950_binding.cu | 93 ++ ...d_act_mxfp8_e62d5f6dda3083f31950_kernel.cu | 1103 +++++++++++++++++ 4 files changed, 1282 insertions(+), 14 deletions(-) create mode 100644 flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu create mode 100644 flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu diff --git a/flashinfer/cake_gated_act_mxfp8_backend/build.json b/flashinfer/cake_gated_act_mxfp8_backend/build.json index bda91ca7dc0..962d6ebaf2e 100644 --- a/flashinfer/cake_gated_act_mxfp8_backend/build.json +++ b/flashinfer/cake_gated_act_mxfp8_backend/build.json @@ -1578,6 +1578,29 @@ "csrc/cake_gated_act_mxfp8_dd19cc9df9f6cf060a97_binding.cu" ] }, + { + "arch": "sm_107a", + "closure_sha256": "70a788a2fa82a216b95b2a657014b191869170665c385789d0104cff922ee185", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_gated_act_mxfp8_e62d5f6dda3083f31950_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gated_act_mxfp8_e62d5f6dda3083f31950_sm_107a.so", + "sha256": { + "csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu": "b43d9d0b40c2f387f9cbb87700442e72351f90cdd81960b509d491e07d29fa7b", + "csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu": "cee181a75c1de648acf411ddfea8f981086fae4972eedfa186000692463687c9" + }, + "sources": [ + "csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu", + "csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu" + ] + }, { "arch": "sm_100a", "closure_sha256": "868e35d8508dd8fb540376ef9f63d5165730e90306e1c231642357cff205a818", diff --git a/flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py b/flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py index 8b52adc7c9e..4926bf785ae 100644 --- a/flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py +++ b/flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py @@ -1314,12 +1314,10 @@ "route": "both", "select_rule": None, }, - # The forward, rowwise and colwise: the route's fallback on sm_107a; on - # sm_103a it runs when M * K >= 2^21 and no record of a higher threshold - # admits the call. - "cake_gated_act_mxfp8_0520c14fde53ef7bfea7": { + # The forward, rowwise and colwise, with 4 CTAs per SM and the called band + # path: the route's fallback on sm_107a (no select_rule). + "cake_gated_act_mxfp8_e62d5f6dda3083f31950": { "arches": [ - "sm_103a", "sm_107a", ], "arg_plan": [ @@ -1340,8 +1338,8 @@ "tile_cols": 128, "tile_rows": 32, }, - "kernel": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7", - "kernel_symbol": "kernel_cake_gated_act_mxfp8_0520c14fde53ef7bfea7", + "kernel": "cake_gated_act_mxfp8_e62d5f6dda3083f31950", + "kernel_symbol": "kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950", "launch": { "block": (256, 1, 1), "cluster": (1, 1, 1), @@ -1349,7 +1347,7 @@ "dynamic_smem_bytes": 8192, "use_pdl": False, }, - "library": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7", + "library": "cake_gated_act_mxfp8_e62d5f6dda3083f31950", "logical": { "K": "K", "M": "M", @@ -1361,12 +1359,7 @@ }, "role": "fwd", "route": "both", - "select_rule": { - "elements": { - "sm_103a": 2097152, - }, - "kind": "elements_at_least", - }, + "select_rule": None, }, # The forward, rowwise and colwise, with 8 CTAs per SM, the packed sigmoid # and the called band path (on sm_100a only): it runs when M * K >= 2^22 @@ -1475,6 +1468,61 @@ "kind": "elements_at_least", }, }, + # The forward, rowwise and colwise (on sm_103a and sm_107a only): it runs + # when M * K >= 2^21 (2^20 on sm_107a) and no record of a higher threshold + # admits the call. + "cake_gated_act_mxfp8_0520c14fde53ef7bfea7": { + "arches": [ + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "gated_input"], + ["buffer", "row_output"], + ["buffer", "col_output"], + ["buffer", "row_scales"], + ["buffer", "col_scales"], + ["parameter", "M"], + ["parameter", "K"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "kind": "tiles", + "tile_cols": 128, + "tile_rows": 32, + }, + "kernel": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7", + "kernel_symbol": "kernel_cake_gated_act_mxfp8_0520c14fde53ef7bfea7", + "launch": { + "block": (256, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 8192, + "use_pdl": False, + }, + "library": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7", + "logical": { + "K": "K", + "M": "M", + "col_output": "col_output", + "col_scales": "col_scales", + "gated_input": "gated_input", + "row_output": "row_output", + "row_scales": "row_scales", + }, + "role": "fwd", + "route": "both", + "select_rule": { + "elements": { + "sm_103a": 2097152, + "sm_107a": 1048576, + }, + "kind": "elements_at_least", + }, + }, # The forward, rowwise and colwise, with 8 CTAs per SM and the packed # sigmoid (on sm_103a only): it runs when M * K >= 2^22 and no record of a # higher threshold admits the call. @@ -3401,6 +3449,7 @@ "cake_gated_act_mxfp8_6902c8391e6f660586a4", ], "fwd_both": [ + "cake_gated_act_mxfp8_e62d5f6dda3083f31950", "cake_gated_act_mxfp8_0520c14fde53ef7bfea7", "cake_gated_act_mxfp8_3b4157aba492acf411db", "cake_gated_act_mxfp8_bbe04581406928a16cbc", diff --git a/flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu b/flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu new file mode 100644 index 00000000000..0d89d6b2a3d --- /dev/null +++ b/flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by Cake: the tvm-ffi host binding of kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950(__nv_bfloat16* __restrict__ gated_input, uint8_t* __restrict__ row_output, uint8_t* __restrict__ col_output, uint8_t* __restrict__ row_scales, uint8_t* __restrict__ col_scales, int M, int K); + +namespace cake_binding_kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void Run(TensorView a_gated_input, TensorView a_row_output, TensorView a_col_output, TensorView a_row_scales, TensorView a_col_scales, int64_t a_M, int64_t a_K, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_gated_input.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_gated_input, "gated_input"); + check_dtype(a_gated_input, DLDataType{kDLBfloat, 16, 1}, "gated_input"); + check_cuda_tensor(a_row_output, "row_output"); + check_dtype(a_row_output, DLDataType{kDLUInt, 8, 1}, "row_output"); + check_cuda_tensor(a_col_output, "col_output"); + check_dtype(a_col_output, DLDataType{kDLUInt, 8, 1}, "col_output"); + check_cuda_tensor(a_row_scales, "row_scales"); + check_dtype(a_row_scales, DLDataType{kDLUInt, 8, 1}, "row_scales"); + check_cuda_tensor(a_col_scales, "col_scales"); + check_dtype(a_col_scales, DLDataType{kDLUInt, 8, 1}, "col_scales"); + TVM_FFI_CHECK(a_M >= -2147483648LL && a_M <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'M' = " << a_M + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_K >= -2147483648LL && a_K <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'K' = " << a_K + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_gated_input, "gated_input"); + check_contiguous(a_row_output, "row_output"); + check_same_device(a_row_output, a_gated_input, "row_output", "gated_input"); + check_contiguous(a_col_output, "col_output"); + check_same_device(a_col_output, a_gated_input, "col_output", "gated_input"); + check_contiguous(a_row_scales, "row_scales"); + check_same_device(a_row_scales, a_gated_input, "row_scales", "gated_input"); + check_contiguous(a_col_scales, "col_scales"); + check_same_device(a_col_scales, a_gated_input, "col_scales", "gated_input"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + void* v_gated_input = DataPtr(a_gated_input); + void* v_row_output = DataPtr(a_row_output); + void* v_col_output = DataPtr(a_col_output); + void* v_row_scales = DataPtr(a_row_scales); + void* v_col_scales = DataPtr(a_col_scales); + int32_t v_M = static_cast(a_M); + int32_t v_K = static_cast(a_K); + void* kernel_args[] = {&v_gated_input, &v_row_output, &v_col_output, &v_row_scales, &v_col_scales, &v_M, &v_K}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(256, 1, 1); + config.dynamicSmemBytes = 8192; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu b/flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu new file mode 100644 index 00000000000..370c86f92ab --- /dev/null +++ b/flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu @@ -0,0 +1,1103 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_TILE_OFF 0 +#define SMEM_TILE_STAGE_BYTES 8192 +#define SMEM_TILE_STRIDE 8192 +#define SMEM_TOTAL 8192 +#define THREADS 256 +#define LAUNCH_MIN_BLOCKS 4 + +#include + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +__device__ __forceinline__ __nv_bfloat162 __as_bf16x2(unsigned int v) { + __nv_bfloat162_raw raw; + raw.x = static_cast(v); + raw.y = static_cast(v >> 16); + return __nv_bfloat162(raw); +} + +__device__ __forceinline__ float2 cake_bf16x2_to_f32x2(__nv_bfloat162 v) { + const unsigned int w = *reinterpret_cast(&v); + return make_float2(__uint_as_float(w << 16), __uint_as_float(w & 0xffff0000u)); +} + +__device__ __forceinline__ float mul_f32_rn_noftz(float a, float b) { + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __noinline__ unsigned int band_pair(__nv_bfloat16* src, int index, int k) +{ + unsigned int _vec_load_0[1]; + { + _vec_load_0[0] = *reinterpret_cast(src + index + 0); + } + unsigned int _vec_load_1[1]; + { + _vec_load_1[0] = *reinterpret_cast(src + (index + k) + 0); + } + float2 cvt_f32 = cake_bf16x2_to_f32x2(__as_bf16x2(_vec_load_0[0])); + float2 cvt_f32_1 = cake_bf16x2_to_f32x2(__as_bf16x2(_vec_load_1[0])); + float e0 = approx_exp2((-cvt_f32.x) * 1.442695f); + float e1 = approx_exp2((-cvt_f32.y) * 1.442695f); + float d0 = 1.0f + e0; + float d1 = 1.0f + e1; + float r0 = approx_rcp(d0); + float r1 = approx_rcp(d1); + float _fma_0; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_0) : "f"(-d0), "f"(r0), "f"(1.0f)); + float fma; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma) : "f"(r0), "f"(_fma_0), "f"(r0)); + float _fma_1; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_1) : "f"(-d1), "f"(r1), "f"(1.0f)); + float fma_1; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1) : "f"(r1), "f"(_fma_1), "f"(r1)); + float t0 = fma; + float t1 = fma_1; + float d4_8a; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(d4_8a) : "f"(d0), "f"(0.25f)); + float r4_8a = approx_rcp(d4_8a); + float _fma_3; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_3) : "f"(-d4_8a), "f"(r4_8a), "f"(1.0f)); + float _fma_2; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_2) : "f"(r4_8a), "f"(_fma_3), "f"(r4_8a)); + float big_8a; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(big_8a) : "f"(mul_f32_rn_noftz(_fma_2, 1.2676506e+30f)), "f"(1.4073749e+14f)); + float _add_4; + asm("add.rn.f32 %0, %1, %2;" : "=f"(_add_4) : "f"(big_8a), "f"(8.388608e+06f)); + float whole_8a; + asm("sub.rn.f32 %0, %1, %2;" : "=f"(whole_8a) : "f"(_add_4), "f"(8.388608e+06f)); + float above_8a; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(above_8a) : "f"(-d4_8a), "f"(mul_f32_rn_noftz(mul_f32_rn_noftz(big_8a, 7.888609e-31f), 7.1054274e-15f)), "f"(1.0f)); + float _sub_6; + asm("sub.rn.f32 %0, %1, %2;" : "=f"(_sub_6) : "f"(big_8a), "f"(whole_8a)); + float _abs_5; + asm("abs.f32 %0, %1;" : "=f"(_abs_5) : "f"(_sub_6)); + float _add_7; + asm("add.rn.f32 %0, %1, %2;" : "=f"(_add_7) : "f"(big_8a), "f"(0.5f)); + float _sub_8; + asm("sub.rn.f32 %0, %1, %2;" : "=f"(_sub_8) : "f"(big_8a), "f"(0.5f)); + float rounded_8a = ((_abs_5 == 0.5f) ? ((above_8a > 0.0f) ? _add_7 : _sub_8) : whole_8a); + float u0 = ((d0 >= 8.507059e+37f) ? mul_f32_rn_noftz(mul_f32_rn_noftz(rounded_8a, 7.888609e-31f), 1.7763568e-15f) : t0); + float d4_8a_0; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(d4_8a_0) : "f"(d1), "f"(0.25f)); + float r4_8a_0 = approx_rcp(d4_8a_0); + float _fma_10; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_10) : "f"(-d4_8a_0), "f"(r4_8a_0), "f"(1.0f)); + float _fma_9; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_9) : "f"(r4_8a_0), "f"(_fma_10), "f"(r4_8a_0)); + float big_8a_0; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(big_8a_0) : "f"(mul_f32_rn_noftz(_fma_9, 1.2676506e+30f)), "f"(1.4073749e+14f)); + float _add_11; + asm("add.rn.f32 %0, %1, %2;" : "=f"(_add_11) : "f"(big_8a_0), "f"(8.388608e+06f)); + float whole_8a_0; + asm("sub.rn.f32 %0, %1, %2;" : "=f"(whole_8a_0) : "f"(_add_11), "f"(8.388608e+06f)); + float above_8a_0; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(above_8a_0) : "f"(-d4_8a_0), "f"(mul_f32_rn_noftz(mul_f32_rn_noftz(big_8a_0, 7.888609e-31f), 7.1054274e-15f)), "f"(1.0f)); + float _sub_13; + asm("sub.rn.f32 %0, %1, %2;" : "=f"(_sub_13) : "f"(big_8a_0), "f"(whole_8a_0)); + float _abs_12; + asm("abs.f32 %0, %1;" : "=f"(_abs_12) : "f"(_sub_13)); + float _add_14; + asm("add.rn.f32 %0, %1, %2;" : "=f"(_add_14) : "f"(big_8a_0), "f"(0.5f)); + float _sub_15; + asm("sub.rn.f32 %0, %1, %2;" : "=f"(_sub_15) : "f"(big_8a_0), "f"(0.5f)); + float rounded_8a_0 = ((_abs_12 == 0.5f) ? ((above_8a_0 > 0.0f) ? _add_14 : _sub_15) : whole_8a_0); + float u1 = ((d1 >= 8.507059e+37f) ? mul_f32_rn_noftz(mul_f32_rn_noftz(rounded_8a_0, 7.888609e-31f), 1.7763568e-15f) : t1); + float2 _f32x2_mul_a_16 = mul_f32x2_noftz(cvt_f32, make_float2(max_noftz(u0, 0.0f), max_noftz(u1, 0.0f))); + float2 mul; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul) : "l"(*(const unsigned long long*)&_f32x2_mul_a_16), "l"(*(const unsigned long long*)&cvt_f32_1)); + unsigned int band_pair_out = __as_u32(__float22bfloat162_rn(make_float2(mul.x, mul.y))); + return band_pair_out; +} + +extern "C" { + +__global__ __launch_bounds__(256, LAUNCH_MIN_BLOCKS) void +kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950(__nv_bfloat16* __restrict__ gated_input, uint8_t* __restrict__ row_output, uint8_t* __restrict__ col_output, uint8_t* __restrict__ row_scales, uint8_t* __restrict__ col_scales, int M, int K) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* tile = reinterpret_cast(smem_raw + 0); + const int tile_addr = smem + 0; + + // === Task calls (declaration order) === + int half = tid & 1; + int blk = tid >> 1 & 3; + int row = tid >> 3; + int bx = blockIdx.x; + int by = blockIdx.y; + int grow = by * 32 + row; + int col = bx * 128 + blk * 32 + half * 16; + int gate_index = grow * (2 * K) + col; + unsigned int words[8]; + { + uint32_t _iv_0_0; + uint32_t _iv_0_1; + uint32_t _iv_0_2; + uint32_t _iv_0_3; + uint32_t _iv_0_4; + uint32_t _iv_0_5; + uint32_t _iv_0_6; + uint32_t _iv_0_7; + asm volatile("ld.global.L2::cache_hint.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8], %9;" + : "=r"(_iv_0_0), "=r"(_iv_0_1), "=r"(_iv_0_2), "=r"(_iv_0_3), "=r"(_iv_0_4), "=r"(_iv_0_5), "=r"(_iv_0_6), "=r"(_iv_0_7) : "l"((const void*)(gated_input + gate_index + (0))), "l"(0x12F0000000000000ULL) : "memory"); + words[0 + 0] = (uint32_t)_iv_0_0; + words[0 + 1] = (uint32_t)_iv_0_1; + words[0 + 2] = (uint32_t)_iv_0_2; + words[0 + 3] = (uint32_t)_iv_0_3; + words[0 + 4] = (uint32_t)_iv_0_4; + words[0 + 5] = (uint32_t)_iv_0_5; + words[0 + 6] = (uint32_t)_iv_0_6; + words[0 + 7] = (uint32_t)_iv_0_7; + } + unsigned int words_0[8]; + { + uint32_t _iv_1_0; + uint32_t _iv_1_1; + uint32_t _iv_1_2; + uint32_t _iv_1_3; + uint32_t _iv_1_4; + uint32_t _iv_1_5; + uint32_t _iv_1_6; + uint32_t _iv_1_7; + asm volatile("ld.global.L2::cache_hint.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8], %9;" + : "=r"(_iv_1_0), "=r"(_iv_1_1), "=r"(_iv_1_2), "=r"(_iv_1_3), "=r"(_iv_1_4), "=r"(_iv_1_5), "=r"(_iv_1_6), "=r"(_iv_1_7) : "l"((const void*)(gated_input + (gate_index + K) + (0))), "l"(0x12F0000000000000ULL) : "memory"); + words_0[0 + 0] = (uint32_t)_iv_1_0; + words_0[0 + 1] = (uint32_t)_iv_1_1; + words_0[0 + 2] = (uint32_t)_iv_1_2; + words_0[0 + 3] = (uint32_t)_iv_1_3; + words_0[0 + 4] = (uint32_t)_iv_1_4; + words_0[0 + 5] = (uint32_t)_iv_1_5; + words_0[0 + 6] = (uint32_t)_iv_1_6; + words_0[0 + 7] = (uint32_t)_iv_1_7; + } + unsigned int packed[8]; + __nv_bfloat162 lowest = __as_bf16x2(words[0]); + lowest = __hmin2_nan(lowest, __as_bf16x2(words[1])); + lowest = __hmin2_nan(lowest, __as_bf16x2(words[2])); + lowest = __hmin2_nan(lowest, __as_bf16x2(words[3])); + lowest = __hmin2_nan(lowest, __as_bf16x2(words[4])); + lowest = __hmin2_nan(lowest, __as_bf16x2(words[5])); + lowest = __hmin2_nan(lowest, __as_bf16x2(words[6])); + lowest = __hmin2_nan(lowest, __as_bf16x2(words[7])); + float2 _cvt_f32_0 = cake_bf16x2_to_f32x2(lowest); + float _min_3; + asm("min.NaN.f32 %0, %1, %2;" : "=f"(_min_3) : "f"(_cvt_f32_0.x), "f"(_cvt_f32_0.y)); + int _warp_vote_2 = __all_sync(0xFFFFFFFF, _min_3 > -87.0f); + if (_warp_vote_2 != 0) { + float2 cvt_f32 = cake_bf16x2_to_f32x2(__as_bf16x2(words[0])); + float2 cvt_f32_1 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[0])); + float e0 = approx_exp2((-cvt_f32.x) * 1.442695f); + float e1 = approx_exp2((-cvt_f32.y) * 1.442695f); + float d0 = 1.0f + e0; + float d1 = 1.0f + e1; + float r0 = approx_rcp(d0); + float r1 = approx_rcp(d1); + float _fma_5; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_5) : "f"(-d0), "f"(r0), "f"(1.0f)); + float _fma_4; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_4) : "f"(r0), "f"(_fma_5), "f"(r0)); + float _fma_7; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_7) : "f"(-d1), "f"(r1), "f"(1.0f)); + float _fma_6; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_6) : "f"(r1), "f"(_fma_7), "f"(r1)); + float2 _f32x2_mul_a_8 = mul_f32x2_noftz(cvt_f32, make_float2(_fma_4, _fma_6)); + float2 mul; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul) : "l"(*(const unsigned long long*)&_f32x2_mul_a_8), "l"(*(const unsigned long long*)&cvt_f32_1)); + packed[0] = __as_u32(__float22bfloat162_rn(make_float2(mul.x, mul.y))); + float2 cvt_f32_0 = cake_bf16x2_to_f32x2(__as_bf16x2(words[1])); + float2 cvt_f32_1_0 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[1])); + float e0_0 = approx_exp2((-cvt_f32_0.x) * 1.442695f); + float e1_0 = approx_exp2((-cvt_f32_0.y) * 1.442695f); + float d0_0 = 1.0f + e0_0; + float d1_0 = 1.0f + e1_0; + float r0_0 = approx_rcp(d0_0); + float r1_0 = approx_rcp(d1_0); + float _fma_10; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_10) : "f"(-d0_0), "f"(r0_0), "f"(1.0f)); + float _fma_9; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_9) : "f"(r0_0), "f"(_fma_10), "f"(r0_0)); + float _fma_12; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_12) : "f"(-d1_0), "f"(r1_0), "f"(1.0f)); + float _fma_11; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_11) : "f"(r1_0), "f"(_fma_12), "f"(r1_0)); + float2 _f32x2_mul_a_13 = mul_f32x2_noftz(cvt_f32_0, make_float2(_fma_9, _fma_11)); + float2 mul_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_0) : "l"(*(const unsigned long long*)&_f32x2_mul_a_13), "l"(*(const unsigned long long*)&cvt_f32_1_0)); + packed[1] = __as_u32(__float22bfloat162_rn(make_float2(mul_0.x, mul_0.y))); + float2 cvt_f32_2 = cake_bf16x2_to_f32x2(__as_bf16x2(words[2])); + float2 cvt_f32_1_1 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[2])); + float e0_1 = approx_exp2((-cvt_f32_2.x) * 1.442695f); + float e1_1 = approx_exp2((-cvt_f32_2.y) * 1.442695f); + float d0_1 = 1.0f + e0_1; + float d1_1 = 1.0f + e1_1; + float r0_1 = approx_rcp(d0_1); + float r1_1 = approx_rcp(d1_1); + float _fma_15; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_15) : "f"(-d0_1), "f"(r0_1), "f"(1.0f)); + float _fma_14; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_14) : "f"(r0_1), "f"(_fma_15), "f"(r0_1)); + float _fma_17; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_17) : "f"(-d1_1), "f"(r1_1), "f"(1.0f)); + float _fma_16; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_16) : "f"(r1_1), "f"(_fma_17), "f"(r1_1)); + float2 _f32x2_mul_a_18 = mul_f32x2_noftz(cvt_f32_2, make_float2(_fma_14, _fma_16)); + float2 mul_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_1) : "l"(*(const unsigned long long*)&_f32x2_mul_a_18), "l"(*(const unsigned long long*)&cvt_f32_1_1)); + packed[2] = __as_u32(__float22bfloat162_rn(make_float2(mul_1.x, mul_1.y))); + float2 cvt_f32_3 = cake_bf16x2_to_f32x2(__as_bf16x2(words[3])); + float2 cvt_f32_1_2 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[3])); + float e0_2 = approx_exp2((-cvt_f32_3.x) * 1.442695f); + float e1_2 = approx_exp2((-cvt_f32_3.y) * 1.442695f); + float d0_2 = 1.0f + e0_2; + float d1_2 = 1.0f + e1_2; + float r0_2 = approx_rcp(d0_2); + float r1_2 = approx_rcp(d1_2); + float _fma_20; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_20) : "f"(-d0_2), "f"(r0_2), "f"(1.0f)); + float _fma_19; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_19) : "f"(r0_2), "f"(_fma_20), "f"(r0_2)); + float _fma_22; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_22) : "f"(-d1_2), "f"(r1_2), "f"(1.0f)); + float _fma_21; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_21) : "f"(r1_2), "f"(_fma_22), "f"(r1_2)); + float2 _f32x2_mul_a_23 = mul_f32x2_noftz(cvt_f32_3, make_float2(_fma_19, _fma_21)); + float2 mul_2; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_2) : "l"(*(const unsigned long long*)&_f32x2_mul_a_23), "l"(*(const unsigned long long*)&cvt_f32_1_2)); + packed[3] = __as_u32(__float22bfloat162_rn(make_float2(mul_2.x, mul_2.y))); + float2 cvt_f32_4 = cake_bf16x2_to_f32x2(__as_bf16x2(words[4])); + float2 cvt_f32_1_3 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[4])); + float e0_3 = approx_exp2((-cvt_f32_4.x) * 1.442695f); + float e1_3 = approx_exp2((-cvt_f32_4.y) * 1.442695f); + float d0_3 = 1.0f + e0_3; + float d1_3 = 1.0f + e1_3; + float r0_3 = approx_rcp(d0_3); + float r1_3 = approx_rcp(d1_3); + float _fma_25; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_25) : "f"(-d0_3), "f"(r0_3), "f"(1.0f)); + float _fma_24; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_24) : "f"(r0_3), "f"(_fma_25), "f"(r0_3)); + float _fma_27; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_27) : "f"(-d1_3), "f"(r1_3), "f"(1.0f)); + float _fma_26; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_26) : "f"(r1_3), "f"(_fma_27), "f"(r1_3)); + float2 _f32x2_mul_a_28 = mul_f32x2_noftz(cvt_f32_4, make_float2(_fma_24, _fma_26)); + float2 mul_3; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_3) : "l"(*(const unsigned long long*)&_f32x2_mul_a_28), "l"(*(const unsigned long long*)&cvt_f32_1_3)); + packed[4] = __as_u32(__float22bfloat162_rn(make_float2(mul_3.x, mul_3.y))); + float2 cvt_f32_5 = cake_bf16x2_to_f32x2(__as_bf16x2(words[5])); + float2 cvt_f32_1_4 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[5])); + float e0_4 = approx_exp2((-cvt_f32_5.x) * 1.442695f); + float e1_4 = approx_exp2((-cvt_f32_5.y) * 1.442695f); + float d0_4 = 1.0f + e0_4; + float d1_4 = 1.0f + e1_4; + float r0_4 = approx_rcp(d0_4); + float r1_4 = approx_rcp(d1_4); + float _fma_30; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_30) : "f"(-d0_4), "f"(r0_4), "f"(1.0f)); + float _fma_29; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_29) : "f"(r0_4), "f"(_fma_30), "f"(r0_4)); + float _fma_32; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_32) : "f"(-d1_4), "f"(r1_4), "f"(1.0f)); + float _fma_31; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_31) : "f"(r1_4), "f"(_fma_32), "f"(r1_4)); + float2 _f32x2_mul_a_33 = mul_f32x2_noftz(cvt_f32_5, make_float2(_fma_29, _fma_31)); + float2 mul_4; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_4) : "l"(*(const unsigned long long*)&_f32x2_mul_a_33), "l"(*(const unsigned long long*)&cvt_f32_1_4)); + packed[5] = __as_u32(__float22bfloat162_rn(make_float2(mul_4.x, mul_4.y))); + float2 cvt_f32_6 = cake_bf16x2_to_f32x2(__as_bf16x2(words[6])); + float2 cvt_f32_1_5 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[6])); + float e0_5 = approx_exp2((-cvt_f32_6.x) * 1.442695f); + float e1_5 = approx_exp2((-cvt_f32_6.y) * 1.442695f); + float d0_5 = 1.0f + e0_5; + float d1_5 = 1.0f + e1_5; + float r0_5 = approx_rcp(d0_5); + float r1_5 = approx_rcp(d1_5); + float _fma_35; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_35) : "f"(-d0_5), "f"(r0_5), "f"(1.0f)); + float _fma_34; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_34) : "f"(r0_5), "f"(_fma_35), "f"(r0_5)); + float _fma_37; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_37) : "f"(-d1_5), "f"(r1_5), "f"(1.0f)); + float _fma_36; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_36) : "f"(r1_5), "f"(_fma_37), "f"(r1_5)); + float2 _f32x2_mul_a_38 = mul_f32x2_noftz(cvt_f32_6, make_float2(_fma_34, _fma_36)); + float2 mul_5; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_5) : "l"(*(const unsigned long long*)&_f32x2_mul_a_38), "l"(*(const unsigned long long*)&cvt_f32_1_5)); + packed[6] = __as_u32(__float22bfloat162_rn(make_float2(mul_5.x, mul_5.y))); + float2 cvt_f32_7 = cake_bf16x2_to_f32x2(__as_bf16x2(words[7])); + float2 cvt_f32_1_6 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[7])); + float e0_6 = approx_exp2((-cvt_f32_7.x) * 1.442695f); + float e1_6 = approx_exp2((-cvt_f32_7.y) * 1.442695f); + float d0_6 = 1.0f + e0_6; + float d1_6 = 1.0f + e1_6; + float r0_6 = approx_rcp(d0_6); + float r1_6 = approx_rcp(d1_6); + float _fma_40; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_40) : "f"(-d0_6), "f"(r0_6), "f"(1.0f)); + float _fma_39; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_39) : "f"(r0_6), "f"(_fma_40), "f"(r0_6)); + float _fma_42; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_42) : "f"(-d1_6), "f"(r1_6), "f"(1.0f)); + float _fma_41; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_41) : "f"(r1_6), "f"(_fma_42), "f"(r1_6)); + float2 _f32x2_mul_a_43 = mul_f32x2_noftz(cvt_f32_7, make_float2(_fma_39, _fma_41)); + float2 mul_6; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_6) : "l"(*(const unsigned long long*)&_f32x2_mul_a_43), "l"(*(const unsigned long long*)&cvt_f32_1_6)); + packed[7] = __as_u32(__float22bfloat162_rn(make_float2(mul_6.x, mul_6.y))); + } else { + int band = 0; + float2 cvt_f32_8 = cake_bf16x2_to_f32x2(__as_bf16x2(words[0])); + float2 cvt_f32_1_7 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[0])); + float e0_7 = approx_exp2((-cvt_f32_8.x) * 1.442695f); + float e1_7 = approx_exp2((-cvt_f32_8.y) * 1.442695f); + float d0_7 = 1.0f + e0_7; + float d1_7 = 1.0f + e1_7; + float r0_7 = approx_rcp(d0_7); + float r1_7 = approx_rcp(d1_7); + float _fma_44; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_44) : "f"(-d0_7), "f"(r0_7), "f"(1.0f)); + float fma; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma) : "f"(r0_7), "f"(_fma_44), "f"(r0_7)); + float _fma_45; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_45) : "f"(-d1_7), "f"(r1_7), "f"(1.0f)); + float fma_1; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1) : "f"(r1_7), "f"(_fma_45), "f"(r1_7)); + float t0 = fma; + float t1 = fma_1; + band = band | ((fminf(t0, t1) == 0.0f) ? 1 : 0); + float2 _f32x2_mul_a_46 = mul_f32x2_noftz(cvt_f32_8, make_float2(max_noftz(t0, 0.0f), max_noftz(t1, 0.0f))); + float2 mul_7; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_7) : "l"(*(const unsigned long long*)&_f32x2_mul_a_46), "l"(*(const unsigned long long*)&cvt_f32_1_7)); + packed[0] = __as_u32(__float22bfloat162_rn(make_float2(mul_7.x, mul_7.y))); + float2 cvt_f32_8_1 = cake_bf16x2_to_f32x2(__as_bf16x2(words[1])); + float2 cvt_f32_1_7_1 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[1])); + float e0_7_1 = approx_exp2((-cvt_f32_8_1.x) * 1.442695f); + float e1_7_1 = approx_exp2((-cvt_f32_8_1.y) * 1.442695f); + float d0_7_1 = 1.0f + e0_7_1; + float d1_7_1 = 1.0f + e1_7_1; + float r0_7_1 = approx_rcp(d0_7_1); + float r1_7_1 = approx_rcp(d1_7_1); + float _fma_47; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_47) : "f"(-d0_7_1), "f"(r0_7_1), "f"(1.0f)); + float fma_0; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_0) : "f"(r0_7_1), "f"(_fma_47), "f"(r0_7_1)); + float _fma_48; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_48) : "f"(-d1_7_1), "f"(r1_7_1), "f"(1.0f)); + float fma_1_0; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1_0) : "f"(r1_7_1), "f"(_fma_48), "f"(r1_7_1)); + float t0_0 = fma_0; + float t1_0 = fma_1_0; + band = band | ((fminf(t0_0, t1_0) == 0.0f) ? 2 : 0); + float2 _f32x2_mul_a_49 = mul_f32x2_noftz(cvt_f32_8_1, make_float2(max_noftz(t0_0, 0.0f), max_noftz(t1_0, 0.0f))); + float2 mul_7_1; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_7_1) : "l"(*(const unsigned long long*)&_f32x2_mul_a_49), "l"(*(const unsigned long long*)&cvt_f32_1_7_1)); + packed[1] = __as_u32(__float22bfloat162_rn(make_float2(mul_7_1.x, mul_7_1.y))); + float2 cvt_f32_9 = cake_bf16x2_to_f32x2(__as_bf16x2(words[2])); + float2 cvt_f32_1_8 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[2])); + float e0_8 = approx_exp2((-cvt_f32_9.x) * 1.442695f); + float e1_8 = approx_exp2((-cvt_f32_9.y) * 1.442695f); + float d0_8 = 1.0f + e0_8; + float d1_8 = 1.0f + e1_8; + float r0_8 = approx_rcp(d0_8); + float r1_8 = approx_rcp(d1_8); + float _fma_50; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_50) : "f"(-d0_8), "f"(r0_8), "f"(1.0f)); + float fma_2; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_2) : "f"(r0_8), "f"(_fma_50), "f"(r0_8)); + float _fma_51; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_51) : "f"(-d1_8), "f"(r1_8), "f"(1.0f)); + float fma_1_1; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1_1) : "f"(r1_8), "f"(_fma_51), "f"(r1_8)); + float t0_1 = fma_2; + float t1_1 = fma_1_1; + band = band | ((fminf(t0_1, t1_1) == 0.0f) ? 4 : 0); + float2 _f32x2_mul_a_52 = mul_f32x2_noftz(cvt_f32_9, make_float2(max_noftz(t0_1, 0.0f), max_noftz(t1_1, 0.0f))); + float2 mul_8; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_8) : "l"(*(const unsigned long long*)&_f32x2_mul_a_52), "l"(*(const unsigned long long*)&cvt_f32_1_8)); + packed[2] = __as_u32(__float22bfloat162_rn(make_float2(mul_8.x, mul_8.y))); + float2 cvt_f32_10 = cake_bf16x2_to_f32x2(__as_bf16x2(words[3])); + float2 cvt_f32_1_9 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[3])); + float e0_9 = approx_exp2((-cvt_f32_10.x) * 1.442695f); + float e1_9 = approx_exp2((-cvt_f32_10.y) * 1.442695f); + float d0_9 = 1.0f + e0_9; + float d1_9 = 1.0f + e1_9; + float r0_9 = approx_rcp(d0_9); + float r1_9 = approx_rcp(d1_9); + float _fma_53; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_53) : "f"(-d0_9), "f"(r0_9), "f"(1.0f)); + float fma_3; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_3) : "f"(r0_9), "f"(_fma_53), "f"(r0_9)); + float _fma_54; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_54) : "f"(-d1_9), "f"(r1_9), "f"(1.0f)); + float fma_1_2; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1_2) : "f"(r1_9), "f"(_fma_54), "f"(r1_9)); + float t0_2 = fma_3; + float t1_2 = fma_1_2; + band = band | ((fminf(t0_2, t1_2) == 0.0f) ? 8 : 0); + float2 _f32x2_mul_a_55 = mul_f32x2_noftz(cvt_f32_10, make_float2(max_noftz(t0_2, 0.0f), max_noftz(t1_2, 0.0f))); + float2 mul_9; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_9) : "l"(*(const unsigned long long*)&_f32x2_mul_a_55), "l"(*(const unsigned long long*)&cvt_f32_1_9)); + packed[3] = __as_u32(__float22bfloat162_rn(make_float2(mul_9.x, mul_9.y))); + float2 cvt_f32_11 = cake_bf16x2_to_f32x2(__as_bf16x2(words[4])); + float2 cvt_f32_1_10 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[4])); + float e0_10 = approx_exp2((-cvt_f32_11.x) * 1.442695f); + float e1_10 = approx_exp2((-cvt_f32_11.y) * 1.442695f); + float d0_10 = 1.0f + e0_10; + float d1_10 = 1.0f + e1_10; + float r0_10 = approx_rcp(d0_10); + float r1_10 = approx_rcp(d1_10); + float _fma_56; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_56) : "f"(-d0_10), "f"(r0_10), "f"(1.0f)); + float fma_4; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_4) : "f"(r0_10), "f"(_fma_56), "f"(r0_10)); + float _fma_57; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_57) : "f"(-d1_10), "f"(r1_10), "f"(1.0f)); + float fma_1_3; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1_3) : "f"(r1_10), "f"(_fma_57), "f"(r1_10)); + float t0_3 = fma_4; + float t1_3 = fma_1_3; + band = band | ((fminf(t0_3, t1_3) == 0.0f) ? 16 : 0); + float2 _f32x2_mul_a_58 = mul_f32x2_noftz(cvt_f32_11, make_float2(max_noftz(t0_3, 0.0f), max_noftz(t1_3, 0.0f))); + float2 mul_10; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_10) : "l"(*(const unsigned long long*)&_f32x2_mul_a_58), "l"(*(const unsigned long long*)&cvt_f32_1_10)); + packed[4] = __as_u32(__float22bfloat162_rn(make_float2(mul_10.x, mul_10.y))); + float2 cvt_f32_12 = cake_bf16x2_to_f32x2(__as_bf16x2(words[5])); + float2 cvt_f32_1_11 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[5])); + float e0_11 = approx_exp2((-cvt_f32_12.x) * 1.442695f); + float e1_11 = approx_exp2((-cvt_f32_12.y) * 1.442695f); + float d0_11 = 1.0f + e0_11; + float d1_11 = 1.0f + e1_11; + float r0_11 = approx_rcp(d0_11); + float r1_11 = approx_rcp(d1_11); + float _fma_59; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_59) : "f"(-d0_11), "f"(r0_11), "f"(1.0f)); + float fma_5; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_5) : "f"(r0_11), "f"(_fma_59), "f"(r0_11)); + float _fma_60; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_60) : "f"(-d1_11), "f"(r1_11), "f"(1.0f)); + float fma_1_4; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1_4) : "f"(r1_11), "f"(_fma_60), "f"(r1_11)); + float t0_4 = fma_5; + float t1_4 = fma_1_4; + band = band | ((fminf(t0_4, t1_4) == 0.0f) ? 32 : 0); + float2 _f32x2_mul_a_61 = mul_f32x2_noftz(cvt_f32_12, make_float2(max_noftz(t0_4, 0.0f), max_noftz(t1_4, 0.0f))); + float2 mul_11; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_11) : "l"(*(const unsigned long long*)&_f32x2_mul_a_61), "l"(*(const unsigned long long*)&cvt_f32_1_11)); + packed[5] = __as_u32(__float22bfloat162_rn(make_float2(mul_11.x, mul_11.y))); + float2 cvt_f32_13 = cake_bf16x2_to_f32x2(__as_bf16x2(words[6])); + float2 cvt_f32_1_12 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[6])); + float e0_12 = approx_exp2((-cvt_f32_13.x) * 1.442695f); + float e1_12 = approx_exp2((-cvt_f32_13.y) * 1.442695f); + float d0_12 = 1.0f + e0_12; + float d1_12 = 1.0f + e1_12; + float r0_12 = approx_rcp(d0_12); + float r1_12 = approx_rcp(d1_12); + float _fma_62; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_62) : "f"(-d0_12), "f"(r0_12), "f"(1.0f)); + float fma_6; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_6) : "f"(r0_12), "f"(_fma_62), "f"(r0_12)); + float _fma_63; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_63) : "f"(-d1_12), "f"(r1_12), "f"(1.0f)); + float fma_1_5; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1_5) : "f"(r1_12), "f"(_fma_63), "f"(r1_12)); + float t0_5 = fma_6; + float t1_5 = fma_1_5; + band = band | ((fminf(t0_5, t1_5) == 0.0f) ? 64 : 0); + float2 _f32x2_mul_a_64 = mul_f32x2_noftz(cvt_f32_13, make_float2(max_noftz(t0_5, 0.0f), max_noftz(t1_5, 0.0f))); + float2 mul_12; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_12) : "l"(*(const unsigned long long*)&_f32x2_mul_a_64), "l"(*(const unsigned long long*)&cvt_f32_1_12)); + packed[6] = __as_u32(__float22bfloat162_rn(make_float2(mul_12.x, mul_12.y))); + float2 cvt_f32_14 = cake_bf16x2_to_f32x2(__as_bf16x2(words[7])); + float2 cvt_f32_1_13 = cake_bf16x2_to_f32x2(__as_bf16x2(words_0[7])); + float e0_13 = approx_exp2((-cvt_f32_14.x) * 1.442695f); + float e1_13 = approx_exp2((-cvt_f32_14.y) * 1.442695f); + float d0_13 = 1.0f + e0_13; + float d1_13 = 1.0f + e1_13; + float r0_13 = approx_rcp(d0_13); + float r1_13 = approx_rcp(d1_13); + float _fma_65; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_65) : "f"(-d0_13), "f"(r0_13), "f"(1.0f)); + float fma_7; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_7) : "f"(r0_13), "f"(_fma_65), "f"(r0_13)); + float _fma_66; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(_fma_66) : "f"(-d1_13), "f"(r1_13), "f"(1.0f)); + float fma_1_6; + asm("fma.rn.f32 %0, %1, %2, %3;" : "=f"(fma_1_6) : "f"(r1_13), "f"(_fma_66), "f"(r1_13)); + float t0_6 = fma_7; + float t1_6 = fma_1_6; + band = band | ((fminf(t0_6, t1_6) == 0.0f) ? 128 : 0); + float2 _f32x2_mul_a_67 = mul_f32x2_noftz(cvt_f32_14, make_float2(max_noftz(t0_6, 0.0f), max_noftz(t1_6, 0.0f))); + float2 mul_13; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&mul_13) : "l"(*(const unsigned long long*)&_f32x2_mul_a_67), "l"(*(const unsigned long long*)&cvt_f32_1_13)); + packed[7] = __as_u32(__float22bfloat162_rn(make_float2(mul_13.x, mul_13.y))); + if (band != 0) { + unsigned int popc = (unsigned int)__popc(band); + #pragma unroll 1 + for (int _j = 0; _j < popc; _j++) { + int _find_first_set_68 = __ffs(band); + int j = _find_first_set_68 - 1; + band = band & band - 1; + unsigned int pair = band_pair(gated_input, gate_index + 2 * j, K); + packed[0] = ((j == 0) ? pair : packed[0]); + packed[1] = ((j == 1) ? pair : packed[1]); + packed[2] = ((j == 2) ? pair : packed[2]); + packed[3] = ((j == 3) ? pair : packed[3]); + packed[4] = ((j == 4) ? pair : packed[4]); + packed[5] = ((j == 5) ? pair : packed[5]); + packed[6] = ((j == 6) ? pair : packed[6]); + packed[7] = ((j == 7) ? pair : packed[7]); + } + } + } + unsigned int swizzle = (unsigned int)(row >> 2 & 6 | row & 1); + unsigned int row_base = tile_addr + (unsigned int)(row * 256); + unsigned int chunk = (unsigned int)((tid & 7) * 2); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(row_base + ((chunk ^ swizzle) << 4)), "r"(packed[0]), "r"(packed[1]), "r"(packed[2]), "r"(packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(row_base + ((chunk + 1 ^ swizzle) << 4)), "r"(packed[4]), "r"(packed[5]), "r"(packed[6]), "r"(packed[7]) : "memory"); + __nv_bfloat162 amax = __as_bf16x2(0); + __nv_bfloat162 _hxa_70 = amax; + __nv_bfloat162 _hxb_70 = __as_bf16x2(packed[0]); + unsigned int _hxr_70; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_70) : "r"(*reinterpret_cast(&_hxa_70)), "r"(*reinterpret_cast(&_hxb_70))); + __nv_bfloat162 _max_69 = *reinterpret_cast(&_hxr_70); + amax = _max_69; + __nv_bfloat162 _hxa_72 = amax; + __nv_bfloat162 _hxb_72 = __as_bf16x2(packed[1]); + unsigned int _hxr_72; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_72) : "r"(*reinterpret_cast(&_hxa_72)), "r"(*reinterpret_cast(&_hxb_72))); + __nv_bfloat162 _max_71 = *reinterpret_cast(&_hxr_72); + amax = _max_71; + __nv_bfloat162 _hxa_74 = amax; + __nv_bfloat162 _hxb_74 = __as_bf16x2(packed[2]); + unsigned int _hxr_74; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_74) : "r"(*reinterpret_cast(&_hxa_74)), "r"(*reinterpret_cast(&_hxb_74))); + __nv_bfloat162 _max_73 = *reinterpret_cast(&_hxr_74); + amax = _max_73; + __nv_bfloat162 _hxa_76 = amax; + __nv_bfloat162 _hxb_76 = __as_bf16x2(packed[3]); + unsigned int _hxr_76; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_76) : "r"(*reinterpret_cast(&_hxa_76)), "r"(*reinterpret_cast(&_hxb_76))); + __nv_bfloat162 _max_75 = *reinterpret_cast(&_hxr_76); + amax = _max_75; + __nv_bfloat162 _hxa_78 = amax; + __nv_bfloat162 _hxb_78 = __as_bf16x2(packed[4]); + unsigned int _hxr_78; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_78) : "r"(*reinterpret_cast(&_hxa_78)), "r"(*reinterpret_cast(&_hxb_78))); + __nv_bfloat162 _max_77 = *reinterpret_cast(&_hxr_78); + amax = _max_77; + __nv_bfloat162 _hxa_80 = amax; + __nv_bfloat162 _hxb_80 = __as_bf16x2(packed[5]); + unsigned int _hxr_80; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_80) : "r"(*reinterpret_cast(&_hxa_80)), "r"(*reinterpret_cast(&_hxb_80))); + __nv_bfloat162 _max_79 = *reinterpret_cast(&_hxr_80); + amax = _max_79; + __nv_bfloat162 _hxa_82 = amax; + __nv_bfloat162 _hxb_82 = __as_bf16x2(packed[6]); + unsigned int _hxr_82; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_82) : "r"(*reinterpret_cast(&_hxa_82)), "r"(*reinterpret_cast(&_hxb_82))); + __nv_bfloat162 _max_81 = *reinterpret_cast(&_hxr_82); + amax = _max_81; + __nv_bfloat162 _hxa_84 = amax; + __nv_bfloat162 _hxb_84 = __as_bf16x2(packed[7]); + unsigned int _hxr_84; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_84) : "r"(*reinterpret_cast(&_hxa_84)), "r"(*reinterpret_cast(&_hxb_84))); + __nv_bfloat162 _max_83 = *reinterpret_cast(&_hxr_84); + amax = _max_83; + unsigned int mag = __as_u32(amax) & 2147450879; + unsigned int peer = __shfl_xor_sync(0xFFFFFFFF, mag, 1); + __nv_bfloat162 both = __hmax2_nan(__as_bf16x2(mag), __as_bf16x2(peer)); + unsigned int folded = __as_u32(__hmax2_nan(both, __as_bf16x2(__as_u32(both) >> 16))); + unsigned int amax_bits = (folded & 65535) << 16; + int scale_i32 = (int)(amax_bits + 2031616 >> 23) - 8; + if (scale_i32 < 0) { + scale_i32 = 0; + } + if ((amax_bits & 2139095040) == 2139095040) { + scale_i32 = 255; + } + unsigned int scale = (unsigned int)scale_i32; + unsigned int inv = 254 - scale << 7; + if (scale == 255) { + inv = 32704; + } + unsigned int inv2 = inv | inv << 16; + int scale_col = bx * 4 + blk; + int row_scale_index = ((grow >> 7) * (K >> 7) + (scale_col >> 2)) * 512 + (grow & 31) * 16 + (grow >> 5 & 3) * 4 + (scale_col & 3); + if (half == 0) { + *(reinterpret_cast(row_scales + row_scale_index) + (0)) = (unsigned char)(scale); + } + unsigned int row_words[4]; + uint32_t bf16x2_mul; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul) : "r"(packed[0]), "r"(inv2)); + uint16_t e4m3x2_bf16x2; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2) : "r"(bf16x2_mul)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1) : "r"(packed[1]), "r"(inv2)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_86; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_86) : "r"(bf16x2_mul_1)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_86) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_85; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_85) : "h"(e4m3x2_bf16x2), "h"(_cvt_satfinite_e4m3x2_bf16x2_86)); + row_words[0] = _mov_b32_pack_u16x2_85; + uint32_t bf16x2_mul_0; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_0) : "r"(packed[2]), "r"(inv2)); + uint16_t e4m3x2_bf16x2_0; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2_0) : "r"(bf16x2_mul_0)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_0); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2_0) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1_0; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1_0) : "r"(packed[3]), "r"(inv2)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_88; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_88) : "r"(bf16x2_mul_1_0)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1_0); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_88) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_87; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_87) : "h"(e4m3x2_bf16x2_0), "h"(_cvt_satfinite_e4m3x2_bf16x2_88)); + row_words[1] = _mov_b32_pack_u16x2_87; + uint32_t bf16x2_mul_2; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_2) : "r"(packed[4]), "r"(inv2)); + uint16_t e4m3x2_bf16x2_1; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2_1) : "r"(bf16x2_mul_2)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_2); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2_1) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1_1; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1_1) : "r"(packed[5]), "r"(inv2)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_90; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_90) : "r"(bf16x2_mul_1_1)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1_1); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_90) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_89; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_89) : "h"(e4m3x2_bf16x2_1), "h"(_cvt_satfinite_e4m3x2_bf16x2_90)); + row_words[2] = _mov_b32_pack_u16x2_89; + uint32_t bf16x2_mul_3; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_3) : "r"(packed[6]), "r"(inv2)); + uint16_t e4m3x2_bf16x2_2; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2_2) : "r"(bf16x2_mul_3)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_3); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2_2) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1_2; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1_2) : "r"(packed[7]), "r"(inv2)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_92; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_92) : "r"(bf16x2_mul_1_2)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1_2); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_92) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_91; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_91) : "h"(e4m3x2_bf16x2_2), "h"(_cvt_satfinite_e4m3x2_bf16x2_92)); + row_words[3] = _mov_b32_pack_u16x2_91; + reinterpret_cast(row_output + (grow * K + col))[0] = reinterpret_cast(row_words + 0)[0]; + __syncthreads(); + int lane_1 = tid & 31; + int warp_1 = tid >> 5; + int k8 = lane_1 & 7; + int srow = (k8 >> 1) * 8 + (lane_1 >> 3) * 2 + (k8 & 1); + unsigned int sswizzle = (unsigned int)(srow >> 2 & 6 | srow & 1); + unsigned int cols[8]; + unsigned int strip_chunk0_1_0 = (unsigned int)(warp_1 * 2); + unsigned int strip_src0_1_0 = tile_addr + (unsigned int)(srow * 256) + ((strip_chunk0_1_0 ^ sswizzle) << 4); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(cols[0]), "=r"(cols[1]), "=r"(cols[2]), "=r"(cols[3]) + : "r"(strip_src0_1_0) + : "memory"); + unsigned int strip_chunk0_1_1 = (unsigned int)(warp_1 * 2 + 1); + unsigned int strip_src0_1_1 = tile_addr + (unsigned int)(srow * 256) + ((strip_chunk0_1_1 ^ sswizzle) << 4); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(cols[4]), "=r"(cols[5]), "=r"(cols[6]), "=r"(cols[7]) + : "r"(strip_src0_1_1) + : "memory"); + __nv_bfloat162 amax_a = __as_bf16x2(0); + __nv_bfloat162 amax_b = __as_bf16x2(0); + __nv_bfloat162 _hxa_94 = amax_a; + __nv_bfloat162 _hxb_94 = __as_bf16x2(cols[0]); + unsigned int _hxr_94; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_94) : "r"(*reinterpret_cast(&_hxa_94)), "r"(*reinterpret_cast(&_hxb_94))); + __nv_bfloat162 _max_93 = *reinterpret_cast(&_hxr_94); + amax_a = _max_93; + __nv_bfloat162 _hxa_96 = amax_b; + __nv_bfloat162 _hxb_96 = __as_bf16x2(cols[4]); + unsigned int _hxr_96; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_96) : "r"(*reinterpret_cast(&_hxa_96)), "r"(*reinterpret_cast(&_hxb_96))); + __nv_bfloat162 _max_95 = *reinterpret_cast(&_hxr_96); + amax_b = _max_95; + __nv_bfloat162 _hxa_98 = amax_a; + __nv_bfloat162 _hxb_98 = __as_bf16x2(cols[1]); + unsigned int _hxr_98; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_98) : "r"(*reinterpret_cast(&_hxa_98)), "r"(*reinterpret_cast(&_hxb_98))); + __nv_bfloat162 _max_97 = *reinterpret_cast(&_hxr_98); + amax_a = _max_97; + __nv_bfloat162 _hxa_100 = amax_b; + __nv_bfloat162 _hxb_100 = __as_bf16x2(cols[5]); + unsigned int _hxr_100; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_100) : "r"(*reinterpret_cast(&_hxa_100)), "r"(*reinterpret_cast(&_hxb_100))); + __nv_bfloat162 _max_99 = *reinterpret_cast(&_hxr_100); + amax_b = _max_99; + __nv_bfloat162 _hxa_102 = amax_a; + __nv_bfloat162 _hxb_102 = __as_bf16x2(cols[2]); + unsigned int _hxr_102; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_102) : "r"(*reinterpret_cast(&_hxa_102)), "r"(*reinterpret_cast(&_hxb_102))); + __nv_bfloat162 _max_101 = *reinterpret_cast(&_hxr_102); + amax_a = _max_101; + __nv_bfloat162 _hxa_104 = amax_b; + __nv_bfloat162 _hxb_104 = __as_bf16x2(cols[6]); + unsigned int _hxr_104; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_104) : "r"(*reinterpret_cast(&_hxa_104)), "r"(*reinterpret_cast(&_hxb_104))); + __nv_bfloat162 _max_103 = *reinterpret_cast(&_hxr_104); + amax_b = _max_103; + __nv_bfloat162 _hxa_106 = amax_a; + __nv_bfloat162 _hxb_106 = __as_bf16x2(cols[3]); + unsigned int _hxr_106; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_106) : "r"(*reinterpret_cast(&_hxa_106)), "r"(*reinterpret_cast(&_hxb_106))); + __nv_bfloat162 _max_105 = *reinterpret_cast(&_hxr_106); + amax_a = _max_105; + __nv_bfloat162 _hxa_108 = amax_b; + __nv_bfloat162 _hxb_108 = __as_bf16x2(cols[7]); + unsigned int _hxr_108; + asm("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(_hxr_108) : "r"(*reinterpret_cast(&_hxa_108)), "r"(*reinterpret_cast(&_hxb_108))); + __nv_bfloat162 _max_107 = *reinterpret_cast(&_hxr_108); + amax_b = _max_107; + unsigned int mag_a = __as_u32(amax_a) & 2147450879; + unsigned int mag_b = __as_u32(amax_b) & 2147450879; + unsigned int lows_ab = mag_a & 65535 | mag_b << 16; + unsigned int highs_ab = mag_a >> 16 | mag_b & 4294901760u; + __nv_bfloat162 col_max = __hmax2_nan(__as_bf16x2(lows_ab), __as_bf16x2(highs_ab)); + unsigned int peer1 = __shfl_xor_sync(0xFFFFFFFF, __as_u32(col_max), 1); + col_max = __hmax2_nan(col_max, __as_bf16x2(peer1)); + unsigned int peer2 = __shfl_xor_sync(0xFFFFFFFF, __as_u32(col_max), 2); + col_max = __hmax2_nan(col_max, __as_bf16x2(peer2)); + unsigned int col_max_bits = __as_u32(col_max); + unsigned int amax_bits_0 = (col_max_bits & 65535) << 16; + int scale_i32_0 = (int)(amax_bits_0 + 2031616 >> 23) - 8; + if (scale_i32_0 < 0) { + scale_i32_0 = 0; + } + if ((amax_bits_0 & 2139095040) == 2139095040) { + scale_i32_0 = 255; + } + unsigned int scale_0 = (unsigned int)scale_i32_0; + unsigned int inv_0 = 254 - scale_0 << 7; + if (scale_0 == 255) { + inv_0 = 32704; + } + unsigned int inv2_0 = inv_0 | inv_0 << 16; + unsigned int amax_bits_1 = (col_max_bits >> 16 & 65535) << 16; + int scale_i32_1 = (int)(amax_bits_1 + 2031616 >> 23) - 8; + if (scale_i32_1 < 0) { + scale_i32_1 = 0; + } + if ((amax_bits_1 & 2139095040) == 2139095040) { + scale_i32_1 = 255; + } + unsigned int scale_1 = (unsigned int)scale_i32_1; + unsigned int inv_1 = 254 - scale_1 << 7; + if (scale_1 == 255) { + inv_1 = 32704; + } + unsigned int inv2_1 = inv_1 | inv_1 << 16; + int gcol_a = bx * 128 + warp_1 * 16 + (lane_1 >> 2); + int gcol_b = gcol_a + 8; + int out_row = by * 32 + (lane_1 & 3) * 8; + unsigned int words_a[2]; + unsigned int words_b[2]; + uint32_t bf16x2_mul_4; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_4) : "r"(cols[0]), "r"(inv2_0)); + uint16_t e4m3x2_bf16x2_3; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2_3) : "r"(bf16x2_mul_4)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_4); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2_3) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1_3; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1_3) : "r"(cols[1]), "r"(inv2_0)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_110; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_110) : "r"(bf16x2_mul_1_3)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1_3); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_110) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_109; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_109) : "h"(e4m3x2_bf16x2_3), "h"(_cvt_satfinite_e4m3x2_bf16x2_110)); + words_a[0] = _mov_b32_pack_u16x2_109; + uint32_t bf16x2_mul_5; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_5) : "r"(cols[2]), "r"(inv2_0)); + uint16_t e4m3x2_bf16x2_4; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2_4) : "r"(bf16x2_mul_5)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_5); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2_4) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1_4; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1_4) : "r"(cols[3]), "r"(inv2_0)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_112; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_112) : "r"(bf16x2_mul_1_4)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1_4); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_112) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_111; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_111) : "h"(e4m3x2_bf16x2_4), "h"(_cvt_satfinite_e4m3x2_bf16x2_112)); + words_a[1] = _mov_b32_pack_u16x2_111; + uint32_t bf16x2_mul_6; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_6) : "r"(cols[4]), "r"(inv2_1)); + uint16_t e4m3x2_bf16x2_5; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2_5) : "r"(bf16x2_mul_6)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_6); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2_5) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1_5; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1_5) : "r"(cols[5]), "r"(inv2_1)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_114; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_114) : "r"(bf16x2_mul_1_5)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1_5); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_114) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_113; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_113) : "h"(e4m3x2_bf16x2_5), "h"(_cvt_satfinite_e4m3x2_bf16x2_114)); + words_b[0] = _mov_b32_pack_u16x2_113; + uint32_t bf16x2_mul_7; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_7) : "r"(cols[6]), "r"(inv2_1)); + uint16_t e4m3x2_bf16x2_6; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(e4m3x2_bf16x2_6) : "r"(bf16x2_mul_7)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_7); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(e4m3x2_bf16x2_6) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t bf16x2_mul_1_6; + asm volatile("mul.rn.bf16x2 %0, %1, %2;" : "=r"(bf16x2_mul_1_6) : "r"(cols[7]), "r"(inv2_1)); + uint16_t _cvt_satfinite_e4m3x2_bf16x2_116; + #if !defined(CAKE_CVT_SATFINITE_E4M3X2_BF16X2_FALLBACK) && defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm("cvt.rn.satfinite.e4m3x2.bf16x2 %0, %1;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_116) : "r"(bf16x2_mul_1_6)); + #else + { + const uint32_t _bf16x2 = (uint32_t)(bf16x2_mul_1_6); + const float _lo = __uint_as_float(_bf16x2 << 16); + const float _hi = __uint_as_float(_bf16x2 & 0xFFFF0000u); + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(_cvt_satfinite_e4m3x2_bf16x2_116) : "f"(_hi), "f"(_lo)); + } + #endif + uint32_t _mov_b32_pack_u16x2_115; + asm("mov.b32 %0, {%1, %2};" : "=r"(_mov_b32_pack_u16x2_115) : "h"(e4m3x2_bf16x2_6), "h"(_cvt_satfinite_e4m3x2_bf16x2_116)); + words_b[1] = _mov_b32_pack_u16x2_115; + reinterpret_cast(col_output + (gcol_a * M + out_row))[0] = reinterpret_cast(words_a + 0)[0]; + reinterpret_cast(col_output + (gcol_b * M + out_row))[0] = reinterpret_cast(words_b + 0)[0]; + if ((lane_1 & 3) == 0) { + int row_tiles = M >> 7; + int col_scale_a = ((gcol_a >> 7) * row_tiles + (by >> 2)) * 512 + (gcol_a & 31) * 16 + (gcol_a >> 5 & 3) * 4 + (by & 3); + int col_scale_b = ((gcol_b >> 7) * row_tiles + (by >> 2)) * 512 + (gcol_b & 31) * 16 + (gcol_b >> 5 & 3) * 4 + (by & 3); + *(reinterpret_cast(col_scales + col_scale_a) + (0)) = (unsigned char)(scale_0); + *(reinterpret_cast(col_scales + col_scale_b) + (0)) = (unsigned char)(scale_1); + } +} + +} // extern "C"