[Core] direct indexing on self.block_table_np in compute_slot_mapping #22940

linzebing · 2025-08-14T22:00:36Z

Purpose

Streamline slot mapping computation by replacing Torch tensor flattening and conversion with direct NumPy indexing via ravel, eliminating redundant copies and conversions

Test Plan

# vLLM Serving
VLLM_USE_V1=1 vllm serve facebook/opt-125m \
    --swap-space 16 \
    --disable-log-requests \
    --host :: \
    --dtype float16

# Capture traces
VLLM_USE_V1=1 vllm bench serve \
    --dataset-name random \
    --model facebook/opt-125m \
    --served-model-name facebook/opt-125m \
    --random-input-len 700 \
    --random-output-len 1 \
    --endpoint /v1/completions \
    --ignore-eos \
    --host localhost \
    --port 8000 \
    --num-prompts 100 \
    --profile

Also ran throughput benchmark test:

CUDA_VISIBLE_DEVICES=7 VLLM_USE_V1=1 python3 benchmarks/benchmark_throughput.py \
    --model facebook/opt-125m \
    --backend vllm \
    --input-len 800 \
    --output-len 75 \
    --num-prompts 30000

Test Result

Reduced compute_slot_mapping from 400+μs to 15-30μs.
Throughput improved 2.07% for opt-125m with input=800 and output=75

Before:
Throughput: 268.45 requests/s, 234888.85 total tokens/s, 20133.98 output tokens/s

After:
Throughput: 274.01 requests/s, 239750.19 total tokens/s, 20550.72 output tokens/s

(Optional) Documentation Update

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.

github-actions · 2025-08-14T22:00:43Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Code Review

This pull request introduces a performance optimization in compute_slot_mapping by replacing PyTorch tensor operations with direct NumPy indexing. The change from self.get_cpu_tensor().flatten()[...].numpy() to self.block_table_np.ravel()[...] is more direct, avoids unnecessary function calls and tensor conversions, and leverages NumPy's efficiency for this operation. The provided benchmarks confirm the significant performance improvement. The change is correct and well-justified.

linzebing · 2025-08-14T22:07:07Z

One-liner with significant perf improvement, @heheda12345 , @LucasWilkinson @houseroad and @njhill , can you take a look?

njhill

Wow, nice!

Signed-off-by: linzebing <linzebing1995@gmail.com>

Jialin · 2025-08-15T15:01:02Z

Nice! TIL

Jialin · 2025-08-15T15:07:19Z

vllm/v1/worker/block_table.py

        block_table_indices = (req_indices * self.max_num_blocks_per_req +
                               positions // self.block_size)
-        block_table_cpu = self.get_cpu_tensor()
-        block_numbers = block_table_cpu.flatten()[block_table_indices].numpy()


IIUC, the win is mainly coming from avoid tensor to np copy, while the tensor.flatten and np.ravel should have a similar performance?

yes tensor.flatten and np.ravel are similar. Looks like the main overhead is torch tensor's indexing is slow; the numpy conversion is relatively cheap

torch tensor's indexing is slow

Interesting. Ideally, the indexing cost should be similar as well. I guess the difference might come from

Actually indexing performance gap

Layout difference, so torch.Tensor invoked a copy while np still creating a view.

Jialin · 2025-08-15T15:08:27Z

vllm/v1/worker/block_table.py

        # block_size.
        block_table_indices = (req_indices * self.max_num_blocks_per_req +
                               positions // self.block_size)
-        block_table_cpu = self.get_cpu_tensor()


Maybe we could go and check other get_cpu_tensors reference to see if there's any other opportunities :)

Rest of the callsites are from tpu_model_runner.py, I don't spot anything obvious yet.

facebook-github-bot · 2025-08-15T20:16:48Z

@linzebing has imported this pull request. If you are a Meta employee, you can view this in D80367371.

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com> Signed-off-by: Yiwen Chen <yiwen66@berkeley.edu>

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com>

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com> Signed-off-by: Duncan Moss <djm.moss@gmail.com>

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com>

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com> Signed-off-by: Xiao Yu <xiao.yu@amd.com>

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com>

linzebing requested review from WoosukKwon, alexm-redhat, comaniac, njhill, robertgshaw2-redhat and ywang96 as code owners August 14, 2025 22:00

mergify bot added the v1 label Aug 14, 2025

gemini-code-assist bot reviewed Aug 14, 2025

View reviewed changes

linzebing changed the title ~~[Core] direct indexing on self.block_table_np~~ [Core] direct indexing on self.block_table_np for compute_slot_mapping Aug 14, 2025

linzebing changed the title ~~[Core] direct indexing on self.block_table_np for compute_slot_mapping~~ [Core] direct indexing on self.block_table_np in compute_slot_mapping Aug 14, 2025

njhill approved these changes Aug 14, 2025

View reviewed changes

njhill added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 14, 2025

[Core] direct indexing on self.block_table_np

ccc8065

Signed-off-by: linzebing <linzebing1995@gmail.com>

linzebing force-pushed the block_table branch from e2f773f to ccc8065 Compare August 15, 2025 12:43

Jialin approved these changes Aug 15, 2025

View reviewed changes

njhill merged commit 6e67077 into vllm-project:main Aug 15, 2025
39 checks passed

yiliu30 pushed a commit to yiliu30/vllm-fork that referenced this pull request Aug 19, 2025

[Core] direct indexing on self.block_table_np in compute_slot_mapping (…

8d44579

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com>

divakar-amd pushed a commit to divakar-amd/vllm_upstream that referenced this pull request Aug 20, 2025

[Core] direct indexing on self.block_table_np in compute_slot_mapping (…

0a09281

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com>

djmmoss pushed a commit to djmmoss/vllm that referenced this pull request Aug 21, 2025

[Core] direct indexing on self.block_table_np in compute_slot_mapping (…

612eab5

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com> Signed-off-by: Duncan Moss <djm.moss@gmail.com>

epwalsh pushed a commit to epwalsh/vllm that referenced this pull request Aug 28, 2025

[Core] direct indexing on self.block_table_np in compute_slot_mapping (…

f0b079f

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com>

xiao-llm pushed a commit to xiao-llm/vllm that referenced this pull request Aug 28, 2025

[Core] direct indexing on self.block_table_np in compute_slot_mapping (…

609b9c1

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com> Signed-off-by: Xiao Yu <xiao.yu@amd.com>

zhewenl pushed a commit to zhewenl/vllm that referenced this pull request Aug 28, 2025

[Core] direct indexing on self.block_table_np in compute_slot_mapping (…

70f2e88

…vllm-project#22940) Signed-off-by: linzebing <linzebing1995@gmail.com>

Uh oh!

[Core] direct indexing on self.block_table_np in compute_slot_mapping #22940

[Core] direct indexing on self.block_table_np in compute_slot_mapping #22940

Uh oh!

Conversation

linzebing commented Aug 14, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Aug 14, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

linzebing commented Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Jialin commented Aug 15, 2025

Uh oh!

Jialin Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

linzebing Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

Jialin Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

Jialin Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

linzebing Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

facebook-github-bot commented Aug 15, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

linzebing commented Aug 14, 2025 •

edited by github-actions bot

Loading

linzebing commented Aug 14, 2025 •

edited

Loading