Shortfin llm beam search

[stbaione]

Initial implementation of beam_search for the LLM server. Putting it up as a draft for now to get some feedback on the code, and because it needs unit/integration tests. Too big/important of a change to merge without it.

The beam_search specific logic is contained with beam_manager.py, while generate.py just contains logic for managing the InferenceExecRequests and orchestrating the overall flow.

This also keeps all of the logic above the Batcher level, which minimized changes a lot from what I previously tried.

At a high level, the idea is that we:

1. Run prefill normally. Get the kvcache initialized and obtain first token.
2. Create n_beams InferenceExecRequests. Here we replicate the initial req that we used for prefill, including replicating the KVCaches pages.
3. Group each of these reqs under a BeamGroup. This is a helper class that handles actually performing beam_search token selection, and tracking the reqs.
4. Submit all reqs to batch, and wait for them all to finish.
5. For each req select the top_k tokens, based on cumulative log probability.
6. Select the overall top_k tokens, from all reqs.
7. Update our beams. Do any replication/beam collapses if needed.
8. If a req generates an eos_token, add it to the set of completed_reqs.
9. Repeat until either, max_completion_tokens is reached, or all beams generated an eos_token
10. When we return, we either select the highest cumulative_probability beam, or return all beams, depending on the request params.

Selecting top-k

Extra attention in this section would be appreciated.

For each beam:

<begin_loop>

1. Obtain logits from decode invocation.
2. Apply log_softmax to logits. By taking the log of softmax, we can use addition to track the cumulative probabilities, instead of multiplication with the raw probabilities. If we did multiplication with raw probabilities, our cumulative probabilities will become smaller and smaller, until we lose precision. source (search for 'From products to addition`)
3. Select the top_k values and tokens from the -1 axis. Track the cumulative log probs for each possible token.
   <end_loop>

We then return the top possible selections, in sorted order, based on which potential tokens would yield the beams with the highest probability.

[rsuderman]

We shouldn't be copying every page but only incomplete pages. I.e. the final page if it is incomplete. The rest should be tracked with the retain count system.

@renxida should be able to point to it

[stbaione]

Messaged @renxida and asked if he can provide context for the page copying. I think I see what you're saying though. If a page is full, we just need to read from it, not write to it. So, we should be able to share the full pages, which is all except the last one. That last non-full one is where we may actually see differences, so needs to be copied?

[rsuderman]

Most importantly when moving we want to use a topk algorithm and not sort and select topk. This is easy 20k+ elements for most models so even simple top-k implementations will outperform sorting the whole array. This can be included in the TODO above.

[stbaione]

This avoids sorting the 20k+ elements in the array, unless I missed something somewhere.

I use argpartition to ensure that the k last elements of the array are the largest. This runs in O(n) time. It doesn't guarantee that those k last elements are sorted though.
source

From there, I slice off those maximum k indices. Then I do an index view into logits to obtain topk_values.

So, len(topk_values) == k.

This argsort is just sorting the k top values (i.e. 4 top values), not values of the entire array.

Will double-check the shape in debugger

[stbaione]

Yeah, double-checked, len(topk_values) == k

[rsuderman]

This should be just handled by a default loader. The export json is just an example IR, our dataclass loaders should be resilient to missing optional values.

[rsuderman]

Beam search is a service level control so there is no benefit to including it during export.

[stbaione]

Will remove all reference in sharktank

[rsuderman]

These functions should be abstracted into a helper componsite class instead of integrated into generate. The idea should be that you create the class for the type of search you do.

E.g.

class DecodeManager():
  def __init__():
    pass
   
   def decode():
      ...
 
class GreedyDecoder():
    def decode():
       ....
       
       
class BeamSearchDecoder():
    def decode():
       ....

[rsuderman]

Essentially we want this class to define how someone would write different sampling methods and we can just dependency inject the different behavior into the generator.

[stbaione]

Great idea, will make this refactor

[rsuderman]

It concerns me that we see result_logits being retained / surviving. If we do want it to survive this buffer should be held via retainment and not copied.

[stbaione]

Could you elaborate on the issue and diff between retainment and coped? I think there's some background/context that I'm missing

[rsuderman]

Post selecting the K beams you want to continue down you should normalize the running decode. Float addition can accumulate error so if you find the minimum value and subtract that from all continued beams we can continue down the hypothesis train. There will need to be some correction for different length hypothesis to accommodate for them but it just means tracking what the accumulated normalization was.