Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
change
sglang
Commits
7151194b
Unverified
Commit
7151194b
authored
Jun 24, 2025
by
Ke Bao
Committed by
GitHub
Jun 24, 2025
Browse files
Remove cumsum_buffer initilization (#7439)
parent
2ed68d7a
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
7 additions
and
2 deletions
+7
-2
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
+7
-2
No files found.
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View file @
7151194b
...
@@ -750,9 +750,11 @@ def moe_align_block_size(
...
@@ -750,9 +750,11 @@ def moe_align_block_size(
by block_size for proper block matrix operations.
by block_size for proper block matrix operations.
"""
"""
max_num_tokens_padded
=
topk_ids
.
numel
()
+
num_experts
*
(
block_size
-
1
)
max_num_tokens_padded
=
topk_ids
.
numel
()
+
num_experts
*
(
block_size
-
1
)
sorted_ids
,
cumsum_buffer
=
init_sorted_ids_and_cumsum_buffer
(
sorted_ids
=
torch
.
empty
(
max_num_tokens_padded
,
topk_ids
.
numel
(),
num_experts
,
topk_ids
.
device
(
max_num_tokens_padded
,
),
dtype
=
torch
.
int32
,
device
=
topk_ids
.
device
)
)
sorted_ids
.
fill_
(
topk_ids
.
numel
())
max_num_m_blocks
=
triton
.
cdiv
(
max_num_tokens_padded
,
block_size
)
max_num_m_blocks
=
triton
.
cdiv
(
max_num_tokens_padded
,
block_size
)
expert_ids
=
torch
.
empty
(
expert_ids
=
torch
.
empty
(
(
max_num_m_blocks
,),
dtype
=
torch
.
int32
,
device
=
topk_ids
.
device
(
max_num_m_blocks
,),
dtype
=
torch
.
int32
,
device
=
topk_ids
.
device
...
@@ -768,6 +770,9 @@ def moe_align_block_size(
...
@@ -768,6 +770,9 @@ def moe_align_block_size(
num_tokens_post_pad
,
num_tokens_post_pad
,
)
)
else
:
else
:
cumsum_buffer
=
torch
.
empty
(
(
num_experts
+
1
,),
dtype
=
torch
.
int32
,
device
=
topk_ids
.
device
)
token_cnts_buffer
=
torch
.
empty
(
token_cnts_buffer
=
torch
.
empty
(
(
num_experts
+
1
)
*
num_experts
,
(
num_experts
+
1
)
*
num_experts
,
dtype
=
torch
.
int32
,
dtype
=
torch
.
int32
,
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment