[Perf] Fix TRTLLM ragged prefill perf regression (#55499)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Wei Zhao committed
Sep 9, 2026 at 16:38 UTC
8c87c333b84c85908b1d11f0044457692277c6f3
6 files changed
+51
-4
docker/Dockerfile
+1
-1
@@ -960,7 +960,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
960
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
961
# https://docs.flashinfer.ai/installation.html
962
# From versions.json: .flashinfer.version
963
-ARG FLASHINFER_VERSION=0.6.18
963
+ARG FLASHINFER_VERSION=0.6.18.post1
964
RUN --mount=type=cache,target=/opt/uv/cache \
965
if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \
966
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
docker/versions.json
+1
-1
@@ -71,7 +71,7 @@
71
"default": "true"
72
},
73
"FLASHINFER_VERSION": {
74
- "default": "0.6.18"
74
+ "default": "0.6.18.post1"
75
},
76
"GDRCOPY_CUDA_VERSION": {
77
"default": "12.8"
requirements/cuda.txt
+2
-2
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
14
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
15
# install_requires so the published wheel does not carry an unresolvable pin
16
--extra-index-url https://flashinfer.ai/whl/
17
-flashinfer-python==0.6.18
18
-flashinfer-cubin==0.6.18
17
+flashinfer-python==0.6.18.post1
18
+flashinfer-cubin==0.6.18.post1
19
apache-tvm-ffi==0.1.11
20
tilelang==0.1.12
21
nvidia-cudnn-frontend>=1.19.1
tests/v1/attention/test_mla_context_chunks.py
+4
@@ -57,6 +57,9 @@ def build_chunked_context(
57
([1024, 1024, 1024], 1024),
58
# Unaligned tails.
59
([37, 1000, 5], 1024),
60
+ # A zero-context request between active requests must not become an
61
+ # empty TRTLLM ragged-attention row.
62
+ ([4, 0, 2], 1024),
63
([1], 1024),
64
],
65
)
@@ -74,6 +77,7 @@ def test_chunks_gather_every_context_row_exactly_once(context_lens, workspace_si
77
gathered: dict[int, list[tuple[int, int]]] = {}
78
previous_request_start = -1
79
for chunk in metadata.chunks:
80
+ assert chunk.all_rows_active
81
assert chunk.num_context_tokens <= workspace_size
82
assert chunk.num_context_tokens == int(chunk.cu_seq_lens[-1])
83
assert chunk.token_to_seq.shape[0] == chunk.num_context_tokens
vllm/model_executor/layers/attention/mla_attention.py
+4
@@ -1547,6 +1547,7 @@ class MLACommonPrefillMetadata:
1547
max_seq_len: int
1548
seq_lens: torch.Tensor
1549
token_to_seq: torch.Tensor
1550
+ all_rows_active: bool
1551
1552
# for mla DCP
1553
padded_local_seq_lens: list[int] | None = None
@@ -2124,6 +2125,8 @@ def build_mla_chunked_context_metadata(
2125
max_seq_len=max(plan.seq_lens),
2126
seq_lens=seq_lens_cpu[request_slice],
2127
token_to_seq=token_to_seq[token_slice],
2128
+ all_rows_active=all(length > 0 for length in query_lens)
2129
+ and all(length > 0 for length in plan.seq_lens),
2130
num_local_context_tokens=local_token_slice.stop - local_token_slice.start,
2131
)
2132
if use_dcp:
@@ -2529,6 +2532,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2532
output_dtype=self.model_config.dtype,
2533
q_data_type=self.q_data_type,
2534
prefill_backend=self._prefill_backend,
2535
+ query_lens_cpu=prefill_query_lens_cpu,
2536
)
2537
2538
self._prefill_backend.prepare_metadata(prefill_metadata)
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+39
@@ -91,6 +91,26 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
91
self._query_seq_lens = (
92
prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
93
)
94
+ query_lens_cpu = prefill_metadata.query_lens_cpu
95
+ if query_lens_cpu is None:
96
+ raise ValueError("TRTLLM ragged prefill requires CPU query lengths")
97
+ if query_lens_cpu.device.type != "cpu":
98
+ raise ValueError("TRTLLM ragged prefill query lengths must be on CPU")
99
+ if query_lens_cpu.ndim != 1 or query_lens_cpu.numel() == 0:
100
+ raise ValueError(
101
+ "TRTLLM ragged prefill requires a non-empty 1D query-length tensor"
102
+ )
103
+ min_query_len = int(torch.min(query_lens_cpu).item())
104
+ if min_query_len < 0:
105
+ raise ValueError("TRTLLM ragged prefill query lengths must be non-negative")
106
+ self._has_active_rows = min_query_len > 0
107
+ if not self._has_active_rows:
108
+ has_mixed_rows = bool(torch.any(query_lens_cpu > 0).item())
109
+ if has_mixed_rows:
110
+ raise ValueError(
111
+ "TRTLLM ragged prefill contains mixed active and empty query "
112
+ f"rows: query_lens={query_lens_cpu.tolist()}"
113
+ )
114
115
def supports_out(self) -> bool:
116
# Output head dim is v.shape[-1] == v_head_dim, so `out` is unpadded.
@@ -116,6 +136,18 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
136
dtype=self._prefill_metadata.output_dtype,
137
)
138
139
+ if not self._has_active_rows:
140
+ out.zero_()
141
+ if return_softmax_lse:
142
+ lse = torch.empty(
143
+ self.num_heads,
144
+ q.shape[0],
145
+ dtype=torch.float32,
146
+ device=q.device,
147
+ )
148
+ return out, lse.fill_(-float("inf"))
149
+ return out
150
+
151
ret = trtllm_ragged_attention_deepseek(
152
query=q,
153
key=k,
@@ -135,6 +167,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
167
is_causal=True,
168
return_lse=return_softmax_lse,
169
out=out,
170
+ skip_all_rows_active_check=True,
171
)
172
173
if isinstance(ret, tuple):
@@ -152,6 +185,11 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
185
) -> tuple[torch.Tensor, torch.Tensor]:
186
from flashinfer.prefill import trtllm_ragged_attention_deepseek
187
188
+ if not chunk.all_rows_active:
189
+ raise ValueError(
190
+ "TRTLLM ragged context prefill contains an empty query or KV row"
191
+ )
192
+
193
if out is None:
194
out = torch.empty(
195
q.shape[0],
@@ -180,6 +218,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
218
is_causal=False,
219
return_lse=True,
220
out=out,
221
+ skip_all_rows_active_check=True,
222
)
223
224
# Convert from (q_len, num_heads) to (num_heads, q_len)