[NPU] GLM-5 optimize with fused kernels

[cen121212]

Motivation

Optimize the inference performance of GLM5 on NPU

Modifications

1 Add caching for sin_cos in rotary embedding

2 Replace native operators with Triton kernels

3 Processing quantized weights to improve acceptance rate

Accuracy Tests

dataset	version	metric	mode	vllm-api-general-chat
aime2025	3fb7e8	accuracy	gen	93.33

Benchmarking and Profiling

64K+1K with 90% cache

╒══════════════════════════╤═════════╤═════════════════╤═════════════════╤═════════════════╤═════════════════╤═════════════════╤═════════════════╤═════════════════╤═════╕
│ Performance Parameters   │ Stage   │ Average         │ Min             │ Max             │ Median          │ P75             │ P90             │ P99             │  N  │
╞══════════════════════════╪═════════╪═════════════════╪═════════════════╪═════════════════╪═════════════════╪═════════════════╪═════════════════╪═════════════════╪═════╡
│ E2EL                     │ total   │ 34577.8 ms      │ 26469.4 ms      │ 54932.2 ms      │ 32519.4 ms      │ 34941.8 ms      │ 44525.6 ms      │ 53725.7 ms      │ 128 │
├──────────────────────────┼─────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────┤
│ TTFT                     │ total   │ 12049.5 ms      │ 2654.4 ms       │ 31667.5 ms      │ 9990.9 ms       │ 11186.6 ms      │ 23833.4 ms      │ 31028.8 ms      │ 128 │
├──────────────────────────┼─────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────┤
│ TPOT                     │ total   │ 22.0 ms         │ 19.6 ms         │ 27.3 ms         │ 21.6 ms         │ 23.0 ms         │ 24.0 ms         │ 26.2 ms         │ 128 │
├──────────────────────────┼─────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────┤
│ ITL                      │ total   │ 73.7 ms         │ 0.0 ms          │ 385.7 ms        │ 82.6 ms         │ 83.0 ms         │ 83.6 ms         │ 226.4 ms        │ 128 │
├──────────────────────────┼─────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────┤
│ InputTokens              │ total   │ 64002.2578      │ 64000.0         │ 64005.0         │ 64002.0         │ 64003.0         │ 64004.0         │ 64004.73        │ 128 │
├──────────────────────────┼─────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────┤
│ OutputTokens             │ total   │ 1024.0          │ 1024.0          │ 1024.0          │ 1024.0          │ 1024.0          │ 1024.0          │ 1024.0          │ 128 │
├──────────────────────────┼─────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────────────────┼─────┤
│ OutputTokenThroughput    │ total   │ 30.3386 token/s │ 18.6412 token/s │ 38.6862 token/s │ 31.4889 token/s │ 32.9899 token/s │ 34.2058 token/s │ 36.9653 token/s │ 128 │
╘══════════════════════════╧═════════╧═════════════════╧═════════════════╧═════════════════╧═════════════════╧═════════════════╧═════════════════╧═════════════════╧═════╛
╒══════════════════════════╤═════════╤════════════════════╕
│ Common Metric            │ Stage   │ Value              │
╞══════════════════════════╪═════════╪════════════════════╡
│ Benchmark Duration       │ total   │ 156606.5015 ms     │
├──────────────────────────┼─────────┼────────────────────┤
│ Total Requests           │ total   │ 128                │
├──────────────────────────┼─────────┼────────────────────┤
│ Failed Requests          │ total   │ 0                  │
├──────────────────────────┼─────────┼────────────────────┤
│ Success Requests         │ total   │ 128                │
├──────────────────────────┼─────────┼────────────────────┤
│ Concurrency              │ total   │ 28.2616            │
├──────────────────────────┼─────────┼────────────────────┤
│ Max Concurrency          │ total   │ 32                 │
├──────────────────────────┼─────────┼────────────────────┤
│ Request Throughput       │ total   │ 0.8173 req/s       │
├──────────────────────────┼─────────┼────────────────────┤
│ Total Input Tokens       │ total   │ 8192289            │
├──────────────────────────┼─────────┼────────────────────┤
│ Prefill Token Throughput │ total   │ 5311.609 token/s   │
├──────────────────────────┼─────────┼────────────────────┤
│ Total Generated Tokens   │ total   │ 131072             │
├──────────────────────────┼─────────┼────────────────────┤
│ Input Token Throughput   │ total   │ 52311.2956 token/s │
├──────────────────────────┼─────────┼────────────────────┤
│ Output Token Throughput  │ total   │ 836.9512 token/s   │
├──────────────────────────┼─────────┼────────────────────┤
│ Total Token Throughput   │ total   │ 53148.2468 token/s │

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.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @cen121212, 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 focuses on enhancing the performance of attention mechanisms, particularly for DeepSeekV2 models, when running on NPU hardware. It achieves this by incorporating custom fused NPU kernels that streamline the processing of QKV projections and optimize the application of Rotary Position Embeddings, aiming to improve overall efficiency.

