[ROCm][Perf] Remove AITER paged-MQA outputs guard for DeepSeek-V4 (#55808)

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -284,68 +284,6 @@ def _launch_sparse_decode_reduce(
     return out
 
 
-@torch.inference_mode()
-def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None:
-    from vllm._aiter_ops import rocm_aiter_ops
-    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
-
-    device = torch.device("cuda")
-
-    class FakeWorkspaceManager:
-        def get_simultaneous(self, *shapes_and_dtypes):
-            return [
-                torch.empty(shape, dtype=dtype, device=device)
-                for shape, dtype in shapes_and_dtypes
-            ]
-
-    def fake_paged_mqa_logits(
-        q_fp8,
-        kv_cache_fp8,
-        weights,
-        out_logits,
-        context_lens,
-        block_tables,
-        max_seq_len,
-        **kwargs,
-    ):
-        del (
-            q_fp8,
-            kv_cache_fp8,
-            weights,
-            context_lens,
-            block_tables,
-            max_seq_len,
-            kwargs,
-        )
-        out_logits.fill_(float("nan"))
-
-    monkeypatch.setattr(mod, "_ON_GFX942", False)
-    monkeypatch.setattr(mod, "_ON_GFX950", True)
-    monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
-    monkeypatch.setattr(
-        mod,
-        "paged_mqa_logits_module",
-        lambda: SimpleNamespace(deepgemm_fp8_paged_mqa_logits=fake_paged_mqa_logits),
-    )
-    monkeypatch.setattr(
-        mod, "current_workspace_manager", lambda: FakeWorkspaceManager()
-    )
-
-    q_fp8 = torch.empty((1, 1, 1, 1), dtype=torch.uint8, device=device)
-    kv_cache_fp8 = torch.empty((1, 1, 1, 5), dtype=torch.uint8, device=device)
-    logits = mod.rocm_fp8_paged_mqa_logits(
-        q_fp8,
-        kv_cache_fp8,
-        torch.empty((1, 1), dtype=torch.float32, device=device),
-        torch.ones(1, dtype=torch.int32, device=device),
-        torch.zeros((1, 1), dtype=torch.int32, device=device),
-        torch.empty(0, dtype=torch.int32, device=device),
-        1,
-    )
-
-    assert not torch.isnan(logits).any()
-
-
 @torch.inference_mode()
 def test_compute_global_topk_ragged_indices_and_indptr() -> None:
     from vllm.models.deepseek_v4.amd.rocm import (

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -657,7 +657,6 @@ def rocm_fp8_paged_mqa_logits(
                 KVBlockSize=block_size,
                 WavePerEU=2,
             )
-            out_logits.nan_to_num_(float("-inf"))
             return out_logits
         deepgemm_fp8_paged_mqa_logits_stage1 = (
             aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1