EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix flashinfer version (#576)

Browse Source
This commit is contained in:
Yueyang Pan
2024-07-02 07:08:39 +02:00
committed by GitHub
parent 26294b2f3d
commit d9ac639202
2 changed files with 5 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/layers/radix_attention.py
View File

@@ -30,12 +30,8 @@ class RadixAttention(nn.Module):
self.prefill_forward = self.prefill_forward_flashinfer
self.extend_forward = self.prefill_forward_flashinfer
self.decode_forward = self.decode_forward_flashinfer
# flashinfer only accepts a boolean logit_cap argument
if logit_cap > 0:
assert logit_cap == 30
self.logit_cap = True
else:
self.logit_cap = False
# flashinfer now accepts float logit_cap argument
self.logit_cap = logit_cap if logit_cap > 0 else 0
else:
self.prefill_forward = self.prefill_forward_triton
self.extend_forward = self.extend_forward_triton
@@ -110,7 +106,7 @@ class RadixAttention(nn.Module):
o = input_metadata.flashinfer_prefill_wrapper.forward(
q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
input_metadata.token_to_kv_pool.kv_data[self.layer_id],
logits_cap=self.logit_cap,
logits_soft_cap=self.logit_cap,
)
return o.view(-1, self.tp_q_head_num * self.head_dim)
@@ -121,7 +117,7 @@ class RadixAttention(nn.Module):
o = input_metadata.flashinfer_decode_wrapper.forward(
q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
input_metadata.token_to_kv_pool.kv_data[self.layer_id],
logits_cap=self.logit_cap,
logits_soft_cap=self.logit_cap,
)
return o.view(-1, self.tp_q_head_num * self.head_dim)