@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2e98406048

Fork 0

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, 09:45 UTC

[Refactor] Improve indexer decode path metadata preparation (#38865)

Yongye Zhu committed Apr 8, 2026 at 23:49 UTC 2e98406048798c45eb623360de2b46825ef74848

4 files changed +162 -102

csrc/sampler.cu

+24 -9

@@ -564,8 +564,9 @@ template <int kNumThreadsPerBlock, bool useRadixSort,

564 bool multipleBlocksPerRow = false, bool mergeBlocks = false>

565 static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(

566 const float* logits, const int* seqLens, int* outIndices, int stride0,

567 - int stride1, const int topK, int next_n, float* outLogits = nullptr,

568 - const int numBlocksToMerge = 0, const int* indices = nullptr) {

567 + int stride1, const int topK, int next_n, int seqLensIs2D = 0,

568 + float* outLogits = nullptr, const int numBlocksToMerge = 0,

569 + const int* indices = nullptr) {

570 // The number of bins in the histogram.

571 static constexpr int kNumBins = 2048;

572

@@ -574,8 +575,16 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(

575

576 // The range of logits within the row.

577 int rowStart = 0;

577 - int seq_len = seqLens[rowIdx / next_n];

578 - int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1);

578 + int batch_idx = rowIdx / next_n;

579 + int next_n_idx = rowIdx % next_n;

580 + // seqLensIs2D=0: 1D seqLens — all rows in a batch share the same seq_len;

581 + // kernel computes per-row effective length via offset.

582 + // seqLensIs2D=1: 2D seqLens — each logit row has its own pre-computed

583 + // effective length (flat index rowIdx = b*next_n + j maps

584 + // directly to seqLens[b, j] in C-contiguous layout).

585 + int seq_len = seqLensIs2D ? seqLens[rowIdx] : seqLens[batch_idx];

586 + int rowEnd =

587 + seqLensIs2D ? max(0, seq_len) : max(0, seq_len - next_n + next_n_idx + 1);

588

589 // Local pointers to this block

590 if constexpr (!multipleBlocksPerRow && !mergeBlocks) {

@@ -653,6 +662,11 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,

662 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

663 const auto numColumns = logits.size(1);

664

665 + // True if seqLens is 2D (B, next_n): each logit row has its own pre-computed

666 + // effective seq_len. False if seqLens is 1D (B,): all rows in a batch share

667 + // the same seq_len and the kernel computes the per-row offset itself.

668 + int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0;

669 +

670 if (numColumns < kSortingAlgorithmThreshold) {

671 // Use insertion sort

672 vllm::topKPerRowDecode<kNumThreadsPerBlock, false>

@@ -660,7 +674,7 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,

674 logits.data_ptr<float>(), seqLens.data_ptr<int>(),

675 indices.data_ptr<int>(), static_cast<int>(stride0),

676 static_cast<int>(stride1), static_cast<int>(topK),

663 - static_cast<int>(next_n));

677 + static_cast<int>(next_n), seqLensIs2D);

678 } else if (numColumns < kSplitWorkThreshold) {

679 // From this threshold, use radix sort instead

680 vllm::topKPerRowDecode<kNumThreadsPerBlock, true>

@@ -668,7 +682,7 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,

682 logits.data_ptr<float>(), seqLens.data_ptr<int>(),

683 indices.data_ptr<int>(), static_cast<int>(stride0),

684 static_cast<int>(stride1), static_cast<int>(topK),

671 - static_cast<int>(next_n));

685 + static_cast<int>(next_n), seqLensIs2D);

686 } else {

687 // Long sequences are run in two steps

688 constexpr auto multipleBlocksPerRowConfig = 10;

@@ -686,15 +700,16 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,

700 logits.data_ptr<float>(), seqLens.data_ptr<int>(),

701 outIndicesAux.data_ptr<int>(), static_cast<int>(stride0),

702 static_cast<int>(stride1), static_cast<int>(topK),

689 - static_cast<int>(next_n), outLogitsAux.data_ptr<float>());

703 + static_cast<int>(next_n), seqLensIs2D,

704 + outLogitsAux.data_ptr<float>());

705

706 constexpr int kNumThreadsPerBlockMerge = 1024;

707 vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>

