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
OpenDAS
ktransformers
Commits
9037bf30
Unverified
Commit
9037bf30
authored
Apr 09, 2025
by
Atream
Committed by
GitHub
Apr 09, 2025
Browse files
Merge pull request #1090 from kvcache-ai/Atream-patch-1
Update attention.py
parents
94476ce5
3b9e16ce
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
2 additions
and
1 deletion
+2
-1
ktransformers/operators/attention.py
ktransformers/operators/attention.py
+2
-1
No files found.
ktransformers/operators/attention.py
View file @
9037bf30
...
...
@@ -435,6 +435,7 @@ class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention):
kv_len_arr
=
torch
.
tensor
([
position_ids
[
0
,
-
1
].
item
()
+
1
],
dtype
=
torch
.
int32
,
device
=
self
.
device
)
self
.
mla_wrapper
.
plan
(
qo_indptr
,
None
,
None
,
kv_len_arr
,
None
,
self
.
num_heads
,
self
.
kv_lora_rank
,
self
.
qk_rope_head_dim
,
...
...
@@ -849,4 +850,4 @@ class flashinfer_attn(BaseInjectedModule, DeepseekV2Attention):
attn_output
=
attn_output
.
transpose
(
0
,
1
)
attn_output
=
attn_output
.
reshape
(
q_len
,
self
.
num_heads
*
self
.
v_head_dim
)
attn_output
=
self
.
o_proj
(
attn_output
,
num_tokens_tensors
)
return
attn_output
\ No newline at end of file
return
attn_output
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