CPU: map changes from developing branch in sgl-kernel #6833

yanbing-j · 2025-06-03T08:45:27Z

Motivation

This PR is to map changes from developing branch in sgl-kernel, including:

optimize decode perf by mapping algo from FlashMLA
optimize fp8 performance.
Add API of fuse q_a_proj and kv_a_proj
Support non-contiguous input in norm
Update topk API along with SGLang main branch

Modifications

Checklist

Format your code according to the Code Formatting with Pre-Commit.
Add unit tests as outlined in the Running Unit Tests.
Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

gemini-code-assist

Hello @yanbing-j, 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!

Summary of Changes

Hello! Gemini here, providing a summary of this pull request. The primary goal of this PR, as stated in the description, is to optimize the performance of the decode operation by incorporating algorithms inspired by FlashMLA. The changes focus on the CPU kernel implementations for attention, particularly for decode and extend phases. This involves introducing VNNI packing for KV caches and restructuring the attention kernels to leverage these packed formats and improve data locality and computation efficiency. The PR also includes some minor optimizations to utility functions and refactors parts of the decode attention logic into smaller, more focused functions.

Highlights

Performance Optimization: Implements performance optimizations for the decode and extend attention kernels on CPU by adopting techniques from FlashMLA, specifically VNNI packing.
VNNI Packing: Adds new utility functions (pack_vnni_Nx32, pack_vnni_Kx32, and updates pack_vnni, pack_vnni2) to convert KV cache data into the VNNI format, which is beneficial for vectorized and matrix multiplication operations on modern CPUs (like those with AVX512 VNNI support).
MLA Kernel: Introduces a dedicated kernel (decode_attention_mla_kernel_impl) specifically optimized for the Multi-Layer Attention (MLA) case (where num_heads_kv == 1), utilizing the new VNNI packing and a different attention calculation flow.
Code Refactoring: Refactors the main decode attention logic by extracting common steps like KV buffer updates (decode_set_kv_buffer) and accumulation across KV splits (decode_accumulate_kv_splits) into separate functions, improving modularity.
AVX512 Utilities: Adds AVX512-specific transpose utility functions (transpose_16x16_32bit, transpose_2x32_16bit) in vec.h to support the VNNI packing implementations.

Changelog

Click here to see the changelog

sgl-kernel/csrc/cpu/decode.cpp
- Added include for gemm.h.
- Implemented AVX512-specific VNNI packing functions (pack_vnni_Nx32, pack_vnni).
- Refactored fill_stub and copy_stub with constexpr int kVecSize and #pragma GCC unroll 4.
- Added a new copy_stub template for copying from float to scalar_t with a block size.
- Extracted KV buffer update logic into decode_set_kv_buffer.
- Extracted accumulation across KV splits logic into decode_accumulate_kv_splits.
- Introduced decode_attention_mla_kernel_impl for MLA optimization.
- Modified decode_attention_kernel_impl (MHA) to use the new refactored functions.
- Modified decode_attention_grouped_kernel_impl (GQA/MQA) to use the new refactored functions and adjusted logit cap loop bounds.
- Updated decode_attention_cpu to call decode_set_kv_buffer, check for MLA, and dispatch to the appropriate kernel (decode_attention_mla_kernel_impl or decode_attention_grouped_kernel_impl).
- Allocated temporary buffer for MLA kernel.
sgl-kernel/csrc/cpu/extend.cpp
- Implemented AVX512-specific VNNI packing functions (pack_vnni_Nx32, pack_vnni_Kx32).
- Updated pack_vnni and pack_vnni2 to use the new AVX512 packing functions when available.
- Refactored fill_stub and copy_stub with constexpr int kVecSize and #pragma GCC unroll 4.
sgl-kernel/csrc/cpu/vec.h
- Added AVX512 transpose utility functions (transpose_16x16_32bit, transpose_2x32_16bit).

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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.

