diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 2f1e8d2793..e8cfa70d6a 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/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 ( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 58bef917a2..46e1ee0596 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/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