From a7555fc997f36ef3f0ed53a7632cece1e93e6314 Mon Sep 17 00:00:00 2001
From: Chi McIsaac <153383231+qimcis@users.noreply.github.com>
Date: Fri, 22 May 2026 03:18:49 -0400
Subject: [PATCH] [diffusion] fix: fix MOVA DAC bf16 on ROCm (#25674)

---
 .../multimodal_gen/runtime/models/vaes/dac.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/dac.py b/python/sglang/multimodal_gen/runtime/models/vaes/dac.py
index ee4750f3d..a944bcb99 100644
--- a/python/sglang/multimodal_gen/runtime/models/vaes/dac.py
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/dac.py
@@ -20,9 +20,7 @@ from sglang.multimodal_gen.runtime.models.vaes.common import (
 )
 
 
-# Scripting this brings model speed up 1.4x
-@torch.jit.script
-def snake(x, alpha):
+def _snake(x, alpha):
     shape = x.shape
     x = x.reshape(shape[0], shape[1], -1)
     x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
@@ -30,12 +28,27 @@ def snake(x, alpha):
     return x
 
 
+# Scripting this brings model speed up 1.4x
+snake = torch.jit.script(_snake)
+
+
+# ROCm HIPRTC can fail to compile the scripted bf16 Snake kernel.
+def _should_use_eager_snake_on_rocm_bf16(x: torch.Tensor, alpha: torch.Tensor) -> bool:
+    return (
+        torch.version.hip is not None
+        and (x.is_cuda or alpha.is_cuda)
+        and (x.dtype == torch.bfloat16 or alpha.dtype == torch.bfloat16)
+    )
+
+
 class Snake1d(nn.Module):
     def __init__(self, channels):
         super().__init__()
         self.alpha = nn.Parameter(torch.ones(1, channels, 1))
 
     def forward(self, x):
+        if _should_use_eager_snake_on_rocm_bf16(x, self.alpha):
+            return _snake(x, self.alpha)
         return snake(x, self.alpha)