[Bugfix][Attention] Ignore empty MLA context chunks during merge

tests/kernels/attention/test_merge_attn_states.py

@@ -11,6 +11,9 @@
     scaled_fp8_quant,
 )
 from vllm.platforms import current_platform
+from vllm.v1.attention.ops.triton_merge_attn_states import (
+    mask_empty_context_lse,
+)
 from vllm.v1.attention.ops.triton_merge_attn_states import (
     merge_attn_states as merge_attn_states_triton,
 )
@@ -73,6 +76,24 @@ def merge_attn_states_torch(
 all_case_info: list[tuple] = []
 
 
+def test_mask_empty_context_lse() -> None:
+    query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
+    query_start_loc = torch.cat(
+        (torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
+    ).cuda()
+    context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
+    context_start_loc = torch.cat(
+        (torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
+    ).cuda()
+    lse = torch.randn(4, 165, device="cuda")
+    expected = lse.clone()
+    expected[:, 33:164] = float("-inf")
+
+    mask_empty_context_lse(lse, query_start_loc, context_start_loc)
+
+    torch.testing.assert_close(lse, expected)
+
+
 def generate_markdown_table():
     global all_case_info
     table_header = (

vllm/model_executor/layers/attention/mla_attention.py

@@ -277,6 +277,7 @@
 from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
+from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context_lse
 from vllm.v1.attention.selector import get_attn_backend
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
@@ -1342,6 +1343,7 @@ class ChunkedContextMetadata:
         workspace: torch.Tensor
         token_to_seq: torch.Tensor
         chunk_total_token: list[int]
+        has_empty_context: list[bool]
 
         # for mla DCP
         padded_local_chunk_seq_lens: list[list[int]] | None = None
@@ -1551,6 +1553,7 @@ def build_mla_chunked_context_metadata(
     )
     chunk_seq_lens = chunk_ends - chunk_starts
     chunk_seq_lens.clamp_(min=0)
+    has_empty_context = torch.any(chunk_seq_lens == 0, dim=1).tolist()
 
     cu_seq_lens_cpu = torch.zeros(
         num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
@@ -1629,6 +1632,7 @@ def build_mla_chunked_context_metadata(
             token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
             chunk_total_token=chunk_total_token.tolist(),
             workspace=chunked_prefill_workspace,
+            has_empty_context=has_empty_context,
             prefill_tokens_with_context=prefill_tokens_with_context,
             padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
             local_context_lens_allranks=local_context_lens_allranks.tolist(),
@@ -1651,6 +1655,7 @@ def build_mla_chunked_context_metadata(
             token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
             chunk_total_token=chunk_total_token,
             workspace=chunked_prefill_workspace,
+            has_empty_context=has_empty_context,
             prefill_tokens_with_context=prefill_tokens_with_context,
         )
 
@@ -2238,6 +2243,12 @@ def _compute_prefill_context(
                     v=v,
                 )
             )
+            if prefill_metadata.chunked_context.has_empty_context[i]:
+                mask_empty_context_lse(
+                    attn_softmax_lse,
+                    prefill_metadata.query_start_loc,
+                    prefill_metadata.chunked_context.cu_seq_lens[i],
+                )
 
             if output is None:
                 output = attn_output
@@ -2388,6 +2399,12 @@ def _context_parallel_compute_prefill_context(
                     v=v,
                 )
             )
+            if prefill_metadata.chunked_context.has_empty_context[i]:
+                mask_empty_context_lse(
+                    attn_softmax_lse,
+                    prefill_metadata.query_start_loc,
+                    prefill_metadata.chunked_context.cu_seq_lens[i],
+                )
 
             if output is None:
                 output = attn_output

vllm/v1/attention/ops/triton_merge_attn_states.py

@@ -9,6 +9,95 @@
 float8_info = torch.finfo(current_platform.fp8_dtype())
 
 
+def mask_empty_context_lse(
+    lse: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    context_start_loc: torch.Tensor,
+) -> None:
+    """Set LSE to -inf for queries whose context chunk is empty."""
+    num_heads, num_tokens = lse.shape
+    num_reqs = query_start_loc.shape[0] - 1
+    block_size = 128
+    # Reserve the worst-case number of request-local blocks.
+    num_query_blocks = num_tokens // block_size + num_reqs
+    mask_empty_context_lse_kernel[(num_query_blocks,)](
+        lse,
+        query_start_loc,
+        context_start_loc,
+        lse.stride(0),
+        lse.stride(1),
+        num_reqs,
+        NUM_HEADS=num_heads,
+        BLOCK_SIZE=block_size,
+        BLOCK_HEADS=8,
+        num_warps=8,
+    )
+
+
+@triton.jit
+def mask_empty_context_lse_kernel(
+    lse,
+    query_start_loc,
+    context_start_loc,
+    lse_head_stride,
+    lse_token_stride,
+    num_reqs,
+    NUM_HEADS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCK_HEADS: tl.constexpr,
+):
+    query_block_idx = tl.program_id(0)
+
+    lanes = tl.arange(0, 32)
+    chunk_start = 0
+    req_idx = 0
+    req_idx_found = False
+    while (chunk_start < num_reqs) & (not req_idx_found):
+        req_offsets = chunk_start + lanes
+        req_mask = req_offsets < num_reqs
+        query_starts = tl.load(query_start_loc + req_offsets, mask=req_mask)
+        # Assume the worst-case number of blocks for each request.
+        req_block_starts = query_starts // BLOCK_SIZE + req_offsets
+        matched_idx = tl.sum(
+            (req_mask & (req_block_starts <= query_block_idx)).to(tl.int32)
+        )
+        # matched_idx == 32 means the match is past this warp chunk.
+        req_idx = chunk_start + matched_idx - 1
+        req_idx_found = matched_idx < 32
+        chunk_start += 32
+
+    query_start = tl.load(query_start_loc + req_idx)
+    query_end = tl.load(query_start_loc + req_idx + 1)
+    query_len = query_end - query_start
+    req_first_block = query_start // BLOCK_SIZE + req_idx
+    block_in_req = query_block_idx - req_first_block
+    token_offset = block_in_req * BLOCK_SIZE
+    if token_offset >= query_len:
+        return
+
+    context_start = tl.load(context_start_loc + req_idx)
+    context_end = tl.load(context_start_loc + req_idx + 1)
+    if context_start != context_end:
+        return
+
+    token_offsets = token_offset + tl.arange(0, BLOCK_SIZE)
+    token_indices = query_start + token_offsets
+    token_lse_offsets = token_indices * lse_token_stride
+    valid_tokens = token_offsets < query_len
+    head_offsets = tl.arange(0, BLOCK_HEADS)
+    for head_start in range(0, NUM_HEADS, BLOCK_HEADS):
+        head_indices = head_start + head_offsets
+        lse_ptrs = (
+            lse + head_indices[:, None] * lse_head_stride + token_lse_offsets[None, :]
+        )
+        valid_heads = head_indices < NUM_HEADS
+        tl.store(
+            lse_ptrs,
+            float("-inf"),
+            mask=valid_heads[:, None] & valid_tokens[None, :],
+        )
+
+
 # Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
 # can be used to combine partial attention results (in the split-KV case)
 def merge_attn_states(