[ROCm] Fix DeepSeek V4 indexer numerics and coverage (#50803)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Aug 20, 2026 at 04:34 UTC
6b68db441e7689676d52a64400f367f06276f2bb
3 files changed
+42
-4
tests/kernels/test_compressor_kv_cache.py
+16
-1
@@ -812,7 +812,22 @@ def _reference_kv_compress_norm_rope(
812
813
@pytest.mark.parametrize("num_tokens", [1, 7, 32])
814
@pytest.mark.parametrize("kv_block_size", [16, 32])
815
-@pytest.mark.parametrize("use_fp4", [False, True])
815
+@pytest.mark.parametrize(
816
+ "use_fp4",
817
+ [
818
+ False,
819
+ pytest.param(
820
+ True,
821
+ marks=pytest.mark.skipif(
822
+ not (
823
+ current_platform.is_cuda()
824
+ and current_platform.is_device_capability_family(100)
825
+ ),
826
+ reason="MXFP4 indexer cache requires an SM100-family GPU",
827
+ ),
828
+ ),
829
+ ],
830
+)
831
def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
832
"""Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
833
HEAD_DIM = 128
tests/kernels/test_fused_indexer_q_rope_quant.py
+17
-1
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
24
per_token_group_quant_fp8,
25
)
26
from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant
27
+from vllm.platforms import current_platform
28
from vllm.utils.import_utils import has_cutedsl
29
30
HEAD_DIM = 128
@@ -126,7 +127,22 @@ def _reference(
127
128
@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023])
129
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
129
-@pytest.mark.parametrize("use_fp4", [False, True])
130
+@pytest.mark.parametrize(
131
+ "use_fp4",
132
+ [
133
+ False,
134
+ pytest.param(
135
+ True,
136
+ marks=pytest.mark.skipif(
137
+ not (
138
+ current_platform.is_cuda()
139
+ and current_platform.is_device_capability_family(100)
140
+ ),
141
+ reason="MXFP4 indexer cache requires an SM100-family GPU",
142
+ ),
143
+ ),
144
+ ],
145
+)
146
@pytest.mark.parametrize("use_cutedsl", [False, True])
147
@torch.inference_mode()
148
def test_fused_indexer_q_rope_quant_matches_unfused(
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+9
-2
@@ -91,6 +91,7 @@ def _fused_indexer_q_rope_quant_kernel(
91
index_weights_out_stride,
92
FP8_MAX: tl.constexpr = 448.0,
93
USE_FNUZ: tl.constexpr = False,
94
+ USE_EXPLICIT_FMA: tl.constexpr = False,
95
):
96
# Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding
97
# + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the
@@ -118,8 +119,13 @@ def _fused_indexer_q_rope_quant_kernel(
119
rot_base = base_ptr + INDEX_Q_NOPE_DIM
120
x_even = tl.load(rot_base + half_offset * 2).to(tl.float32)
121
x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32)
121
- r_even = x_even * cos - x_odd * sin
122
- r_odd = x_odd * cos + x_even * sin
122
+ if USE_EXPLICIT_FMA:
123
+ # Match HIP rotary_embedding contraction before bf16 materialization.
124
+ r_even = tl.fma(x_even, cos, -(x_odd * sin))
125
+ r_odd = tl.fma(x_odd, cos, x_even * sin)
126
+ else:
127
+ r_even = x_even * cos - x_odd * sin
128
+ r_odd = x_odd * cos + x_even * sin
129
130
# Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax.
131
# Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py).
@@ -467,6 +473,7 @@ def fused_indexer_q_rope_quant(
473
index_weights_out.stride(0),
474
FP8_MAX=fp8_max,
475
USE_FNUZ=use_fnuz,
476
+ USE_EXPLICIT_FMA=current_platform.is_rocm(),
477
num_warps=1, # TODO: Tune this
478
)
479
return index_q_fp8, index_weights_out