A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 17:15 UTC
[Bugfix][MLA] Add logits size budget to sparse indexer prefill chunking (#36178)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Apr 1, 2026 at 00:15 UTC
eb47454987452c7ee10253839b0c7289c3352c60
4 files changed
+191
-31
tests/v1/attention/test_sparse_mla_backends.py
+76
@@ -42,6 +42,7 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import (
42
FlashMLASparseBackend,
43
triton_convert_req_index_to_global_index,
44
)
45
+from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks
46
from vllm.v1.attention.backends.utils import split_prefill_chunks
47
from vllm.v1.attention.ops import flashmla
48
@@ -716,6 +717,81 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected):
717
assert out == expected
718
719
720
+@pytest.mark.parametrize(
721
+ "seq_lens,query_lens,workspace_size,max_logits_bytes,expected",
722
+ [
723
+ # Logits constraint triggers split (M*N exceeds budget)
724
+ # req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000
725
+ # req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250
726
+ (
727
+ torch.tensor([100, 100, 100]),
728
+ torch.tensor([10, 10, 10]),
729
+ 1000, # workspace allows all
730
+ 5000, # 1250 float32 elems -> forces split
731
+ [
732
+ (slice(0, 1), slice(0, 10)),
733
+ (slice(1, 2), slice(0, 10)),
734
+ (slice(2, 3), slice(0, 10)),
735
+ ],
736
+ ),
737
+ # Both constraints satisfied - all fit in one chunk
738
+ (
739
+ torch.tensor([10, 10, 10]),
740
+ torch.tensor([5, 5, 5]),
741
+ 100,
742
+ 10000, # 2500 elems, M*N = 15*30 = 450 < 2500
743
+ [(slice(0, 3), slice(0, 15))],
744
+ ),
745
+ # Workspace constraint triggers first
746
+ (
747
+ torch.tensor([50, 50, 50]),
748
+ torch.tensor([1, 1, 1]),
749
+ 50, # workspace only fits one at a time
750
+ 1000000, # logits budget is huge
751
+ [
752
+ (slice(0, 1), slice(0, 1)),
753
+ (slice(1, 2), slice(0, 1)),
754
+ (slice(2, 3), slice(0, 1)),
755
+ ],
756
+ ),
757
+ # Greedy filling: first two fit, third doesn't
758
+ # req0: M=5, N=10 -> 50 elems
759
+ # req0+1: M=10, N=20 -> 200 elems <= 250
760
+ # req0+1+2: M=15, N=30 -> 450 elems > 250
761
+ (
762
+ torch.tensor([10, 10, 10]),
763
+ torch.tensor([5, 5, 5]),
764
+ 100,
765
+ 1000, # 250 elems
766
+ [(slice(0, 2), slice(0, 10)), (slice(2, 3), slice(0, 5))],
767
+ ),
768
+ ],
769
+)
770
+def test_split_indexer_prefill_chunks(
771
+ seq_lens, query_lens, workspace_size, max_logits_bytes, expected
772
+):
773
+ out = split_indexer_prefill_chunks(
774
+ seq_lens,
775
+ query_lens,
776
+ workspace_size,
777
+ max_logits_bytes,
778
+ )
779
+ assert out == expected
780
+
781
+
782
+def test_split_indexer_prefill_chunks_single_request_overflow():
783
+ """Test that single request exceeding budget is sub-chunked on query dim."""
784
+ seq_lens = torch.tensor([1000, 50])
785
+ query_lens = torch.tensor([100, 5])
786
+
787
+ out = split_indexer_prefill_chunks(seq_lens, query_lens, 2000, 1000)
788
+ # max_logits_elems = 250, N=1000 -> max_q = 1 -> 100 query sub-chunks
789
+ expected = [(slice(0, 1), slice(i, i + 1)) for i in range(100)]
790
+ # req1: M=5, N=50 -> 250 elems fits budget
791
+ expected.append((slice(1, 2), slice(0, 5)))
792
+ assert out == expected
793
+
794
+
795
def test_triton_convert_returns_valid_counts():
796
"""Test that return_valid_counts correctly counts non-negative indices."""
797
device = torch.device("cuda")
vllm/envs.py
+7
@@ -55,6 +55,7 @@ if TYPE_CHECKING:
55
VLLM_CPU_INT4_W4A8: bool = True
56
VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
57
VLLM_XLA_CHECK_RECOMPILATION: bool = False
58
+ VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512
59
VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
60
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
61
VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
@@ -861,6 +862,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
862
),
863
# Enable SPMD mode for TPU backend.
864
"VLLM_XLA_USE_SPMD": lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))),
865
+ # Maximum size (in MB) for logits tensor in sparse MLA indexer prefill chunks.
866
+ # Bounds the [M, N] float32 logits tensor to prevent CUDA OOM.
867
+ # Default: 512 MB
868
+ "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB": lambda: int(
869
+ os.getenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512")
870
+ ),
871
# If set, the OpenAI API server will stay alive even after the underlying
872
# AsyncLLMEngine errors and stops serving requests
873
"VLLM_KEEP_ALIVE_ON_ENGINE_DEATH": lambda: bool(
vllm/model_executor/layers/sparse_attn_indexer.py
+19
-7
@@ -4,6 +4,7 @@
4
5
import torch
6
7
+import vllm.envs as envs
8
from vllm._aiter_ops import rocm_aiter_ops
9
from vllm.forward_context import get_forward_context
10
from vllm.logger import init_logger
@@ -51,6 +52,14 @@ def sparse_attn_indexer(
52
((total_seq_lens, head_dim), torch.float8_e4m3fn),
53
((total_seq_lens, 4), torch.uint8),
54
)
55
+
56
+ # Dummy allocation to simulate for peak logits tensor memory during inference.
57
+ # FP8 elements so elements == bytes
58
+ max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
59
+ _ = torch.empty(
60
+ max_logits_elems, dtype=torch.uint8, device=hidden_states.device
61
+ )
62
+
63
return sparse_attn_indexer_fake(
64
hidden_states,
65
k_cache_prefix,
@@ -101,13 +110,16 @@ def sparse_attn_indexer(
110
for chunk in prefill_metadata.chunks:
111
k_fp8 = k_fp8_full[: chunk.total_seq_lens]
112
k_scale = k_scale_full[: chunk.total_seq_lens]
104
- ops.cp_gather_indexer_k_quant_cache(
105
- kv_cache,
106
- k_fp8,
107
- k_scale,
108
- chunk.block_table,
109
- chunk.cu_seq_lens,
110
- )
113
+
114
+ if not chunk.skip_kv_gather:
115
+ ops.cp_gather_indexer_k_quant_cache(
116
+ kv_cache,
117
+ k_fp8,
118
+ k_scale,
119
+ chunk.block_table,
120
+ chunk.cu_seq_lens,
121
+ )
122
+
123
logits = fp8_mqa_logits(
124
q_fp8[chunk.token_start : chunk.token_end],
125
(k_fp8, k_scale.view(torch.float32).flatten()),
vllm/v1/attention/backends/mla/indexer.py
+89
-24
@@ -4,6 +4,7 @@ from dataclasses import dataclass
4
5
import torch
6
7
+import vllm.envs as envs
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
@@ -22,7 +23,6 @@ from vllm.v1.attention.backend import (
23
)
24
from vllm.v1.attention.backends.utils import (
25
split_decodes_and_prefills,
25
- split_prefill_chunks,
26
)
27
from vllm.v1.kv_cache_interface import AttentionSpec
28
from vllm.v1.worker.cp_utils import get_total_cp_world_size
@@ -30,6 +30,55 @@ from vllm.v1.worker.cp_utils import get_total_cp_world_size
30
logger = init_logger(__name__)
31
32
33
+def split_indexer_prefill_chunks(
34
+ seq_lens_cpu: torch.Tensor,
35
+ query_lens_cpu: torch.Tensor,
36
+ workspace_size: int,
37
+ max_logits_bytes: int,
38
+ request_offset: int = 0,
39
+) -> list[tuple[slice, slice]]:
40
+ """
41
+ Split prefill requests into chunks for the sparse indexer, respecting:
42
+ - N constraint: total_seq_lens <= workspace_size (existing O(N) workspace)
43
+ - Logits constraint: M * N * 4 <= max_logits_bytes
44
+
45
+ When a single request-level chunk still exceeds the logits budget,
46
+ sub-chunks on the query dimension (M) to bound peak memory.
47
+
48
+ Returns list of (req_slice, query_slice) tuples.
49
+ """
50
+ chunks: list[tuple[slice, slice]] = []
51
+ n = len(seq_lens_cpu)
52
+ max_logits_elems = max_logits_bytes // 4
53
+ end = 0
54
+
55
+ while end < n:
56
+ start, chunk_m, chunk_n = end, 0, 0
57
+
58
+ while end < n:
59
+ q, s = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
60
+ new_m, new_n = chunk_m + q, chunk_n + s
61
+ if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
62
+ chunk_m, chunk_n = new_m, new_n
63
+ end += 1
64
+ else:
65
+ break
66
+
67
+ # A single request can exceed the budget, requiring sub-chunking
68
+ # on the query dimension.
69
+ if end == start:
70
+ chunk_m, chunk_n = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
71
+ end += 1
72
+
73
+ req_slice = slice(start + request_offset, end + request_offset)
74
+ max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m
75
+ for q_off in range(0, chunk_m, max_q):
76
+ sub_m = min(max_q, chunk_m - q_off)
77
+ chunks.append((req_slice, slice(q_off, q_off + sub_m)))
78
+
79
+ return chunks
80
+
81
+
82
class DeepseekV32IndexerBackend(AttentionBackend):
83
@staticmethod
84
def get_name() -> str:
@@ -81,6 +130,7 @@ class DeepseekV32IndexerPrefillChunkMetadata:
130
token_start: int
131
token_end: int
132
num_reqs: int
133
+ skip_kv_gather: bool = False
134
135
136
@dataclass
@@ -271,43 +321,51 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
321
)
322
323
def build_one_prefill_chunk(
274
- self, reqs_start, reqs_end, query_start_loc_cpu, seq_lens_cpu, block_table
275
- ):
324
+ self,
325
+ req_slice: slice,
326
+ query_slice: slice,
327
+ query_start_loc_cpu,
328
+ seq_lens_cpu,
329
+ block_table,
330
+ skip_kv_gather: bool = False,
331
+ ) -> DeepseekV32IndexerPrefillChunkMetadata:
332
prefill_query_start_loc = (
277
- query_start_loc_cpu[reqs_start : reqs_end + 1]
278
- - query_start_loc_cpu[reqs_start]
333
+ query_start_loc_cpu[req_slice.start : req_slice.stop + 1]
334
+ - query_start_loc_cpu[req_slice.start]
335
)
336
cu_seqlen_ks, cu_seqlen_ke = kv_spans_from_batches(
281
- prefill_query_start_loc, seq_lens_cpu[reqs_start:reqs_end], self.device
337
+ prefill_query_start_loc, seq_lens_cpu[req_slice], self.device
338
+ )
339
+ token_start = query_start_loc_cpu[req_slice.start].item()
340
+ total_seq_lens = seq_lens_cpu[req_slice].sum()
341
+ num_reqs = req_slice.stop - req_slice.start
342
+ seq_idx = torch.arange(0, num_reqs, dtype=torch.int32)
343
+ token_to_seq = torch.repeat_interleave(seq_idx, seq_lens_cpu[req_slice]).to(
344
+ self.device
345
)
283
- token_start = query_start_loc_cpu[reqs_start].item()
284
- token_end = query_start_loc_cpu[reqs_end].item()
285
- total_seq_lens = seq_lens_cpu[reqs_start:reqs_end].sum()
286
- seq_idx = torch.arange(0, reqs_end - reqs_start, dtype=torch.int32)
287
- token_to_seq = torch.repeat_interleave(
288
- seq_idx, seq_lens_cpu[reqs_start:reqs_end]
289
- ).to(self.device)
346
assert total_seq_lens <= self.max_prefill_buffer_size
347
cu_seq_lens = (
348
torch.cat(
349
[
350
torch.zeros(1, dtype=torch.int32),
295
- seq_lens_cpu[reqs_start:reqs_end].cumsum(dim=0),
351
+ seq_lens_cpu[req_slice].cumsum(dim=0),
352
]
353
)
354
.to(torch.int32)
355
.to(self.device)
356
)
357
+
358
return DeepseekV32IndexerPrefillChunkMetadata(
302
- cu_seqlen_ks=cu_seqlen_ks,
303
- cu_seqlen_ke=cu_seqlen_ke,
359
+ cu_seqlen_ks=cu_seqlen_ks[query_slice],
360
+ cu_seqlen_ke=cu_seqlen_ke[query_slice],
361
cu_seq_lens=cu_seq_lens,
362
token_to_seq=token_to_seq,
363
total_seq_lens=total_seq_lens,
307
- block_table=block_table[reqs_start:reqs_end],
308
- token_start=token_start,
309
- token_end=token_end,
310
- num_reqs=reqs_end - reqs_start,
364
+ block_table=block_table[req_slice],
365
+ token_start=token_start + query_slice.start,
366
+ token_end=token_start + query_slice.stop,
367
+ num_reqs=num_reqs,
368
+ skip_kv_gather=skip_kv_gather,
369
)
370
371
def build(
@@ -333,20 +391,27 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
391
392
prefill_metadata = None
393
if num_prefills > 0:
336
- chunk_seq_ids = split_prefill_chunks(
394
+ prefill_query_lens_cpu = torch.diff(
395
+ query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1]
396
+ )
397
+ max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
398
+ chunk_specs = split_indexer_prefill_chunks(
399
common_attn_metadata.seq_lens_cpu[num_decodes:],
400
+ prefill_query_lens_cpu,
401
self.max_prefill_buffer_size,
402
+ max_logits_bytes,
403
request_offset=num_decodes,
404
)
405
chunks = [
406
self.build_one_prefill_chunk(
343
- reqs_start,
344
- reqs_end,
407
+ req_slice,
408
+ query_slice,
409
query_start_loc_cpu,
410
common_attn_metadata.seq_lens_cpu,
411
common_attn_metadata.block_table_tensor,
412
+ skip_kv_gather=query_slice.start > 0,
413
)
349
- for reqs_start, reqs_end in chunk_seq_ids
414
+ for req_slice, query_slice in chunk_specs
415
]
416
prefill_metadata = DeepseekV32IndexerPrefillMetadata(
417
chunks=chunks,