cuda_kernels.cu File Reference

CUDA kernel implementations for TinyLlama GPU operations. More...

#include "cuda_kernels.h"
#include "determinism.h"
#include "logger.h"
#include "model_macros.h"
#include <stdint.h>
#include <stdio.h>
#include <algorithm>
#include <cmath>
#include <iostream>
#include <cfloat>
#include <cstdlib>
#include <string>

Include dependency graph for cuda_kernels.cu:

Go to the source code of this file.

Functions
static bool	kernel_debug_copies ()
__device__ float	atomicMaxFloat (float *address, float val)
__device__ float	atomicAddFloat (float *address, float val)
__device__ float	bf16_to_float32_device (uint16_t bf16_raw)
	Converts BF16 to FP32 on the device.
__device__ uint16_t	float32_to_bf16_device (float fp32_val)
__global__ void	rmsnorm_sum_squares_kernel (const float *x, float *partial_sums, int n)
	Kernel for computing sum of squares in RMS normalization.
__global__ void	rmsnorm_apply_kernel (const float *x, const float *weight, float *out, int n, float inv_norm_factor)
	Kernel for applying RMS normalization.
__global__ void	reduce_partial_sums_kernel (const float *partial_sums, float *total_sum_sq_out, int num_partial_sums)
void	rmsnorm_vector_cuda (const float *x_dev, const float *weight_dev, float *out_dev, int n, float eps, cudaStream_t stream)
	Implementation of RMS normalization with device pointers.
void	rmsnorm_vector_cuda (const std::vector< float > &x_in_host, const std::vector< float > &weight_host, std::vector< float > &out_host, int n, float eps)
void	matvec_f32_f32_cuda (cublasHandle_t handle, const float *mat_f32_dev, const float *vec_f32_dev, float *out_f32_dev, int rows, int cols, cudaStream_t stream)
	Implementation of matrix-vector multiplication with FP32.
void	matvec_bf16_f32_cuda (cublasHandle_t handle, const uint16_t *mat_bf16_dev, const float *vec_f32_dev, float *out_f32_dev, int rows, int cols, bool use_tensor_cores, cudaStream_t stream)
	Matrix-vector multiplication with BF16 matrix and FP32 vector/output on GPU.
void	matvec_f32_f32_cuda (cublasHandle_t handle, const std::vector< float > &mat_f32_host, const std::vector< float > &vec_f32_host, std::vector< float > &out_f32_host, int rows, int cols)
__global__ void	silu_kernel (const float *x, float *out, int n)
void	silu_cuda (const std::vector< float > &x_host, std::vector< float > &out_host, int n)
__global__ void	softmax_find_max_kernel (const float *x, float *partial_max, int n)
__global__ void	softmax_exp_sum_kernel (const float *x, float *partial_sums, int n, float max_val)
__global__ void	softmax_normalize_kernel (const float *x, float *out, int n, float max_val, float inv_sum)
void	softmax_vector_cuda (const std::vector< float > &x_host, std::vector< float > &out_host, int n)
__global__ void	swiglu_kernel (const float *gate, const float *up, float *out, int n)
void	swiglu_cuda (const float *gate_dev, const float *up_dev, float *out_dev, int n, cudaStream_t stream)
__global__ void	swiglu_batch_cuda_kernel (float *d_out_batch, const float *d_gate_act_batch, const float *d_up_act_batch, int num_tokens, int intermediate_size)
void	swiglu_batch_cuda (float *d_out_batch, const float *d_gate_act_batch, const float *d_up_act_batch, int num_tokens, int intermediate_size, cudaStream_t stream)
__global__ void	rope_kernel (float *x, int num_heads, int head_dim, const float *all_freqs_cis_base, int pos, bool use_adjacent_pairing)
void	rope_cuda (float *x_dev, int num_heads, int head_dim, const float *all_freqs_cis_dev_base, int pos, bool use_adjacent_pairing, cudaStream_t stream)
__global__ void	rope_batch_kernel_cuda (float *x_batch_tensor, int num_tokens, int num_heads_in_tensor, int head_dim, const float *d_all_freqs_cis_base, int start_pos_offset, bool use_adjacent_pairing)
void	rope_batch_cuda (float *d_q_batch, float *d_k_batch, const float *d_all_freqs_cis_base, int num_tokens, int num_q_heads, int num_kv_heads, int head_dim, int start_pos_in_kv_cache, bool use_adjacent_pairing, cudaStream_t stream)
__global__ void	attention_kernel (const float *Q_current, const float *K_layer_cache_base, const float *V_layer_cache_base, float *out, int current_seq_len, int head_dim, float scale, int cache_num_kv_heads, int num_q_heads)
void	attention_cuda (const float *Q_current_dev, const float *K_layer_cache_base, const float *V_layer_cache_base, float *out_dev, int num_heads, int current_seq_len, int head_dim, float scale, int cache_max_seq_len, int cache_num_kv_heads, cudaStream_t stream)
__global__ void	attention_selective_dequant_kernel (const float *__restrict__ Q_current_dev, const int8_t *__restrict__ K_quantized_cache_base, const int8_t *__restrict__ V_quantized_cache_base, const float *__restrict__ K_scales_cache_base, const float *__restrict__ V_scales_cache_base, float *__restrict__ selective_k_dequant_buffer, float *__restrict__ selective_v_dequant_buffer, float *__restrict__ out_dev, int num_heads, int current_seq_len, int head_dim, float scale, int cache_max_seq_len, int cache_num_kv_heads)
void	attention_cuda_selective_dequant (const float *Q_current_dev, const int8_t *K_quantized_cache_base, const int8_t *V_quantized_cache_base, const float *K_scales_cache_base, const float *V_scales_cache_base, float *selective_k_dequant_buffer, float *selective_v_dequant_buffer, float *out_dev, int num_heads, int current_seq_len, int head_dim, float scale, int cache_max_seq_len, int cache_num_kv_heads, cudaStream_t stream)
__global__ void	add_vectors_kernel (const float *a, const float *b, float *result, int n)
void	add_vectors_cuda (const float *a_dev, const float *b_dev, float *result_dev, int n, cudaStream_t stream)
__global__ void	add_residual_kernel (const float *matvec_out, const float *residual, float *result, int n)
void	add_residual_cuda (const float *matvec_out_dev, const float *residual_dev, float *result_dev, int n, cudaStream_t stream)
__global__ void	add_residual_batch_cuda_kernel (float *d_output_batch, const float *d_input_a_batch, const float *d_input_b_batch, int num_tokens, int hidden_size)
void	add_residual_batch_cuda (float *d_output_batch, const float *d_input_a_batch, const float *d_input_b_batch, int num_tokens, int hidden_size, cudaStream_t stream)
__global__ void	update_kv_cache_kernel (float *cache_base_ptr, const float *current_kv_vector, int pos, int kv_head_idx, int max_seq_len, int num_kv_heads, int head_dim)
void	update_kv_cache_cuda (float *cache_base_ptr, const float *current_kv_head_vector, int pos, int kv_head_idx, int max_seq_len, int num_kv_heads, int head_dim, cudaStream_t stream)
__global__ void	rope_and_update_kv_cache_kernel (float *cache_base_ptr, const float *kv_vector_head, const float *all_freqs_cis_base, int pos, int kv_head_idx, int max_seq_len, int num_kv_heads, int head_dim)
void	rope_and_update_kv_cache_cuda (float *cache_base_ptr, const float *kv_vector_head, const float *all_freqs_cis_base, int pos, int kv_head_idx, int max_seq_len, int num_kv_heads, int head_dim, cudaStream_t stream)
__global__ void	lookup_embedding_kernel (const void *__restrict__ table_dev, float *__restrict__ output_dev, int token_id, int hidden_size, int vocab_size, bool is_bf16)
void	lookup_embedding_cuda (const void *table_dev, float *output_dev, int token_id, int hidden_size, int vocab_size, bool is_bf16, cudaStream_t stream)
__global__ void	convert_bf16_to_fp32_kernel (const uint16_t *__restrict__ bf16_in, float *__restrict__ fp32_out, size_t n_elements)
__global__ void	rmsnorm_optimized_kernel (const float *__restrict__ x, const float *__restrict__ weight, float *__restrict__ out, int n, float eps)
	Optimized RMS normalization using simple warp primitives and shared memory.
__global__ void	softmax_optimized_kernel (const float *__restrict__ x, float *__restrict__ out, int n)
	Optimized softmax using shared memory and warp primitives.
__global__ void	attention_optimized_kernel (const float *__restrict__ Q_current, const float *__restrict__ K_layer_cache_base, const float *__restrict__ V_layer_cache_base, float *__restrict__ out, int current_seq_len, int head_dim, float scale, int cache_num_kv_heads, int num_q_heads, const int *step_dev)
	Corrected and optimized attention kernel.
__global__ void	rmsnorm_optimized_det_kernel (const float *__restrict__ x, const float *__restrict__ weight, float *__restrict__ out, int n, float eps)
__global__ void	attention_optimized_det_kernel (const float *__restrict__ Q_current, const float *__restrict__ K_layer_cache_base, const float *__restrict__ V_layer_cache_base, float *__restrict__ out, int current_seq_len, int head_dim, float scale, int cache_num_kv_heads, int num_q_heads, const int *step_dev)
void	rmsnorm_vector_cuda_optimized (const float *x_dev, const float *weight_dev, float *out_dev, int n, float eps, cudaStream_t stream)
	Optimized wrappers using the new high-performance kernels.
void	softmax_vector_cuda_optimized (const float *x_dev, float *out_dev, int n, cudaStream_t stream)
int	decode_attention_kind ()
void	decode_attention_tiles (int *min_chunk, int *max_chunks)
void	attention_cuda_optimized (const float *Q_current_dev, const float *K_layer_cache_base, const float *V_layer_cache_base, float *out_dev, int num_heads, int current_seq_len, int head_dim, float scale, int cache_max_seq_len, int cache_num_kv_heads, cudaStream_t stream)
__global__ void	convert_fp32_to_bf16_kernel (const float *__restrict__ fp32_in, uint16_t *__restrict__ bf16_out, size_t n_elements)
void	convert_fp32_to_bf16_cuda (const float *fp32_in_dev, uint16_t *bf16_out_dev, size_t n_elements, cudaStream_t stream)
void	gemm_bf16_bf16_cuda (cublasHandle_t handle, bool transa_user, bool transb_user, int m_user, int n_user, int k_user, const float *alpha_user, const uint16_t *A_bf16_user, int lda_user, const uint16_t *B_bf16_user, int ldb_user, const float *beta_user, uint16_t *C_bf16_user, int ldc_user, cudaStream_t stream)
void	gemm_f32_to_bf16_f32_cuda (cublasHandle_t handle, bool transa_user, bool transb_user, int m_user, int n_user, int k_user, const float *alpha_user, const float *A_f32_user, int lda_user, const uint16_t *B_bf16_user, int ldb_user, const float *beta_user, float *C_f32_user, int ldc_user, cudaStream_t stream)
void	convert_bf16_to_fp32_cuda (const uint16_t *bf16_in_dev, float *fp32_out_dev, size_t n_elements, cudaStream_t stream)
__global__ void	decode_embed_kernel (const float *__restrict__ table, float *__restrict__ out, const int *__restrict__ step_dev, int hidden_size, int vocab_size)
void	decode_embed_cuda (const float *table_dev, float *out_dev, const int *step_dev, int hidden_size, int vocab_size, cudaStream_t stream)
__global__ void	rope_devpos_kernel (float *x, int num_heads, int head_dim, const float *all_freqs_cis_base, const int *step_dev, bool use_adjacent_pairing)
void	rope_devpos_cuda (float *x_dev, int num_heads, int head_dim, const float *all_freqs_cis_dev_base, const int *step_dev, bool use_adjacent_pairing, cudaStream_t stream)
__global__ void	kv_write_devpos_kernel (float *k_cache, float *v_cache, const float *__restrict__ k, const float *__restrict__ v, const int *step_dev, int row)
void	kv_write_devpos_cuda (float *k_cache_dev, float *v_cache_dev, const float *k_dev, const float *v_dev, const int *step_dev, int row, cudaStream_t stream)
void	attention_decode_at_cuda (const float *Q_current_dev, const float *K_layer_cache_base, const float *V_layer_cache_base, float *out_dev, int num_heads, int seq_len, int head_dim, float scale, int cache_max_seq_len, int cache_num_kv_heads, cudaStream_t stream)
void	attention_devpos_cuda (const float *Q_current_dev, const float *K_layer_cache_base, const float *V_layer_cache_base, float *out_dev, int num_heads, int head_dim, float scale, int cache_max_seq_len, int cache_num_kv_heads, const int *step_dev, cudaStream_t stream)
void	gemm_rm_f32_cuda (cublasHandle_t handle, const float *X, const float *W, float *Y, int T, int O, int I, cudaStream_t stream)
size_t	moe_ffn_scratch_floats (int T, int d, int d_ff, int E, int k)
void	moe_ffn_cuda (cublasHandle_t handle, const float *x, int T, int d, int d_ff, int E, int k, bool renorm, const float *router, const float *gate_exps, const float *up_exps, const float *down_exps, float *scratch, float *out, cudaStream_t stream)

Detailed Description

CUDA kernel implementations for TinyLlama GPU operations.

This file contains the actual CUDA kernel implementations for the GPU-accelerated operations declared in cuda_kernels.h. The implementations are optimized for NVIDIA GPUs and use CUDA-specific features for maximum performance.

Key implementation details:

• Uses shared memory for efficient reduction operations
• Implements parallel reduction patterns for normalization
• Leverages cuBLAS for matrix operations
• Supports both FP32 and BF16 formats
• Includes error checking and memory management

Definition in file cuda_kernels.cu.

Function Documentation

add_residual_batch_cuda()

void add_residual_batch_cuda	(	float *	d_output_batch,
		const float *	d_input_a_batch,
		const float *	d_input_b_batch,
		int	num_tokens,
		int	hidden_size,
		cudaStream_t	stream
	)

Definition at line 1274 of file cuda_kernels.cu.

                         {
 
 
    if (num_tokens == 0 || hidden_size == 0) {
        Logger::info("[ADD_RESIDUAL_BATCH_CUDA_SKIP] num_tokens or hidden_size is 0. Nothing to do.");
        Logger::info("[ADD_RESIDUAL_BATCH_CUDA_EXIT] Skipped operation.");
        return; // Nothing to do
    }
 
    // Log initial input values (first token, first few elements)
    int log_count_elements = std::min(3, hidden_size);
    if (num_tokens > 0 && hidden_size > 0) {
        if (d_input_a_batch) {
            std::vector<float> h_input_a_sample(log_count_elements);
            // Input A is [num_tokens, hidden_size]. Offset for T0 is 0.
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_input_a_sample.data(), d_input_a_batch, log_count_elements * sizeof(float), cudaMemcpyDeviceToHost, stream));
        }
        if (d_input_b_batch) {
            std::vector<float> h_input_b_sample(log_count_elements);
            // Input B is [num_tokens, hidden_size]. Offset for T0 is 0.
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_input_b_sample.data(), d_input_b_batch, log_count_elements * sizeof(float), cudaMemcpyDeviceToHost, stream));
        }
    }
 
    const int threads_per_block = 256; 
    dim3 grid_dim(num_tokens);
    dim3 block_dim(threads_per_block);
 
    add_residual_batch_cuda_kernel<<<grid_dim, block_dim, 0, stream>>>(
        d_output_batch,
        d_input_a_batch,
        d_input_b_batch,
        num_tokens,
        hidden_size
    );
    gpuErrchk(cudaGetLastError());
 
    // Log output values
    if (num_tokens > 0 && hidden_size > 0 && d_output_batch) {
        std::vector<float> h_output_sample(log_count_elements);
        // Output is [num_tokens, hidden_size]. Offset for T0 is 0.
        if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_output_sample.data(), d_output_batch, log_count_elements * sizeof(float), cudaMemcpyDeviceToHost, stream));
    }
}

