夜雨聆风BaseSamplingLayer::forward
1. skip_all, return 2. skip_any, copy logits 3. invokeBatchApplyTemperaturePenalty 4. invokeBatchApplyRepetitionPenalty 5. invokeMinLengthPenalty 6. runSampling
template <typename T>void BaseSamplingLayer<T>::forward(DecodingOutputParams& outputs, ForwardParams const& params){ auto* logits = params.logits.template getPtr<T>();#define ALL_OF(p_, sz_, dt_, v_) (std::all_of(p_, p_ + sz_, [&](dt_ b) { return b == v_; })) bool* skip_decode = skip_decode_ + ite * local_batch_size; if (ALL_OF(skip_decode, local_batch_size, bool, true)) { // No sample in the current batch to do TopX sampling. return; } skip_any_ = std::any_of(skip_decode, skip_decode + local_batch_size, [](bool b) { return b; }); if (skip_any_) { // A TopX Sampling layer directly changes the logit values. In case of // skip_any==true, meaning topk and topp layers will run simultaneously for // a batch in the same step. We copy the logits to an internal buffer, not // affecting the other sampling layers. TLLM_CHECK(params.logits.size() == local_batch_size * vocab_size_padded_); cudaD2Dcpy(runtime_logits_buf_, logits, params.logits.size()); logits = runtime_logits_buf_; } if (embedding_bias != nullptr || !ALL_OF(std::begin(mTemperature) + ite * local_batch_size, local_batch_size, float, 1.0f)) { invokeBatchApplyTemperaturePenalty(logits, embedding_bias, temperature_buf_ + ite * local_batch_size, local_batch_size, vocab_size_, vocab_size_padded_, stream_); } if (step > 1 && repetition_penalty_type_ != RepetitionPenaltyType::None) { if (!ALL_OF(std::begin(mRepetitionPenalty) + ite * local_batch_size, local_batch_size, float, default_value)) { invokeBatchApplyRepetitionPenalty(logits, repetition_penalty_buf_ + ite * local_batch_size, outputs.output_ids_ptr.template getPtr<const int*>(), outputs.sequence_length->getPtr<const int>(), batch_size, local_batch_size, vocab_size_padded_, input_lengths, repetition_penalty_type_, params.max_seq_len, stream_); } } invokeMinLengthPenalty(logits, min_lengths_buf_ + ite * local_batch_size, end_ids, outputs.sequence_length->getPtr<const int>(), input_lengths, local_batch_size, vocab_size_padded_, stream_); runSampling(outputs, params); if (is_free_buffer_after_forward_) { freeBuffer(); }}invokeApplyTemperaturePenalty
1.float, 超出 vocab_size 的部分 logits 为 -MXT_T_VAL2.half, 只处理 vocab_size 内的数。当 vocab_size %20, vocab_size_padd %2 0 时有效。3. temperature 有一个 1e-6f 的epsion值
// TODO Add half2 implementationtemplate <typename T>__global__ void applyTemperaturePenalty(T* logits, const T* bias, const float temperature_inverse, const int m, const int vocab_size, const int vocab_size_padd){ const bool IS_FP16 = std::is_same<T, half>::value; const T MAX_T_VAL = (IS_FP16) ? 65504.F : FLT_MAX; for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < m * vocab_size_padd; index += blockDim.x * gridDim.x) { T bias_val = bias == nullptr ? (T) (0.0f) : bias[index % vocab_size_padd]; if (index % vocab_size_padd < vocab_size) { logits[index] = (logits[index] + bias_val) * (T) temperature_inverse; } else { logits[index] = -MAX_T_VAL; } }}template <>__global__ void applyTemperaturePenalty(half2* logits, const half2* bias, const float temperature_inverse, const int batch_size, const int vocab_size, const int vocab_size_padded){ assert(vocab_size % 2 == 0); assert(vocab_size_padded % 2 == 0); const half2 mask_val = __float2half2_rn(-65504.0f); const half2 temp_inv = __float2half2_rn(temperature_inverse); const int half_vocab_size = vocab_size / 2; const int half_vocab_size_padded = vocab_size_padded / 2; for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < batch_size * half_vocab_size_padded; index += blockDim.x * gridDim.x) { int vocab_idx = index % half_vocab_size_padded; half2 logit = vocab_idx < half_vocab_size ? __ldg(&logits[index]) : mask_val; if (vocab_idx < half_vocab_size) { if (bias != nullptr) { logit = __hadd2(logit, bias[vocab_idx]); } logits[index] = __hmul2(logit, temp_inv); } }}template <typename T>void invokeApplyTemperaturePenalty(T* logits, const T* bias, const float temperature, const int batch_size, const int vocab_size, const int vocab_size_padd, cudaStream_t stream){ dim3 grid(min(vocab_size_padd, 1024)); dim3 block(min(batch_size * vocab_size_padd / block.x, 65536)); const T temperature_inverse = (T) (1.f / (temperature + 1e-6f)); if (std::is_same<T, half>::value && vocab_size % 2 == 0 && vocab_size_padd % 2 == 0) { applyTemperaturePenalty<<<grid, block, 0, stream>>>(reinterpret_cast<half2*>(logits), reinterpret_cast<const half2*>(bias), temperature_inverse, batch_size, vocab_size, vocab_size_padd); } else { applyTemperaturePenalty<T> <<<grid, block, 0, stream>>>(logits, bias, temperature_inverse, batch_size, vocab_size, vocab_size_padd); }}template void invokeApplyTemperaturePenalty(float* logits, const float* bias, const float temperature, const int batch_size, const int vocab_size, const int vocab_size_padd, cudaStream_t stream);template void invokeApplyTemperaturePenalty(half* logits, const half* bias, const float temperature, const int batch_size, const int vocab_size, const int vocab_size_padd, cudaStream_t stream);batchApplyTemperaturePenalty
template <typename T>__global__ void batchApplyTemperaturePenalty(T* logits, const T* bias, const float* temperatures, const int batch_size, const int vocab_size, const int vocab_size_padd){ // TODO: Add macro or device function to get MAX_T_VAL. const bool IS_FP16 = std::is_same<T, half>::value; const T MAX_T_VAL = (IS_FP16) ? 65504.F : FLT_MAX; extern __shared__ float inv_temperatures[]; if (threadIdx.x < batch_size) { inv_temperatures[threadIdx.x] = 1.0f / (temperatures[threadIdx.x] + 1e-6f); } __syncthreads(); for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < batch_size * vocab_size_padd; index += blockDim.x * gridDim.x) { int batch_idx = index / vocab_size_padd; int vocab_idx = index % vocab_size_padd; T logit = (vocab_idx < vocab_size) ? logits[index] : -MAX_T_VAL; if (vocab_idx < vocab_size) { if (bias != nullptr) { logit += bias[vocab_idx]; } logit *= inv_temperatures[batch_idx]; } logits[index] = logit; }}__global__ void batchApplyTemperaturePenalty_h2(half2* logits, const half2* bias, const float* temperatures, const int batch_size, const int vocab_size, const int vocab_size_padded){ assert(vocab_size % 2 == 0); assert(vocab_size_padded % 2 == 0); extern __shared__ half2 h2_inv_temperatures[]; if (threadIdx.x < batch_size) { h2_inv_temperatures[threadIdx.x] = __float2half2_rn(1.f / (temperatures[threadIdx.x] + 1e-6f)); } __syncthreads(); const half2 mask_val = __float2half2_rn(-65504.0f); const int half_vocab_size = vocab_size / 2; const int half_vocab_size_padded = vocab_size_padded / 2; for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < batch_size * half_vocab_size_padded; index += blockDim.x * gridDim.x) { int batch_idx = index / half_vocab_size_padded; int vocab_idx = index % half_vocab_size_padded; half2 logit = vocab_idx < half_vocab_size ? __ldg(&logits[index]) : mask_val; if (vocab_idx < half_vocab_size) { if (bias != nullptr) { logit = __hadd2(logit, bias[vocab_idx]); } logits[index] = __hmul2(logit, h2_inv_temperatures[batch_idx]); } }}template <typename T>void invokeBatchApplyTemperaturePenalty(T* logits, const T* bias, const float* temperatures, const int batch_size, const int vocab_size, const int vocab_size_padd, cudaStream_t stream){ TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); dim3 grid(min(vocab_size_padd, 1024)); dim3 block(min(batch_size * vocab_size_padd / block.x, 65536)); if (std::is_same<T, half>::value && vocab_size % 2 == 0 && vocab_size_padd % 2 == 0) { size_t smem_size = sizeof(half2) * batch_size; batchApplyTemperaturePenalty_h2<<<grid, block, smem_size, stream>>>(reinterpret_cast<half2*>(logits), reinterpret_cast<const half2*>(bias), temperatures, batch_size, vocab_size, vocab_size_padd); } else { size_t smem_size = sizeof(float) * batch_size; batchApplyTemperaturePenalty<T> <<<grid, block, smem_size, stream>>>(logits, bias, temperatures, batch_size, vocab_size, vocab_size_padd); }}template void invokeBatchApplyTemperaturePenalty(float* logits, const float* bias, const float* temperatures, const int batch_size, const int vocab_size, const int vocab_size_padd, cudaStream_t stream);template void invokeBatchApplyTemperaturePenalty(half* logits, const half* bias, const float* temperatures, const int batch_size, const int vocab_size, const int vocab_size_padd, cudaStream_t stream);batchApplyRepetitionPenalty
template <typename T>void invokeBatchApplyRepetitionPenalty(T* logits, const float* penalties, const int** output_ids, const int* sequence_lengths, const int batch_size, const int local_batch_size, const int vocab_size, const int* input_lengths, RepetitionPenaltyType penalty_type, int max_seq_len, cudaStream_t stream){ // Inputs // logits [local_batch_size, vocab_size] : logit values. // penalties [local_batch_size] : repetition penalty factors. // output_ids int**, [bs] array, each array has [1, max_seq_len] // sequence_lengths int*, [bs] // input_lengths [local_batch_size], input lengths TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); dim3 grid(min(max_seq_len, 1024)); dim3 block(batch_size); size_t smem_size = max_seq_len * (sizeof(float) + sizeof(int)); if (penalty_type == RepetitionPenaltyType::Additive) { batchApplyRepetitionPenalty<T, RepetitionPenaltyType::Additive><<<grid, block, smem_size, stream>>>( logits, penalties, output_ids, sequence_lengths, batch_size, vocab_size, input_lengths, max_seq_len); } else if (penalty_type == RepetitionPenaltyType::Multiplicative) { batchApplyRepetitionPenalty<T, RepetitionPenaltyType::Multiplicative><<<grid, block, smem_size, stream>>>( logits, penalties, output_ids, sequence_lengths, batch_size, vocab_size, input_lengths, max_seq_len); } else if (penalty_type == RepetitionPenaltyType::None) { // do nothing }}template void invokeBatchApplyRepetitionPenalty(float* logits, const float* penalties, const int** output_ids, const int* sequence_lengths, const int batch_size, const int local_batch_size, const int vocab_size, const int* input_lengths, RepetitionPenaltyType penalty_type, int max_seq_len, cudaStream_t stream);template void invokeBatchApplyRepetitionPenalty(half* logits, const float* penalties, const int** output_ids, const int* sequence_lengths, const int batch_size, const int local_batch_size, const int vocab_size, const int* input_lengths, RepetitionPenaltyType penalty_type, int max_seq_len, cudaStream_t stream);template <typename T, RepetitionPenaltyType penalty_type>__global__ void batchApplyRepetitionPenalty(T* logits, const float* penalties, const int** output_ids, const int* sequence_lengths, const int batch_size, const int vocab_size, const int* input_lengths, const int max_seq_len){ extern __shared__ float penalty_logits[]; // max_seq_len * (sizeof(float) + sizeof(int)); int* penalty_indices = (int*) (penalty_logits + max_seq_len); // [max_seq_len ,] const int batch_idx = blockIdx.x; // dim for B const float penalty = penalties[batch_idx]; const int current_step = sequence_lengths[batch_idx]; logits += batch_idx * vocab_size; // logits.shape = [B, V] // Phase 1. Find indices to penalize and keep the penalized values. // A vocab id can appear multiple times but should be penalized once. // threadIdx.x for T(current_sequnce_length) for (int index = threadIdx.x; index < current_step; index += blockDim.x) { // output_ids shape: (beam_size, input_len + output_len), [beamsize, max_seq_len] // blockIdx.y is dim for BeamSize， 这里blockIdx.y=0, blockDim.y=1 int penalty_index = output_ids[batch_idx][blockIdx.y * max_seq_len + index]; assert(penalty_index < vocab_size); penalty_indices[index] = penalty_index; // penalty_index in V float logit = (float) logits[penalty_index]; if (penalty_type == RepetitionPenaltyType::Additive) { penalty_logits[index] = logit - penalty; } else if (penalty_type == RepetitionPenaltyType::Multiplicative) { penalty_logits[index] = logit < 0.0f ? logit * penalty : logit / penalty; } else if (penalty_type == RepetitionPenaltyType::None) { penalty_logits[index] = logit; } else { // Unsupported type assert(false); } } if (blockDim.x > 32) // not in wrap, wrap size is 32 { __syncthreads(); } // Phase 2. Replace a logit value by the penalized one. for (int index = threadIdx.x; index < current_step; index += blockDim.x) { logits[penalty_indices[index]] = penalty_logits[index]; }}batchApplyMinLengthPenalty
template <typename T>void invokeMinLengthPenalty(T* logits, const int* min_lengths, const int* end_ids, const int* sequnece_lengths, const int* input_lengths, const int batch_size, const int vocab_size_padded, cudaStream_t stream){ const int block_size = min(batch_size, 1024); // B const int grid_size = (batch_size + block_size - 1) / block_size; batchApplyMinLengthPenalty<<<grid_size, block_size, 0, stream>>>( logits, min_lengths, end_ids, sequnece_lengths, input_lengths, vocab_size_padded);}template void invokeMinLengthPenalty(float* logits, const int* min_lengths, const int* end_ids, const int* sequnece_lengths, const int* input_lengths, const int batch_size, const int vocab_size_padded, cudaStream_t stream);template void invokeMinLengthPenalty(half* logits, const int* min_lengths, const int* end_ids, const int* sequnece_lengths, const int* input_lengths, const int batch_size, const int vocab_size_padded, cudaStream_t stream);template <typename T>__global__ void batchApplyMinLengthPenalty(T* logits, const int* min_lengths, const int* end_ids, const int* sequence_lengths, const int* input_lengths, const int vocab_size_padded){ int bid = threadIdx.x + blockIdx.x * blockDim.x; // batch index auto const input_length{input_lengths == nullptr ? 0 : input_lengths[bid]}; // We need +1 because sequence_lengths = num_gen_tokens + input_length - 1, which is equal to the length of k/v // caches. // sequence_lengths = num_gen_tokens + input_length - 1 这里的-1应该不包含最后的end_id // 这里的+1，应该是指当前step的长度。 if (sequence_lengths[bid] + 1 - input_length < min_lengths[bid]) { // 根据数据类型T选择一个极小的负数（-FLT_MAX或half类型的最小值-65504.0f）。在Softmax函数中，这个值会使得对应的概率接近于零。 T mask_val = (std::is_same<T, half>::value) ? -65504.0f : -FLT_MAX; // 通过bid * vocab_size_padded + end_ids[bid]精确找到当前批次样本bid对应的end_id在logits数组中的位置，并将其替换为掩码值。 // 小于min_lengths[bid] 的时候，ends_ids[bid]位置的prob为0，不采样，鼓励多生成。 logits[bid * vocab_size_padded + end_ids[bid]] = mask_val; }}参考文献
• https://github.com/NVIDIA/TensorRT-LLM/blob/v0.5.0/cpp/tensorrt_llm/kernels/samplingPenaltyKernels.cu
