[Fix] Use current Paddle device for Triton cache allocation

[SigureMo]

Summary

Register a Paddle compat override for torch.empty so calls with the generic device="cuda" argument allocate on the current Paddle device. This prevents Triton's autotune cache allocation from selecting an unsupported default place in Paddle multi-card processes.

The behavior change is limited to Paddle compat calls whose device is exactly "cuda"; other arguments and device values are forwarded unchanged. No KDA kernels or numerical paths are modified.

Test plan

• Directly called triton.runtime.driver.active.get_empty_cache_for_benchmark() under Paddle compat on B30Z; it returned a paddle.int32 tensor on the current Place(gpu:0).
• Full test suite was not run per request.

Benchmark / NCU

• Not applicable; no kernel code or performance path was changed.

Breaking changes

• None.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 126c5d776e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Guard the Paddle-only compatibility hook

When Paddle is not separately installed—as with every documented [cuda], [rocm], [xpu], [npu], and [cpu] installation, since pyproject.toml does not declare it—this top-level import raises ModuleNotFoundError before any FLA API can load; this is reproducible with a plain import fla. Install the compatibility override only when running through Paddle compatibility rather than making Paddle mandatory for every backend.

AGENTS.md reference: AGENTS.md:L40-L41

Useful? React with 👍 / 👎.

[Copilot]

Pull request overview

This PR modifies fla package initialization to patch Paddle’s Torch-compat torch.empty behavior so that calls using device="cuda" allocate on the current Paddle device, avoiding incorrect device placement in multi-GPU Paddle processes (notably affecting Triton autotune cache allocations).

Changes:

• Adds a Torch-compat override for torch.empty that drops device="cuda" and forwards allocation to paddle.empty.
• Registers the override via paddle.compat.proxy._extend_torch_proxy_overrides(...) during fla import.

Suppressed comments (2)

fla/init.py:24

• This introduces a global behavior change at import time (overriding Paddle compat’s torch.empty), but there’s no automated test coverage ensuring:

• device="cuda" allocates on the current Paddle device, and
• the override remains compatible with common torch.empty(...) call patterns.

Given the size of the existing pytest suite (including tests/paddle/), it would be good to add a small unit/integration test that exercises the override (with a GPU skip if needed).

paddle.compat.proxy._extend_torch_proxy_overrides(
    {
        "torch.empty": paddle.compat.proxy.RawOverriddenAttribute(

fla/init.py:26

• This registers an import-time patch via private Paddle APIs (paddle.compat.proxy._extend_torch_proxy_overrides and RawOverriddenAttribute). If those symbols are missing or renamed in a different Paddle version, importing fla will crash even when users don’t rely on the compat layer.

It would be safer to guard this registration with hasattr(...) (or a narrow try/except AttributeError) so the package remains importable when the compat proxy override mechanism is unavailable.

paddle.compat.proxy._extend_torch_proxy_overrides(
    {
        "torch.empty": paddle.compat.proxy.RawOverriddenAttribute(
            _torch_compat_empty
        ),

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.