Update sgl-kernel UTs for activation/topk/norm/rope kernels#6452
Merged
zhyncs merged 13 commits intosgl-project:mainfrom May 23, 2025
Merged
Update sgl-kernel UTs for activation/topk/norm/rope kernels#6452zhyncs merged 13 commits intosgl-project:mainfrom
zhyncs merged 13 commits intosgl-project:mainfrom
Conversation
mingfeima
requested changes
May 20, 2025
Collaborator
mingfeima
left a comment
mingfeima
left a comment
There was a problem hiding this comment.
merge test_grouped_topk.py and test_biased_grouped_topk.py into one, name it after test_topk.py. We only focus on DS v2 and v3 right now, later on we need to add more types of topk kernels.
test/srt/cpu/test_rope.py
Outdated
| def _forward_ref(self, positions, query, key, cos_sin_cache, offsets=None): | ||
| self.rotary_dim = 64 | ||
| self.head_size = 64 | ||
| self.is_neox_style = False |
Collaborator
There was a problem hiding this comment.
do we support is_neox in the C++ kernels?
mingfeima
requested changes
May 21, 2025
test/srt/cpu/test_activation.py
Outdated
Comment on lines
+23
to
+25
| def test_activation(self): | ||
| self._run_single_test([128, 22016], torch.bfloat16, "cpu") | ||
| self._run_single_test([129, 22016], torch.float16, "cpu") |
Collaborator
There was a problem hiding this comment.
This test case won't take too much time, suggest to use itertools.product to cover more combinations.
test/srt/cpu/test_norm.py
Outdated
Comment on lines
+33
to
+37
| def _run_single_test(self, shape, dtype, device="cuda"): | ||
|
|
||
| x = torch.randn(shape, dtype=dtype).to(device=device) | ||
| hidden_size = x.size(-1) | ||
| weight = torch.randn(hidden_size, dtype=dtype).to(device=device) |
Collaborator
There was a problem hiding this comment.
Suggested change
| def _run_single_test(self, shape, dtype, device="cuda"): | |
| x = torch.randn(shape, dtype=dtype).to(device=device) | |
| hidden_size = x.size(-1) | |
| weight = torch.randn(hidden_size, dtype=dtype).to(device=device) | |
| def _run_single_test(self, shape, dtype): | |
| x = torch.randn(shape, dtype=dtype) | |
| hidden_size = x.size(-1) | |
| weight = torch.randn(hidden_size, dtype=dtype) |
test/srt/cpu/test_norm.py
Outdated
Comment on lines
+46
to
+47
| # TEST: fused_add_rmsnorm | ||
| # flashinfer writes x and residual inplaced |
Collaborator
There was a problem hiding this comment.
Suggested change
| # TEST: fused_add_rmsnorm | |
| # flashinfer writes x and residual inplaced |
test/srt/cpu/test_norm.py
Outdated
Comment on lines
+50
to
+51
| residual = torch.randn(shape, dtype=dtype).to(device=device) | ||
| ref_residual = residual.clone() |
Collaborator
There was a problem hiding this comment.
Suggested change
| residual = torch.randn(shape, dtype=dtype).to(device=device) | |
| ref_residual = residual.clone() | |
| residual = torch.randn(shape, dtype=dtype) | |
| ref_residual = residual.clone() |
test/srt/cpu/test_norm.py
Outdated
Comment on lines
+62
to
+65
| def test_norm(self): | ||
| self._run_single_test([4096, 4096], torch.bfloat16, "cpu") | ||
| self._run_single_test([1024, 4096], torch.bfloat16, "cpu") | ||
| self._run_single_test([1024, 4096 + 13], torch.float16, "cpu") |
Collaborator
There was a problem hiding this comment.
use itertools.product and remove 'cpu'
test/srt/cpu/test_rope.py
Outdated
Comment on lines
+24
to
+26
| self.rotary_dim = 64 | ||
| self.head_size = 64 | ||
| self.is_neox_style = False |
Collaborator
There was a problem hiding this comment.
make these input parameters, rotary_dim, head_size and is_neox_stype
|
|
||
|
|
||
| # This is used by the Deepseek-V2 model | ||
| class TestGroupedTopK(CustomTestCase): |
Collaborator
There was a problem hiding this comment.
can we import that native impl from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/topk.py
test/srt/cpu/test_topk.py
Outdated
Comment on lines
+90
to
+99
| def _biased_grouped_topk( | ||
| self, | ||
| hidden_states: torch.Tensor, | ||
| gating_output: torch.Tensor, | ||
| correction_bias: torch.Tensor, | ||
| topk: int, | ||
| renormalize: bool, | ||
| num_expert_group: int = 0, | ||
| topk_group: int = 0, | ||
| ): |
Collaborator
There was a problem hiding this comment.
same as above, can we import the native impl from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/topk.py
mingfeima
requested changes
May 23, 2025
test/srt/cpu/test_rope.py
Outdated
Comment on lines
+14
to
+38
| def _forward_ref( | ||
| self, | ||
| positions, | ||
| query, | ||
| key, | ||
| cos_sin_cache, | ||
| rotary_dim, | ||
| head_size, | ||
| is_neox_style, | ||
| offsets=None, | ||
| ): | ||
| query_rot = query[..., :rotary_dim] | ||
| key_rot = key[..., :rotary_dim] | ||
| if rotary_dim < head_size: | ||
| query_pass = query[..., rotary_dim:] | ||
| key_pass = key[..., rotary_dim:] | ||
|
|
||
| cos_sin = cos_sin_cache[ | ||
| torch.add(positions, offsets) if offsets is not None else positions | ||
| ] | ||
| cos, sin = cos_sin.chunk(2, dim=-1) | ||
| if is_neox_style: | ||
| # shape [batch_size, seq_len]. | ||
| cos = cos.repeat(1, 1, 2).unsqueeze(-2) | ||
| sin = sin.repeat(1, 1, 2).unsqueeze(-2) |
Collaborator
There was a problem hiding this comment.
| # This is used by the Deepseek-V2 model | ||
| class TestGroupedTopK(CustomTestCase): | ||
| def _run_single_test(self, M, E, G, topk, topk_group, renormalize, dtype): | ||
|
|
Collaborator
There was a problem hiding this comment.
you may set a seed here as topk is unstable sort.
test/srt/cpu/utils.py
Outdated
Comment on lines
+151
to
+163
|
|
||
|
|
||
| def _rotate_neox(x: torch.Tensor) -> torch.Tensor: | ||
| x1 = x[..., : x.shape[-1] // 2] | ||
| x2 = x[..., x.shape[-1] // 2 :] | ||
| return torch.cat((-x2, x1), dim=-1) | ||
|
|
||
|
|
||
| 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) |
mingfeima
approved these changes
May 23, 2025
zhyncs
reviewed
May 23, 2025
test/srt/cpu/test_activation.py
Outdated
|
|
||
| import torch | ||
| import torch.nn.functional as F | ||
| from sgl_kernel.common_ops import silu_and_mul_cpu as silu_and_mul |
Collaborator
There was a problem hiding this comment.
Contributor
Author
There was a problem hiding this comment.
@zhyncs Thanks, has updated according to #6404 (comment).
Collaborator
|
|
Collaborator
|
6 tasks
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
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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 to add UTs for activation/topk/morm/rope kernels.
Modifications
Checklist