From f2304d68c3bddbdc5666d97eeb270912f869086c Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 27 Mar 2026 23:24:07 +0100 Subject: [PATCH] feat(bf16): add warp shuffle + reduction BF16 wrappers to common header MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - bf16_shfl_xor() — warp shuffle XOR for __nv_bfloat16 - bf16_shfl_down() — warp shuffle down - bf16_warp_sum() — full warp-level sum reduction - bf16_warp_max() — full warp-level max reduction These hide the mandatory float cast (no native BF16 shuffle HW) behind a clean BF16 API. Kernels use bf16_warp_sum(val) instead of manual __bfloat162float → __shfl_xor_sync → __float2bfloat16. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../cuda_pipeline/common_device_functions.cuh | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 49e8f7aa1..e6b1f235a 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -56,6 +56,27 @@ __device__ __forceinline__ __nv_bfloat16 bf16(float x) { return __float2bfloat16(x); } +/** BF16 warp shuffle — wraps __shfl_xor_sync for __nv_bfloat16. */ +__device__ __forceinline__ __nv_bfloat16 bf16_shfl_xor(unsigned mask, __nv_bfloat16 val, int offset) { + return __float2bfloat16(__shfl_xor_sync(mask, __bfloat162float(val), offset)); +} +/** BF16 warp shuffle down. */ +__device__ __forceinline__ __nv_bfloat16 bf16_shfl_down(unsigned mask, __nv_bfloat16 val, int offset) { + return __float2bfloat16(__shfl_down_sync(mask, __bfloat162float(val), offset)); +} +/** BF16 warp-level sum reduction (16 lanes). */ +__device__ __forceinline__ __nv_bfloat16 bf16_warp_sum(__nv_bfloat16 val) { + for (int offset = 16; offset > 0; offset >>= 1) + val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset); + return val; +} +/** BF16 warp-level max reduction. */ +__device__ __forceinline__ __nv_bfloat16 bf16_warp_max(__nv_bfloat16 val) { + for (int offset = 16; offset > 0; offset >>= 1) + val = bf16_fmax(val, bf16_shfl_xor(0xFFFFFFFF, val, offset)); + return val; +} + /** BF16 leaky ReLU — operates on raw bf16, returns bf16. */ __device__ __forceinline__ __nv_bfloat16 leaky_relu_bf16(__nv_bfloat16 x) { return (x > bf16_zero()) ? x : bf16(0.01f) * x;