References Logger::info(), j, and kernel_debug_copies().

add_residual_batch_cuda_kernel()

__global__ void add_residual_batch_cuda_kernel	(	float *	d_output_batch,
		const float *	d_input_a_batch,
		const float *	d_input_b_batch,
		int	num_tokens,
		int	hidden_size
	)

Definition at line 1250 of file cuda_kernels.cu.

                     {
 
    int token_idx = blockIdx.x;
    if (token_idx >= num_tokens) {
        return;
    }
 
    // Base pointers for the current token
    const float* current_token_input_a_ptr = d_input_a_batch + (size_t)token_idx * hidden_size;
    const float* current_token_input_b_ptr = d_input_b_batch + (size_t)token_idx * hidden_size;
    float* current_token_output_ptr = d_output_batch + (size_t)token_idx * hidden_size;
 
    // Threads in the block parallelize over hidden_size
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        current_token_output_ptr[i] = current_token_input_a_ptr[i] + current_token_input_b_ptr[i];
    }
}

References j.

add_residual_cuda()

void add_residual_cuda	(	const float *	matvec_out_dev,
		const float *	residual_dev,
		float *	result_dev,
		int	n,
		cudaStream_t	stream
	)

Definition at line 1239 of file cuda_kernels.cu.

                                                                      {
  const int threads_per_block = 256;
  const int num_blocks = (n + threads_per_block - 1) / threads_per_block;
 
  add_residual_kernel<<<num_blocks, threads_per_block, 0, stream>>>(
      matvec_out_dev, residual_dev, result_dev, n);
  gpuErrchk(cudaGetLastError());
}

References j, and n.

add_residual_kernel()

__global__ void add_residual_kernel	(	const float *	matvec_out,
		const float *	residual,
		float *	result,
		int	n
	)

Definition at line 1230 of file cuda_kernels.cu.

                                           {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) {
    result[i] = matvec_out[i] + residual[i];
  }
}

References j, and n.

add_vectors_cuda()

void add_vectors_cuda	(	const float *	a_dev,
		const float *	b_dev,
		float *	result_dev,
		int	n,
		cudaStream_t	stream
	)

Definition at line 1220 of file cuda_kernels.cu.

                                                  {
  const int threads_per_block = 256;
  const int num_blocks = (n + threads_per_block - 1) / threads_per_block;
 
  add_vectors_kernel<<<num_blocks, threads_per_block, 0, stream>>>(
      a_dev, b_dev, result_dev, n);
  gpuErrchk(cudaGetLastError());
}

References j, and n.

add_vectors_kernel()

__global__ void add_vectors_kernel	(	const float *	a,
		const float *	b,
		float *	result,
		int	n
	)

Definition at line 1212 of file cuda_kernels.cu.

                                                         {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) {
    result[i] = a[i] + b[i];
  }
}

References a, j, and n.

atomicAddFloat()

__device__ float atomicAddFloat	(	float *	address,
		float	val
	)

Definition at line 81 of file cuda_kernels.cu.

                                                           {
    int* address_as_int = (int*)address;
    int old = *address_as_int, assumed;
    do {
        assumed = old;
        old = atomicCAS(address_as_int, assumed,
                        __float_as_int(val + __int_as_float(assumed)));
    } while (assumed != old);
    return __int_as_float(old);
}

References j.

Referenced by attention_optimized_kernel(), and softmax_optimized_kernel().

atomicMaxFloat()

__device__ float atomicMaxFloat	(	float *	address,
		float	val
	)

Definition at line 70 of file cuda_kernels.cu.

                                                           {
    int* address_as_int = (int*)address;
    int old = *address_as_int, assumed;
    do {
        assumed = old;
        old = atomicCAS(address_as_int, assumed,
                        __float_as_int(fmaxf(val, __int_as_float(assumed))));
    } while (assumed != old);
    return __int_as_float(old);
}

References j.

Referenced by attention_optimized_kernel(), and softmax_optimized_kernel().

attention_cuda()

void attention_cuda	(	const float *	Q_current_dev,
		const float *	K_layer_cache_base,
		const float *	V_layer_cache_base,
		float *	out_dev,
		int	num_heads,
		int	current_seq_len,
		int	head_dim,
		float	scale,
		int	cache_max_seq_len,
		int	cache_num_kv_heads,
		cudaStream_t	stream
	)

Definition at line 1093 of file cuda_kernels.cu.

                                         {
  // Grid: num_heads blocks, each block handles one attention head
  // Block: head_dim threads, each thread handles one output dimension
  dim3 grid_dim(num_heads);
  dim3 block_dim(head_dim);
  // scores[current_seq_len] + dot_product_terms[head_dim] (the kernel's extern shared layout)
  const size_t shm = (size_t)(current_seq_len + head_dim) * sizeof(float);
 
  attention_kernel<<<grid_dim, block_dim, shm, stream>>>(
      Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev,
      current_seq_len, head_dim, scale, cache_num_kv_heads, num_heads);
  gpuErrchk(cudaGetLastError());
}

References j, and scale.

attention_cuda_optimized()

void attention_cuda_optimized	(	const float *	Q_current_dev,
		const float *	K_layer_cache_base,
		const float *	V_layer_cache_base,
		float *	out_dev,
		int	num_heads,
		int	current_seq_len,
		int	head_dim,
		float	scale,
		int	cache_max_seq_len,
		int	cache_num_kv_heads,
		cudaStream_t	stream
	)

Definition at line 3319 of file cuda_kernels.cu.

                                                  {
    if (current_seq_len == 0) return;
    if (decode_attn_device(Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev, num_heads,
                          current_seq_len, head_dim, scale, current_seq_len, cache_num_kv_heads, nullptr, stream))
        return;
    const int threads_per_block = 256;
    size_t shared_mem_size = (cache_max_seq_len + 1) * sizeof(float);
 
    if (ember::deterministic()) {
        attention_optimized_det_kernel<<<num_heads, threads_per_block, shared_mem_size, stream>>>(
            Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev,
            current_seq_len, head_dim, scale, cache_num_kv_heads, num_heads, nullptr);
        gpuErrchk(cudaGetLastError());
        return;
    }
    attention_optimized_kernel<<<num_heads, threads_per_block, shared_mem_size, stream>>>(
        Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev,
        current_seq_len, head_dim, scale, cache_num_kv_heads, num_heads, nullptr);
    gpuErrchk(cudaGetLastError());
}

References ember::deterministic(), j, and scale.

Referenced by attention_decode_at_cuda().

attention_cuda_selective_dequant()

void attention_cuda_selective_dequant	(	const float *	Q_current_dev,
		const int8_t *	K_quantized_cache_base,
		const int8_t *	V_quantized_cache_base,
		const float *	K_scales_cache_base,
		const float *	V_scales_cache_base,
		float *	selective_k_dequant_buffer,
		float *	selective_v_dequant_buffer,
		float *	out_dev,
		int	num_heads,
		int	current_seq_len,
		int	head_dim,
		float	scale,
		int	cache_max_seq_len,
		int	cache_num_kv_heads,
		cudaStream_t	stream
	)

Definition at line 1178 of file cuda_kernels.cu.

                                                          {
    
    if (num_heads == 0 || current_seq_len == 0 || head_dim == 0) {
        Logger::warning("[ATTENTION_SELECTIVE_DEQUANT] Invalid parameters: num_heads=" + 
                       std::to_string(num_heads) + ", current_seq_len=" + std::to_string(current_seq_len) + 
                       ", head_dim=" + std::to_string(head_dim));
        return;
    }
    
    // Grid: num_heads blocks, each block handles one attention head
    // Block: head_dim threads, each thread handles one output dimension  
    dim3 grid_dim(num_heads);
    dim3 block_dim(head_dim);
    
    attention_selective_dequant_kernel<<<grid_dim, block_dim, 0, stream>>>(
        Q_current_dev, K_quantized_cache_base, V_quantized_cache_base,
        K_scales_cache_base, V_scales_cache_base,
        selective_k_dequant_buffer, selective_v_dequant_buffer,
        out_dev, num_heads, current_seq_len, head_dim, scale,
        cache_max_seq_len, cache_num_kv_heads);
    
    gpuErrchk(cudaGetLastError());
}

References j, scale, and Logger::warning().

attention_decode_at_cuda()

void attention_decode_at_cuda	(	const float *	Q_current_dev,
		const float *	K_layer_cache_base,
		const float *	V_layer_cache_base,
		float *	out_dev,
		int	num_heads,
		int	seq_len,
		int	head_dim,
		float	scale,
		int	cache_max_seq_len,
		int	cache_num_kv_heads,
		cudaStream_t	stream
	)

Definition at line 3600 of file cuda_kernels.cu.

                                                   {
  if (!decode_attn_device(Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev, num_heads, seq_len, head_dim,
                         scale, cache_max_seq_len, cache_num_kv_heads, nullptr, stream))
    attention_cuda_optimized(Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev, num_heads, seq_len,
                             head_dim, scale, cache_max_seq_len, cache_num_kv_heads, stream);
}

References attention_cuda_optimized(), j, and scale.

attention_devpos_cuda()

void attention_devpos_cuda	(	const float *	Q_current_dev,
		const float *	K_layer_cache_base,
		const float *	V_layer_cache_base,
		float *	out_dev,
		int	num_heads,
		int	head_dim,
		float	scale,
		int	cache_max_seq_len,
		int	cache_num_kv_heads,
		const int *	step_dev,
		cudaStream_t	stream
	)

Definition at line 3610 of file cuda_kernels.cu.

                                                {
  if (decode_attn_device(Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev, num_heads,
                        /*seq_len=*/0, head_dim, scale, cache_max_seq_len, cache_num_kv_heads, step_dev, stream))
    return;
  const size_t shared_mem_size = (cache_max_seq_len + 1) * sizeof(float);
  if (ember::deterministic()) {
    attention_optimized_det_kernel<<<num_heads, 256, shared_mem_size, stream>>>(
        Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev, /*current_seq_len=*/0, head_dim, scale,
        cache_num_kv_heads, num_heads, step_dev);
    gpuErrchk(cudaGetLastError());
    return;
  }
  attention_optimized_kernel<<<num_heads, 256, shared_mem_size, stream>>>(
      Q_current_dev, K_layer_cache_base, V_layer_cache_base, out_dev, /*current_seq_len=*/0, head_dim, scale,
      cache_num_kv_heads, num_heads, step_dev);
  gpuErrchk(cudaGetLastError());
}

References ember::deterministic(), j, and scale.

attention_kernel()

__global__ void attention_kernel	(	const float *	Q_current,
		const float *	K_layer_cache_base,
		const float *	V_layer_cache_base,
		float *	out,
		int	current_seq_len,
		int	head_dim,
		float	scale,
		int	cache_num_kv_heads,
		int	num_q_heads
	)

Definition at line 981 of file cuda_kernels.cu.

                                                                          {
 
 
  int q_head_idx = blockIdx.x;
  int d_idx = threadIdx.x; // dimension index for this thread
 
  if (q_head_idx >= num_q_heads) return; // Should not happen with correct launch
 
  // Determine the corresponding KV head for this Q head (for GQA/MQA)
  int heads_per_kv = num_q_heads / cache_num_kv_heads;
  int kv_head_idx = q_head_idx / heads_per_kv;
 
  // Shared memory for scores (one score per k_pos) and for dot product reduction
  extern __shared__ float shared_data[];
  float* scores = shared_data; // First part of shared_data: size current_seq_len
  float* dot_product_terms = &shared_data[current_seq_len]; // Second part: size head_dim (blockDim.x)
 
  // Pointer to the current query head's data
  const float* q_head_ptr = Q_current + q_head_idx * head_dim;
 
  for (int k_pos = 0; k_pos < current_seq_len; ++k_pos) {
    // Pointer to the k_pos-th key vector for the relevant kv_head
    size_t k_vec_offset = (size_t)k_pos * cache_num_kv_heads * head_dim +
                          (size_t)kv_head_idx * head_dim;
    const float* k_vec_ptr = K_layer_cache_base + k_vec_offset;
 
    dot_product_terms[d_idx] = q_head_ptr[d_idx] * k_vec_ptr[d_idx];
    __syncthreads(); // Ensure all terms are written before reduction
 
    // Reduction in shared memory (e.g., by thread 0 or tree-based)
    if (d_idx == 0) {
      float current_dot = 0.0f;
      for (int r = 0; r < head_dim; ++r) {
        current_dot += dot_product_terms[r];
      }
      scores[k_pos] = current_dot * scale;
    }
    __syncthreads(); 
  }
 
  float thread_max_score = -INFINITY;
  for (int i = d_idx; i < current_seq_len; i += blockDim.x) { // blockDim.x is head_dim
    if (scores[i] > thread_max_score) {
      thread_max_score = scores[i];
    }
  }
  dot_product_terms[d_idx] = thread_max_score; // Repurpose dot_product_terms shared memory
  __syncthreads();
 
  float block_max_score = -INFINITY;
  if (d_idx == 0) {
    for (int r = 0; r < head_dim; ++r) { // head_dim might be > current_seq_len for first few tokens
      if (dot_product_terms[r] > block_max_score) { // Check if dot_product_terms[r] was validly written
        block_max_score = dot_product_terms[r];
      }
    }
  }
  __syncthreads();
  if(d_idx == 0) dot_product_terms[0] = block_max_score; // Store it.
  __syncthreads();
  block_max_score = dot_product_terms[0]; // All threads read it.
 
 
  // Calculate exp scores and sum
  float thread_exp_sum = 0.0f;
  for (int i = d_idx; i < current_seq_len; i += blockDim.x) {
    float prob = expf(scores[i] - block_max_score);
    scores[i] = prob; // Update scores in-place with exp(val - max)
    thread_exp_sum += prob;
  }
  // Reduce thread_exp_sum across the block
  dot_product_terms[d_idx] = thread_exp_sum; // Repurpose again
  __syncthreads();
 
  float block_exp_sum = 0.0f;
  if (d_idx == 0) {
    for (int r = 0; r < head_dim; ++r) { // Similar concern as max_score reduction range
      block_exp_sum += dot_product_terms[r];
    }
  }
  __syncthreads();
  // Store and reload block_exp_sum for all threads
  if(d_idx == 0) dot_product_terms[0] = block_exp_sum;
  __syncthreads();
  block_exp_sum = dot_product_terms[0];
 
 
  float inv_sum = 1.0f / (block_exp_sum + 1e-9f); // Add epsilon for stability
 
  // Normalize scores in shared memory
  for (int i = d_idx; i < current_seq_len; i += blockDim.x) {
    scores[i] *= inv_sum;
  }
  __syncthreads(); // Ensure all scores are normalized before weighted sum
 
  double weighted_val_d = 0.0; // Use double for accumulation
  for (int k_pos = 0; k_pos < current_seq_len; ++k_pos) {
    size_t v_vec_offset = (size_t)k_pos * cache_num_kv_heads * head_dim +
                          (size_t)kv_head_idx * head_dim;
    const float* v_vec_ptr = V_layer_cache_base + v_vec_offset;
    
    weighted_val_d += static_cast<double>(scores[k_pos]) * static_cast<double>(v_vec_ptr[d_idx]);
  }
  
  // Write the final result for this thread's dimension d_idx
  out[q_head_idx * head_dim + d_idx] = static_cast<float>(weighted_val_d);
}

