[ROCm] Use backend-default dot precision for ReplaySSM (#49909)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

.buildkite/test-amd.yaml

@@ -169,20 +169,6 @@ steps:
   - pip install helion==1.1.0
   - pytest -v -s kernels/helion/
 
-- label: Kernels Mamba Test # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - csrc/mamba/
-  - tests/kernels/mamba
-  - vllm/model_executor/layers/mamba/ops
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s kernels/mamba
-
 #------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Test (Other CPU) # TBD
@@ -1595,6 +1581,21 @@ steps:
   commands:
   - pytest -v -s kernels/test_kda.py
 
+- label: Kernels Mamba Test # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  dind: false
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/mamba/
+  - tests/kernels/mamba
+  - vllm/model_executor/layers/mamba/ops
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/mamba
+
 - label: Kernels MoE Test %N # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]

vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py

@@ -6,14 +6,16 @@ import torch
 
 from vllm.model_executor.layers.mamba.ops.mamba_ssm import convert_rs_fp16x2, softplus
 from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 
 
 @triton.heuristics(
     {
-        "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]
-        is not None
+        "HAS_STATE_BATCH_INDICES": lambda args: (
+            args["state_batch_indices_ptr"] is not None
+        )
     }
 )
 @triton.heuristics(
@@ -121,8 +123,9 @@ def _replayssm_output_only_precompute_kernel(
 @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})
 @triton.heuristics(
     {
-        "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]
-        is not None
+        "HAS_STATE_BATCH_INDICES": lambda args: (
+            args["state_batch_indices_ptr"] is not None
+        )
     }
 )
 @triton.heuristics(
@@ -210,6 +213,7 @@ def _replayssm_output_only_kernel(
     NF_NDS: tl.constexpr,
     FL_DSTATE_TILE: tl.constexpr,
     FL_NDS: tl.constexpr,
+    DOT_INPUT_PRECISION: tl.constexpr,
     USE_RS_ROUNDING: tl.constexpr,
     PHILOX_ROUNDS: tl.constexpr,
     # heuristic-computed
@@ -409,12 +413,15 @@ def _replayssm_output_only_kernel(
             B_all_dot = tl.where(
                 offs_k_dot[:, None] == write_pos, B_cur_tile[None, :], B_all_dot
             )
-            # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on
-            # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag.
+            # CUDA uses tf32x3 to keep fp32 parity with the elementwise
+            # baseline. Other platforms use their backend default; bf16/fp16
+            # inputs are unaffected by this flag.
             B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to(
                 x_ptr.dtype.element_ty
             )
-            delta_state = tl.dot(x_all_ty, B_scaled, input_precision="tf32x3")
+            delta_state = tl.dot(
+                x_all_ty, B_scaled, input_precision=DOT_INPUT_PRECISION
+            )
             state_ptrs = (
                 state_ptr
                 + offs_m[:, None] * stride_state_dim
@@ -574,6 +581,9 @@ def selective_state_update_replayssm_output_only(
     nf_nds = triton.cdiv(bs_dstate, nf_dstate_tile)
     fl_dstate_tile = max(16, min(fl_tile, bs_dstate))
     fl_nds = triton.cdiv(bs_dstate, fl_dstate_tile)
+    # AMD Triton does not support tf32x3, so use its backend default. CUDA
+    # retains tf32x3 to preserve fp32 parity with the elementwise baseline.
+    dot_input_precision = None if current_platform.is_rocm() else "tf32x3"
 
     grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads)
     z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0)
@@ -698,6 +708,7 @@ def selective_state_update_replayssm_output_only(
             nf_nds,
             fl_dstate_tile,
             fl_nds,
+            dot_input_precision,
             enable_stochastic_rounding,
             cache_philox_rounds,
             num_warps=num_warps,