mamba3 utils.py: PTX inline-asm helpers (cos_approx/sin_approx/tanh_approx) with NVPTX 'f' constraints cannot compile on ROCm/AMDGPU

Summary

mamba_ssm/ops/triton/mamba3/utils.py implements its fast trig/gating approximations (cos_approx, sin_approx, tanh_approx, sech2_approx, sigmoid_approx) with NVIDIA PTX inline assembly via tl.inline_asm_elementwise(..., constraints="=f,f"). The f constraint is NVPTX register-class notation; LLVM's AMDGPU backend defines no such class and cannot allocate it. Any Mamba-3 kernel that compiles these helpers therefore hard-fails on ROCm at first Triton compile with:

error: couldn't allocate output register for constraint 'f'

Environment

• mamba-ssm 2.3.2.post1, source-built (MAMBA_FORCE_BUILD=TRUE)
• torch 2.13.0+rocm7.2, Triton 3.7.1
• AMD Radeon RX 6700 XT (gfx1031; HSA_OVERRIDE_GFX_VERSION=10.3.0)

Reproduction (100%, two minimal kernels)

The failure reproduces through the module API (mamba_ssm.modules.mamba3.Mamba3.forward → mamba3_siso_combined) and directly with a kernel that only touches the helpers. Kernel A below fails 100%; kernel B (identical, builtins only) compiles and runs clean on the same box.

Kernel A — PTX helpers (reproduces the error 100%):

import torch, triton, triton.language as tl
from mamba_ssm.ops.triton.mamba3.utils import cos_approx, sin_approx

@triton.jit
def k_cos(x_ptr, o_ptr, N: tl.constexpr):
    offs = tl.arange(0, N)
    x = tl.load(x_ptr + offs)
    tl.store(o_ptr + offs, cos_approx(x) + sin_approx(x))

x = torch.randn(64, device="cuda", dtype=torch.float32)
o = torch.zeros(64, device="cuda", dtype=torch.float32)
k_cos[(1,)](x, o, 64)   # -> "couldn't allocate output register for constraint 'f'"

Kernel B — builtins (compiles and runs clean on ROCm):

import torch, triton, triton.language as tl

@triton.jit
def k_builtin(x_ptr, o_ptr, N: tl.constexpr):
    offs = tl.arange(0, N)
    x = tl.load(x_ptr + offs)
    tl.store(o_ptr + offs, tl.sin(x) + tl.cos(x))

x = torch.randn(64, device="cuda", dtype=torch.float32)
o = torch.zeros(64, device="cuda", dtype=torch.float32)
k_builtin[(1,)](x, o, 64)   # OK

Root cause

utils.py lines ~14–52 emit PTX opcodes (cos.approx.f32 $0, $1;, sin.approx.f32 $0, $1;, tanh.approx.f32 $0, $1;) with constraints="=f,f". Per Triton's inline_asm_elementwise contract, the assembly string "must match target's assembly format" — NVPTX text is being fed to a backend that must emit AMDGCN, and f is not a valid AMDGPU register-constraint letter. This is a source-level construct, not a packaging issue (verified against a MAMBA_FORCE_BUILD=TRUE source build), and unrelated to the causal-conv1d source-build fix.

Fix request

Replace the five helpers with ROCm-portable builtins (tl.cos/tl.sin; tanh/sech2/sigmoid via tl.sigmoid or libdevice), or gate the inline-asm path on the backend (NVIDIA-only) with a builtin fallback for ROCm.

One numerical-equivalence gate must ship with the fix: cos.approx.f32/sin.approx.f32 and tl.cos/tl.sin are different functions (PTX fast-approx vs target hardware cos). The swap must be verified numerically within fp16-accumulation tolerance on both targets (NVIDIA + ROCm) — not just "compiles" — using kernels A/B above with an fp32 reference, as the fix's own verification artifact.

Scope

All five helpers in utils.py carry the same construct (cos/sin are hit first, in the RoPE; tanh/sech2/sigmoid in gating). The fix should cover the file. Until this lands, the Mamba-3 SISO path cannot compile on ROCm at all.