Highlights

• NPU Kernel Integration for DeepSeekV2 Attention: Integrated a new fused_split_qk_norm NPU kernel to optimize the splitting of QKV (Query, Key, Value) projections and apply normalization within the DeepSeekV2 attention mechanism, replacing previous manual operations.
• Rotary Embedding Optimization: Optimized the calculation of sine and cosine values for Rotary Position Embeddings (RoPE) by caching them globally after the first layer, preventing redundant computations in subsequent layers.
• Fused RoPE and QK for NPU: Introduced a fused_rope_qk_mqa NPU kernel for efficient application of Rotary Position Embeddings combined with Query-Key operations, specifically for NPU hardware under certain conditions.

🧠 New Feature in Public Preview: You can now enable Memory 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.

Changelog

• python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
  • Imported fused_split_qk_norm from sgl_kernel_npu.norm.fused_split_qk_norm.
  • Replaced manual QKV splitting and Q-normalization with a call to the new fused_split_qk_norm function.
  • Removed redundant latent_cache splitting and kv_a_layernorm operations, as they are now handled by the fused kernel.
• python/sglang/srt/layers/attention/nsa/nsa_indexer.py
  • Added global SIN and COS variables to cache rotary embedding values.
  • Modified forward_npu to calculate cos and sin only for layer_id == 0 and reuse cached values for subsequent layers.
• python/sglang/srt/layers/rotary_embedding.py
  • Imported fused_rope_qk_mqa from sgl_kernel_npu.norm.fused_split_qk_norm.
  • Added a conditional call to fused_rope_qk_mqa for NPU-specific RoPE application when query dimensions meet certain criteria.

Activity

• The pull request is currently marked as 'WIP' (Work In Progress), indicating active development.
• The description uses a generic template, suggesting that specific details about motivation, modifications, and testing are yet to be filled in by the author.

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 introduces several NPU-specific optimizations, primarily by using fused kernels for operations like QK normalization and rotary embeddings. The changes aim to improve performance on NPU hardware.

My review focuses on ensuring the new code is robust, maintainable, and free of potential issues like race conditions. I've identified a few areas for improvement:

• A hardcoded epsilon value in a fused kernel call should be replaced with a value from the model's configuration to improve robustness.
• The use of global variables for caching is not thread-safe and should be replaced with a safer mechanism, such as caching on the forward_batch object.
• A magic number used for an NPU-specific optimization should be defined as a named constant for better readability.

Overall, the direction of using fused kernels is good for performance. Addressing these points will make the implementation more solid.

[gemini-code-assist]

Using global variables SIN and COS for caching is not thread-safe and can introduce race conditions in a concurrent environment, leading to incorrect results. It is highly recommended to remove these global variables and use a safer caching mechanism. See my other comment for a suggested implementation that caches on the forward_batch object.

[gemini-code-assist]

Using global variables for caching is not thread-safe and can lead to race conditions. A better approach is to cache these values on the forward_batch object, which is scoped to a single request. This ensures that concurrent requests do not interfere with each other's cached values.

        if not hasattr(forward_batch, "npu_indexer_sin_cos_cache"):
            cos_sin = self.rotary_emb.cos_sin_cache[positions]
            cos, sin = cos_sin.chunk(2, dim=-1)
            cos = cos.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
            sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
            forward_batch.npu_indexer_sin_cos_cache = (sin, cos)
        sin, cos = forward_batch.npu_indexer_sin_cos_cache

[gemini-code-assist]

The eps value is hardcoded to 1e-6. It's better to use the epsilon value from the layer normalization module to ensure consistency and maintainability. This makes the code more robust if the model's configuration changes.

Suggested change

	eps=1e-6,
	eps=m.q_a_layernorm.variance_epsilon,

[gemini-code-assist]

The magic number 65535 should be replaced with a named constant (e.g., NPU_FUSED_ROPE_MAX_ELEMENTS = 65535) defined at the top of the file with other NPU-related constants. This improves readability and maintainability.

[sglang-npu-bot]

/tag-and-rerun-ci

[sglang-npu-bot]

/tag-and-rerun-ci

[sglang-npu-bot]

/tag-and-rerun-ci

[iforgetmyname]

/tag-and-rerun-ci

[cen121212]

/tag-and-rerun-ci

[cen121212]

/tag-and-rerun-ci

[cen121212]

/tag-and-rerun-ci

[cen121212]

/tag-and-rerun-ci