[Op] DeepSeekV3.2 support bmm_transpose operator

[ZYang6263]

What this PR does / why we need it?

DeepSeekV3.2 support bmm_transpose operator.

Does this PR introduce any user-facing change?

How was this patch tested?

• vLLM version: v0.12.0
• vLLM main: vllm-project/vllm@ad32e3e

[gemini-code-assist]

Code Review

This pull request introduces a new bmm_transpose operator to replace torch_npu.npu_transpose_batchmatmul in _v_up_proj. While the Python-side change is straightforward, my review of the C++ implementation for the new operator (torch.ops._C_ascend.batch_matmul_transpose) revealed critical issues related to out-of-bounds memory access and multi-device safety. These issues, detailed in the line comment, could lead to runtime errors and must be addressed before this change is merged.

[gemini-code-assist]

The C++ implementation of the batch_matmul_transpose operator has a couple of critical issues that could lead to runtime errors.

1. Potential for out-of-bounds access: In csrc/batch_matmul_transpose/op_host/batch_matmul_transpose.h, the tiling data for different batch sizes is cached in a static array global_tiling_data of size MAX_CAPTURE_NUM (1024). The index into this array, batchIdx, is derived from the number of tokens (opShape.m). If the number of tokens is greater than 1024, which is common during prefill, this will result in an out-of-bounds access and a runtime error.

   // csrc/batch_matmul_transpose/op_host/batch_matmul_transpose.h:104
   int32_t batchIdx = opShape.m - 1;
   ...
   if (batchIdx >= 0 && batchIdx < MAX_CAPTURE_NUM) { // MAX_CAPTURE_NUM is 1024
       ...
   } else {
       TORCH_CHECK(false, "batchIdx is out of range: ", batchIdx);
   }

2. Not safe for multi-device execution: The global_tiling_data is a static variable, meaning it's initialized only once. Its device is set to the device of the input tensor from the first call. In a multi-GPU environment where workers on different devices might call this operator, this will cause device mismatch errors for any call not on the initial device.

   // csrc/batch_matmul_transpose/op_host/batch_matmul_transpose.h:106
   static auto global_tiling_data = at::empty(
       {tilingSize * MAX_CAPTURE_NUM}, at::TensorOptions().dtype(at::kByte).device(tensor_a.options().device()));

   A potential solution for the multi-device issue is to use a thread-safe map from device index to the tiling data tensor, for example using std::map with a std::mutex.

Given these issues, the underlying implementation of this operator needs to be revised before it can be safely used.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[hust17yixuan]

this ops may have some problem in high level cann version, considering merge lately

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.