References j, out, and scale.

attention_optimized_det_kernel()

__global__ void attention_optimized_det_kernel	(	const float *__restrict__	Q_current,
		const float *__restrict__	K_layer_cache_base,
		const float *__restrict__	V_layer_cache_base,
		float *__restrict__	out,
		int	current_seq_len,
		int	head_dim,
		float	scale,
		int	cache_num_kv_heads,
		int	num_q_heads,
		const int *	step_dev
	)

Definition at line 2987 of file cuda_kernels.cu.

                                                                    {
    if (step_dev) current_seq_len = step_dev[1] + 1;
 
    const int head_idx = blockIdx.x;
    const int tid_in_block = threadIdx.x;
    const int block_dim_x = blockDim.x;
 
    if (head_idx >= num_q_heads) return;
 
    const int heads_per_kv = num_q_heads / cache_num_kv_heads;
    const int kv_head_idx = head_idx / heads_per_kv;
 
    extern __shared__ float s_data[];
    float* s_scores = s_data;                 // [current_seq_len]
    __shared__ float s_warp_val[32];          // one partial per warp
    __shared__ float s_block_val;             // the block-wide max, then sum
    const int num_warps = (block_dim_x + warpSize - 1) / warpSize;
 
    const float* q_vec = Q_current + head_idx * head_dim;
 
    // --- 1. Q.K^T scores ---
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        const float* k_vec = K_layer_cache_base + ((size_t)k_idx * cache_num_kv_heads + kv_head_idx) * head_dim;
        float score = 0.0f;
        for (int d = 0; d < head_dim; ++d) {
            score += q_vec[d] * k_vec[d];
        }
        s_scores[k_idx] = score * scale;
    }
    __syncthreads();
 
    // --- 2a. max score: warp maxima, then in warp order ---
    float thread_max_score = -INFINITY;
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        thread_max_score = fmaxf(thread_max_score, s_scores[k_idx]);
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) thread_max_score = fmaxf(thread_max_score, __shfl_down_sync(0xFFFFFFFF, thread_max_score, offset));
    if ((tid_in_block % warpSize) == 0) s_warp_val[tid_in_block / warpSize] = thread_max_score;
    __syncthreads();
    if (tid_in_block == 0) {
        float m = -INFINITY;
        for (int w = 0; w < num_warps; ++w) m = fmaxf(s_warp_val[w], m);
        s_block_val = m;
    }
    __syncthreads();
    float max_score = s_block_val;
    __syncthreads();  // every thread has read the max before the slots are reused
 
    // --- 2b. exp and their sum: warp sums, then added in warp order ---
    float thread_sum_exp = 0.0f;
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        float val = expf(s_scores[k_idx] - max_score);
        s_scores[k_idx] = val;
        thread_sum_exp += val;
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) thread_sum_exp += __shfl_down_sync(0xFFFFFFFF, thread_sum_exp, offset);
    if ((tid_in_block % warpSize) == 0) s_warp_val[tid_in_block / warpSize] = thread_sum_exp;
    __syncthreads();
    if (tid_in_block == 0) {
        float s = 0.0f;
        for (int w = 0; w < num_warps; ++w) s = s_warp_val[w] + s;
        s_block_val = s;
    }
    __syncthreads();
    float sum_exp = s_block_val;
    float inv_sum_exp = (sum_exp == 0.0f || sum_exp < 1e-9f) ? 0.0f : 1.0f / sum_exp;
 
    // --- 2c. normalise ---
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        s_scores[k_idx] *= inv_sum_exp;
    }
    __syncthreads();
 
    // --- 3. P V ---
    float* out_for_this_head = out + head_idx * head_dim;
    for (int d_out = tid_in_block; d_out < head_dim; d_out += block_dim_x) {
        float weighted_val = 0.0f;
        for (int k_idx = 0; k_idx < current_seq_len; ++k_idx) {
            const float* v_vec = V_layer_cache_base + ((size_t)k_idx * cache_num_kv_heads + kv_head_idx) * head_dim;
            weighted_val += s_scores[k_idx] * v_vec[d_out];
        }
        out_for_this_head[d_out] = weighted_val;
    }
}

References j, out, and scale.

attention_optimized_kernel()

__global__ void attention_optimized_kernel	(	const float *__restrict__	Q_current,
		const float *__restrict__	K_layer_cache_base,
		const float *__restrict__	V_layer_cache_base,
		float *__restrict__	out,
		int	current_seq_len,
		int	head_dim,
		float	scale,
		int	cache_num_kv_heads,
		int	num_q_heads,
		const int *	step_dev
	)

Corrected and optimized attention kernel.

Uses shared memory for scores and performs a full softmax operation correctly. Each block computes attention for one head.

Definition at line 2839 of file cuda_kernels.cu.

                                                               {
    // Graph decode: the length comes from device memory (step_dev[1] = pos).
    if (step_dev) current_seq_len = step_dev[1] + 1;
 
    const int head_idx = blockIdx.x; // Current Q head index this block is working on
    const int tid_in_block = threadIdx.x;
    const int block_dim_x = blockDim.x; // Threads per block (e.g., 256)
 
    if (head_idx >= num_q_heads) return;
 
    // Map current Q head to its corresponding KV cache head index (for GQA/MQA)
    const int heads_per_kv = num_q_heads / cache_num_kv_heads;
    const int kv_head_idx = head_idx / heads_per_kv;
 
    // Dynamic shared memory. Layout:
    // s_scores[current_seq_len]: to store Q.K^T results
    // s_reduction_val[1]: for block-wide max and sum_exp
    extern __shared__ float s_data[];
    float* s_scores = s_data;
    float* s_reduction_val = &s_data[current_seq_len]; // Single float for block-wide reduction
 
    // Pointer to the current Q vector for this specific head
    // Assuming Q_current is already offset for the specific token if in batch generation.
    // Here, Q_current is expected to be just [head_idx * head_dim] for the current token's Q vector.
    const float* q_vec = Q_current + head_idx * head_dim;
 
    // --- 1. Compute Q.K^T scores and store in s_scores ---
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        // Pointer to the k_idx-th K vector in the cache for the relevant kv_head
        const float* k_vec = K_layer_cache_base + ((size_t)k_idx * cache_num_kv_heads + kv_head_idx) * head_dim;
        float score = 0.0f;
        for (int d = 0; d < head_dim; ++d) {
            score += q_vec[d] * k_vec[d];
        }
        s_scores[k_idx] = score * scale;
    }
    __syncthreads(); // Ensure all scores are computed and stored in s_scores
 
    // --- 2. Softmax calculation on s_scores ---
    // a. Find max score in s_scores (block-wide reduction)
    if (tid_in_block == 0) *s_reduction_val = -INFINITY;
    __syncthreads();
 
    float thread_max_score = -INFINITY;
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        thread_max_score = fmaxf(thread_max_score, s_scores[k_idx]);
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) thread_max_score = fmaxf(thread_max_score, __shfl_down_sync(0xFFFFFFFF, thread_max_score, offset));
    if ((tid_in_block % warpSize) == 0) atomicMaxFloat(s_reduction_val, thread_max_score);
    __syncthreads();
    float max_score = *s_reduction_val;
    __syncthreads();  // every thread has read the max before thread 0 reuses the slot
 
    // b. Subtract max, compute exp, and sum exp (block-wide reduction for sum)
    if (tid_in_block == 0) *s_reduction_val = 0.0f;
    __syncthreads();
 
    float thread_sum_exp = 0.0f;
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        float val = expf(s_scores[k_idx] - max_score);
        s_scores[k_idx] = val; // Store exp(score - max)
        thread_sum_exp += val;
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) thread_sum_exp += __shfl_down_sync(0xFFFFFFFF, thread_sum_exp, offset);
    if ((tid_in_block % warpSize) == 0) atomicAddFloat(s_reduction_val, thread_sum_exp);
    __syncthreads();
    float sum_exp = *s_reduction_val;
    float inv_sum_exp = (sum_exp == 0.0f || sum_exp < 1e-9f) ? 0.0f : 1.0f / sum_exp; // add small epsilon for stability
 
    // c. Normalize s_scores (they now contain softmax probabilities)
    for (int k_idx = tid_in_block; k_idx < current_seq_len; k_idx += block_dim_x) {
        s_scores[k_idx] *= inv_sum_exp;
    }
    __syncthreads();
 
    // --- 3. Multiply by V and accumulate to output ---
    // Output for this head: out_for_this_head[d_out]
    float* out_for_this_head = out + head_idx * head_dim;
 
    for (int d_out = tid_in_block; d_out < head_dim; d_out += block_dim_x) {
        float weighted_val = 0.0f;
        for (int k_idx = 0; k_idx < current_seq_len; ++k_idx) { // All threads iterate full k_idx to sum contributions
            const float* v_vec = V_layer_cache_base + ((size_t)k_idx * cache_num_kv_heads + kv_head_idx) * head_dim;
            weighted_val += s_scores[k_idx] * v_vec[d_out];
        }
        out_for_this_head[d_out] = weighted_val; // Each thread writes to its assigned output dimension part
    }
}

References atomicAddFloat(), atomicMaxFloat(), j, out, and scale.

attention_selective_dequant_kernel()

__global__ void attention_selective_dequant_kernel	(	const float *__restrict__	Q_current_dev,
		const int8_t *__restrict__	K_quantized_cache_base,
		const int8_t *__restrict__	V_quantized_cache_base,
		const float *__restrict__	K_scales_cache_base,
		const float *__restrict__	V_scales_cache_base,
		float *__restrict__	selective_k_dequant_buffer,
		float *__restrict__	selective_v_dequant_buffer,
		float *__restrict__	out_dev,
		int	num_heads,
		int	current_seq_len,
		int	head_dim,
		float	scale,
		int	cache_max_seq_len,
		int	cache_num_kv_heads
	)

Definition at line 1112 of file cuda_kernels.cu.

                                                                {
    
    int head_idx = blockIdx.x;
    int dim_idx = threadIdx.x;
    
    if (head_idx >= num_heads || dim_idx >= head_dim) return;
    
    // Map query head to corresponding KV head (for GQA/MQA)
    int kv_heads_per_q_group = num_heads / cache_num_kv_heads;
    int kv_head_idx = head_idx / kv_heads_per_q_group;
    
    // Get current query vector for this head and dimension
    float q_val = Q_current_dev[head_idx * head_dim + dim_idx];
    
    float attn_out = 0.0f;
    
    // Selective dequantization and attention computation
    for (int t = 0; t < current_seq_len; ++t) {
        // Calculate offset for this token and KV head in the quantized cache
        size_t cache_offset_base = (size_t)t * cache_num_kv_heads * head_dim + (size_t)kv_head_idx * head_dim;
        size_t scale_offset = (size_t)t * cache_num_kv_heads + kv_head_idx;
        
        // Dequantize only the specific elements we need for this head/dimension
        if (dim_idx == 0) {
            // Only thread 0 in each block dequantizes the entire head for this token
            float k_scale = K_scales_cache_base[scale_offset];
            float v_scale = V_scales_cache_base[scale_offset];
            
            for (int d = 0; d < head_dim; ++d) {
                // Dequantize K values for this token/head
                int8_t k_quant = K_quantized_cache_base[cache_offset_base + d];
                selective_k_dequant_buffer[t * head_dim + d] = k_scale * (float)k_quant;
                
                // Dequantize V values for this token/head  
                int8_t v_quant = V_quantized_cache_base[cache_offset_base + d];
                selective_v_dequant_buffer[t * head_dim + d] = v_scale * (float)v_quant;
            }
        }
        
        __syncthreads(); // Ensure dequantization is complete before using values
        
        // Compute attention score for this token
        float k_val = selective_k_dequant_buffer[t * head_dim + dim_idx];
        float score = q_val * k_val * scale;
        
        // Apply attention score to value
        float v_val = selective_v_dequant_buffer[t * head_dim + dim_idx];
        attn_out += score * v_val;
        
        __syncthreads(); // Ensure all threads are done before next iteration
    }
    
    // Write output
    out_dev[head_idx * head_dim + dim_idx] = attn_out;
}

References j, and scale.

bf16_to_float32_device()

__device__ float bf16_to_float32_device ( uint16_t  bf16_raw )

inline

Converts BF16 to FP32 on the device.

This device function converts Brain Floating Point (BF16) values to FP32. It uses CUDA's native BF16 support on architectures >= 800, and falls back to a manual conversion on older architectures.

Parameters

bf16_raw

Raw BF16 value as uint16_t

Returns

Converted FP32 value

Definition at line 102 of file cuda_kernels.cu.

                                                                  {
#if __CUDA_ARCH__ >= 800
  __nv_bfloat16 bf16_val;
 
  memcpy(&bf16_val, &bf16_raw, sizeof(uint16_t));
  return __bfloat162float(bf16_val);
#else
 
  unsigned int bits = ((unsigned int)bf16_raw) << 16;
  float result;
  memcpy(&result, &bits, sizeof(float));
  return result;
#endif
}

References j.

Referenced by convert_bf16_to_fp32_kernel(), and lookup_embedding_kernel().

convert_bf16_to_fp32_cuda()

void convert_bf16_to_fp32_cuda	(	const uint16_t *	bf16_in_dev,
		float *	fp32_out_dev,
		size_t	n_elements,
		cudaStream_t	stream
	)

Definition at line 3515 of file cuda_kernels.cu.

                                                                       {
    if (n_elements == 0) return;
    
    const int threads = 256;
    const int blocks = (n_elements + threads - 1) / threads;
    
    convert_bf16_to_fp32_kernel<<<blocks, threads, 0, stream>>>(bf16_in_dev, fp32_out_dev, n_elements);
    CUDA_CHECK(cudaGetLastError());
}

References j.

Referenced by gemm_f32_to_bf16_f32_cuda().

convert_bf16_to_fp32_kernel()

__global__ void convert_bf16_to_fp32_kernel	(	const uint16_t *__restrict__	bf16_in,
		float *__restrict__	fp32_out,
		size_t	n_elements
	)

Definition at line 1498 of file cuda_kernels.cu.

                                                               {
  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < n_elements) {
    fp32_out[idx] = bf16_to_float32_device(bf16_in[idx]);
  }
}

References bf16_to_float32_device(), and j.

convert_fp32_to_bf16_cuda()

void convert_fp32_to_bf16_cuda	(	const float *	fp32_in_dev,
		uint16_t *	bf16_out_dev,
		size_t	n_elements,
		cudaStream_t	stream
	)

Definition at line 3357 of file cuda_kernels.cu.

                                                                       {
    if (n_elements == 0) return;
    
    const int threads = 256;
    const int blocks = (n_elements + threads - 1) / threads;
    
    convert_fp32_to_bf16_kernel<<<blocks, threads, 0, stream>>>(fp32_in_dev, bf16_out_dev, n_elements);
    CUDA_CHECK(cudaGetLastError());
}

References j.

Referenced by gemm_f32_to_bf16_f32_cuda().

convert_fp32_to_bf16_kernel()

