[Feature][Operator] Add DeepSeek V4.1 sparse attention operators

[GDzhu01]

What this PR does / why we need it?

Adds the native and Triton operator layer required by DeepSeek V4.1 on Ascend:

• Quant Lightning Indexer V2 and metadata operators
• Sparse Flash MLA and metadata operators
• compressor, index preparation, query quantization, Engram INT8, and required MoE operator updates
• operator documentation, native examples, static checks, and operator-level unit/e2e tests

This is the operator-only first part of the DeepSeek V4.1 enablement. Framework/model integration is intentionally split into dependent PR #16423.

Does this PR introduce any user-facing change?

Yes. It exposes the native operator capabilities required by DeepSeek V4.1 serving, but does not register or enable the model integration by itself.

How was this patch tested?

• operator static checks: 8 passed
• focused related unit tests: 342 passed, 1 skipped
• broad related unit tests: 855 passed, 12 skipped
• clean native rebuild on Ascend A3: passed
• final single-card NPU operator regression: 359 passed
• pre-commit scoped checks: Ruff, Ruff format, codespell, typos, clang-format, markdownlint, and forbidden-import check passed

End-to-end model validation is recorded in dependent PR #16423.

• vLLM main: vllm-project/vllm@a97dacb

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request adds the necessary operator infrastructure to support the DeepSeek V4.1 model on Ascend NPU hardware. It focuses on implementing high-performance native and Triton kernels for sparse attention and quantization operations. This is the first phase of the model enablement, providing the required operator layer while deferring full framework integration to a subsequent pull request.

Highlights

• DeepSeek V4.1 Operator Support: Introduced native and Triton operator layers for DeepSeek V4.1 on Ascend, including Quant Lightning Indexer V2, Sparse Flash MLA, and various MoE-related operators.
• Operator Testing and Validation: Added comprehensive operator-level unit and end-to-end tests, along with documentation and static checks to ensure operator stability.
• Performance Benchmarking: Included new benchmark scripts for indexer index preparation and query quantization to optimize operator performance on NPU.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

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

---

Tip

💡 Consider Linking a Related Issue or RFC

Your PR title contains the [Feature] tag, indicating a bug fix or new feature.

Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:

• Fixes #<issue_number>
• Closes #<issue_number>
• Resolves #<issue_number>
• Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

[gemini-code-assist]

Code Review

This pull request introduces a new two-level TopK candidate selection mechanism for the QuantLightningIndexerV2 operator, specifically optimized for the Ascend910B (arch22) architecture. The changes include adding support for TND layout, non-contiguous key strides, and an output index offset feature. The PR also includes comprehensive test assets and infrastructure updates to support batch consistency testing and performance profiling. My review identified several critical issues, including incorrect API usage for DataCopyPad, a compilation error due to a variable naming mismatch in the invocation macro, missing type-fixing wrappers for special quantization modes, and a potential crash in the InferShape implementation when handling optional outputs.

[gemini-code-assist]

In Ascend C, DataCopyPad expects DataCopyPadParams (or DataCopyPadExtParams), whereas DataCopy expects DataCopyExtParams. Since dataCopyParams is declared as DataCopyExtParams, calling DataCopyPad with it will result in a compilation error or undefined behavior due to struct layout mismatch. Please use AscendC::DataCopy instead of AscendC::DataCopyPad when copying the tail.

            if (tailSize > 0) {
                AscendC::DataCopyExtParams dataCopyParams;
                dataCopyParams.blockCount = 1;
                dataCopyParams.blockLen = tailSize * sizeof(T);
                dataCopyParams.srcStride = 0;
                dataCopyParams.dstStride = 0;
                AscendC::DataCopy(outGm[gmOffset], popBuffer, dataCopyParams);
            }

[gemini-code-assist]

The parameter in the quant_lightning_indexer_v2 function signature is named workspace, but the macro INVOKE_LI_CANDIDATE_OP_IMPL passes user to op.Init. This will cause a compilation error because user is not declared in this scope. Please change user to workspace.

#define INVOKE_LI_CANDIDATE_OP_IMPL(templateClass, ...) \
    do { \
        templateClass<QLIV2Type<__VA_ARGS__>> op; \
        op.Init(query, key, weights, queryScale, keyScale, cuSeqlensQ, cuSeqlensK, sequsedQ, sequsedK, cmpResidualK, \
                blockTable, outputIdxOffset, metadata, candidateTopkIndex, sparseIndices, sparseValues, \
                candidateTopkIndexOut, workspace, tiling_data, &tPipe); \
        op.Process(); \
    } while (0)

