Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "d6c80e0cf7f63000a5d3e5d561ab349c697e2641"
Unverified Commit 014cab4d authored by Yineng Zhang's avatar Yineng Zhang Committed by GitHub
Browse files

update forward_return_lse (#3425)

parent 4d2dbeac
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
python/sglang/srt/layers/attention/flashinfer_backend.py
View file @ 014cab4d
...@@ -409,9 +409,9 @@ class FlashInferAttnBackend(AttentionBackend): ...@@ -409,9 +409,9 @@ class FlashInferAttnBackend(AttentionBackend):
) )
else: else:
o1, s1 = self.prefill_wrapper_ragged.forward_return_lse( o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), q.view(-1, layer.tp_q_head_num, layer.head_dim),
k.contiguous().view(-1, layer.tp_k_head_num, layer.head_dim), k.view(-1, layer.tp_k_head_num, layer.head_dim),
v.contiguous().view(-1, layer.tp_v_head_num, layer.head_dim), v.view(-1, layer.tp_v_head_num, layer.head_dim),
causal=True, causal=True,
sm_scale=layer.scaling, sm_scale=layer.scaling,
logits_soft_cap=logits_soft_cap, logits_soft_cap=logits_soft_cap,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment