From 6b68db441e7689676d52a64400f367f06276f2bb Mon Sep 17 00:00:00 2001
From: Andreas Karatzas <akaratza@amd.com>
Date: Thu, 20 Aug 2026 04:34:00 -0500
Subject: [PATCH] [ROCm] Fix DeepSeek V4 indexer numerics and coverage (#50803)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>
---
 tests/kernels/test_compressor_kv_cache.py      | 17 ++++++++++++++++-
 .../kernels/test_fused_indexer_q_rope_quant.py | 18 +++++++++++++++++-
 .../deepseek_v4/common/ops/fused_indexer_q.py  | 11 +++++++++--
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py
index b300d0ebec..c1049eb658 100644
--- a/tests/kernels/test_compressor_kv_cache.py
+++ b/tests/kernels/test_compressor_kv_cache.py
@@ -812,7 +812,22 @@ def _reference_kv_compress_norm_rope(
 
 @pytest.mark.parametrize("num_tokens", [1, 7, 32])
 @pytest.mark.parametrize("kv_block_size", [16, 32])
-@pytest.mark.parametrize("use_fp4", [False, True])
+@pytest.mark.parametrize(
+    "use_fp4",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=pytest.mark.skipif(
+                not (
+                    current_platform.is_cuda()
+                    and current_platform.is_device_capability_family(100)
+                ),
+                reason="MXFP4 indexer cache requires an SM100-family GPU",
+            ),
+        ),
+    ],
+)
 def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
     """Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
     HEAD_DIM = 128
diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py
index 6114b7efd6..5b568c938e 100644
--- a/tests/kernels/test_fused_indexer_q_rope_quant.py
+++ b/tests/kernels/test_fused_indexer_q_rope_quant.py
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
 )
 from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant
+from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_cutedsl
 
 HEAD_DIM = 128
@@ -126,7 +127,22 @@ def _reference(
 
 @pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023])
 @pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("use_fp4", [False, True])
+@pytest.mark.parametrize(
+    "use_fp4",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=pytest.mark.skipif(
+                not (
+                    current_platform.is_cuda()
+                    and current_platform.is_device_capability_family(100)
+                ),
+                reason="MXFP4 indexer cache requires an SM100-family GPU",
+            ),
+        ),
+    ],
+)
 @pytest.mark.parametrize("use_cutedsl", [False, True])
 @torch.inference_mode()
 def test_fused_indexer_q_rope_quant_matches_unfused(
diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
index 3ec0c2b29d..9ae34de1a0 100644
--- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
@@ -91,6 +91,7 @@ def _fused_indexer_q_rope_quant_kernel(
     index_weights_out_stride,
     FP8_MAX: tl.constexpr = 448.0,
     USE_FNUZ: tl.constexpr = False,
+    USE_EXPLICIT_FMA: tl.constexpr = False,
 ):
     # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding
     # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the
@@ -118,8 +119,13 @@ def _fused_indexer_q_rope_quant_kernel(
     rot_base = base_ptr + INDEX_Q_NOPE_DIM
     x_even = tl.load(rot_base + half_offset * 2).to(tl.float32)
     x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32)
-    r_even = x_even * cos - x_odd * sin
-    r_odd = x_odd * cos + x_even * sin
+    if USE_EXPLICIT_FMA:
+        # Match HIP rotary_embedding contraction before bf16 materialization.
+        r_even = tl.fma(x_even, cos, -(x_odd * sin))
+        r_odd = tl.fma(x_odd, cos, x_even * sin)
+    else:
+        r_even = x_even * cos - x_odd * sin
+        r_odd = x_odd * cos + x_even * sin
 
     # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax.
     # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py).
@@ -467,6 +473,7 @@ def fused_indexer_q_rope_quant(
             index_weights_out.stride(0),
             FP8_MAX=fp8_max,
             USE_FNUZ=use_fnuz,
+            USE_EXPLICIT_FMA=current_platform.is_rocm(),
             num_warps=1,  # TODO: Tune this
         )
     return index_q_fp8, index_weights_out