From 060b5f61dcb103ad86f091364b180985fa5c6867 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Wed, 22 Jul 2026 04:23:48 +0200 Subject: [PATCH] [Bugfix][Attention] Ignore empty MLA context chunks during merge (#49294) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .../configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml | 1 + .../configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml | 1 + .../attention/test_merge_attn_states.py | 56 +++++++++ .../layers/attention/mla_attention.py | 19 +++ .../attention/ops/triton_merge_attn_states.py | 119 ++++++++++++++++++ 5 files changed, 196 insertions(+) diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml index 13d71af20e..14b4f2c0d1 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml +++ b/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 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml index b21d0b12f0..1a96afe46d 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml +++ b/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 diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index 40af84887a..d5c9a1f474 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -11,6 +11,9 @@ from vllm._custom_ops import ( 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 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16] 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 = ( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 14662e9fc0..884c4ca4e6 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -277,6 +277,7 @@ from vllm.v1.attention.backends.utils import ( 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 MLACommonPrefillMetadata: 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 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): 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 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): 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 diff --git a/vllm/v1/attention/ops/triton_merge_attn_states.py b/vllm/v1/attention/ops/triton_merge_attn_states.py index 14a52ada97..ca06c2970b 100644 --- a/vllm/v1/attention/ops/triton_merge_attn_states.py +++ b/vllm/v1/attention/ops/triton_merge_attn_states.py @@ -9,6 +9,118 @@ from vllm.triton_utils import tl, triton 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))