Add min_p sampling support (#14)

* feat: added min_p support from request to CUDA logit processor

fix: masksData offset for min_p handling

fix: min_p request constraint handling

chore: adjusted variable casing and comments

wip: move min_p as parameter to add_request

refactor: moved min_p log space conversion into async_exec

fix: added lnMinPs from entries to PostProcessorFn

* fix: dummy grammar, min_p in RequestParams, code cleanups

refactor: added min_p to RequestParams

chore: code cleanup (unnecessary import & serde default default)

Author: DwieDave

llgtrt/src/async_exec.rs

@@ -62,6 +62,7 @@ struct ReqData {
     // it seems to create them one by one
     // this array keeps track of assignment of req_id to llg state
     llg_infos: Vec<ConstraintInfo>,
+    min_p: f32,
     prompt_len: usize,
     is_run: bool,
 }
@@ -90,6 +91,7 @@ struct PendingSeq {
     prompt_len: usize,
     is_run: bool,
     entry: TlcLogitsEntry,
+    min_p: f32,
     stop: bool,
     // setting this will stop the sequence with given error
     error: Option<String>,
@@ -145,6 +147,11 @@ impl PendingSeq {
         let mask = step_res.sample_mask.as_ref().expect("No mask");
         self.entry.out_mask_pointer = copy_mask(mask);
         self.entry.temperature = llg.temperature;
+        self.entry.ln_min_p = if self.min_p > 0.0 {
+            self.min_p.ln()
+        } else {
+            -f32::MAX
+        };
 
         Ok(())
     }
@@ -160,6 +167,7 @@ impl PendingSeq {
             prompt_len: rd.prompt_len,
             entry: entry.clone(),
             stop: false,
+            min_p: rd.min_p,
             error: None,
             is_run: rd.is_run,
         }
@@ -276,6 +284,7 @@ extern "C" fn logits_processor(logits: *mut TlcLogitsEntry, num_logits: u32) {
             let entry = &mut entries[ps.entry_idx];
             entry.out_mask_pointer = ps.entry.out_mask_pointer;
             entry.temperature = ps.entry.temperature;
+            entry.ln_min_p = ps.entry.ln_min_p;
             let mut llg = ps.llg;
             if let Some(rd) = exec.req_data.get_mut(&entry.client_req_id()) {
                 if rd.logs.is_empty() {
@@ -448,6 +457,7 @@ impl AsyncExecutor {
                 llgs: llgs.into_iter().map(Some).collect(),
                 llg_infos: vec![],
                 prompt_len,
+                min_p: init.params.min_p,
                 logs: String::new(),
                 is_run,
             },

llgtrt/src/routes/completions.rs

@@ -103,6 +103,7 @@ fn req_params_from_openai(params: &CommonCreateParams) -> Result<RequestParams>
     let mut r = RequestParams {
         temperature: params.temperature,
         top_p: params.top_p,
+        min_p: params.min_p,
         max_new_tokens: params
             .max_completion_tokens
             .unwrap_or_else(|| params.max_tokens.unwrap_or(16)) as u32,
@@ -282,7 +283,16 @@ fn llg_grammar(params: &CommonCreateParams) -> Result<Option<TopLevelGrammar>> {
             log::debug!("using Lark grammar");
             lark_to_llguidance(lark_grammar)?
         }
-        _ => return Ok(None),
+        _ => {
+            if params.min_p > 0.0 {
+                // Returning a Dummy-grammar to enforce logit processing when min_p is set
+                let grm = TopLevelGrammar::from_regex(llguidance::api::RegexNode::Regex(
+                    r"(\n|.)*".to_string(),
+                ));
+                return Ok(Some(grm));
+            }
+            return Ok(None);
+        }
     };
     Ok(Some(grm))
 }

llgtrt/src/routes/openai.rs

@@ -165,6 +165,9 @@ pub struct CommonCreateParams {
     /// tokens comprising the top 10% probability mass are considered.
     #[serde(default = "default_top_p")]
     pub top_p: f32,
+    /// Filters out tokens with probability less than min_p multiplied by the probability of the most likely token
+    #[serde(default)]
+    pub min_p: f32,
     /// A unique identifier representing your end-user, which can help OpenAI to monitor and detect
     /// abuse.
     #[allow(dead_code)]

trtllm-c/logits.cpp

@@ -1,3 +1,4 @@
+#include <cfloat>
 #include <stdexcept>
 #include <string>
 #include <cmath>
@@ -34,7 +35,7 @@ void* tlc_alloc_logit_data(int32_t mask_stride_, int32_t max_batch_size_)
     assert(max_batch_size > 0);
     assert(mask_stride % 4 == 0);
 
-    size_t hd_size = max_batch_size * sizeof(int64_t) * 4;
+    size_t hd_size = max_batch_size * sizeof(int64_t) * 5;
     size_t sz2 = hd_size + max_batch_size * mask_stride;
     masks_size = sz2;
     if (cudaHostAlloc(&masksData, sz2, cudaHostAllocDefault))
@@ -46,7 +47,7 @@ void* tlc_alloc_logit_data(int32_t mask_stride_, int32_t max_batch_size_)
 
 float* tlc_mask_fraction_ptr()
 {
-    return (float*) ((uint8_t*) masksData + max_batch_size * sizeof(int64_t) * 3);
+    return (float*) ((uint8_t*) masksData + max_batch_size * sizeof(int64_t) * 4);
 }
 
 #define MAX_BATCH_SIZE 128
@@ -109,6 +110,7 @@ static void logitsPostProcessorFn(std::vector<tle::IdType> const& reqIds, std::v
         entry._num_tokens = tokens[i].get()[0].size();
         entry.out_mask_pointer = nullptr;
         entry.temperature = 1.0f;
+        entry.ln_min_p = -FLT_MAX;
         entries.push_back(entry);
 
         // auto shape = logits[i].getShape();
@@ -130,8 +132,10 @@ static void logitsPostProcessorFn(std::vector<tle::IdType> const& reqIds, std::v
     int64_t* logitPtrs = (int64_t*) masksData;
     int64_t* masksOffsets = logitPtrs + batchSize;
     float* temperatures = (float*) (logitPtrs + 2 * batchSize);
+    float* lnMinPs = (float*) (logitPtrs + 3 * batchSize);
 
     int64_t temperatures_offset = (uint8_t*) temperatures - (uint8_t*) masksData;
+    int64_t ln_min_p_offset = (uint8_t*) lnMinPs - (uint8_t*) masksData;
     int64_t mask_fractions_offset = (uint8_t*) tlc_mask_fraction_ptr() - (uint8_t*) masksData;
 
     int64_t* cudaLogitPtrs = (int64_t*) cudaMasksData;
@@ -187,6 +191,8 @@ static void logitsPostProcessorFn(std::vector<tle::IdType> const& reqIds, std::v
 
         masksOffsets[dp] = mask_offset;
         temperatures[dp] = entries[i].temperature;
+        lnMinPs[dp] = entries[i].ln_min_p;
+
 
         if (mask_offset > max_offset)
             max_offset = mask_offset;
@@ -201,8 +207,8 @@ static void logitsPostProcessorFn(std::vector<tle::IdType> const& reqIds, std::v
     if (dp > 0)
     {
         cudaMemcpyAsync(cudaMasksData, masksData, max_offset + mask_stride, cudaMemcpyHostToDevice, stream);
-        mask_logits_ext(cudaLogitPtrs, cudaMasksOffsets, mask_fractions_offset, temperatures_offset, dp, nVocab,
-            mask_stride / 4, tp, stream);
+        mask_logits_ext(cudaLogitPtrs, cudaMasksOffsets, mask_fractions_offset, temperatures_offset, ln_min_p_offset,
+            dp, nVocab, mask_stride / 4, tp, stream);
         cudaMemcpyAsync((uint8_t*) masksData + mask_fractions_offset, (uint8_t*) cudaMasksData + mask_fractions_offset,
             dp * sizeof(float), cudaMemcpyDeviceToHost, stream);
 

trtllm-c/mask_logits.cu

@@ -62,7 +62,7 @@ __inline__ __device__ void blockReduceMax2(T& val, int& idx, T flt_max)
 
 template <typename T>
 __global__ void mask_logits_kernel(T** logit_ptrs, int64_t* mask_offsets, size_t batch_size, size_t n_vocab,
-    size_t mask_stride, float* temperatures, T flt_max, float* mask_fractions)
+    size_t mask_stride, float* temperatures, float* ln_min_p, T flt_max, float* mask_fractions)
 {
     auto const batch_idx = blockIdx.x;
     auto logits_ptr = logit_ptrs[batch_idx];
@@ -135,6 +135,10 @@ __global__ void mask_logits_kernel(T** logit_ptrs, int64_t* mask_offsets, size_t
             else
             {
                 logit_adjusted = (logit - s_max_val_allowed) * beta;
+                if ((float) logit_adjusted < ln_min_p[batch_idx])
+                {
+                    logit_adjusted = -flt_max;
+                }
             }
         }
 
@@ -154,6 +158,7 @@ void mask_logits_ext(int64_t* d_logit_ptrs, // in,out [batch_size]
     int64_t* d_mask_offsets,                // in [int32_t,mask_stride], [batch_size]
     int64_t mask_fractions_offset,          // out, float, [batch_size]
     int64_t temperature_offset,             // in, float, [batch_size]; can be 0.0f for argmax
+    int64_t ln_min_p_offset,                // in, float, [batch_size]; log_e(min_p) for min_p > 0.0f, -FLT_MAX otherwise
     size_t batch_size,                      // current batch size
     size_t n_vocab,                         // vocab size
     size_t mask_stride,                     // n_vocab / 32 or thereabouts
@@ -167,10 +172,11 @@ void mask_logits_ext(int64_t* d_logit_ptrs, // in,out [batch_size]
 
     float* mask_fractions = reinterpret_cast<float*>((uint8_t*) d_logit_ptrs + mask_fractions_offset);
     float* temperatures = reinterpret_cast<float*>((uint8_t*) d_logit_ptrs + temperature_offset);
+    float* ln_min_ps = reinterpret_cast<float*>((uint8_t*) d_logit_ptrs + ln_min_p_offset);
 
 #define LAUNCH_KERNEL(T, m)                                                                                            \
-    mask_logits_kernel<T><<<grid, block, 0, stream>>>(                                                                 \
-        (T**) d_logit_ptrs, d_mask_offsets, batch_size, n_vocab, mask_stride, temperatures, m, mask_fractions)
+    mask_logits_kernel<T><<<grid, block, 0, stream>>>((T**) d_logit_ptrs, d_mask_offsets, batch_size, n_vocab,         \
+        mask_stride, temperatures, ln_min_ps, m, mask_fractions)
 
     switch (tp)
     {

trtllm-c/mask_logits.h

@@ -10,6 +10,7 @@ void mask_logits_ext(int64_t* d_logit_ptrs, // in,out [batch_size]
     int64_t* d_mask_offsets,                // in [int32_t,mask_stride], [batch_size]
     int64_t mask_fractions_offset,          // out, float, [batch_size]
     int64_t temperature_offset,             // in, float, [batch_size]; can be 0.0f for argmax
+    int64_t ln_min_p_offset,                // in, float, [batch_size]; log_e(min_p) for min_p > 0.0f, -FLT_MAX otherwise
     size_t batch_size,                      // current batch size
     size_t n_vocab,                         // vocab size
     size_t mask_stride,                     // n_vocab / 32 or thereabouts

trtllm-c/tlc.h

@@ -23,6 +23,8 @@ extern "C"
         uint32_t _num_tokens;
         // set by the callback (initially 1.0)
         float temperature;
+        // set by the callback (initially -FLT_MAX)
+        float ln_min_p;
         // set by the callback (initially NULL)
         uint32_t* out_mask_pointer;
     } TlcLogitsEntry;
@@ -116,6 +118,7 @@ extern "C"
         uint32_t eos_token_id;
         float temperature;
         float top_p;
+        float min_p;
         float frequency_penalty;
         float presence_penalty;
         float priority;
@@ -183,4 +186,4 @@ extern "C"
 }
 #endif
 
-#endif // TLC_H
+#endif // TLC_H

trtllm_rs/src/tlc.rs

@@ -55,6 +55,7 @@ impl Default for RequestParams {
             num_return_sequences: 1,
             temperature: f32::NAN,
             top_p: 1.0,
+            min_p: 0.0,
             presence_penalty: 0.0,
             frequency_penalty: 0.0,
             top_k: 0,