iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][Attention] Ignore empty MLA context chunks during merge (#49294)

Browse Source 
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Lucas Wilkinson
2026-07-22 02:23:48 +00:00
committed by GitHub
co-authored by OpenAI Codex
parent ec59c1579f
commit 060b5f61dc
5 changed files with 196 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+1
View File
@@ -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
+1
View File
@@ -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
+56
View File
@@ -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 = (
vllm/model_executor/layers/attention/mla_attention.py
+19
View File
@@ -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
vllm/v1/attention/ops/triton_merge_attn_states.py
+119
View File
@@ -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))