[Bugfix][Attention] Ignore empty MLA context chunks during merge (#49294)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Lucas Wilkinson committed
Jul 22, 2026 at 04:23 UTC
060b5f61dcb103ad86f091364b180985fa5c6867
5 files changed
+196
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+1
@@ -6,6 +6,7 @@ max_concurrency: 100
6
server_args: >-
7
--enforce-eager
8
--max-model-len 4096
9
+ --max-num-batched-tokens 32768
10
--safetensors-load-strategy prefetch
11
--moe-backend flashinfer_cutlass
12
--prefill-context-parallel-size 4
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+1
@@ -6,6 +6,7 @@ max_concurrency: 100
6
server_args: >-
7
--enforce-eager
8
--max-model-len 4096
9
+ --max-num-batched-tokens 32768
10
--safetensors-load-strategy prefetch
11
--moe-backend flashinfer_cutlass
12
--tensor-parallel-size 2
tests/kernels/attention/test_merge_attn_states.py
+56
@@ -11,6 +11,9 @@ from vllm._custom_ops import (
11
scaled_fp8_quant,
12
)
13
from vllm.platforms import current_platform
14
+from vllm.v1.attention.ops.triton_merge_attn_states import (
15
+ mask_empty_context,
16
+)
17
from vllm.v1.attention.ops.triton_merge_attn_states import (
18
merge_attn_states as merge_attn_states_triton,
19
)
@@ -73,6 +76,59 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
76
all_case_info: list[tuple] = []
77
78
79
+def test_mask_empty_context() -> None:
80
+ query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
81
+ query_start_loc = torch.cat(
82
+ (torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
83
+ ).cuda()
84
+ context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
85
+ context_start_loc = torch.cat(
86
+ (torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
87
+ ).cuda()
88
+ num_heads, num_tokens, head_dim = 4, 165, 16
89
+ lse = torch.randn(num_heads, num_tokens, device="cuda")
90
+ output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
91
+ # Empty-context rows carry undefined (possibly non-finite) attention output.
92
+ output[33:164] = float("nan")
93
+
94
+ expected_lse = lse.clone()
95
+ expected_lse[:, 33:164] = float("-inf")
96
+ expected_output = output.clone()
97
+ expected_output[33:164] = 0.0
98
+
99
+ mask_empty_context(lse, output, query_start_loc, context_start_loc)
100
+
101
+ torch.testing.assert_close(lse, expected_lse)
102
+ torch.testing.assert_close(output, expected_output)
103
+
104
+
105
+@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
106
+@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
107
+def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
108
+ """When a token is empty on both sides (both LSE -inf), the 0/0 softmax
109
+ scales must not surface as NaN in the merged output."""
110
+ num_tokens, num_heads, head_size = 6, 8, 128
111
+ prefix_output = torch.zeros(
112
+ num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
113
+ )
114
+ prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
115
+ suffix_output = torch.zeros(
116
+ num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
117
+ )
118
+ suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
119
+
120
+ # Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
121
+ # their outputs and set both LSEs to -inf).
122
+ empty = slice(2, 4)
123
+ prefix_lse[:, empty] = float("-inf")
124
+ suffix_lse[:, empty] = float("-inf")
125
+
126
+ output = torch.empty_like(prefix_output)
127
+ merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
128
+
129
+ assert not output.isnan().any()
130
+
131
+
132
def generate_markdown_table():
133
global all_case_info
134
table_header = (
vllm/model_executor/layers/attention/mla_attention.py
+19
@@ -277,6 +277,7 @@ from vllm.v1.attention.backends.utils import (
277
from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
278
from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
279
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
280
+from vllm.v1.attention.ops.triton_merge_attn_states import mask_empty_context
281
from vllm.v1.attention.selector import get_attn_backend
282
from vllm.v1.kv_cache_interface import (
283
AttentionSpec,
@@ -1383,6 +1384,7 @@ class MLACommonPrefillMetadata:
1384
workspace: torch.Tensor
1385
token_to_seq: torch.Tensor
1386
chunk_total_token: list[int]
1387
+ has_empty_context: list[bool]
1388
1389
# for mla DCP
1390
padded_local_chunk_seq_lens: list[list[int]] | None = None
@@ -1592,6 +1594,7 @@ def build_mla_chunked_context_metadata(
1594
)
1595
chunk_seq_lens = chunk_ends - chunk_starts
1596
chunk_seq_lens.clamp_(min=0)
1597
+ has_empty_context = torch.any(chunk_seq_lens == 0, dim=1).tolist()
1598
1599
cu_seq_lens_cpu = torch.zeros(
1600
num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
@@ -1670,6 +1673,7 @@ def build_mla_chunked_context_metadata(
1673
token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
1674
chunk_total_token=chunk_total_token.tolist(),
1675
workspace=chunked_prefill_workspace,
1676
+ has_empty_context=has_empty_context,
1677
prefill_tokens_with_context=prefill_tokens_with_context,
1678
padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
1679
local_context_lens_allranks=local_context_lens_allranks.tolist(),
@@ -1692,6 +1696,7 @@ def build_mla_chunked_context_metadata(
1696
token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
1697
chunk_total_token=chunk_total_token,
1698
workspace=chunked_prefill_workspace,
1699
+ has_empty_context=has_empty_context,
1700
prefill_tokens_with_context=prefill_tokens_with_context,
1701
)
1702
@@ -2279,6 +2284,13 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2284
v=v,
2285
)
2286
)
2287
+ if prefill_metadata.chunked_context.has_empty_context[i]:
2288
+ mask_empty_context(
2289
+ attn_softmax_lse,
2290
+ attn_output,
2291
+ prefill_metadata.query_start_loc,
2292
+ prefill_metadata.chunked_context.cu_seq_lens[i],
2293
+ )
2294
2295
if output is None:
2296
output = attn_output
@@ -2429,6 +2441,13 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2441
v=v,
2442
)
2443
)
2444
+ if prefill_metadata.chunked_context.has_empty_context[i]:
2445
+ mask_empty_context(
2446
+ attn_softmax_lse,
2447
+ attn_output,
2448
+ prefill_metadata.query_start_loc,
2449
+ prefill_metadata.chunked_context.cu_seq_lens[i],
2450
+ )
2451
2452
if output is None:
2453
output = attn_output
vllm/v1/attention/ops/triton_merge_attn_states.py
+119
@@ -9,6 +9,118 @@ from vllm.triton_utils import tl, triton
9
float8_info = torch.finfo(current_platform.fp8_dtype())
10
11
12
+def mask_empty_context(
13
+ lse: torch.Tensor,
14
+ output: torch.Tensor,
15
+ query_start_loc: torch.Tensor,
16
+ context_start_loc: torch.Tensor,
17
+) -> None:
18
+ """Neutralize context chunks that cover no keys before merging.
19
+
20
+ A prefill query whose context chunk is empty attended to no keys, so its
21
+ partial attention is undefined: the backend leaves the output rows as
22
+ uninitialized scratch (which may hold NaN/Inf) even when it reports an LSE
23
+ of -inf. Sanitize both here so ``merge_attn_states`` can stay generic:
24
+ force the LSE to -inf (zero softmax weight) and zero the undefined output
25
+ rows (so a zero weight cannot combine with NaN/Inf). Emptiness is derived
26
+ from the context offsets, not from the -inf LSE, so no merge kernel has to
27
+ reason about undefined partials.
28
+
29
+ Args:
30
+ lse: Chunk log-sum-exp, shape [num_heads, num_tokens].
31
+ output: Chunk attention output, shape [num_tokens, num_heads, ...].
32
+ query_start_loc: Prefill query cumulative offsets, shape [num_reqs + 1].
33
+ context_start_loc: Chunk context cumulative offsets,
34
+ shape [num_reqs + 1]; an empty chunk has a zero-length span.
35
+ """
36
+ num_heads, num_tokens = lse.shape
37
+ num_reqs = query_start_loc.shape[0] - 1
38
+ block_size = 128
39
+ # Reserve the worst-case number of request-local blocks.
40
+ num_query_blocks = num_tokens // block_size + num_reqs
41
+ is_empty = torch.zeros(num_tokens, dtype=torch.bool, device=lse.device)
42
+ mask_empty_context_kernel[(num_query_blocks,)](
43
+ lse,
44
+ is_empty,
45
+ query_start_loc,
46
+ context_start_loc,
47
+ lse.stride(0),
48
+ lse.stride(1),
49
+ num_reqs,
50
+ NUM_HEADS=num_heads,
51
+ BLOCK_SIZE=block_size,
52
+ BLOCK_HEADS=8,
53
+ num_warps=8,
54
+ )
55
+ output.masked_fill_(is_empty[:, None, None], 0.0)
56
+
57
+
58
+@triton.jit
59
+def mask_empty_context_kernel(
60
+ lse,
61
+ is_empty,
62
+ query_start_loc,
63
+ context_start_loc,
64
+ lse_head_stride,
65
+ lse_token_stride,
66
+ num_reqs,
67
+ NUM_HEADS: tl.constexpr,
68
+ BLOCK_SIZE: tl.constexpr,
69
+ BLOCK_HEADS: tl.constexpr,
70
+):
71
+ query_block_idx = tl.program_id(0)
72
+
73
+ lanes = tl.arange(0, 32)
74
+ chunk_start = 0
75
+ req_idx = 0
76
+ req_idx_found = False
77
+ while (chunk_start < num_reqs) & (not req_idx_found):
78
+ req_offsets = chunk_start + lanes
79
+ req_mask = req_offsets < num_reqs
80
+ query_starts = tl.load(query_start_loc + req_offsets, mask=req_mask)
81
+ # Assume the worst-case number of blocks for each request.
82
+ req_block_starts = query_starts // BLOCK_SIZE + req_offsets
83
+ matched_idx = tl.sum(
84
+ (req_mask & (req_block_starts <= query_block_idx)).to(tl.int32)
85
+ )
86
+ # matched_idx == 32 means the match is past this warp chunk.
87
+ req_idx = chunk_start + matched_idx - 1
88
+ req_idx_found = matched_idx < 32
89
+ chunk_start += 32
90
+
91
+ query_start = tl.load(query_start_loc + req_idx)
92
+ query_end = tl.load(query_start_loc + req_idx + 1)
93
+ query_len = query_end - query_start
94
+ req_first_block = query_start // BLOCK_SIZE + req_idx
95
+ block_in_req = query_block_idx - req_first_block
96
+ token_offset = block_in_req * BLOCK_SIZE
97
+ if token_offset >= query_len:
98
+ return
99
+
100
+ context_start = tl.load(context_start_loc + req_idx)
101
+ context_end = tl.load(context_start_loc + req_idx + 1)
102
+ if context_start != context_end:
103
+ return
104
+
105
+ token_offsets = token_offset + tl.arange(0, BLOCK_SIZE)
106
+ token_indices = query_start + token_offsets
107
+ token_lse_offsets = token_indices * lse_token_stride
108
+ valid_tokens = token_offsets < query_len
109
+ tl.store(is_empty + token_indices, True, mask=valid_tokens)
110
+ head_offsets = tl.arange(0, BLOCK_HEADS)
111
+ for head_start in range(0, NUM_HEADS, BLOCK_HEADS):
112
+ head_indices = head_start + head_offsets
113
+ lse_ptrs = (
114
+ lse + head_indices[:, None] * lse_head_stride + token_lse_offsets[None, :]
115
+ )
116
+ valid_heads = head_indices < NUM_HEADS
117
+ tl.store(
118
+ lse_ptrs,
119
+ float("-inf"),
120
+ mask=valid_heads[:, None] & valid_tokens[None, :],
121
+ )
122
+
123
+
124
# Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
125
# can be used to combine partial attention results (in the split-KV case)
126
def merge_attn_states(
@@ -136,6 +248,9 @@ def merge_attn_states_kernel(
248
249
if OUTPUT_LSE:
250
out_lse = tl.log(out_se) + max_lse
251
+ # Both sides empty (max_lse == -inf) => undefined merge; keep -inf so
252
+ # downstream merges continue to treat the token as empty.
253
+ out_lse = tl.where(max_lse == float("-inf"), float("-inf"), out_lse)
254
tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse)
255
256
p_out = tl.load(
@@ -159,6 +274,10 @@ def merge_attn_states_kernel(
274
p_scale = p_se / out_se
275
s_scale = s_se / out_se
276
out = p_out * p_scale + s_out * s_scale
277
+ # If both sides are empty (max_lse == -inf) the scales are 0/0 = NaN; emit
278
+ # zeros rather than NaN. Callers with empty chunks (see mask_empty_context)
279
+ # zero those inputs, so this only guards the fully-undefined corner.
280
+ out = tl.where(max_lse == float("-inf"), 0.0, out)
281
282
if USE_FP8:
283
out = out * (1.0 / tl.load(output_scale))