From 3bba4d2da8ce6321e485b002763f29dd15799a04 Mon Sep 17 00:00:00 2001 From: sherylll Date: Thu, 8 Oct 2026 17:11:16 -0700 Subject: [PATCH 1/5] Add fp8 variant --- .../neighbors/detail/cagra/graph_shared.cu | 2 +- cpp/src/neighbors/detail/nn_descent.cuh | 189 ++++--------- .../neighbors/detail/promoted_i4_warp_mma.cuh | 249 ++++++++++++++++++ .../quantize/detail/bbq_distance.cuh | 12 +- 4 files changed, 311 insertions(+), 141 deletions(-) create mode 100644 cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index 8dcb4d8b6f..d65711fc21 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -179,7 +179,7 @@ __device__ __forceinline__ float bbq_row_distance( namespace bbq = cuvs::preprocessing::quantize::bbq; const uint32_t raw = bbq::code_inner_product(quantizer_document, quantizer_query, row_document, row_query); - return bbq::bbq_calculate_metric(raw, + return bbq::bbq_calculate_metric(static_cast(raw), bbq::get_dequant_factors(quantizer_document, row_document), bbq::get_dequant_factors(quantizer_query, row_query), quantizer_document, diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 37528e0461..64f7a7919e 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -8,6 +8,7 @@ #include "ann_utils.cuh" #include "neighbors_device_intrinsics.cuh" #include "nn_descent_gnnd.hpp" +#include "promoted_i4_warp_mma.cuh" #include "../../core/nvtx.hpp" #include "../../core/omp_wrapper.hpp" @@ -1508,7 +1509,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) // the min-search loops below only visit idx_in_list < new_size. if (row0 < new_size) { const Index_t doc_id0 = new_neighbors[row0]; - s_distances[distance0] = bbq_calculate_metric(acc0[k], + s_distances[distance0] = bbq_calculate_metric(static_cast(acc0[k]), s_document_factors[row0], my_col_factors, dataset_document, @@ -1521,7 +1522,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) if (row0 + 1 < new_size) { const Index_t doc_id1 = new_neighbors[row0 + 1]; s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = - bbq_calculate_metric(acc1[k], + bbq_calculate_metric(static_cast(acc1[k]), s_document_factors[row0 + 1], my_col_factors, dataset_document, @@ -1624,7 +1625,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) // otherwise bounded, and this cell is never read downstream either way). if (row0 < new_size) { const Index_t doc_id0 = new_neighbors[row0]; - s_distances[distance0] = bbq_calculate_metric(acc0_old[k], + s_distances[distance0] = bbq_calculate_metric(static_cast(acc0_old[k]), s_document_factors[row0], my_old_col_factors, dataset_document, @@ -1637,7 +1638,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) if (row0 + 1 < new_size) { const Index_t doc_id1 = new_neighbors[row0 + 1]; s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = - bbq_calculate_metric(acc1_old[k], + bbq_calculate_metric(static_cast(acc1_old[k]), s_document_factors[row0 + 1], my_old_col_factors, dataset_document, @@ -1682,34 +1683,22 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) } } -// Masks one nibble half of a packed-4b fragment into `dst`, normalised to 0..15. The high half is -// shifted down so both halves reach the MMA at the same scale and can share one accumulator. -template -__device__ __forceinline__ void select_packed_nibble(const FragT& src, FragT& dst) -{ - constexpr int kWords = FragT::num_elements / 4; - static_assert(FragT::num_elements % 4 == 0, "fragment must be a whole number of u32 words"); - const auto* w = reinterpret_cast(src.x); - auto* d = reinterpret_cast(dst.x); -#pragma unroll - for (int i = 0; i < kWords; ++i) { - d[i] = High ? ((w[i] >> 4) & 0x0f0f0f0fu) : (w[i] & 0x0f0f0f0fu); - } -} +// Selects fp8 (e4m3) mma.sync instead of u8 wmma for the BBQ local join below. fp8 mma.sync +// needs sm_89+, so lower archs keep the wmma path even when this is set. +#ifndef CUVS_NND_BBQ_FP8_MMA +#define CUVS_NND_BBQ_FP8_MMA 0 +#endif +#if CUVS_NND_BBQ_FP8_MMA && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) +#define CUVS_NND_BBQ_USE_FP8_MMA 1 +#else +#define CUVS_NND_BBQ_USE_FP8_MMA 0 +#endif // int4 tensor-core BBQ local join, covering both the symmetric (one quantizer, self-join) and -// asymmetric (two quantizers) cases. Modeled on local_join_kernel_wmma: nvcuda::wmma fragments -// (u4 x u4 -> s32, shape m8n8k32) replace the popc/dp4a inner product; the accumulator lives in -// registers across the whole K reduction and is stored to s_distances once per (row-tile, -// col-tile), not accumulated into shared memory every step like the scalar kernel. -// -// Warp tiling: num_warps = BLOCK_SIZE/32 warps arranged as a WARPS_PER_DIM x WARPS_PER_DIM square -// grid (WARPS_PER_DIM=4 so 4x4=16=num_warps), each warp owning a (MAX_NUM_BI_SAMPLES/WARPS_PER_DIM) -// region of the MAX_NUM_BI_SAMPLES x MAX_NUM_BI_SAMPLES output matrix, same as -// local_join_kernel_wmma's WMMA_M=N=16 warp assignment. Since int4 MMA tiles are MMA_M x MMA_N -// (8x8, the only shape nvcuda::wmma exposes for u4), each warp covers its region via a -// SUB_PER_DIM x SUB_PER_DIM grid of native tiles instead of a single call -- SUB_PER_DIM is a -// forced consequence of (MAX_NUM_BI_SAMPLES/MMA_M) / WARPS_PER_DIM, not an arbitrary choice. +// asymmetric (two quantizers) cases. Modeled on local_join_kernel_wmma: a warp-level i4 MMA (see +// promoted_i4_warp_mma.cuh) replaces the popc/dp4a inner product; the accumulator lives in +// registers across the whole K reduction and is stored to s_distances once per phase, not +// accumulated into shared memory every step like the scalar kernel. template = 750) - using namespace nvcuda; - constexpr int MMA_M = 16; - constexpr int MMA_N = 16; - // A k-step reads 16 packed bytes = 32 codes; each nibble half is a separate MMA over MMA_K. - constexpr int MMA_K = 16; - // num_warps = BLOCK_SIZE/32 = 16, arranged as a square WARPS_PER_DIM x WARPS_PER_DIM grid since - // 4*4=16 matches exactly; the static_assert is what actually enforces this holds for the - // current BLOCK_SIZE, WARPS_PER_DIM itself isn't derived (no trivial constexpr integer sqrt). + // The warps form a square WARPS_PER_DIM x WARPS_PER_DIM grid, each owning a WARP_TILE x WARP_TILE + // block of the MAX_NUM_BI_SAMPLES x MAX_NUM_BI_SAMPLES output; warp_mma covers that block with + // whatever native MMA shape it uses. WARPS_PER_DIM isn't derived (no trivial constexpr integer + // sqrt), so the static_assert ties it to BLOCK_SIZE. constexpr int WARPS_PER_DIM = 4; static_assert(WARPS_PER_DIM * WARPS_PER_DIM == BLOCK_SIZE / raft::warp_size(), "warp grid must be square and match num_warps = BLOCK_SIZE/32"); - // Each warp owns a WARP_TILE x WARP_TILE region of the MAX_NUM_BI_SAMPLES x MAX_NUM_BI_SAMPLES - // output matrix. TILES_PER_DIM is how many native MMA_M x MMA_N tiles span one output dimension; - // SUB_PER_DIM (native tiles per warp per dim) is a forced consequence of TILES_PER_DIM / - // WARPS_PER_DIM, not an arbitrary choice -- it's 2 here only because 8/4=2 for these particular - // MAX_NUM_BI_SAMPLES/MMA_M/WARPS_PER_DIM values. - static_assert(MAX_NUM_BI_SAMPLES % MMA_M == 0 && MMA_M == MMA_N, - "MAX_NUM_BI_SAMPLES must divide evenly into square MMA_MxMMA_N tiles"); - constexpr int TILES_PER_DIM = MAX_NUM_BI_SAMPLES / MMA_M; - static_assert(TILES_PER_DIM % WARPS_PER_DIM == 0, - "warps must evenly tile the native MMA tiles in each output dimension"); - constexpr int SUB_PER_DIM = TILES_PER_DIM / WARPS_PER_DIM; - constexpr int WARP_TILE = SUB_PER_DIM * MMA_M; + constexpr int WARP_TILE = MAX_NUM_BI_SAMPLES / WARPS_PER_DIM; + static_assert(MAX_NUM_BI_SAMPLES % WARPS_PER_DIM == 0, "warps must evenly tile the output"); constexpr int BBQ_ROW_BYTES = 128; constexpr int MMA_PAD = 16; - static_assert(MMA_PAD % 16 == 0, "row padding must preserve 16-byte IMMA row alignment"); - // Strides are in fragment elements (u8), so plain byte counts, though the data is packed 4b. - constexpr int K_STEPS_PER_TILE = BBQ_ROW_BYTES / MMA_K; - constexpr int KSTEP_BYTES = MMA_K; - constexpr int ROW_STRIDE_U8 = BBQ_ROW_BYTES + MMA_PAD; + // Strides are in bytes, though the data is packed 4b. + constexpr int ROW_STRIDE = BBQ_ROW_BYTES + MMA_PAD; +#if CUVS_NND_BBQ_USE_FP8_MMA + using warp_mma = i4_warp_mma_fp8; +#else + using warp_mma = i4_warp_mma_u8; +#endif + static_assert(BBQ_ROW_BYTES % warp_mma::KSTEP_BYTES == 0, "a tile must be whole k-steps"); constexpr int MMA_STORE_STRIDE = SKEWED_MAX_NUM_BI_SAMPLES; @@ -1775,15 +1749,14 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) // operand: the `new` list in phase 1 and the `old` list in phase 2, query quantizer. Under // SelfJoin, phase 1 leaves s_col_vec untouched and both fragments read s_row_vec. __shared__ int s_list[MAX_NUM_BI_SAMPLES * 2]; - __shared__ __align__(16) uint8_t s_row_vec[MAX_NUM_BI_SAMPLES][BBQ_ROW_BYTES + MMA_PAD]; + __shared__ __align__(16) uint8_t s_row_vec[MAX_NUM_BI_SAMPLES][ROW_STRIDE]; // s_col_vec aliases s_distances's memory instead of its own array: their live ranges never // overlap (col_vec fully consumed by the kk-loop before distances are stored; distances fully // consumed by the min-search loop before the next phase re-stages col_vec). __shared__ __align__(16) DistData_t s_distances[MAX_NUM_BI_SAMPLES * MMA_STORE_STRIDE]; - static_assert(sizeof(s_distances) >= MAX_NUM_BI_SAMPLES * (BBQ_ROW_BYTES + MMA_PAD), + static_assert(sizeof(s_distances) >= MAX_NUM_BI_SAMPLES * ROW_STRIDE, "s_col_vec aliases s_distances's memory and must fit inside it"); - auto(*s_col_vec)[BBQ_ROW_BYTES + MMA_PAD] = - reinterpret_cast(s_distances); + auto(*s_col_vec)[ROW_STRIDE] = reinterpret_cast(s_distances); __shared__ int s_unique_counter[2]; // Document-side (row axis) dequant factors -- see the identical buffer and full rationale in // local_join_kernel_bbq_simt. Staged once below, reused unchanged by both phases (rows are @@ -1836,25 +1809,21 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) const int warp_id_x = warp_id % WARPS_PER_DIM; // One phase: accumulate s_row_vec (rows, `new` list) against col_buf over the whole K range, - // then store the accumulators to s_distances (converted to final float distances in place right - // after). col_buf is s_row_vec itself when phase 1 is a - // self-join, otherwise s_col_vec. col_neighbors/col_size select which list the B operand stages. - // alias_tag: staging skip and b_frag source + // store the raw sums to s_distances, then convert them to final float distances in place. + // col_buf is s_row_vec itself when phase 1 is a self-join, otherwise s_col_vec. + // col_neighbors/col_size select which list the B operand stages. // row_resident: the A operand is already staged from a previous phase, which is only true // when the whole row fits one tile (n_tiles == 1) so nothing overwrote it. auto run_phase = [&](const Index_t* col_neighbors, int col_size, bool row_resident) { const bool alias_col = (DocumentLayout == QueryLayout) && (col_neighbors == new_neighbors); - wmma::fragment c_frag[SUB_PER_DIM][SUB_PER_DIM]; - for (int msub = 0; msub < SUB_PER_DIM; ++msub) { - for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { - wmma::fill_fragment(c_frag[msub][nsub], 0); - } - } - + const auto* col_buf = alias_col ? s_row_vec : s_col_vec; + const int row0 = warp_id_y * WARP_TILE; + const int col0 = warp_id_x * WARP_TILE; // Lists are usually much shorter than MAX_NUM_BI_SAMPLES, so a warp whose whole block starts - // past either extent only produces cells the epilogue discards -- skip its MMA and its store. - // Staging is cooperative, so this must not skip either __syncthreads() in the step loop. - const bool warp_active = warp_id_y * WARP_TILE < new_size && warp_id_x * WARP_TILE < col_size; + // past either extent only produces cells the metric pass discards -- skip its MMA and its + // store. Staging is cooperative, so this must not skip either __syncthreads() in the step loop. + const bool warp_active = row0 < new_size && col0 < col_size; + warp_mma mma; for (int step = 0; step < n_tiles; ++step) { if (!row_resident) { @@ -1866,64 +1835,17 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) s_col_vec, dataset_query, col_neighbors, col_size, step, query_row_bytes); } __syncthreads(); - - // a_frag depends only on (msub, kk); b_frag depends only on (nsub, kk) -- load each once - // per kk and reuse across the other sub-tile index, instead of reloading redundantly inside - // a full msub x nsub x kk cross product. - // Deliberately not #pragma unroll'd: full unrolling here keeps more fragment live ranges - // simultaneous, driving register pressure up -- letting the compiler pick reduces that at - // the cost of some intra-warp ILP. - for (int kk = 0; warp_active && kk < K_STEPS_PER_TILE; ++kk) { - wmma::fragment - a_frag[SUB_PER_DIM]; - wmma::fragment - b_frag[SUB_PER_DIM]; - const auto* col_buf = alias_col ? s_row_vec : s_col_vec; // compile-time select - for (int msub = 0; msub < SUB_PER_DIM; ++msub) { - const int row0 = warp_id_y * WARP_TILE + msub * MMA_M; - wmma::load_matrix_sync(a_frag[msub], s_row_vec[row0] + kk * KSTEP_BYTES, ROW_STRIDE_U8); - } - for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { - const int col0 = warp_id_x * WARP_TILE + nsub * MMA_N; - wmma::load_matrix_sync(b_frag[nsub], col_buf[col0] + kk * KSTEP_BYTES, ROW_STRIDE_U8); - } - // One scratch fragment per operand, reused across the halves. - std::remove_reference_t a_half[SUB_PER_DIM]; - std::remove_reference_t b_half[SUB_PER_DIM]; - auto accumulate_half = [&](auto high_tag) { - constexpr bool High = decltype(high_tag)::value; - for (int msub = 0; msub < SUB_PER_DIM; ++msub) { - select_packed_nibble(a_frag[msub], a_half[msub]); - } - for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { - select_packed_nibble(b_frag[nsub], b_half[nsub]); - } - for (int msub = 0; msub < SUB_PER_DIM; ++msub) { - for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { - wmma::mma_sync(c_frag[msub][nsub], a_half[msub], b_half[nsub], c_frag[msub][nsub]); - } - } - }; - accumulate_half(std::false_type{}); - accumulate_half(std::true_type{}); - } + if (warp_active) { mma.template accumulate(s_row_vec, col_buf, row0, col0); } __syncthreads(); } - for (int msub = 0; warp_active && msub < SUB_PER_DIM; ++msub) { - const int row0 = warp_id_y * WARP_TILE + msub * MMA_M; - for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { - const int col0 = warp_id_x * WARP_TILE + nsub * MMA_N; - wmma::store_matrix_sync( - reinterpret_cast(s_distances) + row0 * MMA_STORE_STRIDE + col0, - c_frag[msub][nsub], - MMA_STORE_STRIDE, - wmma::mem_row_major); - } - } + // Storing first ends the accumulators' live range before the metric code runs, instead of + // carrying both at once, and spreads the metric over every thread rather than the active warps. + using raw_t = typename warp_mma::raw_t; + auto* raw_view = reinterpret_cast(s_distances); + if (warp_active) { mma.template store(raw_view, row0, col0); } __syncthreads(); - // Converts store_matrix_sync's raw int32 dot products into final float distances, in place // col = i % MAX_NUM_BI_SAMPLES is invariant across a thread's own iterations // (BLOCK_SIZE is a multiple of MAX_NUM_BI_SAMPLES), so this thread's column factors are // fetched once and reused below @@ -1933,9 +1855,8 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) my_col_factors = get_dequant_factors(dataset_query, col_neighbors[my_col]); } - // Cell (row, col) is read (as raw int, via this alias) and written (as the final float) at - // the same address by the same thread - auto* raw_view = reinterpret_cast(s_distances); + // Cell (row, col) is read (as raw_t, via raw_view) and written (as the final float) at the + // same address by the same thread constexpr int total_cells = MAX_NUM_BI_SAMPLES * MAX_NUM_BI_SAMPLES; for (int i = tx; i < total_cells; i += BLOCK_SIZE) { const int row = i / MAX_NUM_BI_SAMPLES; @@ -1944,7 +1865,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) const int distance0 = row * MMA_STORE_STRIDE + col; const Index_t doc_id = new_neighbors[row]; const Index_t query_id = col_neighbors[col]; - const uint32_t raw = static_cast(raw_view[distance0]); + const float raw = static_cast(raw_view[distance0]); s_distances[distance0] = bbq_calculate_metric(raw, s_document_factors[row], my_col_factors, diff --git a/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh b/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh new file mode 100644 index 0000000000..c4ce31ad50 --- /dev/null +++ b/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh @@ -0,0 +1,249 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +#include +#include + +// Warp-level tensor-core dot products of i4 codes (0..15, two per byte) staged in shared memory as +// rows of RowStride bytes. A warp owns a WarpTile x WarpTile block of the output: accumulate() adds +// the dot products over one staged K range, store() writes the raw sums as raw_t. The u8 (wmma) +// and fp8 (mma.sync) variants share this interface, so callers can switch between them. + +namespace cuvs::neighbors::nn_descent::detail { + +// Masks one nibble half of a packed-i4 fragment into `dst`, normalised to 0..15. The high half is +// shifted down so both halves reach the MMA at the same scale and can share one accumulator. +template +__device__ __forceinline__ void select_packed_nibble(const FragT& src, FragT& dst) +{ + constexpr int kWords = FragT::num_elements / 4; + static_assert(FragT::num_elements % 4 == 0, "fragment must be a whole number of u32 words"); + const auto* w = reinterpret_cast(src.x); + auto* d = reinterpret_cast(dst.x); +#pragma unroll + for (int i = 0; i < kWords; ++i) { + d[i] = High ? ((w[i] >> 4) & 0x0f0f0f0fu) : (w[i] & 0x0f0f0f0fu); + } +} + +// The staged bytes are packed i4; the MMA reads them as u8 and splits the nibbles here rather than +// leaving it to ptxas. u4 MMA is native only through sm_89, and past that ptxas emits this same +// split out of line, once per MMA. Inlining it keeps the IMMAs in the loop body on every arch, and +// the 16x16x16 tile u8 allows drops SUB_M/SUB_N to 1. +template +struct i4_warp_mma_u8 { + using raw_t = int; + + static constexpr int MMA_M = 16; + static constexpr int MMA_N = 16; + // A k-step reads 16 packed bytes = 32 codes; each nibble half is a separate MMA over MMA_K. + static constexpr int MMA_K = 16; + static constexpr int KSTEP_BYTES = MMA_K; + static constexpr int SUB_M = WarpTile / MMA_M; + static constexpr int SUB_N = WarpTile / MMA_N; + static_assert(WarpTile % MMA_M == 0 && WarpTile % MMA_N == 0, + "the warp tile must be a whole number of native MMA tiles"); + static_assert(RowStride % 16 == 0, "row stride must preserve 16-byte IMMA row alignment"); + + using a_frag_t = nvcuda::wmma:: + fragment; + using b_frag_t = nvcuda::wmma:: + fragment; + nvcuda::wmma::fragment c[SUB_M][SUB_N]; + + __device__ __forceinline__ i4_warp_mma_u8() + { + for (int msub = 0; msub < SUB_M; ++msub) { + for (int nsub = 0; nsub < SUB_N; ++nsub) { + nvcuda::wmma::fill_fragment(c[msub][nsub], 0); + } + } + } + + // a_frag depends only on (msub, kk); b_frag depends only on (nsub, kk) -- load each once per kk + // and reuse across the other sub-tile index, instead of reloading redundantly inside a full + // msub x nsub x kk cross product. + // Kept rolled: unrolling the k-steps keeps more fragment live ranges simultaneous and costs a + // block per SM in registers, for some intra-warp ILP. + template + __device__ __forceinline__ void accumulate(const uint8_t (*a)[RowStride], + const uint8_t (*b)[RowStride], + int row0, + int col0) + { +#pragma unroll 1 + for (int kk = 0; kk < RowBytes / KSTEP_BYTES; ++kk) { + a_frag_t a_frag[SUB_M]; + b_frag_t b_frag[SUB_N]; + for (int msub = 0; msub < SUB_M; ++msub) { + nvcuda::wmma::load_matrix_sync( + a_frag[msub], a[row0 + msub * MMA_M] + kk * KSTEP_BYTES, RowStride); + } + for (int nsub = 0; nsub < SUB_N; ++nsub) { + nvcuda::wmma::load_matrix_sync( + b_frag[nsub], b[col0 + nsub * MMA_N] + kk * KSTEP_BYTES, RowStride); + } + // One scratch fragment per operand, reused across the halves. + a_frag_t a_half[SUB_M]; + b_frag_t b_half[SUB_N]; + auto accumulate_half = [&](auto high_tag) { + constexpr bool High = decltype(high_tag)::value; + for (int msub = 0; msub < SUB_M; ++msub) { + select_packed_nibble(a_frag[msub], a_half[msub]); + } + for (int nsub = 0; nsub < SUB_N; ++nsub) { + select_packed_nibble(b_frag[nsub], b_half[nsub]); + } + for (int msub = 0; msub < SUB_M; ++msub) { + for (int nsub = 0; nsub < SUB_N; ++nsub) { + nvcuda::wmma::mma_sync(c[msub][nsub], a_half[msub], b_half[nsub], c[msub][nsub]); + } + } + }; + accumulate_half(std::false_type{}); + accumulate_half(std::true_type{}); + } + } + + template + __device__ __forceinline__ void store(raw_t* dst, int row0, int col0) const + { + for (int msub = 0; msub < SUB_M; ++msub) { + for (int nsub = 0; nsub < SUB_N; ++nsub) { + nvcuda::wmma::store_matrix_sync(dst + (row0 + msub * MMA_M) * Stride + col0 + nsub * MMA_N, + c[msub][nsub], + Stride, + nvcuda::wmma::mem_row_major); + } + } + } +}; + +// fp8 mma.sync needs sm_89+. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) +// Expands the eight 4-bit codes packed in w to e4m3 bytes of the same value (0..15 are exact in +// e4m3): lo gets the codes in w's low 16 bits, hi the codes in its high 16 bits. prmt tables hold +// only 8 bytes, so 0..7 and 8..15 use separate tables; a selector with bit 3 set yields the +// selected byte's sign, 0 here, which zeroes the lanes the other table owns. prmt reads only the +// selector's low 16 bits, hence the shift for hi. +__device__ __forceinline__ void nibbles_to_e4m3(uint32_t w, uint32_t& lo, uint32_t& hi) +{ + constexpr uint32_t k0_3 = 0x44403800u; + constexpr uint32_t k4_7 = 0x4e4c4a48u; + constexpr uint32_t k8_11 = 0x53525150u; + constexpr uint32_t k12_15 = 0x57565554u; + auto lookup = [](uint32_t sel) { + uint32_t below_8, from_8; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(below_8) : "r"(k0_3), "r"(k4_7), "r"(sel)); + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(from_8) : "r"(k8_11), "r"(k12_15), "r"(sel ^ 0x8888u)); + return below_8 | from_8; + }; + lo = lookup(w); + hi = lookup(w >> 16); +} + +// c += a * b over one m16n8k32 tile. The tensor core's fp8 accumulator keeps far fewer bits than +// f32, so chaining c through it drops low bits of the integer sums. Each MMA starts from zero +// instead -- one tile's sum fits exactly -- and the running sum is kept with ordinary f32 adds. +__device__ __forceinline__ void mma_e4m3_m16n8k32(float (&c)[4], + const uint32_t (&a)[4], + const uint32_t (&b)[2]) +{ + float d[4]; + asm( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %10, %10, %10};" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(0.0f)); +#pragma unroll + for (int i = 0; i < 4; ++i) { + c[i] += d[i]; + } +} + +// Feeds the 4-bit codes to mma.sync as e4m3, expanded in registers from plain 32-bit LDS. +template +struct i4_warp_mma_fp8 { + using raw_t = float; + + static constexpr int MMA_M = 16; + static constexpr int MMA_N = 8; + // A k-step reads 16 packed bytes = 32 codes, expanded to e4m3 and consumed by one MMA. + static constexpr int KSTEP_BYTES = 16; + static constexpr int SUB_M = WarpTile / MMA_M; + static constexpr int SUB_N = WarpTile / MMA_N; + static_assert(WarpTile % MMA_M == 0 && WarpTile % MMA_N == 0, + "the warp tile must be a whole number of native MMA tiles"); + static_assert(RowStride % 4 == 0, "rows must stay 4-byte aligned for the u32 fragment loads"); + + float c[SUB_M][SUB_N][4] = {}; + + // m16n8k32 fragment coordinates: `group` is the A row (and +8) and the B column this lane holds; + // `quad` picks its 4-byte slice of the k-step. A and B read the same slice, and expanding it to + // (low 16 bits -> k 0..15, high -> k 16..31) the same way keeps their k orders aligned. + // Kept rolled for the same register reason as the u8 path. + template + __device__ __forceinline__ void accumulate(const uint8_t (*a)[RowStride], + const uint8_t (*b)[RowStride], + int row0, + int col0) + { + const int lane = threadIdx.x % raft::warp_size(); + const int group = lane / 4; + const int quad = lane % 4; +#pragma unroll 1 + for (int kk = 0; kk < RowBytes / KSTEP_BYTES; ++kk) { + const int byte0 = kk * KSTEP_BYTES + quad * 4; + uint32_t a_frag[SUB_M][4]; + uint32_t b_frag[SUB_N][2]; + for (int msub = 0; msub < SUB_M; ++msub) { + const int row = row0 + msub * MMA_M + group; + const uint32_t r0 = *reinterpret_cast(a[row] + byte0); + const uint32_t r8 = *reinterpret_cast(a[row + 8] + byte0); + nibbles_to_e4m3(r0, a_frag[msub][0], a_frag[msub][2]); + nibbles_to_e4m3(r8, a_frag[msub][1], a_frag[msub][3]); + } + for (int nsub = 0; nsub < SUB_N; ++nsub) { + const uint32_t w = + *reinterpret_cast(b[col0 + nsub * MMA_N + group] + byte0); + nibbles_to_e4m3(w, b_frag[nsub][0], b_frag[nsub][1]); + } + for (int msub = 0; msub < SUB_M; ++msub) { + for (int nsub = 0; nsub < SUB_N; ++nsub) { + mma_e4m3_m16n8k32(c[msub][nsub], a_frag[msub], b_frag[nsub]); + } + } + } + } + + // Lane (group, quad) holds columns quad*2, +1 of rows group and group+8. + template + __device__ __forceinline__ void store(raw_t* dst, int row0, int col0) const + { + static_assert(Stride % 2 == 0, "float2 stores need even row strides"); + const int lane = threadIdx.x % raft::warp_size(); + const int group = lane / 4; + const int quad = lane % 4; + for (int msub = 0; msub < SUB_M; ++msub) { + for (int nsub = 0; nsub < SUB_N; ++nsub) { + for (int h = 0; h < 2; ++h) { + const int row = row0 + msub * MMA_M + group + h * 8; + const int col = col0 + nsub * MMA_N + quad * 2; + *reinterpret_cast(dst + row * Stride + col) = + make_float2(c[msub][nsub][h * 2], c[msub][nsub][h * 2 + 1]); + } + } + } + } +}; +#endif + +} // namespace cuvs::neighbors::nn_descent::detail diff --git a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh index 4d426d5454..7850a050c7 100644 --- a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh +++ b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh @@ -466,13 +466,14 @@ get_dequant_factors(const quantizer_view& quantizer, int64_t row) } // Converts one raw BBQ dot product into a final (post-epilogue) float distance, given both -// operands' precomputed dequant factors. Evaluated exactly once per matrix cell +// operands' precomputed dequant factors. Evaluated exactly once per matrix cell. `raw` is float so +// float accumulators (fp8 MMA) pass straight through; integer-accumulating callers cast. // dim/centroid_norm_sq/row_norm both come directly from quantizer_document/quantizer_query rather // than being passed separately, since every caller reads them the same id-indexed way -- row_norm // is only read for CosineExpanded (skipped entirely otherwise). template __device__ __forceinline__ float bbq_calculate_metric( - uint32_t raw, + float raw, const bbq_dequant_factors& doc_factors, const bbq_dequant_factors& query_factors, const quantizer_view& quantizer_document, @@ -485,10 +486,9 @@ __device__ __forceinline__ float bbq_calculate_metric( constexpr bool can_postprocess_dist = std::is_same_v; const float dim = static_cast(quantizer_document.dim()); - const float centered = dim * doc_factors.lower * query_factors.lower + - query_factors.lower * doc_factors.sum_delta + - doc_factors.lower * query_factors.sum_delta + - doc_factors.delta * query_factors.delta * static_cast(raw); + const float centered = + dim * doc_factors.lower * query_factors.lower + query_factors.lower * doc_factors.sum_delta + + doc_factors.lower * query_factors.sum_delta + doc_factors.delta * query_factors.delta * raw; const float corrections = doc_factors.corrections + query_factors.corrections; float d; if (metric == cuvs::distance::DistanceType::L2Expanded || From 77241477417aae52d2a8b611b7805f5e862aa277 Mon Sep 17 00:00:00 2001 From: sherylll Date: Mon, 12 Oct 2026 21:12:02 -0700 Subject: [PATCH 2/5] Add arch guard --- cpp/src/neighbors/detail/nn_descent.cuh | 29 ++++++++++++------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 64f7a7919e..a472a9c9ab 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1683,17 +1683,6 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) } } -// Selects fp8 (e4m3) mma.sync instead of u8 wmma for the BBQ local join below. fp8 mma.sync -// needs sm_89+, so lower archs keep the wmma path even when this is set. -#ifndef CUVS_NND_BBQ_FP8_MMA -#define CUVS_NND_BBQ_FP8_MMA 0 -#endif -#if CUVS_NND_BBQ_FP8_MMA && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) -#define CUVS_NND_BBQ_USE_FP8_MMA 1 -#else -#define CUVS_NND_BBQ_USE_FP8_MMA 0 -#endif - // int4 tensor-core BBQ local join, covering both the symmetric (one quantizer, self-join) and // asymmetric (two quantizers) cases. Modeled on local_join_kernel_wmma: a warp-level i4 MMA (see // promoted_i4_warp_mma.cuh) replaces the popc/dp4a inner product; the accumulator lives in @@ -1705,7 +1694,16 @@ template , typename DistEpilogue_t> -RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) +RAFT_KERNEL +#ifdef __CUDA_ARCH__ +// Same minBlocksPerMultiprocessor = 4 arches as local_join_kernel_wmma +#if (__CUDA_ARCH__) == 700 || (__CUDA_ARCH__) == 800 || (__CUDA_ARCH__) == 900 || \ + (__CUDA_ARCH__) == 1000 +__launch_bounds__(BLOCK_SIZE, 4) +#else +__launch_bounds__(BLOCK_SIZE) +#endif +#endif local_join_kernel_bbq_wmma(const Index_t* graph_new, const Index_t* rev_graph_new, const int2* sizes_new, @@ -1736,7 +1734,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) constexpr int MMA_PAD = 16; // Strides are in bytes, though the data is packed 4b. constexpr int ROW_STRIDE = BBQ_ROW_BYTES + MMA_PAD; -#if CUVS_NND_BBQ_USE_FP8_MMA +#if __CUDA_ARCH__ >= 1030 using warp_mma = i4_warp_mma_fp8; #else using warp_mma = i4_warp_mma_u8; @@ -1848,11 +1846,12 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) // col = i % MAX_NUM_BI_SAMPLES is invariant across a thread's own iterations // (BLOCK_SIZE is a multiple of MAX_NUM_BI_SAMPLES), so this thread's column factors are - // fetched once and reused below + // fetched once and reused below. In a self-join they are the staged document factors. const int my_col = tx % MAX_NUM_BI_SAMPLES; bbq_dequant_factors my_col_factors{}; if (my_col < col_size) { - my_col_factors = get_dequant_factors(dataset_query, col_neighbors[my_col]); + my_col_factors = alias_col ? s_document_factors[my_col] + : get_dequant_factors(dataset_query, col_neighbors[my_col]); } // Cell (row, col) is read (as raw_t, via raw_view) and written (as the final float) at the From b1e8f455da626f667790a7f7e6c63b4f71367566 Mon Sep 17 00:00:00 2001 From: sherylll Date: Tue, 13 Oct 2026 00:32:17 -0700 Subject: [PATCH 3/5] Remove unrolled code_inner_product variants --- .../neighbors/detail/cagra/graph_shared.cu | 9 +- cpp/src/neighbors/detail/nn_descent.cuh | 140 +++---- .../quantize/detail/bbq_distance.cuh | 366 ++++-------------- 3 files changed, 135 insertions(+), 380 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index d65711fc21..48ff2960c9 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -176,9 +176,12 @@ __device__ __forceinline__ float bbq_row_distance( int64_t row_document, int64_t row_query) { - namespace bbq = cuvs::preprocessing::quantize::bbq; - const uint32_t raw = - bbq::code_inner_product(quantizer_document, quantizer_query, row_document, row_query); + namespace bbq = cuvs::preprocessing::quantize::bbq; + const uint32_t raw = bbq::code_inner_product(&quantizer_document.codes(row_document, 0), + quantizer_document.layout, + &quantizer_query.codes(row_query, 0), + quantizer_query.layout, + bbq::get_encoded_row_length(quantizer_document)); return bbq::bbq_calculate_metric(static_cast(raw), bbq::get_dequant_factors(quantizer_document, row_document), bbq::get_dequant_factors(quantizer_query, row_query), diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index a472a9c9ab..0358ada2ea 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1326,7 +1326,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) constexpr int QUERY_ROW_BYTES = 128; constexpr int BBQ_PAD = alignof(uint32_t); - // The document buffer is normally only the A operand (two rows broadcast across a warp), so it + // The document buffer is normally only the A operand (one row broadcast across a warp), so it // needs no bank-conflict pad. Under SelfJoin it doubles as the B operand in phase 1 (32 // consecutive columns at one byte offset), so it needs the same skew the query buffer gets. constexpr int DOC_PAD = SelfJoin ? BBQ_PAD : 0; @@ -1343,7 +1343,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) // Document-side (row axis) dequant factors, indexed by list position -- shared by both phases // below since rows are always the `new_neighbors`/document side in both, so this is staged once // and never restaged. At a given accumulation step, up to MAX_NUM_BI_SAMPLES threads share the - // same row0 (see the pair_idx assignment below), so caching avoids up to 64x redundant scattered + // same row (see the pair_idx assignment below), so caching avoids up to 64x redundant scattered // global reads __shared__ bbq_dequant_factors s_document_factors[MAX_NUM_BI_SAMPLES]; @@ -1418,9 +1418,8 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) constexpr int query_plane_tile_u32 = query_plane_tile / 4; const int n_tiles = raft::ceildiv(plane_bytes, plane_tile); - // with NUM_SAMPLES=32, BLOCK_SIZE=256, pairs_per_thread = 32 * 64 / 512 = 4 - constexpr int num_row_pairs = MAX_NUM_BI_SAMPLES / 2; - constexpr int num_pairs = num_row_pairs * MAX_NUM_BI_SAMPLES; + // One (row, col) cell per pair: pairs_per_thread = 64 * 64 / 512 = 8. + constexpr int num_pairs = MAX_NUM_BI_SAMPLES * MAX_NUM_BI_SAMPLES; constexpr int pairs_per_thread = (num_pairs + BLOCK_SIZE - 1) / BLOCK_SIZE; // Every thread's column is invariant across its pairs_per_thread iterations, so fetch it once @@ -1431,8 +1430,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) my_col_factors = get_dequant_factors(dataset_query, new_neighbors[my_col]); } - uint32_t acc0[pairs_per_thread] = {}; - uint32_t acc1[pairs_per_thread] = {}; + uint32_t acc[pairs_per_thread] = {}; for (int step = 0; step < n_tiles; ++step) { const bool last_tile = (step == n_tiles - 1); const int num_load = last_tile ? plane_bytes - step * plane_tile : plane_tile; @@ -1474,9 +1472,9 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) for (int k = 0; k < pairs_per_thread; ++k) { const int pair_idx = tx + k * BLOCK_SIZE; if (pair_idx >= num_pairs) continue; - const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; - const int col = pair_idx % MAX_NUM_BI_SAMPLES; - if (col >= new_size) continue; + const int row = pair_idx / MAX_NUM_BI_SAMPLES; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (row >= new_size || col >= new_size) continue; // Phase 1 B operand: s_doc_vec under SelfJoin (see the staging note above). const uint8_t* row_b; if constexpr (SelfJoin) { @@ -1484,13 +1482,8 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) } else { row_b = s_query_vec[col]; } - cuvs::preprocessing::quantize::bbq::bbq_code_inner_product_2x1( - s_doc_vec[row0], s_doc_vec[row0 + 1], row_b, acc0[k], acc1[k]); + acc[k] += cuvs::preprocessing::quantize::bbq::code_inner_product( + s_doc_vec[row], DocumentLayout, row_b, QueryLayout, doc_row_bytes); } __syncthreads(); } @@ -1498,40 +1491,22 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) for (int k = 0; k < pairs_per_thread; ++k) { const int pair_idx = tx + k * BLOCK_SIZE; if (pair_idx >= num_pairs) continue; - const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; - const int col = pair_idx % MAX_NUM_BI_SAMPLES; - if (col >= new_size) continue; - const int distance0 = row0 * SKEWED_MAX_NUM_BI_SAMPLES + col; - const Index_t query_id = new_neighbors[col]; - // row0 (unlike row0 + 1 below) isn't otherwise bounded by new_size -- gate it explicitly so a - // stale/garbage new_neighbors[row0] past the real list never reaches a dataset_document lookup - // (row_norm, for CosineExpanded); the cell is provably never read downstream either way, since - // the min-search loops below only visit idx_in_list < new_size. - if (row0 < new_size) { - const Index_t doc_id0 = new_neighbors[row0]; - s_distances[distance0] = bbq_calculate_metric(static_cast(acc0[k]), - s_document_factors[row0], - my_col_factors, - dataset_document, - dataset_query, - metric, - dist_epilogue, - doc_id0, - query_id); - } - if (row0 + 1 < new_size) { - const Index_t doc_id1 = new_neighbors[row0 + 1]; - s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = - bbq_calculate_metric(static_cast(acc1[k]), - s_document_factors[row0 + 1], - my_col_factors, - dataset_document, - dataset_query, - metric, - dist_epilogue, - doc_id1, - query_id); - } + const int row = pair_idx / MAX_NUM_BI_SAMPLES; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + // Gating row keeps a stale new_neighbors[row] past the real list from reaching a + // dataset_document lookup (row_norm, for CosineExpanded); the min-search loops below only + // visit idx_in_list < new_size, so the cell is never read either way. + if (row >= new_size || col >= new_size) continue; + s_distances[row * SKEWED_MAX_NUM_BI_SAMPLES + col] = + bbq_calculate_metric(static_cast(acc[k]), + s_document_factors[row], + my_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + new_neighbors[row], + new_neighbors[col]); } __syncthreads(); @@ -1562,8 +1537,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) my_old_col_factors = get_dequant_factors(dataset_query, old_neighbors[my_old_col]); } - uint32_t acc0_old[pairs_per_thread] = {}; - uint32_t acc1_old[pairs_per_thread] = {}; + uint32_t acc_old[pairs_per_thread] = {}; for (int step = 0; step < n_tiles; ++step) { const bool last_tile = (step == n_tiles - 1); const int num_load = last_tile ? plane_bytes - step * plane_tile : plane_tile; @@ -1599,16 +1573,11 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) for (int k = 0; k < pairs_per_thread; ++k) { const int pair_idx = tx + k * BLOCK_SIZE; if (pair_idx >= num_pairs) continue; - const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; - const int col = pair_idx % MAX_NUM_BI_SAMPLES; - if (col >= old_size) continue; - cuvs::preprocessing::quantize::bbq::bbq_code_inner_product_2x1( - s_doc_vec[row0], s_doc_vec[row0 + 1], s_query_vec[col], acc0_old[k], acc1_old[k]); + const int row = pair_idx / MAX_NUM_BI_SAMPLES; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (row >= new_size || col >= old_size) continue; + acc_old[k] += cuvs::preprocessing::quantize::bbq::code_inner_product( + s_doc_vec[row], DocumentLayout, s_query_vec[col], QueryLayout, doc_row_bytes); } __syncthreads(); } @@ -1616,38 +1585,19 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) for (int k = 0; k < pairs_per_thread; ++k) { const int pair_idx = tx + k * BLOCK_SIZE; if (pair_idx >= num_pairs) continue; - const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; - const int col = pair_idx % MAX_NUM_BI_SAMPLES; - if (col >= old_size) continue; - const int distance0 = row0 * SKEWED_MAX_NUM_BI_SAMPLES + col; - const Index_t query_id = old_neighbors[col]; - // See the identical row0 < new_size gate in phase 1 above -- same reasoning (row0 isn't - // otherwise bounded, and this cell is never read downstream either way). - if (row0 < new_size) { - const Index_t doc_id0 = new_neighbors[row0]; - s_distances[distance0] = bbq_calculate_metric(static_cast(acc0_old[k]), - s_document_factors[row0], - my_old_col_factors, - dataset_document, - dataset_query, - metric, - dist_epilogue, - doc_id0, - query_id); - } - if (row0 + 1 < new_size) { - const Index_t doc_id1 = new_neighbors[row0 + 1]; - s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = - bbq_calculate_metric(static_cast(acc1_old[k]), - s_document_factors[row0 + 1], - my_old_col_factors, - dataset_document, - dataset_query, - metric, - dist_epilogue, - doc_id1, - query_id); - } + const int row = pair_idx / MAX_NUM_BI_SAMPLES; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (row >= new_size || col >= old_size) continue; + s_distances[row * SKEWED_MAX_NUM_BI_SAMPLES + col] = + bbq_calculate_metric(static_cast(acc_old[k]), + s_document_factors[row], + my_old_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + new_neighbors[row], + old_neighbors[col]); } __syncthreads(); diff --git a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh index 7850a050c7..bc7ad1d134 100644 --- a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh +++ b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh @@ -50,41 +50,69 @@ _RAFT_HOST_DEVICE constexpr int get_code_planes(const bbq_code_layout layout) } // -------------------------------------------------------------------------- -// Code inner products (1x1) -// Raw byte-level dot products over one row pair, per code layout. +// Code promotion +// Widens a narrower layout to 4-bit width, so an asymmetric pair can meet in one format. +// -------------------------------------------------------------------------- + +// Promotes one native word of dense packed_1b codes (1 bit/value, 8 values/byte, MSB-first: +// the value at position 8*byte+i sits at bit (7-i)) into four 4-bit-width, packed_4b-style +// output words, packed into one uint4 for a single vectorized store. +// +// Branch-free SWAR: extract the 4 (2-bit) fields of each byte lane-wise across all 4 native +// bytes at once, spread each field 0-3 into a nibble value lane-wise, then transpose the 4 +// resulting field-words into the 4 per-native-byte output words with chained __byte_perm pairs. +__device__ __forceinline__ uint4 packed_1b_to_4b(uint32_t native_word) +{ + const uint32_t spread0 = ((native_word >> 3) & 0x10101010u) | ((native_word >> 6) & 0x01010101u); + const uint32_t spread1 = ((native_word >> 1) & 0x10101010u) | ((native_word >> 4) & 0x01010101u); + const uint32_t spread2 = ((native_word << 1) & 0x10101010u) | ((native_word >> 2) & 0x01010101u); + const uint32_t spread3 = ((native_word << 3) & 0x10101010u) | (native_word & 0x01010101u); + const uint32_t t0 = __byte_perm(spread0, spread1, 0x5140); + const uint32_t t1 = __byte_perm(spread0, spread1, 0x7362); + const uint32_t t2 = __byte_perm(spread2, spread3, 0x5140); + const uint32_t t3 = __byte_perm(spread2, spread3, 0x7362); + return uint4{__byte_perm(t0, t2, 0x5410), + __byte_perm(t0, t2, 0x7632), + __byte_perm(t1, t3, 0x5410), + __byte_perm(t1, t3, 0x7632)}; +} + +// -------------------------------------------------------------------------- +// Code inner products +// Raw byte-level dot products over one document row and one query row. // -------------------------------------------------------------------------- /** - * Cross-plane binary inner product over `Planes` bit planes, shifting each (i, j) plane pair by - * i + j. `Planes == 1` is a plain binary product, so this covers packed_1b as well as the - * transposed layouts. The 4-byte body needs both rows 4-byte aligned; the byte tail handles a - * stripe whose length is not a multiple of 4. + * Cross-plane binary inner product over two bit-sliced rows whose plane counts need not match, + * shifting each (i, j) plane pair by i + j. Equal counts are the symmetric case, and a single + * plane is a plain binary product, so this covers every bit-sliced pair including packed_1b. + * + * All three bit-sliced layouts store ceildiv(dim, 8) bytes per plane, which is what lets one + * @p plane_bytes describe both operands. */ -template -__device__ __forceinline__ uint32_t code_inner_product_transposed(const uint8_t* row_a, - const uint8_t* row_b, - size_t n_bytes, - uint32_t result = 0) +__device__ __forceinline__ uint32_t code_inner_product_planes(const uint8_t* row_document, + int document_planes, + const uint8_t* row_query, + int query_planes, + size_t plane_bytes, + uint32_t result = 0) { - const size_t stripe = n_bytes / Planes; - // Plane p starts at row + p * stripe and is read as uint32_t words, so a stripe that is not a - // multiple of 4 misaligns every plane past the first -- the byte tail below only covers a ragged - // stripe *length*, not a ragged stripe *offset*. Planes == 1 has no offset and so is exempt. - assert(Planes == 1 || stripe % sizeof(uint32_t) == 0); + // Planes past the first start at a multiple of plane_bytes and are read as uint32_t words. + assert((document_planes == 1 && query_planes == 1) || plane_bytes % sizeof(uint32_t) == 0); #pragma unroll - for (int i = 0; i < Planes; ++i) { + for (int i = 0; i < document_planes; ++i) { #pragma unroll - for (int j = 0; j < Planes; ++j) { - const uint8_t* a = row_a + i * stripe; - const uint8_t* b = row_b + j * stripe; + for (int j = 0; j < query_planes; ++j) { + const uint8_t* a = row_document + i * plane_bytes; + const uint8_t* b = row_query + j * plane_bytes; uint32_t partial = 0; size_t k = 0; #pragma unroll 4 - for (; k + sizeof(uint32_t) <= stripe; k += sizeof(uint32_t)) { + for (; k + sizeof(uint32_t) <= plane_bytes; k += sizeof(uint32_t)) { partial += __popc(*reinterpret_cast(a + k) & *reinterpret_cast(b + k)); } - for (; k < stripe; ++k) { + for (; k < plane_bytes; ++k) { partial += __popc(static_cast(a[k] & b[k])); } result += partial << (i + j); @@ -145,111 +173,6 @@ __device__ __forceinline__ uint32_t code_inner_product_packed_8b(const uint8_t* return result; } -/** - * Integer inner product between two encoded rows. - * - * The uint32_t result bounds every BBQ layout to 66,050 dimensions: the worst case is - * `packed_8b`, where `dim * 255 * 255` must not exceed UINT32_MAX. - */ -__device__ __forceinline__ uint32_t code_inner_product(const uint8_t* row_a, - const uint8_t* row_b, - const bbq_code_layout layout, - const size_t n_bytes, - uint32_t result = 0) -{ - switch (layout) { - case bbq_code_layout::packed_1b: - return code_inner_product_transposed<1>(row_a, row_b, n_bytes, result); - case bbq_code_layout::transposed_2b: - return code_inner_product_transposed<2>(row_a, row_b, n_bytes, result); - case bbq_code_layout::packed_4b: - return code_inner_product_packed_4b(row_a, row_b, n_bytes, result); - case bbq_code_layout::transposed_4b: - return code_inner_product_transposed<4>(row_a, row_b, n_bytes, result); - case bbq_code_layout::packed_8b: - return code_inner_product_packed_8b(row_a, row_b, n_bytes, result); - case bbq_code_layout::packed_7b: - default: - return code_inner_product_packed_8b( - row_a, - row_b, - n_bytes, - result, - static_cast((uint32_t{1} << get_bit_width(layout)) - 1)); - } -} - -// -------------------------------------------------------------------------- -// Code promotion -// Widens a narrower layout to 4-bit width, so an asymmetric pair can meet in one format. -// -------------------------------------------------------------------------- - -// Promotes one native word of dense packed_1b codes (1 bit/value, 8 values/byte, MSB-first: -// the value at position 8*byte+i sits at bit (7-i)) into four 4-bit-width, packed_4b-style -// output words, packed into one uint4 for a single vectorized store. -// -// Branch-free SWAR: extract the 4 (2-bit) fields of each byte lane-wise across all 4 native -// bytes at once, spread each field 0-3 into a nibble value lane-wise, then transpose the 4 -// resulting field-words into the 4 per-native-byte output words with chained __byte_perm pairs. -__device__ __forceinline__ uint4 packed_1b_to_4b(uint32_t native_word) -{ - const uint32_t spread0 = ((native_word >> 3) & 0x10101010u) | ((native_word >> 6) & 0x01010101u); - const uint32_t spread1 = ((native_word >> 1) & 0x10101010u) | ((native_word >> 4) & 0x01010101u); - const uint32_t spread2 = ((native_word << 1) & 0x10101010u) | ((native_word >> 2) & 0x01010101u); - const uint32_t spread3 = ((native_word << 3) & 0x10101010u) | (native_word & 0x01010101u); - const uint32_t t0 = __byte_perm(spread0, spread1, 0x5140); - const uint32_t t1 = __byte_perm(spread0, spread1, 0x7362); - const uint32_t t2 = __byte_perm(spread2, spread3, 0x5140); - const uint32_t t3 = __byte_perm(spread2, spread3, 0x7362); - return uint4{__byte_perm(t0, t2, 0x5410), - __byte_perm(t0, t2, 0x7632), - __byte_perm(t1, t3, 0x5410), - __byte_perm(t1, t3, 0x7632)}; -} - -// -------------------------------------------------------------------------- -// Cross-layout inner products (1x1) -// One document row against one query row encoded in a different layout. -// -------------------------------------------------------------------------- - -/** - * Cross-plane binary inner product over two bit-sliced rows whose plane counts need not match, - * shifting each (i, j) plane pair by i + j. The runtime-width counterpart of - * code_inner_product_transposed, which the local-join kernels reach with both counts known at - * compile time; a single plane is a plain binary product, so this covers packed_1b too. - * - * All three bit-sliced layouts store ceildiv(dim, 8) bytes per plane, which is what lets one - * @p plane_bytes describe both operands. - */ -__device__ __forceinline__ uint32_t code_inner_product_planes(const uint8_t* row_document, - int document_planes, - const uint8_t* row_query, - int query_planes, - size_t plane_bytes, - uint32_t result = 0) -{ - // Planes past the first start at a multiple of plane_bytes and are read as uint32_t words. - assert((document_planes == 1 && query_planes == 1) || plane_bytes % sizeof(uint32_t) == 0); - for (int i = 0; i < document_planes; ++i) { - for (int j = 0; j < query_planes; ++j) { - const uint8_t* a = row_document + i * plane_bytes; - const uint8_t* b = row_query + j * plane_bytes; - uint32_t partial = 0; - size_t k = 0; -#pragma unroll 4 - for (; k + sizeof(uint32_t) <= plane_bytes; k += sizeof(uint32_t)) { - partial += __popc(*reinterpret_cast(a + k) & - *reinterpret_cast(b + k)); - } - for (; k < plane_bytes; ++k) { - partial += __popc(static_cast(a[k] & b[k])); - } - result += partial << (i + j); - } - } - return result; -} - /** * packed_1b document against a packed_4b query. The document is promoted to 4-bit width one * native word at a time -- 32 dimensions, which is exactly the four query words covering the same @@ -279,167 +202,46 @@ __device__ __forceinline__ uint32_t code_inner_product_1b_x_packed_4b(const uint } /** - * Integer inner product between a document row and a query row, each read through its own - * quantizer. The two layouts may differ: the supported (document, query) pairs are the ones - * nn-descent's local join dispatches on, and passing the same view twice is the symmetric case. - */ -template -__device__ __forceinline__ uint32_t -code_inner_product(const quantizer_view& quantizer_document, - const quantizer_view& quantizer_query, - int64_t row_document, - int64_t row_query) -{ - const uint8_t* document = &quantizer_document.codes(row_document, 0); - const uint8_t* query = &quantizer_query.codes(row_query, 0); - - if (quantizer_document.layout == quantizer_query.layout) { - return code_inner_product( - document, query, quantizer_document.layout, get_encoded_row_length(quantizer_document)); - } - if (quantizer_document.layout == bbq_code_layout::packed_1b && - quantizer_query.layout == bbq_code_layout::packed_4b) { - return code_inner_product_1b_x_packed_4b( - document, query, get_encoded_row_length(quantizer_document)); - } - // Every remaining supported pair is bit-sliced on both sides: (packed_1b, transposed_2b), - // (packed_1b, transposed_4b) and (transposed_2b, transposed_4b). - assert(quantizer_document.layout == bbq_code_layout::packed_1b || - quantizer_document.layout == bbq_code_layout::transposed_2b); - assert(quantizer_query.layout == bbq_code_layout::transposed_2b || - quantizer_query.layout == bbq_code_layout::transposed_4b); - return code_inner_product_planes(document, - get_code_planes(quantizer_document.layout), - query, - get_code_planes(quantizer_query.layout), - (quantizer_document.dim() + 7) / 8); -} - -// -------------------------------------------------------------------------- -// Fused inner products (2x1) -// Two left rows against a shared right operand, for the local-join inner loop. -// -------------------------------------------------------------------------- - -/** - * Two cross-plane inner products over statically known document and query plane counts. - * Not asymmetric-specific: a symmetric pair is just document_planes == query_planes, which is how - * the 1x1 and 2t x 2t self-joins are computed. + * Integer inner product between a document row and a query row, each in its own layout. The query + * layout picks the form, since each supported pair is fixed by it: a packed_4b query meets a + * packed_4b document or a promoted packed_1b one, packed_7b/8b only themselves, and every other + * query is bit-sliced against a bit-sliced document. The local join calls this with constant + * layouts and row width, so once inlined the dispatch and the byte tails fold away. + * + * The uint32_t result bounds every BBQ layout to 66,050 dimensions: the worst case is + * `packed_8b`, where `dim * 255 * 255` must not exceed UINT32_MAX. */ -template -__device__ inline void code_inner_product_planes_2x1(const uint8_t* row_a0, - const uint8_t* row_a1, - const uint8_t* row_b, - uint32_t& total0, - uint32_t& total1) +__device__ __forceinline__ uint32_t code_inner_product(const uint8_t* row_document, + const bbq_code_layout document_layout, + const uint8_t* row_query, + const bbq_code_layout query_layout, + const size_t document_row_bytes) { - constexpr size_t document_plane_stride = document_row_bytes / document_planes; - constexpr size_t query_plane_stride = query_row_bytes / query_planes; - // Both operands are stepped by their own plane stride and then read as uint32_t, so both - // strides -- not just the query's -- must be 4-byte aligned. - static_assert(query_plane_stride % sizeof(uint32_t) == 0); - static_assert(document_plane_stride % sizeof(uint32_t) == 0); -#pragma unroll - for (int p_query = 0; p_query < query_planes; ++p_query) { -#pragma unroll - for (int p_document = 0; p_document < document_planes; ++p_document) { - const uint8_t* a0 = row_a0 + p_document * document_plane_stride; - const uint8_t* a1 = row_a1 + p_document * document_plane_stride; - const uint8_t* b = row_b + p_query * query_plane_stride; - uint32_t partial0 = 0; - uint32_t partial1 = 0; -#pragma unroll 4 - for (size_t i = 0; i < query_plane_stride; i += sizeof(uint32_t)) { - const auto wa0 = *reinterpret_cast(a0 + i); - const auto wa1 = *reinterpret_cast(a1 + i); - const auto wb = *reinterpret_cast(b + i); - partial0 += __popc(wa0 & wb); - partial1 += __popc(wa1 & wb); - } - total0 += partial0 << (p_document + p_query); - total1 += partial1 << (p_document + p_query); + switch (query_layout) { + case bbq_code_layout::packed_4b: + return document_layout == bbq_code_layout::packed_1b + ? code_inner_product_1b_x_packed_4b(row_document, row_query, document_row_bytes) + : code_inner_product_packed_4b(row_document, row_query, document_row_bytes); + case bbq_code_layout::packed_7b: + case bbq_code_layout::packed_8b: + assert(document_layout == query_layout); + return code_inner_product_packed_8b( + row_document, + row_query, + document_row_bytes, + 0, + static_cast((uint32_t{1} << get_bit_width(query_layout)) - 1)); + default: { + const int document_planes = get_code_planes(document_layout); + return code_inner_product_planes(row_document, + document_planes, + row_query, + get_code_planes(query_layout), + document_row_bytes / document_planes); } } } -template -__device__ inline void code_inner_product_packed_4b_symmetric_2x1(const uint8_t* row_a0, - const uint8_t* row_a1, - const uint8_t* row_b, - uint32_t& total0, - uint32_t& total1) -{ - static_assert(n_bytes % sizeof(uint32_t) == 0); - constexpr uint32_t nibble_mask = 0x0F0F0F0Fu; -#pragma unroll 4 - for (size_t i = 0; i < n_bytes; i += sizeof(uint32_t)) { - const auto a0 = *reinterpret_cast(row_a0 + i); - const auto a1 = *reinterpret_cast(row_a1 + i); - const auto b = *reinterpret_cast(row_b + i); - const auto b_low = b & nibble_mask; - const auto b_high = (b >> 4) & nibble_mask; - total0 = __dp4a(a0 & nibble_mask, b_low, total0); - total0 = __dp4a((a0 >> 4) & nibble_mask, b_high, total0); - total1 = __dp4a(a1 & nibble_mask, b_low, total1); - total1 = __dp4a((a1 >> 4) & nibble_mask, b_high, total1); - } -} - -template -__device__ inline void code_inner_product_packed_8b_2x1(const uint8_t* row_a0, - const uint8_t* row_a1, - const uint8_t* row_b, - uint32_t& total0, - uint32_t& total1, - uint8_t code_mask = 0xFFu) -{ - static_assert(n_bytes % sizeof(uint32_t) == 0); - const uint32_t word_mask = uint32_t{code_mask} * 0x01010101u; -#pragma unroll 4 - for (size_t i = 0; i < n_bytes; i += sizeof(uint32_t)) { - const auto a0 = *reinterpret_cast(row_a0 + i) & word_mask; - const auto a1 = *reinterpret_cast(row_a1 + i) & word_mask; - const auto b = *reinterpret_cast(row_b + i) & word_mask; - total0 = __dp4a(a0, b, total0); - total1 = __dp4a(a1, b, total1); - } -} - -// Selects the SIMT inner product for a (document, query) layout pair: bit-sliced layouts go to the -// cross-plane popc, densely-packed ones to dp4a. dp4a needs both operands in the same packing, so -// those forms apply whenever the two layouts match. packed_7b and packed_8b reach this from -// GNND::local_join; packed_4b does not (symmetric packed_4b goes to the wmma kernel), but is kept -// as a SIMT reference point. -template -__device__ __forceinline__ void bbq_code_inner_product_2x1(const uint8_t* row_a0, - const uint8_t* row_a1, - const uint8_t* row_b, - uint32_t& total0, - uint32_t& total1) -{ - namespace bbq = cuvs::preprocessing::quantize::bbq; - if constexpr (DocumentLayout == QueryLayout && DocumentLayout == bbq_code_layout::packed_4b) { - bbq::code_inner_product_packed_4b_symmetric_2x1( - row_a0, row_a1, row_b, total0, total1); - } else if constexpr (DocumentLayout == QueryLayout && - (DocumentLayout == bbq_code_layout::packed_8b || - DocumentLayout == bbq_code_layout::packed_7b)) { - // packed_7b is packed_8b with the top bit masked off, matching code_inner_product's - // (1 << bits) - 1 mask for the same two layouts. - constexpr uint8_t code_mask = DocumentLayout == bbq_code_layout::packed_7b ? 0x7Fu : 0xFFu; - bbq::code_inner_product_packed_8b_2x1( - row_a0, row_a1, row_b, total0, total1, code_mask); - } else { - bbq:: - code_inner_product_planes_2x1( - row_a0, row_a1, row_b, total0, total1); - } -} - // -------------------------------------------------------------------------- // Dequantization // Turns a raw code inner product into a final float distance. From 1bcf72aeb293dd46e4adeab28369ebbbaa924b28 Mon Sep 17 00:00:00 2001 From: sherylll Date: Wed, 14 Oct 2026 02:53:03 -0700 Subject: [PATCH 4/5] Use a scaling trick for int4->fp8 conversion --- .../neighbors/detail/promoted_i4_warp_mma.cuh | 71 +++++++++---------- 1 file changed, 35 insertions(+), 36 deletions(-) diff --git a/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh b/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh index c4ce31ad50..cf6c1b823f 100644 --- a/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh +++ b/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh @@ -5,7 +5,7 @@ #pragma once -#include +#include #include @@ -129,30 +129,23 @@ struct i4_warp_mma_u8 { // fp8 mma.sync needs sm_89+. #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) -// Expands the eight 4-bit codes packed in w to e4m3 bytes of the same value (0..15 are exact in -// e4m3): lo gets the codes in w's low 16 bits, hi the codes in its high 16 bits. prmt tables hold -// only 8 bytes, so 0..7 and 8..15 use separate tables; a selector with bit 3 set yields the -// selected byte's sign, 0 here, which zeroes the lanes the other table owns. prmt reads only the -// selector's low 16 bits, hence the shift for hi. -__device__ __forceinline__ void nibbles_to_e4m3(uint32_t w, uint32_t& lo, uint32_t& hi) -{ - constexpr uint32_t k0_3 = 0x44403800u; - constexpr uint32_t k4_7 = 0x4e4c4a48u; - constexpr uint32_t k8_11 = 0x53525150u; - constexpr uint32_t k12_15 = 0x57565554u; - auto lookup = [](uint32_t sel) { - uint32_t below_8, from_8; - asm("prmt.b32 %0, %1, %2, %3;" : "=r"(below_8) : "r"(k0_3), "r"(k4_7), "r"(sel)); - asm("prmt.b32 %0, %1, %2, %3;" : "=r"(from_8) : "r"(k8_11), "r"(k12_15), "r"(sel ^ 0x8888u)); - return below_8 | from_8; - }; - lo = lookup(w); - hi = lookup(w >> 16); -} +// A 4-bit code c (0..15) masked into its own byte and read as e4m3 is c / kE4m3CodeScale: +// c byte e4m3 value +// 0 0x00 0 +// 1 0x01 1 / 512 +// ... +// 15 0x0F 15 / 512 +// so splitting a packed word into its low and high nibbles, one code per byte, is the whole +// expansion to e4m3, and every product carries 1 / kE4m3CodeScale^2, removed in store(). +// Codes 1..7 are e4m3 subnormals (exponent bits 0000). Hardware that flushed subnormal MMA inputs +// to zero would silently drop them; we accept that risk because e4m3 defines these values as part +// of its (already small) range and the fp8 MMA paths honor them. nvcc's -ftz does not apply: it +// covers f32 arithmetic, and every f32 value here is a normal number. +constexpr int kE4m3CodeScale = 512; // c += a * b over one m16n8k32 tile. The tensor core's fp8 accumulator keeps far fewer bits than -// f32, so chaining c through it drops low bits of the integer sums. Each MMA starts from zero -// instead -- one tile's sum fits exactly -- and the running sum is kept with ordinary f32 adds. +// f32, so chaining c through it drops low bits of the sums. Each MMA starts from zero instead -- +// one tile's sum fits exactly -- and the running sum is kept with ordinary f32 adds. __device__ __forceinline__ void mma_e4m3_m16n8k32(float (&c)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) @@ -188,7 +181,7 @@ struct i4_warp_mma_fp8 { // m16n8k32 fragment coordinates: `group` is the A row (and +8) and the B column this lane holds; // `quad` picks its 4-byte slice of the k-step. A and B read the same slice, and expanding it to - // (low 16 bits -> k 0..15, high -> k 16..31) the same way keeps their k orders aligned. + // (low nibbles -> k 0..15, high nibbles -> k 16..31) the same way keeps their k orders aligned. // Kept rolled for the same register reason as the u8 path. template __device__ __forceinline__ void accumulate(const uint8_t (*a)[RowStride], @@ -196,9 +189,10 @@ struct i4_warp_mma_fp8 { int row0, int col0) { - const int lane = threadIdx.x % raft::warp_size(); - const int group = lane / 4; - const int quad = lane % 4; + constexpr uint32_t kLowNibbles = 0x0f0f0f0fu; + const int lane = threadIdx.x % raft::warp_size(); + const int group = lane / 4; + const int quad = lane % 4; #pragma unroll 1 for (int kk = 0; kk < RowBytes / KSTEP_BYTES; ++kk) { const int byte0 = kk * KSTEP_BYTES + quad * 4; @@ -208,13 +202,16 @@ struct i4_warp_mma_fp8 { const int row = row0 + msub * MMA_M + group; const uint32_t r0 = *reinterpret_cast(a[row] + byte0); const uint32_t r8 = *reinterpret_cast(a[row + 8] + byte0); - nibbles_to_e4m3(r0, a_frag[msub][0], a_frag[msub][2]); - nibbles_to_e4m3(r8, a_frag[msub][1], a_frag[msub][3]); + a_frag[msub][0] = r0 & kLowNibbles; + a_frag[msub][2] = (r0 >> 4) & kLowNibbles; + a_frag[msub][1] = r8 & kLowNibbles; + a_frag[msub][3] = (r8 >> 4) & kLowNibbles; } for (int nsub = 0; nsub < SUB_N; ++nsub) { const uint32_t w = *reinterpret_cast(b[col0 + nsub * MMA_N + group] + byte0); - nibbles_to_e4m3(w, b_frag[nsub][0], b_frag[nsub][1]); + b_frag[nsub][0] = w & kLowNibbles; + b_frag[nsub][1] = (w >> 4) & kLowNibbles; } for (int msub = 0; msub < SUB_M; ++msub) { for (int nsub = 0; nsub < SUB_N; ++nsub) { @@ -224,21 +221,23 @@ struct i4_warp_mma_fp8 { } } - // Lane (group, quad) holds columns quad*2, +1 of rows group and group+8. + // Lane (group, quad) holds columns quad*2, +1 of rows group and group+8. Runs once, after every + // K chunk is accumulated, so this is where the codes' 1 / kE4m3CodeScale^2 comes off. template __device__ __forceinline__ void store(raw_t* dst, int row0, int col0) const { static_assert(Stride % 2 == 0, "float2 stores need even row strides"); - const int lane = threadIdx.x % raft::warp_size(); - const int group = lane / 4; - const int quad = lane % 4; + constexpr float code_scale_sq = kE4m3CodeScale * kE4m3CodeScale; + const int lane = threadIdx.x % raft::warp_size(); + const int group = lane / 4; + const int quad = lane % 4; for (int msub = 0; msub < SUB_M; ++msub) { for (int nsub = 0; nsub < SUB_N; ++nsub) { for (int h = 0; h < 2; ++h) { const int row = row0 + msub * MMA_M + group + h * 8; const int col = col0 + nsub * MMA_N + quad * 2; - *reinterpret_cast(dst + row * Stride + col) = - make_float2(c[msub][nsub][h * 2], c[msub][nsub][h * 2 + 1]); + *reinterpret_cast(dst + row * Stride + col) = make_float2( + c[msub][nsub][h * 2] * code_scale_sq, c[msub][nsub][h * 2 + 1] * code_scale_sq); } } } From 843ddc9b0eef06b3371af8a6d4c97a58decffe10 Mon Sep 17 00:00:00 2001 From: sherylll Date: Thu, 15 Oct 2026 19:17:34 -0700 Subject: [PATCH 5/5] Drop FP8 wmma variant --- .../neighbors/detail/cagra/graph_shared.cu | 2 +- cpp/src/neighbors/detail/nn_descent.cuh | 12 +- .../neighbors/detail/promoted_i4_warp_mma.cuh | 121 +----------------- .../quantize/detail/bbq_distance.cuh | 12 +- 4 files changed, 12 insertions(+), 135 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index 48ff2960c9..af9251b6b6 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -182,7 +182,7 @@ __device__ __forceinline__ float bbq_row_distance( &quantizer_query.codes(row_query, 0), quantizer_query.layout, bbq::get_encoded_row_length(quantizer_document)); - return bbq::bbq_calculate_metric(static_cast(raw), + return bbq::bbq_calculate_metric(raw, bbq::get_dequant_factors(quantizer_document, row_document), bbq::get_dequant_factors(quantizer_query, row_query), quantizer_document, diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 0358ada2ea..8dd3b4b89c 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1498,7 +1498,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) // visit idx_in_list < new_size, so the cell is never read either way. if (row >= new_size || col >= new_size) continue; s_distances[row * SKEWED_MAX_NUM_BI_SAMPLES + col] = - bbq_calculate_metric(static_cast(acc[k]), + bbq_calculate_metric(acc[k], s_document_factors[row], my_col_factors, dataset_document, @@ -1589,7 +1589,7 @@ RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) const int col = pair_idx % MAX_NUM_BI_SAMPLES; if (row >= new_size || col >= old_size) continue; s_distances[row * SKEWED_MAX_NUM_BI_SAMPLES + col] = - bbq_calculate_metric(static_cast(acc_old[k]), + bbq_calculate_metric(acc_old[k], s_document_factors[row], my_old_col_factors, dataset_document, @@ -1684,11 +1684,7 @@ __launch_bounds__(BLOCK_SIZE) constexpr int MMA_PAD = 16; // Strides are in bytes, though the data is packed 4b. constexpr int ROW_STRIDE = BBQ_ROW_BYTES + MMA_PAD; -#if __CUDA_ARCH__ >= 1030 - using warp_mma = i4_warp_mma_fp8; -#else - using warp_mma = i4_warp_mma_u8; -#endif + using warp_mma = i4_warp_mma_u8; static_assert(BBQ_ROW_BYTES % warp_mma::KSTEP_BYTES == 0, "a tile must be whole k-steps"); constexpr int MMA_STORE_STRIDE = SKEWED_MAX_NUM_BI_SAMPLES; @@ -1814,7 +1810,7 @@ __launch_bounds__(BLOCK_SIZE) const int distance0 = row * MMA_STORE_STRIDE + col; const Index_t doc_id = new_neighbors[row]; const Index_t query_id = col_neighbors[col]; - const float raw = static_cast(raw_view[distance0]); + const auto raw = static_cast(raw_view[distance0]); s_distances[distance0] = bbq_calculate_metric(raw, s_document_factors[row], my_col_factors, diff --git a/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh b/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh index cf6c1b823f..fb2e4f288a 100644 --- a/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh +++ b/cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh @@ -14,8 +14,7 @@ // Warp-level tensor-core dot products of i4 codes (0..15, two per byte) staged in shared memory as // rows of RowStride bytes. A warp owns a WarpTile x WarpTile block of the output: accumulate() adds -// the dot products over one staged K range, store() writes the raw sums as raw_t. The u8 (wmma) -// and fp8 (mma.sync) variants share this interface, so callers can switch between them. +// the dot products over one staged K range, store() writes the raw sums as raw_t. namespace cuvs::neighbors::nn_descent::detail { @@ -127,122 +126,4 @@ struct i4_warp_mma_u8 { } }; -// fp8 mma.sync needs sm_89+. -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) -// A 4-bit code c (0..15) masked into its own byte and read as e4m3 is c / kE4m3CodeScale: -// c byte e4m3 value -// 0 0x00 0 -// 1 0x01 1 / 512 -// ... -// 15 0x0F 15 / 512 -// so splitting a packed word into its low and high nibbles, one code per byte, is the whole -// expansion to e4m3, and every product carries 1 / kE4m3CodeScale^2, removed in store(). -// Codes 1..7 are e4m3 subnormals (exponent bits 0000). Hardware that flushed subnormal MMA inputs -// to zero would silently drop them; we accept that risk because e4m3 defines these values as part -// of its (already small) range and the fp8 MMA paths honor them. nvcc's -ftz does not apply: it -// covers f32 arithmetic, and every f32 value here is a normal number. -constexpr int kE4m3CodeScale = 512; - -// c += a * b over one m16n8k32 tile. The tensor core's fp8 accumulator keeps far fewer bits than -// f32, so chaining c through it drops low bits of the sums. Each MMA starts from zero instead -- -// one tile's sum fits exactly -- and the running sum is kept with ordinary f32 adds. -__device__ __forceinline__ void mma_e4m3_m16n8k32(float (&c)[4], - const uint32_t (&a)[4], - const uint32_t (&b)[2]) -{ - float d[4]; - asm( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %10, %10, %10};" - : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(0.0f)); -#pragma unroll - for (int i = 0; i < 4; ++i) { - c[i] += d[i]; - } -} - -// Feeds the 4-bit codes to mma.sync as e4m3, expanded in registers from plain 32-bit LDS. -template -struct i4_warp_mma_fp8 { - using raw_t = float; - - static constexpr int MMA_M = 16; - static constexpr int MMA_N = 8; - // A k-step reads 16 packed bytes = 32 codes, expanded to e4m3 and consumed by one MMA. - static constexpr int KSTEP_BYTES = 16; - static constexpr int SUB_M = WarpTile / MMA_M; - static constexpr int SUB_N = WarpTile / MMA_N; - static_assert(WarpTile % MMA_M == 0 && WarpTile % MMA_N == 0, - "the warp tile must be a whole number of native MMA tiles"); - static_assert(RowStride % 4 == 0, "rows must stay 4-byte aligned for the u32 fragment loads"); - - float c[SUB_M][SUB_N][4] = {}; - - // m16n8k32 fragment coordinates: `group` is the A row (and +8) and the B column this lane holds; - // `quad` picks its 4-byte slice of the k-step. A and B read the same slice, and expanding it to - // (low nibbles -> k 0..15, high nibbles -> k 16..31) the same way keeps their k orders aligned. - // Kept rolled for the same register reason as the u8 path. - template - __device__ __forceinline__ void accumulate(const uint8_t (*a)[RowStride], - const uint8_t (*b)[RowStride], - int row0, - int col0) - { - constexpr uint32_t kLowNibbles = 0x0f0f0f0fu; - const int lane = threadIdx.x % raft::warp_size(); - const int group = lane / 4; - const int quad = lane % 4; -#pragma unroll 1 - for (int kk = 0; kk < RowBytes / KSTEP_BYTES; ++kk) { - const int byte0 = kk * KSTEP_BYTES + quad * 4; - uint32_t a_frag[SUB_M][4]; - uint32_t b_frag[SUB_N][2]; - for (int msub = 0; msub < SUB_M; ++msub) { - const int row = row0 + msub * MMA_M + group; - const uint32_t r0 = *reinterpret_cast(a[row] + byte0); - const uint32_t r8 = *reinterpret_cast(a[row + 8] + byte0); - a_frag[msub][0] = r0 & kLowNibbles; - a_frag[msub][2] = (r0 >> 4) & kLowNibbles; - a_frag[msub][1] = r8 & kLowNibbles; - a_frag[msub][3] = (r8 >> 4) & kLowNibbles; - } - for (int nsub = 0; nsub < SUB_N; ++nsub) { - const uint32_t w = - *reinterpret_cast(b[col0 + nsub * MMA_N + group] + byte0); - b_frag[nsub][0] = w & kLowNibbles; - b_frag[nsub][1] = (w >> 4) & kLowNibbles; - } - for (int msub = 0; msub < SUB_M; ++msub) { - for (int nsub = 0; nsub < SUB_N; ++nsub) { - mma_e4m3_m16n8k32(c[msub][nsub], a_frag[msub], b_frag[nsub]); - } - } - } - } - - // Lane (group, quad) holds columns quad*2, +1 of rows group and group+8. Runs once, after every - // K chunk is accumulated, so this is where the codes' 1 / kE4m3CodeScale^2 comes off. - template - __device__ __forceinline__ void store(raw_t* dst, int row0, int col0) const - { - static_assert(Stride % 2 == 0, "float2 stores need even row strides"); - constexpr float code_scale_sq = kE4m3CodeScale * kE4m3CodeScale; - const int lane = threadIdx.x % raft::warp_size(); - const int group = lane / 4; - const int quad = lane % 4; - for (int msub = 0; msub < SUB_M; ++msub) { - for (int nsub = 0; nsub < SUB_N; ++nsub) { - for (int h = 0; h < 2; ++h) { - const int row = row0 + msub * MMA_M + group + h * 8; - const int col = col0 + nsub * MMA_N + quad * 2; - *reinterpret_cast(dst + row * Stride + col) = make_float2( - c[msub][nsub][h * 2] * code_scale_sq, c[msub][nsub][h * 2 + 1] * code_scale_sq); - } - } - } - } -}; -#endif - } // namespace cuvs::neighbors::nn_descent::detail diff --git a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh index bc7ad1d134..3a73791379 100644 --- a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh +++ b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh @@ -268,14 +268,13 @@ get_dequant_factors(const quantizer_view& quantizer, int64_t row) } // Converts one raw BBQ dot product into a final (post-epilogue) float distance, given both -// operands' precomputed dequant factors. Evaluated exactly once per matrix cell. `raw` is float so -// float accumulators (fp8 MMA) pass straight through; integer-accumulating callers cast. +// operands' precomputed dequant factors. Evaluated exactly once per matrix cell // dim/centroid_norm_sq/row_norm both come directly from quantizer_document/quantizer_query rather // than being passed separately, since every caller reads them the same id-indexed way -- row_norm // is only read for CosineExpanded (skipped entirely otherwise). template __device__ __forceinline__ float bbq_calculate_metric( - float raw, + uint32_t raw, const bbq_dequant_factors& doc_factors, const bbq_dequant_factors& query_factors, const quantizer_view& quantizer_document, @@ -288,9 +287,10 @@ __device__ __forceinline__ float bbq_calculate_metric( constexpr bool can_postprocess_dist = std::is_same_v; const float dim = static_cast(quantizer_document.dim()); - const float centered = - dim * doc_factors.lower * query_factors.lower + query_factors.lower * doc_factors.sum_delta + - doc_factors.lower * query_factors.sum_delta + doc_factors.delta * query_factors.delta * raw; + const float centered = dim * doc_factors.lower * query_factors.lower + + query_factors.lower * doc_factors.sum_delta + + doc_factors.lower * query_factors.sum_delta + + doc_factors.delta * query_factors.delta * static_cast(raw); const float corrections = doc_factors.corrections + query_factors.corrections; float d; if (metric == cuvs::distance::DistanceType::L2Expanded ||