fix: replace batched_count_greater_than to avoid dynamic shape TypeError on HPU

[kamil-kaczor]

Summary

Upstream vLLM decorates batched_count_greater_than with @torch.compile(dynamic=True), which causes Habana's recipe_compiler to raise TypeError: Cannot convert symbols to int when processing symbolic shapes. Additionally, mark_unbacked in the caller (gather_logprobs) prevents dynamic=False from being a viable alternative.

Fix

Replace with a plain (uncompiled) version of the same function. The patching is deferred to load_general_plugins time via a hook on vllm.plugins.load_general_plugins, because importing vllm.v1.sample.sampler during early plugin registration triggers a heavy import chain that interferes with platform initialisation.

Why deferred patching?

• Importing vllm.v1.sample.sampler during apply() (called from register()) triggers a heavy import chain that resets platform detection, causing Device string must not be empty.
• The patching hooks into load_general_plugins which runs in every process (parent + EngineCore subprocess) after the platform is ready.
• sampler.py uses from ... import batched_count_greater_than which creates a module-level global resolved via LOAD_GLOBAL at call time, so patching the module attribute works.

Testing

• test_skip_tokenizer_initialization PASSES
• test_engine_args (3 tests) PASS
• Inference with logprobs=5 produces correct output

[github-actions]

🚧 CI Blocked

The main CI workflow was not started for the following reason:

Your branch is behind the base branch. Please merge or rebase to get the latest changes.

[Copilot]

Pull request overview

This PR updates the HPU plugin’s runtime monkey-patches to avoid a Habana HPU recipe_compiler failure caused by upstream @torch.compile(dynamic=True) on vllm.v1.sample.ops.logprobs.batched_count_greater_than, by replacing it with an uncompiled equivalent and deferring the patch until after platform initialization.

Changes:

• Add an HPU-safe (non-compiled) replacement for batched_count_greater_than.
• Defer patch application by wrapping vllm.plugins.load_general_plugins so the heavy vllm.v1.sample.* import chain runs after platform init.
• Update the patch module documentation and slightly refactor existing monkey-patch assignments for readability.

[github-actions]

🚧 CI Blocked

The main CI workflow was not started for the following reason:

Your branch is behind the base branch. Please merge or rebase to get the latest changes.

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
54f548e9e58087f0155e4e164e416ad7efdfde6d

[github-actions]

✅ CI Passed

All checks passed successfully against the following vllm commit:
54f548e9e58087f0155e4e164e416ad7efdfde6d