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, 01:00 UTC
[MLA] Optimize mla indexer prepare uniform decode for MTP > 1 (#39458)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Apr 16, 2026 at 16:27 UTC
bf9a5ddb24af910f53e7d20305045010d7471072
1 file changed
+100
-43
vllm/v1/attention/backends/mla/indexer.py
+100
-43
@@ -8,6 +8,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
11
+from vllm.triton_utils import tl, triton
12
from vllm.utils.deep_gemm import (
13
get_paged_mqa_logits_metadata,
14
has_deep_gemm,
@@ -30,6 +31,40 @@ from vllm.v1.worker.cp_utils import get_total_cp_world_size
31
logger = init_logger(__name__)
32
33
34
+@triton.jit
35
+def _prepare_uniform_decode_kernel(
36
+ seq_lens_ptr,
37
+ decode_seq_lens_ptr,
38
+ block_table_ptr,
39
+ block_table_stride,
40
+ expanded_block_table_ptr,
41
+ expanded_bt_stride,
42
+ decode_lens_ptr,
43
+ max_decode_len,
44
+ BLOCK_SIZE: tl.constexpr,
45
+):
46
+ idx = tl.program_id(0)
47
+ req_id = idx // max_decode_len
48
+ local_idx = idx % max_decode_len
49
+
50
+ # Compute number of KVs attended to by this token.
51
+ seq_len = tl.load(seq_lens_ptr + req_id)
52
+ per_token_seq_len = seq_len - max_decode_len + local_idx + 1
53
+ tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
54
+
55
+ # Copy block table row.
56
+ src = block_table_ptr + req_id * block_table_stride
57
+ dst = expanded_block_table_ptr + idx * expanded_bt_stride
58
+ for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
59
+ off = i + tl.arange(0, BLOCK_SIZE)
60
+ mask = off < expanded_bt_stride
61
+ src_block = tl.load(src + off, mask=mask)
62
+ tl.store(dst + off, src_block, mask=mask)
63
+
64
+ # All reqs now have decode_len = 1.
65
+ tl.store(decode_lens_ptr + idx, 1)
66
+
67
+
68
def split_indexer_prefill_chunks(
69
seq_lens_cpu: torch.Tensor,
70
query_lens_cpu: torch.Tensor,
@@ -405,52 +440,75 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
440
Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).
441
seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, next_n) for native MTP.
442
"""
443
+ min_decode_len = int(decode_lens_cpu.min().item())
444
if not use_native and max_decode_len > 1:
445
assert self.decode_seq_lens_buffer.dim() == 1
410
- # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
411
- # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
412
- # The context lengths are therefore
413
- # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
414
-
415
- # 3 + 1 + 4 + 0 = 8
416
- actual_expanded = int(decode_lens_cpu.sum().item())
417
-
418
- # Fuse expanded_base and expanded_starts into a single repeat_interleave:
419
- # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
420
- # where context_start[b] = seq_lens[b] - decode_lens[b].
421
- # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
422
- # expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4]
423
- # result = [8, 9, 10, 7, 9, 10, 11, 12]
424
- expanded_offsets = torch.repeat_interleave(
425
- seq_lens - decode_lens - query_start_loc,
426
- decode_lens,
427
- output_size=actual_expanded,
428
- )
446
+ if min_decode_len == max_decode_len:
447
+ # Uniform decode lengths.
448
+ num_decode_tokens = num_decodes * max_decode_len
449
+ _prepare_uniform_decode_kernel[(num_decode_tokens,)](
450
+ seq_lens,
451
+ self.decode_seq_lens_buffer,
452
+ block_table,
453
+ block_table.stride(0),
454
+ self.expanded_block_table_buffer,
455
+ self.expanded_block_table_buffer.stride(0),
456
+ self.decode_lens_buffer,
457
+ max_decode_len,
458
+ BLOCK_SIZE=1024,
459
+ )
460
+ self.decode_seq_lens_buffer[num_decode_tokens:] = 0
461
+ seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
462
+ block_table = self.expanded_block_table_buffer[:num_decode_tokens]
463
+ decode_lens = self.decode_lens_buffer[:num_decode_tokens]
464
+ return seq_lens, block_table, decode_lens, num_decode_tokens, False
465
+ else:
466
+ # Variable decode lengths.
467
+ # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
468
+ # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
469
+ # The context lengths are therefore
470
+ # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
471
+
472
+ # 3 + 1 + 4 + 0 = 8
473
+ actual_expanded = int(decode_lens_cpu.sum().item())
474
+
475
+ # Fuse expanded_base and expanded_starts into a single
476
+ # repeat_interleave:
477
+ # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
478
+ # where context_start[b] = seq_lens[b] - decode_lens[b].
479
+ # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
480
+ # expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4]
481
+ # result = [8, 9, 10, 7, 9, 10, 11, 12]
482
+ expanded_offsets = torch.repeat_interleave(
483
+ seq_lens - decode_lens - query_start_loc,
484
+ decode_lens,
485
+ output_size=actual_expanded,
486
+ )
487
430
- # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
431
- self.decode_seq_lens_buffer[:actual_expanded] = (
432
- expanded_offsets + self.arange_buffer[:actual_expanded] + 1
433
- )
434
- self.decode_seq_lens_buffer[actual_expanded:] = 0
435
- seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
436
-
437
- # Give each of the flattened entries the same block table row as the
438
- # original request.
439
- self.expanded_block_table_buffer[:actual_expanded] = (
440
- torch.repeat_interleave(
441
- block_table, decode_lens, dim=0, output_size=actual_expanded
488
+ # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
489
+ self.decode_seq_lens_buffer[:actual_expanded] = (
490
+ expanded_offsets + self.arange_buffer[:actual_expanded] + 1
491
)
443
- )
444
- if actual_expanded < num_decode_tokens:
445
- self.expanded_block_table_buffer[
446
- actual_expanded:num_decode_tokens, 0
447
- ] = 0
448
- block_table = self.expanded_block_table_buffer[:num_decode_tokens]
449
-
450
- # All reqs now have decode_len=1
451
- self.decode_lens_buffer[:num_decode_tokens] = 1
452
- decode_lens = self.decode_lens_buffer[:num_decode_tokens]
453
- return seq_lens, block_table, decode_lens, num_decode_tokens, False
492
+ self.decode_seq_lens_buffer[actual_expanded:] = 0
493
+ seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
494
+
495
+ # Give each of the flattened entries the same block table row as the
496
+ # original request.
497
+ self.expanded_block_table_buffer[:actual_expanded] = (
498
+ torch.repeat_interleave(
499
+ block_table, decode_lens, dim=0, output_size=actual_expanded
500
+ )
501
+ )
502
+ if actual_expanded < num_decode_tokens:
503
+ self.expanded_block_table_buffer[
504
+ actual_expanded:num_decode_tokens, 0
505
+ ] = 0
506
+ block_table = self.expanded_block_table_buffer[:num_decode_tokens]
507
+
508
+ # All reqs now have decode_len=1
509
+ self.decode_lens_buffer[:num_decode_tokens] = 1
510
+ decode_lens = self.decode_lens_buffer[:num_decode_tokens]
511
+ return seq_lens, block_table, decode_lens, num_decode_tokens, False
512
else:
513
# Native path: plain decode (next_n==1) or spec decode
514
# with 2D per-token context lengths (next_n > 1).
@@ -459,7 +517,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
517
# decode_len < next_n due to padding or short prefills), the simple
518
# reshape in sparse_attn_indexer won't work. Use pack_seq_triton
519
# (requires_padding) instead.
462
- min_decode_len = int(decode_lens_cpu.min().item())
520
requires_padding = min_decode_len != max_decode_len
521
if use_native and next_n > 1:
522
assert self.decode_seq_lens_buffer.dim() == 2