CPU cores hum fast,
Data packed, attention cast,
Decode speeds increase.

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 significant optimizations for decode performance, particularly targeting AVX512-capable CPUs by mapping algorithms inspired by FlashMLA. The changes involve substantial refactoring of the attention kernels, introduction of VNNI packing for keys and values, and a new MLA (Masked Language Attention) kernel.

The refactoring separates KV cache updates, attention computation, and cross-KV-split accumulation, which generally improves modularity. The use of AVX512 intrinsics and manual loop unrolling in fill_stub and copy_stub functions are good performance practices.

However, there are a few areas that need attention, primarily a potential correctness issue related to logit_cap application in the GQA/MQA kernel and some assumptions that should be clarified or asserted for robustness. It would also be beneficial to see benchmark results to quantify the performance improvements, as indicated in the PR checklist.

Summary of Findings

Potential Correctness Issue with logit_cap in GQA/MQA Kernel: In sgl-kernel/csrc/cpu/decode.cpp, the application of logit_cap in decode_attention_grouped_kernel_impl (line 1305) processes BLOCK_H * BLOCK_N elements, while the valid computed scores are BLOCK_H * n_size. If n_size < BLOCK_N and logit_cap != 0, this could lead to incorrect results by processing stale data. This was flagged as high severity.
Assumptions on Head Size Divisibility for AVX512 Paths: Several AVX512-specific packing functions (e.g., pack_vnni in decode.cpp and extend.cpp) assume that head sizes (K, Kv, or N in pack_vnni2) are multiples of 32, with comments like "no remainder". If these assumptions are violated, it could lead to incorrect behavior. Adding assertions or ensuring upstream checks would improve robustness. This was flagged as medium severity.
Clarity of Comments: A comment in decode.cpp (line 39, pack_vnni_Nx32) "skip 64 elems for deepseek" is unclear and could benefit from more explanation. This was flagged as medium severity.
MLA Kernel logit_cap Limitation: The new MLA kernel in decode.cpp (line 1042) explicitly checks that logit_cap == 0.f. This limitation should be noted, as it might affect applicability if logit_cap is needed. This was flagged as medium severity.
Missing Benchmark Results: The PR aims to optimize performance, but the checklist item for providing benchmark results is unchecked. These results are important for validating the impact of the changes. (Not commented in-line as it's a general PR concern).

Merge Readiness

This pull request introduces substantial and complex optimizations for CPU decoding performance. While the overall direction and many of the low-level optimizations (like VNNI packing and loop unrolling) are promising, there's a high-severity potential correctness issue regarding logit_cap application in the GQA/MQA kernel that must be addressed before merging. Additionally, several medium-severity concerns regarding assumptions about data dimensions and comment clarity should be reviewed.

Given the nature and extent of these changes, thorough testing, especially for edge cases and configurations where logit_cap != 0, is crucial. It would also be very helpful to see benchmark results demonstrating the performance gains.

I am unable to approve pull requests. Based on this review, I recommend that the identified high-severity issue be resolved, and the medium-severity points be considered, before this PR is merged. Further review by others would also be beneficial.

sgl-kernel/csrc/cpu/decode.cpp

sgl-kernel/csrc/cpu/extend.cpp

mingfeima

i think we can add all the diffs from our developing branch to this one.

mingfeima · 2025-06-04T05:19:16Z

@yanbing-j to simplify upstreaming efforts, put all the diffs from https://github.com/mingfeima/sglang/tree/cpu_opt_ww11 in sgl-kernel to this one.

change the PR title to be something like "CPU: map changes from devleoping branch in sgl-kernel"
list what has these commits done in the PR commit message.
since the new commit may have multiple authors, you need to change the authors in the original commits to this PR (see below), otherwise they won't be shown once this is landed.
update the test cases for MLA, you may refer to https://github.com/mingfeima/sgl-cpu-tests/blob/main/test_mla.py

### first use rebase to merge the commits
git rebase -i HEAD~N

### change the author of commit after rebasing
git commit --amend --author="New Name <new_email@example.com>"
git rebase --continue

mingfeima · 2025-06-09T05:08:10Z

sgl-kernel/csrc/cpu/decode.cpp

+  // check if we have MLA here
+  void* k_buffer_data = k_buffer.data_ptr();
+  void* v_buffer_data = v_buffer.data_ptr();
+  const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1);


