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

[LucasWilkinson]

Summary

FIX for #49334
Alternative to #49196

• mark zero-context MLA attention states with -inf LSE before merging
• use a single in-place Triton kernel driven by the existing ragged query and context offsets
• resolve query-block ownership 32 request boundaries at a time with a warp-local reduction, amortized across all attention heads
• launch the mask only for prefill chunks containing an empty context row
• avoid per-token context-length or token-to-request metadata

Wide PCP batches can split cached context into workspace-sized chunks. A later chunk can contain no KV tokens for some prefill rows. The backend may return a finite LSE for those empty rows, causing merge_attn_states to treat undefined output as valid attention and corrupt the prefill result.

Validation

Focused kernel test:

CUDA_VISIBLE_DEVICES=0 \
.venv/bin/python -m pytest \
  tests/kernels/attention/test_merge_attn_states.py \
  -k test_mask_empty_context_lse -q

Result: 1 passed.

All file-scoped pre-commit hooks passed, including Ruff, formatting, mypy, SPDX, configuration validation, and forbidden-import checks.

Warmed kernel microbenchmark on an NVIDIA B300 with 64 attention heads:

Shape	Warp-parallel Triton
2K balanced	12.51 us
32K balanced	12.57 us
131K skewed	13.18 us

The kernel preserves token-proportional load balancing without adding token-to-request metadata.

Runtime configuration (4 GPUs):

CUDA_VISIBLE_DEVICES=0,1,2,3 \
VLLM_USE_V2_MODEL_RUNNER=1 \
.venv/bin/vllm serve nvidia/GLM-5.2-NVFP4 \
  --enforce-eager \
  --max-model-len 4096 \
  --max-num-batched-tokens 32768 \
  --safetensors-load-strategy prefetch \
  --moe-backend flashinfer_cutlass \
  --tensor-parallel-size 1 \
  --prefill-context-parallel-size 4 \
  --enable-expert-parallel \
  --kv-cache-dtype fp8 \
  --trust-remote-code

100-sample GSM8K smoke evaluation:

.venv/bin/python tests/evals/gsm8k/gsm8k_eval.py \
  --num-questions 100 \
  --num-shots 5 \
  --max-tokens 256 \
  --max-concurrency 100 \
  --port 8000

Deployment	Before	After	Invalid after
TP1+PCP4+EP4	1/100	93/100	0/100

The Triton implementation completed the evaluation in 47.862 seconds. For comparison, disabling prefix caching scored 96/100, consistent with the bug being in cached-context merging rather than sampling.

A 300-sample, three-wave run with the same 32K token budget scored 279/300 (93.0%) with no invalid responses. The PCP eval YAMLs now pin that validated token budget instead of falling back to the 8K default.

AI assistance

AI assistance was used to diagnose and implement this draft. This PR should remain a draft until the human submitter reviews every changed line, confirms the test and evaluation results, and can explain and defend the change end-to-end.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[MatthewBonanni]

LGTM, thanks!