[ROCm][Perf] Remove AITER paged-MQA outputs guard for DeepSeek-V4 (#55808)
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Shanshan Shen committed
Sep 8, 2026 at 19:58 UTC
25047604fe5518ecb99929edc3fe2f5e76bc9f4a
2 files changed
-63
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
-62
@@ -284,68 +284,6 @@ def _launch_sparse_decode_reduce(
284
return out
285
286
287
-@torch.inference_mode()
288
-def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None:
289
- from vllm._aiter_ops import rocm_aiter_ops
290
- from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
291
-
292
- device = torch.device("cuda")
293
-
294
- class FakeWorkspaceManager:
295
- def get_simultaneous(self, *shapes_and_dtypes):
296
- return [
297
- torch.empty(shape, dtype=dtype, device=device)
298
- for shape, dtype in shapes_and_dtypes
299
- ]
300
-
301
- def fake_paged_mqa_logits(
302
- q_fp8,
303
- kv_cache_fp8,
304
- weights,
305
- out_logits,
306
- context_lens,
307
- block_tables,
308
- max_seq_len,
309
- **kwargs,
310
- ):
311
- del (
312
- q_fp8,
313
- kv_cache_fp8,
314
- weights,
315
- context_lens,
316
- block_tables,
317
- max_seq_len,
318
- kwargs,
319
- )
320
- out_logits.fill_(float("nan"))
321
-
322
- monkeypatch.setattr(mod, "_ON_GFX942", False)
323
- monkeypatch.setattr(mod, "_ON_GFX950", True)
324
- monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
325
- monkeypatch.setattr(
326
- mod,
327
- "paged_mqa_logits_module",
328
- lambda: SimpleNamespace(deepgemm_fp8_paged_mqa_logits=fake_paged_mqa_logits),
329
- )
330
- monkeypatch.setattr(
331
- mod, "current_workspace_manager", lambda: FakeWorkspaceManager()
332
- )
333
-
334
- q_fp8 = torch.empty((1, 1, 1, 1), dtype=torch.uint8, device=device)
335
- kv_cache_fp8 = torch.empty((1, 1, 1, 5), dtype=torch.uint8, device=device)
336
- logits = mod.rocm_fp8_paged_mqa_logits(
337
- q_fp8,
338
- kv_cache_fp8,
339
- torch.empty((1, 1), dtype=torch.float32, device=device),
340
- torch.ones(1, dtype=torch.int32, device=device),
341
- torch.zeros((1, 1), dtype=torch.int32, device=device),
342
- torch.empty(0, dtype=torch.int32, device=device),
343
- 1,
344
- )
345
-
346
- assert not torch.isnan(logits).any()
347
-
348
-
287
@torch.inference_mode()
288
def test_compute_global_topk_ragged_indices_and_indptr() -> None:
289
from vllm.models.deepseek_v4.amd.rocm import (
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
-1
@@ -657,7 +657,6 @@ def rocm_fp8_paged_mqa_logits(
657
KVBlockSize=block_size,
658
WavePerEU=2,
659
)
660
- out_logits.nan_to_num_(float("-inf"))
660
return out_logits
661
deepgemm_fp8_paged_mqa_logits_stage1 = (
662
aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1