__global__ void convert_fp32_to_bf16_kernel	(	const float *__restrict__	fp32_in,
		uint16_t *__restrict__	bf16_out,
		size_t	n_elements
	)

Definition at line 3347 of file cuda_kernels.cu.

                                                               {
    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n_elements) {
        bf16_out[idx] = float32_to_bf16_device(fp32_in[idx]);
    }
}

References float32_to_bf16_device(), and j.

decode_attention_kind()

int decode_attention_kind

(

)

Definition at line 3297 of file cuda_kernels.cu.

                            {
    static const int kind = [] {
        const char* e = std::getenv("EMBER_DECODE_ATTN");
        const std::string s = e ? e : "";
        return s == "legacy" ? 2 : s == "split" ? 1 : 0;
    }();
    return kind;
}

References j, and kind.

decode_attention_tiles()

void decode_attention_tiles	(	int *	min_chunk,
		int *	max_chunks
	)

Definition at line 3306 of file cuda_kernels.cu.

                                                             {
    static const std::pair<int, int> t = [] {
        std::pair<int, int> d{ember_tiled::kTiledMinChunk, ember_tiled::kTiledMaxChunks};
        const char* e = std::getenv("EMBER_DECODE_ATTN_TILE");
        int a = 0, b = 0;
        if (e && std::sscanf(e, "%d,%d", &a, &b) == 2 && a >= 1 && b >= 1 && b <= ember_tiled::kTiledMaxChunks)
            d = {a, b};
        return d;
    }();
    *min_chunk = t.first;
    *max_chunks = t.second;
}

References a, j, ember_tiled::kTiledMaxChunks, and ember_tiled::kTiledMinChunk.

decode_embed_cuda()

void decode_embed_cuda	(	const float *	table_dev,
		float *	out_dev,
		const int *	step_dev,
		int	hidden_size,
		int	vocab_size,
		cudaStream_t	stream
	)

Definition at line 3541 of file cuda_kernels.cu.

                                                            {
  decode_embed_kernel<<<(hidden_size + 255) / 256, 256, 0, stream>>>(table_dev, out_dev, step_dev, hidden_size,
                                                                     vocab_size);
  gpuErrchk(cudaGetLastError());
}

References decode_embed_kernel(), and j.

decode_embed_kernel()

__global__ void decode_embed_kernel	(	const float *__restrict__	table,
		float *__restrict__	out,
		const int *__restrict__	step_dev,
		int	hidden_size,
		int	vocab_size
	)

Definition at line 3533 of file cuda_kernels.cu.

                                                                                                       {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= hidden_size) return;
  const int tok = step_dev[0];
  out[i] = (tok >= 0 && tok < vocab_size) ? table[(size_t)tok * hidden_size + i] : 0.0f;
}

References j, out, and tok.

Referenced by decode_embed_cuda().

float32_to_bf16_device()

__device__ uint16_t float32_to_bf16_device ( float  fp32_val )

inline

Definition at line 118 of file cuda_kernels.cu.

                                                                  {
#if __CUDA_ARCH__ >= 800
    __nv_bfloat16 bf16_val = __float2bfloat16(fp32_val);
    uint16_t bf16_raw;
    memcpy(&bf16_raw, &bf16_val, sizeof(uint16_t));
    return bf16_raw;
#else
    // Manual conversion for older architectures
    unsigned int bits;
    memcpy(&bits, &fp32_val, sizeof(float));
    // Extract upper 16 bits for BF16
    return (uint16_t)(bits >> 16);
#endif
}

References j.

Referenced by convert_fp32_to_bf16_kernel().

gemm_bf16_bf16_cuda()

void gemm_bf16_bf16_cuda	(	cublasHandle_t	handle,
		bool	transa_user,
		bool	transb_user,
		int	m_user,
		int	n_user,
		int	k_user,
		const float *	alpha_user,
		const uint16_t *	A_bf16_user,
		int	lda_user,
		const uint16_t *	B_bf16_user,
		int	ldb_user,
		const float *	beta_user,
		uint16_t *	C_bf16_user,
		int	ldc_user,
		cudaStream_t	stream
	)

Definition at line 3369 of file cuda_kernels.cu.

                                              {
    
    cublasStatus_t status = cublasSetStream(handle, stream);
    if (status != CUBLAS_STATUS_SUCCESS) {
        Logger::error("cublasSetStream failed in gemm_bf16_bf16_cuda with error: " + std::to_string(status));
        throw std::runtime_error("cublasSetStream failed");
    }
 
    // Log parameters for debugging
    Logger::info("[GEMM_BF16_BF16] M=" + std::to_string(m_user) + ", N=" + std::to_string(n_user) + 
                 ", K=" + std::to_string(k_user) + ", LDA=" + std::to_string(lda_user) + 
                 ", LDB=" + std::to_string(ldb_user) + ", LDC=" + std::to_string(ldc_user));
 
    cublasOperation_t opA_cublas, opB_cublas;
    int M_cublas, N_cublas, K_cublas;
    const uint16_t *A_cublas_ptr, *B_cublas_ptr;
    int LDA_cublas, LDB_cublas;
 
    // Apply the same column-major transformation as gemm_f32_f32_cuda
    M_cublas = n_user;
    N_cublas = m_user;
    K_cublas = k_user;
 
    A_cublas_ptr = B_bf16_user;    // B_user data is the first matrix for cuBLAS call
    LDA_cublas = (transb_user) ? k_user : ldb_user;
 
    B_cublas_ptr = A_bf16_user;    // A_user data is the second matrix for cuBLAS call
    LDB_cublas = (transa_user) ? m_user : lda_user;
 
    // Determine cuBLAS operations
    if (!transa_user && !transb_user) {
        opA_cublas = CUBLAS_OP_N; // for B_user
        opB_cublas = CUBLAS_OP_N; // for A_user
    } else if (transa_user && !transb_user) {
        opA_cublas = CUBLAS_OP_N; // for B_user
        opB_cublas = CUBLAS_OP_T; // for A_user (A^T)
    } else if (!transa_user && transb_user) {
        opA_cublas = CUBLAS_OP_T; // for B_user (B^T)
        opB_cublas = CUBLAS_OP_N; // for A_user
    } else { // transa_user && transb_user
        opA_cublas = CUBLAS_OP_T; // for B_user (B^T)
        opB_cublas = CUBLAS_OP_T; // for A_user (A^T)
    }
    
    int LDC_cublas = ldc_user;
 
    // Use cublasGemmEx with BF16 Tensor Cores
    status = cublasGemmEx(handle, 
                          opA_cublas, opB_cublas, 
                          M_cublas, N_cublas, K_cublas, 
                          alpha_user, 
                          A_cublas_ptr, CUDA_R_16BF, LDA_cublas, 
                          B_cublas_ptr, CUDA_R_16BF, LDB_cublas, 
                          beta_user, 
                          C_bf16_user, CUDA_R_16BF, LDC_cublas, 
                          CUBLAS_COMPUTE_32F, // Use FP32 for accumulation for stability
                          CUBLAS_GEMM_DEFAULT_TENSOR_OP); // Enable Tensor Cores
 
    if (status != CUBLAS_STATUS_SUCCESS) {
        Logger::error("cublasGemmEx BF16 failed in gemm_bf16_bf16_cuda with error: " + std::to_string(status));
        Logger::error("GEMM params: transa=" + std::to_string(transa_user) + " transb=" + std::to_string(transb_user) + 
                     " m=" + std::to_string(m_user) + " n=" + std::to_string(n_user) + " k=" + std::to_string(k_user));
        throw std::runtime_error("cublasGemmEx BF16 failed");
    }
    
    Logger::info("[GEMM_BF16_BF16] SUCCESS with Tensor Cores");
}

References Logger::error(), Logger::info(), and j.

gemm_f32_to_bf16_f32_cuda()

void gemm_f32_to_bf16_f32_cuda	(	cublasHandle_t	handle,
		bool	transa_user,
		bool	transb_user,
		int	m_user,
		int	n_user,
		int	k_user,
		const float *	alpha_user,
		const float *	A_f32_user,
		int	lda_user,
		const uint16_t *	B_bf16_user,
		int	ldb_user,
		const float *	beta_user,
		float *	C_f32_user,
		int	ldc_user,
		cudaStream_t	stream
	)

Definition at line 3446 of file cuda_kernels.cu.

                                                    {
    
    cublasStatus_t stream_status = cublasSetStream(handle, stream);
    if (stream_status != CUBLAS_STATUS_SUCCESS) {
        Logger::error("cublasSetStream failed in gemm_f32_to_bf16_f32_cuda with error: " + std::to_string(stream_status));
        throw std::runtime_error("cublasSetStream failed");
    }
    
    // Convert FP32 input A to BF16
    size_t A_size = (size_t)m_user * k_user;
    uint16_t* A_bf16_temp;
    CUDA_CHECK(cudaMalloc(&A_bf16_temp, A_size * sizeof(uint16_t)));
    convert_fp32_to_bf16_cuda(A_f32_user, A_bf16_temp, A_size, stream);
    
    // Allocate BF16 output buffer
    size_t C_size = (size_t)m_user * n_user;
    uint16_t* C_bf16_temp;
    CUDA_CHECK(cudaMalloc(&C_bf16_temp, C_size * sizeof(uint16_t)));
    
    // Convert alpha and beta to BF16 for host-side pointers
    uint16_t alpha_bf16_host = float32_to_bf16_host(*alpha_user);
    uint16_t beta_bf16_host = float32_to_bf16_host(*beta_user);
    
    // DIRECT parameter mapping for cublasGemmEx
    cublasOperation_t opA = transa_user ? CUBLAS_OP_T : CUBLAS_OP_N;
    cublasOperation_t opB = transb_user ? CUBLAS_OP_T : CUBLAS_OP_N;
    
    // A is A_bf16_temp (converted from A_f32_user)
    // B is B_bf16_user (already BF16)
    // C is C_bf16_temp (will be converted back to C_f32_user)
    cublasStatus_t status = cublasGemmEx(handle,
        opA, opB,                
        m_user, n_user, k_user, 
        &alpha_bf16_host, // Pointer to BF16 host alpha
        A_bf16_temp, CUDA_R_16BF, lda_user,
        B_bf16_user, CUDA_R_16BF, ldb_user,
        &beta_bf16_host,  // Pointer to BF16 host beta
        C_bf16_temp, CUDA_R_16BF, ldc_user,
        CUBLAS_COMPUTE_16F, // Compute in BF16 (A_16*B_16 -> C_16)
        CUBLAS_GEMM_DEFAULT_TENSOR_OP
    );
    
    if (status != CUBLAS_STATUS_SUCCESS) {
        Logger::error("cublasGemmEx BF16 (A_bf16*B_bf16->C_bf16) failed in gemm_f32_to_bf16_f32_cuda with error: " + std::to_string(status));
        Logger::error("GEMM params: transA=" + std::to_string(opA) + " transB=" + std::to_string(opB) + 
                     " m=" + std::to_string(m_user) + " n=" + std::to_string(n_user) + " k=" + std::to_string(k_user) + 
                     " lda=" + std::to_string(lda_user) + " ldb=" + std::to_string(ldb_user) + " ldc=" + std::to_string(ldc_user));
        CUDA_CHECK(cudaFree(A_bf16_temp));
        CUDA_CHECK(cudaFree(C_bf16_temp));
        throw std::runtime_error("BF16 Tensor Core GEMM failed");
    }
    
    // Convert BF16 result back to FP32
    convert_bf16_to_fp32_cuda(C_bf16_temp, C_f32_user, C_size, stream);
    
    // Cleanup temporary buffers
    CUDA_CHECK(cudaFree(A_bf16_temp));
    CUDA_CHECK(cudaFree(C_bf16_temp));
}

References convert_bf16_to_fp32_cuda(), convert_fp32_to_bf16_cuda(), Logger::error(), and j.

gemm_rm_f32_cuda()

void gemm_rm_f32_cuda	(	cublasHandle_t	handle,
		const float *	X,
		const float *	W,
		float *	Y,
		int	T,
		int	O,
		int	I,
		cudaStream_t	stream
	)

Definition at line 3759 of file cuda_kernels.cu.

                                           {
  cublasSetStream(handle, stream);
  moe_gemm_rm(handle, X, W, Y, T, O, I);
}

References j.

kernel_debug_copies()

static bool kernel_debug_copies ( )

static

Definition at line 64 of file cuda_kernels.cu.

                                  {
  static const bool on = std::getenv("EMBER_KERNEL_DEBUG") != nullptr;
  return on;
}

References j.

Referenced by add_residual_batch_cuda(), rope_batch_cuda(), and swiglu_batch_cuda().

kv_write_devpos_cuda()

void kv_write_devpos_cuda	(	float *	k_cache_dev,
		float *	v_cache_dev,
		const float *	k_dev,
		const float *	v_dev,
		const int *	step_dev,
		int	row,
		cudaStream_t	stream
	)

Definition at line 3593 of file cuda_kernels.cu.

                                                                             {
  kv_write_devpos_kernel<<<(row + 255) / 256, 256, 0, stream>>>(k_cache_dev, v_cache_dev, k_dev, v_dev, step_dev,
                                                               row);
  gpuErrchk(cudaGetLastError());
}

References j, and kv_write_devpos_kernel().

kv_write_devpos_kernel()

__global__ void kv_write_devpos_kernel	(	float *	k_cache,
		float *	v_cache,
		const float *__restrict__	k,
		const float *__restrict__	v,
		const int *	step_dev,
		int	row
	)

Definition at line 3584 of file cuda_kernels.cu.

                                                                                                  {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= row) return;
  const size_t off = (size_t)step_dev[1] * row + i;
  k_cache[off] = k[i];
  v_cache[off] = v[i];
}

References j, k, off, and v.

Referenced by kv_write_devpos_cuda().

lookup_embedding_cuda()

void lookup_embedding_cuda	(	const void *	table_dev,
		float *	output_dev,
		int	token_id,
		int	hidden_size,
		int	vocab_size,
		bool	is_bf16,
		cudaStream_t	stream
	)

Definition at line 1474 of file cuda_kernels.cu.

                                                              {
  if (!table_dev || !output_dev) {
    Logger::error("lookup_embedding_cuda: Received null device pointer(s).");
 
    return;
  }
  if (hidden_size <= 0 || vocab_size <= 0) {
    Logger::error("lookup_embedding_cuda: Invalid hidden_size or vocab_size.");
    return;
  }
 
  int threads_per_block = 256;
 
  int blocks_per_grid =
      (hidden_size + threads_per_block - 1) / threads_per_block;
 
  lookup_embedding_kernel<<<blocks_per_grid, threads_per_block, 0, stream>>>(
      table_dev, output_dev, token_id, hidden_size, vocab_size, is_bf16);
 
  gpuErrchk(cudaGetLastError());
}

References Logger::error(), and j.

lookup_embedding_kernel()

__global__ void lookup_embedding_kernel	(	const void *__restrict__	table_dev,
		float *__restrict__	output_dev,
		int	token_id,
		int	hidden_size,
		int	vocab_size,
		bool	is_bf16
	)

Definition at line 1445 of file cuda_kernels.cu.

                                                                      {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
 
  if (idx >= hidden_size) {
    return;
  }
 
  if (token_id < 0 || token_id >= vocab_size) {
    output_dev[idx] = 0.0f;
    return;
  }
  size_t offset = (size_t)token_id * hidden_size + idx;
 
  if (is_bf16) {
    const uint16_t* table_bf16 = static_cast<const uint16_t*>(table_dev);
 
    uint16_t val_bf16 = table_bf16[offset];
 
    output_dev[idx] = bf16_to_float32_device(val_bf16);
  } else {
    const float* table_f32 = static_cast<const float*>(table_dev);
 
    output_dev[idx] = table_f32[offset];
  }
}

