Add kernels for Longcat2

[LinyuanLi0046]

Motivation

This PR adds several Ascend NPU kernels required by LongCat Pro inference, mainly covering over-embedding n-gram id generation, OE token table update, and MLP lightning indexer support.

Changes

• Add compute_n_gram_ids kernel.
  Computes n-gram ids for LongCat Pro over-embedding.
• Add update_oe_token_table kernel.
  Updates the OE token table for batched requests.
• Add mlp_lightning_indexer kernel.
  Support LongCat Indexer, which supports kv block length and local/init windows.
• Add Python tests.

Testing

Added the following test files:

• tests/python/sgl_kernel_npu/test_compute_n_gram_ids.py
• tests/python/sgl_kernel_npu/test_update_oe_token_table.py
• tests/python/sgl_kernel_npu/test_mlp_lightning_indexer.py

[gemini-code-assist]

Code Review

This pull request introduces three new NPU operators for sglang: compute_n_gram_ids, update_oe_token_table, and mlp_lightning_indexer, along with their corresponding host tiling, kernel implementations, CMake configurations, PyTorch bindings, and Python tests. The code review identified several critical issues that need to be addressed: in update_oe_token_table, BLOCK_SIZE must be corrected to 8 to prevent out-of-bounds memory access, and the kernel must handle empty ignore_tokens gracefully to avoid runtime crashes and redundant operations. Additionally, static tiling caches in both compute_n_gram_ids and mlp_lightning_indexer require mutex synchronization to prevent race conditions in multi-threaded environments. Finally, a performance bottleneck in compute_n_gram_ids caused by direct global memory reads in nested loops should be optimized using Unified Buffer copies, and unused duplicate tiling files should be deleted.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

BLOCK_SIZE is defined as 32 (elements), but the host tiling code aligns 'ub_factor' to 32 bytes (which is 8 elements of int32_t). If 'count' is close to 'ub_factor', aligning 'count' to 32 elements can result in an 'alignCnt' that exceeds 'ub_factor', causing out-of-bounds reads/writes on the LocalTensors. Since Ascend C DataCopy requires 32-byte alignment, the element alignment should be 8.

    constexpr static int64_t BLOCK_SIZE = 8;

[gemini-code-assist]

If 'ignore_tokens' is empty, 'ignoreTokenNum_' is 0, which causes 'pipe_->InitBuffer' to allocate 0 bytes for 'isIgnoreQue_'. This is invalid and causes runtime crashes on the NPU. Ensure a minimum allocation size of at least 1 element.

        pipe_->InitBuffer(isIgnoreQue_, 1, std::max(ignoreTokenNum_, 1U) * sizeof(int32_t));

[gemini-code-assist]

Bypass loading and dequeuing ignore tokens when 'ignoreTokenNum_' is 0 to avoid empty DataCopy and queue operations.

        if (ignoreTokenNum_ > 0) {\n            LoadIgnore();\n        }\n        auto reqLensTensor = reqLensQue_.DeQue<int32_t>();\n        LocalTensor<int32_t> isIgnoreTensor;\n        if (ignoreTokenNum_ > 0) {\n            isIgnoreTensor = isIgnoreQue_.DeQue<int32_t>();\n        }

[gemini-code-assist]

Only free 'isIgnoreTensor' if it was allocated (i.e., when 'ignoreTokenNum_ > 0').

        if (ignoreTokenNum_ > 0) {\n            isIgnoreQue_.FreeTensor(isIgnoreTensor);\n        }

[gemini-code-assist]

Bypass the ignore token comparison loop entirely when 'ignoreTokenNum_ == 0' to avoid redundant operations and potential issues with unallocated tensors.

    __aicore__ inline void IsIgnoreToken(const LocalTensor<int32_t> &isIgnoreTensor, uint32_t count)\n    {\n        auto tokenTensor = tokenQue_.DeQue<int32_t>();\n        auto alignCnt = CeilDiv(count, BLOCK_SIZE) * BLOCK_SIZE;\n        if (ignoreTokenNum_ > 0) {\n            auto minusTensor = minusQue_.Get<int32_t>();\n            Duplicate(minusTensor, static_cast<int32_t>(-1), count);\n            PipeBarrier<PIPE_V>();\n            for (int64_t idx = 0; idx < ignoreTokenNum_; idx++) {

[gemini-code-assist]

The static tiling cache variables are accessed and modified without synchronization. In a multi-threaded PyTorch environment, this can lead to race conditions and memory corruption. Protect the cache access with a local static mutex.

at::Tensor GetOrCreateCachedTilingTensor(const ComputeNGramIdsTilingData &tiling_data,\n                                         const at::Tensor &device_anchor)\n{\n    static std::mutex tiling_cache_mutex;\n    std::lock_guard<std::mutex> lock(tiling_cache_mutex);\n    int32_t tiling_size =

[gemini-code-assist]

The static tiling cache variables are accessed and modified without synchronization. Protect the cache access with a local static mutex to ensure thread safety.

at::Tensor GetOrCreateCachedTilingTensor(const LITilingData &tilingData, const at::Tensor &device_anchor)\n{\n    static std::mutex tiling_cache_mutex;\n    std::lock_guard<std::mutex> lock(tiling_cache_mutex);\n    const int64_t tiling_size = static_cast<int64_t>(sizeof(LITilingData));

[gemini-code-assist]

Reading from 'oeTokenTableGm' directly using 'GetValue' inside nested loops bypasses the Unified Buffer (UB) and triggers high-latency global memory reads for every single token. This will severely bottleneck kernel performance on the NPU. Consider copying the relevant slice of 'oeTokenTableGm' for the current request into a LocalTensor (UB) using 'DataCopy' before starting the token loop.

[gemini-code-assist]

The files 'lightning_indexer_tiling.cpp' and 'lightning_indexer_tiling.h' under 'csrc/mlp_lightning_indexer/op_host/tiling/' are completely unused and not compiled by 'csrc/CMakeLists.txt' (which compiles 'mlp_lightning_indexer_tiling.cpp' instead). These files appear to be leftover duplicates from copying the 'lightning_indexer' operator. Please delete them to keep the codebase clean and maintainable.