[ROCm][Perf][Bugfix] DSv4 indexer: use platform FP8 dtype (fnuz) for Q-quant on gfx942 (#46730)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
Jul 1, 2026 at 12:33 UTC
aa8bb5562ebe435bb19276d094a021e9721f5bce
1 file changed
+19
-8
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+19
-8
@@ -3,6 +3,7 @@
3
4
import torch
5
6
+from vllm.platforms import current_platform
7
from vllm.triton_utils import tl, triton
8
from vllm.utils.import_utils import has_cutedsl
9
@@ -88,6 +89,8 @@ def _fused_indexer_q_rope_quant_kernel(
89
index_weights_head_scale,
90
index_weights_out_ptr,
91
index_weights_out_stride,
92
+ FP8_MAX: tl.constexpr = 448.0,
93
+ USE_FNUZ: tl.constexpr = False,
94
):
95
# Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding
96
# + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the
@@ -128,26 +131,28 @@ def _fused_indexer_q_rope_quant_kernel(
131
nope_offset = tl.arange(0, INDEX_Q_NOPE_DIM)
132
x_nope = tl.load(base_ptr + nope_offset).to(tl.float32)
133
amax = tl.maximum(amax, tl.max(tl.abs(x_nope)))
131
- index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
134
+ index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
135
index_q_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(index_q_scale)))
136
134
- # Store quantized values to index_q_fp8
137
+ # Store quantized values to index_q_fp8. FNUZ (e4m3fnuz) on gfx942, OCP
138
+ # (e4m3fn) elsewhere -- matches the K cache.
139
+ fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
140
fp8_base_ptr = (
141
index_q_fp8_ptr + tok_idx * index_q_fp8_stride0 + head_idx * index_q_fp8_stride1
142
)
143
if INDEX_Q_NOPE_DIM > 0:
144
tl.store(
145
fp8_base_ptr + nope_offset,
141
- tl.div_rn(x_nope, index_q_scale).to(tl.float8e4nv),
146
+ tl.div_rn(x_nope, index_q_scale).to(fp8_dtype),
147
)
148
fp8_rot_base = fp8_base_ptr + INDEX_Q_NOPE_DIM
149
tl.store(
150
fp8_rot_base + half_offset * 2,
146
- tl.div_rn(r_even, index_q_scale).to(tl.float8e4nv),
151
+ tl.div_rn(r_even, index_q_scale).to(fp8_dtype),
152
)
153
tl.store(
154
fp8_rot_base + half_offset * 2 + 1,
150
- tl.div_rn(r_odd, index_q_scale).to(tl.float8e4nv),
155
+ tl.div_rn(r_odd, index_q_scale).to(fp8_dtype),
156
)
157
158
# FP8 weight-fold contract:
@@ -299,8 +304,9 @@ def fused_indexer_q_rope_quant(
304
Weight-fold semantics (important — the two paths differ):
305
306
FP8 path (use_fp4=False, default):
302
- q_fp8 : (T, H, HEAD_DIM) float8_e4m3fn, per-token-per-head
303
- scalar scale (NOT stored — folded into weights below)
307
+ q_fp8 : (T, H, HEAD_DIM) platform fp8 (e4m3fnuz on gfx942,
308
+ e4m3fn elsewhere); per-token-per-head scalar scale
309
+ (NOT stored — folded into weights below)
310
weights_out = weights * q_scale * softmax_scale * head_scale
311
Rationale: a single per-token q_scale is a scalar the downstream FP8
312
logits kernel would otherwise multiply in. Folding it into `weights`
@@ -397,7 +403,10 @@ def fused_indexer_q_rope_quant(
403
index_q_scale.view(torch.int32).squeeze(-1),
404
), index_weights_out
405
400
- index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
406
+ fp8_dtype = current_platform.fp8_dtype()
407
+ use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
408
+ fp8_max = 224.0 if use_fnuz else 448.0
409
+ index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
410
if has_cutedsl():
411
# lazily import, otherwise some tests fail due to CUDA driver init failure.
412
from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
@@ -433,6 +442,8 @@ def fused_indexer_q_rope_quant(
442
index_weights_head_scale,
443
index_weights_out,
444
index_weights_out.stride(0),
445
+ FP8_MAX=fp8_max,
446
+ USE_FNUZ=use_fnuz,
447
num_warps=1, # TODO: Tune this
448
)
449
return index_q_fp8, index_weights_out