References bf16_to_float32_device(), and j.

matvec_bf16_f32_cuda()

void matvec_bf16_f32_cuda	(	cublasHandle_t	handle,
		const uint16_t *	mat_bf16_dev,
		const float *	vec_f32_dev,
		float *	out_f32_dev,
		int	rows,
		int	cols,
		bool	use_tensor_cores,
		cudaStream_t	stream
	)

Matrix-vector multiplication with BF16 matrix and FP32 vector/output on GPU.

This function performs matrix-vector multiplication (mat * vec = out) where the matrix is in BF16 format, and the vector and output are in FP32 format. It can use BF16 Tensor Cores via cublasGemmEx if supported and enabled. Otherwise, it falls back to converting the BF16 matrix to FP32 and then using cublasSgemm.

Matrix 'mat_bf16_dev' is assumed to be in row-major format with 'rows' rows and 'cols' columns. Vector 'vec_f32_dev' is assumed to have 'cols' elements. Output 'out_f32_dev' will have 'rows' elements.

Parameters

handle	cuBLAS handle
mat_bf16_dev	BF16 matrix (device pointer, row-major, rows x cols)
vec_f32_dev	FP32 vector (device pointer, cols elements)
out_f32_dev	FP32 output vector (device pointer, rows elements)
rows	Number of matrix rows
cols	Number of matrix columns
use_tensor_cores	Whether to attempt using BF16 Tensor Cores with cublasGemmEx
stream	CUDA stream (optional)

Definition at line 353 of file cuda_kernels.cu.

                                               {
  cublasStatus_t status;
  const float alpha = 1.0f;
  const float beta = 0.0f;
 
  // For C = A * B (row-major view) where A is mat_bf16_dev (rows x cols) and B is vec_f32_dev (cols x 1)
  // cuBLAS expects column-major. We use A^T for op(A) to simulate row-major A.
  // M_blas = rows (rows of C, rows of op(A))
  // N_blas = 1   (cols of C, cols of op(B))
  // K_blas = cols (cols of op(A), rows of op(B))
  //
  // If op(A) is CUBLAS_OP_T, then A is passed as K_blas x M_blas (cols x rows) column-major. LDA = K_blas = cols.
  // If op(B) is CUBLAS_OP_N, then B is passed as K_blas x N_blas (cols x 1) column-major. LDB = K_blas = cols.
  // C is M_blas x N_blas (rows x 1) column-major. LDC = M_blas = rows.
 
  // Disable Tensor Cores for matrix-vector operations (N=1) as they are not supported
  // by cublasGemmEx with mixed precision BF16->FP32 (confirmed by standalone test)
  if (use_tensor_cores) {
    Logger::info("[CUDA_KERNEL] matvec_bf16_f32_cuda: Tensor Cores requested, but disabling for matrix-vector (N=1) operation. Mixed precision BF16->FP32 with N=1 returns CUBLAS_STATUS_NOT_SUPPORTED in cublasGemmEx.");
    use_tensor_cores = false;
  }
 
  if (use_tensor_cores) {
    // Set the cuBLAS stream explicitly
    cublasStatus_t stream_status = cublasSetStream(handle, stream);
    if (stream_status != CUBLAS_STATUS_SUCCESS) {
      Logger::warning("[CUDA_KERNEL] cublasSetStream failed with status: " + std::to_string(stream_status) + ". Falling back to FP32 Sgemm.");
      use_tensor_cores = false;
    } else {
      // Logger::debug("[CUDA_KERNEL] matvec_bf16_f32_cuda attempting cublasGemmEx with Tensor Cores.");
      // Log all parameters for debugging
      Logger::info("[CUDA_KERNEL] cublasGemmEx parameters: M=" + std::to_string(rows) + 
                   ", N=1, K=" + std::to_string(cols) + 
                   ", LDA=" + std::to_string(cols) + 
                   ", LDB=" + std::to_string(cols) + 
                   ", LDC=" + std::to_string(rows) +
                   ", mat_ptr=" + std::to_string(reinterpret_cast<uintptr_t>(mat_bf16_dev)) +
                   ", vec_ptr=" + std::to_string(reinterpret_cast<uintptr_t>(vec_f32_dev)) +
                   ", out_ptr=" + std::to_string(reinterpret_cast<uintptr_t>(out_f32_dev)));
      
      // A is mat_bf16_dev (CUDA_R_16BF)
      // B is vec_f32_dev (CUDA_R_32F)
      // C is out_f32_dev (CUDA_R_32F)
      // Computation type CUDA_R_32F
      status = cublasGemmEx(handle,
                            CUBLAS_OP_T,        // Transpose A (mat_bf16_dev)
                            CUBLAS_OP_N,        // No transpose B (vec_f32_dev)
                            rows,               // M
                            1,                  // N
                            cols,               // K
                            &alpha,             // Alpha
                            mat_bf16_dev,       // A
                            CUDA_R_16BF,        // Atype
                            cols,               // LDA (K_blas for A^T)
                            vec_f32_dev,        // B
                            CUDA_R_32F,         // Btype
                            cols,               // LDB (K_blas for B_NoTrans)
                            &beta,              // Beta
                            out_f32_dev,        // C
                            CUDA_R_32F,         // Ctype
                            rows,               // LDC (M_blas)
                            CUBLAS_COMPUTE_32F, // Compute type (use FP32 for accumulation for stability with BF16 inputs)
                                                // For pure BF16 compute (if B was also BF16), could use CUBLAS_COMPUTE_32F_FAST_16BF
                                                // Or CUBLAS_GEMM_DEFAULT_TENSOR_OP for algo to let cuBLAS pick
                            CUBLAS_GEMM_DEFAULT_TENSOR_OP); // Algorithm
    }
    if (status != CUBLAS_STATUS_SUCCESS) {
      Logger::warning("[CUDA_KERNEL] cublasGemmEx with Tensor Cores failed for matvec_bf16_f32_cuda: " + std::to_string(status) + ". Falling back to FP32 Sgemm.");
      // Fallback to FP32 path
      use_tensor_cores = false; // Ensure fallback path is taken
    } else {
      // Logger::debug("[CUDA_KERNEL] cublasGemmEx with Tensor Cores SUCCEEDED for matvec_bf16_f32_cuda.");
      gpuErrchk(cudaGetLastError());
      return; // Success with Tensor Cores
    }
  }
 
  // Fallback path: Convert BF16 matrix to FP32 and use cublasSgemm
  if (!use_tensor_cores) {
    // Logger::debug("[CUDA_KERNEL] matvec_bf16_f32_cuda using fallback: BF16->FP32 conversion then Sgemm.");
    // Set the cuBLAS stream for fallback path too
    cublasStatus_t stream_status = cublasSetStream(handle, stream);
    if (stream_status != CUBLAS_STATUS_SUCCESS) {
      throw std::runtime_error("cublasSetStream failed in matvec_bf16_f32_cuda fallback: " + std::to_string(stream_status));
    }
    
    float* mat_fp32_dev = nullptr;
    size_t mat_size_elements = static_cast<size_t>(rows) * cols;
    gpuErrchk(cudaMalloc(&mat_fp32_dev, mat_size_elements * sizeof(float)));
 
    const int threads = 256;
    const int blocks = (mat_size_elements + threads - 1) / threads;
    convert_bf16_to_fp32_kernel<<<blocks, threads, 0, stream>>>(mat_bf16_dev, mat_fp32_dev, mat_size_elements);
    gpuErrchk(cudaGetLastError());
    // No explicit sync needed here if stream is managed correctly by caller or if Sgemm uses the same stream.
 
    status = cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,
                           rows, 1, cols,
                           &alpha,
                           mat_fp32_dev, cols,   // A (mat_fp32_dev), LDA = K_blas = cols
                           vec_f32_dev, cols,    // B (vec_f32_dev), LDB = K_blas = cols
                           &beta,
                           out_f32_dev, rows);   // C (out_f32_dev), LDC = M_blas = rows
    
    gpuErrchk(cudaFree(mat_fp32_dev));
    if (status != CUBLAS_STATUS_SUCCESS) {
      throw std::runtime_error("cuBLAS Sgemm failed in matvec_bf16_f32_cuda fallback: " + std::to_string(status));
    }
    gpuErrchk(cudaGetLastError());
  }
}

References Logger::info(), j, and Logger::warning().

matvec_f32_f32_cuda() [1/2]

void matvec_f32_f32_cuda	(	cublasHandle_t	handle,
		const float *	mat_f32_dev,
		const float *	vec_f32_dev,
		float *	out_f32_dev,
		int	rows,
		int	cols,
		cudaStream_t	stream
	)

Implementation of matrix-vector multiplication with FP32.

This function performs matrix-vector multiplication using cuBLAS, optimized for FP32 precision. It handles the cuBLAS setup and error checking.

Parameters

handle	cuBLAS handle
mat_f32_dev	Matrix in row-major format (device pointer)
vec_f32_dev	Vector to multiply (device pointer)
out_f32_dev	Output vector (device pointer)
rows	Number of matrix rows
cols	Number of matrix columns
stream	CUDA stream for asynchronous execution

Definition at line 306 of file cuda_kernels.cu.

                                                                  {
  // Existing implementation of matvec_f32_f32_cuda
  // C = A^T * B where A is mat_f32_dev (passed as KxM) and B is vec_f32_dev (KxN)
  // M_blas = rows, N_blas = 1, K_blas = cols
  // A is mat_f32_dev, LDA = cols (K_blas)
  // B is vec_f32_dev, LDB = cols (K_blas)
  // C is out_f32_dev, LDC = rows (M_blas)
  const float alpha = 1.0f;
  const float beta = 0.0f;
  cublasStatus_t status = cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, 
                                     rows, 1, cols, 
                                     &alpha, 
                                     mat_f32_dev, cols, 
                                     vec_f32_dev, cols, 
                                     &beta, 
                                     out_f32_dev, rows);
  if (status != CUBLAS_STATUS_SUCCESS) {
    throw std::runtime_error("cuBLAS Sgemm failed in matvec_f32_f32_cuda: " +
                             std::to_string(status));
  }
  gpuErrchk(cudaGetLastError());
}

References j.

Referenced by matvec_f32_f32_cuda().

matvec_f32_f32_cuda() [2/2]

void matvec_f32_f32_cuda	(	cublasHandle_t	handle,
		const std::vector< float > &	mat_f32_host,
		const std::vector< float > &	vec_f32_host,
		std::vector< float > &	out_f32_host,
		int	rows,
		int	cols
	)

Definition at line 470 of file cuda_kernels.cu.

                                                                             {
  if (mat_f32_host.size() != (size_t)rows * cols) {
    throw std::runtime_error(
        "matvec_f32_f32_cuda (host/host/host): mat size mismatch.");
  }
  if (vec_f32_host.size() != (size_t)cols) {
    throw std::runtime_error(
        "matvec_f32_f32_cuda (host/host/host): vec size mismatch.");
  }
  out_f32_host.resize(rows);
 
  float* mat_f32_dev = nullptr;
  float* vec_f32_dev = nullptr;
  float* out_f32_dev = nullptr;
 
  gpuErrchk(cudaMalloc(&mat_f32_dev, mat_f32_host.size() * sizeof(float)));
  gpuErrchk(cudaMalloc(&vec_f32_dev, vec_f32_host.size() * sizeof(float)));
  gpuErrchk(cudaMalloc(&out_f32_dev, out_f32_host.size() * sizeof(float)));
 
  gpuErrchk(cudaMemcpy(mat_f32_dev, mat_f32_host.data(),
                       mat_f32_host.size() * sizeof(float),
                       cudaMemcpyHostToDevice));
  gpuErrchk(cudaMemcpy(vec_f32_dev, vec_f32_host.data(),
                       vec_f32_host.size() * sizeof(float),
                       cudaMemcpyHostToDevice));
 
  matvec_f32_f32_cuda(handle, mat_f32_dev, vec_f32_dev, out_f32_dev, rows, cols,
                      0);
 
  gpuErrchk(cudaDeviceSynchronize());
 
  gpuErrchk(cudaMemcpy(out_f32_host.data(), out_f32_dev,
                       out_f32_host.size() * sizeof(float),
                       cudaMemcpyDeviceToHost));
 
  gpuErrchk(cudaFree(mat_f32_dev));
  gpuErrchk(cudaFree(vec_f32_dev));
  gpuErrchk(cudaFree(out_f32_dev));
}

References j, and matvec_f32_f32_cuda().

moe_ffn_cuda()

void moe_ffn_cuda	(	cublasHandle_t	handle,
		const float *	x,
		int	T,
		int	d,
		int	d_ff,
		int	E,
		int	k,
		bool	renorm,
		const float *	router,
		const float *	gate_exps,
		const float *	up_exps,
		const float *	down_exps,
		float *	scratch,
		float *	out,
		cudaStream_t	stream
	)

Definition at line 3770 of file cuda_kernels.cu.

                                                                   {
  if (moe_use_routed(T, k) && E <= 256 && k <= 64) {
    const int tk = T * k;
    const float** Ag = reinterpret_cast<const float**>(scratch);
    const float** Au = Ag + tk;
    const float** Ad = Au + tk;
    const float** Bx = Ad + tk;
    const float** Bg = Bx + tk;
    float** Cg = reinterpret_cast<float**>(scratch) + 5 * (size_t)tk;
    float** Cu = Cg + tk;
    float** Cy = Cu + tk;
    float* logits = reinterpret_cast<float*>(Cy + tk);
    float* gsel = logits + (size_t)T * E;
    float* g = gsel + tk;
    float* u = g + (size_t)tk * d_ff;
    float* y = u + (size_t)tk * d_ff;
    cublasSetStream(handle, stream);
    moe_gemm_rm(handle, x, router, logits, T, E, d);
    moe_route_sel_kernel<<<T, 32, 0, stream>>>(logits, E, k, renorm ? 1 : 0, x, d, d_ff, gate_exps, up_exps,
                                               down_exps, gsel, g, u, y, Ag, Au, Ad, Bx, Bg, Cg, Cu, Cy);
    const float one = 1.0f, zero = 0.0f;
    // Each item: row-major [rows][cols] weight times a column vector (see moe_gemm_rm).
    if (cublasSgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, d_ff, 1, d, &one, Ag, d, Bx, d, &zero, Cg, d_ff, tk) !=
            CUBLAS_STATUS_SUCCESS ||
        cublasSgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, d_ff, 1, d, &one, Au, d, Bx, d, &zero, Cu, d_ff, tk) !=
            CUBLAS_STATUS_SUCCESS)
      throw std::runtime_error("moe_ffn_cuda: cublasSgemmBatched failed");
    const size_t n = (size_t)tk * d_ff;
    moe_swiglu_kernel<<<(unsigned)((n + 255) / 256), 256, 0, stream>>>(g, u, n);
    if (cublasSgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, d, 1, d_ff, &one, Ad, d_ff, Bg, d_ff, &zero, Cy, d, tk) !=
        CUBLAS_STATUS_SUCCESS)
      throw std::runtime_error("moe_ffn_cuda: cublasSgemmBatched failed");
    moe_combine_sel_kernel<<<T, 128, 0, stream>>>(y, gsel, out, d, k);
    gpuErrchk(cudaGetLastError());
    return;
  }
  float* logits = scratch;
  float* gates = logits + (size_t)T * E;
  float* sel = gates + (size_t)T * E;
  float* g = sel + (size_t)T * E;
  float* u = g + (size_t)T * d_ff;
  float* y = u + (size_t)T * d_ff;
  cublasSetStream(handle, stream);
  moe_gemm_rm(handle, x, router, logits, T, E, d);
  moe_route_kernel<<<T, 32, 0, stream>>>(logits, gates, sel, E, k, renorm ? 1 : 0);
  const size_t n = (size_t)T * d_ff;
  for (int e = 0; e < E; ++e) {
    moe_gemm_rm(handle, x, gate_exps + (size_t)e * d_ff * d, g, T, d_ff, d);
    moe_gemm_rm(handle, x, up_exps + (size_t)e * d_ff * d, u, T, d_ff, d);
    moe_swiglu_kernel<<<(unsigned)((n + 255) / 256), 256, 0, stream>>>(g, u, n);
    moe_gemm_rm(handle, g, down_exps + (size_t)e * d * d_ff, y + (size_t)e * T * d, T, d, d_ff);
  }
  moe_combine_kernel<<<T, 128, 0, stream>>>(y, gates, sel, out, T, d, E);
  gpuErrchk(cudaGetLastError());
}

