[ROCm][Bugfix] Sanitize AITER paged-MQA logits before sparse top-k for DeepSeek-V4 (#49714)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Jul 29, 2026 at 02:08 UTC
6c7e679f048dc6123caecc3985150766e455ff22
2 files changed
+65
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+64
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from types import SimpleNamespace
5
+
6
import pytest
7
import torch
8
@@ -196,6 +198,68 @@ def _ragged_from_rows(
198
)
199
200
201
+@torch.inference_mode()
202
+def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None:
203
+ from vllm._aiter_ops import rocm_aiter_ops
204
+ from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
205
+
206
+ device = torch.device("cuda")
207
+
208
+ class FakeWorkspaceManager:
209
+ def get_simultaneous(self, *shapes_and_dtypes):
210
+ return [
211
+ torch.empty(shape, dtype=dtype, device=device)
212
+ for shape, dtype in shapes_and_dtypes
213
+ ]
214
+
215
+ def fake_paged_mqa_logits(
216
+ q_fp8,
217
+ kv_cache_fp8,
218
+ weights,
219
+ out_logits,
220
+ context_lens,
221
+ block_tables,
222
+ max_seq_len,
223
+ **kwargs,
224
+ ):
225
+ del (
226
+ q_fp8,
227
+ kv_cache_fp8,
228
+ weights,
229
+ context_lens,
230
+ block_tables,
231
+ max_seq_len,
232
+ kwargs,
233
+ )
234
+ out_logits.fill_(float("nan"))
235
+
236
+ monkeypatch.setattr(mod, "_ON_GFX942", False)
237
+ monkeypatch.setattr(mod, "_ON_GFX950", True)
238
+ monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
239
+ monkeypatch.setattr(
240
+ mod,
241
+ "paged_mqa_logits_module",
242
+ lambda: SimpleNamespace(deepgemm_fp8_paged_mqa_logits=fake_paged_mqa_logits),
243
+ )
244
+ monkeypatch.setattr(
245
+ mod, "current_workspace_manager", lambda: FakeWorkspaceManager()
246
+ )
247
+
248
+ q_fp8 = torch.empty((1, 1, 1, 1), dtype=torch.uint8, device=device)
249
+ kv_cache_fp8 = torch.empty((1, 1, 1, 5), dtype=torch.uint8, device=device)
250
+ logits = mod.rocm_fp8_paged_mqa_logits(
251
+ q_fp8,
252
+ kv_cache_fp8,
253
+ torch.empty((1, 1), dtype=torch.float32, device=device),
254
+ torch.ones(1, dtype=torch.int32, device=device),
255
+ torch.zeros((1, 1), dtype=torch.int32, device=device),
256
+ torch.empty(0, dtype=torch.int32, device=device),
257
+ 1,
258
+ )
259
+
260
+ assert not torch.isnan(logits).any()
261
+
262
+
263
@torch.inference_mode()
264
def test_compute_global_topk_ragged_indices_and_indptr() -> None:
265
from vllm.models.deepseek_v4.amd.rocm import (
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+1
@@ -442,6 +442,7 @@ def rocm_fp8_paged_mqa_logits(
442
KVBlockSize=block_size,
443
WavePerEU=2,
444
)
445
+ out_logits.nan_to_num_(float("-inf"))
446
return out_logits
447
deepgemm_fp8_paged_mqa_logits_stage1 = (
448
aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1