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

support fp8 kvcache for hybrid attn backend on GPT-OSS (#9783)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Rain Jiang
2025-09-01 12:17:12 -07:00
committed by GitHub
parent 598c0bc19d
commit 9db8025376
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/srt/models/gpt_oss.py
View File

@@ -193,8 +193,9 @@ class GptOssSparseMoeBlock(nn.Module):
return ans return ans
def _enable_fused_set_kv_buffer(): def _enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
return _is_cuda """Enable fused set_kv_buffer only on CUDA with bfloat16 KV cache."""
return _is_cuda and forward_batch.token_to_kv_pool.dtype == torch.bfloat16
# TODO maybe move to a model-common utils # TODO maybe move to a model-common utils
@@ -341,7 +342,7 @@ class GptOssAttention(nn.Module):
layer=self.attn, layer=self.attn,
forward_batch=forward_batch, forward_batch=forward_batch,
) )
if _enable_fused_set_kv_buffer() if _enable_fused_set_kv_buffer(forward_batch)
else None else None
), ),
) )
@@ -355,7 +356,7 @@ class GptOssAttention(nn.Module):
attn_output = self.attn( attn_output = self.attn(
*inner_state, *inner_state,
sinks=self.sinks, sinks=self.sinks,
save_kv_cache=not _enable_fused_set_kv_buffer(), save_kv_cache=not _enable_fused_set_kv_buffer(forward_batch),
) )
output, _ = self.o_proj(attn_output) output, _ = self.o_proj(attn_output)
return output return output