708 <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(

709 outLogitsAux.data_ptr<float>(), seqLens.data_ptr<int>(),

710 indices.data_ptr<int>(), multipleBlocksPerRowConfig * topK, 1,

696 - static_cast<int>(topK), static_cast<int>(next_n), nullptr,

697 - multipleBlocksPerRowConfig, outIndicesAux.data_ptr<int>());

711 + static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,

712 + nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr<int>());

713 }

714 }

715

csrc/topk.cu

+4 -2

@@ -21,13 +21,15 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,

21 TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");

22 TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");

23 TORCH_CHECK(logits.dim() == 2, "logits must be 2D");

24 - TORCH_CHECK(lengths.dim() == 1, "lengths must be 1D");

24 + TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,

25 + "lengths must be 1D or 2D");

26 + TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");

27 TORCH_CHECK(output.dim() == 2, "output must be 2D");

29 const int64_t num_rows = logits.size(0);

30 const int64_t stride = logits.size(1);

30 - TORCH_CHECK(lengths.size(0) == num_rows, "lengths size mismatch");

32 + TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");

33 TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,

34 "output size mismatch");

35 namespace P = vllm::persistent;

vllm/model_executor/layers/sparse_attn_indexer.py

+7 -16

@@ -183,13 +183,15 @@ def sparse_attn_indexer(

183 # TODO: move and optimize below logic with triton kernels

184 batch_size = padded_q_fp8_decode_tokens.shape[0]

185 next_n = padded_q_fp8_decode_tokens.shape[1]

186 - assert batch_size == decode_metadata.seq_lens.shape[0]

186 num_padded_tokens = batch_size * next_n

187 + seq_lens = decode_metadata.seq_lens[:batch_size]

188 + # seq_lens is (B, next_n) for native spec decode, (B,) otherwise.

189 + # fp8_paged_mqa_logits and all topk kernels accept both shapes.

190 logits = fp8_paged_mqa_logits(

191 padded_q_fp8_decode_tokens,

192 kv_cache,

193 weights[:num_padded_tokens],

192 - decode_metadata.seq_lens,

194 + seq_lens,

195 decode_metadata.block_table,

196 decode_metadata.schedule_metadata,

197 max_model_len=max_model_len,

@@ -198,17 +200,6 @@ def sparse_attn_indexer(

200 num_rows = logits.shape[0]

201 topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]

202

201 - if next_n == 1:

202 - lengths = decode_metadata.seq_lens

203 - else:

204 - # (bs,) -> (bs, 1) + (next_n,) -> (bs, next_n) -> (bs * next_n,)

205 - lengths = (

206 - decode_metadata.seq_lens.unsqueeze(1)

207 - - next_n

208 - + 1

209 - + decode_metadata.offsets

210 - ).flatten()

211 -

203 if current_platform.is_cuda():

204 workspace_manager = current_workspace_manager()

205 (topk_workspace,) = workspace_manager.get_simultaneous(

@@ -216,7 +207,7 @@ def sparse_attn_indexer(

207 )

208 torch.ops._C.persistent_topk(

209 logits,

219 - lengths,

210 + decode_metadata.seq_lens,

211 topk_indices,

212 topk_workspace,

213 topk_tokens,

@@ -227,7 +218,7 @@ def sparse_attn_indexer(

218 ops.top_k_per_row_decode(

219 logits,

220 next_n,

230 - decode_metadata.seq_lens,

221 + seq_lens,

222 topk_indices,

223 num_rows,

224 logits.stride(0),

@@ -238,7 +229,7 @@ def sparse_attn_indexer(

229 torch.ops._C.top_k_per_row_decode(

230 logits,

231 next_n,

241 - decode_metadata.seq_lens,

232 + seq_lens,

233 topk_indices,

234 num_rows,

235 logits.stride(0),

vllm/v1/attention/backends/mla/indexer.py

+127 -75

@@ -141,11 +141,14 @@ class DeepseekV32IndexerPrefillMetadata:

141 @dataclass

142 class DeepSeekV32IndexerDecodeMetadata:

143 block_table: torch.Tensor

144 + # seq_lens: per-token effective context lengths.

145 + # - flatten path / plain decode: 1D (batch_size,)

146 + # - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1

147 + # Both fp8_paged_mqa_logits and the topk kernels accept both shapes.

148 seq_lens: torch.Tensor

149 decode_lens: torch.Tensor

150 requires_padding: bool

151 schedule_metadata: torch.Tensor

148 - offsets: torch.Tensor | None # Precomputed offsets for speculative decoding

152

153

154 @dataclass

@@ -283,21 +286,31 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

286 sm_count = num_compute_units(self.device.index)

287 self.num_sms = sm_count

288

286 - self.decode_lens_buffer = torch.empty(

287 - (scheduler_config.max_num_batched_tokens,),

288 - dtype=torch.int32,

289 - device=self.device,

290 - )

289 self.offsets_buffer = torch.arange(

290 next_n, device=self.device, dtype=torch.int32

291 )

294 - self.arange_buffer = torch.arange(

295 - scheduler_config.max_num_seqs * next_n,

292 + self.decode_lens_buffer = torch.zeros(

293 + (scheduler_config.max_num_batched_tokens,),

294 dtype=torch.int32,

295 device=self.device,

296 )

299 - self.expanded_seq_lens_buffer = torch.zeros(

300 - (scheduler_config.max_num_batched_tokens,),

297 + if not self.use_flattening and next_n > 1:

298 + # Native MTP: 2D buffer for per-token seq_lens.

299 + # Flattening path is never used, so no expanded_seq_lens_buffer.

300 + self.decode_seq_lens_buffer = torch.zeros(

301 + (scheduler_config.max_num_seqs, next_n),

302 + dtype=torch.int32,

303 + device=self.device,

304 + )

305 + else:

306 + # Flattening or no MTP: 1D buffer for expanded per-token seq_lens.

307 + self.decode_seq_lens_buffer = torch.zeros(

308 + (scheduler_config.max_num_batched_tokens,),

309 + dtype=torch.int32,

310 + device=self.device,

311 + )

312 + self.arange_buffer = torch.arange(

313 + scheduler_config.max_num_seqs * next_n,

314 dtype=torch.int32,

315 device=self.device,

316 )

@@ -367,6 +380,96 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

380 skip_kv_gather=skip_kv_gather,

381 )

382

383 + def _prepare_decode_tensors(

384 + self,

385 + seq_lens: torch.Tensor,

386 + block_table: torch.Tensor,

387 + decode_lens: torch.Tensor,

388 + decode_lens_cpu: torch.Tensor,

389 + query_start_loc: torch.Tensor,

390 + num_decodes: int,

391 + num_decode_tokens: int,

392 + use_native: bool,

393 + next_n: int,

394 + max_decode_len: int,

395 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:

396 + """Expand seq_lens/block_table/decode_lens for the decode kernels.

397 +

398 + Flatten path (not use_native, max_decode_len > 1):

399 + Each multi-token decode request is expanded into individual

400 + single-token entries so the kernel always sees next_n=1.

401 +

402 + Native path (use_native or max_decode_len == 1):

403 + Plain decode or spec-decode with 2D per-token context lengths.

404 +

405 + Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).

406 + seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, next_n) for native MTP.

407 + """

408 + if not use_native and max_decode_len > 1:

409 + 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 + )

429 +

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

442 + )

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

454 + else:

455 + # Native path: plain decode (next_n==1) or spec decode

456 + # with 2D per-token context lengths (next_n > 1).

457 + #

458 + # When decode_lens are not truly uniform (e.g. some requests have

459 + # decode_len < next_n due to padding or short prefills), the simple

460 + # reshape in sparse_attn_indexer won't work. Use pack_seq_triton

461 + # (requires_padding) instead.

462 + min_decode_len = int(decode_lens_cpu.min().item())

463 + requires_padding = min_decode_len != max_decode_len

464 + if use_native and next_n > 1:

465 + assert self.decode_seq_lens_buffer.dim() == 2

466 + # (B, next_n): token j attends to L - next_n + j + 1 KV tokens

467 + self.decode_seq_lens_buffer[:num_decodes] = (

468 + seq_lens.unsqueeze(1) - next_n + 1 + self.offsets_buffer

469 + )

470 + seq_lens = self.decode_seq_lens_buffer[:num_decodes]

471 + return seq_lens, block_table, decode_lens, num_decodes, requires_padding

472 +

473 def build(

474 self,

475 common_prefix_len: int,

@@ -434,68 +537,20 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

537 next_n = 1 + self.num_speculative_tokens

538 use_native = not self.use_flattening and max_decode_len == next_n

539

437 - if use_native and next_n > 1:

438 - offsets = self.offsets_buffer

439 - elif max_decode_len > 1:

440 - # Flatten multi-token decode requests into single-token

441 - # batch entries, expanding seq_lens and block tables so

442 - # the kernel always sees next_n=1.

443 -

444 - # Also handles the edge case where use_flattening=False

445 - # but max_decode_len != next_n (e.g. a batch containing some

446 - # short prefills (q_len < next_n) and no true decodes).

447 -

448 - # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is

449 - # padding) and decode_lens [3, 1, 4, 0] in the below example comments.

450 - # The context lengths are therefore

451 - # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].

452 -

453 - # 3 + 1 + 4 + 0 = 8

454 - actual_expanded = int(decode_lens_cpu.sum().item())

455 -

456 - # [7, 6, 8, 0] -> [7, 7, 7, 6, 8, 8, 8, 8]

457 - expanded_base = torch.repeat_interleave(

458 - seq_lens - decode_lens, decode_lens, output_size=actual_expanded

459 - )

460 -

461 - # [0, 3, 4, 8] -> [0, 0, 0, 3, 4, 4, 4, 4]

462 - expanded_starts = torch.repeat_interleave(

463 - common_attn_metadata.query_start_loc[:num_decodes],

464 - decode_lens,

465 - output_size=actual_expanded,

540 + seq_lens, block_table, decode_lens, batch_size, requires_padding = (

541 + self._prepare_decode_tensors(

542 + seq_lens=seq_lens,

543 + block_table=block_table,

544 + decode_lens=decode_lens,

545 + decode_lens_cpu=decode_lens_cpu,

546 + query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],

547 + num_decodes=num_decodes,

548 + num_decode_tokens=num_decode_tokens,

549 + use_native=use_native,

550 + next_n=next_n,

551 + max_decode_len=max_decode_len,

552 )

467 -

468 - # [0, 1, 2, 0, 0, 1, 2, 3]

469 - positions_within = (

470 - self.arange_buffer[:actual_expanded] - expanded_starts

471 - )

472 -

473 - # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space

474 - self.expanded_seq_lens_buffer[:actual_expanded] = (

475 - expanded_base + positions_within + 1

476 - )

477 - self.expanded_seq_lens_buffer[actual_expanded:] = 0

478 - seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]

479 -

480 - # Give each of the flattened entries the same block table row as the

481 - # original request.

482 - self.expanded_block_table_buffer[:actual_expanded] = (

483 - torch.repeat_interleave(

484 - block_table, decode_lens, dim=0, output_size=actual_expanded

485 - )

486 - )

487 - if actual_expanded < num_decode_tokens:

488 - self.expanded_block_table_buffer[

489 - actual_expanded:num_decode_tokens, 0

490 - ] = 0

491 - block_table = self.expanded_block_table_buffer[:num_decode_tokens]

492 -

493 - # All reqs now have decode_len=1

494 - self.decode_lens_buffer[:num_decode_tokens] = 1

495 - decode_lens = self.decode_lens_buffer[:num_decode_tokens]

496 - offsets = None

497 - else:

498 - offsets = None

553 + )

554

555 # DeepGEMM is required for the paged MQA logits on CUDA devices

556 if current_platform.is_cuda() and has_deep_gemm():

@@ -509,9 +564,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

564 block_table=block_table,

565 seq_lens=seq_lens,

566 decode_lens=decode_lens,

512 - requires_padding=False,

567 + requires_padding=requires_padding,

568 schedule_metadata=self.scheduler_metadata_buffer,

514 - offsets=offsets,

569 )

570

571 attn_metadata = DeepseekV32IndexerMetadata(

@@ -531,6 +585,4 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

585 decode=decode_metadata,

586 )

587

534 - # if get_tensor_model_parallel_rank() == 0:

535 - # logger.info(f"attn_metadata: {attn_metadata}")

588 return attn_metadata