Add custom logits processor API

[EricLBuehler]

New sampling order:

1. Apply penalties from sampling_params
2. ⭐ Apply these custom logits processors sequentially
3. Apply temperature and softmax
4. Sample the next token (topk, topp, minp, etc)

@Dan-wanna-M would this API work? You can find an example of using it here. Basically, you provide a Vec of Arc<dyn Fn(&Tensor, &[u32]) -> Result<Tensor> + Send + Sync>, which are applied.

Please let me know if this works for you or what changes would be necessary, and I can merge this PR. kbnf looks super cool!

[github-actions]

Code Metrics Report

  ===============================================================================
 Language            Files        Lines         Code     Comments       Blanks
===============================================================================
 C Header                2           35           28            0            7
 Dockerfile              1           34           25            0            9
 Happy                   1          442          369            0           73
 JSON                   11          102          101            0            1
 Python                 46         2018         1718           62          238
 TOML                   20          617          544           11           62
 YAML                    1            9            8            1            0
-------------------------------------------------------------------------------
 Jupyter Notebooks       4            0            0            0            0
 |- Markdown             2           77           32           31           14
 |- Python               2          196          169            1           26
 (Total)                            273          201           32           40
-------------------------------------------------------------------------------
 Markdown               28         1960            0         1481          479
 |- BASH                 5          101           98            0            3
 |- JSON                 1           12           12            0            0
 |- Python               5           92           82            0           10
 |- Rust                 6          408          365           19           24
 |- TOML                 2           75           63            0           12
 (Total)                           2648          620         1500          528
-------------------------------------------------------------------------------
 Rust                  195        59957        54396         1100         4461
 |- Markdown           100          906           13          843           50
 (Total)                          60863        54409         1943         4511
===============================================================================
 Total                 310        65174        57189         2655         5330
===============================================================================
  

[Dan-wanna-M]

@EricLBuehler One thing I am pondering is that kbnf is stateful, which means to use Arc<dyn Fn()> I need to use Mutex<T>. This requirement by itself is fine, but I am wondering its implications. For example, if logits processors will be directly shared across threads and executed in parallel, then some very bad stuff probably will happen.

[EricLBuehler]

@EricLBuehler One thing I am pondering is that kbnf is stateful, which means to use Arc<dyn Fn()> I need to use Mutex. This requirement by itself is fine, but I am wondering its implications. For example, if logits processors will be directly shared across threads and executed in parallel, then some very bad stuff probably will happen.

@Dan-wanna-M yes, we parallelize the sampling process. Perhaps we could change the API to accept some generic type which has some method to do the logits processing? Then, you could contain the state in that type in Mutex.

[EricLBuehler]

@Dan-wanna-M I just updated it so now the API is:

pub trait CustomLogitsProcessor: Send + Sync {
    /// Logits and sequence context (prompt and generated tokens), returning modified tokens.
    fn apply(&self, logits: &Tensor, context: &[u32]) -> Result<Tensor>;
}

In the example, we have a custom struct:

struct ThresholdLogitsProcessor {
    threshold: f64,
}

impl CustomLogitsProcessor for ThresholdLogitsProcessor {
    fn apply(&self, logits: &Tensor, _context: &[u32]) -> Result<Tensor> {
        // Mask is 1 for true, 0 for false.
        let mask = logits.ge(self.threshold)?;
        logits.broadcast_mul(&mask.to_dtype(logits.dtype())?)
    }
}

Would this work if you could put shared state into the custom struct w/ a mutex?

[Dan-wanna-M]

@Dan-wanna-M I just updated it so now the API is:

pub trait CustomLogitsProcessor: Send + Sync {
    /// Logits and sequence context (prompt and generated tokens), returning modified tokens.
    fn apply(&self, logits: &Tensor, context: &[u32]) -> Result<Tensor>;
}

In the example, we have a custom struct:

struct ThresholdLogitsProcessor {
    threshold: f64,
}

impl CustomLogitsProcessor for ThresholdLogitsProcessor {
    fn apply(&self, logits: &Tensor, _context: &[u32]) -> Result<Tensor> {
        // Mask is 1 for true, 0 for false.
        let mask = logits.ge(self.threshold)?;
        logits.broadcast_mul(&mask.to_dtype(logits.dtype())?)
    }
}

Would this work if you could put shared state into the custom struct w/ a mutex?

I spent some time rethinking it and I think as long as one generation only uses one logits processor then it should work(So a batch of N generation will need N separate logits processor)

[EricLBuehler]

@Dan-wanna-M I just updated it so now the API is:

pub trait CustomLogitsProcessor: Send + Sync {
    /// Logits and sequence context (prompt and generated tokens), returning modified tokens.
    fn apply(&self, logits: &Tensor, context: &[u32]) -> Result<Tensor>;
}

In the example, we have a custom struct:

struct ThresholdLogitsProcessor {
    threshold: f64,
}

impl CustomLogitsProcessor for ThresholdLogitsProcessor {
    fn apply(&self, logits: &Tensor, _context: &[u32]) -> Result<Tensor> {
        // Mask is 1 for true, 0 for false.
        let mask = logits.ge(self.threshold)?;
        logits.broadcast_mul(&mask.to_dtype(logits.dtype())?)
    }
}

Would this work if you could put shared state into the custom struct w/ a mutex?

I spent some time rethinking it and I think as long as one generation only uses one logits processor then it should work(So a batch of N generation will need N separate logits processor)

Sounds good, and this API implements CustomLogitsProcessor for the corresponding closure type! I'll merge this if there are no other problems.

[Dan-wanna-M]

@Dan-wanna-M I just updated it so now the API is:

pub trait CustomLogitsProcessor: Send + Sync {
    /// Logits and sequence context (prompt and generated tokens), returning modified tokens.
    fn apply(&self, logits: &Tensor, context: &[u32]) -> Result<Tensor>;
}

In the example, we have a custom struct:

struct ThresholdLogitsProcessor {
    threshold: f64,
}

impl CustomLogitsProcessor for ThresholdLogitsProcessor {
    fn apply(&self, logits: &Tensor, _context: &[u32]) -> Result<Tensor> {
        // Mask is 1 for true, 0 for false.
        let mask = logits.ge(self.threshold)?;
        logits.broadcast_mul(&mask.to_dtype(logits.dtype())?)
    }
}

Would this work if you could put shared state into the custom struct w/ a mutex?

I spent some time rethinking it and I think as long as one generation only uses one logits processor then it should work(So a batch of N generation will need N separate logits processor)

Sounds good, and this API implements CustomLogitsProcessor for the corresponding closure type! I'll merge this if there are no other problems.

Yes it should be mergeable now

[EricLBuehler]

Sounds good @Dan-wanna-M! Just merged it.