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
00b19f19
Unverified
Commit
00b19f19
authored
Sep 01, 2024
by
Byron Hsu
Committed by
GitHub
Sep 01, 2024
Browse files
[triton] Remove the zero initialization of qk_acc by directly writing the result (#1288)
parent
6cb32ef9
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
2 additions
and
5 deletions
+2
-5
python/sglang/srt/layers/extend_attention.py
python/sglang/srt/layers/extend_attention.py
+2
-5
No files found.
python/sglang/srt/layers/extend_attention.py
View file @
00b19f19
...
@@ -127,8 +127,7 @@ def _fwd_kernel(
...
@@ -127,8 +127,7 @@ def _fwd_kernel(
)
)
k
=
tl
.
load
(
K_Buffer
+
offs_buf_k
,
mask
=
mask_n
[
None
,
:],
other
=
0.0
)
k
=
tl
.
load
(
K_Buffer
+
offs_buf_k
,
mask
=
mask_n
[
None
,
:],
other
=
0.0
)
qk
=
tl
.
zeros
([
BLOCK_M
,
BLOCK_N
],
dtype
=
tl
.
float32
)
qk
=
tl
.
dot
(
q
.
to
(
k
.
dtype
),
k
)
qk
+=
tl
.
dot
(
q
.
to
(
k
.
dtype
),
k
)
if
BLOCK_DPE
>
0
:
if
BLOCK_DPE
>
0
:
offs_kpe
=
(
offs_kpe
=
(
offs_kv_loc
[
None
,
:]
*
stride_buf_kbs
offs_kv_loc
[
None
,
:]
*
stride_buf_kbs
...
@@ -179,9 +178,7 @@ def _fwd_kernel(
...
@@ -179,9 +178,7 @@ def _fwd_kernel(
)
)
k
=
tl
.
load
(
K_Extend
+
offs_k
,
mask
=
mask_n
[
None
,
:],
other
=
0.0
)
k
=
tl
.
load
(
K_Extend
+
offs_k
,
mask
=
mask_n
[
None
,
:],
other
=
0.0
)
qk
=
tl
.
zeros
([
BLOCK_M
,
BLOCK_N
],
dtype
=
tl
.
float32
)
qk
=
tl
.
dot
(
q
,
k
,
out_dtype
=
tl
.
float32
)
qk
+=
tl
.
dot
(
q
,
k
)
if
BLOCK_DPE
>
0
:
if
BLOCK_DPE
>
0
:
offs_kpe
=
(
offs_kpe
=
(
(
cur_seq_extend_start_contiguous
+
start_n
+
offs_n
[
None
,
:])
(
cur_seq_extend_start_contiguous
+
start_n
+
offs_n
[
None
,
:])
...
...
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