let's put tighter restrict here:

Suggested change

const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1);

const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1) && (head_size == head_size_v + 64);

fix gemini code review

Thank you for the suggestion, @mingfeima. Adding the condition (head_size == head_size_v + 64) provides a more specific check for the MLA condition. This tighter restriction is likely important for ensuring the correct kernel is used based on the specific head size configuration, which aligns with the PR's goal of optimizing decode performance by mapping algorithms like FlashMLA. This seems like a valuable refinement.

mingfeima

almost done, just some minor issues to address to.

sgl-kernel/csrc/cpu/extend.cpp

mingfeima · 2025-06-09T05:15:47Z

sgl-kernel/csrc/cpu/qkv_proj.cpp

@@ -624,3 +629,69 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(

  return std::make_tuple(q_input, k_input, v_input);
 }
+
+std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qkv_proj_with_rope(


not a very good naming ...

Suggested change

std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qkv_proj_with_rope(

std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(

mingfeima · 2025-06-09T05:21:18Z

sgl-kernel/csrc/cpu/qkv_proj.cpp

+  CHECK_EQ(qkv_a_proj_weight.size(0), q_lora_rank + kv_lora_rank + qk_rope_head_dim);
+  CHECK_EQ(qkv_a_proj_weight.size(1), get_row_size(hidden_size, use_int8_w8a8));
+
+  at::Tensor q_a_proj_weight = qkv_a_proj_weight.narrow(0, 0, q_lora_rank).contiguous();


use torch.split

and since we are spliting dim0, you don't need .contiguous(). The result tensors should be contiguous anyway, you may double check that.

mingfeima · 2025-06-09T05:21:46Z

sgl-kernel/csrc/cpu/torch_extension_cpu.cpp

@@ -182,6 +182,26 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(
    bool is_vnni,
    std::optional<std::vector<int64_t>> block_size);

+std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qkv_proj_with_rope(


change the func name.

mingfeima · 2025-06-09T05:23:45Z

test/srt/cpu/test_norm.py

@@ -37,7 +37,8 @@ def _forward_native(

    def _norm_test(self, m, n, dtype):

-        x = torch.randn([m, n], dtype=dtype)
+        x = torch.randn([m, 2 * n], dtype=dtype)
+        x = x[..., :n]  # Make x non-contiguous


write something below in the utils.py

def make_non_contiguous(Tensor: x) -> Tensor

sgl-kernel/csrc/cpu/qkv_proj.cpp

decode.cpp: update the parallel scheme for DeepSeek MLA remove logit cap for mla kernel use brgemm for MLA and vectorized packing to vnni format

The bug is due to Btmp1 is not allocated by at::empty which may contain NaN, and Btmp1 is padded when sending to brgemm. So we need to init Btmp1 to a whatever value for each thread only once.

1. brgemm impl: move brgemm out of inner loop 2. avx512 impl: move scaling out of inner loop 3. fp8_scaled_mm: change BLOCK_M to 128 to reduce access to B 4. cvt_fp8_bf16: ignore NaN handling ``` Comparing: True max_diff = 0.01562, asum = 10.562, bsum = 10.375 gemm_bf16(native): 89.812 us, gemm_fp8(opt): 124.585 us Comparing: True max_diff = 0.01562, asum = -32.500, bsum = -32.750 gemm_bf16(native): 83.805 us, gemm_fp8(opt): 125.586 us Comparing: True max_diff = 0.01562, asum = -35.750, bsum = -36.500 gemm_bf16(native): 89.579 us, gemm_fp8(opt): 151.284 us Comparing: True max_diff = 0.03125, asum = 4512.000, bsum = 4512.000 gemm_bf16(native): 262.104 us, gemm_fp8(opt): 615.823 us ``` ``` Comparing: True max_diff = 0.01562, asum = 10.562, bsum = 10.375 gemm_bf16(native): 86.403 us, gemm_fp8(opt): 95.792 us Comparing: True max_diff = 0.01562, asum = -32.500, bsum = -32.750 gemm_bf16(native): 84.178 us, gemm_fp8(opt): 100.573 us Comparing: True max_diff = 0.01562, asum = -35.750, bsum = -36.500 gemm_bf16(native): 90.365 us, gemm_fp8(opt): 114.198 us Comparing: True max_diff = 0.03125, asum = 4512.000, bsum = 4512.000 gemm_bf16(native): 267.053 us, gemm_fp8(opt): 404.231 us ```

``` gemm_bf16(native): 255.233 us, gemm_fp8(opt): 140.526 us, gemm_int8(opt): 99.790 us, gemm_bf16(opt): 188.522 us gemm_bf16(native): 254.283 us, gemm_fp8(opt): 125.441 us, gemm_int8(opt): 99.161 us, gemm_bf16(opt): 189.518 us ```

Co-authored-by: mingfeima <mingfei.ma@intel.com>

Merge branch 'sgl_20250610_sync_tag047 of git@code.alipay.com:Theta/SGLang.git into main https://code.alipay.com/Theta/SGLang/pull_requests/52 Reviewed-by: 剑川 <jianchuan.gys@antgroup.com> * [Bugfix] Fix slice operation when chunk size mismatch (sgl-project#6697) * [Bugfix] Fix ChatCompletion endpoint of mini_lb when stream is set (sgl-project#6703) * [CI] Fix setup of disaggregation with different tp (sgl-project#6706) * [PD] Remove Unnecessary Exception Handling for FastQueue.get() (sgl-project#6712) * Fuse routed_scaling_factor in DeepSeek (sgl-project#6710) * Overlap two kernels in DeepSeek with communication (sgl-project#6711) * Minor refactor two-batch overlap (sgl-project#6682) * Speed up when having padding tokens two-batch overlap (sgl-project#6668) * [Feature] Support Flashinfer fp8 blockwise GEMM kernel on Blackwell (sgl-project#6479) * Fix LoRA bench (sgl-project#6719) * temp * Fix PP for Qwen3 MoE (sgl-project#6709) * [feat] triton kernel for get_last_loc (sgl-project#6676) * [fix] more mem for draft_extend cuda_graph (sgl-project#6726) * [PD] bug fix: Update status if nixl receiver send a a dummy req. (sgl-project#6720) * Tune memory arguments on B200 (sgl-project#6718) * Add DeepSeek-R1-0528 function call chat template (sgl-project#6725) * refactor(tool call): Fix BaseFormatDetector tool_index issue and refactor `parse_streaming_increment` (sgl-project#6715) * Add draft extend CUDA graph for Triton backend (sgl-project#6705) * refactor apply_w8a8_block_fp8_linear in fp (sgl-project#6545) * [PD] Support completion endpoint (sgl-project#6729) * PD Rust LB (PO2) (sgl-project#6437) * Super tiny enable sole usage of expert distribution metrics and update doc (sgl-project#6680) * Support picking variants of EPLB algorithms (sgl-project#6728) * Support tuning DeepEP configs (sgl-project#6742) * [test] add ut and bm for get_last_loc (sgl-project#6746) * Fix mem_fraction_static for AMD CI (sgl-project#6748) * [fix][RL] Fix DeepSeekV3ForCausalLM.post_load_weights for multiple update weight (sgl-project#6265) * Improve EPLB logical to physical dispatch map (sgl-project#6727) * Update DeepSeek-R1-0528 function call chat template (sgl-project#6765) * [PD] Optimize time out logic and add env var doc for mooncake (sgl-project#6761) * Fix aiohttp 'Chunk too big' in bench_serving (sgl-project#6737) * Support sliding window in triton backend (sgl-project#6509) * Fix shared experts fusion error (sgl-project#6289) * Fix one bug in the grouped-gemm triton kernel (sgl-project#6772) * update llama4 chat template and pythonic parser (sgl-project#6679) * feat(tool call): Enhance Llama32Detector for improved JSON parsing in non-stream (sgl-project#6784) * Support token-level quantization for EP MoE (sgl-project#6782) * Temporarily lower mmlu threshold for triton sliding window backend (sgl-project#6785) * ci: relax test_function_call_required (sgl-project#6786) * Add intel_amx backend for Radix Attention for CPU (sgl-project#6408) * Fix incorrect LoRA weight loading for fused gate_up_proj (sgl-project#6734) * fix(PD-disaggregation): Can not get local ip (sgl-project#6792) * [FIX] mmmu bench serving result display error (sgl-project#6525) (sgl-project#6791) * Bump torch to 2.7.0 (sgl-project#6788) * chore: bump sgl-kernel v0.1.5 (sgl-project#6794) * Improve profiler and integrate profiler in bench_one_batch_server (sgl-project#6787) * chore: upgrade sgl-kernel v0.1.5 (sgl-project#6795) * [Minor] Always append newline after image token when parsing chat message (sgl-project#6797) * Update CI tests for Llama4 models (sgl-project#6421) * [Feat] Enable PDL automatically on Hopper architecture (sgl-project#5981) * chore: update blackwell docker (sgl-project#6800) * misc: cache is_hopper_arch (sgl-project#6799) * Remove contiguous before Flashinfer groupwise fp8 gemm (sgl-project#6804) * Correctly abort the failed grammar requests & Improve the handling of abort (sgl-project#6803) * [EP] Add cuda kernel for moe_ep_pre_reorder (sgl-project#6699) * Add draft extend CUDA graph for flashinfer backend (sgl-project#6805) * Refactor CustomOp to avoid confusing bugs (sgl-project#5382) * Tiny log prefill time (sgl-project#6780) * Tiny fix EPLB assertion about rebalancing period and recorder window size (sgl-project#6813) * Add simple utility to dump tensors for debugging (sgl-project#6815) * Fix profiles do not have consistent names (sgl-project#6811) * Speed up rebalancing when using non-static dispatch algorithms (sgl-project#6812) * [1/2] Add Kernel support for Cutlass based Fused FP4 MoE (sgl-project#6093) * [Router] Fix k8s Service Discovery (sgl-project#6766) * Add CPU optimized kernels for topk and rope fusions (sgl-project#6456) * fix new_page_count_next_decode (sgl-project#6671) * Fix wrong weight reference in dynamic EPLB (sgl-project#6818) * Minor add metrics to expert location updater (sgl-project#6816) * [Refactor] Rename `n_share_experts_fusion` as `num_fused_shared_experts` (sgl-project#6735) * [FEAT] Add transformers backend support (sgl-project#5929) * [fix] recover auto-dispatch for rmsnorm and rope (sgl-project#6745) * fix ep_moe_reorder kernel bugs (sgl-project#6858) * [Refactor] Multimodal data processing for VLM (sgl-project#6659) * Decoder-only Scoring API (sgl-project#6460) * feat: add dp-rank to KV events (sgl-project#6852) * Set `num_fused_shared_experts` as `num_shared_experts` when shared_experts fusion is not disabled (sgl-project#6736) * Fix one missing arg in DeepEP (sgl-project#6878) * Support LoRA in TestOpenAIVisionServer and fix fused kv_proj loading bug. (sgl-project#6861) * support 1 shot allreduce in 1-node and 2-node using mscclpp (sgl-project#6277) * Fix Qwen3MoE missing token padding optimization (sgl-project#6820) * Tiny update error hints (sgl-project#6846) * Support layerwise rebalancing experts (sgl-project#6851) * Tiny allow profiler API to auto create directory (sgl-project#6865) * Support Blackwell DeepEP docker images (sgl-project#6868) * [EP] Add cuda kernel for moe_ep_post_reorder (sgl-project#6837) * [theta]merge 0605 * oai: fix openAI client error with single request via batch api (sgl-project#6170) * [PD] Fix potential perf spike caused by tracker gc and optimize doc (sgl-project#6764) * Use deepgemm instead of triton for fused_qkv_a_proj_with_mqa (sgl-project#6890) * [CUTLASS-FP4-MOE] Introduce CutlassMoEParams class for easy initialization of Cutlass Grouped Gems Metadata (sgl-project#6887) * bugfix(OAI): Fix image_data processing for jinja chat templates (sgl-project#6877) * [CPU] enable CI for PRs, add Dockerfile and auto build task (sgl-project#6458) * AITER backend extension and workload optimizations (sgl-project#6838) * [theta]merge * [theta]merge * [Feature] Support Flashinfer fmha on Blackwell (sgl-project#6930) * Fix a bug in abort & Improve docstrings for abort (sgl-project#6931) * Tiny support customize DeepEP max dispatch tokens per rank (sgl-project#6934) * Sync the changes on cuda graph runners (sgl-project#6932) * [PD] Optimize transfer queue forward logic for dummy rank (sgl-project#6922) * [Refactor] image data process in bench_serving (sgl-project#6879) * [fix] logical_to_all_physical_map index 256 is out of bounds in EP parallel. (sgl-project#6767) * Add triton fused moe kernel config for E=257 on B200 (sgl-project#6939) * [sgl-kernel] update deepgemm (sgl-project#6942) * chore: bump sgl-kernel v0.1.6 (sgl-project#6943) * Minor compile fused topk (sgl-project#6944) * [Bugfix] pipeline parallelism and Eagle Qwen2 (sgl-project#6910) * Tiny re-introduce profile id logging (sgl-project#6912) * Add triton version as a fused_moe_triton config search key to avoid performace decrease in different Triton version (sgl-project#5955) * reduce torch.zeros overhead in moe align block size kernel (sgl-project#6369) * chore: upgrade sgl-kernel v0.1.6 (sgl-project#6945) * add fbgemm moe grouped gemm kernel benchmark (sgl-project#6924) * [Docker] Add docker file for SGL Router (sgl-project#6915) * Disabling mixed chunked prefill when eagle is enabled (sgl-project#6874) * Add canary for EPLB rebalancing (sgl-project#6895) * Refactor global_server_args_dict (sgl-project#6866) * Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220) * Update server timeout time in AMD CI. (sgl-project#6953) * [misc] add is_cpu() (sgl-project#6950) * Add H20 fused MoE kernel tuning configs for DeepSeek-R1/V3 (sgl-project#6885) * Add a CUDA kernel for fusing mapping and weighted sum for MoE. (sgl-project#6916) * chore: bump sgl-kernel v0.1.6.post1 (sgl-project#6955) * chore: upgrade sgl-kernel v0.1.6.post1 (sgl-project#6957) * [DeepseekR1-FP4] Add Support for nvidia/DeepSeekR1-FP4 model (sgl-project#6853) * Revert "Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220)" (sgl-project#6968) * [AMD] Add more tests to per-commit-amd (sgl-project#6926) * chore: bump sgl-kernel v0.1.7 (sgl-project#6963) * Slightly improve the sampler to skip unnecessary steps (sgl-project#6956) * rebase h20 fused_moe config (sgl-project#6966) * Fix CI and triton moe Configs (sgl-project#6974) * Remove unnecessary kernels of num_token_non_padded (sgl-project#6965) * Extend cuda graph capture bs for B200 (sgl-project#6937) * Fuse routed scaling factor in deepseek (sgl-project#6970) * Sync cuda graph runners (sgl-project#6976) * Fix draft extend ut stability with flush cache (sgl-project#6979) * Fix triton sliding window test case (sgl-project#6981) * Fix expert distribution dumping causes OOM (sgl-project#6967) * Minor remove one kernel for DeepSeek (sgl-project#6977) * [perf][sgl-kernel] extend cutlass_mla_decode to support num_head < 128 (sgl-project#6929) * Enable more unit tests for AMD CI. (sgl-project#6983) * Use torch.compile to fuse flash attention decode metadata preparation (sgl-project#6973) * Eliminate stream sync to speed up LoRA batch init (sgl-project#6960) * support qwen3 emebedding (sgl-project#6990) * Fix torch profiler bugs for bench_offline_throughput.py (sgl-project#6557) * chore: upgrade flashinfer v0.2.6.post1 jit (sgl-project#6958) * cleanup tmp dir (sgl-project#7007) * chore: update pr test xeon (sgl-project#7008) * Fix cutlass MLA gets almost zero accuracy (sgl-project#6998) * Update amd nightly models CI. (sgl-project#6992) * feat: add direct routing strategy to DP worker (sgl-project#6884) * Fallback to lower triton version for unfound fused moe configs (sgl-project#7013) * Fix torchvision version for Blackwell (sgl-project#7015) * Simplify prepare_extend_after_decode (sgl-project#6987) * Migrate to assertEqual (sgl-project#6741) * Fix torch version in blackwell dockerfile (sgl-project#7017) * chore: update pr test xeon (sgl-project#7018) * Update default settings for blackwell (sgl-project#7023) * Support both approximate and exact expert distribution collection (sgl-project#6964) * Add decode req pool (sgl-project#6980) * [theta]merge 0610 * [theta]merge 0610 * [CI] Add CI workflow for sgl-router docker build (sgl-project#7027) * Fix fused_moe triton configs (sgl-project#7029) * CPU: map changes from developing branch in sgl-kernel (sgl-project#6833) * chore: bump v0.4.7 (sgl-project#7038) * Update README.md (sgl-project#7040)

gemini-code-assist bot reviewed Jun 3, 2025

View reviewed changes

gemini-code-assist bot suggested changes Jun 3, 2025

View reviewed changes

mingfeima requested changes Jun 4, 2025

View reviewed changes

mingfeima added cpu cpu backend performance optimization intel sgl-kernel labels Jun 4, 2025

yanbing-j changed the title ~~optimize decode perf by mapping algo from FlashMLA~~ CPU: map changes from devleoping branch in sgl-kernel, including FlashMLA and fp8 optimization Jun 4, 2025

yanbing-j force-pushed the yanbing/port_flashmla branch 3 times, most recently from ba5dbda to 4edb7b0 Compare June 6, 2025 06:42

chunyuan-w mentioned this pull request Jun 6, 2025

[CPU] [BF16] Call fused_experts_cpu, weight_packed_linear and bmm_cpu kernel in DeepSeek model #6641

Merged

yanbing-j changed the title ~~CPU: map changes from devleoping branch in sgl-kernel, including FlashMLA and fp8 optimization~~ CPU: map changes from devleoping branch in sgl-kernel Jun 6, 2025

yanbing-j requested a review from mingfeima June 6, 2025 07:42

yanbing-j changed the title ~~CPU: map changes from devleoping branch in sgl-kernel~~ CPU: map changes from developing branch in sgl-kernel Jun 9, 2025

mingfeima reviewed Jun 9, 2025

View reviewed changes

mingfeima requested changes Jun 9, 2025

View reviewed changes

yanbing-j force-pushed the yanbing/port_flashmla branch from dfbaf22 to f0e0030 Compare June 9, 2025 09:05

yanbing-j requested a review from mingfeima June 9, 2025 09:05

mingfeima requested changes Jun 10, 2025

View reviewed changes

sgl-kernel/csrc/cpu/qkv_proj.cpp Show resolved Hide resolved

sgl-kernel/csrc/cpu/qkv_proj.cpp Show resolved Hide resolved

sgl-kernel/csrc/cpu/qkv_proj.cpp Show resolved Hide resolved

mingfeima mentioned this pull request Jun 10, 2025

Update python API of activation, topk, norm and rope and remove vllm dependency #6614

Merged

6 tasks

mingfeima approved these changes Jun 10, 2025

View reviewed changes

mingfeima marked this pull request as ready for review June 10, 2025 02:37

mingfeima requested review from zhyncs, ispobock, HandH1998, BBuf, yizhang2077, merrymercy and FlamingoPg as code owners June 10, 2025 02:37

mingfeima requested a review from Ying1123 as a code owner June 10, 2025 02:37

yanbing-j force-pushed the yanbing/port_flashmla branch 2 times, most recently from 4ea8722 to 2eb21aa Compare June 10, 2025 06:25

mingfeima approved these changes Jun 10, 2025

View reviewed changes

zhyncs approved these changes Jun 10, 2025

View reviewed changes

zhyncs merged commit fcde67b into sgl-project:main Jun 10, 2025
5 of 51 checks passed

mingfeima and others added 13 commits June 10, 2025 10:19

Optimize decode perf by mapping algo from FlashMLA

253e3ba

decode.cpp: update the parallel scheme for DeepSeek MLA remove logit cap for mla kernel use brgemm for MLA and vectorized packing to vnni format

bug fix: fix NaN issue in decode MLA

d11b8b8

The bug is due to Btmp1 is not allocated by at::empty which may contain NaN, and Btmp1 is padded when sending to brgemm. So we need to init Btmp1 to a whatever value for each thread only once.

extend.cpp: vectorize pack_vnni and pack_vnni2

9954a6e

gemm_fp8.cpp: improve performance by enabling prefetch

8983b41

``` gemm_bf16(native): 255.233 us, gemm_fp8(opt): 140.526 us, gemm_int8(opt): 99.790 us, gemm_bf16(opt): 188.522 us gemm_bf16(native): 254.283 us, gemm_fp8(opt): 125.441 us, gemm_int8(opt): 99.161 us, gemm_bf16(opt): 189.518 us ```

interface.cpp, shm.cpp: remove compilation warnings

fe040f3

gemm_fp8.cpp: fix NaN issue when K is not 128x, e.g. 576

7820632

Add FlashMLA test and update test_qkv_proj_with_rope.py

622bffc

Add API of fused q_a_proj and kv_a_proj and update UTs

1e7429a

Make norm work with non-contiguous input

02921f6

Update based on comments

bf82ad0

Update topk API

2eb21aa

Merge branch 'main' into yanbing/port_flashmla

5f87eef

almaslof pushed a commit to mpashkovskii/sglang that referenced this pull request Jun 11, 2025

CPU: map changes from developing branch in sgl-kernel (sgl-project#6833)

3d7f5b8

Co-authored-by: mingfeima <mingfei.ma@intel.com>

jianan-gu pushed a commit to jianan-gu/sglang that referenced this pull request Jun 12, 2025

CPU: map changes from developing branch in sgl-kernel (sgl-project#6833)

ddfeca2

Co-authored-by: mingfeima <mingfei.ma@intel.com>

	const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1);
	const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1) && (head_size == head_size_v + 64);

	std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qkv_proj_with_rope(
	std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(

CPU: map changes from developing branch in sgl-kernel #6833

CPU: map changes from developing branch in sgl-kernel #6833

Uh oh!

Conversation

yanbing-j commented Jun 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Changelog

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Summary of Findings

Merge Readiness

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mingfeima left a comment

Choose a reason for hiding this comment

Uh oh!

mingfeima commented Jun 4, 2025

Uh oh!

mingfeima Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

mingfeima Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

mingfeima left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

mingfeima Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

mingfeima Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

mingfeima Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

mingfeima Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yanbing-j commented Jun 3, 2025 •

edited

Loading