[ROCm] Use backend-default dot precision for ReplaySSM (#49909)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 27, 2026 at 05:11 UTC
30fbd055379ce4c5f26fcece6cfc90c5d8596f59
2 files changed
+33
-21
.buildkite/test-amd.yaml
+15
-14
@@ -169,20 +169,6 @@ steps:
169
- pip install helion==1.1.0
170
- pytest -v -s kernels/helion/
171
172
-- label: Kernels Mamba Test # TBD
173
- timeout_in_minutes: 180
174
- mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
175
- agent_pool: mi250_1
176
- optional: true
177
- working_dir: "/vllm-workspace/tests"
178
- source_file_dependencies:
179
- - csrc/mamba/
180
- - tests/kernels/mamba
181
- - vllm/model_executor/layers/mamba/ops
182
- - vllm/platforms/rocm.py
183
- commands:
184
- - pytest -v -s kernels/mamba
185
-
172
#------------------------------------------------------ mi250 · models / basic -------------------------------------------------------#
173
174
- label: Basic Models Test (Other CPU) # TBD
@@ -1595,6 +1581,21 @@ steps:
1581
commands:
1582
- pytest -v -s kernels/test_kda.py
1583
1584
+- label: Kernels Mamba Test # TBD
1585
+ timeout_in_minutes: 180
1586
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
1587
+ dind: false
1588
+ agent_pool: mi300_1
1589
+ optional: true
1590
+ working_dir: "/vllm-workspace/tests"
1591
+ source_file_dependencies:
1592
+ - csrc/mamba/
1593
+ - tests/kernels/mamba
1594
+ - vllm/model_executor/layers/mamba/ops
1595
+ - vllm/platforms/rocm.py
1596
+ commands:
1597
+ - pytest -v -s kernels/mamba
1598
+
1599
- label: Kernels MoE Test %N # TBD
1600
timeout_in_minutes: 180
1601
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py
+18
-7
@@ -6,14 +6,16 @@ import torch
6
7
from vllm.model_executor.layers.mamba.ops.mamba_ssm import convert_rs_fp16x2, softplus
8
from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config
9
+from vllm.platforms import current_platform
10
from vllm.triton_utils import tl, triton
11
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
12
13
14
@triton.heuristics(
15
{
15
- "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]
16
- is not None
16
+ "HAS_STATE_BATCH_INDICES": lambda args: (
17
+ args["state_batch_indices_ptr"] is not None
18
+ )
19
}
20
)
21
@triton.heuristics(
@@ -121,8 +123,9 @@ def _replayssm_output_only_precompute_kernel(
123
@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})
124
@triton.heuristics(
125
{
124
- "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]
125
- is not None
126
+ "HAS_STATE_BATCH_INDICES": lambda args: (
127
+ args["state_batch_indices_ptr"] is not None
128
+ )
129
}
130
)
131
@triton.heuristics(
@@ -210,6 +213,7 @@ def _replayssm_output_only_kernel(
213
NF_NDS: tl.constexpr,
214
FL_DSTATE_TILE: tl.constexpr,
215
FL_NDS: tl.constexpr,
216
+ DOT_INPUT_PRECISION: tl.constexpr,
217
USE_RS_ROUNDING: tl.constexpr,
218
PHILOX_ROUNDS: tl.constexpr,
219
# heuristic-computed
@@ -409,12 +413,15 @@ def _replayssm_output_only_kernel(
413
B_all_dot = tl.where(
414
offs_k_dot[:, None] == write_pos, B_cur_tile[None, :], B_all_dot
415
)
412
- # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on
413
- # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag.
416
+ # CUDA uses tf32x3 to keep fp32 parity with the elementwise
417
+ # baseline. Other platforms use their backend default; bf16/fp16
418
+ # inputs are unaffected by this flag.
419
B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to(
420
x_ptr.dtype.element_ty
421
)
417
- delta_state = tl.dot(x_all_ty, B_scaled, input_precision="tf32x3")
422
+ delta_state = tl.dot(
423
+ x_all_ty, B_scaled, input_precision=DOT_INPUT_PRECISION
424
+ )
425
state_ptrs = (
426
state_ptr
427
+ offs_m[:, None] * stride_state_dim
@@ -574,6 +581,9 @@ def selective_state_update_replayssm_output_only(
581
nf_nds = triton.cdiv(bs_dstate, nf_dstate_tile)
582
fl_dstate_tile = max(16, min(fl_tile, bs_dstate))
583
fl_nds = triton.cdiv(bs_dstate, fl_dstate_tile)
584
+ # AMD Triton does not support tf32x3, so use its backend default. CUDA
585
+ # retains tf32x3 to preserve fp32 parity with the elementwise baseline.
586
+ dot_input_precision = None if current_platform.is_rocm() else "tf32x3"
587
588
grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads)
589
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(
708
nf_nds,
709
fl_dstate_tile,
710
fl_nds,
711
+ dot_input_precision,
712
enable_stochastic_rounding,
713
cache_philox_rounds,
714
num_warps=num_warps,