From 317fa852675325bc9ada7320fc536a2807d6a8fb Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 20:27:03 -0700 Subject: [PATCH 01/14] Optimize Blackwell softmax with Loom kernels --- csrc/blackwell_softmax.cu | 169 ++++++ flashinfer/aot.py | 3 + flashinfer/jit/blackwell_softmax.py | 29 + flashinfer/sampling.py | 51 +- include/flashinfer/blackwell_softmax.cuh | 725 +++++++++++++++++++++++ tests/utils/test_sampling.py | 19 + 6 files changed, 987 insertions(+), 9 deletions(-) create mode 100644 csrc/blackwell_softmax.cu create mode 100644 flashinfer/jit/blackwell_softmax.py create mode 100644 include/flashinfer/blackwell_softmax.cuh diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu new file mode 100644 index 00000000000..a2a4992dc3f --- /dev/null +++ b/csrc/blackwell_softmax.cu @@ -0,0 +1,169 @@ +/* + * 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. + */ +#include +#include + +#include +#include +#include + +#include "tvm_ffi_utils.h" + +using namespace flashinfer; +using tvm::ffi::Optional; + +namespace { + +constexpr int kBootstrapThreads = 256; +constexpr int kRowwiseThreads = 512; +constexpr int kMaxSplits = 64; +constexpr size_t kDynamicSmemBytes = 128; + +enum class ParameterKind : int { + kNone = 0, + kScalar = 1, + kPerRow = 2, +}; + +bool use_rowwise_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind) { + const bool small_low_row = rows <= 32 && vocab_size <= 16384; + const bool dense_aligned_mid_row = + rows > 128 && rows <= 384 && vocab_size >= 24576 && vocab_size <= 256000 && + vocab_size % 4 == 0 && parameter_kind == ParameterKind::kNone; + const bool dense_aligned_high_row_narrow = + rows > 384 && rows <= 1024 && vocab_size >= 24576 && vocab_size <= 32000 && + vocab_size % 4 == 0 && parameter_kind == ParameterKind::kNone; + const bool measured_large_odd = rows > 128 && rows <= 512 && vocab_size >= 24576 && + vocab_size <= 131072 && vocab_size % 4 != 0; + return small_low_row || dense_aligned_mid_row || dense_aligned_high_row_narrow || + measured_large_odd; +} + +cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temperature_arr, + float temperature_val, ParameterKind parameter_kind, + uint32_t rows, uint32_t vocab_size, void* workspace, + size_t workspace_bytes, cudaStream_t stream) { + if (rows == 0 || vocab_size == 0 || + static_cast(rows) * vocab_size > + static_cast(std::numeric_limits::max())) { + return cudaErrorNotSupported; + } + + float* parameter = temperature_arr != nullptr ? temperature_arr : logits; + int rows_i = static_cast(rows); + int vocab_size_i = static_cast(vocab_size); + int parameter_kind_i = static_cast(parameter_kind); + + if (use_rowwise_kernel(rows, vocab_size, parameter_kind)) { + void* args[] = {&logits, ¶meter, &output, &rows_i, + &vocab_size_i, ¶meter_kind_i, &temperature_val}; + return cudaLaunchKernel( + reinterpret_cast(kernel_flashinfer_blackwell_softmax_followup_rowwise), + dim3(rows), dim3(kRowwiseThreads), args, kDynamicSmemBytes, stream); + } + + int active_blocks_per_sm = 0; + cudaError_t status = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &active_blocks_per_sm, kernel_flashinfer_blackwell_softmax_bootstrap_seed, + kBootstrapThreads, kDynamicSmemBytes); + if (status != cudaSuccess) { + return status; + } + if (active_blocks_per_sm <= 0) { + return cudaErrorNotSupported; + } + + int sm_count = 0; + int device = 0; + if ((status = cudaGetDevice(&device)) != cudaSuccess || + (status = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device)) != + cudaSuccess) { + return status; + } + + const uint64_t tiles_per_row = ceil_div(static_cast(vocab_size), + static_cast(kBootstrapThreads)); + const uint64_t cooperative_capacity = + static_cast(active_blocks_per_sm) * static_cast(sm_count); + const uint64_t provisional_grid = + std::max(1, std::min(static_cast(rows) * tiles_per_row, + cooperative_capacity)); + const uint64_t split_capacity = + provisional_grid >= rows ? provisional_grid / static_cast(rows) : 1; + const int splits = static_cast( + std::max(1, std::min({tiles_per_row, kMaxSplits, split_capacity}))); + const uint64_t grid = + std::max(1, std::min(provisional_grid, + static_cast(rows) * splits)); + const size_t partial_count = static_cast(rows) * static_cast(splits); + const size_t required_workspace_bytes = partial_count * 2 * sizeof(float); + if (required_workspace_bytes > workspace_bytes) { + return cudaErrorNotSupported; + } + + auto* partial_max = static_cast(workspace); + auto* partial_sum = partial_max + partial_count; + int splits_i = splits; + void* args[] = {&logits, ¶meter, &output, &partial_max, + &partial_sum, &rows_i, &vocab_size_i, &splits_i, + ¶meter_kind_i, &temperature_val}; + return cudaLaunchCooperativeKernel( + reinterpret_cast(kernel_flashinfer_blackwell_softmax_bootstrap_seed), + dim3(static_cast(grid)), dim3(kBootstrapThreads), args, kDynamicSmemBytes, + stream); +} + +} // namespace + +void blackwell_softmax(TensorView workspace_buffer, TensorView logits, TensorView output, + Optional maybe_temperature_arr, double temperature_val, + bool enable_pdl, bool temperature_is_none) { + CHECK_INPUT(workspace_buffer); + CHECK_INPUT(logits); + CHECK_INPUT(output); + CHECK_DIM(2, logits); + + const auto rows = static_cast(logits.size(0)); + const auto vocab_size = static_cast(logits.size(1)); + const bool has_temperature_arr = maybe_temperature_arr.has_value(); + const ParameterKind parameter_kind = + has_temperature_arr ? ParameterKind::kPerRow + : (temperature_is_none ? ParameterKind::kNone : ParameterKind::kScalar); + + ffi::CUDADeviceGuard device_guard(logits.device().device_id); + auto stream = get_stream(logits.device()); + auto* logits_ptr = static_cast(logits.data_ptr()); + auto* output_ptr = static_cast(output.data_ptr()); + auto* temperature_ptr = + has_temperature_arr ? static_cast(maybe_temperature_arr.value().data_ptr()) : nullptr; + const size_t workspace_bytes = + get_element_size(workspace_buffer) * workspace_buffer.size(0); + + cudaError_t status = + launch_blackwell_softmax(logits_ptr, output_ptr, temperature_ptr, + static_cast(temperature_val), parameter_kind, rows, + vocab_size, workspace_buffer.data_ptr(), workspace_bytes, stream); + if (status == cudaErrorNotSupported) { + status = sampling::OnlineSoftmax( + logits_ptr, output_ptr, rows, vocab_size, temperature_ptr, + static_cast(temperature_val), workspace_buffer.data_ptr(), workspace_bytes, + enable_pdl, stream); + } + TVM_FFI_ICHECK(status == cudaSuccess) + << "Blackwell Softmax failed with error code " << cudaGetErrorString(status); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(softmax, blackwell_softmax); diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 04e71dd11af..06803cbf2f9 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -108,6 +108,7 @@ from .jit.page import gen_page_module from .jit.quantization import gen_quantization_module from .jit.rope import gen_rope_module +from .jit.blackwell_softmax import gen_blackwell_softmax_module from .jit.sampling import gen_sampling_module from .jit.spdlog import gen_spdlog_module from .jit.moe_utils import gen_moe_utils_module @@ -634,6 +635,8 @@ def gen_all_modules( gen_sampling_module(), gen_topk_module(), ] + if has_sm100 or has_sm103: + jit_specs.append(gen_blackwell_softmax_module()) # Fused RMSNorm+SiLU: pre-compile all LUT configs (SM100+ only) if has_sm100: for C in _SUPPORTED_C: diff --git a/flashinfer/jit/blackwell_softmax.py b/flashinfer/jit/blackwell_softmax.py new file mode 100644 index 00000000000..de29d5fed53 --- /dev/null +++ b/flashinfer/jit/blackwell_softmax.py @@ -0,0 +1,29 @@ +""" +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. +""" + +from . import env as jit_env +from .core import JitSpec, current_compilation_context, gen_jit_spec + + +def gen_blackwell_softmax_module() -> JitSpec: + nvcc_flags = current_compilation_context.get_nvcc_flags_list( + supported_major_versions=[10] + ) + return gen_jit_spec( + "blackwell_softmax", + [jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax.cu"], + extra_cuda_cflags=nvcc_flags, + ) diff --git a/flashinfer/sampling.py b/flashinfer/sampling.py index 7e83574c650..923fb36000b 100644 --- a/flashinfer/sampling.py +++ b/flashinfer/sampling.py @@ -20,6 +20,7 @@ import torch from .api_logging import flashinfer_api +from .jit.blackwell_softmax import gen_blackwell_softmax_module from .jit.sampling import gen_sampling_module from .trace.templates.sampling import ( chain_speculative_sampling_trace, @@ -63,6 +64,17 @@ def get_seed_and_offset( return int(seed), int(offset) +@functools.cache +def get_blackwell_softmax_module(): + return gen_blackwell_softmax_module().build_and_load() + + +@functools.cache +def _supports_blackwell_softmax(device_index: int) -> bool: + major, minor = torch.cuda.get_device_capability(device_index) + return major == 10 and minor in (0, 3) + + @functools.cache def get_sampling_module(): module = gen_sampling_module().build_and_load() @@ -74,20 +86,35 @@ def softmax( maybe_temperature_arr: Optional[torch.Tensor], temperature_val: float, enable_pdl: bool, + temperature_is_none: bool, ) -> torch.Tensor: logits = logits.float() probs = torch.empty_like(logits, device=logits.device) maybe_temperature_arr = ( maybe_temperature_arr.float() if maybe_temperature_arr is not None else None ) - module.softmax( - workspace_buffer, - logits, - probs, - maybe_temperature_arr, - temperature_val, - enable_pdl, - ) + device_index = logits.device.index + if device_index is None: + device_index = torch.cuda.current_device() + if _supports_blackwell_softmax(device_index): + get_blackwell_softmax_module().softmax( + workspace_buffer, + logits, + probs, + maybe_temperature_arr, + temperature_val, + enable_pdl, + temperature_is_none, + ) + else: + module.softmax( + workspace_buffer, + logits, + probs, + maybe_temperature_arr, + temperature_val, + enable_pdl, + ) return probs @register_fake_op("flashinfer::softmax") @@ -97,6 +124,7 @@ def _fake_softmax( maybe_temperature_arr: Optional[torch.Tensor], temperature_val: float, enable_pdl: bool, + temperature_is_none: bool, ) -> torch.Tensor: return torch.empty_like(logits, device=logits.device, dtype=torch.float32) @@ -779,6 +807,7 @@ def softmax( [0.1724, 0.2719, 0.1991, 0.1465, 0.2101]], device='cuda:0') """ workspace_buffer = _get_cache_buf("softmax_workspace", 1024 * 1024, logits.device) + temperature_is_none = temperature is None if temperature is None: temperature = 1.0 @@ -787,7 +816,11 @@ def softmax( enable_pdl = device_support_pdl(logits.device) return get_sampling_module().softmax( - workspace_buffer, logits, *_to_tensor_scalar_tuple(temperature), enable_pdl + workspace_buffer, + logits, + *_to_tensor_scalar_tuple(temperature), + enable_pdl, + temperature_is_none, ) diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh new file mode 100644 index 00000000000..de38bfb21b3 --- /dev/null +++ b/include/flashinfer/blackwell_softmax.cuh @@ -0,0 +1,725 @@ +/* + * 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. + */ +#pragma once + +#include +#include +#include + +// Generated by Loom from Cake commit 25dba320359dc009daf7039d1778a8721681bba8. +// The sm_100a and sm_103a outputs are byte-identical (payload SHA-256: +// d6e46c2babcd74565ce80df3009e3cf2751f773c483b58f0d4676c41bf0dfd44). +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#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; +} + +#include + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + +__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; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_REDUCE_SMEM_OFF 0 +#define SMEM_REDUCE_SMEM_STAGE_BYTES 32 +#define SMEM_REDUCE_SMEM_STRIDE 32 +#define SMEM_TOTAL 128 +#define THREADS 256 + +extern "C" { + +__global__ __launch_bounds__(256, 2) void +kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, float* __restrict__ partial_max, float* __restrict__ partial_sum, int rows, int vocab_size, int splits, int parameter_kind, float scalar_temperature) +{ + 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; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* reduce_smem = reinterpret_cast(smem_raw + 0); + const int reduce_smem_addr = smem + 0; + + // === Task calls (dependency order) === + int total_tasks = rows * splits; + for (int task = bid; task < total_tasks; task += num_bids) { + int row = task / splits; + int split = task - row * splits; + float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); + int row_base = row * vocab_size; + int start_col = split * 256 + tid; + int col_stride = splits * 256; + int vector_groups = vocab_size / 8; + int first_group = split * vector_groups / splits; + int last_group = (split + 1) * vector_groups / splits; + float local_max = -CUDART_INF_F; + float local_sum = 0.0f; + if (row_base % 8 == 0) { + for (int group = first_group + tid; group < last_group; group += 256) { + int col = group * 8; + unsigned long long index = (unsigned long long)(row_base + col); + float _vec_load_0[8]; + { + unsigned _ldv8_0_0; + unsigned _ldv8_0_1; + unsigned _ldv8_0_2; + unsigned _ldv8_0_3; + unsigned _ldv8_0_4; + unsigned _ldv8_0_5; + unsigned _ldv8_0_6; + unsigned _ldv8_0_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_0_0), "=r"(_ldv8_0_1), "=r"(_ldv8_0_2), "=r"(_ldv8_0_3), "=r"(_ldv8_0_4), "=r"(_ldv8_0_5), "=r"(_ldv8_0_6), "=r"(_ldv8_0_7) : "l"((const void*)(x + (index))) : "memory"); + _vec_load_0[0 + 0] = __uint_as_float(_ldv8_0_0); + _vec_load_0[0 + 1] = __uint_as_float(_ldv8_0_1); + _vec_load_0[0 + 2] = __uint_as_float(_ldv8_0_2); + _vec_load_0[0 + 3] = __uint_as_float(_ldv8_0_3); + _vec_load_0[0 + 4] = __uint_as_float(_ldv8_0_4); + _vec_load_0[0 + 5] = __uint_as_float(_ldv8_0_5); + _vec_load_0[0 + 6] = __uint_as_float(_ldv8_0_6); + _vec_load_0[0 + 7] = __uint_as_float(_ldv8_0_7); + } + float group_max = -CUDART_INF_F; + for (int j = 0; j < 8; j++) { + float scaled = _vec_load_0[j] / temperature; + float _max_0 = max_noftz(group_max, scaled); + group_max = _max_0; + } + if (splits == 1) { + if (group_max > local_max) { + float _exp2_0 = approx_exp2((local_max - group_max) * 1.4426950408889634f); + local_sum *= _exp2_0; + local_max = group_max; + } + if (group_max > -CUDART_INF_F) { + for (int j_1 = 0; j_1 < 8; j_1++) { + float scaled_1 = _vec_load_0[j_1] / temperature; + if (scaled_1 > -CUDART_INF_F) { + float _exp2_1 = approx_exp2((scaled_1 - local_max) * 1.4426950408889634f); + local_sum += _exp2_1; + } + } + } + } else { + float _max_1 = max_noftz(local_max, group_max); + local_max = _max_1; + } + } + int tail_col = vector_groups * 8 + tid; + if (split == splits - 1 && tail_col < vocab_size) { + unsigned long long tail_index = (unsigned long long)(row_base + tail_col); + float tail_scaled = x[tail_index] / temperature; + if (splits == 1) { + if (tail_scaled > local_max) { + float _exp2_2 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); + local_sum *= _exp2_2; + local_max = tail_scaled; + } + if (tail_scaled > -CUDART_INF_F) { + float _exp2_3 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_3; + } + } else { + float _max_2 = max_noftz(local_max, tail_scaled); + local_max = _max_2; + } + } + } else { + for (int col_1 = start_col; col_1 < vocab_size; col_1 += col_stride) { + unsigned long long index_1 = (unsigned long long)(row_base + col_1); + float scaled_2 = x[index_1] / temperature; + if (splits == 1) { + if (scaled_2 > local_max) { + float _exp2_4 = approx_exp2((local_max - scaled_2) * 1.4426950408889634f); + local_sum *= _exp2_4; + local_max = scaled_2; + } + if (scaled_2 > -CUDART_INF_F) { + float _exp2_5 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); + local_sum += _exp2_5; + } + } else { + float _max_3 = max_noftz(local_max, scaled_2); + local_max = _max_3; + } + } + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float warp_max = _warp_reduce_0; + if (lane == 0) { + reduce_smem[warp] = warp_max; + } + __syncthreads(); + float warp_partial_max = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_1 = warp_partial_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); + float block_max = _warp_reduce_1; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = block_max; + } + } + __syncthreads(); + float cta_max = reduce_smem[0]; + __syncthreads(); + if (splits == 1) { + if (local_max > -CUDART_INF_F) { + float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); + local_sum *= _exp2_6; + } + } else if (cta_max > -CUDART_INF_F) { + if (row_base % 8 == 0) { + for (int group_1 = first_group + tid; group_1 < last_group; group_1 += 256) { + int col_2 = group_1 * 8; + unsigned long long index_2 = (unsigned long long)(row_base + col_2); + float _vec_load_1[8]; + { + unsigned _ldv8_1_0; + unsigned _ldv8_1_1; + unsigned _ldv8_1_2; + unsigned _ldv8_1_3; + unsigned _ldv8_1_4; + unsigned _ldv8_1_5; + unsigned _ldv8_1_6; + unsigned _ldv8_1_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) : "l"((const void*)(x + (index_2))) : "memory"); + _vec_load_1[0 + 0] = __uint_as_float(_ldv8_1_0); + _vec_load_1[0 + 1] = __uint_as_float(_ldv8_1_1); + _vec_load_1[0 + 2] = __uint_as_float(_ldv8_1_2); + _vec_load_1[0 + 3] = __uint_as_float(_ldv8_1_3); + _vec_load_1[0 + 4] = __uint_as_float(_ldv8_1_4); + _vec_load_1[0 + 5] = __uint_as_float(_ldv8_1_5); + _vec_load_1[0 + 6] = __uint_as_float(_ldv8_1_6); + _vec_load_1[0 + 7] = __uint_as_float(_ldv8_1_7); + } + for (int j_2 = 0; j_2 < 8; j_2++) { + float scaled_3 = _vec_load_1[j_2] / temperature; + float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_7; + } + } + int tail_col_1 = vector_groups * 8 + tid; + if (split == splits - 1 && tail_col_1 < vocab_size) { + unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col_1); + float tail_scaled_1 = x[tail_index_1] / temperature; + float _exp2_8 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_8; + } + } else { + for (int col_3 = start_col; col_3 < vocab_size; col_3 += col_stride) { + unsigned long long index_3 = (unsigned long long)(row_base + col_3); + float scaled_4 = x[index_3] / temperature; + float _exp2_9 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_9; + } + } + } + float _warp_reduce_2 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); + float warp_sum = _warp_reduce_2; + if (lane == 0) { + reduce_smem[warp] = warp_sum; + } + __syncthreads(); + float warp_partial_sum = ((lane < 8) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_3 = warp_partial_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); + float block_sum = _warp_reduce_3; + __syncthreads(); + if (splits == 1) { + float _rcp_0 = approx_rcp(block_sum); + float inv_sum = _rcp_0; + if (row_base % 8 == 0) { + for (int group_2 = first_group + tid; group_2 < last_group; group_2 += 256) { + int col_4 = group_2 * 8; + unsigned long long index_4 = (unsigned long long)(row_base + col_4); + float _vec_load_2[8]; + { + unsigned _ldv8_2_0; + unsigned _ldv8_2_1; + unsigned _ldv8_2_2; + unsigned _ldv8_2_3; + unsigned _ldv8_2_4; + unsigned _ldv8_2_5; + unsigned _ldv8_2_6; + unsigned _ldv8_2_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) : "l"((const void*)(x + (index_4))) : "memory"); + _vec_load_2[0 + 0] = __uint_as_float(_ldv8_2_0); + _vec_load_2[0 + 1] = __uint_as_float(_ldv8_2_1); + _vec_load_2[0 + 2] = __uint_as_float(_ldv8_2_2); + _vec_load_2[0 + 3] = __uint_as_float(_ldv8_2_3); + _vec_load_2[0 + 4] = __uint_as_float(_ldv8_2_4); + _vec_load_2[0 + 5] = __uint_as_float(_ldv8_2_5); + _vec_load_2[0 + 6] = __uint_as_float(_ldv8_2_6); + _vec_load_2[0 + 7] = __uint_as_float(_ldv8_2_7); + } + for (int j_3 = 0; j_3 < 8; j_3++) { + float scaled_5 = _vec_load_2[j_3] / temperature; + float _exp2_10 = approx_exp2((scaled_5 - cta_max) * 1.4426950408889634f); + _vec_load_2[j_3] = _exp2_10 * inv_sum; + } + { + unsigned _stv8_3_0 = __float_as_uint(_vec_load_2[0 + 0]); + unsigned _stv8_3_1 = __float_as_uint(_vec_load_2[0 + 1]); + unsigned _stv8_3_2 = __float_as_uint(_vec_load_2[0 + 2]); + unsigned _stv8_3_3 = __float_as_uint(_vec_load_2[0 + 3]); + unsigned _stv8_3_4 = __float_as_uint(_vec_load_2[0 + 4]); + unsigned _stv8_3_5 = __float_as_uint(_vec_load_2[0 + 5]); + unsigned _stv8_3_6 = __float_as_uint(_vec_load_2[0 + 6]); + unsigned _stv8_3_7 = __float_as_uint(_vec_load_2[0 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(output + (index_4))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); + } + } + int tail_col_2 = vector_groups * 8 + tid; + if (tail_col_2 < vocab_size) { + unsigned long long tail_index_2 = (unsigned long long)(row_base + tail_col_2); + float tail_scaled_2 = x[tail_index_2] / temperature; + float _exp2_11 = approx_exp2((tail_scaled_2 - cta_max) * 1.4426950408889634f); + output[tail_index_2] = _exp2_11 * inv_sum; + } + } else { + for (int col_5 = start_col; col_5 < vocab_size; col_5 += col_stride) { + unsigned long long index_5 = (unsigned long long)(row_base + col_5); + float scaled_6 = x[index_5] / temperature; + float _exp2_12 = approx_exp2((scaled_6 - cta_max) * 1.4426950408889634f); + output[index_5] = _exp2_12 * inv_sum; + } + } + __syncthreads(); + } else if (warp == 0) { + if (elect_sync()) { + unsigned long long partial_index = (unsigned long long)task; + partial_max[partial_index] = cta_max; + partial_sum[partial_index] = block_sum; + } + } + } + if (splits > 1) { + __threadfence(); + cooperative_groups::this_grid().sync(); + } + int phase_2_tasks = ((splits > 1) ? total_tasks : 0); + for (int task_1 = bid; task_1 < phase_2_tasks; task_1 += num_bids) { + int row_1 = task_1 / splits; + int split_1 = task_1 - row_1 * splits; + float local_max_1 = -CUDART_INF_F; + for (int partial_split = tid; partial_split < splits; partial_split += 256) { + unsigned long long partial_index_1 = (unsigned long long)row_1 * (unsigned long long)splits + (unsigned long long)partial_split; + float _max_4 = max_noftz(local_max_1, partial_max[partial_index_1]); + local_max_1 = _max_4; + } + float _warp_reduce_4 = local_max_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_4 = max_noftz(_warp_reduce_4, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_4, offset)); + float warp_max_1 = _warp_reduce_4; + if (lane == 0) { + reduce_smem[warp] = warp_max_1; + } + __syncthreads(); + float warp_partial_max_1 = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_5 = warp_partial_max_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_5 = max_noftz(_warp_reduce_5, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_5, offset)); + float merged_max = _warp_reduce_5; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = merged_max; + } + } + __syncthreads(); + float global_max = reduce_smem[0]; + __syncthreads(); + float local_sum_1 = 0.0f; + for (int partial_split_1 = tid; partial_split_1 < splits; partial_split_1 += 256) { + unsigned long long partial_index_2 = (unsigned long long)row_1 * (unsigned long long)splits + (unsigned long long)partial_split_1; + float split_max = partial_max[partial_index_2]; + if (split_max > -CUDART_INF_F) { + float _exp2_13 = approx_exp2((split_max - global_max) * 1.4426950408889634f); + local_sum_1 += partial_sum[partial_index_2] * _exp2_13; + } + } + float _warp_reduce_6 = local_sum_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_6 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_6, offset); + float warp_sum_1 = _warp_reduce_6; + if (lane == 0) { + reduce_smem[warp] = warp_sum_1; + } + __syncthreads(); + float warp_partial_sum_1 = ((lane < 8) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_7 = warp_partial_sum_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_7 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_7, offset); + float merged_sum = _warp_reduce_7; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + float _rcp_1 = approx_rcp(merged_sum); + reduce_smem[0] = _rcp_1; + } + } + __syncthreads(); + float inv_sum_1 = reduce_smem[0]; + __syncthreads(); + float temperature_1 = ((parameter_kind == 2) ? parameter[row_1] : scalar_temperature); + int row_base_1 = row_1 * vocab_size; + int start_col_1 = split_1 * 256 + tid; + int col_stride_1 = splits * 256; + int vector_groups_1 = vocab_size / 8; + int first_group_1 = split_1 * vector_groups_1 / splits; + int last_group_1 = (split_1 + 1) * vector_groups_1 / splits; + if (row_base_1 % 8 == 0) { + for (int group_3 = first_group_1 + tid; group_3 < last_group_1; group_3 += 256) { + int col_6 = group_3 * 8; + unsigned long long index_6 = (unsigned long long)(row_base_1 + col_6); + float _vec_load_3[8]; + { + unsigned _ldv8_4_0; + unsigned _ldv8_4_1; + unsigned _ldv8_4_2; + unsigned _ldv8_4_3; + unsigned _ldv8_4_4; + unsigned _ldv8_4_5; + unsigned _ldv8_4_6; + unsigned _ldv8_4_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_4_0), "=r"(_ldv8_4_1), "=r"(_ldv8_4_2), "=r"(_ldv8_4_3), "=r"(_ldv8_4_4), "=r"(_ldv8_4_5), "=r"(_ldv8_4_6), "=r"(_ldv8_4_7) : "l"((const void*)(x + (index_6))) : "memory"); + _vec_load_3[0 + 0] = __uint_as_float(_ldv8_4_0); + _vec_load_3[0 + 1] = __uint_as_float(_ldv8_4_1); + _vec_load_3[0 + 2] = __uint_as_float(_ldv8_4_2); + _vec_load_3[0 + 3] = __uint_as_float(_ldv8_4_3); + _vec_load_3[0 + 4] = __uint_as_float(_ldv8_4_4); + _vec_load_3[0 + 5] = __uint_as_float(_ldv8_4_5); + _vec_load_3[0 + 6] = __uint_as_float(_ldv8_4_6); + _vec_load_3[0 + 7] = __uint_as_float(_ldv8_4_7); + } + for (int j_4 = 0; j_4 < 8; j_4++) { + float scaled_7 = _vec_load_3[j_4] / temperature_1; + float _exp2_14 = approx_exp2((scaled_7 - global_max) * 1.4426950408889634f); + _vec_load_3[j_4] = _exp2_14 * inv_sum_1; + } + { + unsigned _stv8_5_0 = __float_as_uint(_vec_load_3[0 + 0]); + unsigned _stv8_5_1 = __float_as_uint(_vec_load_3[0 + 1]); + unsigned _stv8_5_2 = __float_as_uint(_vec_load_3[0 + 2]); + unsigned _stv8_5_3 = __float_as_uint(_vec_load_3[0 + 3]); + unsigned _stv8_5_4 = __float_as_uint(_vec_load_3[0 + 4]); + unsigned _stv8_5_5 = __float_as_uint(_vec_load_3[0 + 5]); + unsigned _stv8_5_6 = __float_as_uint(_vec_load_3[0 + 6]); + unsigned _stv8_5_7 = __float_as_uint(_vec_load_3[0 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(output + (index_6))), "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3), "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) : "memory"); + } + } + int tail_col_3 = vector_groups_1 * 8 + tid; + if (split_1 == splits - 1 && tail_col_3 < vocab_size) { + unsigned long long tail_index_3 = (unsigned long long)(row_base_1 + tail_col_3); + float tail_scaled_3 = x[tail_index_3] / temperature_1; + float _exp2_15 = approx_exp2((tail_scaled_3 - global_max) * 1.4426950408889634f); + output[tail_index_3] = _exp2_15 * inv_sum_1; + } + } else { + for (int col_7 = start_col_1; col_7 < vocab_size; col_7 += col_stride_1) { + unsigned long long index_7 = (unsigned long long)(row_base_1 + col_7); + float scaled_8 = x[index_7] / temperature_1; + float _exp2_16 = approx_exp2((scaled_8 - global_max) * 1.4426950408889634f); + output[index_7] = _exp2_16 * inv_sum_1; + } + } + } +} + +} // extern "C" + +#undef LOOM_INF +#undef NUM_MAIN_STAGES +#undef SMEM_REDUCE_SMEM_OFF +#undef SMEM_REDUCE_SMEM_STAGE_BYTES +#undef SMEM_REDUCE_SMEM_STRIDE +#undef SMEM_TOTAL +#undef THREADS +#undef reduce_smem_addr + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_REDUCE_SMEM_OFF 0 +#define SMEM_REDUCE_SMEM_STAGE_BYTES 64 +#define SMEM_REDUCE_SMEM_STRIDE 64 +#define SMEM_TOTAL 128 +#define THREADS 512 + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_flashinfer_blackwell_softmax_followup_rowwise(float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, int rows, int vocab_size, int parameter_kind, float scalar_temperature) +{ + 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; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* reduce_smem = reinterpret_cast(smem_raw + 0); + const int reduce_smem_addr = smem + 0; + + // === Task calls (dependency order) === + int row = bid; + float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); + int row_base = row * vocab_size; + int row_misalignment = row_base % 4; + int aligned_begin = (4 - row_misalignment) % 4; + int _min_0 = ((aligned_begin) < (vocab_size) ? (aligned_begin) : (vocab_size)); + aligned_begin = _min_0; + int vector_groups = (vocab_size - aligned_begin) / 4; + int aligned_end = aligned_begin + vector_groups * 4; + float local_max = -CUDART_INF_F; + float local_sum = 0.0f; + for (int col = tid; col < aligned_begin; col += 512) { + unsigned long long index = (unsigned long long)(row_base + col); + float scaled = x[index] / temperature; + if (scaled > local_max) { + float _exp2_0 = approx_exp2((local_max - scaled) * 1.4426950408889634f); + local_sum *= _exp2_0; + local_max = scaled; + } + if (scaled > -CUDART_INF_F) { + float _exp2_1 = approx_exp2((scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_1; + } + } + for (int group = tid; group < vector_groups; group += 512) { + int col_1 = aligned_begin + group * 4; + unsigned long long index_1 = (unsigned long long)(row_base + col_1); + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(x + index_1); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float group_max = -CUDART_INF_F; + for (int j = 0; j < 4; j++) { + float scaled_1 = _vec_load_0[j] / temperature; + float _max_0 = max_noftz(group_max, scaled_1); + group_max = _max_0; + } + if (group_max > local_max) { + float _exp2_2 = approx_exp2((local_max - group_max) * 1.4426950408889634f); + local_sum *= _exp2_2; + local_max = group_max; + } + if (group_max > -CUDART_INF_F) { + for (int j_1 = 0; j_1 < 4; j_1++) { + float scaled_2 = _vec_load_0[j_1] / temperature; + if (scaled_2 > -CUDART_INF_F) { + float _exp2_3 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); + local_sum += _exp2_3; + } + } + } + } + int tail_col = aligned_end + tid; + if (tail_col < vocab_size) { + unsigned long long tail_index = (unsigned long long)(row_base + tail_col); + float tail_scaled = x[tail_index] / temperature; + if (tail_scaled > local_max) { + float _exp2_4 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); + local_sum *= _exp2_4; + local_max = tail_scaled; + } + if (tail_scaled > -CUDART_INF_F) { + float _exp2_5 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_5; + } + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float warp_max = _warp_reduce_0; + if (lane == 0) { + reduce_smem[warp] = warp_max; + } + __syncthreads(); + float warp_partial_max = ((lane < 16) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_1 = warp_partial_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); + float block_max = _warp_reduce_1; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = block_max; + } + } + __syncthreads(); + float cta_max = reduce_smem[0]; + __syncthreads(); + if (local_max > -CUDART_INF_F) { + float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); + local_sum *= _exp2_6; + } + float _warp_reduce_2 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); + float warp_sum = _warp_reduce_2; + if (lane == 0) { + reduce_smem[warp] = warp_sum; + } + __syncthreads(); + float warp_partial_sum = ((lane < 16) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_3 = warp_partial_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); + float block_sum = _warp_reduce_3; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + float _rcp_0 = approx_rcp(block_sum); + reduce_smem[0] = _rcp_0; + } + } + __syncthreads(); + float inv_sum = reduce_smem[0]; + for (int col_2 = tid; col_2 < aligned_begin; col_2 += 512) { + unsigned long long index_2 = (unsigned long long)(row_base + col_2); + float scaled_3 = x[index_2] / temperature; + float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); + output[index_2] = _exp2_7 * inv_sum; + } + for (int group_1 = tid; group_1 < vector_groups; group_1 += 512) { + int col_3 = aligned_begin + group_1 * 4; + unsigned long long index_3 = (unsigned long long)(row_base + col_3); + float _vec_load_1[4]; + { + float4 _v4 = *reinterpret_cast(x + index_3); + _vec_load_1[0 + 0] = _v4.x; + _vec_load_1[0 + 1] = _v4.y; + _vec_load_1[0 + 2] = _v4.z; + _vec_load_1[0 + 3] = _v4.w; + } + for (int j_2 = 0; j_2 < 4; j_2++) { + float scaled_4 = _vec_load_1[j_2] / temperature; + float _exp2_8 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); + _vec_load_1[j_2] = _exp2_8 * inv_sum; + } + { + float4 _v4 = make_float4(_vec_load_1[0 + 0], _vec_load_1[0 + 1], _vec_load_1[0 + 2], _vec_load_1[0 + 3]); + *reinterpret_cast(output + index_3) = _v4; + } + } + if (tail_col < vocab_size) { + unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col); + float tail_scaled_1 = x[tail_index_1] / temperature; + float _exp2_9 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); + output[tail_index_1] = _exp2_9 * inv_sum; + } +} + +} // extern "C" + +#undef LOOM_INF +#undef NUM_MAIN_STAGES +#undef SMEM_REDUCE_SMEM_OFF +#undef SMEM_REDUCE_SMEM_STAGE_BYTES +#undef SMEM_REDUCE_SMEM_STRIDE +#undef SMEM_TOTAL +#undef THREADS +#undef reduce_smem_addr diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index 4d1fc548960..2f55bbee0f5 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -76,6 +76,25 @@ def test_softmax( assert torch.allclose(probs, probs_ref, atol=1e-5) +@pytest.mark.parametrize( + "batch_size,vocab_size", + [ + (129, 24576), # aligned mid-row route + (385, 24577), # odd-stride route + ], +) +def test_softmax_blackwell_rowwise_routes(batch_size, vocab_size): + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("Loom Softmax routes require SM100 or SM103") + + torch.manual_seed(42) + logits = torch.randn(batch_size, vocab_size, device="cuda") + probs = flashinfer.sampling.softmax(logits, temperature=None, enable_pdl=False) + probs_ref = torch.softmax(logits, dim=-1) + + assert torch.allclose(probs, probs_ref, atol=1e-5) + + @pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) @pytest.mark.parametrize( "distribution", From 6e025b138e83e2ccdabd1f64303eaf83e7844cad Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 20:39:01 -0700 Subject: [PATCH 02/14] Use CUDA types in generated Blackwell header --- include/flashinfer/blackwell_softmax.cuh | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index de38bfb21b3..42590526b79 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -17,21 +17,15 @@ #include #include +#include #include +#include +#include + // Generated by Loom from Cake commit 25dba320359dc009daf7039d1778a8721681bba8. // The sm_100a and sm_103a outputs are byte-identical (payload SHA-256: // d6e46c2babcd74565ce80df3009e3cf2751f773c483b58f0d4676c41bf0dfd44). -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; - -typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; - -#include __device__ __forceinline__ int make_warp_uniform(int x) { int result; From 210ec16856d18ae676ed0bbb8f141bba811abd7d Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 20:45:04 -0700 Subject: [PATCH 03/14] Keep Blackwell softmax dispatch separate --- flashinfer/sampling.py | 92 +++++++++++++++++++++++++----------- tests/utils/test_sampling.py | 24 +++++++--- 2 files changed, 82 insertions(+), 34 deletions(-) diff --git a/flashinfer/sampling.py b/flashinfer/sampling.py index 923fb36000b..cd722142f6f 100644 --- a/flashinfer/sampling.py +++ b/flashinfer/sampling.py @@ -65,8 +65,48 @@ def get_seed_and_offset( @functools.cache -def get_blackwell_softmax_module(): - return gen_blackwell_softmax_module().build_and_load() +def get_blackwell_softmax_op(): + module = gen_blackwell_softmax_module().build_and_load() + + @register_custom_op( + "flashinfer::blackwell_softmax", mutates_args=("workspace_buffer",) + ) + def blackwell_softmax( + workspace_buffer: torch.Tensor, + logits: torch.Tensor, + maybe_temperature_arr: Optional[torch.Tensor], + temperature_val: float, + enable_pdl: bool, + temperature_is_none: bool, + ) -> torch.Tensor: + logits = logits.float() + probs = torch.empty_like(logits, device=logits.device) + maybe_temperature_arr = ( + maybe_temperature_arr.float() if maybe_temperature_arr is not None else None + ) + module.softmax( + workspace_buffer, + logits, + probs, + maybe_temperature_arr, + temperature_val, + enable_pdl, + temperature_is_none, + ) + return probs + + @register_fake_op("flashinfer::blackwell_softmax") + def _fake_blackwell_softmax( + workspace_buffer: torch.Tensor, + logits: torch.Tensor, + maybe_temperature_arr: Optional[torch.Tensor], + temperature_val: float, + enable_pdl: bool, + temperature_is_none: bool, + ) -> torch.Tensor: + return torch.empty_like(logits, device=logits.device, dtype=torch.float32) + + return blackwell_softmax @functools.cache @@ -86,35 +126,20 @@ def softmax( maybe_temperature_arr: Optional[torch.Tensor], temperature_val: float, enable_pdl: bool, - temperature_is_none: bool, ) -> torch.Tensor: logits = logits.float() probs = torch.empty_like(logits, device=logits.device) maybe_temperature_arr = ( maybe_temperature_arr.float() if maybe_temperature_arr is not None else None ) - device_index = logits.device.index - if device_index is None: - device_index = torch.cuda.current_device() - if _supports_blackwell_softmax(device_index): - get_blackwell_softmax_module().softmax( - workspace_buffer, - logits, - probs, - maybe_temperature_arr, - temperature_val, - enable_pdl, - temperature_is_none, - ) - else: - module.softmax( - workspace_buffer, - logits, - probs, - maybe_temperature_arr, - temperature_val, - enable_pdl, - ) + module.softmax( + workspace_buffer, + logits, + probs, + maybe_temperature_arr, + temperature_val, + enable_pdl, + ) return probs @register_fake_op("flashinfer::softmax") @@ -124,7 +149,6 @@ def _fake_softmax( maybe_temperature_arr: Optional[torch.Tensor], temperature_val: float, enable_pdl: bool, - temperature_is_none: bool, ) -> torch.Tensor: return torch.empty_like(logits, device=logits.device, dtype=torch.float32) @@ -815,12 +839,24 @@ def softmax( if enable_pdl is None: enable_pdl = device_support_pdl(logits.device) + temperature_args = _to_tensor_scalar_tuple(temperature) + device_index = logits.device.index + if device_index is None: + device_index = torch.cuda.current_device() + if _supports_blackwell_softmax(device_index): + return get_blackwell_softmax_op()( + workspace_buffer, + logits, + *temperature_args, + enable_pdl, + temperature_is_none, + ) + return get_sampling_module().softmax( workspace_buffer, logits, - *_to_tensor_scalar_tuple(temperature), + *temperature_args, enable_pdl, - temperature_is_none, ) diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index 2f55bbee0f5..8e6516f8d5f 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -77,20 +77,32 @@ def test_softmax( @pytest.mark.parametrize( - "batch_size,vocab_size", + "batch_size,vocab_size,temperature_kind", [ - (129, 24576), # aligned mid-row route - (385, 24577), # odd-stride route + (1, 32000, "none"), # cooperative route + (256, 32000, "none"), # rowwise route + (1, 111, "scalar"), # scalar-temperature rowwise route + (1, 111, "per_row"), # per-row-temperature rowwise route ], ) -def test_softmax_blackwell_rowwise_routes(batch_size, vocab_size): +def test_softmax_blackwell_routes(batch_size, vocab_size, temperature_kind): if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): pytest.skip("Loom Softmax routes require SM100 or SM103") torch.manual_seed(42) logits = torch.randn(batch_size, vocab_size, device="cuda") - probs = flashinfer.sampling.softmax(logits, temperature=None, enable_pdl=False) - probs_ref = torch.softmax(logits, dim=-1) + if temperature_kind == "none": + temperature = None + probs_ref = torch.softmax(logits, dim=-1) + elif temperature_kind == "scalar": + temperature = 0.5 + probs_ref = torch.softmax(logits / temperature, dim=-1) + else: + temperature = torch.full((batch_size,), 0.5, device="cuda") + probs_ref = torch.softmax(logits / temperature[:, None], dim=-1) + probs = flashinfer.sampling.softmax( + logits, temperature=temperature, enable_pdl=False + ) assert torch.allclose(probs, probs_ref, atol=1e-5) From 6a8e024b28d844083f3ad183f3e8afae8ce146bd Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 20:56:41 -0700 Subject: [PATCH 04/14] Format Blackwell softmax sources --- csrc/blackwell_softmax.cu | 63 +- include/flashinfer/blackwell_softmax.cuh | 1204 +++++++++++----------- 2 files changed, 639 insertions(+), 628 deletions(-) diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu index a2a4992dc3f..b68401726af 100644 --- a/csrc/blackwell_softmax.cu +++ b/csrc/blackwell_softmax.cu @@ -13,11 +13,10 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include -#include - #include #include +#include +#include #include #include "tvm_ffi_utils.h" @@ -40,12 +39,12 @@ enum class ParameterKind : int { bool use_rowwise_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind) { const bool small_low_row = rows <= 32 && vocab_size <= 16384; - const bool dense_aligned_mid_row = - rows > 128 && rows <= 384 && vocab_size >= 24576 && vocab_size <= 256000 && - vocab_size % 4 == 0 && parameter_kind == ParameterKind::kNone; - const bool dense_aligned_high_row_narrow = - rows > 384 && rows <= 1024 && vocab_size >= 24576 && vocab_size <= 32000 && - vocab_size % 4 == 0 && parameter_kind == ParameterKind::kNone; + const bool dense_aligned_mid_row = rows > 128 && rows <= 384 && vocab_size >= 24576 && + vocab_size <= 256000 && vocab_size % 4 == 0 && + parameter_kind == ParameterKind::kNone; + const bool dense_aligned_high_row_narrow = rows > 384 && rows <= 1024 && vocab_size >= 24576 && + vocab_size <= 32000 && vocab_size % 4 == 0 && + parameter_kind == ParameterKind::kNone; const bool measured_large_odd = rows > 128 && rows <= 512 && vocab_size >= 24576 && vocab_size <= 131072 && vocab_size % 4 != 0; return small_low_row || dense_aligned_mid_row || dense_aligned_high_row_narrow || @@ -68,7 +67,7 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper int parameter_kind_i = static_cast(parameter_kind); if (use_rowwise_kernel(rows, vocab_size, parameter_kind)) { - void* args[] = {&logits, ¶meter, &output, &rows_i, + void* args[] = {&logits, ¶meter, &output, &rows_i, &vocab_size_i, ¶meter_kind_i, &temperature_val}; return cudaLaunchKernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_followup_rowwise), @@ -77,8 +76,8 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper int active_blocks_per_sm = 0; cudaError_t status = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &active_blocks_per_sm, kernel_flashinfer_blackwell_softmax_bootstrap_seed, - kBootstrapThreads, kDynamicSmemBytes); + &active_blocks_per_sm, kernel_flashinfer_blackwell_softmax_bootstrap_seed, kBootstrapThreads, + kDynamicSmemBytes); if (status != cudaSuccess) { return status; } @@ -94,20 +93,18 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper return status; } - const uint64_t tiles_per_row = ceil_div(static_cast(vocab_size), - static_cast(kBootstrapThreads)); + const uint64_t tiles_per_row = + ceil_div(static_cast(vocab_size), static_cast(kBootstrapThreads)); const uint64_t cooperative_capacity = static_cast(active_blocks_per_sm) * static_cast(sm_count); - const uint64_t provisional_grid = - std::max(1, std::min(static_cast(rows) * tiles_per_row, - cooperative_capacity)); + const uint64_t provisional_grid = std::max( + 1, std::min(static_cast(rows) * tiles_per_row, cooperative_capacity)); const uint64_t split_capacity = provisional_grid >= rows ? provisional_grid / static_cast(rows) : 1; const int splits = static_cast( std::max(1, std::min({tiles_per_row, kMaxSplits, split_capacity}))); - const uint64_t grid = - std::max(1, std::min(provisional_grid, - static_cast(rows) * splits)); + const uint64_t grid = std::max( + 1, std::min(provisional_grid, static_cast(rows) * splits)); const size_t partial_count = static_cast(rows) * static_cast(splits); const size_t required_workspace_bytes = partial_count * 2 * sizeof(float); if (required_workspace_bytes > workspace_bytes) { @@ -117,13 +114,11 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper auto* partial_max = static_cast(workspace); auto* partial_sum = partial_max + partial_count; int splits_i = splits; - void* args[] = {&logits, ¶meter, &output, &partial_max, - &partial_sum, &rows_i, &vocab_size_i, &splits_i, - ¶meter_kind_i, &temperature_val}; + void* args[] = {&logits, ¶meter, &output, &partial_max, &partial_sum, + &rows_i, &vocab_size_i, &splits_i, ¶meter_kind_i, &temperature_val}; return cudaLaunchCooperativeKernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_bootstrap_seed), - dim3(static_cast(grid)), dim3(kBootstrapThreads), args, kDynamicSmemBytes, - stream); + dim3(static_cast(grid)), dim3(kBootstrapThreads), args, kDynamicSmemBytes, stream); } } // namespace @@ -149,18 +144,16 @@ void blackwell_softmax(TensorView workspace_buffer, TensorView logits, TensorVie auto* output_ptr = static_cast(output.data_ptr()); auto* temperature_ptr = has_temperature_arr ? static_cast(maybe_temperature_arr.value().data_ptr()) : nullptr; - const size_t workspace_bytes = - get_element_size(workspace_buffer) * workspace_buffer.size(0); + const size_t workspace_bytes = get_element_size(workspace_buffer) * workspace_buffer.size(0); - cudaError_t status = - launch_blackwell_softmax(logits_ptr, output_ptr, temperature_ptr, - static_cast(temperature_val), parameter_kind, rows, - vocab_size, workspace_buffer.data_ptr(), workspace_bytes, stream); + cudaError_t status = launch_blackwell_softmax( + logits_ptr, output_ptr, temperature_ptr, static_cast(temperature_val), parameter_kind, + rows, vocab_size, workspace_buffer.data_ptr(), workspace_bytes, stream); if (status == cudaErrorNotSupported) { - status = sampling::OnlineSoftmax( - logits_ptr, output_ptr, rows, vocab_size, temperature_ptr, - static_cast(temperature_val), workspace_buffer.data_ptr(), workspace_bytes, - enable_pdl, stream); + status = sampling::OnlineSoftmax(logits_ptr, output_ptr, rows, vocab_size, + temperature_ptr, static_cast(temperature_val), + workspace_buffer.data_ptr(), workspace_bytes, + enable_pdl, stream); } TVM_FFI_ICHECK(status == cudaSuccess) << "Blackwell Softmax failed with error code " << cudaGetErrorString(status); diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index 42590526b79..13d6d682cbb 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -15,58 +15,56 @@ */ #pragma once +#include +#include + #include #include #include #include -#include -#include - // Generated by Loom from Cake commit 25dba320359dc009daf7039d1778a8721681bba8. // The sm_100a and sm_103a outputs are byte-identical (payload SHA-256: // d6e46c2babcd74565ce80df3009e3cf2751f773c483b58f0d4676c41bf0dfd44). __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; + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); + return result; } -#include - #include +#include __device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; } __device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; + 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; + 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; + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; } #define LOOM_INF CUDART_INF_F @@ -79,635 +77,655 @@ __device__ __forceinline__ float max_noftz(float a, float b) { extern "C" { -__global__ __launch_bounds__(256, 2) void -kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, float* __restrict__ partial_max, float* __restrict__ partial_sum, int rows, int vocab_size, int splits, int parameter_kind, float scalar_temperature) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; +__global__ __launch_bounds__(256, 2) void kernel_flashinfer_blackwell_softmax_bootstrap_seed( + float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, + float* __restrict__ partial_max, float* __restrict__ partial_sum, int rows, int vocab_size, + int splits, int parameter_kind, float scalar_temperature) { + 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; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - const int bid = blockIdx.x; - const int num_bids = gridDim.x; + const int bid = blockIdx.x; + const int num_bids = gridDim.x; - // Kernel setup ops - float* reduce_smem = reinterpret_cast(smem_raw + 0); - const int reduce_smem_addr = smem + 0; + // Kernel setup ops + float* reduce_smem = reinterpret_cast(smem_raw + 0); + const int reduce_smem_addr = smem + 0; - // === Task calls (dependency order) === - int total_tasks = rows * splits; - for (int task = bid; task < total_tasks; task += num_bids) { - int row = task / splits; - int split = task - row * splits; - float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); - int row_base = row * vocab_size; - int start_col = split * 256 + tid; - int col_stride = splits * 256; - int vector_groups = vocab_size / 8; - int first_group = split * vector_groups / splits; - int last_group = (split + 1) * vector_groups / splits; - float local_max = -CUDART_INF_F; - float local_sum = 0.0f; - if (row_base % 8 == 0) { - for (int group = first_group + tid; group < last_group; group += 256) { - int col = group * 8; - unsigned long long index = (unsigned long long)(row_base + col); - float _vec_load_0[8]; - { - unsigned _ldv8_0_0; - unsigned _ldv8_0_1; - unsigned _ldv8_0_2; - unsigned _ldv8_0_3; - unsigned _ldv8_0_4; - unsigned _ldv8_0_5; - unsigned _ldv8_0_6; - unsigned _ldv8_0_7; - asm volatile( - "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_0_0), "=r"(_ldv8_0_1), "=r"(_ldv8_0_2), "=r"(_ldv8_0_3), "=r"(_ldv8_0_4), "=r"(_ldv8_0_5), "=r"(_ldv8_0_6), "=r"(_ldv8_0_7) : "l"((const void*)(x + (index))) : "memory"); - _vec_load_0[0 + 0] = __uint_as_float(_ldv8_0_0); - _vec_load_0[0 + 1] = __uint_as_float(_ldv8_0_1); - _vec_load_0[0 + 2] = __uint_as_float(_ldv8_0_2); - _vec_load_0[0 + 3] = __uint_as_float(_ldv8_0_3); - _vec_load_0[0 + 4] = __uint_as_float(_ldv8_0_4); - _vec_load_0[0 + 5] = __uint_as_float(_ldv8_0_5); - _vec_load_0[0 + 6] = __uint_as_float(_ldv8_0_6); - _vec_load_0[0 + 7] = __uint_as_float(_ldv8_0_7); - } - float group_max = -CUDART_INF_F; - for (int j = 0; j < 8; j++) { - float scaled = _vec_load_0[j] / temperature; - float _max_0 = max_noftz(group_max, scaled); - group_max = _max_0; - } - if (splits == 1) { - if (group_max > local_max) { - float _exp2_0 = approx_exp2((local_max - group_max) * 1.4426950408889634f); - local_sum *= _exp2_0; - local_max = group_max; - } - if (group_max > -CUDART_INF_F) { - for (int j_1 = 0; j_1 < 8; j_1++) { - float scaled_1 = _vec_load_0[j_1] / temperature; - if (scaled_1 > -CUDART_INF_F) { - float _exp2_1 = approx_exp2((scaled_1 - local_max) * 1.4426950408889634f); - local_sum += _exp2_1; - } - } - } - } else { - float _max_1 = max_noftz(local_max, group_max); - local_max = _max_1; - } - } - int tail_col = vector_groups * 8 + tid; - if (split == splits - 1 && tail_col < vocab_size) { - unsigned long long tail_index = (unsigned long long)(row_base + tail_col); - float tail_scaled = x[tail_index] / temperature; - if (splits == 1) { - if (tail_scaled > local_max) { - float _exp2_2 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); - local_sum *= _exp2_2; - local_max = tail_scaled; - } - if (tail_scaled > -CUDART_INF_F) { - float _exp2_3 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); - local_sum += _exp2_3; - } - } else { - float _max_2 = max_noftz(local_max, tail_scaled); - local_max = _max_2; - } - } - } else { - for (int col_1 = start_col; col_1 < vocab_size; col_1 += col_stride) { - unsigned long long index_1 = (unsigned long long)(row_base + col_1); - float scaled_2 = x[index_1] / temperature; - if (splits == 1) { - if (scaled_2 > local_max) { - float _exp2_4 = approx_exp2((local_max - scaled_2) * 1.4426950408889634f); - local_sum *= _exp2_4; - local_max = scaled_2; - } - if (scaled_2 > -CUDART_INF_F) { - float _exp2_5 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); - local_sum += _exp2_5; - } - } else { - float _max_3 = max_noftz(local_max, scaled_2); - local_max = _max_3; - } - } - } - float _warp_reduce_0 = local_max; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); - float warp_max = _warp_reduce_0; - if (lane == 0) { - reduce_smem[warp] = warp_max; - } - __syncthreads(); - float warp_partial_max = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); - float _warp_reduce_1 = warp_partial_max; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); - float block_max = _warp_reduce_1; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - reduce_smem[0] = block_max; - } - } - __syncthreads(); - float cta_max = reduce_smem[0]; - __syncthreads(); - if (splits == 1) { - if (local_max > -CUDART_INF_F) { - float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); - local_sum *= _exp2_6; - } - } else if (cta_max > -CUDART_INF_F) { - if (row_base % 8 == 0) { - for (int group_1 = first_group + tid; group_1 < last_group; group_1 += 256) { - int col_2 = group_1 * 8; - unsigned long long index_2 = (unsigned long long)(row_base + col_2); - float _vec_load_1[8]; - { - unsigned _ldv8_1_0; - unsigned _ldv8_1_1; - unsigned _ldv8_1_2; - unsigned _ldv8_1_3; - unsigned _ldv8_1_4; - unsigned _ldv8_1_5; - unsigned _ldv8_1_6; - unsigned _ldv8_1_7; - asm volatile( - "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) : "l"((const void*)(x + (index_2))) : "memory"); - _vec_load_1[0 + 0] = __uint_as_float(_ldv8_1_0); - _vec_load_1[0 + 1] = __uint_as_float(_ldv8_1_1); - _vec_load_1[0 + 2] = __uint_as_float(_ldv8_1_2); - _vec_load_1[0 + 3] = __uint_as_float(_ldv8_1_3); - _vec_load_1[0 + 4] = __uint_as_float(_ldv8_1_4); - _vec_load_1[0 + 5] = __uint_as_float(_ldv8_1_5); - _vec_load_1[0 + 6] = __uint_as_float(_ldv8_1_6); - _vec_load_1[0 + 7] = __uint_as_float(_ldv8_1_7); - } - for (int j_2 = 0; j_2 < 8; j_2++) { - float scaled_3 = _vec_load_1[j_2] / temperature; - float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); - local_sum += _exp2_7; - } - } - int tail_col_1 = vector_groups * 8 + tid; - if (split == splits - 1 && tail_col_1 < vocab_size) { - unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col_1); - float tail_scaled_1 = x[tail_index_1] / temperature; - float _exp2_8 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); - local_sum += _exp2_8; - } - } else { - for (int col_3 = start_col; col_3 < vocab_size; col_3 += col_stride) { - unsigned long long index_3 = (unsigned long long)(row_base + col_3); - float scaled_4 = x[index_3] / temperature; - float _exp2_9 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); - local_sum += _exp2_9; - } - } - } - float _warp_reduce_2 = local_sum; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); - float warp_sum = _warp_reduce_2; - if (lane == 0) { - reduce_smem[warp] = warp_sum; - } - __syncthreads(); - float warp_partial_sum = ((lane < 8) ? reduce_smem[lane] : 0.0f); - float _warp_reduce_3 = warp_partial_sum; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); - float block_sum = _warp_reduce_3; - __syncthreads(); - if (splits == 1) { - float _rcp_0 = approx_rcp(block_sum); - float inv_sum = _rcp_0; - if (row_base % 8 == 0) { - for (int group_2 = first_group + tid; group_2 < last_group; group_2 += 256) { - int col_4 = group_2 * 8; - unsigned long long index_4 = (unsigned long long)(row_base + col_4); - float _vec_load_2[8]; - { - unsigned _ldv8_2_0; - unsigned _ldv8_2_1; - unsigned _ldv8_2_2; - unsigned _ldv8_2_3; - unsigned _ldv8_2_4; - unsigned _ldv8_2_5; - unsigned _ldv8_2_6; - unsigned _ldv8_2_7; - asm volatile( - "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) : "l"((const void*)(x + (index_4))) : "memory"); - _vec_load_2[0 + 0] = __uint_as_float(_ldv8_2_0); - _vec_load_2[0 + 1] = __uint_as_float(_ldv8_2_1); - _vec_load_2[0 + 2] = __uint_as_float(_ldv8_2_2); - _vec_load_2[0 + 3] = __uint_as_float(_ldv8_2_3); - _vec_load_2[0 + 4] = __uint_as_float(_ldv8_2_4); - _vec_load_2[0 + 5] = __uint_as_float(_ldv8_2_5); - _vec_load_2[0 + 6] = __uint_as_float(_ldv8_2_6); - _vec_load_2[0 + 7] = __uint_as_float(_ldv8_2_7); - } - for (int j_3 = 0; j_3 < 8; j_3++) { - float scaled_5 = _vec_load_2[j_3] / temperature; - float _exp2_10 = approx_exp2((scaled_5 - cta_max) * 1.4426950408889634f); - _vec_load_2[j_3] = _exp2_10 * inv_sum; - } - { - unsigned _stv8_3_0 = __float_as_uint(_vec_load_2[0 + 0]); - unsigned _stv8_3_1 = __float_as_uint(_vec_load_2[0 + 1]); - unsigned _stv8_3_2 = __float_as_uint(_vec_load_2[0 + 2]); - unsigned _stv8_3_3 = __float_as_uint(_vec_load_2[0 + 3]); - unsigned _stv8_3_4 = __float_as_uint(_vec_load_2[0 + 4]); - unsigned _stv8_3_5 = __float_as_uint(_vec_load_2[0 + 5]); - unsigned _stv8_3_6 = __float_as_uint(_vec_load_2[0 + 6]); - unsigned _stv8_3_7 = __float_as_uint(_vec_load_2[0 + 7]); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(output + (index_4))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); - } - } - int tail_col_2 = vector_groups * 8 + tid; - if (tail_col_2 < vocab_size) { - unsigned long long tail_index_2 = (unsigned long long)(row_base + tail_col_2); - float tail_scaled_2 = x[tail_index_2] / temperature; - float _exp2_11 = approx_exp2((tail_scaled_2 - cta_max) * 1.4426950408889634f); - output[tail_index_2] = _exp2_11 * inv_sum; - } - } else { - for (int col_5 = start_col; col_5 < vocab_size; col_5 += col_stride) { - unsigned long long index_5 = (unsigned long long)(row_base + col_5); - float scaled_6 = x[index_5] / temperature; - float _exp2_12 = approx_exp2((scaled_6 - cta_max) * 1.4426950408889634f); - output[index_5] = _exp2_12 * inv_sum; - } - } - __syncthreads(); - } else if (warp == 0) { - if (elect_sync()) { - unsigned long long partial_index = (unsigned long long)task; - partial_max[partial_index] = cta_max; - partial_sum[partial_index] = block_sum; - } - } - } - if (splits > 1) { - __threadfence(); - cooperative_groups::this_grid().sync(); - } - int phase_2_tasks = ((splits > 1) ? total_tasks : 0); - for (int task_1 = bid; task_1 < phase_2_tasks; task_1 += num_bids) { - int row_1 = task_1 / splits; - int split_1 = task_1 - row_1 * splits; - float local_max_1 = -CUDART_INF_F; - for (int partial_split = tid; partial_split < splits; partial_split += 256) { - unsigned long long partial_index_1 = (unsigned long long)row_1 * (unsigned long long)splits + (unsigned long long)partial_split; - float _max_4 = max_noftz(local_max_1, partial_max[partial_index_1]); - local_max_1 = _max_4; - } - float _warp_reduce_4 = local_max_1; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_4 = max_noftz(_warp_reduce_4, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_4, offset)); - float warp_max_1 = _warp_reduce_4; - if (lane == 0) { - reduce_smem[warp] = warp_max_1; - } - __syncthreads(); - float warp_partial_max_1 = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); - float _warp_reduce_5 = warp_partial_max_1; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_5 = max_noftz(_warp_reduce_5, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_5, offset)); - float merged_max = _warp_reduce_5; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - reduce_smem[0] = merged_max; - } - } - __syncthreads(); - float global_max = reduce_smem[0]; - __syncthreads(); - float local_sum_1 = 0.0f; - for (int partial_split_1 = tid; partial_split_1 < splits; partial_split_1 += 256) { - unsigned long long partial_index_2 = (unsigned long long)row_1 * (unsigned long long)splits + (unsigned long long)partial_split_1; - float split_max = partial_max[partial_index_2]; - if (split_max > -CUDART_INF_F) { - float _exp2_13 = approx_exp2((split_max - global_max) * 1.4426950408889634f); - local_sum_1 += partial_sum[partial_index_2] * _exp2_13; - } - } - float _warp_reduce_6 = local_sum_1; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_6 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_6, offset); - float warp_sum_1 = _warp_reduce_6; - if (lane == 0) { - reduce_smem[warp] = warp_sum_1; - } - __syncthreads(); - float warp_partial_sum_1 = ((lane < 8) ? reduce_smem[lane] : 0.0f); - float _warp_reduce_7 = warp_partial_sum_1; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_7 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_7, offset); - float merged_sum = _warp_reduce_7; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - float _rcp_1 = approx_rcp(merged_sum); - reduce_smem[0] = _rcp_1; - } - } - __syncthreads(); - float inv_sum_1 = reduce_smem[0]; - __syncthreads(); - float temperature_1 = ((parameter_kind == 2) ? parameter[row_1] : scalar_temperature); - int row_base_1 = row_1 * vocab_size; - int start_col_1 = split_1 * 256 + tid; - int col_stride_1 = splits * 256; - int vector_groups_1 = vocab_size / 8; - int first_group_1 = split_1 * vector_groups_1 / splits; - int last_group_1 = (split_1 + 1) * vector_groups_1 / splits; - if (row_base_1 % 8 == 0) { - for (int group_3 = first_group_1 + tid; group_3 < last_group_1; group_3 += 256) { - int col_6 = group_3 * 8; - unsigned long long index_6 = (unsigned long long)(row_base_1 + col_6); - float _vec_load_3[8]; - { - unsigned _ldv8_4_0; - unsigned _ldv8_4_1; - unsigned _ldv8_4_2; - unsigned _ldv8_4_3; - unsigned _ldv8_4_4; - unsigned _ldv8_4_5; - unsigned _ldv8_4_6; - unsigned _ldv8_4_7; - asm volatile( - "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_4_0), "=r"(_ldv8_4_1), "=r"(_ldv8_4_2), "=r"(_ldv8_4_3), "=r"(_ldv8_4_4), "=r"(_ldv8_4_5), "=r"(_ldv8_4_6), "=r"(_ldv8_4_7) : "l"((const void*)(x + (index_6))) : "memory"); - _vec_load_3[0 + 0] = __uint_as_float(_ldv8_4_0); - _vec_load_3[0 + 1] = __uint_as_float(_ldv8_4_1); - _vec_load_3[0 + 2] = __uint_as_float(_ldv8_4_2); - _vec_load_3[0 + 3] = __uint_as_float(_ldv8_4_3); - _vec_load_3[0 + 4] = __uint_as_float(_ldv8_4_4); - _vec_load_3[0 + 5] = __uint_as_float(_ldv8_4_5); - _vec_load_3[0 + 6] = __uint_as_float(_ldv8_4_6); - _vec_load_3[0 + 7] = __uint_as_float(_ldv8_4_7); - } - for (int j_4 = 0; j_4 < 8; j_4++) { - float scaled_7 = _vec_load_3[j_4] / temperature_1; - float _exp2_14 = approx_exp2((scaled_7 - global_max) * 1.4426950408889634f); - _vec_load_3[j_4] = _exp2_14 * inv_sum_1; - } - { - unsigned _stv8_5_0 = __float_as_uint(_vec_load_3[0 + 0]); - unsigned _stv8_5_1 = __float_as_uint(_vec_load_3[0 + 1]); - unsigned _stv8_5_2 = __float_as_uint(_vec_load_3[0 + 2]); - unsigned _stv8_5_3 = __float_as_uint(_vec_load_3[0 + 3]); - unsigned _stv8_5_4 = __float_as_uint(_vec_load_3[0 + 4]); - unsigned _stv8_5_5 = __float_as_uint(_vec_load_3[0 + 5]); - unsigned _stv8_5_6 = __float_as_uint(_vec_load_3[0 + 6]); - unsigned _stv8_5_7 = __float_as_uint(_vec_load_3[0 + 7]); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(output + (index_6))), "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3), "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) : "memory"); - } - } - int tail_col_3 = vector_groups_1 * 8 + tid; - if (split_1 == splits - 1 && tail_col_3 < vocab_size) { - unsigned long long tail_index_3 = (unsigned long long)(row_base_1 + tail_col_3); - float tail_scaled_3 = x[tail_index_3] / temperature_1; - float _exp2_15 = approx_exp2((tail_scaled_3 - global_max) * 1.4426950408889634f); - output[tail_index_3] = _exp2_15 * inv_sum_1; - } - } else { - for (int col_7 = start_col_1; col_7 < vocab_size; col_7 += col_stride_1) { - unsigned long long index_7 = (unsigned long long)(row_base_1 + col_7); - float scaled_8 = x[index_7] / temperature_1; - float _exp2_16 = approx_exp2((scaled_8 - global_max) * 1.4426950408889634f); - output[index_7] = _exp2_16 * inv_sum_1; - } - } - } -} - -} // extern "C" - -#undef LOOM_INF -#undef NUM_MAIN_STAGES -#undef SMEM_REDUCE_SMEM_OFF -#undef SMEM_REDUCE_SMEM_STAGE_BYTES -#undef SMEM_REDUCE_SMEM_STRIDE -#undef SMEM_TOTAL -#undef THREADS -#undef reduce_smem_addr - -#define LOOM_INF CUDART_INF_F -#define NUM_MAIN_STAGES 1 -#define SMEM_REDUCE_SMEM_OFF 0 -#define SMEM_REDUCE_SMEM_STAGE_BYTES 64 -#define SMEM_REDUCE_SMEM_STRIDE 64 -#define SMEM_TOTAL 128 -#define THREADS 512 - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_flashinfer_blackwell_softmax_followup_rowwise(float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, int rows, int vocab_size, int parameter_kind, float scalar_temperature) -{ - 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; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // Kernel setup ops - float* reduce_smem = reinterpret_cast(smem_raw + 0); - const int reduce_smem_addr = smem + 0; - - // === Task calls (dependency order) === - int row = bid; + // === Task calls (dependency order) === + int total_tasks = rows * splits; + for (int task = bid; task < total_tasks; task += num_bids) { + int row = task / splits; + int split = task - row * splits; float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); int row_base = row * vocab_size; - int row_misalignment = row_base % 4; - int aligned_begin = (4 - row_misalignment) % 4; - int _min_0 = ((aligned_begin) < (vocab_size) ? (aligned_begin) : (vocab_size)); - aligned_begin = _min_0; - int vector_groups = (vocab_size - aligned_begin) / 4; - int aligned_end = aligned_begin + vector_groups * 4; + int start_col = split * 256 + tid; + int col_stride = splits * 256; + int vector_groups = vocab_size / 8; + int first_group = split * vector_groups / splits; + int last_group = (split + 1) * vector_groups / splits; float local_max = -CUDART_INF_F; float local_sum = 0.0f; - for (int col = tid; col < aligned_begin; col += 512) { + if (row_base % 8 == 0) { + for (int group = first_group + tid; group < last_group; group += 256) { + int col = group * 8; unsigned long long index = (unsigned long long)(row_base + col); - float scaled = x[index] / temperature; - if (scaled > local_max) { - float _exp2_0 = approx_exp2((local_max - scaled) * 1.4426950408889634f); - local_sum *= _exp2_0; - local_max = scaled; - } - if (scaled > -CUDART_INF_F) { - float _exp2_1 = approx_exp2((scaled - local_max) * 1.4426950408889634f); - local_sum += _exp2_1; - } - } - for (int group = tid; group < vector_groups; group += 512) { - int col_1 = aligned_begin + group * 4; - unsigned long long index_1 = (unsigned long long)(row_base + col_1); - float _vec_load_0[4]; + float _vec_load_0[8]; { - float4 _v4 = *reinterpret_cast(x + index_1); - _vec_load_0[0 + 0] = _v4.x; - _vec_load_0[0 + 1] = _v4.y; - _vec_load_0[0 + 2] = _v4.z; - _vec_load_0[0 + 3] = _v4.w; + unsigned _ldv8_0_0; + unsigned _ldv8_0_1; + unsigned _ldv8_0_2; + unsigned _ldv8_0_3; + unsigned _ldv8_0_4; + unsigned _ldv8_0_5; + unsigned _ldv8_0_6; + unsigned _ldv8_0_7; + asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_0_0), "=r"(_ldv8_0_1), "=r"(_ldv8_0_2), "=r"(_ldv8_0_3), + "=r"(_ldv8_0_4), "=r"(_ldv8_0_5), "=r"(_ldv8_0_6), "=r"(_ldv8_0_7) + : "l"((const void*)(x + (index))) + : "memory"); + _vec_load_0[0 + 0] = __uint_as_float(_ldv8_0_0); + _vec_load_0[0 + 1] = __uint_as_float(_ldv8_0_1); + _vec_load_0[0 + 2] = __uint_as_float(_ldv8_0_2); + _vec_load_0[0 + 3] = __uint_as_float(_ldv8_0_3); + _vec_load_0[0 + 4] = __uint_as_float(_ldv8_0_4); + _vec_load_0[0 + 5] = __uint_as_float(_ldv8_0_5); + _vec_load_0[0 + 6] = __uint_as_float(_ldv8_0_6); + _vec_load_0[0 + 7] = __uint_as_float(_ldv8_0_7); } float group_max = -CUDART_INF_F; - for (int j = 0; j < 4; j++) { - float scaled_1 = _vec_load_0[j] / temperature; - float _max_0 = max_noftz(group_max, scaled_1); - group_max = _max_0; + for (int j = 0; j < 8; j++) { + float scaled = _vec_load_0[j] / temperature; + float _max_0 = max_noftz(group_max, scaled); + group_max = _max_0; } - if (group_max > local_max) { - float _exp2_2 = approx_exp2((local_max - group_max) * 1.4426950408889634f); - local_sum *= _exp2_2; + if (splits == 1) { + if (group_max > local_max) { + float _exp2_0 = approx_exp2((local_max - group_max) * 1.4426950408889634f); + local_sum *= _exp2_0; local_max = group_max; - } - if (group_max > -CUDART_INF_F) { - for (int j_1 = 0; j_1 < 4; j_1++) { - float scaled_2 = _vec_load_0[j_1] / temperature; - if (scaled_2 > -CUDART_INF_F) { - float _exp2_3 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); - local_sum += _exp2_3; - } + } + if (group_max > -CUDART_INF_F) { + for (int j_1 = 0; j_1 < 8; j_1++) { + float scaled_1 = _vec_load_0[j_1] / temperature; + if (scaled_1 > -CUDART_INF_F) { + float _exp2_1 = approx_exp2((scaled_1 - local_max) * 1.4426950408889634f); + local_sum += _exp2_1; + } } + } + } else { + float _max_1 = max_noftz(local_max, group_max); + local_max = _max_1; } - } - int tail_col = aligned_end + tid; - if (tail_col < vocab_size) { + } + int tail_col = vector_groups * 8 + tid; + if (split == splits - 1 && tail_col < vocab_size) { unsigned long long tail_index = (unsigned long long)(row_base + tail_col); float tail_scaled = x[tail_index] / temperature; - if (tail_scaled > local_max) { - float _exp2_4 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); - local_sum *= _exp2_4; + if (splits == 1) { + if (tail_scaled > local_max) { + float _exp2_2 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); + local_sum *= _exp2_2; local_max = tail_scaled; + } + if (tail_scaled > -CUDART_INF_F) { + float _exp2_3 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_3; + } + } else { + float _max_2 = max_noftz(local_max, tail_scaled); + local_max = _max_2; } - if (tail_scaled > -CUDART_INF_F) { - float _exp2_5 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); + } + } else { + for (int col_1 = start_col; col_1 < vocab_size; col_1 += col_stride) { + unsigned long long index_1 = (unsigned long long)(row_base + col_1); + float scaled_2 = x[index_1] / temperature; + if (splits == 1) { + if (scaled_2 > local_max) { + float _exp2_4 = approx_exp2((local_max - scaled_2) * 1.4426950408889634f); + local_sum *= _exp2_4; + local_max = scaled_2; + } + if (scaled_2 > -CUDART_INF_F) { + float _exp2_5 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); local_sum += _exp2_5; + } + } else { + float _max_3 = max_noftz(local_max, scaled_2); + local_max = _max_3; } + } } float _warp_reduce_0 = local_max; - #pragma unroll +#pragma unroll for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + _warp_reduce_0 = + max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); float warp_max = _warp_reduce_0; if (lane == 0) { - reduce_smem[warp] = warp_max; + reduce_smem[warp] = warp_max; } __syncthreads(); - float warp_partial_max = ((lane < 16) ? reduce_smem[lane] : -CUDART_INF_F); + float warp_partial_max = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); float _warp_reduce_1 = warp_partial_max; - #pragma unroll +#pragma unroll for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); + _warp_reduce_1 = + max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); float block_max = _warp_reduce_1; __syncthreads(); if (warp == 0) { - if (elect_sync()) { - reduce_smem[0] = block_max; - } + if (elect_sync()) { + reduce_smem[0] = block_max; + } } __syncthreads(); float cta_max = reduce_smem[0]; __syncthreads(); - if (local_max > -CUDART_INF_F) { + if (splits == 1) { + if (local_max > -CUDART_INF_F) { float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); local_sum *= _exp2_6; + } + } else if (cta_max > -CUDART_INF_F) { + if (row_base % 8 == 0) { + for (int group_1 = first_group + tid; group_1 < last_group; group_1 += 256) { + int col_2 = group_1 * 8; + unsigned long long index_2 = (unsigned long long)(row_base + col_2); + float _vec_load_1[8]; + { + unsigned _ldv8_1_0; + unsigned _ldv8_1_1; + unsigned _ldv8_1_2; + unsigned _ldv8_1_3; + unsigned _ldv8_1_4; + unsigned _ldv8_1_5; + unsigned _ldv8_1_6; + unsigned _ldv8_1_7; + asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), + "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) + : "l"((const void*)(x + (index_2))) + : "memory"); + _vec_load_1[0 + 0] = __uint_as_float(_ldv8_1_0); + _vec_load_1[0 + 1] = __uint_as_float(_ldv8_1_1); + _vec_load_1[0 + 2] = __uint_as_float(_ldv8_1_2); + _vec_load_1[0 + 3] = __uint_as_float(_ldv8_1_3); + _vec_load_1[0 + 4] = __uint_as_float(_ldv8_1_4); + _vec_load_1[0 + 5] = __uint_as_float(_ldv8_1_5); + _vec_load_1[0 + 6] = __uint_as_float(_ldv8_1_6); + _vec_load_1[0 + 7] = __uint_as_float(_ldv8_1_7); + } + for (int j_2 = 0; j_2 < 8; j_2++) { + float scaled_3 = _vec_load_1[j_2] / temperature; + float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_7; + } + } + int tail_col_1 = vector_groups * 8 + tid; + if (split == splits - 1 && tail_col_1 < vocab_size) { + unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col_1); + float tail_scaled_1 = x[tail_index_1] / temperature; + float _exp2_8 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_8; + } + } else { + for (int col_3 = start_col; col_3 < vocab_size; col_3 += col_stride) { + unsigned long long index_3 = (unsigned long long)(row_base + col_3); + float scaled_4 = x[index_3] / temperature; + float _exp2_9 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_9; + } + } } float _warp_reduce_2 = local_sum; - #pragma unroll +#pragma unroll for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); + _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); float warp_sum = _warp_reduce_2; if (lane == 0) { - reduce_smem[warp] = warp_sum; + reduce_smem[warp] = warp_sum; } __syncthreads(); - float warp_partial_sum = ((lane < 16) ? reduce_smem[lane] : 0.0f); + float warp_partial_sum = ((lane < 8) ? reduce_smem[lane] : 0.0f); float _warp_reduce_3 = warp_partial_sum; - #pragma unroll +#pragma unroll for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); + _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); float block_sum = _warp_reduce_3; __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - float _rcp_0 = approx_rcp(block_sum); - reduce_smem[0] = _rcp_0; + if (splits == 1) { + float _rcp_0 = approx_rcp(block_sum); + float inv_sum = _rcp_0; + if (row_base % 8 == 0) { + for (int group_2 = first_group + tid; group_2 < last_group; group_2 += 256) { + int col_4 = group_2 * 8; + unsigned long long index_4 = (unsigned long long)(row_base + col_4); + float _vec_load_2[8]; + { + unsigned _ldv8_2_0; + unsigned _ldv8_2_1; + unsigned _ldv8_2_2; + unsigned _ldv8_2_3; + unsigned _ldv8_2_4; + unsigned _ldv8_2_5; + unsigned _ldv8_2_6; + unsigned _ldv8_2_7; + asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), + "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) + : "l"((const void*)(x + (index_4))) + : "memory"); + _vec_load_2[0 + 0] = __uint_as_float(_ldv8_2_0); + _vec_load_2[0 + 1] = __uint_as_float(_ldv8_2_1); + _vec_load_2[0 + 2] = __uint_as_float(_ldv8_2_2); + _vec_load_2[0 + 3] = __uint_as_float(_ldv8_2_3); + _vec_load_2[0 + 4] = __uint_as_float(_ldv8_2_4); + _vec_load_2[0 + 5] = __uint_as_float(_ldv8_2_5); + _vec_load_2[0 + 6] = __uint_as_float(_ldv8_2_6); + _vec_load_2[0 + 7] = __uint_as_float(_ldv8_2_7); + } + for (int j_3 = 0; j_3 < 8; j_3++) { + float scaled_5 = _vec_load_2[j_3] / temperature; + float _exp2_10 = approx_exp2((scaled_5 - cta_max) * 1.4426950408889634f); + _vec_load_2[j_3] = _exp2_10 * inv_sum; + } + { + unsigned _stv8_3_0 = __float_as_uint(_vec_load_2[0 + 0]); + unsigned _stv8_3_1 = __float_as_uint(_vec_load_2[0 + 1]); + unsigned _stv8_3_2 = __float_as_uint(_vec_load_2[0 + 2]); + unsigned _stv8_3_3 = __float_as_uint(_vec_load_2[0 + 3]); + unsigned _stv8_3_4 = __float_as_uint(_vec_load_2[0 + 4]); + unsigned _stv8_3_5 = __float_as_uint(_vec_load_2[0 + 5]); + unsigned _stv8_3_6 = __float_as_uint(_vec_load_2[0 + 6]); + unsigned _stv8_3_7 = __float_as_uint(_vec_load_2[0 + 7]); + asm volatile("st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" ::"l"( + (void*)(output + (index_4))), + "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), + "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) + : "memory"); + } + } + int tail_col_2 = vector_groups * 8 + tid; + if (tail_col_2 < vocab_size) { + unsigned long long tail_index_2 = (unsigned long long)(row_base + tail_col_2); + float tail_scaled_2 = x[tail_index_2] / temperature; + float _exp2_11 = approx_exp2((tail_scaled_2 - cta_max) * 1.4426950408889634f); + output[tail_index_2] = _exp2_11 * inv_sum; + } + } else { + for (int col_5 = start_col; col_5 < vocab_size; col_5 += col_stride) { + unsigned long long index_5 = (unsigned long long)(row_base + col_5); + float scaled_6 = x[index_5] / temperature; + float _exp2_12 = approx_exp2((scaled_6 - cta_max) * 1.4426950408889634f); + output[index_5] = _exp2_12 * inv_sum; } + } + __syncthreads(); + } else if (warp == 0) { + if (elect_sync()) { + unsigned long long partial_index = (unsigned long long)task; + partial_max[partial_index] = cta_max; + partial_sum[partial_index] = block_sum; + } } + } + if (splits > 1) { + __threadfence(); + cooperative_groups::this_grid().sync(); + } + int phase_2_tasks = ((splits > 1) ? total_tasks : 0); + for (int task_1 = bid; task_1 < phase_2_tasks; task_1 += num_bids) { + int row_1 = task_1 / splits; + int split_1 = task_1 - row_1 * splits; + float local_max_1 = -CUDART_INF_F; + for (int partial_split = tid; partial_split < splits; partial_split += 256) { + unsigned long long partial_index_1 = (unsigned long long)row_1 * (unsigned long long)splits + + (unsigned long long)partial_split; + float _max_4 = max_noftz(local_max_1, partial_max[partial_index_1]); + local_max_1 = _max_4; + } + float _warp_reduce_4 = local_max_1; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_4 = + max_noftz(_warp_reduce_4, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_4, offset)); + float warp_max_1 = _warp_reduce_4; + if (lane == 0) { + reduce_smem[warp] = warp_max_1; + } + __syncthreads(); + float warp_partial_max_1 = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_5 = warp_partial_max_1; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_5 = + max_noftz(_warp_reduce_5, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_5, offset)); + float merged_max = _warp_reduce_5; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = merged_max; + } + } + __syncthreads(); + float global_max = reduce_smem[0]; __syncthreads(); - float inv_sum = reduce_smem[0]; - for (int col_2 = tid; col_2 < aligned_begin; col_2 += 512) { - unsigned long long index_2 = (unsigned long long)(row_base + col_2); - float scaled_3 = x[index_2] / temperature; - float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); - output[index_2] = _exp2_7 * inv_sum; + float local_sum_1 = 0.0f; + for (int partial_split_1 = tid; partial_split_1 < splits; partial_split_1 += 256) { + unsigned long long partial_index_2 = (unsigned long long)row_1 * (unsigned long long)splits + + (unsigned long long)partial_split_1; + float split_max = partial_max[partial_index_2]; + if (split_max > -CUDART_INF_F) { + float _exp2_13 = approx_exp2((split_max - global_max) * 1.4426950408889634f); + local_sum_1 += partial_sum[partial_index_2] * _exp2_13; + } } - for (int group_1 = tid; group_1 < vector_groups; group_1 += 512) { - int col_3 = aligned_begin + group_1 * 4; - unsigned long long index_3 = (unsigned long long)(row_base + col_3); - float _vec_load_1[4]; + float _warp_reduce_6 = local_sum_1; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_6 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_6, offset); + float warp_sum_1 = _warp_reduce_6; + if (lane == 0) { + reduce_smem[warp] = warp_sum_1; + } + __syncthreads(); + float warp_partial_sum_1 = ((lane < 8) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_7 = warp_partial_sum_1; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_7 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_7, offset); + float merged_sum = _warp_reduce_7; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + float _rcp_1 = approx_rcp(merged_sum); + reduce_smem[0] = _rcp_1; + } + } + __syncthreads(); + float inv_sum_1 = reduce_smem[0]; + __syncthreads(); + float temperature_1 = ((parameter_kind == 2) ? parameter[row_1] : scalar_temperature); + int row_base_1 = row_1 * vocab_size; + int start_col_1 = split_1 * 256 + tid; + int col_stride_1 = splits * 256; + int vector_groups_1 = vocab_size / 8; + int first_group_1 = split_1 * vector_groups_1 / splits; + int last_group_1 = (split_1 + 1) * vector_groups_1 / splits; + if (row_base_1 % 8 == 0) { + for (int group_3 = first_group_1 + tid; group_3 < last_group_1; group_3 += 256) { + int col_6 = group_3 * 8; + unsigned long long index_6 = (unsigned long long)(row_base_1 + col_6); + float _vec_load_3[8]; { - float4 _v4 = *reinterpret_cast(x + index_3); - _vec_load_1[0 + 0] = _v4.x; - _vec_load_1[0 + 1] = _v4.y; - _vec_load_1[0 + 2] = _v4.z; - _vec_load_1[0 + 3] = _v4.w; + unsigned _ldv8_4_0; + unsigned _ldv8_4_1; + unsigned _ldv8_4_2; + unsigned _ldv8_4_3; + unsigned _ldv8_4_4; + unsigned _ldv8_4_5; + unsigned _ldv8_4_6; + unsigned _ldv8_4_7; + asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_4_0), "=r"(_ldv8_4_1), "=r"(_ldv8_4_2), "=r"(_ldv8_4_3), + "=r"(_ldv8_4_4), "=r"(_ldv8_4_5), "=r"(_ldv8_4_6), "=r"(_ldv8_4_7) + : "l"((const void*)(x + (index_6))) + : "memory"); + _vec_load_3[0 + 0] = __uint_as_float(_ldv8_4_0); + _vec_load_3[0 + 1] = __uint_as_float(_ldv8_4_1); + _vec_load_3[0 + 2] = __uint_as_float(_ldv8_4_2); + _vec_load_3[0 + 3] = __uint_as_float(_ldv8_4_3); + _vec_load_3[0 + 4] = __uint_as_float(_ldv8_4_4); + _vec_load_3[0 + 5] = __uint_as_float(_ldv8_4_5); + _vec_load_3[0 + 6] = __uint_as_float(_ldv8_4_6); + _vec_load_3[0 + 7] = __uint_as_float(_ldv8_4_7); } - for (int j_2 = 0; j_2 < 4; j_2++) { - float scaled_4 = _vec_load_1[j_2] / temperature; - float _exp2_8 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); - _vec_load_1[j_2] = _exp2_8 * inv_sum; + for (int j_4 = 0; j_4 < 8; j_4++) { + float scaled_7 = _vec_load_3[j_4] / temperature_1; + float _exp2_14 = approx_exp2((scaled_7 - global_max) * 1.4426950408889634f); + _vec_load_3[j_4] = _exp2_14 * inv_sum_1; } { - float4 _v4 = make_float4(_vec_load_1[0 + 0], _vec_load_1[0 + 1], _vec_load_1[0 + 2], _vec_load_1[0 + 3]); - *reinterpret_cast(output + index_3) = _v4; + unsigned _stv8_5_0 = __float_as_uint(_vec_load_3[0 + 0]); + unsigned _stv8_5_1 = __float_as_uint(_vec_load_3[0 + 1]); + unsigned _stv8_5_2 = __float_as_uint(_vec_load_3[0 + 2]); + unsigned _stv8_5_3 = __float_as_uint(_vec_load_3[0 + 3]); + unsigned _stv8_5_4 = __float_as_uint(_vec_load_3[0 + 4]); + unsigned _stv8_5_5 = __float_as_uint(_vec_load_3[0 + 5]); + unsigned _stv8_5_6 = __float_as_uint(_vec_load_3[0 + 6]); + unsigned _stv8_5_7 = __float_as_uint(_vec_load_3[0 + 7]); + asm volatile("st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" ::"l"( + (void*)(output + (index_6))), + "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3), + "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) + : "memory"); } + } + int tail_col_3 = vector_groups_1 * 8 + tid; + if (split_1 == splits - 1 && tail_col_3 < vocab_size) { + unsigned long long tail_index_3 = (unsigned long long)(row_base_1 + tail_col_3); + float tail_scaled_3 = x[tail_index_3] / temperature_1; + float _exp2_15 = approx_exp2((tail_scaled_3 - global_max) * 1.4426950408889634f); + output[tail_index_3] = _exp2_15 * inv_sum_1; + } + } else { + for (int col_7 = start_col_1; col_7 < vocab_size; col_7 += col_stride_1) { + unsigned long long index_7 = (unsigned long long)(row_base_1 + col_7); + float scaled_8 = x[index_7] / temperature_1; + float _exp2_16 = approx_exp2((scaled_8 - global_max) * 1.4426950408889634f); + output[index_7] = _exp2_16 * inv_sum_1; + } + } + } +} + +} // extern "C" + +#undef LOOM_INF +#undef NUM_MAIN_STAGES +#undef SMEM_REDUCE_SMEM_OFF +#undef SMEM_REDUCE_SMEM_STAGE_BYTES +#undef SMEM_REDUCE_SMEM_STRIDE +#undef SMEM_TOTAL +#undef THREADS +#undef reduce_smem_addr + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_REDUCE_SMEM_OFF 0 +#define SMEM_REDUCE_SMEM_STAGE_BYTES 64 +#define SMEM_REDUCE_SMEM_STRIDE 64 +#define SMEM_TOTAL 128 +#define THREADS 512 + +extern "C" { + +__global__ __launch_bounds__(512, 1) void kernel_flashinfer_blackwell_softmax_followup_rowwise( + float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, int rows, + int vocab_size, int parameter_kind, float scalar_temperature) { + 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; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* reduce_smem = reinterpret_cast(smem_raw + 0); + const int reduce_smem_addr = smem + 0; + + // === Task calls (dependency order) === + int row = bid; + float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); + int row_base = row * vocab_size; + int row_misalignment = row_base % 4; + int aligned_begin = (4 - row_misalignment) % 4; + int _min_0 = ((aligned_begin) < (vocab_size) ? (aligned_begin) : (vocab_size)); + aligned_begin = _min_0; + int vector_groups = (vocab_size - aligned_begin) / 4; + int aligned_end = aligned_begin + vector_groups * 4; + float local_max = -CUDART_INF_F; + float local_sum = 0.0f; + for (int col = tid; col < aligned_begin; col += 512) { + unsigned long long index = (unsigned long long)(row_base + col); + float scaled = x[index] / temperature; + if (scaled > local_max) { + float _exp2_0 = approx_exp2((local_max - scaled) * 1.4426950408889634f); + local_sum *= _exp2_0; + local_max = scaled; + } + if (scaled > -CUDART_INF_F) { + float _exp2_1 = approx_exp2((scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_1; + } + } + for (int group = tid; group < vector_groups; group += 512) { + int col_1 = aligned_begin + group * 4; + unsigned long long index_1 = (unsigned long long)(row_base + col_1); + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(x + index_1); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float group_max = -CUDART_INF_F; + for (int j = 0; j < 4; j++) { + float scaled_1 = _vec_load_0[j] / temperature; + float _max_0 = max_noftz(group_max, scaled_1); + group_max = _max_0; + } + if (group_max > local_max) { + float _exp2_2 = approx_exp2((local_max - group_max) * 1.4426950408889634f); + local_sum *= _exp2_2; + local_max = group_max; + } + if (group_max > -CUDART_INF_F) { + for (int j_1 = 0; j_1 < 4; j_1++) { + float scaled_2 = _vec_load_0[j_1] / temperature; + if (scaled_2 > -CUDART_INF_F) { + float _exp2_3 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); + local_sum += _exp2_3; + } + } + } + } + int tail_col = aligned_end + tid; + if (tail_col < vocab_size) { + unsigned long long tail_index = (unsigned long long)(row_base + tail_col); + float tail_scaled = x[tail_index] / temperature; + if (tail_scaled > local_max) { + float _exp2_4 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); + local_sum *= _exp2_4; + local_max = tail_scaled; + } + if (tail_scaled > -CUDART_INF_F) { + float _exp2_5 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_5; + } + } + float _warp_reduce_0 = local_max; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float warp_max = _warp_reduce_0; + if (lane == 0) { + reduce_smem[warp] = warp_max; + } + __syncthreads(); + float warp_partial_max = ((lane < 16) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_1 = warp_partial_max; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); + float block_max = _warp_reduce_1; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = block_max; + } + } + __syncthreads(); + float cta_max = reduce_smem[0]; + __syncthreads(); + if (local_max > -CUDART_INF_F) { + float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); + local_sum *= _exp2_6; + } + float _warp_reduce_2 = local_sum; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); + float warp_sum = _warp_reduce_2; + if (lane == 0) { + reduce_smem[warp] = warp_sum; + } + __syncthreads(); + float warp_partial_sum = ((lane < 16) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_3 = warp_partial_sum; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); + float block_sum = _warp_reduce_3; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + float _rcp_0 = approx_rcp(block_sum); + reduce_smem[0] = _rcp_0; + } + } + __syncthreads(); + float inv_sum = reduce_smem[0]; + for (int col_2 = tid; col_2 < aligned_begin; col_2 += 512) { + unsigned long long index_2 = (unsigned long long)(row_base + col_2); + float scaled_3 = x[index_2] / temperature; + float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); + output[index_2] = _exp2_7 * inv_sum; + } + for (int group_1 = tid; group_1 < vector_groups; group_1 += 512) { + int col_3 = aligned_begin + group_1 * 4; + unsigned long long index_3 = (unsigned long long)(row_base + col_3); + float _vec_load_1[4]; + { + float4 _v4 = *reinterpret_cast(x + index_3); + _vec_load_1[0 + 0] = _v4.x; + _vec_load_1[0 + 1] = _v4.y; + _vec_load_1[0 + 2] = _v4.z; + _vec_load_1[0 + 3] = _v4.w; + } + for (int j_2 = 0; j_2 < 4; j_2++) { + float scaled_4 = _vec_load_1[j_2] / temperature; + float _exp2_8 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); + _vec_load_1[j_2] = _exp2_8 * inv_sum; } - if (tail_col < vocab_size) { - unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col); - float tail_scaled_1 = x[tail_index_1] / temperature; - float _exp2_9 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); - output[tail_index_1] = _exp2_9 * inv_sum; + { + float4 _v4 = make_float4(_vec_load_1[0 + 0], _vec_load_1[0 + 1], _vec_load_1[0 + 2], + _vec_load_1[0 + 3]); + *reinterpret_cast(output + index_3) = _v4; } + } + if (tail_col < vocab_size) { + unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col); + float tail_scaled_1 = x[tail_index_1] / temperature; + float _exp2_9 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); + output[tail_index_1] = _exp2_9 * inv_sum; + } } -} // extern "C" +} // extern "C" #undef LOOM_INF #undef NUM_MAIN_STAGES From 2093dcb9a466a97c85ddae1e63a06753ab3647b9 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 23:27:51 -0700 Subject: [PATCH 05/14] Integrate final Blackwell softmax dispatcher --- csrc/blackwell_softmax.cu | 33 +- csrc/blackwell_softmax_bootstrap.cu | 519 ++++++++++++++++ csrc/blackwell_softmax_rowwise.cu | 263 ++++++++ csrc/blackwell_softmax_warp.cu | 133 +++++ flashinfer/jit/blackwell_softmax.py | 7 +- include/flashinfer/blackwell_softmax.cuh | 727 +---------------------- tests/utils/test_sampling.py | 4 +- 7 files changed, 965 insertions(+), 721 deletions(-) create mode 100644 csrc/blackwell_softmax_bootstrap.cu create mode 100644 csrc/blackwell_softmax_rowwise.cu create mode 100644 csrc/blackwell_softmax_warp.cu diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu index b68401726af..f1173ea0cd0 100644 --- a/csrc/blackwell_softmax.cu +++ b/csrc/blackwell_softmax.cu @@ -28,8 +28,12 @@ namespace { constexpr int kBootstrapThreads = 256; constexpr int kRowwiseThreads = 512; +constexpr int kWarpThreads = 128; +constexpr int kWarpRowsPerCta = 4; constexpr int kMaxSplits = 64; -constexpr size_t kDynamicSmemBytes = 128; +constexpr size_t kBootstrapDynamicSmemBytes = 128; +constexpr size_t kRowwiseDynamicSmemBytes = 128; +constexpr size_t kWarpDynamicSmemBytes = 0; enum class ParameterKind : int { kNone = 0, @@ -37,14 +41,21 @@ enum class ParameterKind : int { kPerRow = 2, }; +bool use_warp_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind) { + return rows <= 128 && vocab_size <= 257 && + (parameter_kind == ParameterKind::kScalar || + parameter_kind == ParameterKind::kPerRow); +} + bool use_rowwise_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind) { const bool small_low_row = rows <= 32 && vocab_size <= 16384; const bool dense_aligned_mid_row = rows > 128 && rows <= 384 && vocab_size >= 24576 && vocab_size <= 256000 && vocab_size % 4 == 0 && parameter_kind == ParameterKind::kNone; const bool dense_aligned_high_row_narrow = rows > 384 && rows <= 1024 && vocab_size >= 24576 && - vocab_size <= 32000 && vocab_size % 4 == 0 && - parameter_kind == ParameterKind::kNone; + vocab_size <= 64000 && vocab_size % 4 == 0 && + (parameter_kind == ParameterKind::kNone || + vocab_size <= 32000); const bool measured_large_odd = rows > 128 && rows <= 512 && vocab_size >= 24576 && vocab_size <= 131072 && vocab_size % 4 != 0; return small_low_row || dense_aligned_mid_row || dense_aligned_high_row_narrow || @@ -66,18 +77,27 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper int vocab_size_i = static_cast(vocab_size); int parameter_kind_i = static_cast(parameter_kind); + if (use_warp_kernel(rows, vocab_size, parameter_kind)) { + void* args[] = {&logits, ¶meter, &output, &rows_i, + &vocab_size_i, ¶meter_kind_i, &temperature_val}; + return cudaLaunchKernel( + reinterpret_cast(kernel_flashinfer_blackwell_softmax_followup_warp), + dim3(ceil_div(rows, static_cast(kWarpRowsPerCta))), dim3(kWarpThreads), args, + kWarpDynamicSmemBytes, stream); + } + if (use_rowwise_kernel(rows, vocab_size, parameter_kind)) { void* args[] = {&logits, ¶meter, &output, &rows_i, &vocab_size_i, ¶meter_kind_i, &temperature_val}; return cudaLaunchKernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_followup_rowwise), - dim3(rows), dim3(kRowwiseThreads), args, kDynamicSmemBytes, stream); + dim3(rows), dim3(kRowwiseThreads), args, kRowwiseDynamicSmemBytes, stream); } int active_blocks_per_sm = 0; cudaError_t status = cudaOccupancyMaxActiveBlocksPerMultiprocessor( &active_blocks_per_sm, kernel_flashinfer_blackwell_softmax_bootstrap_seed, kBootstrapThreads, - kDynamicSmemBytes); + kBootstrapDynamicSmemBytes); if (status != cudaSuccess) { return status; } @@ -118,7 +138,8 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper &rows_i, &vocab_size_i, &splits_i, ¶meter_kind_i, &temperature_val}; return cudaLaunchCooperativeKernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_bootstrap_seed), - dim3(static_cast(grid)), dim3(kBootstrapThreads), args, kDynamicSmemBytes, stream); + dim3(static_cast(grid)), dim3(kBootstrapThreads), args, + kBootstrapDynamicSmemBytes, stream); } } // namespace diff --git a/csrc/blackwell_softmax_bootstrap.cu b/csrc/blackwell_softmax_bootstrap.cu new file mode 100644 index 00000000000..4543241e9ff --- /dev/null +++ b/csrc/blackwell_softmax_bootstrap.cu @@ -0,0 +1,519 @@ +/* + * 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 Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. +// Exact generated payload: 24102 bytes, +// sha256:46049370dbd905ff7234d414745d197f883157c8edc01d83230562b4dff5f862. +// The sm_100a and sm_103a payloads are byte-identical. + +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#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 LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_REDUCE_SMEM_OFF 0 +#define SMEM_REDUCE_SMEM_STAGE_BYTES 32 +#define SMEM_REDUCE_SMEM_STRIDE 32 +#define SMEM_TOTAL 128 +#define THREADS 256 + +#include +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__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; +} + +extern "C" { + +__global__ __launch_bounds__(256, 2) void +kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, float* __restrict__ partial_max, float* __restrict__ partial_sum, int rows, int vocab_size, int splits, int parameter_kind, float scalar_temperature) +{ + 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; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* reduce_smem = reinterpret_cast(smem_raw + 0); + const int reduce_smem_addr = smem + 0; + + // === Task calls (dependency order) === + int total_tasks = rows * splits; + for (int task = bid; task < total_tasks; task += num_bids) { + int row = task / splits; + int split = task - row * splits; + float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); + int row_base = row * vocab_size; + int start_col = split * 256 + tid; + int col_stride = splits * 256; + int vector_groups = vocab_size / 8; + int first_group = split * vector_groups / splits; + int last_group = (split + 1) * vector_groups / splits; + float local_max = -CUDART_INF_F; + float local_sum = 0.0f; + if (row_base % 8 == 0) { + for (int group = first_group + tid; group < last_group; group += 256) { + int col = group * 8; + unsigned long long index = (unsigned long long)(row_base + col); + float _vec_load_0[8]; + { + unsigned _ldv8_0_0; + unsigned _ldv8_0_1; + unsigned _ldv8_0_2; + unsigned _ldv8_0_3; + unsigned _ldv8_0_4; + unsigned _ldv8_0_5; + unsigned _ldv8_0_6; + unsigned _ldv8_0_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_0_0), "=r"(_ldv8_0_1), "=r"(_ldv8_0_2), "=r"(_ldv8_0_3), "=r"(_ldv8_0_4), "=r"(_ldv8_0_5), "=r"(_ldv8_0_6), "=r"(_ldv8_0_7) : "l"((const void*)(x + (index))) : "memory"); + _vec_load_0[0 + 0] = __uint_as_float(_ldv8_0_0); + _vec_load_0[0 + 1] = __uint_as_float(_ldv8_0_1); + _vec_load_0[0 + 2] = __uint_as_float(_ldv8_0_2); + _vec_load_0[0 + 3] = __uint_as_float(_ldv8_0_3); + _vec_load_0[0 + 4] = __uint_as_float(_ldv8_0_4); + _vec_load_0[0 + 5] = __uint_as_float(_ldv8_0_5); + _vec_load_0[0 + 6] = __uint_as_float(_ldv8_0_6); + _vec_load_0[0 + 7] = __uint_as_float(_ldv8_0_7); + } + float group_max = -CUDART_INF_F; + for (int j = 0; j < 8; j++) { + float scaled = _vec_load_0[j] / temperature; + float _max_0 = max_noftz(group_max, scaled); + group_max = _max_0; + } + if (splits == 1) { + if (group_max > local_max) { + float _exp2_0 = approx_exp2((local_max - group_max) * 1.4426950408889634f); + local_sum *= _exp2_0; + local_max = group_max; + } + if (group_max > -CUDART_INF_F) { + for (int j_1 = 0; j_1 < 8; j_1++) { + float scaled_1 = _vec_load_0[j_1] / temperature; + if (scaled_1 > -CUDART_INF_F) { + float _exp2_1 = approx_exp2((scaled_1 - local_max) * 1.4426950408889634f); + local_sum += _exp2_1; + } + } + } + } else { + float _max_1 = max_noftz(local_max, group_max); + local_max = _max_1; + } + } + int tail_col = vector_groups * 8 + tid; + if (split == splits - 1 && tail_col < vocab_size) { + unsigned long long tail_index = (unsigned long long)(row_base + tail_col); + float tail_scaled = x[tail_index] / temperature; + if (splits == 1) { + if (tail_scaled > local_max) { + float _exp2_2 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); + local_sum *= _exp2_2; + local_max = tail_scaled; + } + if (tail_scaled > -CUDART_INF_F) { + float _exp2_3 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_3; + } + } else { + float _max_2 = max_noftz(local_max, tail_scaled); + local_max = _max_2; + } + } + } else { + for (int col_1 = start_col; col_1 < vocab_size; col_1 += col_stride) { + unsigned long long index_1 = (unsigned long long)(row_base + col_1); + float scaled_2 = x[index_1] / temperature; + if (splits == 1) { + if (scaled_2 > local_max) { + float _exp2_4 = approx_exp2((local_max - scaled_2) * 1.4426950408889634f); + local_sum *= _exp2_4; + local_max = scaled_2; + } + if (scaled_2 > -CUDART_INF_F) { + float _exp2_5 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); + local_sum += _exp2_5; + } + } else { + float _max_3 = max_noftz(local_max, scaled_2); + local_max = _max_3; + } + } + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float warp_max = _warp_reduce_0; + if (lane == 0) { + reduce_smem[warp] = warp_max; + } + __syncthreads(); + float warp_partial_max = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_1 = warp_partial_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); + float block_max = _warp_reduce_1; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = block_max; + } + } + __syncthreads(); + float cta_max = reduce_smem[0]; + __syncthreads(); + if (splits == 1) { + if (local_max > -CUDART_INF_F) { + float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); + local_sum *= _exp2_6; + } + } else if (cta_max > -CUDART_INF_F) { + if (row_base % 8 == 0) { + for (int group_1 = first_group + tid; group_1 < last_group; group_1 += 256) { + int col_2 = group_1 * 8; + unsigned long long index_2 = (unsigned long long)(row_base + col_2); + float _vec_load_1[8]; + { + unsigned _ldv8_1_0; + unsigned _ldv8_1_1; + unsigned _ldv8_1_2; + unsigned _ldv8_1_3; + unsigned _ldv8_1_4; + unsigned _ldv8_1_5; + unsigned _ldv8_1_6; + unsigned _ldv8_1_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) : "l"((const void*)(x + (index_2))) : "memory"); + _vec_load_1[0 + 0] = __uint_as_float(_ldv8_1_0); + _vec_load_1[0 + 1] = __uint_as_float(_ldv8_1_1); + _vec_load_1[0 + 2] = __uint_as_float(_ldv8_1_2); + _vec_load_1[0 + 3] = __uint_as_float(_ldv8_1_3); + _vec_load_1[0 + 4] = __uint_as_float(_ldv8_1_4); + _vec_load_1[0 + 5] = __uint_as_float(_ldv8_1_5); + _vec_load_1[0 + 6] = __uint_as_float(_ldv8_1_6); + _vec_load_1[0 + 7] = __uint_as_float(_ldv8_1_7); + } + for (int j_2 = 0; j_2 < 8; j_2++) { + float scaled_3 = _vec_load_1[j_2] / temperature; + float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_7; + } + } + int tail_col_1 = vector_groups * 8 + tid; + if (split == splits - 1 && tail_col_1 < vocab_size) { + unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col_1); + float tail_scaled_1 = x[tail_index_1] / temperature; + float _exp2_8 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_8; + } + } else { + for (int col_3 = start_col; col_3 < vocab_size; col_3 += col_stride) { + unsigned long long index_3 = (unsigned long long)(row_base + col_3); + float scaled_4 = x[index_3] / temperature; + float _exp2_9 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); + local_sum += _exp2_9; + } + } + } + float _warp_reduce_2 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); + float warp_sum = _warp_reduce_2; + if (lane == 0) { + reduce_smem[warp] = warp_sum; + } + __syncthreads(); + float warp_partial_sum = ((lane < 8) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_3 = warp_partial_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); + float block_sum = _warp_reduce_3; + __syncthreads(); + if (splits == 1) { + float _rcp_0 = approx_rcp(block_sum); + float inv_sum = _rcp_0; + if (row_base % 8 == 0) { + for (int group_2 = first_group + tid; group_2 < last_group; group_2 += 256) { + int col_4 = group_2 * 8; + unsigned long long index_4 = (unsigned long long)(row_base + col_4); + float _vec_load_2[8]; + { + unsigned _ldv8_2_0; + unsigned _ldv8_2_1; + unsigned _ldv8_2_2; + unsigned _ldv8_2_3; + unsigned _ldv8_2_4; + unsigned _ldv8_2_5; + unsigned _ldv8_2_6; + unsigned _ldv8_2_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) : "l"((const void*)(x + (index_4))) : "memory"); + _vec_load_2[0 + 0] = __uint_as_float(_ldv8_2_0); + _vec_load_2[0 + 1] = __uint_as_float(_ldv8_2_1); + _vec_load_2[0 + 2] = __uint_as_float(_ldv8_2_2); + _vec_load_2[0 + 3] = __uint_as_float(_ldv8_2_3); + _vec_load_2[0 + 4] = __uint_as_float(_ldv8_2_4); + _vec_load_2[0 + 5] = __uint_as_float(_ldv8_2_5); + _vec_load_2[0 + 6] = __uint_as_float(_ldv8_2_6); + _vec_load_2[0 + 7] = __uint_as_float(_ldv8_2_7); + } + for (int j_3 = 0; j_3 < 8; j_3++) { + float scaled_5 = _vec_load_2[j_3] / temperature; + float _exp2_10 = approx_exp2((scaled_5 - cta_max) * 1.4426950408889634f); + _vec_load_2[j_3] = _exp2_10 * inv_sum; + } + { + unsigned _stv8_3_0 = __float_as_uint(_vec_load_2[0 + 0]); + unsigned _stv8_3_1 = __float_as_uint(_vec_load_2[0 + 1]); + unsigned _stv8_3_2 = __float_as_uint(_vec_load_2[0 + 2]); + unsigned _stv8_3_3 = __float_as_uint(_vec_load_2[0 + 3]); + unsigned _stv8_3_4 = __float_as_uint(_vec_load_2[0 + 4]); + unsigned _stv8_3_5 = __float_as_uint(_vec_load_2[0 + 5]); + unsigned _stv8_3_6 = __float_as_uint(_vec_load_2[0 + 6]); + unsigned _stv8_3_7 = __float_as_uint(_vec_load_2[0 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(output + (index_4))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); + } + } + int tail_col_2 = vector_groups * 8 + tid; + if (tail_col_2 < vocab_size) { + unsigned long long tail_index_2 = (unsigned long long)(row_base + tail_col_2); + float tail_scaled_2 = x[tail_index_2] / temperature; + float _exp2_11 = approx_exp2((tail_scaled_2 - cta_max) * 1.4426950408889634f); + output[tail_index_2] = _exp2_11 * inv_sum; + } + } else { + for (int col_5 = start_col; col_5 < vocab_size; col_5 += col_stride) { + unsigned long long index_5 = (unsigned long long)(row_base + col_5); + float scaled_6 = x[index_5] / temperature; + float _exp2_12 = approx_exp2((scaled_6 - cta_max) * 1.4426950408889634f); + output[index_5] = _exp2_12 * inv_sum; + } + } + __syncthreads(); + } else if (warp == 0) { + if (elect_sync()) { + unsigned long long partial_index = (unsigned long long)task; + partial_max[partial_index] = cta_max; + partial_sum[partial_index] = block_sum; + } + } + } + if (splits > 1) { + __threadfence(); + cooperative_groups::this_grid().sync(); + } + int phase_2_tasks = ((splits > 1) ? total_tasks : 0); + for (int task_1 = bid; task_1 < phase_2_tasks; task_1 += num_bids) { + int row_1 = task_1 / splits; + int split_1 = task_1 - row_1 * splits; + float local_max_1 = -CUDART_INF_F; + for (int partial_split = tid; partial_split < splits; partial_split += 256) { + unsigned long long partial_index_1 = (unsigned long long)row_1 * (unsigned long long)splits + (unsigned long long)partial_split; + float _max_4 = max_noftz(local_max_1, partial_max[partial_index_1]); + local_max_1 = _max_4; + } + float _warp_reduce_4 = local_max_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_4 = max_noftz(_warp_reduce_4, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_4, offset)); + float warp_max_1 = _warp_reduce_4; + if (lane == 0) { + reduce_smem[warp] = warp_max_1; + } + __syncthreads(); + float warp_partial_max_1 = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_5 = warp_partial_max_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_5 = max_noftz(_warp_reduce_5, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_5, offset)); + float merged_max = _warp_reduce_5; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = merged_max; + } + } + __syncthreads(); + float global_max = reduce_smem[0]; + __syncthreads(); + float local_sum_1 = 0.0f; + for (int partial_split_1 = tid; partial_split_1 < splits; partial_split_1 += 256) { + unsigned long long partial_index_2 = (unsigned long long)row_1 * (unsigned long long)splits + (unsigned long long)partial_split_1; + float split_max = partial_max[partial_index_2]; + if (split_max > -CUDART_INF_F) { + float _exp2_13 = approx_exp2((split_max - global_max) * 1.4426950408889634f); + local_sum_1 += partial_sum[partial_index_2] * _exp2_13; + } + } + float _warp_reduce_6 = local_sum_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_6 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_6, offset); + float warp_sum_1 = _warp_reduce_6; + if (lane == 0) { + reduce_smem[warp] = warp_sum_1; + } + __syncthreads(); + float warp_partial_sum_1 = ((lane < 8) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_7 = warp_partial_sum_1; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_7 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_7, offset); + float merged_sum = _warp_reduce_7; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + float _rcp_1 = approx_rcp(merged_sum); + reduce_smem[0] = _rcp_1; + } + } + __syncthreads(); + float inv_sum_1 = reduce_smem[0]; + __syncthreads(); + float temperature_1 = ((parameter_kind == 2) ? parameter[row_1] : scalar_temperature); + int row_base_1 = row_1 * vocab_size; + int start_col_1 = split_1 * 256 + tid; + int col_stride_1 = splits * 256; + int vector_groups_1 = vocab_size / 8; + int first_group_1 = split_1 * vector_groups_1 / splits; + int last_group_1 = (split_1 + 1) * vector_groups_1 / splits; + if (row_base_1 % 8 == 0) { + for (int group_3 = first_group_1 + tid; group_3 < last_group_1; group_3 += 256) { + int col_6 = group_3 * 8; + unsigned long long index_6 = (unsigned long long)(row_base_1 + col_6); + float _vec_load_3[8]; + { + unsigned _ldv8_4_0; + unsigned _ldv8_4_1; + unsigned _ldv8_4_2; + unsigned _ldv8_4_3; + unsigned _ldv8_4_4; + unsigned _ldv8_4_5; + unsigned _ldv8_4_6; + unsigned _ldv8_4_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_4_0), "=r"(_ldv8_4_1), "=r"(_ldv8_4_2), "=r"(_ldv8_4_3), "=r"(_ldv8_4_4), "=r"(_ldv8_4_5), "=r"(_ldv8_4_6), "=r"(_ldv8_4_7) : "l"((const void*)(x + (index_6))) : "memory"); + _vec_load_3[0 + 0] = __uint_as_float(_ldv8_4_0); + _vec_load_3[0 + 1] = __uint_as_float(_ldv8_4_1); + _vec_load_3[0 + 2] = __uint_as_float(_ldv8_4_2); + _vec_load_3[0 + 3] = __uint_as_float(_ldv8_4_3); + _vec_load_3[0 + 4] = __uint_as_float(_ldv8_4_4); + _vec_load_3[0 + 5] = __uint_as_float(_ldv8_4_5); + _vec_load_3[0 + 6] = __uint_as_float(_ldv8_4_6); + _vec_load_3[0 + 7] = __uint_as_float(_ldv8_4_7); + } + for (int j_4 = 0; j_4 < 8; j_4++) { + float scaled_7 = _vec_load_3[j_4] / temperature_1; + float _exp2_14 = approx_exp2((scaled_7 - global_max) * 1.4426950408889634f); + _vec_load_3[j_4] = _exp2_14 * inv_sum_1; + } + { + unsigned _stv8_5_0 = __float_as_uint(_vec_load_3[0 + 0]); + unsigned _stv8_5_1 = __float_as_uint(_vec_load_3[0 + 1]); + unsigned _stv8_5_2 = __float_as_uint(_vec_load_3[0 + 2]); + unsigned _stv8_5_3 = __float_as_uint(_vec_load_3[0 + 3]); + unsigned _stv8_5_4 = __float_as_uint(_vec_load_3[0 + 4]); + unsigned _stv8_5_5 = __float_as_uint(_vec_load_3[0 + 5]); + unsigned _stv8_5_6 = __float_as_uint(_vec_load_3[0 + 6]); + unsigned _stv8_5_7 = __float_as_uint(_vec_load_3[0 + 7]); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(output + (index_6))), "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3), "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) : "memory"); + } + } + int tail_col_3 = vector_groups_1 * 8 + tid; + if (split_1 == splits - 1 && tail_col_3 < vocab_size) { + unsigned long long tail_index_3 = (unsigned long long)(row_base_1 + tail_col_3); + float tail_scaled_3 = x[tail_index_3] / temperature_1; + float _exp2_15 = approx_exp2((tail_scaled_3 - global_max) * 1.4426950408889634f); + output[tail_index_3] = _exp2_15 * inv_sum_1; + } + } else { + for (int col_7 = start_col_1; col_7 < vocab_size; col_7 += col_stride_1) { + unsigned long long index_7 = (unsigned long long)(row_base_1 + col_7); + float scaled_8 = x[index_7] / temperature_1; + float _exp2_16 = approx_exp2((scaled_8 - global_max) * 1.4426950408889634f); + output[index_7] = _exp2_16 * inv_sum_1; + } + } + } +} + +} // extern "C" + +// End exact generated payload. diff --git a/csrc/blackwell_softmax_rowwise.cu b/csrc/blackwell_softmax_rowwise.cu new file mode 100644 index 00000000000..712b78ab6bf --- /dev/null +++ b/csrc/blackwell_softmax_rowwise.cu @@ -0,0 +1,263 @@ +/* + * 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 Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. +// Exact generated payload: 8720 bytes, +// sha256:e47f82923112e143173f32dbb145986f8c9a1503fd2bb71f22f5c065769162e3. +// The sm_100a and sm_103a payloads are byte-identical. + +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#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 LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_REDUCE_SMEM_OFF 0 +#define SMEM_REDUCE_SMEM_STAGE_BYTES 64 +#define SMEM_REDUCE_SMEM_STRIDE 64 +#define SMEM_TOTAL 128 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__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; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_flashinfer_blackwell_softmax_followup_rowwise(float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, int rows, int vocab_size, int parameter_kind, float scalar_temperature) +{ + 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; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* reduce_smem = reinterpret_cast(smem_raw + 0); + const int reduce_smem_addr = smem + 0; + + // === Task calls (dependency order) === + int row = bid; + float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); + int row_base = row * vocab_size; + int row_misalignment = row_base % 4; + int aligned_begin = (4 - row_misalignment) % 4; + int _min_0 = ((aligned_begin) < (vocab_size) ? (aligned_begin) : (vocab_size)); + aligned_begin = _min_0; + int vector_groups = (vocab_size - aligned_begin) / 4; + int aligned_end = aligned_begin + vector_groups * 4; + float local_max = -CUDART_INF_F; + float local_sum = 0.0f; + for (int col = tid; col < aligned_begin; col += 512) { + unsigned long long index = (unsigned long long)(row_base + col); + float scaled = x[index] / temperature; + if (scaled > local_max) { + float _exp2_0 = approx_exp2((local_max - scaled) * 1.4426950408889634f); + local_sum *= _exp2_0; + local_max = scaled; + } + if (scaled > -CUDART_INF_F) { + float _exp2_1 = approx_exp2((scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_1; + } + } + for (int group = tid; group < vector_groups; group += 512) { + int col_1 = aligned_begin + group * 4; + unsigned long long index_1 = (unsigned long long)(row_base + col_1); + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(x + index_1); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float group_max = -CUDART_INF_F; + for (int j = 0; j < 4; j++) { + float scaled_1 = _vec_load_0[j] / temperature; + float _max_0 = max_noftz(group_max, scaled_1); + group_max = _max_0; + } + if (group_max > local_max) { + float _exp2_2 = approx_exp2((local_max - group_max) * 1.4426950408889634f); + local_sum *= _exp2_2; + local_max = group_max; + } + if (group_max > -CUDART_INF_F) { + for (int j_1 = 0; j_1 < 4; j_1++) { + float scaled_2 = _vec_load_0[j_1] / temperature; + if (scaled_2 > -CUDART_INF_F) { + float _exp2_3 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); + local_sum += _exp2_3; + } + } + } + } + int tail_col = aligned_end + tid; + if (tail_col < vocab_size) { + unsigned long long tail_index = (unsigned long long)(row_base + tail_col); + float tail_scaled = x[tail_index] / temperature; + if (tail_scaled > local_max) { + float _exp2_4 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); + local_sum *= _exp2_4; + local_max = tail_scaled; + } + if (tail_scaled > -CUDART_INF_F) { + float _exp2_5 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); + local_sum += _exp2_5; + } + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float warp_max = _warp_reduce_0; + if (lane == 0) { + reduce_smem[warp] = warp_max; + } + __syncthreads(); + float warp_partial_max = ((lane < 16) ? reduce_smem[lane] : -CUDART_INF_F); + float _warp_reduce_1 = warp_partial_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); + float block_max = _warp_reduce_1; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = block_max; + } + } + __syncthreads(); + float cta_max = reduce_smem[0]; + __syncthreads(); + if (local_max > -CUDART_INF_F) { + float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); + local_sum *= _exp2_6; + } + float _warp_reduce_2 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); + float warp_sum = _warp_reduce_2; + if (lane == 0) { + reduce_smem[warp] = warp_sum; + } + __syncthreads(); + float warp_partial_sum = ((lane < 16) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_3 = warp_partial_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); + float block_sum = _warp_reduce_3; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + float _rcp_0 = approx_rcp(block_sum); + reduce_smem[0] = _rcp_0; + } + } + __syncthreads(); + float inv_sum = reduce_smem[0]; + for (int col_2 = tid; col_2 < aligned_begin; col_2 += 512) { + unsigned long long index_2 = (unsigned long long)(row_base + col_2); + float scaled_3 = x[index_2] / temperature; + float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); + output[index_2] = _exp2_7 * inv_sum; + } + for (int group_1 = tid; group_1 < vector_groups; group_1 += 512) { + int col_3 = aligned_begin + group_1 * 4; + unsigned long long index_3 = (unsigned long long)(row_base + col_3); + float _vec_load_1[4]; + { + float4 _v4 = *reinterpret_cast(x + index_3); + _vec_load_1[0 + 0] = _v4.x; + _vec_load_1[0 + 1] = _v4.y; + _vec_load_1[0 + 2] = _v4.z; + _vec_load_1[0 + 3] = _v4.w; + } + for (int j_2 = 0; j_2 < 4; j_2++) { + float scaled_4 = _vec_load_1[j_2] / temperature; + float _exp2_8 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); + _vec_load_1[j_2] = _exp2_8 * inv_sum; + } + { + float4 _v4 = make_float4(_vec_load_1[0 + 0], _vec_load_1[0 + 1], _vec_load_1[0 + 2], _vec_load_1[0 + 3]); + *reinterpret_cast(output + index_3) = _v4; + } + } + if (tail_col < vocab_size) { + unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col); + float tail_scaled_1 = x[tail_index_1] / temperature; + float _exp2_9 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); + output[tail_index_1] = _exp2_9 * inv_sum; + } +} + +} // extern "C" + +// End exact generated payload. diff --git a/csrc/blackwell_softmax_warp.cu b/csrc/blackwell_softmax_warp.cu new file mode 100644 index 00000000000..b56f1b2c1b2 --- /dev/null +++ b/csrc/blackwell_softmax_warp.cu @@ -0,0 +1,133 @@ +/* + * 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 Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. +// Exact generated payload: 3923 bytes, +// sha256:8bc9df2b57d4e6021d0add03110a11f0cbb3bab1218e94d00d92f4bbf597ee9e. +// The sm_100a and sm_103a payloads are byte-identical. + +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#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 LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#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; +} + +extern "C" { + +__global__ __launch_bounds__(128, 4) void +kernel_flashinfer_blackwell_softmax_followup_warp(float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, int rows, int vocab_size, int parameter_kind, float scalar_temperature) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = bid * 4 + warp; + if (row < rows) { + float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); + int row_base = row * vocab_size; + float row_values[9]; + float local_max = -CUDART_INF_F; + for (int row_iteration = 0; row_iteration < 9; row_iteration++) { + int col = lane + row_iteration * 32; + row_values[row_iteration] = -CUDART_INF_F; + if (col < vocab_size) { + unsigned long long index = (unsigned long long)(row_base + col); + float scaled = x[index] / temperature; + row_values[row_iteration] = scaled; + float _max_0 = max_noftz(local_max, scaled); + local_max = _max_0; + } + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + float row_max = _warp_reduce_0; + float local_sum = 0.0f; + for (int row_iteration_1 = 0; row_iteration_1 < 9; row_iteration_1++) { + int col_1 = lane + row_iteration_1 * 32; + if (col_1 < vocab_size) { + float scaled_1 = row_values[row_iteration_1]; + if (scaled_1 > -CUDART_INF_F) { + float _exp2_0 = approx_exp2((scaled_1 - row_max) * 1.4426950408889634f); + local_sum += _exp2_0; + } + } + } + float _warp_reduce_1 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + float row_sum = _warp_reduce_1; + float _rcp_0 = approx_rcp(row_sum); + float inv_sum = _rcp_0; + for (int row_iteration_2 = 0; row_iteration_2 < 9; row_iteration_2++) { + int col_2 = lane + row_iteration_2 * 32; + if (col_2 < vocab_size) { + unsigned long long index_1 = (unsigned long long)(row_base + col_2); + float scaled_2 = row_values[row_iteration_2]; + float _exp2_1 = approx_exp2((scaled_2 - row_max) * 1.4426950408889634f); + output[index_1] = _exp2_1 * inv_sum; + } + } + } +} + +} // extern "C" + +// End exact generated payload. diff --git a/flashinfer/jit/blackwell_softmax.py b/flashinfer/jit/blackwell_softmax.py index de29d5fed53..ada1930e82c 100644 --- a/flashinfer/jit/blackwell_softmax.py +++ b/flashinfer/jit/blackwell_softmax.py @@ -24,6 +24,11 @@ def gen_blackwell_softmax_module() -> JitSpec: ) return gen_jit_spec( "blackwell_softmax", - [jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax.cu"], + [ + jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax.cu", + jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_bootstrap.cu", + jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_rowwise.cu", + jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_warp.cu", + ], extra_cuda_cflags=nvcc_flags, ) diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index 13d6d682cbb..8f4c7cb4bd6 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -15,723 +15,26 @@ */ #pragma once -#include -#include +#include -#include -#include -#include -#include - -// Generated by Loom from Cake commit 25dba320359dc009daf7039d1778a8721681bba8. -// The sm_100a and sm_103a outputs are byte-identical (payload SHA-256: -// d6e46c2babcd74565ce80df3009e3cf2751f773c483b58f0d4676c41bf0dfd44). - -__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; -} - -#include -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - -__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; -} - -#define LOOM_INF CUDART_INF_F -#define NUM_MAIN_STAGES 1 -#define SMEM_REDUCE_SMEM_OFF 0 -#define SMEM_REDUCE_SMEM_STAGE_BYTES 32 -#define SMEM_REDUCE_SMEM_STRIDE 32 -#define SMEM_TOTAL 128 -#define THREADS 256 - -extern "C" { - -__global__ __launch_bounds__(256, 2) void kernel_flashinfer_blackwell_softmax_bootstrap_seed( - float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, - float* __restrict__ partial_max, float* __restrict__ partial_sum, int rows, int vocab_size, - int splits, int parameter_kind, float scalar_temperature) { - 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; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - // Kernel setup ops - float* reduce_smem = reinterpret_cast(smem_raw + 0); - const int reduce_smem_addr = smem + 0; - - // === Task calls (dependency order) === - int total_tasks = rows * splits; - for (int task = bid; task < total_tasks; task += num_bids) { - int row = task / splits; - int split = task - row * splits; - float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); - int row_base = row * vocab_size; - int start_col = split * 256 + tid; - int col_stride = splits * 256; - int vector_groups = vocab_size / 8; - int first_group = split * vector_groups / splits; - int last_group = (split + 1) * vector_groups / splits; - float local_max = -CUDART_INF_F; - float local_sum = 0.0f; - if (row_base % 8 == 0) { - for (int group = first_group + tid; group < last_group; group += 256) { - int col = group * 8; - unsigned long long index = (unsigned long long)(row_base + col); - float _vec_load_0[8]; - { - unsigned _ldv8_0_0; - unsigned _ldv8_0_1; - unsigned _ldv8_0_2; - unsigned _ldv8_0_3; - unsigned _ldv8_0_4; - unsigned _ldv8_0_5; - unsigned _ldv8_0_6; - unsigned _ldv8_0_7; - asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_0_0), "=r"(_ldv8_0_1), "=r"(_ldv8_0_2), "=r"(_ldv8_0_3), - "=r"(_ldv8_0_4), "=r"(_ldv8_0_5), "=r"(_ldv8_0_6), "=r"(_ldv8_0_7) - : "l"((const void*)(x + (index))) - : "memory"); - _vec_load_0[0 + 0] = __uint_as_float(_ldv8_0_0); - _vec_load_0[0 + 1] = __uint_as_float(_ldv8_0_1); - _vec_load_0[0 + 2] = __uint_as_float(_ldv8_0_2); - _vec_load_0[0 + 3] = __uint_as_float(_ldv8_0_3); - _vec_load_0[0 + 4] = __uint_as_float(_ldv8_0_4); - _vec_load_0[0 + 5] = __uint_as_float(_ldv8_0_5); - _vec_load_0[0 + 6] = __uint_as_float(_ldv8_0_6); - _vec_load_0[0 + 7] = __uint_as_float(_ldv8_0_7); - } - float group_max = -CUDART_INF_F; - for (int j = 0; j < 8; j++) { - float scaled = _vec_load_0[j] / temperature; - float _max_0 = max_noftz(group_max, scaled); - group_max = _max_0; - } - if (splits == 1) { - if (group_max > local_max) { - float _exp2_0 = approx_exp2((local_max - group_max) * 1.4426950408889634f); - local_sum *= _exp2_0; - local_max = group_max; - } - if (group_max > -CUDART_INF_F) { - for (int j_1 = 0; j_1 < 8; j_1++) { - float scaled_1 = _vec_load_0[j_1] / temperature; - if (scaled_1 > -CUDART_INF_F) { - float _exp2_1 = approx_exp2((scaled_1 - local_max) * 1.4426950408889634f); - local_sum += _exp2_1; - } - } - } - } else { - float _max_1 = max_noftz(local_max, group_max); - local_max = _max_1; - } - } - int tail_col = vector_groups * 8 + tid; - if (split == splits - 1 && tail_col < vocab_size) { - unsigned long long tail_index = (unsigned long long)(row_base + tail_col); - float tail_scaled = x[tail_index] / temperature; - if (splits == 1) { - if (tail_scaled > local_max) { - float _exp2_2 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); - local_sum *= _exp2_2; - local_max = tail_scaled; - } - if (tail_scaled > -CUDART_INF_F) { - float _exp2_3 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); - local_sum += _exp2_3; - } - } else { - float _max_2 = max_noftz(local_max, tail_scaled); - local_max = _max_2; - } - } - } else { - for (int col_1 = start_col; col_1 < vocab_size; col_1 += col_stride) { - unsigned long long index_1 = (unsigned long long)(row_base + col_1); - float scaled_2 = x[index_1] / temperature; - if (splits == 1) { - if (scaled_2 > local_max) { - float _exp2_4 = approx_exp2((local_max - scaled_2) * 1.4426950408889634f); - local_sum *= _exp2_4; - local_max = scaled_2; - } - if (scaled_2 > -CUDART_INF_F) { - float _exp2_5 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); - local_sum += _exp2_5; - } - } else { - float _max_3 = max_noftz(local_max, scaled_2); - local_max = _max_3; - } - } - } - float _warp_reduce_0 = local_max; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_0 = - max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); - float warp_max = _warp_reduce_0; - if (lane == 0) { - reduce_smem[warp] = warp_max; - } - __syncthreads(); - float warp_partial_max = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); - float _warp_reduce_1 = warp_partial_max; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_1 = - max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); - float block_max = _warp_reduce_1; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - reduce_smem[0] = block_max; - } - } - __syncthreads(); - float cta_max = reduce_smem[0]; - __syncthreads(); - if (splits == 1) { - if (local_max > -CUDART_INF_F) { - float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); - local_sum *= _exp2_6; - } - } else if (cta_max > -CUDART_INF_F) { - if (row_base % 8 == 0) { - for (int group_1 = first_group + tid; group_1 < last_group; group_1 += 256) { - int col_2 = group_1 * 8; - unsigned long long index_2 = (unsigned long long)(row_base + col_2); - float _vec_load_1[8]; - { - unsigned _ldv8_1_0; - unsigned _ldv8_1_1; - unsigned _ldv8_1_2; - unsigned _ldv8_1_3; - unsigned _ldv8_1_4; - unsigned _ldv8_1_5; - unsigned _ldv8_1_6; - unsigned _ldv8_1_7; - asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), - "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) - : "l"((const void*)(x + (index_2))) - : "memory"); - _vec_load_1[0 + 0] = __uint_as_float(_ldv8_1_0); - _vec_load_1[0 + 1] = __uint_as_float(_ldv8_1_1); - _vec_load_1[0 + 2] = __uint_as_float(_ldv8_1_2); - _vec_load_1[0 + 3] = __uint_as_float(_ldv8_1_3); - _vec_load_1[0 + 4] = __uint_as_float(_ldv8_1_4); - _vec_load_1[0 + 5] = __uint_as_float(_ldv8_1_5); - _vec_load_1[0 + 6] = __uint_as_float(_ldv8_1_6); - _vec_load_1[0 + 7] = __uint_as_float(_ldv8_1_7); - } - for (int j_2 = 0; j_2 < 8; j_2++) { - float scaled_3 = _vec_load_1[j_2] / temperature; - float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); - local_sum += _exp2_7; - } - } - int tail_col_1 = vector_groups * 8 + tid; - if (split == splits - 1 && tail_col_1 < vocab_size) { - unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col_1); - float tail_scaled_1 = x[tail_index_1] / temperature; - float _exp2_8 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); - local_sum += _exp2_8; - } - } else { - for (int col_3 = start_col; col_3 < vocab_size; col_3 += col_stride) { - unsigned long long index_3 = (unsigned long long)(row_base + col_3); - float scaled_4 = x[index_3] / temperature; - float _exp2_9 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); - local_sum += _exp2_9; - } - } - } - float _warp_reduce_2 = local_sum; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); - float warp_sum = _warp_reduce_2; - if (lane == 0) { - reduce_smem[warp] = warp_sum; - } - __syncthreads(); - float warp_partial_sum = ((lane < 8) ? reduce_smem[lane] : 0.0f); - float _warp_reduce_3 = warp_partial_sum; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); - float block_sum = _warp_reduce_3; - __syncthreads(); - if (splits == 1) { - float _rcp_0 = approx_rcp(block_sum); - float inv_sum = _rcp_0; - if (row_base % 8 == 0) { - for (int group_2 = first_group + tid; group_2 < last_group; group_2 += 256) { - int col_4 = group_2 * 8; - unsigned long long index_4 = (unsigned long long)(row_base + col_4); - float _vec_load_2[8]; - { - unsigned _ldv8_2_0; - unsigned _ldv8_2_1; - unsigned _ldv8_2_2; - unsigned _ldv8_2_3; - unsigned _ldv8_2_4; - unsigned _ldv8_2_5; - unsigned _ldv8_2_6; - unsigned _ldv8_2_7; - asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), - "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) - : "l"((const void*)(x + (index_4))) - : "memory"); - _vec_load_2[0 + 0] = __uint_as_float(_ldv8_2_0); - _vec_load_2[0 + 1] = __uint_as_float(_ldv8_2_1); - _vec_load_2[0 + 2] = __uint_as_float(_ldv8_2_2); - _vec_load_2[0 + 3] = __uint_as_float(_ldv8_2_3); - _vec_load_2[0 + 4] = __uint_as_float(_ldv8_2_4); - _vec_load_2[0 + 5] = __uint_as_float(_ldv8_2_5); - _vec_load_2[0 + 6] = __uint_as_float(_ldv8_2_6); - _vec_load_2[0 + 7] = __uint_as_float(_ldv8_2_7); - } - for (int j_3 = 0; j_3 < 8; j_3++) { - float scaled_5 = _vec_load_2[j_3] / temperature; - float _exp2_10 = approx_exp2((scaled_5 - cta_max) * 1.4426950408889634f); - _vec_load_2[j_3] = _exp2_10 * inv_sum; - } - { - unsigned _stv8_3_0 = __float_as_uint(_vec_load_2[0 + 0]); - unsigned _stv8_3_1 = __float_as_uint(_vec_load_2[0 + 1]); - unsigned _stv8_3_2 = __float_as_uint(_vec_load_2[0 + 2]); - unsigned _stv8_3_3 = __float_as_uint(_vec_load_2[0 + 3]); - unsigned _stv8_3_4 = __float_as_uint(_vec_load_2[0 + 4]); - unsigned _stv8_3_5 = __float_as_uint(_vec_load_2[0 + 5]); - unsigned _stv8_3_6 = __float_as_uint(_vec_load_2[0 + 6]); - unsigned _stv8_3_7 = __float_as_uint(_vec_load_2[0 + 7]); - asm volatile("st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" ::"l"( - (void*)(output + (index_4))), - "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), - "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) - : "memory"); - } - } - int tail_col_2 = vector_groups * 8 + tid; - if (tail_col_2 < vocab_size) { - unsigned long long tail_index_2 = (unsigned long long)(row_base + tail_col_2); - float tail_scaled_2 = x[tail_index_2] / temperature; - float _exp2_11 = approx_exp2((tail_scaled_2 - cta_max) * 1.4426950408889634f); - output[tail_index_2] = _exp2_11 * inv_sum; - } - } else { - for (int col_5 = start_col; col_5 < vocab_size; col_5 += col_stride) { - unsigned long long index_5 = (unsigned long long)(row_base + col_5); - float scaled_6 = x[index_5] / temperature; - float _exp2_12 = approx_exp2((scaled_6 - cta_max) * 1.4426950408889634f); - output[index_5] = _exp2_12 * inv_sum; - } - } - __syncthreads(); - } else if (warp == 0) { - if (elect_sync()) { - unsigned long long partial_index = (unsigned long long)task; - partial_max[partial_index] = cta_max; - partial_sum[partial_index] = block_sum; - } - } - } - if (splits > 1) { - __threadfence(); - cooperative_groups::this_grid().sync(); - } - int phase_2_tasks = ((splits > 1) ? total_tasks : 0); - for (int task_1 = bid; task_1 < phase_2_tasks; task_1 += num_bids) { - int row_1 = task_1 / splits; - int split_1 = task_1 - row_1 * splits; - float local_max_1 = -CUDART_INF_F; - for (int partial_split = tid; partial_split < splits; partial_split += 256) { - unsigned long long partial_index_1 = (unsigned long long)row_1 * (unsigned long long)splits + - (unsigned long long)partial_split; - float _max_4 = max_noftz(local_max_1, partial_max[partial_index_1]); - local_max_1 = _max_4; - } - float _warp_reduce_4 = local_max_1; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_4 = - max_noftz(_warp_reduce_4, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_4, offset)); - float warp_max_1 = _warp_reduce_4; - if (lane == 0) { - reduce_smem[warp] = warp_max_1; - } - __syncthreads(); - float warp_partial_max_1 = ((lane < 8) ? reduce_smem[lane] : -CUDART_INF_F); - float _warp_reduce_5 = warp_partial_max_1; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_5 = - max_noftz(_warp_reduce_5, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_5, offset)); - float merged_max = _warp_reduce_5; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - reduce_smem[0] = merged_max; - } - } - __syncthreads(); - float global_max = reduce_smem[0]; - __syncthreads(); - float local_sum_1 = 0.0f; - for (int partial_split_1 = tid; partial_split_1 < splits; partial_split_1 += 256) { - unsigned long long partial_index_2 = (unsigned long long)row_1 * (unsigned long long)splits + - (unsigned long long)partial_split_1; - float split_max = partial_max[partial_index_2]; - if (split_max > -CUDART_INF_F) { - float _exp2_13 = approx_exp2((split_max - global_max) * 1.4426950408889634f); - local_sum_1 += partial_sum[partial_index_2] * _exp2_13; - } - } - float _warp_reduce_6 = local_sum_1; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_6 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_6, offset); - float warp_sum_1 = _warp_reduce_6; - if (lane == 0) { - reduce_smem[warp] = warp_sum_1; - } - __syncthreads(); - float warp_partial_sum_1 = ((lane < 8) ? reduce_smem[lane] : 0.0f); - float _warp_reduce_7 = warp_partial_sum_1; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_7 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_7, offset); - float merged_sum = _warp_reduce_7; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - float _rcp_1 = approx_rcp(merged_sum); - reduce_smem[0] = _rcp_1; - } - } - __syncthreads(); - float inv_sum_1 = reduce_smem[0]; - __syncthreads(); - float temperature_1 = ((parameter_kind == 2) ? parameter[row_1] : scalar_temperature); - int row_base_1 = row_1 * vocab_size; - int start_col_1 = split_1 * 256 + tid; - int col_stride_1 = splits * 256; - int vector_groups_1 = vocab_size / 8; - int first_group_1 = split_1 * vector_groups_1 / splits; - int last_group_1 = (split_1 + 1) * vector_groups_1 / splits; - if (row_base_1 % 8 == 0) { - for (int group_3 = first_group_1 + tid; group_3 < last_group_1; group_3 += 256) { - int col_6 = group_3 * 8; - unsigned long long index_6 = (unsigned long long)(row_base_1 + col_6); - float _vec_load_3[8]; - { - unsigned _ldv8_4_0; - unsigned _ldv8_4_1; - unsigned _ldv8_4_2; - unsigned _ldv8_4_3; - unsigned _ldv8_4_4; - unsigned _ldv8_4_5; - unsigned _ldv8_4_6; - unsigned _ldv8_4_7; - asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=r"(_ldv8_4_0), "=r"(_ldv8_4_1), "=r"(_ldv8_4_2), "=r"(_ldv8_4_3), - "=r"(_ldv8_4_4), "=r"(_ldv8_4_5), "=r"(_ldv8_4_6), "=r"(_ldv8_4_7) - : "l"((const void*)(x + (index_6))) - : "memory"); - _vec_load_3[0 + 0] = __uint_as_float(_ldv8_4_0); - _vec_load_3[0 + 1] = __uint_as_float(_ldv8_4_1); - _vec_load_3[0 + 2] = __uint_as_float(_ldv8_4_2); - _vec_load_3[0 + 3] = __uint_as_float(_ldv8_4_3); - _vec_load_3[0 + 4] = __uint_as_float(_ldv8_4_4); - _vec_load_3[0 + 5] = __uint_as_float(_ldv8_4_5); - _vec_load_3[0 + 6] = __uint_as_float(_ldv8_4_6); - _vec_load_3[0 + 7] = __uint_as_float(_ldv8_4_7); - } - for (int j_4 = 0; j_4 < 8; j_4++) { - float scaled_7 = _vec_load_3[j_4] / temperature_1; - float _exp2_14 = approx_exp2((scaled_7 - global_max) * 1.4426950408889634f); - _vec_load_3[j_4] = _exp2_14 * inv_sum_1; - } - { - unsigned _stv8_5_0 = __float_as_uint(_vec_load_3[0 + 0]); - unsigned _stv8_5_1 = __float_as_uint(_vec_load_3[0 + 1]); - unsigned _stv8_5_2 = __float_as_uint(_vec_load_3[0 + 2]); - unsigned _stv8_5_3 = __float_as_uint(_vec_load_3[0 + 3]); - unsigned _stv8_5_4 = __float_as_uint(_vec_load_3[0 + 4]); - unsigned _stv8_5_5 = __float_as_uint(_vec_load_3[0 + 5]); - unsigned _stv8_5_6 = __float_as_uint(_vec_load_3[0 + 6]); - unsigned _stv8_5_7 = __float_as_uint(_vec_load_3[0 + 7]); - asm volatile("st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" ::"l"( - (void*)(output + (index_6))), - "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3), - "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) - : "memory"); - } - } - int tail_col_3 = vector_groups_1 * 8 + tid; - if (split_1 == splits - 1 && tail_col_3 < vocab_size) { - unsigned long long tail_index_3 = (unsigned long long)(row_base_1 + tail_col_3); - float tail_scaled_3 = x[tail_index_3] / temperature_1; - float _exp2_15 = approx_exp2((tail_scaled_3 - global_max) * 1.4426950408889634f); - output[tail_index_3] = _exp2_15 * inv_sum_1; - } - } else { - for (int col_7 = start_col_1; col_7 < vocab_size; col_7 += col_stride_1) { - unsigned long long index_7 = (unsigned long long)(row_base_1 + col_7); - float scaled_8 = x[index_7] / temperature_1; - float _exp2_16 = approx_exp2((scaled_8 - global_max) * 1.4426950408889634f); - output[index_7] = _exp2_16 * inv_sum_1; - } - } - } -} - -} // extern "C" - -#undef LOOM_INF -#undef NUM_MAIN_STAGES -#undef SMEM_REDUCE_SMEM_OFF -#undef SMEM_REDUCE_SMEM_STAGE_BYTES -#undef SMEM_REDUCE_SMEM_STRIDE -#undef SMEM_TOTAL -#undef THREADS -#undef reduce_smem_addr - -#define LOOM_INF CUDART_INF_F -#define NUM_MAIN_STAGES 1 -#define SMEM_REDUCE_SMEM_OFF 0 -#define SMEM_REDUCE_SMEM_STAGE_BYTES 64 -#define SMEM_REDUCE_SMEM_STRIDE 64 -#define SMEM_TOTAL 128 -#define THREADS 512 +// Frozen from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. +// Weave sm_100a and sm_103a output is byte-identical for all three kernels: +// bootstrap: sha256:46049370dbd905ff7234d414745d197f883157c8edc01d83230562b4dff5f862 +// rowwise: sha256:e47f82923112e143173f32dbb145986f8c9a1503fd2bb71f22f5c065769162e3 +// warp: sha256:8bc9df2b57d4e6021d0add03110a11f0cbb3bab1218e94d00d92f4bbf597ee9e extern "C" { -__global__ __launch_bounds__(512, 1) void kernel_flashinfer_blackwell_softmax_followup_rowwise( - float* __restrict__ x, float* __restrict__ parameter, float* __restrict__ output, int rows, - int vocab_size, int parameter_kind, float scalar_temperature) { - 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; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +__global__ void kernel_flashinfer_blackwell_softmax_bootstrap_seed( + float* logits, float* parameter, float* output, float* partial_max, float* partial_sum, int rows, + int vocab_size, int splits, int parameter_kind, float scalar_temperature); - const int bid = blockIdx.x; - const int num_bids = gridDim.x; +__global__ void kernel_flashinfer_blackwell_softmax_followup_rowwise( + float* logits, float* parameter, float* output, int rows, int vocab_size, int parameter_kind, + float scalar_temperature); - // Kernel setup ops - float* reduce_smem = reinterpret_cast(smem_raw + 0); - const int reduce_smem_addr = smem + 0; - - // === Task calls (dependency order) === - int row = bid; - float temperature = ((parameter_kind == 2) ? parameter[row] : scalar_temperature); - int row_base = row * vocab_size; - int row_misalignment = row_base % 4; - int aligned_begin = (4 - row_misalignment) % 4; - int _min_0 = ((aligned_begin) < (vocab_size) ? (aligned_begin) : (vocab_size)); - aligned_begin = _min_0; - int vector_groups = (vocab_size - aligned_begin) / 4; - int aligned_end = aligned_begin + vector_groups * 4; - float local_max = -CUDART_INF_F; - float local_sum = 0.0f; - for (int col = tid; col < aligned_begin; col += 512) { - unsigned long long index = (unsigned long long)(row_base + col); - float scaled = x[index] / temperature; - if (scaled > local_max) { - float _exp2_0 = approx_exp2((local_max - scaled) * 1.4426950408889634f); - local_sum *= _exp2_0; - local_max = scaled; - } - if (scaled > -CUDART_INF_F) { - float _exp2_1 = approx_exp2((scaled - local_max) * 1.4426950408889634f); - local_sum += _exp2_1; - } - } - for (int group = tid; group < vector_groups; group += 512) { - int col_1 = aligned_begin + group * 4; - unsigned long long index_1 = (unsigned long long)(row_base + col_1); - float _vec_load_0[4]; - { - float4 _v4 = *reinterpret_cast(x + index_1); - _vec_load_0[0 + 0] = _v4.x; - _vec_load_0[0 + 1] = _v4.y; - _vec_load_0[0 + 2] = _v4.z; - _vec_load_0[0 + 3] = _v4.w; - } - float group_max = -CUDART_INF_F; - for (int j = 0; j < 4; j++) { - float scaled_1 = _vec_load_0[j] / temperature; - float _max_0 = max_noftz(group_max, scaled_1); - group_max = _max_0; - } - if (group_max > local_max) { - float _exp2_2 = approx_exp2((local_max - group_max) * 1.4426950408889634f); - local_sum *= _exp2_2; - local_max = group_max; - } - if (group_max > -CUDART_INF_F) { - for (int j_1 = 0; j_1 < 4; j_1++) { - float scaled_2 = _vec_load_0[j_1] / temperature; - if (scaled_2 > -CUDART_INF_F) { - float _exp2_3 = approx_exp2((scaled_2 - local_max) * 1.4426950408889634f); - local_sum += _exp2_3; - } - } - } - } - int tail_col = aligned_end + tid; - if (tail_col < vocab_size) { - unsigned long long tail_index = (unsigned long long)(row_base + tail_col); - float tail_scaled = x[tail_index] / temperature; - if (tail_scaled > local_max) { - float _exp2_4 = approx_exp2((local_max - tail_scaled) * 1.4426950408889634f); - local_sum *= _exp2_4; - local_max = tail_scaled; - } - if (tail_scaled > -CUDART_INF_F) { - float _exp2_5 = approx_exp2((tail_scaled - local_max) * 1.4426950408889634f); - local_sum += _exp2_5; - } - } - float _warp_reduce_0 = local_max; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); - float warp_max = _warp_reduce_0; - if (lane == 0) { - reduce_smem[warp] = warp_max; - } - __syncthreads(); - float warp_partial_max = ((lane < 16) ? reduce_smem[lane] : -CUDART_INF_F); - float _warp_reduce_1 = warp_partial_max; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); - float block_max = _warp_reduce_1; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - reduce_smem[0] = block_max; - } - } - __syncthreads(); - float cta_max = reduce_smem[0]; - __syncthreads(); - if (local_max > -CUDART_INF_F) { - float _exp2_6 = approx_exp2((local_max - cta_max) * 1.4426950408889634f); - local_sum *= _exp2_6; - } - float _warp_reduce_2 = local_sum; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); - float warp_sum = _warp_reduce_2; - if (lane == 0) { - reduce_smem[warp] = warp_sum; - } - __syncthreads(); - float warp_partial_sum = ((lane < 16) ? reduce_smem[lane] : 0.0f); - float _warp_reduce_3 = warp_partial_sum; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); - float block_sum = _warp_reduce_3; - __syncthreads(); - if (warp == 0) { - if (elect_sync()) { - float _rcp_0 = approx_rcp(block_sum); - reduce_smem[0] = _rcp_0; - } - } - __syncthreads(); - float inv_sum = reduce_smem[0]; - for (int col_2 = tid; col_2 < aligned_begin; col_2 += 512) { - unsigned long long index_2 = (unsigned long long)(row_base + col_2); - float scaled_3 = x[index_2] / temperature; - float _exp2_7 = approx_exp2((scaled_3 - cta_max) * 1.4426950408889634f); - output[index_2] = _exp2_7 * inv_sum; - } - for (int group_1 = tid; group_1 < vector_groups; group_1 += 512) { - int col_3 = aligned_begin + group_1 * 4; - unsigned long long index_3 = (unsigned long long)(row_base + col_3); - float _vec_load_1[4]; - { - float4 _v4 = *reinterpret_cast(x + index_3); - _vec_load_1[0 + 0] = _v4.x; - _vec_load_1[0 + 1] = _v4.y; - _vec_load_1[0 + 2] = _v4.z; - _vec_load_1[0 + 3] = _v4.w; - } - for (int j_2 = 0; j_2 < 4; j_2++) { - float scaled_4 = _vec_load_1[j_2] / temperature; - float _exp2_8 = approx_exp2((scaled_4 - cta_max) * 1.4426950408889634f); - _vec_load_1[j_2] = _exp2_8 * inv_sum; - } - { - float4 _v4 = make_float4(_vec_load_1[0 + 0], _vec_load_1[0 + 1], _vec_load_1[0 + 2], - _vec_load_1[0 + 3]); - *reinterpret_cast(output + index_3) = _v4; - } - } - if (tail_col < vocab_size) { - unsigned long long tail_index_1 = (unsigned long long)(row_base + tail_col); - float tail_scaled_1 = x[tail_index_1] / temperature; - float _exp2_9 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); - output[tail_index_1] = _exp2_9 * inv_sum; - } -} +__global__ void kernel_flashinfer_blackwell_softmax_followup_warp( + float* logits, float* parameter, float* output, int rows, int vocab_size, int parameter_kind, + float scalar_temperature); } // extern "C" - -#undef LOOM_INF -#undef NUM_MAIN_STAGES -#undef SMEM_REDUCE_SMEM_OFF -#undef SMEM_REDUCE_SMEM_STAGE_BYTES -#undef SMEM_REDUCE_SMEM_STRIDE -#undef SMEM_TOTAL -#undef THREADS -#undef reduce_smem_addr diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index 8e6516f8d5f..57d0fe5b3eb 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -81,8 +81,8 @@ def test_softmax( [ (1, 32000, "none"), # cooperative route (256, 32000, "none"), # rowwise route - (1, 111, "scalar"), # scalar-temperature rowwise route - (1, 111, "per_row"), # per-row-temperature rowwise route + (1, 111, "scalar"), # scalar-temperature warp-packed route + (1, 111, "per_row"), # per-row-temperature warp-packed route ], ) def test_softmax_blackwell_routes(batch_size, vocab_size, temperature_kind): From b82c54b10f0d3d5af021a858adad458f8725fbab Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 30 Jul 2026 06:38:02 -0700 Subject: [PATCH 06/14] Preserve generated Softmax payload formatting --- csrc/blackwell_softmax_bootstrap.cu | 2 ++ csrc/blackwell_softmax_rowwise.cu | 2 ++ csrc/blackwell_softmax_warp.cu | 2 ++ 3 files changed, 6 insertions(+) diff --git a/csrc/blackwell_softmax_bootstrap.cu b/csrc/blackwell_softmax_bootstrap.cu index 4543241e9ff..daa78f3b631 100644 --- a/csrc/blackwell_softmax_bootstrap.cu +++ b/csrc/blackwell_softmax_bootstrap.cu @@ -18,6 +18,7 @@ // sha256:46049370dbd905ff7234d414745d197f883157c8edc01d83230562b4dff5f862. // The sm_100a and sm_103a payloads are byte-identical. +// clang-format off typedef unsigned char uint8_t; typedef unsigned short uint16_t; typedef unsigned int uint32_t; @@ -516,4 +517,5 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* } // extern "C" +// clang-format on // End exact generated payload. diff --git a/csrc/blackwell_softmax_rowwise.cu b/csrc/blackwell_softmax_rowwise.cu index 712b78ab6bf..1b1a3fc93f4 100644 --- a/csrc/blackwell_softmax_rowwise.cu +++ b/csrc/blackwell_softmax_rowwise.cu @@ -18,6 +18,7 @@ // sha256:e47f82923112e143173f32dbb145986f8c9a1503fd2bb71f22f5c065769162e3. // The sm_100a and sm_103a payloads are byte-identical. +// clang-format off typedef unsigned char uint8_t; typedef unsigned short uint16_t; typedef unsigned int uint32_t; @@ -260,4 +261,5 @@ kernel_flashinfer_blackwell_softmax_followup_rowwise(float* __restrict__ x, floa } // extern "C" +// clang-format on // End exact generated payload. diff --git a/csrc/blackwell_softmax_warp.cu b/csrc/blackwell_softmax_warp.cu index b56f1b2c1b2..ec9da376e8d 100644 --- a/csrc/blackwell_softmax_warp.cu +++ b/csrc/blackwell_softmax_warp.cu @@ -18,6 +18,7 @@ // sha256:8bc9df2b57d4e6021d0add03110a11f0cbb3bab1218e94d00d92f4bbf597ee9e. // The sm_100a and sm_103a payloads are byte-identical. +// clang-format off typedef unsigned char uint8_t; typedef unsigned short uint16_t; typedef unsigned int uint32_t; @@ -130,4 +131,5 @@ kernel_flashinfer_blackwell_softmax_followup_warp(float* __restrict__ x, float* } // extern "C" +// clang-format on // End exact generated payload. From 657d533c9ee3f9c24f8ab4287eb94de6ac111761 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 30 Jul 2026 06:42:55 -0700 Subject: [PATCH 07/14] Sync final Loom Softmax dispatcher --- csrc/blackwell_softmax.cu | 9 +++++---- include/flashinfer/blackwell_softmax.cuh | 3 ++- tests/utils/test_sampling.py | 3 +++ 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu index f1173ea0cd0..ec3ec567c43 100644 --- a/csrc/blackwell_softmax.cu +++ b/csrc/blackwell_softmax.cu @@ -53,13 +53,14 @@ bool use_rowwise_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parame vocab_size <= 256000 && vocab_size % 4 == 0 && parameter_kind == ParameterKind::kNone; const bool dense_aligned_high_row_narrow = rows > 384 && rows <= 1024 && vocab_size >= 24576 && - vocab_size <= 64000 && vocab_size % 4 == 0 && - (parameter_kind == ParameterKind::kNone || - vocab_size <= 32000); + vocab_size <= 32000 && vocab_size % 4 == 0; + const bool dense_aligned_high_row_wide = rows > 512 && rows <= 1024 && vocab_size > 32000 && + vocab_size <= 64000 && vocab_size % 4 == 0 && + parameter_kind == ParameterKind::kNone; const bool measured_large_odd = rows > 128 && rows <= 512 && vocab_size >= 24576 && vocab_size <= 131072 && vocab_size % 4 != 0; return small_low_row || dense_aligned_mid_row || dense_aligned_high_row_narrow || - measured_large_odd; + dense_aligned_high_row_wide || measured_large_odd; } cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temperature_arr, diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index 8f4c7cb4bd6..c3594b310d7 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -17,7 +17,8 @@ #include -// Frozen from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. +// Frozen from Cake measured checkpoint 924403cd62ea6dc160ffcea63c944d1e828d393a +// and retained byte-for-byte by MR474 cleanup d77b50ee1e56fb3701a512f1dd2765e6067b1be2. // Weave sm_100a and sm_103a output is byte-identical for all three kernels: // bootstrap: sha256:46049370dbd905ff7234d414745d197f883157c8edc01d83230562b4dff5f862 // rowwise: sha256:e47f82923112e143173f32dbb145986f8c9a1503fd2bb71f22f5c065769162e3 diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index 57d0fe5b3eb..755a5d62a8a 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -81,6 +81,9 @@ def test_softmax( [ (1, 32000, "none"), # cooperative route (256, 32000, "none"), # rowwise route + (512, 64000, "none"), # bootstrap side of the wide-row boundary + (1024, 64000, "none"), # rowwise side of the wide-row boundary + (989, 128256, "per_row"), # full-sweep worst-row profile (1, 111, "scalar"), # scalar-temperature warp-packed route (1, 111, "per_row"), # per-row-temperature warp-packed route ], From bc6f33412fcdc089c8ebdf4dfa9f42b5233551e3 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 30 Jul 2026 06:56:39 -0700 Subject: [PATCH 08/14] Use standard integer types in exported kernels --- csrc/blackwell_softmax_bootstrap.cu | 12 ++---------- csrc/blackwell_softmax_rowwise.cu | 12 ++---------- csrc/blackwell_softmax_warp.cu | 12 ++---------- 3 files changed, 6 insertions(+), 30 deletions(-) diff --git a/csrc/blackwell_softmax_bootstrap.cu b/csrc/blackwell_softmax_bootstrap.cu index daa78f3b631..5db764c5d63 100644 --- a/csrc/blackwell_softmax_bootstrap.cu +++ b/csrc/blackwell_softmax_bootstrap.cu @@ -14,20 +14,12 @@ * limitations under the License. */ // Generated by Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. -// Exact generated payload: 24102 bytes, +// Original generated payload before FlashInfer type adaptation: 24102 bytes, // sha256:46049370dbd905ff7234d414745d197f883157c8edc01d83230562b4dff5f862. // The sm_100a and sm_103a payloads are byte-identical. // clang-format off -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; - -typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; - +#include #include __device__ __forceinline__ int make_warp_uniform(int x) { diff --git a/csrc/blackwell_softmax_rowwise.cu b/csrc/blackwell_softmax_rowwise.cu index 1b1a3fc93f4..04a1475b574 100644 --- a/csrc/blackwell_softmax_rowwise.cu +++ b/csrc/blackwell_softmax_rowwise.cu @@ -14,20 +14,12 @@ * limitations under the License. */ // Generated by Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. -// Exact generated payload: 8720 bytes, +// Original generated payload before FlashInfer type adaptation: 8720 bytes, // sha256:e47f82923112e143173f32dbb145986f8c9a1503fd2bb71f22f5c065769162e3. // The sm_100a and sm_103a payloads are byte-identical. // clang-format off -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; - -typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; - +#include #include __device__ __forceinline__ int make_warp_uniform(int x) { diff --git a/csrc/blackwell_softmax_warp.cu b/csrc/blackwell_softmax_warp.cu index ec9da376e8d..34d587b2520 100644 --- a/csrc/blackwell_softmax_warp.cu +++ b/csrc/blackwell_softmax_warp.cu @@ -14,20 +14,12 @@ * limitations under the License. */ // Generated by Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. -// Exact generated payload: 3923 bytes, +// Original generated payload before FlashInfer type adaptation: 3923 bytes, // sha256:8bc9df2b57d4e6021d0add03110a11f0cbb3bab1218e94d00d92f4bbf597ee9e. // The sm_100a and sm_103a payloads are byte-identical. // clang-format off -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef signed int int32_t; -typedef short int int16_t; - -typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; - +#include #include __device__ __forceinline__ int make_warp_uniform(int x) { From 754aee21132ee79b69d1fb328daca0cdfd671b63 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 30 Jul 2026 15:42:17 -0700 Subject: [PATCH 09/14] style: apply pre-commit formatting --- csrc/blackwell_softmax.cu | 7 +++---- include/flashinfer/blackwell_softmax.cuh | 12 +++++++----- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu index ec3ec567c43..0f14e20bc8c 100644 --- a/csrc/blackwell_softmax.cu +++ b/csrc/blackwell_softmax.cu @@ -43,8 +43,7 @@ enum class ParameterKind : int { bool use_warp_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind) { return rows <= 128 && vocab_size <= 257 && - (parameter_kind == ParameterKind::kScalar || - parameter_kind == ParameterKind::kPerRow); + (parameter_kind == ParameterKind::kScalar || parameter_kind == ParameterKind::kPerRow); } bool use_rowwise_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind) { @@ -139,8 +138,8 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper &rows_i, &vocab_size_i, &splits_i, ¶meter_kind_i, &temperature_val}; return cudaLaunchCooperativeKernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_bootstrap_seed), - dim3(static_cast(grid)), dim3(kBootstrapThreads), args, - kBootstrapDynamicSmemBytes, stream); + dim3(static_cast(grid)), dim3(kBootstrapThreads), args, kBootstrapDynamicSmemBytes, + stream); } } // namespace diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index c3594b310d7..cf19b372954 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -27,15 +27,17 @@ extern "C" { __global__ void kernel_flashinfer_blackwell_softmax_bootstrap_seed( - float* logits, float* parameter, float* output, float* partial_max, float* partial_sum, int rows, - int vocab_size, int splits, int parameter_kind, float scalar_temperature); + float* logits, float* parameter, float* output, float* partial_max, float* partial_sum, + int rows, int vocab_size, int splits, int parameter_kind, float scalar_temperature); __global__ void kernel_flashinfer_blackwell_softmax_followup_rowwise( float* logits, float* parameter, float* output, int rows, int vocab_size, int parameter_kind, float scalar_temperature); -__global__ void kernel_flashinfer_blackwell_softmax_followup_warp( - float* logits, float* parameter, float* output, int rows, int vocab_size, int parameter_kind, - float scalar_temperature); +__global__ void kernel_flashinfer_blackwell_softmax_followup_warp(float* logits, float* parameter, + float* output, int rows, + int vocab_size, + int parameter_kind, + float scalar_temperature); } // extern "C" From b875025b0a3c67e8627452f956285aadc2fb96cd Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 5 Aug 2026 15:01:18 -0700 Subject: [PATCH 10/14] Preserve MR515 Softmax wins in the PR4282 dispatcher --- csrc/blackwell_softmax.cu | 250 ++++++++++++++++++++--- csrc/blackwell_softmax_bootstrap.cu | 72 ++++++- csrc/blackwell_softmax_mr515_exp2.cu | 233 +++++++++++++++++++++ csrc/blackwell_softmax_rowwise.cu | 17 +- csrc/blackwell_softmax_warp.cu | 17 +- flashinfer/jit/blackwell_softmax.py | 1 + flashinfer/sampling.py | 124 ++++++++++- include/flashinfer/blackwell_softmax.cuh | 16 +- tests/utils/test_sampling.py | 138 +++++++++++++ 9 files changed, 823 insertions(+), 45 deletions(-) create mode 100644 csrc/blackwell_softmax_mr515_exp2.cu diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu index 0f14e20bc8c..ad9e39ba748 100644 --- a/csrc/blackwell_softmax.cu +++ b/csrc/blackwell_softmax.cu @@ -29,11 +29,13 @@ namespace { constexpr int kBootstrapThreads = 256; constexpr int kRowwiseThreads = 512; constexpr int kWarpThreads = 128; +constexpr int kMr515Threads = 512; constexpr int kWarpRowsPerCta = 4; constexpr int kMaxSplits = 64; constexpr size_t kBootstrapDynamicSmemBytes = 128; constexpr size_t kRowwiseDynamicSmemBytes = 128; constexpr size_t kWarpDynamicSmemBytes = 0; +constexpr size_t kMr515DynamicSmemBytes = 128; enum class ParameterKind : int { kNone = 0, @@ -41,6 +43,44 @@ enum class ParameterKind : int { kPerRow = 2, }; +// Stable numeric values are intentionally exposed by softmax_route for tests +// and benchmark evidence. Keep this selection in C++ so an API-level +// correctness test cannot silently pass through OnlineSoftmax fallback. +enum class SoftmaxRoute : int64_t { + kFallback = 0, + kWarp = 1, + kRowwise = 2, + kBootstrap = 3, + kMr515V32000T512 = 4, +}; + +bool use_mr515_none_row(uint32_t rows) { + switch (rows) { + case 16: + case 32: + case 64: + case 128: + case 512: + case 1024: + return true; + default: + return false; + } +} + +bool use_mr515_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind, + float temperature_val, bool enable_pdl, int device_major, + int device_minor) { + if (device_major != 10 || device_minor != 3 || vocab_size != 32000) { + return false; + } + if (parameter_kind == ParameterKind::kNone) { + return !enable_pdl && use_mr515_none_row(rows); + } + return parameter_kind == ParameterKind::kScalar && enable_pdl && rows == 64 && + temperature_val == 1.0f; +} + bool use_warp_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind) { return rows <= 128 && vocab_size <= 257 && (parameter_kind == ParameterKind::kScalar || parameter_kind == ParameterKind::kPerRow); @@ -62,36 +102,115 @@ bool use_rowwise_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parame dense_aligned_high_row_wide || measured_large_odd; } -cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temperature_arr, - float temperature_val, ParameterKind parameter_kind, - uint32_t rows, uint32_t vocab_size, void* workspace, - size_t workspace_bytes, cudaStream_t stream) { +SoftmaxRoute select_softmax_route(uint32_t rows, uint32_t vocab_size, + ParameterKind parameter_kind, float temperature_val, + bool enable_pdl, int device_major, int device_minor) { if (rows == 0 || vocab_size == 0 || static_cast(rows) * vocab_size > static_cast(std::numeric_limits::max())) { + return SoftmaxRoute::kFallback; + } + if (use_mr515_kernel(rows, vocab_size, parameter_kind, temperature_val, enable_pdl, + device_major, device_minor)) { + return SoftmaxRoute::kMr515V32000T512; + } + if (use_warp_kernel(rows, vocab_size, parameter_kind)) { + return SoftmaxRoute::kWarp; + } + if (use_rowwise_kernel(rows, vocab_size, parameter_kind)) { + return SoftmaxRoute::kRowwise; + } + return SoftmaxRoute::kBootstrap; +} + +cudaError_t launch_noncooperative_kernel(const void* kernel, dim3 grid, dim3 block, void** args, + size_t dynamic_smem_bytes, bool enable_pdl, + cudaStream_t stream) { + if (!enable_pdl) { + return cudaLaunchKernel(kernel, grid, block, args, dynamic_smem_bytes, stream); + } + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = dynamic_smem_bytes; + config.stream = stream; + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = enable_pdl ? 1 : 0; + config.attrs = &attribute; + config.numAttrs = 1; + return cudaLaunchKernelExC(&config, kernel, args); +} + +cudaError_t launch_cooperative_kernel(const void* kernel, dim3 grid, dim3 block, void** args, + size_t dynamic_smem_bytes, bool enable_pdl, + cudaStream_t stream) { + if (!enable_pdl) { + return cudaLaunchCooperativeKernel(kernel, grid, block, args, dynamic_smem_bytes, stream); + } + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = dynamic_smem_bytes; + config.stream = stream; + cudaLaunchAttribute attributes[2]{}; + attributes[0].id = cudaLaunchAttributeCooperative; + attributes[0].val.cooperative = 1; + attributes[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[1].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attributes; + config.numAttrs = 2; + return cudaLaunchKernelExC(&config, kernel, args); +} + +cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temperature_arr, + float temperature_val, ParameterKind parameter_kind, + uint32_t rows, uint32_t vocab_size, void* workspace, + size_t workspace_bytes, bool enable_pdl, bool is_sm103, + SoftmaxRoute* selected_route, cudaStream_t stream) { + const SoftmaxRoute route = + select_softmax_route(rows, vocab_size, parameter_kind, temperature_val, enable_pdl, + /*device_major=*/10, is_sm103 ? 3 : 0); + *selected_route = route; + if (route == SoftmaxRoute::kFallback) { return cudaErrorNotSupported; } float* parameter = temperature_arr != nullptr ? temperature_arr : logits; int rows_i = static_cast(rows); int vocab_size_i = static_cast(vocab_size); - int parameter_kind_i = static_cast(parameter_kind); + // Keep the existing sm_100a launch path unchanged. The high bit is an + // sm_103a integration-only PDL flag; adapted kernels strip it before + // evaluating the original 0/1/2 parameter-kind ABI. + const bool launch_with_pdl = is_sm103 && enable_pdl; + int parameter_kind_i = static_cast(parameter_kind) | (launch_with_pdl ? 4 : 0); - if (use_warp_kernel(rows, vocab_size, parameter_kind)) { + if (route == SoftmaxRoute::kMr515V32000T512) { + // TEMP_KIND=0 never dereferences this ABI slot. A null dummy tells the + // adapted frozen payload to execute the required PDL wait/signal pair. + float* mr515_temperature = launch_with_pdl ? nullptr : parameter; + void* args[] = {&logits, &mr515_temperature, &output, &temperature_val}; + return launch_noncooperative_kernel( + reinterpret_cast(kernel_mr474_manual_softmax_exp2_t512_vec4), dim3(rows), + dim3(kMr515Threads), args, kMr515DynamicSmemBytes, launch_with_pdl, stream); + } + + if (route == SoftmaxRoute::kWarp) { void* args[] = {&logits, ¶meter, &output, &rows_i, &vocab_size_i, ¶meter_kind_i, &temperature_val}; - return cudaLaunchKernel( + return launch_noncooperative_kernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_followup_warp), dim3(ceil_div(rows, static_cast(kWarpRowsPerCta))), dim3(kWarpThreads), args, - kWarpDynamicSmemBytes, stream); + kWarpDynamicSmemBytes, launch_with_pdl, stream); } - if (use_rowwise_kernel(rows, vocab_size, parameter_kind)) { + if (route == SoftmaxRoute::kRowwise) { void* args[] = {&logits, ¶meter, &output, &rows_i, &vocab_size_i, ¶meter_kind_i, &temperature_val}; - return cudaLaunchKernel( + return launch_noncooperative_kernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_followup_rowwise), - dim3(rows), dim3(kRowwiseThreads), args, kRowwiseDynamicSmemBytes, stream); + dim3(rows), dim3(kRowwiseThreads), args, kRowwiseDynamicSmemBytes, + launch_with_pdl, stream); } int active_blocks_per_sm = 0; @@ -109,7 +228,7 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper int device = 0; if ((status = cudaGetDevice(&device)) != cudaSuccess || (status = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device)) != - cudaSuccess) { + cudaSuccess) { return status; } @@ -136,28 +255,85 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper int splits_i = splits; void* args[] = {&logits, ¶meter, &output, &partial_max, &partial_sum, &rows_i, &vocab_size_i, &splits_i, ¶meter_kind_i, &temperature_val}; - return cudaLaunchCooperativeKernel( + return launch_cooperative_kernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_bootstrap_seed), dim3(static_cast(grid)), dim3(kBootstrapThreads), args, kBootstrapDynamicSmemBytes, - stream); + launch_with_pdl, stream); } -} // namespace +ParameterKind validate_logits_and_temperature(TensorView logits, + Optional maybe_temperature_arr, + bool temperature_is_none) { + CHECK_INPUT(logits); + CHECK_DIM(2, logits); + CHECK_INPUT_TYPE(logits, dl_float32); + TVM_FFI_ICHECK_GT(logits.size(0), 0) << "logits must contain at least one row"; + TVM_FFI_ICHECK_GT(logits.size(1), 0) << "logits must contain at least one vocabulary entry"; + TVM_FFI_ICHECK_LE(static_cast(logits.size(0)), + static_cast(std::numeric_limits::max())); + TVM_FFI_ICHECK_LE(static_cast(logits.size(1)), + static_cast(std::numeric_limits::max())); -void blackwell_softmax(TensorView workspace_buffer, TensorView logits, TensorView output, - Optional maybe_temperature_arr, double temperature_val, - bool enable_pdl, bool temperature_is_none) { + if (maybe_temperature_arr.has_value()) { + const TensorView temperature_arr = maybe_temperature_arr.value(); + CHECK_INPUT(temperature_arr); + CHECK_DIM(1, temperature_arr); + CHECK_INPUT_TYPE(temperature_arr, dl_float32); + CHECK_DEVICE(temperature_arr, logits); + TVM_FFI_ICHECK_EQ(temperature_arr.size(0), logits.size(0)) + << "temperature length must equal logits.size(0)"; + TVM_FFI_ICHECK(!temperature_is_none) + << "temperature_is_none must be false when a temperature tensor is provided"; + return ParameterKind::kPerRow; + } + return temperature_is_none ? ParameterKind::kNone : ParameterKind::kScalar; +} + +void validate_softmax_io(TensorView workspace_buffer, TensorView logits, TensorView output) { CHECK_INPUT(workspace_buffer); - CHECK_INPUT(logits); + CHECK_DIM(1, workspace_buffer); + CHECK_DEVICE(workspace_buffer, logits); CHECK_INPUT(output); - CHECK_DIM(2, logits); + CHECK_DIM(2, output); + CHECK_INPUT_TYPE(output, dl_float32); + CHECK_DEVICE(output, logits); + CHECK_SHAPE(output, logits); + TVM_FFI_ICHECK_NE(output.data_ptr(), logits.data_ptr()) + << "output must be fresh and must not alias logits"; +} + +SoftmaxRoute query_softmax_route(TensorView logits, Optional maybe_temperature_arr, + double temperature_val, bool enable_pdl, + bool temperature_is_none) { + const ParameterKind parameter_kind = + validate_logits_and_temperature(logits, maybe_temperature_arr, temperature_is_none); + ffi::CUDADeviceGuard device_guard(logits.device().device_id); + int device_major = 0; + int device_minor = 0; + cudaError_t status = cudaDeviceGetAttribute(&device_major, cudaDevAttrComputeCapabilityMajor, + logits.device().device_id); + if (status == cudaSuccess) { + status = cudaDeviceGetAttribute(&device_minor, cudaDevAttrComputeCapabilityMinor, + logits.device().device_id); + } + TVM_FFI_ICHECK(status == cudaSuccess) + << "Blackwell Softmax route query failed with error code " << cudaGetErrorString(status); + return select_softmax_route(static_cast(logits.size(0)), + static_cast(logits.size(1)), parameter_kind, + static_cast(temperature_val), enable_pdl, device_major, + device_minor); +} + +void blackwell_softmax_impl(TensorView workspace_buffer, TensorView logits, TensorView output, + Optional maybe_temperature_arr, double temperature_val, + bool enable_pdl, bool temperature_is_none, bool is_sm103) { + const ParameterKind parameter_kind = + validate_logits_and_temperature(logits, maybe_temperature_arr, temperature_is_none); + validate_softmax_io(workspace_buffer, logits, output); const auto rows = static_cast(logits.size(0)); const auto vocab_size = static_cast(logits.size(1)); const bool has_temperature_arr = maybe_temperature_arr.has_value(); - const ParameterKind parameter_kind = - has_temperature_arr ? ParameterKind::kPerRow - : (temperature_is_none ? ParameterKind::kNone : ParameterKind::kScalar); ffi::CUDADeviceGuard device_guard(logits.device().device_id); auto stream = get_stream(logits.device()); @@ -167,17 +343,39 @@ void blackwell_softmax(TensorView workspace_buffer, TensorView logits, TensorVie has_temperature_arr ? static_cast(maybe_temperature_arr.value().data_ptr()) : nullptr; const size_t workspace_bytes = get_element_size(workspace_buffer) * workspace_buffer.size(0); + SoftmaxRoute selected_route = SoftmaxRoute::kFallback; cudaError_t status = launch_blackwell_softmax( logits_ptr, output_ptr, temperature_ptr, static_cast(temperature_val), parameter_kind, - rows, vocab_size, workspace_buffer.data_ptr(), workspace_bytes, stream); - if (status == cudaErrorNotSupported) { + rows, vocab_size, workspace_buffer.data_ptr(), workspace_bytes, enable_pdl, is_sm103, + &selected_route, stream); + // The promoted MR515 route is fail-closed: it must never be hidden by an + // OnlineSoftmax fallback if its actual launch fails. + if (status == cudaErrorNotSupported && selected_route != SoftmaxRoute::kMr515V32000T512) { status = sampling::OnlineSoftmax(logits_ptr, output_ptr, rows, vocab_size, temperature_ptr, static_cast(temperature_val), workspace_buffer.data_ptr(), workspace_bytes, enable_pdl, stream); } TVM_FFI_ICHECK(status == cudaSuccess) - << "Blackwell Softmax failed with error code " << cudaGetErrorString(status); + << "Blackwell Softmax route " << static_cast(selected_route) + << " failed with error code " << cudaGetErrorString(status); +} + +} // namespace + +void blackwell_softmax(TensorView workspace_buffer, TensorView logits, TensorView output, + Optional maybe_temperature_arr, double temperature_val, + bool enable_pdl, bool temperature_is_none, bool is_sm103) { + blackwell_softmax_impl(workspace_buffer, logits, output, maybe_temperature_arr, temperature_val, + enable_pdl, temperature_is_none, is_sm103); +} + +int64_t blackwell_softmax_route(TensorView logits, Optional maybe_temperature_arr, + double temperature_val, bool enable_pdl, + bool temperature_is_none) { + return static_cast(query_softmax_route(logits, maybe_temperature_arr, temperature_val, + enable_pdl, temperature_is_none)); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(softmax, blackwell_softmax); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(softmax_route, blackwell_softmax_route); diff --git a/csrc/blackwell_softmax_bootstrap.cu b/csrc/blackwell_softmax_bootstrap.cu index 5db764c5d63..06880bb583c 100644 --- a/csrc/blackwell_softmax_bootstrap.cu +++ b/csrc/blackwell_softmax_bootstrap.cu @@ -16,7 +16,9 @@ // Generated by Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. // Original generated payload before FlashInfer type adaptation: 24102 bytes, // sha256:46049370dbd905ff7234d414745d197f883157c8edc01d83230562b4dff5f862. -// The sm_100a and sm_103a payloads are byte-identical. +// The raw sm_100a and sm_103a payloads were byte-identical. This integration +// uses standard integer types, adds sm_103a-only PDL wait/trigger control, and +// legalizes each 256-bit inline-PTX access to 2x128-bit on CUDA <12.9. // clang-format off #include @@ -82,6 +84,13 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); const int lane = tid % 32; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + const bool enable_pdl = (parameter_kind & 4) != 0; + parameter_kind &= 3; + if (enable_pdl) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif extern __shared__ __align__(1024) char smem_raw[]; int smem; @@ -122,9 +131,18 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* unsigned _ldv8_0_5; unsigned _ldv8_0_6; unsigned _ldv8_0_7; +#if __CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) asm volatile( "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" : "=r"(_ldv8_0_0), "=r"(_ldv8_0_1), "=r"(_ldv8_0_2), "=r"(_ldv8_0_3), "=r"(_ldv8_0_4), "=r"(_ldv8_0_5), "=r"(_ldv8_0_6), "=r"(_ldv8_0_7) : "l"((const void*)(x + (index))) : "memory"); +#else + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_0_0), "=r"(_ldv8_0_1), "=r"(_ldv8_0_2), "=r"(_ldv8_0_3) : "l"((const void*)(x + (index))) : "memory"); + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_0_4), "=r"(_ldv8_0_5), "=r"(_ldv8_0_6), "=r"(_ldv8_0_7) : "l"((const void*)(x + (index) + 4)) : "memory"); +#endif _vec_load_0[0 + 0] = __uint_as_float(_ldv8_0_0); _vec_load_0[0 + 1] = __uint_as_float(_ldv8_0_1); _vec_load_0[0 + 2] = __uint_as_float(_ldv8_0_2); @@ -243,9 +261,18 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* unsigned _ldv8_1_5; unsigned _ldv8_1_6; unsigned _ldv8_1_7; +#if __CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) asm volatile( "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) : "l"((const void*)(x + (index_2))) : "memory"); +#else + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3) : "l"((const void*)(x + (index_2))) : "memory"); + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) : "l"((const void*)(x + (index_2) + 4)) : "memory"); +#endif _vec_load_1[0 + 0] = __uint_as_float(_ldv8_1_0); _vec_load_1[0 + 1] = __uint_as_float(_ldv8_1_1); _vec_load_1[0 + 2] = __uint_as_float(_ldv8_1_2); @@ -310,9 +337,18 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* unsigned _ldv8_2_5; unsigned _ldv8_2_6; unsigned _ldv8_2_7; +#if __CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) asm volatile( "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3), "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) : "l"((const void*)(x + (index_4))) : "memory"); +#else + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_2_0), "=r"(_ldv8_2_1), "=r"(_ldv8_2_2), "=r"(_ldv8_2_3) : "l"((const void*)(x + (index_4))) : "memory"); + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_2_4), "=r"(_ldv8_2_5), "=r"(_ldv8_2_6), "=r"(_ldv8_2_7) : "l"((const void*)(x + (index_4) + 4)) : "memory"); +#endif _vec_load_2[0 + 0] = __uint_as_float(_ldv8_2_0); _vec_load_2[0 + 1] = __uint_as_float(_ldv8_2_1); _vec_load_2[0 + 2] = __uint_as_float(_ldv8_2_2); @@ -336,9 +372,18 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* unsigned _stv8_3_5 = __float_as_uint(_vec_load_2[0 + 5]); unsigned _stv8_3_6 = __float_as_uint(_vec_load_2[0 + 6]); unsigned _stv8_3_7 = __float_as_uint(_vec_load_2[0 + 7]); +#if __CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) asm volatile( "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" :: "l"((void*)(output + (index_4))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); +#else + asm volatile( + "st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"((void*)(output + (index_4))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3) : "memory"); + asm volatile( + "st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"((void*)(output + (index_4) + 4)), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); +#endif } } int tail_col_2 = vector_groups * 8 + tid; @@ -458,9 +503,18 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* unsigned _ldv8_4_5; unsigned _ldv8_4_6; unsigned _ldv8_4_7; +#if __CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) asm volatile( "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" : "=r"(_ldv8_4_0), "=r"(_ldv8_4_1), "=r"(_ldv8_4_2), "=r"(_ldv8_4_3), "=r"(_ldv8_4_4), "=r"(_ldv8_4_5), "=r"(_ldv8_4_6), "=r"(_ldv8_4_7) : "l"((const void*)(x + (index_6))) : "memory"); +#else + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_4_0), "=r"(_ldv8_4_1), "=r"(_ldv8_4_2), "=r"(_ldv8_4_3) : "l"((const void*)(x + (index_6))) : "memory"); + asm volatile( + "ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ldv8_4_4), "=r"(_ldv8_4_5), "=r"(_ldv8_4_6), "=r"(_ldv8_4_7) : "l"((const void*)(x + (index_6) + 4)) : "memory"); +#endif _vec_load_3[0 + 0] = __uint_as_float(_ldv8_4_0); _vec_load_3[0 + 1] = __uint_as_float(_ldv8_4_1); _vec_load_3[0 + 2] = __uint_as_float(_ldv8_4_2); @@ -484,9 +538,18 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* unsigned _stv8_5_5 = __float_as_uint(_vec_load_3[0 + 5]); unsigned _stv8_5_6 = __float_as_uint(_vec_load_3[0 + 6]); unsigned _stv8_5_7 = __float_as_uint(_vec_load_3[0 + 7]); +#if __CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) asm volatile( "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" :: "l"((void*)(output + (index_6))), "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3), "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) : "memory"); +#else + asm volatile( + "st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"((void*)(output + (index_6))), "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3) : "memory"); + asm volatile( + "st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"((void*)(output + (index_6) + 4)), "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) : "memory"); +#endif } } int tail_col_3 = vector_groups_1 * 8 + tid; @@ -505,9 +568,14 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* } } } +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + if (enable_pdl) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif } } // extern "C" // clang-format on -// End exact generated payload. +// End integrated generated payload. diff --git a/csrc/blackwell_softmax_mr515_exp2.cu b/csrc/blackwell_softmax_mr515_exp2.cu new file mode 100644 index 00000000000..1c89a154a46 --- /dev/null +++ b/csrc/blackwell_softmax_mr515_exp2.cu @@ -0,0 +1,233 @@ +/* + * 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. + */ +// Frozen from Cake commit f120b38798f4481a116127e7cdecbe4ed9d2cf5f. +// Measured kernel commit: ca826e6b97b2cd696e208c842eb4008af2328d65. +// Weave source sha256:ce4bf0aba8a398979b70df863658baa9b3700fa840b4be957997fd0bee64cf0b. +// Original generated payload before FlashInfer type adaptation: 7194 bytes, +// sha256:8ca2513545faf9cb960c6393fb48f4c9f4598a781ad9c929e8d88ebda738a89b. +// This integration uses standard integer types and adds sm_103a-only PDL +// wait/trigger control. The raw payload hash above remains the export identity. +// Specialization: sm_103a, vocab=32000, t512, vec4, TEMP_KIND=0, +// MATERIALIZE_EXP=0, grid=(rows, 1, 1), dynamic shared memory=128 bytes. + +// clang-format off +#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 LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_REDUCE_SMEM_OFF 0 +#define SMEM_REDUCE_SMEM_STAGE_BYTES 128 +#define SMEM_REDUCE_SMEM_STRIDE 128 +#define SMEM_TOTAL 128 +#define THREADS 512 +#define VOCAB_SIZE 32000 +#define VECTORS_PER_THREAD 16 +#define TEMP_KIND 0 +#define MATERIALIZE_EXP 0 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__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 max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + +extern "C" { + +__global__ __launch_bounds__(512) void +kernel_mr474_manual_softmax_exp2_t512_vec4(float* __restrict__ x, float* __restrict__ temperature, float* __restrict__ y, float temperature_scalar) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + const bool enable_pdl = temperature == nullptr; + if (enable_pdl) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + float* reduce_smem = reinterpret_cast(smem_raw + 0); + const int reduce_smem_addr = smem + 0; + + // === Task calls (dependency order) === + float inv_temperature = 1.0f; + { + } + float local_max = -LOOM_INF; + #pragma unroll + for (int tile = 0; tile < VECTORS_PER_THREAD; tile++) { + if ((tile * 512 + tid) * 4 < VOCAB_SIZE) { + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(x + ((unsigned long long)bid * (unsigned long long)VOCAB_SIZE + (unsigned long long)((tile * 512 + tid) * 4)) + 0); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + #pragma unroll + for (int j = 0; j < 4; j++) { + float value = _vec_load_0[j]; + float _max_0 = max_noftz(local_max, value); + local_max = _max_0; + } + } + } + float _warp_reduce_0 = local_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 = max_noftz(_warp_reduce_0, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset)); + local_max = _warp_reduce_0; + if (lane == 0) { + reduce_smem[warp] = local_max; + } + __syncthreads(); + float block_max = ((lane < 16) ? reduce_smem[lane] : -LOOM_INF); + float _warp_reduce_1 = block_max; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 = max_noftz(_warp_reduce_1, __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset)); + block_max = _warp_reduce_1; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = block_max; + } + } + __syncthreads(); + float row_max = reduce_smem[0]; + __syncthreads(); + float local_sum = 0.0f; + #pragma unroll + for (int tile_1 = 0; tile_1 < VECTORS_PER_THREAD; tile_1++) { + if ((tile_1 * 512 + tid) * 4 < VOCAB_SIZE) { + float _vec_load_1[4]; + { + float4 _v4 = *reinterpret_cast(x + ((unsigned long long)bid * (unsigned long long)VOCAB_SIZE + (unsigned long long)((tile_1 * 512 + tid) * 4)) + 0); + _vec_load_1[0 + 0] = _v4.x; + _vec_load_1[0 + 1] = _v4.y; + _vec_load_1[0 + 2] = _v4.z; + _vec_load_1[0 + 3] = _v4.w; + } + float exp_values[4]; + #pragma unroll + for (int j_1 = 0; j_1 < 4; j_1++) { + float value_1 = _vec_load_1[j_1]; + float _exp2_0 = approx_exp2((value_1 - row_max) * 1.4426950408889634f); + float exp_value = ((row_max == -LOOM_INF) ? 0.0f : _exp2_0); + exp_values[j_1] = exp_value; + local_sum += exp_value; + } + } + } + float _warp_reduce_2 = local_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_2 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_2, offset); + local_sum = _warp_reduce_2; + if (lane == 0) { + reduce_smem[warp] = local_sum; + } + __syncthreads(); + float block_sum = ((lane < 16) ? reduce_smem[lane] : 0.0f); + float _warp_reduce_3 = block_sum; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_3 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_3, offset); + block_sum = _warp_reduce_3; + __syncthreads(); + if (warp == 0) { + if (elect_sync()) { + reduce_smem[0] = block_sum; + } + } + __syncthreads(); + float inv_sum = ((reduce_smem[0] == 0.0f) ? 0.0f : 1.0f / reduce_smem[0]); + #pragma unroll + for (int tile_2 = 0; tile_2 < VECTORS_PER_THREAD; tile_2++) { + if ((tile_2 * 512 + tid) * 4 < VOCAB_SIZE) { + float out_values[4]; + { + float _vec_load_3[4]; + { + float4 _v4 = *reinterpret_cast(x + ((unsigned long long)bid * (unsigned long long)VOCAB_SIZE + (unsigned long long)((tile_2 * 512 + tid) * 4)) + 0); + _vec_load_3[0 + 0] = _v4.x; + _vec_load_3[0 + 1] = _v4.y; + _vec_load_3[0 + 2] = _v4.z; + _vec_load_3[0 + 3] = _v4.w; + } + #pragma unroll + for (int j_2 = 0; j_2 < 4; j_2++) { + float value_2 = _vec_load_3[j_2]; + float _exp2_1 = approx_exp2((value_2 - row_max) * 1.4426950408889634f); + out_values[j_2] = ((row_max == -LOOM_INF) ? 0.0f : _exp2_1 * inv_sum); + } + } + { + float4 _v4 = make_float4(out_values[0 + 0], out_values[0 + 1], out_values[0 + 2], out_values[0 + 3]); + *reinterpret_cast(y + ((unsigned long long)bid * (unsigned long long)VOCAB_SIZE + (unsigned long long)((tile_2 * 512 + tid) * 4)) + 0) = _v4; + } + } + } +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + if (enable_pdl) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif +} + +} // extern "C" + +// clang-format on diff --git a/csrc/blackwell_softmax_rowwise.cu b/csrc/blackwell_softmax_rowwise.cu index 04a1475b574..7a347e3a7f6 100644 --- a/csrc/blackwell_softmax_rowwise.cu +++ b/csrc/blackwell_softmax_rowwise.cu @@ -16,7 +16,8 @@ // Generated by Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. // Original generated payload before FlashInfer type adaptation: 8720 bytes, // sha256:e47f82923112e143173f32dbb145986f8c9a1503fd2bb71f22f5c065769162e3. -// The sm_100a and sm_103a payloads are byte-identical. +// The raw sm_100a and sm_103a payloads were byte-identical. This integration +// uses standard integer types and adds sm_103a-only PDL wait/trigger control. // clang-format off #include @@ -81,6 +82,13 @@ kernel_flashinfer_blackwell_softmax_followup_rowwise(float* __restrict__ x, floa const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); const int lane = tid % 32; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + const bool enable_pdl = (parameter_kind & 4) != 0; + parameter_kind &= 3; + if (enable_pdl) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif extern __shared__ __align__(1024) char smem_raw[]; int smem; @@ -249,9 +257,14 @@ kernel_flashinfer_blackwell_softmax_followup_rowwise(float* __restrict__ x, floa float _exp2_9 = approx_exp2((tail_scaled_1 - cta_max) * 1.4426950408889634f); output[tail_index_1] = _exp2_9 * inv_sum; } +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + if (enable_pdl) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif } } // extern "C" // clang-format on -// End exact generated payload. +// End integrated generated payload. diff --git a/csrc/blackwell_softmax_warp.cu b/csrc/blackwell_softmax_warp.cu index 34d587b2520..b8b4bb3fd17 100644 --- a/csrc/blackwell_softmax_warp.cu +++ b/csrc/blackwell_softmax_warp.cu @@ -16,7 +16,8 @@ // Generated by Loom from Cake commit dcece84ec6a568402d0e37fac4f15f61f2cb9741. // Original generated payload before FlashInfer type adaptation: 3923 bytes, // sha256:8bc9df2b57d4e6021d0add03110a11f0cbb3bab1218e94d00d92f4bbf597ee9e. -// The sm_100a and sm_103a payloads are byte-identical. +// The raw sm_100a and sm_103a payloads were byte-identical. This integration +// uses standard integer types and adds sm_103a-only PDL wait/trigger control. // clang-format off #include @@ -63,6 +64,13 @@ kernel_flashinfer_blackwell_softmax_followup_warp(float* __restrict__ x, float* const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); const int lane = tid % 32; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + const bool enable_pdl = (parameter_kind & 4) != 0; + parameter_kind &= 3; + if (enable_pdl) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif const int bid = blockIdx.x; @@ -119,9 +127,14 @@ kernel_flashinfer_blackwell_softmax_followup_warp(float* __restrict__ x, float* } } } +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + if (enable_pdl) { + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + } +#endif } } // extern "C" // clang-format on -// End exact generated payload. +// End integrated generated payload. diff --git a/flashinfer/jit/blackwell_softmax.py b/flashinfer/jit/blackwell_softmax.py index ada1930e82c..96ed8407506 100644 --- a/flashinfer/jit/blackwell_softmax.py +++ b/flashinfer/jit/blackwell_softmax.py @@ -29,6 +29,7 @@ def gen_blackwell_softmax_module() -> JitSpec: jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_bootstrap.cu", jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_rowwise.cu", jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_warp.cu", + jit_env.FLASHINFER_CSRC_DIR / "blackwell_softmax_mr515_exp2.cu", ], extra_cuda_cflags=nvcc_flags, ) diff --git a/flashinfer/sampling.py b/flashinfer/sampling.py index cd722142f6f..13a7479341b 100644 --- a/flashinfer/sampling.py +++ b/flashinfer/sampling.py @@ -15,6 +15,7 @@ """ import functools +from numbers import Real from types import SimpleNamespace from typing import Optional, Tuple, Union import torch @@ -64,9 +65,15 @@ def get_seed_and_offset( return int(seed), int(offset) +@functools.cache +def get_blackwell_softmax_module(): + """Build and retain the raw module used by the API and evidence hooks.""" + return gen_blackwell_softmax_module().build_and_load() + + @functools.cache def get_blackwell_softmax_op(): - module = gen_blackwell_softmax_module().build_and_load() + module = get_blackwell_softmax_module() @register_custom_op( "flashinfer::blackwell_softmax", mutates_args=("workspace_buffer",) @@ -78,6 +85,7 @@ def blackwell_softmax( temperature_val: float, enable_pdl: bool, temperature_is_none: bool, + is_sm103: bool, ) -> torch.Tensor: logits = logits.float() probs = torch.empty_like(logits, device=logits.device) @@ -92,6 +100,7 @@ def blackwell_softmax( temperature_val, enable_pdl, temperature_is_none, + is_sm103, ) return probs @@ -103,18 +112,58 @@ def _fake_blackwell_softmax( temperature_val: float, enable_pdl: bool, temperature_is_none: bool, + is_sm103: bool, ) -> torch.Tensor: return torch.empty_like(logits, device=logits.device, dtype=torch.float32) return blackwell_softmax +_BLACKWELL_SOFTMAX_ROUTE_MR515_V32000_T512 = 4 + + +def _blackwell_softmax_route_for_testing( + logits: torch.Tensor, + temperature: Optional[Union[torch.Tensor, float]] = None, + enable_pdl: bool = False, +) -> int: + """Return the C++ dispatcher route ID without launching a kernel.""" + if not logits.is_cuda: + raise ValueError(f"logits must be a CUDA tensor, got device={logits.device}") + maybe_temperature_arr, temperature_val, temperature_is_none = ( + _validate_softmax_temperature(logits, temperature) + ) + logits_fp32 = logits.float() + temperature_fp32 = ( + maybe_temperature_arr.float() if maybe_temperature_arr is not None else None + ) + return int( + get_blackwell_softmax_module().softmax_route( + logits_fp32, + temperature_fp32, + temperature_val, + enable_pdl, + temperature_is_none, + ) + ) + + +@functools.cache +def _blackwell_softmax_capability(device_index: int) -> Tuple[int, int]: + return torch.cuda.get_device_capability(device_index) + + @functools.cache def _supports_blackwell_softmax(device_index: int) -> bool: - major, minor = torch.cuda.get_device_capability(device_index) + major, minor = _blackwell_softmax_capability(device_index) return major == 10 and minor in (0, 3) +@functools.cache +def _is_sm103(device_index: int) -> bool: + return _blackwell_softmax_capability(device_index) == (10, 3) + + @functools.cache def get_sampling_module(): module = gen_sampling_module().build_and_load() @@ -724,6 +773,54 @@ def _to_tensor_scalar_tuple(x): return (None, x) +_SOFTMAX_TEMPERATURE_DTYPES = ( + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, +) + + +def _validate_softmax_temperature( + logits: torch.Tensor, + temperature: Optional[Union[torch.Tensor, float]], +) -> Tuple[Optional[torch.Tensor], float, bool]: + """Validate and normalize the public Softmax temperature argument.""" + if temperature is None: + return None, 1.0, True + + if isinstance(temperature, torch.Tensor): + if temperature.ndim != 1: + raise ValueError( + "temperature tensor must be 1D with one value per logits row, " + f"got shape {tuple(temperature.shape)}" + ) + expected_rows = logits.size(0) + if temperature.size(0) != expected_rows: + raise ValueError( + "temperature tensor length must equal logits.size(0) " + f"({expected_rows}), got {temperature.size(0)}" + ) + if temperature.device != logits.device: + raise ValueError( + "temperature tensor must be on the same device as logits " + f"({logits.device}), got {temperature.device}" + ) + if temperature.dtype not in _SOFTMAX_TEMPERATURE_DTYPES: + raise ValueError( + "temperature tensor must have dtype float16, bfloat16, float32, " + f"or float64, got {temperature.dtype}" + ) + return temperature, 0.0, False + + if isinstance(temperature, bool) or not isinstance(temperature, Real): + raise TypeError( + "temperature must be None, a real scalar, or a 1D floating-point " + "tensor with one value per logits row" + ) + return None, float(temperature), False + + def _validate_and_convert_seed_offset( seed: Union[int, torch.Tensor], offset: Union[int, torch.Tensor], @@ -797,11 +894,13 @@ def softmax( Parameters ---------- logits : torch.Tensor - Input tensor of logits. + Two-dimensional CUDA tensor of logits. temperature: Optional[Union[torch.Tensor, float]] Either a scalar or a tensor of shape ``(batch_size,)``, representing the temperature for temperature scaling. If a scalar, the same temperature is used for all requests. - If a tensor, each request has its own temperature. + If a tensor, each request has its own temperature. The tensor must be on + the same CUDA device as ``logits`` and have dtype float16, bfloat16, + float32, or float64. enable_pdl : Optional[bool] Whether to enable Programmatic Dependent Launch (PDL) for improved performance on supported hardware. If None (default), PDL will be automatically enabled on devices with compute capability >= 9.0. @@ -830,16 +929,18 @@ def softmax( [0.2401, 0.1707, 0.2249, 0.1664, 0.1979], [0.1724, 0.2719, 0.1991, 0.1465, 0.2101]], device='cuda:0') """ + if not logits.is_cuda: + raise ValueError(f"logits must be a CUDA tensor, got device={logits.device}") + + maybe_temperature_arr, temperature_val, temperature_is_none = ( + _validate_softmax_temperature(logits, temperature) + ) workspace_buffer = _get_cache_buf("softmax_workspace", 1024 * 1024, logits.device) - temperature_is_none = temperature is None - if temperature is None: - temperature = 1.0 # Auto-detect PDL support if not specified if enable_pdl is None: enable_pdl = device_support_pdl(logits.device) - temperature_args = _to_tensor_scalar_tuple(temperature) device_index = logits.device.index if device_index is None: device_index = torch.cuda.current_device() @@ -847,15 +948,18 @@ def softmax( return get_blackwell_softmax_op()( workspace_buffer, logits, - *temperature_args, + maybe_temperature_arr, + temperature_val, enable_pdl, temperature_is_none, + _is_sm103(device_index), ) return get_sampling_module().softmax( workspace_buffer, logits, - *temperature_args, + maybe_temperature_arr, + temperature_val, enable_pdl, ) diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index cf19b372954..e996fca838c 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -17,12 +17,18 @@ #include -// Frozen from Cake measured checkpoint 924403cd62ea6dc160ffcea63c944d1e828d393a -// and retained byte-for-byte by MR474 cleanup d77b50ee1e56fb3701a512f1dd2765e6067b1be2. -// Weave sm_100a and sm_103a output is byte-identical for all three kernels: +// The three original generated payloads came from Cake checkpoint +// 924403cd62ea6dc160ffcea63c944d1e828d393a and MR474 cleanup +// d77b50ee1e56fb3701a512f1dd2765e6067b1be2. Their pre-integration Weave +// sm_100a and sm_103a output was byte-identical: // bootstrap: sha256:46049370dbd905ff7234d414745d197f883157c8edc01d83230562b4dff5f862 // rowwise: sha256:e47f82923112e143173f32dbb145986f8c9a1503fd2bb71f22f5c065769162e3 // warp: sha256:8bc9df2b57d4e6021d0add03110a11f0cbb3bab1218e94d00d92f4bbf597ee9e +// Integration adds standard integer types, sm_103a-only PDL device control, +// and a CUDA <12.9 2x128-bit fallback for bootstrap's 256-bit inline PTX. +// The additive sm_103a 32K specialization comes from Cake f120b38798: +// mr515 exp2 t512/vec4 raw generated payload: +// sha256:8ca2513545faf9cb960c6393fb48f4c9f4598a781ad9c929e8d88ebda738a89b extern "C" { @@ -40,4 +46,8 @@ __global__ void kernel_flashinfer_blackwell_softmax_followup_warp(float* logits, int parameter_kind, float scalar_temperature); +__global__ void kernel_mr474_manual_softmax_exp2_t512_vec4(float* logits, float* temperature, + float* output, + float scalar_temperature); + } // extern "C" diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index 755a5d62a8a..aec41fc38ca 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -38,6 +38,144 @@ def gumbel_noise(shape, device): return gumbel_noise +def test_softmax_rejects_cpu_before_cuda_state(monkeypatch): + def fail_if_called(*args, **kwargs): + pytest.fail("CUDA workspace or capability state was accessed for CPU logits") + + monkeypatch.setattr(flashinfer.sampling, "_get_cache_buf", fail_if_called) + monkeypatch.setattr(flashinfer.sampling, "device_support_pdl", fail_if_called) + monkeypatch.setattr( + flashinfer.sampling, "_supports_blackwell_softmax", fail_if_called + ) + + with pytest.raises(ValueError, match="logits must be a CUDA tensor"): + flashinfer.sampling.softmax(torch.randn(2, 3)) + + +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64] +) +def test_validate_softmax_temperature_accepts_floating_row_tensor(dtype): + logits = torch.empty((3, 5), dtype=torch.float32) + temperature = torch.ones(3, dtype=dtype) + + temperature_arr, temperature_val, temperature_is_none = ( + flashinfer.sampling._validate_softmax_temperature(logits, temperature) + ) + + assert temperature_arr is temperature + assert temperature_val == 0.0 + assert not temperature_is_none + + +@pytest.mark.parametrize( + "temperature,expected_value,is_none", + [(None, 1.0, True), (1, 1.0, False), (0.5, 0.5, False)], +) +def test_validate_softmax_temperature_accepts_scalar( + temperature, expected_value, is_none +): + logits = torch.empty((3, 5), dtype=torch.float32) + + temperature_arr, temperature_val, temperature_is_none = ( + flashinfer.sampling._validate_softmax_temperature(logits, temperature) + ) + + assert temperature_arr is None + assert temperature_val == expected_value + assert temperature_is_none is is_none + + +@pytest.mark.parametrize( + "temperature,error_match", + [ + (torch.tensor(1.0), "must be 1D"), + (torch.ones(3, 1), "must be 1D"), + (torch.ones(2), "length must equal logits.size\\(0\\)"), + (torch.ones(3, dtype=torch.int32), "must have dtype"), + (torch.ones(3, dtype=torch.bool), "must have dtype"), + ], +) +def test_validate_softmax_temperature_rejects_invalid_tensor( + temperature, error_match +): + logits = torch.empty((3, 5), dtype=torch.float32) + + with pytest.raises(ValueError, match=error_match): + flashinfer.sampling._validate_softmax_temperature(logits, temperature) + + +def test_validate_softmax_temperature_rejects_other_device(): + logits = torch.empty((3, 5), dtype=torch.float32) + temperature = torch.ones(3, device="meta") + + with pytest.raises(ValueError, match="same device as logits"): + flashinfer.sampling._validate_softmax_temperature(logits, temperature) + + +@pytest.mark.parametrize("temperature", [True, "1.0", object()]) +def test_validate_softmax_temperature_rejects_non_numeric_scalar(temperature): + logits = torch.empty((3, 5), dtype=torch.float32) + + with pytest.raises(TypeError, match="temperature must be None, a real scalar"): + flashinfer.sampling._validate_softmax_temperature(logits, temperature) + + +@pytest.mark.parametrize("rows", [16, 32, 64, 128, 512, 1024]) +def test_blackwell_mr515_none_route_is_exact_and_correct(rows): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("MR515 hybrid route is restricted to sm_103a") + + logits = torch.randn((rows, 32000), device="cuda", dtype=torch.float32) + route = flashinfer.sampling._blackwell_softmax_route_for_testing( + logits, temperature=None, enable_pdl=False + ) + actual = flashinfer.sampling.softmax(logits, temperature=None, enable_pdl=False) + expected = torch.softmax(logits, dim=-1) + + assert route == flashinfer.sampling._BLACKWELL_SOFTMAX_ROUTE_MR515_V32000_T512 + assert actual.data_ptr() != logits.data_ptr() + torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3) + + +def test_blackwell_mr515_scalar_one_pdl_route_is_exact_and_correct(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("MR515 hybrid route is restricted to sm_103a") + + logits = torch.randn((64, 32000), device="cuda", dtype=torch.float32) + route = flashinfer.sampling._blackwell_softmax_route_for_testing( + logits, temperature=1.0, enable_pdl=True + ) + actual = flashinfer.sampling.softmax(logits, temperature=1.0, enable_pdl=True) + expected = torch.softmax(logits, dim=-1) + + assert route == flashinfer.sampling._BLACKWELL_SOFTMAX_ROUTE_MR515_V32000_T512 + assert actual.data_ptr() != logits.data_ptr() + torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3) + + +@pytest.mark.parametrize( + "rows,temperature,enable_pdl", + [ + (4, None, False), + (256, None, False), + (64, None, True), + (64, 1.0, False), + (64, 0.5, True), + ], +) +def test_blackwell_mr515_route_does_not_interpolate(rows, temperature, enable_pdl): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("MR515 hybrid route is restricted to sm_103a") + + logits = torch.empty((rows, 32000), device="cuda", dtype=torch.float32) + route = flashinfer.sampling._blackwell_softmax_route_for_testing( + logits, temperature=temperature, enable_pdl=enable_pdl + ) + + assert route != flashinfer.sampling._BLACKWELL_SOFTMAX_ROUTE_MR515_V32000_T512 + + @pytest.mark.parametrize("batch_size", [1, 99, 989]) @pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) @pytest.mark.parametrize( From e89d2dddfdcdf26ca43e72d714ca42b9aa1605d4 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 5 Aug 2026 15:28:38 -0700 Subject: [PATCH 11/14] Apply FlashInfer formatting to the hybrid dispatcher --- csrc/blackwell_softmax.cu | 30 ++++++++++-------------- include/flashinfer/blackwell_softmax.cuh | 3 +-- tests/utils/test_sampling.py | 4 +--- 3 files changed, 15 insertions(+), 22 deletions(-) diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu index ad9e39ba748..49dff1d45c1 100644 --- a/csrc/blackwell_softmax.cu +++ b/csrc/blackwell_softmax.cu @@ -69,8 +69,7 @@ bool use_mr515_none_row(uint32_t rows) { } bool use_mr515_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind, - float temperature_val, bool enable_pdl, int device_major, - int device_minor) { + float temperature_val, bool enable_pdl, int device_major, int device_minor) { if (device_major != 10 || device_minor != 3 || vocab_size != 32000) { return false; } @@ -102,16 +101,16 @@ bool use_rowwise_kernel(uint32_t rows, uint32_t vocab_size, ParameterKind parame dense_aligned_high_row_wide || measured_large_odd; } -SoftmaxRoute select_softmax_route(uint32_t rows, uint32_t vocab_size, - ParameterKind parameter_kind, float temperature_val, - bool enable_pdl, int device_major, int device_minor) { +SoftmaxRoute select_softmax_route(uint32_t rows, uint32_t vocab_size, ParameterKind parameter_kind, + float temperature_val, bool enable_pdl, int device_major, + int device_minor) { if (rows == 0 || vocab_size == 0 || static_cast(rows) * vocab_size > static_cast(std::numeric_limits::max())) { return SoftmaxRoute::kFallback; } - if (use_mr515_kernel(rows, vocab_size, parameter_kind, temperature_val, enable_pdl, - device_major, device_minor)) { + if (use_mr515_kernel(rows, vocab_size, parameter_kind, temperature_val, enable_pdl, device_major, + device_minor)) { return SoftmaxRoute::kMr515V32000T512; } if (use_warp_kernel(rows, vocab_size, parameter_kind)) { @@ -209,8 +208,7 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper &vocab_size_i, ¶meter_kind_i, &temperature_val}; return launch_noncooperative_kernel( reinterpret_cast(kernel_flashinfer_blackwell_softmax_followup_rowwise), - dim3(rows), dim3(kRowwiseThreads), args, kRowwiseDynamicSmemBytes, - launch_with_pdl, stream); + dim3(rows), dim3(kRowwiseThreads), args, kRowwiseDynamicSmemBytes, launch_with_pdl, stream); } int active_blocks_per_sm = 0; @@ -228,7 +226,7 @@ cudaError_t launch_blackwell_softmax(float* logits, float* output, float* temper int device = 0; if ((status = cudaGetDevice(&device)) != cudaSuccess || (status = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device)) != - cudaSuccess) { + cudaSuccess) { return status; } @@ -311,17 +309,16 @@ SoftmaxRoute query_softmax_route(TensorView logits, Optional maybe_t int device_major = 0; int device_minor = 0; cudaError_t status = cudaDeviceGetAttribute(&device_major, cudaDevAttrComputeCapabilityMajor, - logits.device().device_id); + logits.device().device_id); if (status == cudaSuccess) { status = cudaDeviceGetAttribute(&device_minor, cudaDevAttrComputeCapabilityMinor, logits.device().device_id); } TVM_FFI_ICHECK(status == cudaSuccess) << "Blackwell Softmax route query failed with error code " << cudaGetErrorString(status); - return select_softmax_route(static_cast(logits.size(0)), - static_cast(logits.size(1)), parameter_kind, - static_cast(temperature_val), enable_pdl, device_major, - device_minor); + return select_softmax_route( + static_cast(logits.size(0)), static_cast(logits.size(1)), parameter_kind, + static_cast(temperature_val), enable_pdl, device_major, device_minor); } void blackwell_softmax_impl(TensorView workspace_buffer, TensorView logits, TensorView output, @@ -371,8 +368,7 @@ void blackwell_softmax(TensorView workspace_buffer, TensorView logits, TensorVie } int64_t blackwell_softmax_route(TensorView logits, Optional maybe_temperature_arr, - double temperature_val, bool enable_pdl, - bool temperature_is_none) { + double temperature_val, bool enable_pdl, bool temperature_is_none) { return static_cast(query_softmax_route(logits, maybe_temperature_arr, temperature_val, enable_pdl, temperature_is_none)); } diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index e996fca838c..6fff8d07a76 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -47,7 +47,6 @@ __global__ void kernel_flashinfer_blackwell_softmax_followup_warp(float* logits, float scalar_temperature); __global__ void kernel_mr474_manual_softmax_exp2_t512_vec4(float* logits, float* temperature, - float* output, - float scalar_temperature); + float* output, float scalar_temperature); } // extern "C" diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index aec41fc38ca..846dbc428bd 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -96,9 +96,7 @@ def test_validate_softmax_temperature_accepts_scalar( (torch.ones(3, dtype=torch.bool), "must have dtype"), ], ) -def test_validate_softmax_temperature_rejects_invalid_tensor( - temperature, error_match -): +def test_validate_softmax_temperature_rejects_invalid_tensor(temperature, error_match): logits = torch.empty((3, 5), dtype=torch.float32) with pytest.raises(ValueError, match=error_match): From 947d9c82e9ef6878849d27d5e2ad7ccf57570b4a Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 5 Aug 2026 15:31:52 -0700 Subject: [PATCH 12/14] Clarify Blackwell softmax PDL routing --- flashinfer/sampling.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/flashinfer/sampling.py b/flashinfer/sampling.py index 13a7479341b..8cbf17aeb3b 100644 --- a/flashinfer/sampling.py +++ b/flashinfer/sampling.py @@ -904,6 +904,9 @@ def softmax( enable_pdl : Optional[bool] Whether to enable Programmatic Dependent Launch (PDL) for improved performance on supported hardware. If None (default), PDL will be automatically enabled on devices with compute capability >= 9.0. + The custom Blackwell softmax routes currently apply PDL on SM103; the + SM100 fast path remains non-PDL, while fallback routes preserve their + existing PDL behavior. Returns ------- probs : torch.Tensor From 082c7f54081625de4d138e3a581a6dab3b1ef225 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 5 Aug 2026 16:51:36 -0700 Subject: [PATCH 13/14] Fix Blackwell softmax edge semantics --- csrc/blackwell_softmax.cu | 3 ++ csrc/blackwell_softmax_bootstrap.cu | 4 ++ csrc/blackwell_softmax_mr515_exp2.cu | 4 +- flashinfer/sampling.py | 6 +++ include/flashinfer/blackwell_softmax.cuh | 2 + tests/utils/test_sampling.py | 47 ++++++++++++++++++++++++ 6 files changed, 65 insertions(+), 1 deletion(-) diff --git a/csrc/blackwell_softmax.cu b/csrc/blackwell_softmax.cu index 49dff1d45c1..244a84197d6 100644 --- a/csrc/blackwell_softmax.cu +++ b/csrc/blackwell_softmax.cu @@ -316,6 +316,9 @@ SoftmaxRoute query_softmax_route(TensorView logits, Optional maybe_t } TVM_FFI_ICHECK(status == cudaSuccess) << "Blackwell Softmax route query failed with error code " << cudaGetErrorString(status); + if (device_major != 10 || (device_minor != 0 && device_minor != 3)) { + return SoftmaxRoute::kFallback; + } return select_softmax_route( static_cast(logits.size(0)), static_cast(logits.size(1)), parameter_kind, static_cast(temperature_val), enable_pdl, device_major, device_minor); diff --git a/csrc/blackwell_softmax_bootstrap.cu b/csrc/blackwell_softmax_bootstrap.cu index 06880bb583c..c7ef67612e4 100644 --- a/csrc/blackwell_softmax_bootstrap.cu +++ b/csrc/blackwell_softmax_bootstrap.cu @@ -377,6 +377,8 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" :: "l"((void*)(output + (index_4))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3), "r"(_stv8_3_4), "r"(_stv8_3_5), "r"(_stv8_3_6), "r"(_stv8_3_7) : "memory"); #else + // This path is used only when splits == 1, so no other block + // reads output while the two 128-bit stores are in flight. asm volatile( "st.global.v4.b32 [%0], {%1, %2, %3, %4};" :: "l"((void*)(output + (index_4))), "r"(_stv8_3_0), "r"(_stv8_3_1), "r"(_stv8_3_2), "r"(_stv8_3_3) : "memory"); @@ -543,6 +545,8 @@ kernel_flashinfer_blackwell_softmax_bootstrap_seed(float* __restrict__ x, float* "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" :: "l"((void*)(output + (index_6))), "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3), "r"(_stv8_5_4), "r"(_stv8_5_5), "r"(_stv8_5_6), "r"(_stv8_5_7) : "memory"); #else + // Phase 2 begins after the grid sync and no later phase reads + // output, so replacing one 256-bit store with two is safe here. asm volatile( "st.global.v4.b32 [%0], {%1, %2, %3, %4};" :: "l"((void*)(output + (index_6))), "r"(_stv8_5_0), "r"(_stv8_5_1), "r"(_stv8_5_2), "r"(_stv8_5_3) : "memory"); diff --git a/csrc/blackwell_softmax_mr515_exp2.cu b/csrc/blackwell_softmax_mr515_exp2.cu index 1c89a154a46..646b2fe67d1 100644 --- a/csrc/blackwell_softmax_mr515_exp2.cu +++ b/csrc/blackwell_softmax_mr515_exp2.cu @@ -84,6 +84,8 @@ kernel_mr474_manual_softmax_exp2_t512_vec4(float* __restrict__ x, float* __restr const int warp = make_warp_uniform(tid / 32); const int lane = tid % 32; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030 + // The integration passes nullptr only to encode PDL for this frozen + // TEMP_KIND==0 specialization; the temperature slot is never dereferenced. const bool enable_pdl = temperature == nullptr; if (enable_pdl) { asm volatile("griddepcontrol.wait;" ::: "memory"); @@ -212,7 +214,7 @@ kernel_mr474_manual_softmax_exp2_t512_vec4(float* __restrict__ x, float* __restr for (int j_2 = 0; j_2 < 4; j_2++) { float value_2 = _vec_load_3[j_2]; float _exp2_1 = approx_exp2((value_2 - row_max) * 1.4426950408889634f); - out_values[j_2] = ((row_max == -LOOM_INF) ? 0.0f : _exp2_1 * inv_sum); + out_values[j_2] = _exp2_1 * inv_sum; } } { diff --git a/flashinfer/sampling.py b/flashinfer/sampling.py index 8cbf17aeb3b..22eaab25d2c 100644 --- a/flashinfer/sampling.py +++ b/flashinfer/sampling.py @@ -119,6 +119,7 @@ def _fake_blackwell_softmax( return blackwell_softmax +_BLACKWELL_SOFTMAX_ROUTE_FALLBACK = 0 _BLACKWELL_SOFTMAX_ROUTE_MR515_V32000_T512 = 4 @@ -130,6 +131,11 @@ def _blackwell_softmax_route_for_testing( """Return the C++ dispatcher route ID without launching a kernel.""" if not logits.is_cuda: raise ValueError(f"logits must be a CUDA tensor, got device={logits.device}") + device_index = logits.device.index + if device_index is None: + device_index = torch.cuda.current_device() + if not _supports_blackwell_softmax(device_index): + return _BLACKWELL_SOFTMAX_ROUTE_FALLBACK maybe_temperature_arr, temperature_val, temperature_is_none = ( _validate_softmax_temperature(logits, temperature) ) diff --git a/include/flashinfer/blackwell_softmax.cuh b/include/flashinfer/blackwell_softmax.cuh index 6fff8d07a76..37cf01411da 100644 --- a/include/flashinfer/blackwell_softmax.cuh +++ b/include/flashinfer/blackwell_softmax.cuh @@ -46,6 +46,8 @@ __global__ void kernel_flashinfer_blackwell_softmax_followup_warp(float* logits, int parameter_kind, float scalar_temperature); +// For this frozen TEMP_KIND==0 specialization only, a null temperature pointer +// is an integration-private signal that PDL is enabled; the kernel never reads it. __global__ void kernel_mr474_manual_softmax_exp2_t512_vec4(float* logits, float* temperature, float* output, float scalar_temperature); diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index 846dbc428bd..c936d446556 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -174,6 +174,53 @@ def test_blackwell_mr515_route_does_not_interpolate(rows, temperature, enable_pd assert route != flashinfer.sampling._BLACKWELL_SOFTMAX_ROUTE_MR515_V32000_T512 +@pytest.mark.parametrize( + "rows,vocab_size,temperature,enable_pdl,expected_route", + [ + (1, 111, 0.5, False, 1), + (256, 32000, None, False, 2), + (1, 32000, None, False, 3), + (64, 32000, None, False, 4), + ], +) +def test_blackwell_softmax_all_negative_infinity_matches_public_semantics( + rows, vocab_size, temperature, enable_pdl, expected_route +): + capability = torch.cuda.get_device_capability() + if capability not in ((10, 0), (10, 3)): + pytest.skip("Blackwell Softmax routes require SM100 or SM103") + if expected_route == 4 and capability != (10, 3): + pytest.skip("MR515 hybrid route is restricted to sm_103a") + + logits = torch.full((rows, vocab_size), -torch.inf, device="cuda") + route = flashinfer.sampling._blackwell_softmax_route_for_testing( + logits, temperature=temperature, enable_pdl=enable_pdl + ) + actual = flashinfer.sampling.softmax( + logits, temperature=temperature, enable_pdl=enable_pdl + ) + expected = torch.softmax(logits, dim=-1) + + assert route == expected_route + assert actual.data_ptr() != logits.data_ptr() + assert torch.isnan(expected).all() + assert torch.isnan(actual).all() + + +def test_blackwell_route_observer_falls_back_off_blackwell(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if torch.cuda.get_device_capability() in ((10, 0), (10, 3)): + pytest.skip("This test exercises the non-Blackwell observer gate") + + logits = torch.empty((1, 111), device="cuda") + route = flashinfer.sampling._blackwell_softmax_route_for_testing( + logits, temperature=0.5, enable_pdl=False + ) + + assert route == flashinfer.sampling._BLACKWELL_SOFTMAX_ROUTE_FALLBACK + + @pytest.mark.parametrize("batch_size", [1, 99, 989]) @pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) @pytest.mark.parametrize( From f0edac69e1c4d299c14fa95b7a0073a172203cf9 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 5 Aug 2026 23:38:21 -0700 Subject: [PATCH 14/14] fix(logits-processor): honor softmax architecture dispatch --- flashinfer/logits_processor/operators.py | 19 +++++++---------- tests/utils/test_sampling.py | 27 ++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 11 deletions(-) diff --git a/flashinfer/logits_processor/operators.py b/flashinfer/logits_processor/operators.py index 4df6eafdd3c..291479c4dee 100644 --- a/flashinfer/logits_processor/operators.py +++ b/flashinfer/logits_processor/operators.py @@ -18,7 +18,10 @@ import torch -from flashinfer.sampling import get_sampling_module +from flashinfer.sampling import ( + get_sampling_module, + softmax as sampling_softmax, +) from flashinfer.utils import _get_cache_buf, device_support_pdl from .op import ParameterizedOp @@ -391,20 +394,14 @@ def __call__(self, tensor: TaggedTensor, **kwargs: Any) -> TaggedTensor: ): raise ValueError("Temperature must be positive float or a tensor array") - workspace_buffer = _get_cache_buf( - "softmax_workspace", 1024 * 1024, tensor.data.device - ) - enable_pdl = self.default_params.get("enable_pdl", None) if enable_pdl is None: enable_pdl = device_support_pdl(tensor.data.device) - probs = get_sampling_module().softmax( - workspace_buffer, - tensor.data, - maybe_temperature_arr, - temperature_val, - enable_pdl, + # Keep fused pipelines on the public dispatch path so architecture- + # specific softmax routes match direct flashinfer.sampling.softmax. + probs = sampling_softmax( + logits=tensor.data, temperature=temperature, enable_pdl=enable_pdl ) return TaggedTensor(probs, output_type) diff --git a/tests/utils/test_sampling.py b/tests/utils/test_sampling.py index c936d446556..2a1a2f01533 100644 --- a/tests/utils/test_sampling.py +++ b/tests/utils/test_sampling.py @@ -293,6 +293,33 @@ def test_softmax_blackwell_routes(batch_size, vocab_size, temperature_kind): assert torch.allclose(probs, probs_ref, atol=1e-5) +@pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) +def test_softmax_blackwell_random_per_row_temperature_contract(vocab_size): + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("Loom Softmax routes require SM100 or SM103") + + torch.manual_seed(42) + rows = 989 + logits = torch.randn(rows, vocab_size, device="cuda") + temperature = torch.rand(rows, device="cuda") + route = flashinfer.sampling._blackwell_softmax_route_for_testing( + logits, temperature=temperature, enable_pdl=True + ) + probs = flashinfer.sampling.softmax( + logits, temperature=temperature, enable_pdl=True + ) + probs_ref = torch.softmax(logits / temperature[:, None], dim=-1) + + assert route != flashinfer.sampling._BLACKWELL_SOFTMAX_ROUTE_FALLBACK + assert probs.data_ptr() != logits.data_ptr() + torch.testing.assert_close( + probs, + probs_ref, + atol=1e-3, + rtol=1e-3, + ) + + @pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) @pytest.mark.parametrize( "distribution",