References j, k, n, and out.

moe_ffn_scratch_floats()

size_t moe_ffn_scratch_floats	(	int	T,
		int	d,
		int	d_ff,
		int	E,
		int	k
	)

Definition at line 3765 of file cuda_kernels.cu.

                                                                    {
  const size_t dense = (size_t)T * E * 3 + (size_t)T * d_ff * 2 + (size_t)E * T * d;
  return std::max(dense, moe_routed_floats(T, d, d_ff, E, k));
}

References j, and k.

reduce_partial_sums_kernel()

__global__ void reduce_partial_sums_kernel	(	const float *	partial_sums,
		float *	total_sum_sq_out,
		int	num_partial_sums
	)

Definition at line 185 of file cuda_kernels.cu.

                                                                                                                     {
    extern __shared__ float sdata[];
    unsigned int tid = threadIdx.x;
 
    float my_sum = 0.0f;
    for (unsigned int i = tid; i < num_partial_sums; i += blockDim.x) {
        my_sum += partial_sums[i];
    }
    sdata[tid] = my_sum;
    __syncthreads();
 
    // Standard parallel reduction in shared memory
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
 
    if (tid == 0) {
        total_sum_sq_out[0] = sdata[0];
    }
}

References j.

Referenced by rmsnorm_vector_cuda().

rmsnorm_apply_kernel()

__global__ void rmsnorm_apply_kernel	(	const float *	x,
		const float *	weight,
		float *	out,
		int	n,
		float	inv_norm_factor
	)

Kernel for applying RMS normalization.

This kernel applies the normalization weights and scaling factor to the input tensor. It's the second step in RMS normalization.

Parameters

x	Input tensor
weight	Normalization weights
out	Output tensor
n	Size of tensors
inv_norm_factor	Inverse of the normalization factor

Definition at line 177 of file cuda_kernels.cu.

                                                                               {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) {
    out[i] = x[i] * inv_norm_factor * weight[i];
  }
}

References j, n, and out.

rmsnorm_optimized_det_kernel()

__global__ void rmsnorm_optimized_det_kernel	(	const float *__restrict__	x,
		const float *__restrict__	weight,
		float *__restrict__	out,
		int	n,
		float	eps
	)

Definition at line 2945 of file cuda_kernels.cu.

                                                               {
    const int tid_global = threadIdx.x + blockIdx.x * blockDim.x;
    const int tid_block = threadIdx.x;
 
    __shared__ float s_warp_sum_sq[32];
    __shared__ float s_block_sum_sq;
 
    float thread_partial_sum_sq = 0.0f;
    for (int i = tid_global; i < n; i += gridDim.x * blockDim.x) {
        float val = x[i];
        thread_partial_sum_sq += val * val;
    }
 
    // Reduce within warp
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        thread_partial_sum_sq += __shfl_down_sync(0xFFFFFFFF, thread_partial_sum_sq, offset);
    }
 
    // Warp leaders store their sums; thread 0 adds them in warp order
    if ((tid_block % warpSize) == 0) {
        s_warp_sum_sq[tid_block / warpSize] = thread_partial_sum_sq;
    }
    __syncthreads();
    if (tid_block == 0) {
        const int num_warps = (blockDim.x + warpSize - 1) / warpSize;
        float total = 0.0f;
        for (int w = 0; w < num_warps; ++w) total += s_warp_sum_sq[w];
        s_block_sum_sq = total;
    }
    __syncthreads();
 
    float total_sum_sq = s_block_sum_sq;
    float rms_scale = rsqrtf(total_sum_sq / n + eps);
 
    for (int i = tid_global; i < n; i += gridDim.x * blockDim.x) {
        out[i] = x[i] * weight[i] * rms_scale;
    }
}

References j, n, and out.

rmsnorm_optimized_kernel()

__global__ void rmsnorm_optimized_kernel	(	const float *__restrict__	x,
		const float *__restrict__	weight,
		float *__restrict__	out,
		int	n,
		float	eps
	)

Optimized RMS normalization using simple warp primitives and shared memory.

This kernel uses warp-level reductions and a shared memory variable for block-wide sum.

Definition at line 2738 of file cuda_kernels.cu.

                                                         {
    const int tid_global = threadIdx.x + blockIdx.x * blockDim.x;
    const int tid_block = threadIdx.x;
 
    __shared__ float s_block_sum_sq;
 
    if (tid_block == 0) {
        s_block_sum_sq = 0.0f;
    }
    __syncthreads();
 
    float thread_partial_sum_sq = 0.0f;
    for (int i = tid_global; i < n; i += gridDim.x * blockDim.x) {
        float val = x[i];
        thread_partial_sum_sq += val * val;
    }
 
    // Reduce within warp
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        thread_partial_sum_sq += __shfl_down_sync(0xFFFFFFFF, thread_partial_sum_sq, offset);
    }
 
    // Warp leaders add to shared memory sum
    if ((tid_block % warpSize) == 0) {
        atomicAdd(&s_block_sum_sq, thread_partial_sum_sq);
    }
    __syncthreads(); // Ensure all warps have added their sums
 
    float total_sum_sq = s_block_sum_sq;
    float rms_scale = rsqrtf(total_sum_sq / n + eps);
 
    for (int i = tid_global; i < n; i += gridDim.x * blockDim.x) {
        out[i] = x[i] * weight[i] * rms_scale;
    }
}

References j, n, and out.

Referenced by rmsnorm_vector_cuda_optimized().

rmsnorm_sum_squares_kernel()

__global__ void rmsnorm_sum_squares_kernel	(	const float *	x,
		float *	partial_sums,
		int	n
	)

Kernel for computing sum of squares in RMS normalization.

This kernel computes the sum of squares of input elements using parallel reduction with shared memory. It's the first step in RMS normalization.

Parameters

x	Input tensor
partial_sums	Output array for partial sums
n	Size of input tensor

Definition at line 143 of file cuda_kernels.cu.

                                                  {
  extern __shared__ float sdata[];
 
  unsigned int tid = threadIdx.x;
  unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
 
  sdata[tid] = (i < n) ? x[i] * x[i] : 0.0f;
  __syncthreads();
 
  for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
      sdata[tid] += sdata[tid + s];
    }
    __syncthreads();
  }
 
  if (tid == 0) {
    partial_sums[blockIdx.x] = sdata[0];
  }
}

References j, and n.

Referenced by rmsnorm_vector_cuda().

rmsnorm_vector_cuda() [1/2]

void rmsnorm_vector_cuda	(	const float *	x_dev,
		const float *	weight_dev,
		float *	out_dev,
		int	n,
		float	eps,
		cudaStream_t	stream
	)

Implementation of RMS normalization with device pointers.

This function implements the complete RMS normalization process:

1. Computes sum of squares using parallel reduction
2. Calculates normalization factor
3. Applies normalization with weights

Parameters

x_dev	Input tensor (device pointer)
weight_dev	Normalization weights (device pointer)
out_dev	Output tensor (device pointer)
n	Size of tensors
eps	Epsilon for numerical stability
stream	CUDA stream for asynchronous execution

Definition at line 224 of file cuda_kernels.cu.

                                              {
  const int threads_per_block = 256; // General purpose threads per block
  int num_blocks_reduce_pass1 = (n + threads_per_block - 1) / threads_per_block;
  size_t shared_mem_size_pass1 = threads_per_block * sizeof(float);
 
  float* partial_sums_dev = nullptr;
  gpuErrchk(cudaMalloc(&partial_sums_dev, num_blocks_reduce_pass1 * sizeof(float)));
 
  rmsnorm_sum_squares_kernel<<<num_blocks_reduce_pass1, threads_per_block,
                               shared_mem_size_pass1, stream>>>(x_dev,
                                                          partial_sums_dev, n);
  gpuErrchk(cudaGetLastError());
 
  float* total_sum_sq_dev = nullptr;
  gpuErrchk(cudaMalloc(&total_sum_sq_dev, sizeof(float)));
 
  const int threads_for_final_reduction = (num_blocks_reduce_pass1 <= 256) ? 256 :
                                          (num_blocks_reduce_pass1 <= 512) ? 512 : 1024; 
                                          // Cap at 1024, ensure kernel handles num_partial_sums correctly if it's less than blockDim.x
  size_t shared_mem_final_reduction = threads_for_final_reduction * sizeof(float);
  
  // Launch a single block for the final reduction of partial_sums_dev
  reduce_partial_sums_kernel<<<1, threads_for_final_reduction, 
                               shared_mem_final_reduction, stream>>>(
                                   partial_sums_dev, total_sum_sq_dev, num_blocks_reduce_pass1);
  gpuErrchk(cudaGetLastError());
  gpuErrchk(cudaFree(partial_sums_dev)); // Free intermediate partial sums
 
  float h_total_ssq = 0.0f;
  gpuErrchk(cudaMemcpyAsync(&h_total_ssq, total_sum_sq_dev, sizeof(float),
                       cudaMemcpyDeviceToHost, stream));
  gpuErrchk(cudaStreamSynchronize(stream)); // Ensure h_total_ssq is ready
  gpuErrchk(cudaFree(total_sum_sq_dev));
 
  double total_ssq_double = static_cast<double>(h_total_ssq);
  total_ssq_double /= n;
  float inv_norm_factor = 1.0f / SAFE_SQRT(static_cast<float>(total_ssq_double) + eps);
 
  int num_blocks_apply = (n + threads_per_block - 1) / threads_per_block;
  rmsnorm_apply_kernel<<<num_blocks_apply, threads_per_block, 0, stream>>>(
      x_dev, weight_dev, out_dev, n, inv_norm_factor);
  gpuErrchk(cudaGetLastError());
}

References j, n, reduce_partial_sums_kernel(), rmsnorm_sum_squares_kernel(), and SAFE_SQRT.

Referenced by rmsnorm_vector_cuda().

rmsnorm_vector_cuda() [2/2]

void rmsnorm_vector_cuda	(	const std::vector< float > &	x_in_host,
		const std::vector< float > &	weight_host,
		std::vector< float > &	out_host,
		int	n,
		float	eps
	)

Definition at line 270 of file cuda_kernels.cu.

                                                                       {
  out_host.resize(n);
  float *x_dev = nullptr, *weight_dev = nullptr, *out_dev = nullptr;
  gpuErrchk(cudaMalloc(&x_dev, n * sizeof(float)));
  gpuErrchk(cudaMalloc(&weight_dev, n * sizeof(float)));
  gpuErrchk(cudaMalloc(&out_dev, n * sizeof(float)));
  gpuErrchk(cudaMemcpy(x_dev, x_in_host.data(), n * sizeof(float),
                       cudaMemcpyHostToDevice));
  gpuErrchk(cudaMemcpy(weight_dev, weight_host.data(), n * sizeof(float),
                       cudaMemcpyHostToDevice));
  rmsnorm_vector_cuda(x_dev, weight_dev, out_dev, n, eps, 0);
  gpuErrchk(cudaDeviceSynchronize());
  gpuErrchk(cudaMemcpy(out_host.data(), out_dev, n * sizeof(float),
                       cudaMemcpyDeviceToHost));
  gpuErrchk(cudaFree(x_dev));
  gpuErrchk(cudaFree(weight_dev));
  gpuErrchk(cudaFree(out_dev));
}

References j, n, and rmsnorm_vector_cuda().

rmsnorm_vector_cuda_optimized()

void rmsnorm_vector_cuda_optimized	(	const float *	x_dev,
		const float *	weight_dev,
		float *	out_dev,
		int	n,
		float	eps,
		cudaStream_t	stream
	)

Optimized wrappers using the new high-performance kernels.

Definition at line 3082 of file cuda_kernels.cu.

                                                      {
    if (n <= 0) return;
    const int block_size = 256;
    // Ensure grid_size is 1 as the kernel is designed for block-wide reduction.
    const int grid_size = 1;
    if (ember::deterministic()) {
        rmsnorm_optimized_det_kernel<<<grid_size, block_size, 0, stream>>>(x_dev, weight_dev, out_dev, n, eps);
        gpuErrchk(cudaGetLastError());
        return;
    }
    rmsnorm_optimized_kernel<<<grid_size, block_size, sizeof(float), stream>>>(
        x_dev, weight_dev, out_dev, n, eps);
    gpuErrchk(cudaGetLastError());
}

References ember::deterministic(), j, n, and rmsnorm_optimized_kernel().

rope_and_update_kv_cache_cuda()

void rope_and_update_kv_cache_cuda	(	float *	cache_base_ptr,
		const float *	kv_vector_head,
		const float *	all_freqs_cis_base,
		int	pos,
		int	kv_head_idx,
		int	max_seq_len,
		int	num_kv_heads,
		int	head_dim,
		cudaStream_t	stream
	)

Definition at line 1412 of file cuda_kernels.cu.

                                                        {
  if (head_dim % 2 != 0) {
    Logger::error("rope_and_update_kv_cache_cuda: head_dim must be even.");
    return;
  }
  if (pos < 0 || pos >= max_seq_len) {
    Logger::error("rope_and_update_kv_cache_cuda: pos out of bounds.");
    return;
  }
  if (kv_head_idx < 0 || kv_head_idx >= num_kv_heads) {
    Logger::error("rope_and_update_kv_cache_cuda: kv_head_idx out of bounds.");
    return;
  }
  if (!kv_vector_head || !all_freqs_cis_base || !cache_base_ptr) {
    Logger::error(
        "rope_and_update_kv_cache_cuda: Received null device pointer(s).");
    return;
  }
 
  int threads_per_block = 128;
  int num_blocks = (head_dim / 2 + threads_per_block - 1) / threads_per_block;
 
  rope_and_update_kv_cache_kernel<<<num_blocks, threads_per_block, 0, stream>>>(
      cache_base_ptr, kv_vector_head, all_freqs_cis_base, pos, kv_head_idx,
      max_seq_len, num_kv_heads, head_dim);
  gpuErrchk(cudaGetLastError());
}

References Logger::error(), j, and pos.

rope_and_update_kv_cache_kernel()

__global__ void rope_and_update_kv_cache_kernel	(	float *	cache_base_ptr,
		const float *	kv_vector_head,
		const float *	all_freqs_cis_base,
		int	pos,
		int	kv_head_idx,
		int	max_seq_len,
		int	num_kv_heads,
		int	head_dim
	)

