[Bugfix][Determinism] Preserve output dtype in batch-invariant mean (#59106)

Signed-off-by: Prudhvivuda <prudhvi12042001@gmail.com>
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

@@ -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