From d7ecab51134d5640b298bed735c7494ade40c142 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Tue, 14 Apr 2026 15:25:09 +0800 Subject: [PATCH] [ROCm]fix(aiter): cast fp8 prefill output back to model dtype (#22626) Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com> --- python/sglang/srt/layers/attention/aiter_backend.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index c5c6baff8..18e760dff 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -2406,6 +2406,12 @@ class AiterAttnBackend(AttentionBackend): v_descale=v_descale, ) + # The fp8bf16 aiter prefill kernel returns bf16 even when the + # model computes in fp16. Cast back so the attention output keeps + # the same dtype as the rest of the model activations. + if o.dtype != self.input_dtype: + o = o.to(self.input_dtype) + return o.view(-1, layer.tp_q_head_num * layer.head_dim) def forward_decode(