From 126c5d776e5363f3447c2a9f1fd279dfad85d6e3 Mon Sep 17 00:00:00 2001
From: huangjiyi <43315610+huangjiyi@users.noreply.github.com>
Date: Tue, 4 Aug 2026 11:19:09 +0800
Subject: [PATCH] [Fix] Use current Paddle device for Triton cache allocation

---
 fla/__init__.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/fla/__init__.py b/fla/__init__.py
index c57600e9a0..5fcfc3d1af 100644
--- a/fla/__init__.py
+++ b/fla/__init__.py
@@ -7,9 +7,26 @@
 
 from pkgutil import extend_path
 
+import paddle
+
 __path__ = extend_path(__path__, __name__)
 __version__ = "0.5.2"
 
+
+def _torch_compat_empty(*args, **kwargs):
+    if kwargs.get("device") == "cuda":
+        del kwargs["device"]
+    return paddle.empty(*args, **kwargs)
+
+
+paddle.compat.proxy._extend_torch_proxy_overrides(
+    {
+        "torch.empty": paddle.compat.proxy.RawOverriddenAttribute(
+            _torch_compat_empty
+        ),
+    }
+)
+
 from fla import modules, ops  # noqa: E402
 
 __all__ = ["modules", "ops"]