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

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml

@@ -6,6 +6,7 @@ max_concurrency: 100
 server_args: >-
   --enforce-eager
   --max-model-len 4096
+  --max-num-batched-tokens 32768
   --safetensors-load-strategy prefetch
   --moe-backend flashinfer_cutlass
   --prefill-context-parallel-size 4

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml

@@ -6,6 +6,7 @@ max_concurrency: 100
 server_args: >-
   --enforce-eager
   --max-model-len 4096
+  --max-num-batched-tokens 32768
   --safetensors-load-strategy prefetch
   --moe-backend flashinfer_cutlass
   --tensor-parallel-size 2

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,
+)
 from vllm.v1.attention.ops.triton_merge_attn_states import (
     merge_attn_states as merge_attn_states_triton,
 )
@@ -73,6 +76,59 @@ def merge_attn_states_torch(
 all_case_info: list[tuple] = []
 
 
+def test_mask_empty_context() -> 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()
+    num_heads, num_tokens, head_dim = 4, 165, 16
+    lse = torch.randn(num_heads, num_tokens, device="cuda")
+    output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
+    # Empty-context rows carry undefined (possibly non-finite) attention output.
+    output[33:164] = float("nan")
+
+    expected_lse = lse.clone()
+    expected_lse[:, 33:164] = float("-inf")
+    expected_output = output.clone()
+    expected_output[33:164] = 0.0
+
+    mask_empty_context(lse, output, query_start_loc, context_start_loc)
+
+    torch.testing.assert_close(lse, expected_lse)
+    torch.testing.assert_close(output, expected_output)
+
+
+@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
+@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
+def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
+    """When a token is empty on both sides (both LSE -inf), the 0/0 softmax
+    scales must not surface as NaN in the merged output."""
+    num_tokens, num_heads, head_size = 6, 8, 128
+    prefix_output = torch.zeros(
+        num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
+    )
+    prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
+    suffix_output = torch.zeros(
+        num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
+    )
+    suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
+
+    # Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
+    # their outputs and set both LSEs to -inf).
+    empty = slice(2, 4)
+    prefix_lse[:, empty] = float("-inf")
+    suffix_lse[:, empty] = float("-inf")
+
+    output = torch.empty_like(prefix_output)
+    merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
+
+    assert not output.isnan().any()
+
+
 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
 from vllm.v1.attention.selector import get_attn_backend
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
@@ -1383,6 +1384,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
@@ -1592,6 +1594,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
@@ -1670,6 +1673,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(),
@@ -1692,6 +1696,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,
         )
 
@@ -2279,6 +2284,13 @@ def _compute_prefill_context(
                     v=v,
                 )
             )
+            if prefill_metadata.chunked_context.has_empty_context[i]:
+                mask_empty_context(
+                    attn_softmax_lse,
+                    attn_output,
+                    prefill_metadata.query_start_loc,
+                    prefill_metadata.chunked_context.cu_seq_lens[i],
+                )
 
             if output is None:
                 output = attn_output
@@ -2429,6 +2441,13 @@ def _context_parallel_compute_prefill_context(
                     v=v,
                 )
             )
+            if prefill_metadata.chunked_context.has_empty_context[i]:
+                mask_empty_context(
+                    attn_softmax_lse,
+                    attn_output,
+                    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,118 @@
 float8_info = torch.finfo(current_platform.fp8_dtype())
 
 
+def mask_empty_context(
+    lse: torch.Tensor,
+    output: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    context_start_loc: torch.Tensor,
+) -> None:
+    """Neutralize context chunks that cover no keys before merging.
+
+    A prefill query whose context chunk is empty attended to no keys, so its
+    partial attention is undefined: the backend leaves the output rows as
+    uninitialized scratch (which may hold NaN/Inf) even when it reports an LSE
+    of -inf. Sanitize both here so ``merge_attn_states`` can stay generic:
+    force the LSE to -inf (zero softmax weight) and zero the undefined output
+    rows (so a zero weight cannot combine with NaN/Inf). Emptiness is derived
+    from the context offsets, not from the -inf LSE, so no merge kernel has to
+    reason about undefined partials.
+
+    Args:
+        lse: Chunk log-sum-exp, shape [num_heads, num_tokens].
+        output: Chunk attention output, shape [num_tokens, num_heads, ...].
+        query_start_loc: Prefill query cumulative offsets, shape [num_reqs + 1].
+        context_start_loc: Chunk context cumulative offsets,
+            shape [num_reqs + 1]; an empty chunk has a zero-length span.
+    """
+    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
+    is_empty = torch.zeros(num_tokens, dtype=torch.bool, device=lse.device)
+    mask_empty_context_kernel[(num_query_blocks,)](
+        lse,
+        is_empty,
+        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,
+    )
+    output.masked_fill_(is_empty[:, None, None], 0.0)
+
+
+@triton.jit
+def mask_empty_context_kernel(
+    lse,
+    is_empty,
+    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
+    tl.store(is_empty + token_indices, True, mask=valid_tokens)
+    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(
@@ -136,6 +248,9 @@ def merge_attn_states_kernel(
 
     if OUTPUT_LSE:
         out_lse = tl.log(out_se) + max_lse
+        # Both sides empty (max_lse == -inf) => undefined merge; keep -inf so
+        # downstream merges continue to treat the token as empty.
+        out_lse = tl.where(max_lse == float("-inf"), float("-inf"), out_lse)
         tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse)
 
     p_out = tl.load(
@@ -159,6 +274,10 @@ def merge_attn_states_kernel(
     p_scale = p_se / out_se
     s_scale = s_se / out_se
     out = p_out * p_scale + s_out * s_scale
+    # If both sides are empty (max_lse == -inf) the scales are 0/0 = NaN; emit
+    # zeros rather than NaN. Callers with empty chunks (see mask_empty_context)
+    # zero those inputs, so this only guards the fully-undefined corner.
+    out = tl.where(max_lse == float("-inf"), 0.0, out)
 
     if USE_FP8:
         out = out * (1.0 / tl.load(output_scale))