nvfp4-megamoe-kernel/dsv4/kernels/cuda/indexer_score_topk.cu

#include <cuda.h>
#include <cuda_fp8.h>
#include <cuda_bf16.h>
#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>
#include <limits>

__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 = nibble & 0x07;
    float val = (float)mag * group_scale * global_scale;
    return sign ? -val : val;
}

__device__ void heap_insert(
    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;
    }
}

__global__ void indexer_score_topk_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 total_groups = n_heads * n_groups;
    int n_bytes = head_dim / 2;
    int total_bytes = n_heads * n_bytes;

    // Per-thread heap in REGISTERS (top_k <= 1024, but for small k this works)
    // Actually, use shared memory with a simple layout
    __shared__ float s_heap_scores[1024];  // max top_k
    __shared__ int32_t s_heap_blocks[1024];
    __shared__ float s_w[64];  // max n_heads
    __shared__ int s_lock;

    // Load w_h
    for (int h = tid; h < n_heads; h += n_threads) {
        s_w[h] = w_h[t * n_heads + h];
    }
    // Init heap
    for (int i = tid; i < top_k; i += n_threads) {
        s_heap_scores[i] = -INFINITY;
        s_heap_blocks[i] = -1;
    }
    if (tid == 0) s_lock = 0;
    __syncthreads();

    // Stream over entries
    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 flat = phys_block * entries_per_block + slot_in_block;

        float gs = key_gscale[phys_block];

        // Compute score
        float score = 0.0f;
        for (int h = 0; h < n_heads; h++) {
            float dot = 0.0f;
            int h_byte_off = h * n_bytes;
            int h_group_off = h * n_groups;
            for (int g = 0; g < n_groups; g++) {
                uint8_t raw_sc = key_scale[flat * total_groups + h_group_off + g];
                __nv_fp8_e4m3 fp8_s;
                fp8_s.__x = raw_sc;
                float grp_s = (float)fp8_s * gs;

                for (int b = 0; b < 8; b++) {
                    uint8_t packed = keys_fp4[flat * total_bytes + h_byte_off + g * 8 + b];
                    float v0 = dequant_fp4_scalar(packed, 0, grp_s, 1.0f);
                    float v1 = dequant_fp4_scalar(packed, 1, grp_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 += s_w[h] * dot;
            }
        }

        // Insert into shared heap (serialized via spinlock)
        while (atomicCAS(&s_lock, 0, 1) != 0) {}
        heap_insert(s_heap_scores, s_heap_blocks, score, s, top_k);
        atomicExch(&s_lock, 0);
    }
    __syncthreads();

    // Sort + write output
    if (tid == 0) {
        for (int i = 0; i < top_k; i++) {
            int best = i;
            for (int j = i + 1; j < top_k; j++) {
                if (s_heap_scores[j] > s_heap_scores[best]) best = j;
            }
            if (best != i) {
                float ts = s_heap_scores[i]; int32_t ti = s_heap_blocks[i];
                s_heap_scores[i] = s_heap_scores[best]; s_heap_blocks[i] = s_heap_blocks[best];
                s_heap_scores[best] = ts; s_heap_blocks[best] = ti;
            }
            topk_indices[t * top_k + i] = s_heap_blocks[i];
        }
    }
}

void indexer_score_topk_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);
    indexer_score_topk_kernel<<<T, 128>>>(
        q_I.data_ptr<float>(), w_h.data_ptr<float>(),
        keys_fp4.data_ptr<uint8_t>(), key_scale.data_ptr<uint8_t>(),
        key_gscale.data_ptr<float>(), block_table.data_ptr<int32_t>(),
        valid_lens.data_ptr<int32_t>(), topk_indices.data_ptr<int32_t>(),
        (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", &indexer_score_topk_cuda);
}