From 30fbd055379ce4c5f26fcece6cfc90c5d8596f59 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 27 Jul 2026 05:11:06 -0500 Subject: [PATCH] [ROCm] Use backend-default dot precision for ReplaySSM (#49909) Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 29 ++++++++++--------- ...tive_state_update_replayssm_output_only.py | 25 +++++++++++----- 2 files changed, 33 insertions(+), 21 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 6e5d115f19..7edc293178 100644 --- a/.buildkite/test-amd.yaml +++ b/.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] diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py index c6046b06a5..5b11bb7a08 100644 --- a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py +++ b/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,