Definition at line 1376 of file cuda_kernels.cu.

                                    {
  int pair_idx = blockIdx.x * blockDim.x + threadIdx.x;
  int dim_half = head_dim / 2;
 
  if (pair_idx >= dim_half) return;
 
  int idx0 = pair_idx;
  int idx1 = pair_idx + dim_half;
  float kv0 = kv_vector_head[idx0];
  float kv1 = kv_vector_head[idx1];
 
  size_t freq_base_offset = (size_t)pos * head_dim + (size_t)pair_idx * 2;
  float cos_val = all_freqs_cis_base[freq_base_offset];
  float sin_val = all_freqs_cis_base[freq_base_offset + 1];
 
  float kv0_rotated = kv0 * cos_val - kv1 * sin_val;
  float kv1_rotated = kv0 * sin_val + kv1 * cos_val;
 
  size_t cache_offset_0 = (size_t)pos * num_kv_heads * head_dim +
                          (size_t)kv_head_idx * head_dim + idx0;
  size_t cache_offset_1 = cache_offset_0 + dim_half;
 
  size_t total_cache_size = (size_t)max_seq_len * num_kv_heads * head_dim;
  if (cache_offset_0 >= total_cache_size ||
      cache_offset_1 >= total_cache_size) {
    return;
  }
 
  cache_base_ptr[cache_offset_0] = kv0_rotated;
  cache_base_ptr[cache_offset_1] = kv1_rotated;
}

References j, and pos.

rope_batch_cuda()

void rope_batch_cuda	(	float *	d_q_batch,
		float *	d_k_batch,
		const float *	d_all_freqs_cis_base,
		int	num_tokens,
		int	num_q_heads,
		int	num_kv_heads,
		int	head_dim,
		int	start_pos_in_kv_cache,
		bool	use_adjacent_pairing,
		cudaStream_t	stream
	)

Definition at line 872 of file cuda_kernels.cu.

                                          {
 
 
    if (head_dim == 0) { 
        Logger::warning("[ROPE_BATCH_CUDA_SKIP] head_dim is 0. Skipping RoPE application.");
        return;
    }
    if (head_dim % 2 != 0) {
        Logger::error("rope_batch_cuda: head_dim must be even. Got " + std::to_string(head_dim));
        Logger::info("[ROPE_BATCH_CUDA_EXIT] Error: head_dim not even.");
        return; 
    }
    if (num_tokens == 0) {
        Logger::info("[ROPE_BATCH_CUDA_EXIT] num_tokens is 0. Nothing to process.");
        return; // Nothing to process
    }
 
    int log_count_elements = std::min(3, head_dim);
 
    // Log initial Q values
    if (d_q_batch && num_q_heads > 0 && head_dim > 0) {
        for (int token_to_log_idx = 0; token_to_log_idx < std::min(num_tokens, 2); ++token_to_log_idx) {
            std::vector<float> h_q_sample_before(log_count_elements);
            size_t q_offset_elements = (size_t)token_to_log_idx * num_q_heads * head_dim;
            
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_q_sample_before.data(), 
                                      d_q_batch + q_offset_elements, 
                                      log_count_elements * sizeof(float), 
                                      cudaMemcpyDeviceToHost, stream));
        }
    }
 
    // Log initial K values
    if (d_k_batch && num_kv_heads > 0 && head_dim > 0) {
        for (int token_to_log_idx = 0; token_to_log_idx < std::min(num_tokens, 2); ++token_to_log_idx) {
            std::vector<float> h_k_sample_before(log_count_elements);
            size_t k_offset_elements = (size_t)token_to_log_idx * num_kv_heads * head_dim;
 
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_k_sample_before.data(), 
                                      d_k_batch + k_offset_elements, 
                                      log_count_elements * sizeof(float), 
                                      cudaMemcpyDeviceToHost, stream));
        }
    }
 
    dim3 block_dim(head_dim / 2);
 
    // Apply RoPE to Q batch
    if (d_q_batch && num_q_heads > 0) {
        dim3 grid_dim_q(num_tokens, num_q_heads);
        rope_batch_kernel_cuda<<<grid_dim_q, block_dim, 0, stream>>>(
            d_q_batch,
            num_tokens,
            num_q_heads,
            head_dim,
            d_all_freqs_cis_base,
            start_pos_in_kv_cache,
            use_adjacent_pairing
        );
        gpuErrchk(cudaGetLastError()); 
    }
 
    // Apply RoPE to K batch
    if (d_k_batch && num_kv_heads > 0) {
        dim3 grid_dim_k(num_tokens, num_kv_heads); 
        rope_batch_kernel_cuda<<<grid_dim_k, block_dim, 0, stream>>>(
            d_k_batch,
            num_tokens,
            num_kv_heads, 
            head_dim,
            d_all_freqs_cis_base,
            start_pos_in_kv_cache,
            use_adjacent_pairing
        );
        gpuErrchk(cudaGetLastError()); 
    }
 
    // Log Q values after RoPE
    if (d_q_batch && num_q_heads > 0 && head_dim > 0) {
        for (int token_to_log_idx = 0; token_to_log_idx < std::min(num_tokens, 2); ++token_to_log_idx) {
            std::vector<float> h_q_sample_after(log_count_elements);
            size_t q_offset_elements = (size_t)token_to_log_idx * num_q_heads * head_dim;
            
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_q_sample_after.data(), 
                                      d_q_batch + q_offset_elements, 
                                      log_count_elements * sizeof(float), 
                                      cudaMemcpyDeviceToHost, stream));
        }
    }
 
    // Log K values after RoPE
    if (d_k_batch && num_kv_heads > 0 && head_dim > 0) {
        for (int token_to_log_idx = 0; token_to_log_idx < std::min(num_tokens, 2); ++token_to_log_idx) {
            std::vector<float> h_k_sample_after(log_count_elements);
            size_t k_offset_elements = (size_t)token_to_log_idx * num_kv_heads * head_dim;
 
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_k_sample_after.data(), 
                                      d_k_batch + k_offset_elements, 
                                      log_count_elements * sizeof(float), 
                                      cudaMemcpyDeviceToHost, stream));
        }
    }
}

References Logger::error(), Logger::info(), j, kernel_debug_copies(), and Logger::warning().

rope_batch_kernel_cuda()

__global__ void rope_batch_kernel_cuda	(	float *	x_batch_tensor,
		int	num_tokens,
		int	num_heads_in_tensor,
		int	head_dim,
		const float *	d_all_freqs_cis_base,
		int	start_pos_offset,
		bool	use_adjacent_pairing
	)

Definition at line 812 of file cuda_kernels.cu.

                               {
 
    // Grid: dim3 grid_dim(num_tokens, num_heads_in_tensor);
    // Block: dim3 block_dim(head_dim / 2);
    int pair_within_head_idx = threadIdx.x; // Thread index within the block (0 to head_dim/2 - 1)
    int head_idx = blockIdx.y;              // Block index in Y dimension maps to head index
    int token_idx = blockIdx.x;             // Block index in X dimension maps to token index within the batch
 
    // Boundary checks
    if (pair_within_head_idx >= (head_dim / 2) || 
        head_idx >= num_heads_in_tensor ||
        token_idx >= num_tokens) {
        return;
    }
 
    // Determine the actual sequence position for RoPE for the current token
    int current_token_seq_pos = start_pos_offset + token_idx;
 
    // Fetch cos and sin values from the precomputed table
    // d_all_freqs_cis_base is effectively [max_seq_len][head_dim/2][2]
    // Access: (pos * (head_dim/2) + pair_idx_in_head) * 2 for cos, and +1 for sin
    size_t freq_pair_offset_in_table = ((size_t)current_token_seq_pos * (head_dim / 2) + pair_within_head_idx) * 2;
    float cos_val = d_all_freqs_cis_base[freq_pair_offset_in_table];
    float sin_val = d_all_freqs_cis_base[freq_pair_offset_in_table + 1];
 
    // Calculate base pointer to the current token's data in x_batch_tensor
    // x_batch_tensor is [num_tokens][num_heads_in_tensor][head_dim]
    float* token_data_start = x_batch_tensor + (size_t)token_idx * num_heads_in_tensor * head_dim;
    // Pointer to the start of the current head's data for this token
    float* head_data_start = token_data_start + (size_t)head_idx * head_dim;
 
    int x0_idx_in_head, x1_idx_in_head;
    if (use_adjacent_pairing) {
        // Adjacent pairing: (x_0, x_1), (x_2, x_3), ... for a head
        // pair_within_head_idx = 0 maps to (x_0, x_1), pair_within_head_idx = 1 maps to (x_2, x_3)
        x0_idx_in_head = pair_within_head_idx * 2;
        x1_idx_in_head = pair_within_head_idx * 2 + 1;
    } else {
        // Split-half pairing: (x_0, x_{D/2}), (x_1, x_{D/2+1}), ... for a head
        // pair_within_head_idx = 0 maps to (x_0, x_{D/2}), pair_within_head_idx = 1 maps to (x_1, x_{D/2+1})
        x0_idx_in_head = pair_within_head_idx;
        x1_idx_in_head = pair_within_head_idx + (head_dim / 2);
    }
 
    float x0_val = head_data_start[x0_idx_in_head];
    float x1_val = head_data_start[x1_idx_in_head];
 
 
    head_data_start[x0_idx_in_head] = x0_val * cos_val - x1_val * sin_val;
    head_data_start[x1_idx_in_head] = x0_val * sin_val + x1_val * cos_val;
}

References j.

rope_cuda()

void rope_cuda	(	float *	x_dev,
		int	num_heads,
		int	head_dim,
		const float *	all_freqs_cis_dev_base,
		int	pos,
		bool	use_adjacent_pairing,
		cudaStream_t	stream
	)

Definition at line 800 of file cuda_kernels.cu.

                                                                                                             {
  int total_pairs = num_heads * (head_dim / 2);
  int threads_per_block = 256;
  int num_blocks = (total_pairs + threads_per_block - 1) / threads_per_block;
 
  rope_kernel<<<num_blocks, threads_per_block, 0, stream>>>(
      x_dev, num_heads, head_dim, all_freqs_cis_dev_base, pos, use_adjacent_pairing);
  gpuErrchk(cudaGetLastError());
}

References j, and pos.

rope_devpos_cuda()

void rope_devpos_cuda	(	float *	x_dev,
		int	num_heads,
		int	head_dim,
		const float *	all_freqs_cis_dev_base,
		const int *	step_dev,
		bool	use_adjacent_pairing,
		cudaStream_t	stream
	)

Definition at line 3573 of file cuda_kernels.cu.

                                                                                           {
  const int total_pairs = num_heads * (head_dim / 2);
  rope_devpos_kernel<<<(total_pairs + 255) / 256, 256, 0, stream>>>(x_dev, num_heads, head_dim,
                                                                   all_freqs_cis_dev_base, step_dev,
                                                                   use_adjacent_pairing);
  gpuErrchk(cudaGetLastError());
}

References j, and rope_devpos_kernel().

rope_devpos_kernel()

__global__ void rope_devpos_kernel	(	float *	x,
		int	num_heads,
		int	head_dim,
		const float *	all_freqs_cis_base,
		const int *	step_dev,
		bool	use_adjacent_pairing
	)

Definition at line 3548 of file cuda_kernels.cu.

                                                                                   {
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int total_pairs = num_heads * (head_dim / 2);
  if (idx >= total_pairs) return;
  const int pos = step_dev[1];
  const int head_idx = idx / (head_dim / 2);
  const int dim_pair_idx = idx % (head_dim / 2);
  const size_t freq_base_offset = ((size_t)pos * (head_dim / 2) + dim_pair_idx) * 2;
  int base_x0, base_x1;
  if (use_adjacent_pairing) {
    base_x0 = head_idx * head_dim + (dim_pair_idx * 2);
    base_x1 = head_idx * head_dim + (dim_pair_idx * 2 + 1);
  } else {
    base_x0 = head_idx * head_dim + dim_pair_idx;
    base_x1 = head_idx * head_dim + dim_pair_idx + (head_dim / 2);
  }
  const float x0 = x[base_x0];
  const float x1 = x[base_x1];
  const float cos_val = all_freqs_cis_base[freq_base_offset];
  const float sin_val = all_freqs_cis_base[freq_base_offset + 1];
  x[base_x0] = x0 * cos_val - x1 * sin_val;
  x[base_x1] = x0 * sin_val + x1 * cos_val;
}

References j, and pos.

Referenced by rope_devpos_cuda().

rope_kernel()

__global__ void rope_kernel	(	float *	x,
		int	num_heads,
		int	head_dim,
		const float *	all_freqs_cis_base,
		int	pos,
		bool	use_adjacent_pairing
	)

Definition at line 767 of file cuda_kernels.cu.

                                                                                                 {
  int idx = blockIdx.x * blockDim.x + threadIdx.x; // Global thread index for pairs
  int total_pairs = num_heads * (head_dim / 2);   // Total (dim_i, dim_{i+1}) or (dim_i, dim_{i+D/2}) pairs across all heads
  if (idx >= total_pairs) return;
 
  int head_idx = idx / (head_dim / 2);        // The specific head this thread works on
  int dim_pair_idx = idx % (head_dim / 2);  // The index of the pair within this head (0 to head_dim/2 - 1)
 
  size_t freq_cos_sin_pair_idx = (size_t)pos * (head_dim / 2) + dim_pair_idx;
  size_t freq_base_offset = freq_cos_sin_pair_idx * 2; 
 
  int base_x0, base_x1;
  if (use_adjacent_pairing) {
    // Adjacent pairing: Rotate x[h*HD + 2*j] with x[h*HD + 2*j + 1]
    base_x0 = head_idx * head_dim + (dim_pair_idx * 2);
    base_x1 = head_idx * head_dim + (dim_pair_idx * 2 + 1);
  } else {
    // Split-half pairing: Rotate x[h*HD + j] with x[h*HD + j + head_dim/2]
    base_x0 = head_idx * head_dim + dim_pair_idx;
    base_x1 = head_idx * head_dim + dim_pair_idx + (head_dim / 2);
  }
 
  float x0 = x[base_x0];
  float x1 = x[base_x1];
  
  float cos_val = all_freqs_cis_base[freq_base_offset];
  float sin_val = all_freqs_cis_base[freq_base_offset + 1];
 
  x[base_x0] = x0 * cos_val - x1 * sin_val;
  x[base_x1] = x0 * sin_val + x1 * cos_val;
}

References j, and pos.

silu_cuda()

void silu_cuda	(	const std::vector< float > &	x_host,
		std::vector< float > &	out_host,
		int	n
	)

Definition at line 521 of file cuda_kernels.cu.

                      {
  if (x_host.size() != n) {
    throw std::runtime_error("SiLU CUDA: Input vector size mismatch.");
  }
  out_host.resize(n);
 
  float* x_dev = nullptr;
  float* out_dev = nullptr;
  gpuErrchk(cudaMalloc(&x_dev, n * sizeof(float)));
  gpuErrchk(cudaMalloc(&out_dev, n * sizeof(float)));
 
  gpuErrchk(cudaMemcpy(x_dev, x_host.data(), n * sizeof(float),
                       cudaMemcpyHostToDevice));
 
  const int threads_per_block = 256;
  int num_blocks = (n + threads_per_block - 1) / threads_per_block;
  silu_kernel<<<num_blocks, threads_per_block>>>(x_dev, out_dev, n);
  gpuErrchk(cudaGetLastError());
 
  gpuErrchk(cudaMemcpy(out_host.data(), out_dev, n * sizeof(float),
                       cudaMemcpyDeviceToHost));
  gpuErrchk(cudaDeviceSynchronize());
 
  gpuErrchk(cudaFree(x_dev));
  gpuErrchk(cudaFree(out_dev));
}

