Add fp8 qkv_proj_with_rope kernel for CPU in sgl-kernel and add UT #6493
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
blzheng
changed the title
blzheng
requested review from
zhyncs,
ispobock,
HandH1998,
BBuf,
yizhang2077,
merrymercy,
yinfan98 and
Ying1123
as code owners
blzheng
force-pushed
the
beilei/fp8_qkvproj
branch
from
c123aa7 to
7fca26f
Compare
mingfeima
requested changes
May 23, 2025
Comment on lines
36
to
57
| def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: | ||
| x1 = x[..., ::2] | ||
| x2 = x[..., 1::2] | ||
| x = torch.stack((-x2, x1), dim=-1) | ||
| return x.flatten(-2) | ||
|
|
||
|
|
||
| def rotary_emb(q_pe, k_pe, pos, cos_sin_cache): | ||
| orig_dtype = q_pe.dtype | ||
| q_pe = q_pe.float() | ||
| k_pe = k_pe.float() | ||
| cos_sin_cache = cos_sin_cache.float() | ||
|
|
||
| query_rot = q_pe[..., :rotary_dim] | ||
| key_rot = k_pe[..., :rotary_dim] | ||
| cos_sin = cos_sin_cache[pos] | ||
| cos, sin = cos_sin.chunk(2, dim=-1) | ||
| cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) | ||
| sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) | ||
| query_rot = query_rot * cos + _rotate_gptj(query_rot) * sin | ||
| key_rot = key_rot * cos + _rotate_gptj(key_rot) * sin | ||
| return query_rot.to(orig_dtype), key_rot.to(orig_dtype) |
There was a problem hiding this comment.
| kva_packed = torch.ops.sgl_kernel.convert_weight_packed(kv_a_proj_weight) | ||
| wkc_packed = torch.ops.sgl_kernel.convert_weight_packed(w_kc) | ||
|
|
||
| q_out, k_out, v_out = torch.ops.sgl_kernel.qkv_proj_with_rope( |
There was a problem hiding this comment.
you can import everything you need at the beggining.
from torch.ops.sgl_kernel import xxx, yyy, zzz
There was a problem hiding this comment.
Custom registered ops cannot be accessed through this kind of import statement. I manually wrapped convert_weight_packed and qkv_proj_with_rope as local variables at the beginning of the file to simplify their usage in the unit tests.
mingfeima
approved these changes
May 23, 2025
zhyncs
approved these changes
May 23, 2025
Layssy
pushed a commit
to Layssy/sglang-iaas
that referenced
this pull request
Jun 9, 2025
xwu-intel
pushed a commit
to xwu-intel/sglang
that referenced
this pull request
Jun 17, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Motivation
This PR is a follow-up on #2807 and #5150 to add fp8 qkv_proj_with_rope kernel for CPU. The bf16 and int8 fused_experts kernel is already added in #5150.
This PR also adds UTs for bf16, int8 and fp8 qkv_proj_with_rope kernels for CPU.
This PR also addresses the issues with the definition and usage of
decode_attention_cpuandextend_attention_cpuon the main branch.Modifications
The main change is the C++ kernels for fp8 qkv_proj_with_rope on CPU: sgl-kernel/csrc/cpu/qkv_proj.cpp
The UTs for qkv_proj_with_rope OPs on CPU: test/srt/cpu/test_qkv_proj_with_rope.py
Checklist