Fix guided sampling with outlines

[tae-su-kim]

This is a rebase of PR #153 to habana_main due to the deprecation of habana_next.

Current habana_main includes guided decoding related code from vllm, and the feature is already there in the openAI api endpoint. However, guided decoding currently fails to run with following error:

...
 File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1535, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1585, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/workspace/codes/vllm/model_executor/layers/sampler.py", line 138, in forward
    sample_results, maybe_sampled_tokens_tensor = _sample(
  File "/workspace/codes/vllm/model_executor/layers/sampler.py", line 711, in _sample
    return _sample_with_torch(
  File "/workspace/codes/vllm/model_executor/layers/sampler.py", line 592, in _sample_with_torch
    sample_results = _greedy_sample(seq_groups, greedy_samples)
  File "/workspace/codes/vllm/model_executor/layers/sampler.py", line 336, in _greedy_sample
    samples_lst = samples.tolist()
RuntimeError: synNodeCreateWithId failed for node: strided_insert with synStatus 1 [Invalid argument]. .

This PR suggests to use masked_fill rather than _add for the masking process of guided decode. With this PR, openai endpoint supports guided decoding. For example,

Input:

payload = {
        "model": "/models/Meta-Llama-3-8B-Instruct",
        "messages": [
            {"role": "user", "content": "reply negatively."}
        ],
        "best_of": best_of,
        "use_beam_search": use_beam_search,
        "temperature": 0.0,
        "top_p": 1.0,
        "guided_regex": "[Pp]ositive format |[Nn]egative format",
}

Output:

{'id': 'cmpl-f3e792eb0197492a8d7eec4bb9916936', 'object': 'chat.completion', 'created': 1722847036, 'model': '/models/Meta-Llama-3-8B-Instruct', 'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'Negative format'}, 'logprobs': None, 'finish_reason': 'stop', 'stop_reason': None}], 'usage': {'prompt_tokens': 14, 'total_tokens': 21, 'completion_tokens': 7}}

[tae-su-kim]

We fixed abnormal latency overhead with commit 6d57c18.

Rough benchmark is as follows:

Version	Prefill Latency	Decode Latency
w/o guided_decode	53.1s	962.6s
commit 4bea4e3	156.9s	1322.5s
commit 6d57c18	53.2s	1149.5s

Setup: llama-3-8b, greedy sampling, max_num_seqs 256, 1k request, QPS -1
prefill_latency: latency(1k random input tokens / 1 output token)
decode_latency: latency(1k random input tokens / 1k output tokens without eos) - prefill_latency

[michalkuligowski]

Ruff static analysis found that "cache" imported at the top of file is now not used, please remove the import

[tae-su-kim]

I am benchmarking the performance of .add versus .masked_fill to determine which has better throughput. I'll resolve the conflict based on the benchmark and fix the ruff issue before merging. I'll let you know once it's completed. Apologies for the delayed response!

[tae-su-kim]

@michalkuligowski I removed the import and resolved the conflict based on benchmark results. Both .add and .masked_fill based implementations significantly degrade decode throughput, so I opted to stick with the current code. If any optimizations are possible, I will open another PR.

[michalkuligowski]

Ruff fails on unsorted imports. BTW is outlines version bump required now after those changes?

[tae-su-kim]

There are some marginal throughput difference, but I think most of the updates in this PR is already here. I will close this PR for now. Thank you!