[Models][Qwen3VL] Speedup fast_pos_embed_interpolate

[lgeiger]

Purpose

Followup on #25337 and #25347. fast_pos_embed_interpolate launches many small cuda ops so this PR slightly simplifies and optimised the implementation.
/cc @ywang96 @Isotr0py

Test Plan & Results

I verified that the new implementation doesn't change the computation.

A quick micro benchmark on an H100 shows a 15% speedup of fast_pos_embed_interpolate and I don't think it reduces readability of the code.

import torch
import torch.nn as nn


class Qwen3_VisionTransformer(nn.Module):
    def __init__(self, hidden_size, num_position_embeddings, spatial_merge_size):
        super().__init__()
        self.spatial_merge_size = spatial_merge_size
        self.num_grid_per_side = int(num_position_embeddings**0.5)
        self.dtype = torch.bfloat16
        self.device = torch.device("cuda:0")
        self.pos_embed = nn.Embedding(
            num_position_embeddings, hidden_size, dtype=self.dtype, device=self.device
        )

    def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:
        ...

    def bench_fast_pos_embed_interpolate(self, grid_thw: list[list[int]]):
        self.fast_pos_embed_interpolate(grid_thw)
        torch.cuda.synchronize()


model = Qwen3_VisionTransformer(
    hidden_size=1152, num_position_embeddings=2304, spatial_merge_size=2
)

grid_thw = [[1, 64, 48], [1, 64, 48], [1, 64, 48], [1, 64, 48], [1, 64, 48]]

model.bench_fast_pos_embed_interpolate(grid_thw)
%timeit model.bench_fast_pos_embed_interpolate(grid_thw)

main:    1.33 ms ± 30.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
This PR: 1.13 ms ± 3.85 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[lgeiger]

weights will already be on self.device so no need to copy it again.

[gemini-code-assist]

Code Review

This pull request introduces several well-reasoned optimizations to the fast_pos_embed_interpolate function in Qwen3VL. The changes, including refactoring weight calculations, vectorizing index computations, and using more efficient PyTorch operations like .sum() instead of unbind() followed by manual addition, are mathematically sound and contribute to the reported 11% performance improvement. The code is now more concise and idiomatic. I've reviewed the changes and found them to be correct and beneficial for performance without sacrificing readability. This is a solid improvement.

[lgeiger]

The grid_thw in the above benchmark came from an actual model request. I am now wondering whether it would make sense to cache the result for each t, h, w which in the above case for 5 images of the same shape would significantly speedup this code. I guess this could cause slightly higher memory usage but we could either put this cache on the CPU or at the very least handle the case where all items in the grid_thw list are the same.

[ywang96]

Thanks for the contribution!

Re: caching t, h, w - I had the same idea when I first cleaned up the code here that we could build a small cache on CPU for this, but I also wonder whether the h2d cost is worth the effort

[lgeiger]

Thanks for the fast review. I also updated the tiling logic in 24b6717 which slightly improves things further. I updated the benchmarks above.