Optimizing for CPUs

[AntonOresten]

While the on-GPU sampling is neat, moving logits to the CPU might still be faster.

One idea would be to have logitsample(f::Function, logits), falling back to logitsamplel(f(logits)), with specialized methods like logitsample(::Top_pk, logits) with better time complexity using a partial sort.

Some rough benchmarks show that logitsample ∘ Top_p(0.5) on 100k logits takes ~2 micro milliseconds on an A6000, which sets an upper limit on inference speed.

[murrellb]

I'm ok with a speed limit of 500000tok/s.

[AntonOresten]

Oops, I meant milliseconds. Seems the GPU I was using was already in use though, and Top_p is closer to 700 microseconds, with Top_nσ being faster at 250 microseconds.

[murrellb]

Fine then I'll have to live with a speed limit of 1428toks/s.