[Platform] Add current_platform.num_compute_units interface

[jikunshang]

Purpose

there are some torch.cuda.get_device_properties().multi_processor_count across vllm code base. we can unify it into current_platform.num_compute_units interface to make it clean and extensible for non-cuda hardware like xpu and npu.

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[dosubot]

Related Documentation

Checked 0 published document(s) in 1 knowledge base(s). No updates required.

^{How did I do? Any feedback?}  

[gemini-code-assist]

Code Review

The pull request introduces a new get_num_sm interface to abstract the retrieval of SM (Streaming Multiprocessor) counts across different platforms (CUDA, ROCm, XPU). This change unifies the way SM counts are obtained, making the codebase cleaner and more extensible for non-CUDA hardware. The existing torch.cuda.get_device_properties().multi_processor_count calls are replaced with the new platform-agnostic interface. This is a good step towards improving platform independence and maintainability.

[njhill]

Thanks @jikunshang! LGTM but wonder if we should name it something more generic since SM is a CUDA term AFAIK

Maybe get_multi_processor_count?

[mergify]

Hi @jikunshang, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[jikunshang]

@njhill seems Compute Unit is unified term across GPU vendor https://chatgpt.com/share/699a69bd-1fd0-8002-aeb8-e12bf149533e
I prefer to rename to get_num_compute_units, thoughts?

[mergify]

Hi @jikunshang, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[njhill]

I prefer to rename to get_num_compute_units, thoughts?

How about just compute_unit_count, since there is already device_count?

Or if not then I think just num_compute_units would be better

[njhill]

Probably doesn't make sense to change in places like this which are cuda-specific anyway

[jikunshang]

I think it's ok to always use this current_platfrom API here. we should avoid/reduce using torch.cuda APIs since we are proposing this RFC(#30679). (though non-cuda platform will never fall into current code path)
my understanding here is it will use property to check cuda capability later. I feel we can also refactor this part into something like current_platform.supported_arch

[njhill]

I just think in places like this we are already using torch.cuda.get_device_properties (line above) so it's better to keep the code consistent (like you say, it can always be refactored for portability in future if/when appropriate).

Similarly if it's in cuda or rocm or xpu-specific files/code then I don't think there's a need to use the current_platform version, and actually maybe better not to since it implies that code could be cross-platform which could be misleading.

[jikunshang]

got it. reverted:)

[jikunshang]

I prefer to rename to get_num_compute_units, thoughts?

How about just compute_unit_count, since there is already device_count?

Or if not then I think just num_compute_units would be better

I feel count is for something countable and you can control use it or not.. while we always want to use all compute units. so let's use num_compute_units :)

[njhill]

Also remove platform_utils.get_cu_count() method now?

[jikunshang]

Also remove platform_utils.get_cu_count() method now?

nice catch, done!

[robertgshaw2-redhat]

did this PR break Distributed Tests (2 GPus)?

• https://buildkite.com/vllm/ci/builds/53087#019c941c-f14a-4166-b576-2101c7621154

[jikunshang]

did this PR break Distributed Tests (2 GPus)?

• https://buildkite.com/vllm/ci/builds/53087#019c941c-f14a-4166-b576-2101c7621154

oh yes. sorry for that and thanks @LucasWilkinson for fixing