feat: openai oss attention sink support with trtllm-gen backend #8825 (#8834)

Co-authored-by: averyhuang <averyh@nvidia.com>
2025-08-06 22:18:27 -04:00
parent 5b6acc1495
commit 6ad6c8c9e6
4 changed files with 37 additions and 22 deletions
--- a/python/sglang/srt/models/gpt_oss.py
+++ b/python/sglang/srt/models/gpt_oss.py
@@ -301,7 +301,7 @@ class GptOssAttention(nn.Module):
        hidden_states, forward_batch, inner_state = intermediate_state
        if inner_state is None:
            return hidden_states
-        attn_output = self.attn(*inner_state, sinks=self.sinks)
+        attn_output = self.attn(*inner_state, sinks=self.sinks.to(torch.float32))
        output, _ = self.o_proj(attn_output)
        return output