From 155488d853a0bc42df227dbfc74005b3fd488e94 Mon Sep 17 00:00:00 2001 From: Prudhvi Vuda <53619858+Prudhvivuda@users.noreply.github.com> Date: Sun, 4 Oct 2026 13:07:46 -0400 Subject: [PATCH] [Bugfix][Determinism] Preserve output dtype in batch-invariant mean (#59106) Signed-off-by: Prudhvivuda Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/determinism/batch_invariant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 6490e69a60..403f475702 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -918,7 +918,7 @@ def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = for d in sorted_dims: result = result.squeeze(d) - return result + return result if dtype is not None else result.to(input.dtype) @triton.jit