HPU: offload logits processing to CPU

[madamczykhabana]

Due to high dynamicity on logits processing it's better to offload it completely to CPU instead of computing it on HPU.

[madamczykhabana]

@tae-su-kim , @huijjj
Please check on your end if that PR helps with structured output performance on HPU. In my local testing I saw performance increase from ~60 tok/sec to ~1100 tok/sec on 1k/1k scenario.

Note that when cherry-picking it to v0.8 use vllm/model_executor/guided_decoding/outlines_logits_processors.py from this PR and discard any previous modifications as git by default will leave with_mark_steps and other changes (and calling it for every sample in a batch hurts performance)

[tae-su-kim]

@madamczykhabana Hi, sorry for the delayed response. We are trying to reproduce the number in both v0.8 branch of ours and habana_main, and current PR seems to fail in habana_main. It collides with commit 7c7714d where MQLLMEngine was introduced.

  File "/workspace/codes/vllm/engine/multiprocessing/client.py", line 552, in _process_request
    lp_bytes = cloudpickle.dumps(logits_processors)
  File "/usr/local/lib/python3.10/dist-packages/cloudpickle/cloudpickle.py", line 1479, in dumps
    cp.dump(obj)
  File "/usr/local/lib/python3.10/dist-packages/cloudpickle/cloudpickle.py", line 1245, in dump
    return super().dump(obj)
_pickle.PicklingError: Can't pickle <functools._lru_cache_wrapper object at 0x7fbc567d5010>: it's not the same object as vllm.model_executor.guided_decoding.outlines_logits_processors.BaseLogitsProcessor._get_mask_tensor

It seems like new multiprocessing engine pickles logit processor information using cloudpickle, but cloudpickle is unable to pickle functions decorated with lru_cache (cloudpipe/cloudpickle#178).

I can confirm that this PR works in v0.8 branch of ours (based on vllm-fork v0.5.3) and improves e2e throghput as follows:
(3.1-8B-Instruct, fixed 1K/1K benchmark) 917 tokens/sec to 1233 tokens/sec
(3.1-8B-Instruct, dynamic 1K/1K benchmark) 59 tokens/sec to 1027 tokens/sec

[madamczykhabana]

@tae-su-kim thanks for the info. I've just pushed a workaround for the pickling error. Please check if it helps on your end.
Also, I'm seeing some weird behavior with structured output on habana_main, even without this PR. I'm still trying to figure out why I see a a big discrepancy in number of generated tokens.

[tae-su-kim]

@madamczykhabana Thanks for the prompt update! I will benchmark with multiprocessing again. Can you elaborate more on the discrepancy in number of generated tokens issue? From our side, we observe similar number of tokens: e.g. 952,171 tokens from A100, 959,993 tokens from Gaudi-2 on 1K/1K dynamic benchmark with json guide. We can share our setup if you want.

[madamczykhabana]

@tae-su-kim Sorry for the delay. The token discrepancy was when I was comparing guided-json to non-guided-json. The PR has been rebased and it's ready for review.

[michalkuligowski]

How many masks are expected at maximu? Since _create_mask_tensor is limited and the _cached fun is not, won't we get into a situation where we will get None result for least recently used (128+)?

[michalkuligowski]

You can ignore my last comment, the lrucache maxsize does not have effect here, since _cached cache keeps objects.
If you reverse the condition in line 39 the else can be replaced with return cache[args]
if args not in cache:
cache[args] = fn(*args)
return cache[args]

[madamczykhabana]

Yeah, but at the same time we'd need to calculate the hash and access the map one additional time.

1. if args in cache
2. cache[args] = fn(*args)
3. return cache[args]

In current version you access it twice:

1. args in cache
2. cache[args] = result

To be honest I didn't measure how big of a perf impact this might be, but I'd like to be on the safe side.