diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index 4289aac1e7..b92f093880 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -223,7 +223,7 @@ def _patch_build_deps(monkeypatch, events=None): on CPU.""" def fake_generate_sparse_seqlen_triton( - query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len + seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len ): return torch.zeros(num_tokens, dtype=torch.int32, device="cpu") @@ -276,6 +276,25 @@ def test_build_populates_mixed_split_fields(monkeypatch): assert md.prefill is None +def test_build_fills_paged_kv_indptr_tail(monkeypatch): + """Rows past num_tokens must read as empty: the tail repeats the total.""" + builder = _make_builder() + _patch_build_deps(monkeypatch) + monkeypatch.setattr( + sparse_mod, + "generate_sparse_seqlen_triton", + lambda *args, **kwargs: torch.tensor([3, 5], dtype=torch.int32), + ) + builder.paged_kv_indptr[1:].fill_(-1) + + md = builder.build( + common_prefix_len=0, common_attn_metadata=_make_common_metadata() + ) + + assert md.paged_kv_indptr.tolist() == [0, 3, 8] + assert builder.paged_kv_indptr.tolist() == [0, 3, 8, 8, 8, 8, 8, 8, 8] + + def test_sink_build_skips_persistent_metadata(monkeypatch): builder = _make_builder() builder._use_persistent_metadata = False @@ -385,10 +404,10 @@ def test_dense_prefill_builds_sparse_metadata_for_decode_subset(monkeypatch): sparse_call: dict[str, torch.Tensor] = {} def fake_generate_sparse_seqlen( - query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q + sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q ): sparse_call.update( - query_lens=query_lens.clone(), + query_lens=torch.diff(sparse_query_start_loc), seq_lens=sparse_seq_lens.clone(), query_start_loc=sparse_query_start_loc.clone(), num_tokens=num_tokens, @@ -471,10 +490,10 @@ def test_long_prefill_shares_full_sparse_metadata_with_decode(monkeypatch): sparse_call: dict[str, torch.Tensor] = {} def fake_generate_sparse_seqlen( - query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q + sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q ): sparse_call.update( - query_lens=query_lens.clone(), + query_lens=torch.diff(sparse_query_start_loc), seq_lens=sparse_seq_lens.clone(), query_start_loc=sparse_query_start_loc.clone(), num_tokens=num_tokens, @@ -699,7 +718,7 @@ def test_sparse_metadata_workload_length_dtype_matrix( monkeypatch.setattr( sparse_mod, "generate_sparse_seqlen_triton", - lambda *args, **kwargs: torch.ones(args[4], dtype=torch.int32), + lambda *args, **kwargs: torch.ones(args[3], dtype=torch.int32), ) get_persistent_metadata = Mock() fake_aiter = _FakeAiter("aiter") diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 037b5432ed..be05fba63c 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -322,16 +322,15 @@ def generate_sparse_seqlen_kernel( def generate_sparse_seqlen_triton( - query_lens: torch.Tensor, seq_lens: torch.Tensor, cu_query_lens: torch.Tensor, topk_token: int, num_tokens: int, max_query_len: int, ): - num_seqs = query_lens.size(0) + num_seqs = cu_query_lens.size(0) - 1 # zero initialize the tensor to make sure invalid positions will be zero - out = torch.zeros([num_tokens], dtype=torch.int32, device=query_lens.device) + out = torch.zeros([num_tokens], dtype=torch.int32, device=cu_query_lens.device) block_size = 64 num_block_per_row = triton.cdiv(max_query_len, block_size) grid = ( @@ -627,10 +626,8 @@ class ROCMAiterMLASparseMetadataBuilder( np_to_pinned_tensor(req_id_per_token), non_blocking=True ) query_start_loc = common_attn_metadata.query_start_loc[: num_reqs + 1] - query_lens = query_start_loc[1:] - query_start_loc[:-1] seq_lens = common_attn_metadata.seq_lens[:num_reqs] sparse_seqlen = generate_sparse_seqlen_triton( - query_lens, seq_lens, query_start_loc, self.topk_tokens, @@ -639,7 +636,9 @@ class ROCMAiterMLASparseMetadataBuilder( ) torch.cumsum(sparse_seqlen, dim=0, out=self.paged_kv_indptr[1 : num_tokens + 1]) - self.paged_kv_indptr[num_tokens + 1 :].fill_(self.paged_kv_indptr[num_tokens]) + # fill_() with a value that aliases the destination clones it first. + indptr_tail = self.paged_kv_indptr[num_tokens + 1 :] + indptr_tail.copy_(self.paged_kv_indptr[num_tokens].expand_as(indptr_tail)) req_id_per_token = self.req_id_per_token_buffer[:num_tokens] qo_indptr = self.qo_indptr[: num_tokens + 1]