[perf] Avoid dtype promotion sync in mamba_get_block_table_tensor

[hl475]

Purpose

From the perf profiling on a model using linear attention, one thing draws the attention is aten::to from somewhere in vllm gpu_model_runner preprocess

For this aten::to,

We can tell it is doing to with dtype=3 (int64)

With CC's help, we found it is from mamba_get_block_table_tensor

• start_indices is int32
• offsets is int64
• indices_to_gather = start_indices.unsqueeze(1) + offsets will perform int32+int64 and sync

Now we made offsets as int32 as well, and then no sync.

With the change, new profiling looks much better

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.

[gemini-code-assist]

Code Review

This pull request introduces a performance optimization within the mamba_get_block_table_tensor function. The change modifies the creation of the offsets tensor to use torch.int32, which avoids a dtype promotion to int64 during the subsequent addition operation. The final result is then explicitly cast to torch.int64 to meet the requirements of torch.gather. This is a sound optimization that should reduce overhead as intended. The implementation appears correct and I found no issues.

[heheda12345]

Nice finding! Thanks for your contribution.

[hl475]

CI failures are probably from #33600

[hl475]

failure from entrypoints/openai/responses/test_harmony.py::test_mcp_code_interpreter_streaming[openai/gpt-oss-20b] should not be relevant - https://buildkite.com/organizations/vllm/analytics/suites/ci-1/tests/e82e9b97-3f4c-89ab-933c-7e077bedf38c?branch=main&query=test_mcp_code_interpreter_streaming is red since yesterday.