[NPU] perf update with kvcache nz & w4a8 quant

[liupeng374]

Motivation

1、Use the nz format for kv cache, thie method accelerates the FIA operator.;
2、Moe's w4a8 uses per-channel quantization;
3、Accelerating preprocessing of MHA in prefill using the npu_interleave_rope operator；
4、bugfix num_token_non_padded_cpu;

Modifications

Use export SGLANG_USE_FIA_NZ=1 to enable FIA NZ, and this feature must be turned on together with mlapo export SGLANG_NPU_USE_MLAPO=1.

Accuracy Tests

Benchmarking and Profiling

FIANZ can speed up tpot by about 2ms, and the optimization of prefill can improve performance by more than 10%.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Summary of Changes

Hello @liupeng374, 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 integrates support for an 'FIA NZ' optimization specifically tailored for Ascend NPUs. The primary goal is to enhance the performance and efficiency of attention operations by modifying KV cache management and Rotary Positional Embedding (RoPE) application within the DeepseekV2 attention mechanism. This is achieved through the introduction of a new cache mode and a refactored RoPE calculation that leverages precomputed values and a unified NPU kernel.

Highlights

• FIA NZ Cache Mode: Introduced a new 'FIA NZ' cache mode for NPU operations, controlled by the SGLANG_USE_FIA_NZ environment variable, enabling conditional logic for KV cache handling and attention preprocessing.
• Refactored KV Cache and RoPE for DeepseekV2: The KV cache management and Rotary Positional Embedding (RoPE) application for DeepseekV2 attention on NPU have been refactored, integrating a new torch_npu.npu_kv_rmsnorm_rope_cache function for unified processing.
• Optimized Rotary Embedding Calculation: The Rotary Positional Embedding calculation now precomputes and caches cos and sin values, providing a new utility method (get_cos_sin_cache) for efficient retrieval and formatting, simplifying the forward_npu method.

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 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 counter productive. 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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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

[gemini-code-assist]

Code Review

This pull request adds support for the NZ format on Ascend NPUs, controlled by the SGLANG_USE_FIA_NZ environment variable. The changes primarily affect attention preprocessing and rotary embeddings to leverage NPU-specific optimizations. My review focuses on ensuring these changes are correctly implemented without introducing regressions. I've identified a critical issue in rotary_embedding.py that could break other models, a likely typo in mla_preprocess.py that could cause runtime errors, and some code duplication that could be refactored for better maintainability. Addressing these points will improve the robustness and quality of the code.

[gemini-code-assist]

The cache_mode is being assigned a tuple with a trailing comma. This is likely a typo, as the torch.ops.npu.npu_kv_rmsnorm_rope_cache operator probably expects a string, similar to other parts of the code. This could cause a runtime error or unexpected behavior.

Suggested change

	cache_mode = ("PA_NZ" if _use_fia_nz else "PA_BNSD",)
	cache_mode = "PA_NZ" if _use_fia_nz else "PA_BNSD"

[gemini-code-assist]

The new logic to compute cos_cached_total and sin_cached_total in _compute_cos_sin_cache, and the new get_cos_sin_cache method, depend on self.mscale. This attribute is not present in the base RotaryEmbedding class or some of its subclasses (e.g., LinearScalingRotaryEmbedding), which will cause an AttributeError for models using them.

This logic appears to be specific to DeepseekScalingRotaryEmbedding. To fix this and prevent breaking other models, please:

1. Revert the changes in RotaryEmbedding._compute_cos_sin_cache.
2. Move the logic for computing self.cos_cached_total and self.sin_cached_total into DeepseekScalingRotaryEmbedding._compute_cos_sin_cache.
3. Move the new methods (get_cos_cached_total, get_sin_cached_total, and get_cos_sin_cache) from RotaryEmbedding to DeepseekScalingRotaryEmbedding.

Here's a suggested implementation for DeepseekScalingRotaryEmbedding._compute_cos_sin_cache:

def _compute_cos_sin_cache(self) -> torch.Tensor:
    inv_freq = self._compute_inv_freq(self.scaling_factor)
    t = torch.arange(
        self.max_position_embeddings * self.scaling_factor,
        device=self.device,
        dtype=torch.float32,
    )
    freqs = torch.einsum("i,j -> ij", t, inv_freq)
    cos = freqs.cos() * self.mscale
    sin = freqs.sin() * self.mscale
    cache = torch.cat((cos, sin), dim=-1)

    emb = torch.cat((freqs, freqs), dim=-1)
    self.cos_cached_total = torch.cos(emb) * self.mscale
    self.sin_cached_total = torch.sin(emb) * self.mscale

    return cache

[gemini-code-assist]

The _use_fia_nz flag is also defined in python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py. To avoid code duplication and improve maintainability, consider defining this flag once in a shared utility module (e.g., sglang.srt.hardware_backend.npu.utils) and importing it where needed.

[gemini-code-assist]

The _use_fia_nz flag is also defined in python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py. To avoid code duplication and improve maintainability, consider defining this flag once in a shared utility module (e.g., sglang.srt.hardware_backend.npu.utils) and importing it where needed.

[iforgetmyname]

/tag-and-rerun-ci

[iforgetmyname]

/rerun-failed-ci