From d770111cb1af73495755878c1357be43ae7679d7 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 2 Jun 2026 18:49:40 +0000 Subject: [PATCH] Remove stale duplicate .cu files from indexer/ subfolder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CUDA loader (dsv4/kernels/cuda/loader.py) resolves all .cu files relative to dsv4/kernels/cuda/. The indexer/ subfolder copies were never loaded — they were dead code that could silently diverge from the canonical copies in cuda/. --- dsv4/kernels/indexer/gather_kv.cu | 106 ------- dsv4/kernels/indexer/indexer_score_topk.cu | 306 --------------------- 2 files changed, 412 deletions(-) delete mode 100644 dsv4/kernels/indexer/gather_kv.cu delete mode 100644 dsv4/kernels/indexer/indexer_score_topk.cu diff --git a/dsv4/kernels/indexer/gather_kv.cu b/dsv4/kernels/indexer/gather_kv.cu deleted file mode 100644 index 77692d91..00000000 --- a/dsv4/kernels/indexer/gather_kv.cu +++ /dev/null @@ -1,106 +0,0 @@ -// gather_kv.cu — Gather selected compressed entries into a dense BF16 tile. -// -// One CTA per (query token, key_group). Each CTA handles a contiguous -// group of top-k entries for one query token. Reads from the FP8/BF16 -// split paged pool via block_table resolution, dequantizes FP8 → BF16, -// concatenates the RoPE half, writes to the dense output. -// -// Pure bandwidth-bound kernel — no MMA, just load-multiply-store. -// The output [T, top_k, head_dim] BF16 tile is what the FMHA kernel -// consumes. Sparsity is hidden in the gather; FMHA sees dense tiles. - -#include -#include -#include -#include -#include - - -__global__ void gather_kv_kernel( - // Inputs - const uint8_t* __restrict__ entries_fp8, // [num_blocks, epb, fp8_dim] - const __nv_bfloat16* __restrict__ entries_rope, // [num_blocks, epb, rope_dim] - const float* __restrict__ inv_scale, // [num_blocks, epb] - const int32_t* __restrict__ topk_indices, // [T, top_k] — compressed entry indices - const int32_t* __restrict__ block_table, // [T, max_logical_blocks] - // Output - __nv_bfloat16* __restrict__ output, // [T, top_k, head_dim] BF16 - // Geometry - int T, int top_k, int entries_per_block, - int head_dim, int rope_dim, int max_logical_blocks -) { - int fp8_dim = head_dim - rope_dim; - - // Each CTA handles one (query_token, topk_entry) pair. - int flat_idx = blockIdx.x; - int t = flat_idx / top_k; - int k = flat_idx % top_k; - if (t >= T) return; - - // Resolve which compressed entry to gather. - int comp_idx = topk_indices[t * top_k + k]; - if (comp_idx < 0) { - // Invalid entry — zero fill. - for (int d = threadIdx.x; d < head_dim; d += blockDim.x) { - output[t * top_k * head_dim + k * head_dim + d] = __float2bfloat16(0.0f); - } - return; - } - - int logical_block = comp_idx / entries_per_block; - int slot_in_block = comp_idx % entries_per_block; - int phys_block = block_table[t * max_logical_blocks + logical_block]; - - int block_entry = phys_block * entries_per_block + slot_in_block; - - // Dequantize and write FP8 half. - float s = inv_scale[block_entry]; - for (int d = threadIdx.x; d < fp8_dim; d += blockDim.x) { - uint8_t raw = entries_fp8[block_entry * fp8_dim + d]; - __nv_fp8_e4m3 fp8_val; - fp8_val.__x = raw; - float dequant = (float)fp8_val * s; - output[t * top_k * head_dim + k * head_dim + d] = __float2bfloat16(dequant); - } - - // Copy BF16 RoPE half. - for (int d = threadIdx.x; d < rope_dim; d += blockDim.x) { - output[t * top_k * head_dim + k * head_dim + fp8_dim + d] - = entries_rope[block_entry * rope_dim + d]; - } -} - - -void gather_kv_cuda( - torch::Tensor entries_fp8, - torch::Tensor entries_rope, - torch::Tensor inv_scale, - torch::Tensor topk_indices, - torch::Tensor block_table, - torch::Tensor output, - int64_t entries_per_block, int64_t rope_dim -) { - int T = topk_indices.size(0); - int top_k = topk_indices.size(1); - int head_dim = entries_fp8.size(2) + entries_rope.size(2); - int max_logical_blocks = block_table.size(1); - - int total_entries = T * top_k; - int threads = 128; - gather_kv_kernel<<>>( - entries_fp8.data_ptr(), - reinterpret_cast(entries_rope.data_ptr()), - inv_scale.data_ptr(), - topk_indices.data_ptr(), - block_table.data_ptr(), - reinterpret_cast<__nv_bfloat16*>(output.data_ptr()), - T, top_k, (int)entries_per_block, - (int)head_dim, (int)rope_dim, max_logical_blocks - ); - C10_CUDA_CHECK(cudaGetLastError()); -} - - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("gather_kv", &gather_kv_cuda, "Gather KV entries into dense tile"); -} diff --git a/dsv4/kernels/indexer/indexer_score_topk.cu b/dsv4/kernels/indexer/indexer_score_topk.cu deleted file mode 100644 index a691e81b..00000000 --- a/dsv4/kernels/indexer/indexer_score_topk.cu +++ /dev/null @@ -1,306 +0,0 @@ -// indexer_score_topk.cu — Fused score + ReLU + weighted-sum + top-k kernel. -// -// CSA Lightning Indexer (paper §2.3.1, eq. 16): -// I[t,s] = Σ_h w_h[t,h] · ReLU(q_I[t,h] · K^IComp[s,h]) -// Selected = TopK(I[t,:], k=csa_top_k) -// -// One CTA per query token. Streams indexer keys from the paged pool, -// computes per-head dot products in FP32, ReLU, weighted sum, top-k. -// -// Top-k strategy: each thread maintains a private top-k in registers -// over its strided slice of entries, then a block-level merge via -// bitonic sort on the shared heap. No in-loop barriers, no spinlocks. -// -// Phase 1 (this file): FP32 dot products via standard CUDA ops. -// Phase 2 (future): swap to FP4 tcgen05 MMA for production throughput. - -#include -#include -#include -#include -#include - -#include - -// FP4 E2M1 magnitude lookup (same as production) -__constant__ float E2M1_LUT[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; - -__device__ __forceinline__ float dequant_fp4_scalar( - uint8_t packed, int lane, - float group_scale, float global_scale -) { - int nibble = (lane == 0) ? (packed & 0x0F) : (packed >> 4); - int sign = (nibble >> 3) & 1; - int mag_bits = nibble & 0x07; - float magnitude = E2M1_LUT[mag_bits]; - float val = magnitude * group_scale * global_scale; - return sign ? -val : val; -} - -// ---- Per-thread local top-k ---- -// Each thread keeps LOCAL_K best scores in registers. -// LOCAL_K is a tuning parameter: larger = more accurate merge, -// smaller = less register pressure. -// For top_k=1024 and 128 threads: LOCAL_K=8 means 128*8=1024 candidates -// for the block-level merge, which is exact. -// For top_k=512 and 128 threads: LOCAL_K=4 gives 512 candidates, also exact. -// If top_k > n_threads * LOCAL_K, the merge is approximate (top-K of -// n_threads*LOCAL_K candidates). Increase LOCAL_K or n_threads to compensate. - -#ifndef INDEXER_LOCAL_K -#define INDEXER_LOCAL_K 8 -#endif - -__device__ __forceinline__ void local_heap_insert( - float* scores, int32_t* blocks, - float score, int32_t block_id, int k -) { - if (score <= scores[0]) return; - scores[0] = score; - blocks[0] = block_id; - // Sift down - int root = 0; - while (root < (k >> 1)) { - int left = 2 * root + 1; - int right = 2 * root + 2; - int smallest = root; - if (left < k && scores[left] < scores[smallest]) smallest = left; - if (right < k && scores[right] < scores[smallest]) smallest = right; - if (smallest == root) break; - float ts = scores[root]; int32_t ti = blocks[root]; - scores[root] = scores[smallest]; blocks[root] = blocks[smallest]; - scores[smallest] = ts; blocks[smallest] = ti; - root = smallest; - } -} - -// ---- Block-level merge: merge n_threads × LOCAL_K candidates ---- -// Each thread writes its local top-k to shared memory, then a single -// thread (or warp) does a final top-k selection from the combined set. -// Total candidates = n_threads * LOCAL_K. -// For top_k <= total_candidates, this is exact. -// For top_k > total_candidates, increase LOCAL_K. - -__device__ __forceinline__ void heap_insert_shared( - float* heap_scores, int32_t* heap_blocks, - float score, int32_t block_id, int k -) { - if (score <= heap_scores[0]) return; - heap_scores[0] = score; - heap_blocks[0] = block_id; - int root = 0; - while (root < (k >> 1)) { - int left = 2 * root + 1; - int right = 2 * root + 2; - int smallest = root; - if (left < k && heap_scores[left] < heap_scores[smallest]) smallest = left; - if (right < k && heap_scores[right] < heap_scores[smallest]) smallest = right; - if (smallest == root) break; - float ts = heap_scores[root]; int32_t ti = heap_blocks[root]; - heap_scores[root] = heap_scores[smallest]; heap_blocks[root] = heap_blocks[smallest]; - heap_scores[smallest] = ts; heap_blocks[smallest] = ti; - root = smallest; - } -} - -// =========================================================================== -// Main kernel -// =========================================================================== - -__global__ void indexer_score_topk_fp32_kernel( - const float* __restrict__ q_I, - const float* __restrict__ w_h, - const uint8_t* __restrict__ keys_fp4, - const uint8_t* __restrict__ key_scale, - const float* __restrict__ key_gscale, - const int32_t* __restrict__ block_table, - const int32_t* __restrict__ valid_lens, - int32_t* __restrict__ topk_indices, - int n_heads, int head_dim, int top_k, - int entries_per_block, int max_logical_blocks -) { - int t = blockIdx.x; - if (t >= gridDim.x) return; - - int tid = threadIdx.x; - int n_threads = blockDim.x; - int num_valid = valid_lens[t]; - int n_groups = head_dim / 16; - int n_bytes = head_dim / 2; - - // ---- Per-thread local top-k in registers ---- - // LOCAL_K entries per thread. Min-heap (root = smallest of local best). - float local_scores[INDEXER_LOCAL_K]; - int32_t local_blocks[INDEXER_LOCAL_K]; - for (int i = 0; i < INDEXER_LOCAL_K; i++) { - local_scores[i] = -INFINITY; - local_blocks[i] = -1; - } - - // ---- Load w_h into shared memory ---- - extern __shared__ char smem[]; - float* smem_w = reinterpret_cast(smem); - // The rest of smem is used for the merge phase (allocated after w_h) - // Layout: [w_h: n_heads floats] [merge_scores: top_k floats] [merge_blocks: top_k ints] - // [per_thread_scores: n_threads * LOCAL_K floats] [per_thread_blocks: n_threads * LOCAL_K ints] - // But we allocate dynamically, so let's compute offsets. - - for (int h = tid; h < n_heads; h += n_threads) { - smem_w[h] = w_h[t * n_heads + h]; - } - __syncthreads(); // safe — outside the strided loop - - // ---- Stream over entries (strided, no barriers) ---- - // Each thread handles entries s = tid, tid+n_threads, tid+2*n_threads, ... - // No __syncthreads() in this loop. No shared heap access. - // Each thread accumulates into its private register heap. - - for (int s = tid; s < num_valid; s += n_threads) { - int logical_block = s / entries_per_block; - int slot_in_block = s % entries_per_block; - int phys_block = block_table[t * max_logical_blocks + logical_block]; - int block_entry = phys_block * entries_per_block + slot_in_block; - - float global_s = key_gscale[phys_block]; - - float score = 0.0f; - for (int h = 0; h < n_heads; h++) { - float dot = 0.0f; - for (int g = 0; g < n_groups; g++) { - uint8_t raw_scale = key_scale[block_entry * n_groups + g]; - __nv_fp8_e4m3 fp8_s; - fp8_s.__x = raw_scale; - float group_s = (float)fp8_s * global_s; - - for (int b = 0; b < 8; b++) { - uint8_t packed = keys_fp4[block_entry * n_bytes + g * 8 + b]; - float v0 = dequant_fp4_scalar(packed, 0, group_s, 1.0f); - float v1 = dequant_fp4_scalar(packed, 1, group_s, 1.0f); - int d0 = g * 16 + 2 * b; - int d1 = d0 + 1; - dot += v0 * q_I[t * n_heads * head_dim + h * head_dim + d0]; - dot += v1 * q_I[t * n_heads * head_dim + h * head_dim + d1]; - } - } - if (dot > 0.0f) { - score += smem_w[h] * dot; - } - } - - // Insert into per-thread local heap (registers, no sync needed) - local_heap_insert(local_scores, local_blocks, score, s, INDEXER_LOCAL_K); - } - - // ---- Block-level merge ---- - // Each thread writes its LOCAL_K candidates to shared memory. - // Then one thread builds the final top-k from all candidates. - // Total candidates = n_threads * LOCAL_K. - // For top_k=1024, n_threads=128, LOCAL_K=8: 1024 candidates, exact merge. - // For top_k=512, n_threads=128, LOCAL_K=4: 512 candidates, exact merge. - - float* merge_scores = smem_w + n_heads; - int32_t* merge_blocks = reinterpret_cast(merge_scores + top_k); - float* per_thread_scores = reinterpret_cast(merge_blocks + top_k); - int32_t* per_thread_blocks = reinterpret_cast(per_thread_scores + n_threads * INDEXER_LOCAL_K); - - // Initialize merge heap - for (int i = tid; i < top_k; i += n_threads) { - merge_scores[i] = -INFINITY; - merge_blocks[i] = -1; - } - - // Write local top-k to per-thread region in shared memory - int my_offset = tid * INDEXER_LOCAL_K; - for (int i = 0; i < INDEXER_LOCAL_K; i++) { - per_thread_scores[my_offset + i] = local_scores[i]; - per_thread_blocks[my_offset + i] = local_blocks[i]; - } - __syncthreads(); // wait for all threads to write their candidates - - // Single thread builds the final top-k from all candidates - // This is O(n_threads * LOCAL_K * log(top_k)) — fast for reasonable sizes. - // For n_threads=128, LOCAL_K=8, top_k=1024: 1024 inserts, ~10K comparisons. - if (tid == 0) { - for (int i = 0; i < n_threads * INDEXER_LOCAL_K; i++) { - if (per_thread_scores[i] > -INFINITY) { - heap_insert_shared(merge_scores, merge_blocks, - per_thread_scores[i], per_thread_blocks[i], top_k); - } - } - } - __syncthreads(); // wait for merge to complete - - // ---- Write top-k indices to global memory ---- - // Sort the merge heap by score descending (selection sort, top_k <= 1024) - if (tid == 0) { - for (int i = 0; i < top_k; i++) { - int best = i; - for (int j = i + 1; j < top_k; j++) { - if (merge_scores[j] > merge_scores[best] || - (merge_scores[j] == merge_scores[best] && - merge_blocks[j] < merge_blocks[best])) { - best = j; - } - } - if (best != i) { - float ts = merge_scores[i]; int32_t ti = merge_blocks[i]; - merge_scores[i] = merge_scores[best]; merge_blocks[i] = merge_blocks[best]; - merge_scores[best] = ts; merge_blocks[best] = ti; - } - topk_indices[t * top_k + i] = merge_blocks[i]; - } - } -} - - -// =========================================================================== -// PyTorch binding -// =========================================================================== - -void indexer_score_topk_fp32_cuda( - torch::Tensor q_I, - torch::Tensor w_h, - torch::Tensor keys_fp4, - torch::Tensor key_scale, - torch::Tensor key_gscale, - torch::Tensor block_table, - torch::Tensor valid_lens, - torch::Tensor topk_indices, - int64_t n_heads, int64_t head_dim, int64_t top_k, - int64_t entries_per_block -) { - int T = q_I.size(0); - int max_logical_blocks = block_table.size(1); - int threads = 128; - - // SMEM layout: - // w_h: n_heads floats - // merge_scores: top_k floats - // merge_blocks: top_k ints - // per_thread_scores: n_threads * INDEXER_LOCAL_K floats - // per_thread_blocks: n_threads * INDEXER_LOCAL_K ints - int smem_bytes = n_heads * sizeof(float) - + top_k * sizeof(float) - + top_k * sizeof(int32_t) - + threads * INDEXER_LOCAL_K * sizeof(float) - + threads * INDEXER_LOCAL_K * sizeof(int32_t); - - indexer_score_topk_fp32_kernel<<>>( - q_I.data_ptr(), - w_h.data_ptr(), - keys_fp4.data_ptr(), - key_scale.data_ptr(), - key_gscale.data_ptr(), - block_table.data_ptr(), - valid_lens.data_ptr(), - topk_indices.data_ptr(), - (int)n_heads, (int)head_dim, (int)top_k, - (int)entries_per_block, max_logical_blocks - ); - C10_CUDA_CHECK(cudaGetLastError()); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("indexer_score_topk_fp32", &indexer_score_topk_fp32_cuda, - "Indexer score + top-k (FP32 dot products, no-deadlock)"); -}