[gemini-code-assist]

In quant_lightning_indexer_v2_torch_adpt.h, query, key, queryDequantScale, and keyDequantScale are passed directly to EXEC_NPU_CMD. However, for special quantization modes (like HIFLOAT8, FLOAT4_E2M1, and MXFP8), their ACL data types must be fixed using FixQLIV2AclDtypes and MakeWrapper (as done in quant_lightning_indexer.cpp), otherwise EXEC_NPU_CMD will use default PyTorch scalar type mappings (e.g., ACL_UINT8 instead of ACL_HIFLOAT8), causing runtime errors or incorrect execution. Please apply the type fixing wrappers before calling EXEC_NPU_CMD.

    auto queryWrapper = at_npu::native::TensorWrapper{query, at_npu::native::ConvertToAclDataType(query.scalar_type())};
    auto keyWrapper = at_npu::native::TensorWrapper{key, at_npu::native::ConvertToAclDataType(key.scalar_type())};
    auto queryScaleWrapper = at_npu::native::TensorWrapper{queryDequantScale, at_npu::native::ConvertToAclDataType(queryDequantScale.scalar_type())};
    auto keyScaleWrapper = at_npu::native::TensorWrapper{keyDequantScale, at_npu::native::ConvertToAclDataType(keyDequantScale.scalar_type())};

    if (quantMode == 4) {
        queryWrapper.dtype = ACL_HIFLOAT8;
        keyWrapper.dtype = ACL_HIFLOAT8;
    } else if (quantMode == 5) {
        queryWrapper.dtype = ACL_FLOAT4_E2M1;
        keyWrapper.dtype = ACL_FLOAT4_E2M1;
    } else if (quantMode == 3) {
        queryScaleWrapper.dtype = ACL_FLOAT8_E8M0;
        keyScaleWrapper.dtype = ACL_FLOAT8_E8M0;
    }

    EXEC_NPU_CMD(aclnnQuantLightningIndexerV2, queryWrapper, keyWrapper, weights, queryScaleWrapper, keyScaleWrapper,
              cuSeqlensQ, cuSeqlensK, sequsedQ, sequsedK, cmpResidualK, blockTable, outputIdxOffset,
              metadata, candidateTopkIndexIn, topk, quantMode, maxSeqlenQ, queryLayoutPtr, keyLayoutPtr, maskMode,
              cmpRatio, returnValue, candidateMode, candidateTopkBlocks, candidateBlockSize, keyStride0,
              keyScaleStride0, sparseIndicesOut, sparseValuesOut, candidateTopkIndexOut);

[gemini-code-assist]

Since candidate_topk_index_out is an optional output, context->GetOutputShape(CANDIDATE_TOPK_INDEX_OUTPUT_INDEX) can return nullptr if the output is not bound in the graph. Using OP_CHECK_NULL_WITH_CONTEXT on it will cause the entire InferShape process to fail with GRAPH_FAILED when the optional output is not bound, which defeats the purpose of it being optional. Please wrap the shape setting in a null check instead of failing.

    if (candidateTopkIndexShape != nullptr) {
        const int32_t *candidate_mode = attrs->GetAttrPointer<int32_t>(ATTR_CANDIDATE_MODE_INDEX);
        uint32_t candidateMode = (candidate_mode != nullptr) ? static_cast<uint32_t>(*candidate_mode) : 3U;
        if (candidateMode == CANDIDATE_MODE_SOURCE) {
            OP_CHECK_IF(inputLayoutQueryPtrStr != "BSND",
                        OP_LOGE("QuantLightningIndexerV2",
                                "candidate_mode=1 only supports layout_q=BSND, but got %s.",
                                inputLayoutQueryPtrStr.c_str()),
                        return GRAPH_FAILED);
            const int64_t *candidate_topk_blocks = attrs->GetAttrPointer<int64_t>(ATTR_CANDIDATE_TOPK_BLOCKS_INDEX);
            int64_t candBlocks = (candidate_topk_blocks != nullptr) ? *candidate_topk_blocks : CANDIDATE_TOPK_BLOCKS_FIX;
            *candidateTopkIndexShape = *sparseIndicesShape;
            candidateTopkIndexShape->SetDim(sparseIndicesShape->GetDimNum() - 1, candBlocks);
        } else {
            candidateTopkIndexShape->SetDimNum(1);
            candidateTopkIndexShape->SetDim(0, 0);
        }
    }