References j, and n.

silu_kernel()

__global__ void silu_kernel	(	const float *	x,
		float *	out,
		int	n
	)

Definition at line 513 of file cuda_kernels.cu.

                                                               {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) {
    float x_val = x[i];
    out[i] = x_val / (1.0f + expf(-x_val));
  }
}

References j, n, and out.

Referenced by ember_ir::ir_silu_cuda().

softmax_exp_sum_kernel()

__global__ void softmax_exp_sum_kernel	(	const float *	x,
		float *	partial_sums,
		int	n,
		float	max_val
	)

Definition at line 570 of file cuda_kernels.cu.

                                                             {
  extern __shared__ float sdata[];
  unsigned int tid = threadIdx.x;
  unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
 
  sdata[tid] = (i < n) ? expf(x[i] - max_val) : 0.0f;
  __syncthreads();
 
  for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
      sdata[tid] += sdata[tid + s];
    }
    __syncthreads();
  }
 
  if (tid == 0) {
    partial_sums[blockIdx.x] = sdata[0];
  }
}

References j, and n.

softmax_find_max_kernel()

__global__ void softmax_find_max_kernel	(	const float *	x,
		float *	partial_max,
		int	n
	)

Definition at line 549 of file cuda_kernels.cu.

                                               {
  extern __shared__ float sdata[];
  unsigned int tid = threadIdx.x;
  unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
 
  sdata[tid] = (i < n) ? x[i] : -INFINITY;
  __syncthreads();
 
  for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
      sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
    }
    __syncthreads();
  }
 
  if (tid == 0) {
    partial_max[blockIdx.x] = sdata[0];
  }
}

References j, and n.

softmax_normalize_kernel()

__global__ void softmax_normalize_kernel	(	const float *	x,
		float *	out,
		int	n,
		float	max_val,
		float	inv_sum
	)

Definition at line 591 of file cuda_kernels.cu.

                                                                       {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) {
    out[i] = expf(x[i] - max_val) * inv_sum;
  }
}

References j, n, and out.

softmax_optimized_kernel()

__global__ void softmax_optimized_kernel	(	const float *__restrict__	x,
		float *__restrict__	out,
		int	n
	)

Optimized softmax using shared memory and warp primitives.

Correctly finds max and sum of exponentials at block level for accurate softmax.

Definition at line 2782 of file cuda_kernels.cu.

                                               {
    const int tid_global = threadIdx.x + blockIdx.x * blockDim.x;
    const int tid_block = threadIdx.x;
 
    __shared__ float s_block_max_val;
    __shared__ float s_block_sum_exp;
 
    // Initialize shared variables by the first thread of the block
    if (tid_block == 0) {
        s_block_max_val = -INFINITY;
        s_block_sum_exp = 0.0f;
    }
    __syncthreads();
 
    // Phase 1: Find max
    float thread_max_val = -INFINITY;
    for (int i = tid_global; i < n; i += gridDim.x * blockDim.x) {
        thread_max_val = fmaxf(thread_max_val, x[i]);
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        thread_max_val = fmaxf(thread_max_val, __shfl_down_sync(0xFFFFFFFF, thread_max_val, offset));
    }
    if ((tid_block % warpSize) == 0) { // Warp leaders
        atomicMaxFloat(&s_block_max_val, thread_max_val);
    }
    __syncthreads();
    float global_max = s_block_max_val;
 
    // Phase 2: Compute sum of exp
    float thread_sum_exp = 0.0f;
    for (int i = tid_global; i < n; i += gridDim.x * blockDim.x) {
        thread_sum_exp += expf(x[i] - global_max);
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        thread_sum_exp += __shfl_down_sync(0xFFFFFFFF, thread_sum_exp, offset);
    }
    if ((tid_block % warpSize) == 0) { // Warp leaders
        atomicAddFloat(&s_block_sum_exp, thread_sum_exp);
    }
    __syncthreads();
    float global_sum_exp = s_block_sum_exp;
 
    // Phase 3: Apply softmax
    float inv_sum = (global_sum_exp == 0.0f) ? 0.0f : 1.0f / global_sum_exp;
    for (int i = tid_global; i < n; i += gridDim.x * blockDim.x) {
        out[i] = expf(x[i] - global_max) * inv_sum;
    }
}

References atomicAddFloat(), atomicMaxFloat(), j, n, and out.

Referenced by softmax_vector_cuda_optimized().

softmax_vector_cuda()

void softmax_vector_cuda	(	const std::vector< float > &	x_host,
		std::vector< float > &	out_host,
		int	n
	)

Definition at line 599 of file cuda_kernels.cu.

                                                            {
  if (x_host.size() != n) {
    throw std::runtime_error("Softmax CUDA: Input vector size mismatch.");
  }
  if (n == 0) {
    out_host.clear();
    return;
  }
  out_host.resize(n);
 
  float* x_dev = nullptr;
  float* out_dev = nullptr;
  float* partial_max_dev = nullptr;
  float* partial_sum_dev = nullptr;
 
  const int threads_per_block = 256;
  int num_blocks = (n + threads_per_block - 1) / threads_per_block;
  size_t shared_mem_size = threads_per_block * sizeof(float);
 
  float* partial_max_host = new float[num_blocks];
  float* partial_sum_host = new float[num_blocks];
 
  gpuErrchk(cudaMalloc(&x_dev, n * sizeof(float)));
  gpuErrchk(cudaMalloc(&out_dev, n * sizeof(float)));
  gpuErrchk(cudaMalloc(&partial_max_dev, num_blocks * sizeof(float)));
  gpuErrchk(cudaMalloc(&partial_sum_dev, num_blocks * sizeof(float)));
 
  gpuErrchk(cudaMemcpy(x_dev, x_host.data(), n * sizeof(float),
                       cudaMemcpyHostToDevice));
 
  softmax_find_max_kernel<<<num_blocks, threads_per_block, shared_mem_size>>>(
      x_dev, partial_max_dev, n);
  gpuErrchk(cudaGetLastError());
  gpuErrchk(cudaMemcpy(partial_max_host, partial_max_dev,
                       num_blocks * sizeof(float), cudaMemcpyDeviceToHost));
 
  float max_val = -INFINITY;
  for (int i = 0; i < num_blocks; ++i) {
    if (partial_max_host[i] > max_val) {
      max_val = partial_max_host[i];
    }
  }
 
  softmax_exp_sum_kernel<<<num_blocks, threads_per_block, shared_mem_size>>>(
      x_dev, partial_sum_dev, n, max_val);
  gpuErrchk(cudaGetLastError());
  gpuErrchk(cudaMemcpy(partial_sum_host, partial_sum_dev,
                       num_blocks * sizeof(float), cudaMemcpyDeviceToHost));
 
  double exp_sum = 0.0;
  for (int i = 0; i < num_blocks; ++i) {
    exp_sum += partial_sum_host[i];
  }
  float inv_sum = 1.0f / static_cast<float>(exp_sum);
 
  softmax_normalize_kernel<<<num_blocks, threads_per_block>>>(x_dev, out_dev, n,
                                                              max_val, inv_sum);
  gpuErrchk(cudaGetLastError());
 
  gpuErrchk(cudaMemcpy(out_host.data(), out_dev, n * sizeof(float),
                       cudaMemcpyDeviceToHost));
 
  delete[] partial_max_host;
  delete[] partial_sum_host;
  gpuErrchk(cudaFree(x_dev));
  gpuErrchk(cudaFree(out_dev));
  gpuErrchk(cudaFree(partial_max_dev));
  gpuErrchk(cudaFree(partial_sum_dev));
}

References j, and n.

softmax_vector_cuda_optimized()

void softmax_vector_cuda_optimized	(	const float *	x_dev,
		float *	out_dev,
		int	n,
		cudaStream_t	stream
	)

Definition at line 3099 of file cuda_kernels.cu.

                                                      {
    if (n <= 0) return;
    const int block_size = 256;
    // Ensure grid_size is 1 as the kernel is designed for block-wide reduction.
    const int grid_size = 1;
    // Shared memory: 2 floats (max_val, sum_exp)
    softmax_optimized_kernel<<<grid_size, block_size, 2 * sizeof(float), stream>>>(
        x_dev, out_dev, n);
    gpuErrchk(cudaGetLastError());
}

References j, n, and softmax_optimized_kernel().

swiglu_batch_cuda()

void swiglu_batch_cuda	(	float *	d_out_batch,
		const float *	d_gate_act_batch,
		const float *	d_up_act_batch,
		int	num_tokens,
		int	intermediate_size,
		cudaStream_t	stream
	)

Definition at line 717 of file cuda_kernels.cu.

                         {
 
 
    if (num_tokens == 0 || intermediate_size == 0) {
        Logger::info("[SWIGLU_BATCH_CUDA_SKIP] num_tokens or intermediate_size is 0. Nothing to do.");
        Logger::info("[SWIGLU_BATCH_CUDA_EXIT] Skipped operation.");
        return; // Nothing to do
    }
 
    // Log initial input values (first token, first few elements)
    int log_count_elements = std::min(3, intermediate_size);
    if (num_tokens > 0 && intermediate_size > 0) {
        if (d_gate_act_batch) {
            std::vector<float> h_gate_sample(log_count_elements);
            // Gate is [num_tokens, intermediate_size]. Offset for T0 is 0.
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_gate_sample.data(), d_gate_act_batch, log_count_elements * sizeof(float), cudaMemcpyDeviceToHost, stream));
        }
        if (d_up_act_batch) {
            std::vector<float> h_up_sample(log_count_elements);
            // Up is [num_tokens, intermediate_size]. Offset for T0 is 0.
            if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_up_sample.data(), d_up_act_batch, log_count_elements * sizeof(float), cudaMemcpyDeviceToHost, stream));
        }
    }
 
    const int threads_per_block = 256;
    dim3 grid_dim(num_tokens);
    dim3 block_dim(threads_per_block);
 
    swiglu_batch_cuda_kernel<<<grid_dim, block_dim, 0, stream>>>(
        d_out_batch,
        d_gate_act_batch,
        d_up_act_batch,
        num_tokens,
        intermediate_size
    );
    gpuErrchk(cudaGetLastError());
 
    // Log output values
    if (num_tokens > 0 && intermediate_size > 0 && d_out_batch) {
        std::vector<float> h_out_sample(log_count_elements);
        // Out is [num_tokens, intermediate_size]. Offset for T0 is 0.
        if (kernel_debug_copies()) gpuErrchk(cudaMemcpyAsync(h_out_sample.data(), d_out_batch, log_count_elements * sizeof(float), cudaMemcpyDeviceToHost, stream));
    }
}

References Logger::info(), j, and kernel_debug_copies().

swiglu_batch_cuda_kernel()

__global__ void swiglu_batch_cuda_kernel	(	float *	d_out_batch,
		const float *	d_gate_act_batch,
		const float *	d_up_act_batch,
		int	num_tokens,
		int	intermediate_size
	)

Definition at line 690 of file cuda_kernels.cu.

                           {
 
    int token_idx = blockIdx.x;
    if (token_idx >= num_tokens) {
        return;
    }
 
    // Base pointers for the current token
    const float* current_token_gate_ptr = d_gate_act_batch + (size_t)token_idx * intermediate_size;
    const float* current_token_up_ptr = d_up_act_batch + (size_t)token_idx * intermediate_size;
    float* current_token_out_ptr = d_out_batch + (size_t)token_idx * intermediate_size;
 
    // Threads in the block parallelize over intermediate_size
    for (int i = threadIdx.x; i < intermediate_size; i += blockDim.x) {
        float gate_val = current_token_gate_ptr[i];
        float silu_gate = gate_val / (1.0f + expf(-gate_val)); // SiLU(x) = x * sigmoid(x)
        float up_val = current_token_up_ptr[i];
        current_token_out_ptr[i] = silu_gate * up_val;
    }
}

References j.

swiglu_cuda()

void swiglu_cuda	(	const float *	gate_dev,
		const float *	up_dev,
		float *	out_dev,
		int	n,
		cudaStream_t	stream
	)

Definition at line 680 of file cuda_kernels.cu.

                                             {
  const int threads_per_block = 256;
  int num_blocks = (n + threads_per_block - 1) / threads_per_block;
  swiglu_kernel<<<num_blocks, threads_per_block, 0, stream>>>(gate_dev, up_dev,
                                                              out_dev, n);
  gpuErrchk(cudaGetLastError());
}

References j, and n.

swiglu_kernel()

__global__ void swiglu_kernel	(	const float *	gate,
		const float *	up,
		float *	out,
		int	n
	)

Definition at line 670 of file cuda_kernels.cu.

                                     {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) {
    float gate_val = gate[i];
    float silu_gate = gate_val / (1.0f + expf(-gate_val));
    out[i] = silu_gate * up[i];
  }
}

References j, n, and out.

update_kv_cache_cuda()

void update_kv_cache_cuda	(	float *	cache_base_ptr,
		const float *	current_kv_head_vector,
		int	pos,
		int	kv_head_idx,
		int	max_seq_len,
		int	num_kv_heads,
		int	head_dim,
		cudaStream_t	stream
	)

Definition at line 1344 of file cuda_kernels.cu.

                                                                               {
  dim3 blockDim(head_dim);
  dim3 gridDim(1);
 
  if (pos < 0 || pos >= max_seq_len) {
    Logger::error("update_kv_cache_cuda: pos out of bounds (" +
                  std::to_string(pos) + " >= " + std::to_string(max_seq_len) +
                  ")");
    return;
  }
 
  if (kv_head_idx < 0 || kv_head_idx >= num_kv_heads) {
    Logger::error("update_kv_cache_cuda: kv_head_idx out of bounds (" +
                  std::to_string(kv_head_idx) +
                  " >= " + std::to_string(num_kv_heads) + ")");
    return;
  }
  if (!current_kv_head_vector) {
    Logger::error("update_kv_cache_cuda: Input K/V vector pointer is null.");
    return;
  }
 
  update_kv_cache_kernel<<<gridDim, blockDim, 0, stream>>>(
      cache_base_ptr, current_kv_head_vector, pos, kv_head_idx, max_seq_len,
      num_kv_heads, head_dim);
  gpuErrchk(cudaGetLastError());
}

References Logger::error(), j, and pos.

update_kv_cache_kernel()

__global__ void update_kv_cache_kernel	(	float *	cache_base_ptr,
		const float *	current_kv_vector,
		int	pos,
		int	kv_head_idx,
		int	max_seq_len,
		int	num_kv_heads,
		int	head_dim
	)

Definition at line 1324 of file cuda_kernels.cu.

                                                                       {
  int d = threadIdx.x;
  if (d >= head_dim) return;
 
  size_t cache_offset = (size_t)pos * num_kv_heads * head_dim +
                        (size_t)kv_head_idx * head_dim + d;
 
  size_t source_offset = d;
 
  size_t total_cache_size = (size_t)max_seq_len * num_kv_heads * head_dim;
  if (cache_offset >= total_cache_size) {
    return;
  }
 
  cache_base_ptr[cache_offset] = current_kv_vector[source_offset];
}

References j, and pos.