add mla_preprocess kernel

[kiscad]

What this PR does / why we need it?

• Adds the mla_preprocess custom kernel to provide an optimized pre-processing operator for Multi-head Latent Attention (MLA) on Ascend NPUs.
• Wires the new kernel into the C++ extension pipeline so vLLM can invoke it directly, cutting Python-side tensor shuffling and memory copies that previously bottlenecked MLA compilation paths.

Does this PR introduce any user-facing change?

• No. The change only introduces a low-level kernel; public APIs and inference behavior remain unchanged.

How was this patch tested?

• Dedicated Ascend kernels are not covered by our CI yet, so no extra automated tests were added. Future MLA-focused regression runs will cover this path.

• vLLM version: v0.10.2

• vLLM main: vllm-project/vllm@releases/v0.11.0

[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.

[gemini-code-assist]

Code Review

This pull request introduces a new mla_preprocess kernel, which involves significant changes across the host-side tiling logic, kernel implementation, and PyTorch bindings. The implementation is complex and adds a substantial amount of new code. My review focuses on critical aspects of correctness, performance, and thread safety. I've identified a critical race condition in the host-side tiling logic that needs to be addressed to prevent data corruption in multi-threaded environments. Additionally, there are opportunities to improve kernel performance by optimizing memory copy operations and to enhance correctness and performance on the host by replacing floating-point calculations with integer arithmetic for tiling parameters.

[gemini-code-assist]

The use of a static local variable global_tiling_data to store tiling information creates a potential race condition. If multiple threads call mla_preprocess_tiling concurrently with the same number of tokens N, they will calculate the same bIndex and attempt to write to the same memory location in global_tiling_data simultaneously. This can lead to corrupted data and unpredictable kernel behavior. To ensure thread safety, this buffer should be managed in a thread-safe manner, for example by allocating it as part of a workspace that is unique per call, rather than using a shared static buffer.

[gemini-code-assist]

The calculation of roundBase uses floating-point functions (pow, ceil, log) for what appears to be an integer calculation to find the next power of two. This approach can be slow and may suffer from floating-point precision issues, potentially leading to incorrect tiling parameters. It is recommended to replace this with integer-based bit manipulation for both performance and correctness. The use of log (natural logarithm) is also likely incorrect if the intent is to work with powers of two; log2 would be the mathematical function to use, but bit manipulation avoids this entirely.

For example, to find the smallest power of two greater than or equal to an integer x, you can use a bit-twiddling algorithm:

// Example of a bit-twiddling algorithm to find the next power of 2
uint32_t next_power_of_2(uint32_t n) {
    if (n == 0) return 1;
    n--;
    n |= n >> 1;
    n |= n >> 2;
    n |= n >> 4;
    n |= n >> 8;
    n |= n >> 16;
    n++;
    return n;
}

// The original line could be refactored to use such a helper function.
uint32_t val = CeilDiv(priFlag ? n_ : m_, CONST_16);
uint32_t next_pow2_val = next_power_of_2(val);
uint32_t roundBase = next_pow2_val * CONST_16;

[gemini-code-assist]

The kernel copies the MlaTilingData structure from global memory to the stack field by field. This results in numerous small, individual memory accesses from global memory, which is highly inefficient and can significantly degrade kernel performance. A single bulk memory copy of the entire MlaTilingData structure into a UB buffer, and then a memcpy to the stack variable (or direct use from UB), would be much more efficient.

Example of a more efficient copy:

MlaTilingData mlaTilingData;
// Allocate a temporary buffer in UB
AscendC::TBuf<AscendC::TPosition::VECCALC> tilingUB;
AscendC::TPipe pipe;
pipe.InitBuffer(tilingUB, sizeof(MlaTilingData));

// Copy from GM to UB
AscendC::DataCopy(tilingUB.Get<uint8_t>(), reinterpret_cast<__gm__ uint8_t*>(tiling), sizeof(MlaTilingData));
pipe.Barrier();

// Copy from UB to stack
memcpy(&mlaTilingData, tilingUB.Get<uint8_t>().GetPhyAddr(), sizeof(MlaTilingData));

[gemini-code-assist]

The mla_preprocess_impl function is called with sin_ptr and cos_ptr passed for both (sin1, cos1) and (sin2, cos2) arguments. This suggests a potential mismatch between the caller's intent and the kernel's expectation, or that the kernel signature is redundant. If sin2 and cos2 are intended to be different from sin1 and cos1, this is a bug. If they are always the same, the kernel signature and the host-side function should be simplified to accept only one set of sin/cos tensors to avoid confusion and potential future errors.