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

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 15, 21:00 UTC

[BugFix] MLA + V1, illegal memory access and accuracy issues (#14253)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Mar 5, 2025 at 20:10 UTC f6bb18fd9a19e5e4fb1991339638fc666d06b27a

7 files changed +334 -161

tests/v1/worker/test_gpu_input_batch.py

+89 -1

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 +import inspect

4 from typing import Optional

6 import numpy as np

@@ -9,7 +10,8 @@ import torch

10 from vllm.sampling_params import SamplingParams

11 from vllm.utils import is_pin_memory_available, make_tensor_with_pad

12 from vllm.v1.sample.metadata import SamplingMetadata

12 -from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

13 +from vllm.v1.worker.gpu_input_batch import (BlockTable, CachedRequestState,

14 + InputBatch)

16 VOCAB_SIZE = 1024

17 NUM_OUTPUT_TOKENS = 20

@@ -20,6 +22,34 @@ CUDA_DEVICES = [

22 MAX_NUM_PROMPT_TOKENS = 64

25 +def _compare_objs(obj1, obj2):

26 + attrs = inspect.getmembers(obj1, lambda a: not (inspect.isroutine(a)))

27 + attr_names = set([

28 + a[0] for a in attrs

29 + if not (a[0].startswith('__') and a[0].endswith('__'))

30 + ])

31 + for attr_name in attr_names:

32 + a = getattr(obj1, attr_name)

33 + b = getattr(obj2, attr_name)

34 +

35 + is_same = False

36 + if isinstance(a, torch.Tensor):

37 + if (a.numel() == 0 or b.numel() == 0):

38 + is_same = (a.numel() == 0 and b.numel() == 0)

39 + elif torch.allclose(a, b):

40 + is_same = True

41 + elif isinstance(a, np.ndarray):

42 + if np.allclose(a, b):

43 + is_same = True

44 + elif isinstance(a, (BlockTable, SamplingMetadata)):

45 + _compare_objs(a, b)

46 + is_same = True # if we make it here must be same

47 + elif a == b:

48 + is_same = True

49 + assert is_same, f"Attribute {attr_name} is different"\

50 + f" in {obj1} and {obj2}: {a} != {b}"

51 +

52 +

53 def _remove_requests(

54 input_batch: InputBatch, batch_size: int,

55 reqs: list[CachedRequestState]) -> tuple[set[str], list[int]]:

@@ -254,3 +284,61 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

284 assert torch.allclose(

285 expected_sampling_metadata.allowed_token_ids_mask,

286 sampling_metadata.allowed_token_ids_mask)

287 +

288 +

289 +@pytest.mark.parametrize("device", CUDA_DEVICES)

290 +@pytest.mark.parametrize("batch_size", [32])

291 +@pytest.mark.parametrize("swap_list", [((0, 1), )])

292 +def test_swap_states_in_input_batch(device: str, batch_size: int,

293 + swap_list: list):

294 + """

295 + Tests the logic for managing sampling metadata in the InputBatch.

296 +

297 + This test involves adding a set of requests to the InputBatch,

298 + followed by removing a subset of them. Afterward, the batch is compacted,

299 + and the `make_sampling_metadata` method is invoked on the batch. The

300 + output of `make_sampling_metadata` is then compared against the expected

301 + results to ensure correctness.

302 + """

303 + input_batch: InputBatch = InputBatch(

304 + max_num_reqs=batch_size,

305 + max_model_len=1024,

306 + max_num_blocks_per_req=10,

307 + device=torch.device(device),

308 + pin_memory=is_pin_memory_available(),

309 + vocab_size=1024,

310 + )

311 + ref_input_batch: InputBatch = InputBatch(

312 + max_num_reqs=batch_size,

313 + max_model_len=1024,

314 + max_num_blocks_per_req=10,

315 + device=torch.device(device),

316 + pin_memory=is_pin_memory_available(),

317 + vocab_size=1024,

318 + )

319 +

320 + reqs: list[CachedRequestState] = []

321 + req_id_reqs = {}

322 + req_id_output_token_ids = {}

323 + # Add requests

324 + for req_index in range(batch_size):

325 + req: CachedRequestState = _construct_cached_request_state(req_index)

326 + input_batch.add_request(req, req_index)

327 + reqs.append(req)

328 + req_id_reqs[req.req_id] = req

329 + req_id_output_token_ids[req.req_id] = req.output_token_ids

330 +

331 + reordered_reqs = reqs.copy()

332 + for swap_pair in swap_list:

333 + reordered_reqs[swap_pair[0]], reordered_reqs[swap_pair[1]] = \

334 + reordered_reqs[swap_pair[1]], reordered_reqs[swap_pair[0]]

335 + input_batch.swap_states(swap_pair[0], swap_pair[1])

336 +

337 + for req_index in range(batch_size):

338 + req = reordered_reqs[req_index]

339 + ref_input_batch.add_request(req, req_index)

340 +

341 + input_batch.refresh_sampling_metadata()

342 + ref_input_batch.refresh_sampling_metadata()

343 +

344 + _compare_objs(input_batch, ref_input_batch)

vllm/v1/attention/backends/flash_attn.py

+2 -2

@@ -100,8 +100,8 @@ class FlashAttentionMetadataBuilder:

100 self.runner = runner

101

102 def reorder_batch(self, input_batch: "InputBatch",

103 - scheduler_output: "SchedulerOutput"):

104 - pass

103 + scheduler_output: "SchedulerOutput") -> bool:

104 + return False

105

106 def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,

107 common_prefix_len: int):

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

+178 -127

@@ -275,17 +275,47 @@ class MLACommonBackend(AttentionBackend):

275

276

277 @dataclass

278 -class MLACommonMetadata:

279 - """Metadata for MLACommon.

278 +class MLACommonPrefillMetadata:

279 + """ Prefill Specific Metadata """

280 +

281 + @dataclass

282 + class ChunkedContextMetadata:

283 + # New for MLA (compared to FlashAttention)

284 + # For handling chunked prefill

285 + cu_seq_lens: torch.Tensor

286 + starts: torch.Tensor

287 + seq_tot: list[int]

288 + max_seq_lens: list[int]

289 + workspace: torch.Tensor

290

281 - NOTE: Please read the comment at the top of the file before trying to

282 - understand this class

283 - """

284 - # New for MLA (compared to FlashAttention)

291 # Input positions for rotrary embeddings since for MLA the rotary

292 # position embeddings are applied inside the attention backend

293 input_positions: torch.Tensor

294 + block_table: torch.Tensor

295 + query_start_loc: torch.Tensor

296 + max_query_len: int

297 + chunked_context: Optional[ChunkedContextMetadata] = None

298 +

299

300 +@dataclass

301 +class MLACommonDecodeMetadata:

302 + # Input positions for rotrary embeddings since for MLA the rotary

303 + # position embeddings are applied inside the attention backend

304 + input_positions: torch.Tensor

305 + block_table: torch.Tensor

306 + seq_lens: torch.Tensor

307 +

308 +

309 +D = TypeVar("D", bound=MLACommonDecodeMetadata)

310 +

311 +

312 +@dataclass

313 +class MLACommonMetadata(Generic[D]):

314 + """Metadata for MLACommon.

315 +

316 + NOTE: Please read the comment at the top of the file before trying to

317 + understand this class

318 + """

319 # NOTE(sang): Definition of context_len, query_len, and seq_len.

320 # |---------- N-1 iteration --------|

321 # |---------------- N iteration ---------------------|

@@ -295,30 +325,23 @@ class MLACommonMetadata:

325 # |-- query_len ---|

326

327 num_actual_tokens: int # Number of tokens excluding padding.

298 - max_query_len: int

328 query_start_loc: torch.Tensor

300 - max_seq_len: int

301 - seq_lens: torch.Tensor

302 - block_table: torch.Tensor

329 slot_mapping: torch.Tensor

330

331 + # New for MLA (compared to FlashAttention)

332 + # For handling prefill decode split

333 + num_decodes: int

334 + num_decode_tokens: int

335 + num_prefills: int

336 +

337 # For logging.

338 num_input_tokens: int = 0 # Number of tokens including padding.

339

340 # The dimension of the attention heads

341 head_dim: Optional[int] = None

342

311 - # New for MLA (compared to FlashAttention)

312 - # For chunked prefill

313 - num_decodes: Optional[int] = None

314 - num_decode_tokens: Optional[int] = None

315 - num_prefills: Optional[int] = None

316 - has_context: bool = False

317 - context_chunk_cu_seq_lens: Optional[torch.Tensor] = None

318 - context_chunk_starts: Optional[torch.Tensor] = None

319 - context_chunk_seq_tot: Optional[list[int]] = None

320 - context_chunk_max_seq_lens: Optional[list[int]] = None

321 - chunked_prefill_workspace: Optional[torch.Tensor] = None

343 + decode: Optional[D] = None

344 + prefill: Optional[MLACommonPrefillMetadata] = None

345

346 def __post_init__(self):

347 supported_head_sizes = MLACommonBackend.get_supported_head_sizes()

@@ -329,10 +352,10 @@ class MLACommonMetadata:

352 f"received {self.head_dim}.")

353

354

332 -T = TypeVar("T", bound=MLACommonMetadata)

355 +M = TypeVar("M", bound=MLACommonMetadata)

356

357

335 -class MLACommonMetadataBuilder(Generic[T]):

358 +class MLACommonMetadataBuilder(Generic[M]):

359 """

360 NOTE: Please read the comment at the top of the file before trying to

361 understand this class

@@ -340,8 +363,9 @@ class MLACommonMetadataBuilder(Generic[T]):

363

364 def __init__(self,

365 runner: "GPUModelRunner",

343 - cls: Optional[type[T]] = None):

344 - self.cls = cls if cls is not None else MLACommonMetadata

366 + metadata_cls: Optional[type[M]] = None):

367 + self.metadata_cls = metadata_cls \

368 + if metadata_cls is not None else MLACommonMetadata

369 self.runner = runner

370 scheduler_config = runner.scheduler_config

371 model_config = runner.model_config

@@ -375,7 +399,7 @@ class MLACommonMetadataBuilder(Generic[T]):

399 self.page_size = self.runner.block_size

400

401 def reorder_batch(self, input_batch: "InputBatch",

378 - scheduler_output: "SchedulerOutput"):

402 + scheduler_output: "SchedulerOutput") -> bool:

403 # We now want to reorder the batch so that the "decode" requests are and

404 # the front and the "prefill" requests are at the using the least amount

405 # swaps possible. (NOTE for now we loosely use "decode" to mean requests

@@ -413,6 +437,7 @@ class MLACommonMetadataBuilder(Generic[T]):

437 num_decodes = len(decodes)

438 num_prefills = len(prefills)

439 first_prefill = 0

440 + modified_batch = False

441

442 for i in range(1, min(num_decodes, num_prefills) + 1):

443 # If the decode is at the "back" of the batch, i, we can swap it

@@ -421,6 +446,7 @@ class MLACommonMetadataBuilder(Generic[T]):

446 input_batch.swap_states(prefills[first_prefill],

447 decodes[num_decodes - i])

448 first_prefill += 1

449 + modified_batch = True

450 else:

451 break

452

@@ -432,10 +458,21 @@ class MLACommonMetadataBuilder(Generic[T]):

458 self._num_decode_tokens = num_decode_tokens

459 self._num_prefill_tokens = num_prefill_tokens

460

461 + return modified_batch

462 +

463 + def _build_decode(self, input_positions: torch.Tensor,

464 + block_table: torch.Tensor, seq_lens: torch.Tensor):

465 + return MLACommonDecodeMetadata(

466 + input_positions=input_positions,

467 + block_table=block_table,

468 + seq_lens=seq_lens,

469 + )

470 +

471 def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,

436 - common_prefix_len: int) -> T:

472 + common_prefix_len: int) -> M:

473 + assert self._num_decodes + self._num_prefills == num_reqs

474 +

475 device = self.runner.device

438 - max_seq_len = self.runner.seq_lens_np[:num_reqs].max()

476 query_start_loc = self.runner.query_start_loc_cpu[:num_reqs + 1].to(

477 device, non_blocking=True)

478 seq_lens = self.runner.seq_lens_cpu[:num_reqs].to(device,

@@ -447,85 +484,103 @@ class MLACommonMetadataBuilder(Generic[T]):

484 input_positions = self.runner.positions_cpu[:num_actual_tokens].to(

485 device, non_blocking=True).long()

486

450 - context_chunk_cu_seq_lens = None

451 - context_chunk_starts = None

452 - context_chunk_seq_tot = None

453 - context_chunk_max_seq_lens = None

454 -

455 - num_computed_tokens_cpu_tensor = \

456 - self.runner.input_batch.num_computed_tokens_cpu_tensor[:num_reqs]

457 - context_lens_tensor = \

458 - num_computed_tokens_cpu_tensor.to(device, non_blocking=True)

459 -

460 - if self.chunked_prefill_enabled and self._num_prefills > 0 \

461 - and context_lens_tensor[self._num_decodes:].max() > 0:

462 - # NOTE: it is recommend you read the `Chunked Prefill` section in

463 - # the comment at the top of the file before trying to understand

464 - # the following code

465 -

466 - self.has_context = True

467 -

468 - num_prefills_with_context = \

469 - (context_lens_tensor[self._num_decodes:] > 0).sum().item()

470 -

471 - # currently we allocate an equal amount of workspace for each

472 - # prefill in the batch, we could probably use a more advanced

473 - # algorithm here and allocate more workspace to prefills with

474 - # longer context lengths

475 - max_context_chunk = \

476 - self.chunked_prefill_workspace_size // num_prefills_with_context

477 -

478 - # align max_context_chunk to page_size by rounding down,

479 - # currently the `gather_cache` kernel cannot handle

480 - # `context_chunk_starts` that are not aligned to page_size

481 - max_context_chunk = round_down(max_context_chunk, self.page_size)

482 - assert max_context_chunk > 0

483 - num_chunks = cdiv(context_lens_tensor.max(), max_context_chunk)

484 -

485 - # if `max_context_chunk = 256`, `num_chunks = 3`, and

486 - # `num_prefills_with_context = 4`, create a tensor that looks like

487 - # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]

488 - context_chunk_starts = \

489 - torch.arange(num_chunks, device=device, dtype=torch.int32) \

490 - .unsqueeze(1).expand(-1, self._num_prefills) \

491 - * max_context_chunk

492 - chunk_ends = torch.min(context_lens_tensor[self._num_decodes:] \

493 - .unsqueeze(0), context_chunk_starts + max_context_chunk)

494 - chunk_seq_lens = (chunk_ends - context_chunk_starts).clamp(min=0)

495 - _context_chunk_cu_seq_lens = chunk_seq_lens.cumsum(dim=1).to(

496 - torch.int32)

497 - zero = torch.zeros(num_chunks, dtype=torch.int32, device=device) \

498 - .unsqueeze(-1)

499 - context_chunk_cu_seq_lens = \

500 - torch.cat([zero, _context_chunk_cu_seq_lens], dim=1)

501 - context_chunk_max_seq_lens = \

502 - chunk_seq_lens.max(dim=1).values.tolist()

503 - context_chunk_seq_tot = chunk_seq_lens.sum(dim=1).tolist()

504 - assert max(context_chunk_seq_tot) <= \

505 - self.chunked_prefill_workspace_size

506 -

507 - return self.cls(

508 - input_positions=input_positions,

487 + prefill_metadata = None

488 + if self._num_prefills > 0:

489 + reqs_start = self._num_decodes # prefill_start

490 + tokens_start = self._num_decode_tokens

491 +

492 + context_lens_cpu = self.runner.input_batch.\

493 + num_computed_tokens_cpu_tensor[reqs_start:num_reqs]

494 + context_lens = context_lens_cpu.to(device, non_blocking=True)

495 +

496 + chunked_context_metadata = None

497 + if self.chunked_prefill_enabled and self._num_prefills > 0 \

498 + and context_lens.max() > 0:

499 + # NOTE: it is recommend you read the `Chunked Prefill` section

500 + # in the comment at the top of the file before trying to

501 + # understand the following code

502 +

503 + num_prefills_with_context = (context_lens > 0).sum().item()

504 +

505 + # currently we allocate an equal amount of workspace for each

506 + # prefill in the batch, we could probably use a more advanced

507 + # algorithm here and allocate more workspace to prefills with

508 + # longer context lengths

509 + max_context_chunk = \

510 + self.chunked_prefill_workspace_size \

511 + // num_prefills_with_context

512 +

513 + # align max_context_chunk to page_size by rounding down,

514 + # currently the `gather_cache` kernel cannot handle

515 + # `context_chunk_starts` that are not aligned to page_size

516 + max_context_chunk = round_down(max_context_chunk,

517 + self.page_size)

518 +

519 + assert max_context_chunk > 0

520 + num_chunks = cdiv(context_lens.max(), max_context_chunk)

521 +

522 + # if `max_context_chunk = 256`, `num_chunks = 3`, and

523 + # `num_prefills_with_context = 4`, create a tensor that looks

524 + # like

525 + # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]

526 + chunk_starts = \

527 + torch.arange(num_chunks, device=device, dtype=torch.int32) \

528 + .unsqueeze(1).expand(-1, self._num_prefills) \

529 + * max_context_chunk

530 + chunk_ends = torch.min(context_lens.unsqueeze(0),

531 + chunk_starts + max_context_chunk)

532 + chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)

533 + _chunk_cu_seq_lens = chunk_seq_lens.cumsum(dim=1).to(

534 + torch.int32)

535 + zero = torch.zeros(num_chunks,

536 + dtype=torch.int32,

537 + device=device).unsqueeze(-1)

538 +

539 + chunked_context_metadata = \

540 + MLACommonPrefillMetadata.ChunkedContextMetadata(

541 + cu_seq_lens=torch.cat(

542 + [zero, _chunk_cu_seq_lens], dim=1),

543 + starts=chunk_starts,

544 + seq_tot=chunk_seq_lens.sum(dim=1).tolist(),

545 + max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

546 + workspace=self.chunked_prefill_workspace,

547 + )

548 +

549 + assert max(chunked_context_metadata.max_seq_lens) <= \

550 + self.chunked_prefill_workspace_size

551 +

552 + prefill_metadata = MLACommonPrefillMetadata(

553 + input_positions=input_positions[tokens_start:],

554 + block_table=block_table[reqs_start:, ...],

555 + query_start_loc=query_start_loc[reqs_start:] -

556 + query_start_loc[reqs_start],

557 + max_query_len=seq_lens[reqs_start:].max().item(),

558 + chunked_context=chunked_context_metadata,

559 + )

560 +

561 + decode_metadata = None

562 + if self._num_decodes > 0:

563 + decode_metadata = self._build_decode(

564 + input_positions=input_positions[:self._num_decode_tokens],

565 + block_table=block_table[:self._num_decodes, ...],

566 + seq_lens=seq_lens[:self._num_decodes],

567 + )

568 +

569 + return self.metadata_cls(

570 num_actual_tokens=num_actual_tokens,

510 - max_query_len=max_query_len,

571 query_start_loc=query_start_loc,

512 - max_seq_len=max_seq_len,

513 - seq_lens=seq_lens,

514 - block_table=block_table,

572 slot_mapping=slot_mapping,

573 head_dim=self.runner.model_config.get_head_size(),

574 # MLACommonMetadata Chunk prefill specific

575 num_decodes=self._num_decodes,

576 num_decode_tokens=self._num_decode_tokens,

577 num_prefills=self._num_prefills,

521 - context_chunk_cu_seq_lens=context_chunk_cu_seq_lens,

522 - context_chunk_starts=context_chunk_starts,

523 - context_chunk_seq_tot=context_chunk_seq_tot,

524 - context_chunk_max_seq_lens=context_chunk_max_seq_lens,

578 + prefill=prefill_metadata,

579 + decode=decode_metadata,

580 )

581

582

528 -class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

583 +class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

584 """

585 NOTE: Please read the comment at the top of the file before trying to

586 understand this class

@@ -798,28 +853,24 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

853 kv_c_and_k_pe_cache: torch.Tensor,

854 attn_metadata: MLACommonMetadata,

855 ):

801 - assert attn_metadata.num_prefills is not None

802 - assert attn_metadata.context_chunk_seq_tot is not None

803 - assert attn_metadata.context_chunk_cu_seq_lens is not None

804 - assert attn_metadata.context_chunk_starts is not None

805 - assert attn_metadata.context_chunk_max_seq_lens is not None

856 + assert attn_metadata.prefill is not None

857 + prefill_metadata = attn_metadata.prefill

858 + assert prefill_metadata.chunked_context is not None

859

860 output = None

808 - iters = len(attn_metadata.context_chunk_seq_tot)

809 -

810 - assert attn_metadata.chunked_prefill_workspace is not None

811 - workspace = attn_metadata.chunked_prefill_workspace

861 + iters = len(prefill_metadata.chunked_context.seq_tot)

862 + workspace = prefill_metadata.chunked_context.workspace

863

864 for i in range(iters):

814 - toks = attn_metadata.context_chunk_seq_tot[i]

865 + toks = prefill_metadata.chunked_context.seq_tot[i]

866

867 ops.gather_cache(

868 src_cache=kv_c_and_k_pe_cache,

869 dst=workspace,

819 - block_table=attn_metadata.block_table,

820 - cu_seq_lens=attn_metadata.context_chunk_cu_seq_lens[i],

870 + block_table=prefill_metadata.block_table,

871 + cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],

872 batch_size=attn_metadata.num_prefills,

822 - seq_starts=attn_metadata.context_chunk_starts[i],

873 + seq_starts=prefill_metadata.chunked_context.starts[i],

874 )

875

876 kv_c_normed = workspace[:toks]\

@@ -845,10 +896,10 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

896 q=q,

897 k=k,

898 v=v_padded,

848 - cu_seqlens_q=attn_metadata.query_start_loc,

849 - cu_seqlens_k=attn_metadata.context_chunk_cu_seq_lens[i],

850 - max_seqlen_q=attn_metadata.max_query_len,

851 - max_seqlen_k=attn_metadata.context_chunk_max_seq_lens[i],

899 + cu_seqlens_q=prefill_metadata.query_start_loc,

900 + cu_seqlens_k=prefill_metadata.chunked_context.cu_seq_lens[i],

901 + max_seqlen_q=prefill_metadata.max_query_len,

902 + max_seqlen_k=prefill_metadata.chunked_context.max_seq_lens[i],

903 softmax_scale=self.scale,

904 causal=False, # Context is unmasked

905 return_softmax_lse=True,

@@ -881,7 +932,9 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

932 kv_c_and_k_pe_cache: torch.Tensor,

933 attn_metadata: MLACommonMetadata,

934 ) -> torch.Tensor:

884 - has_context = attn_metadata.has_context

935 + assert attn_metadata.prefill is not None

936 +

937 + has_context = attn_metadata.prefill.chunked_context is not None

938 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\

939 -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)

940 k_nope, v = kv_nope\

@@ -898,10 +951,10 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

951 q=q,

952 k=k,

953 v=v_padded,

901 - cu_seqlens_q=attn_metadata.query_start_loc,

902 - cu_seqlens_k=attn_metadata.query_start_loc,

903 - max_seqlen_q=attn_metadata.max_query_len,

904 - max_seqlen_k=attn_metadata.max_seq_len,

954 + cu_seqlens_q=attn_metadata.prefill.query_start_loc,

955 + cu_seqlens_k=attn_metadata.prefill.query_start_loc,

956 + max_seqlen_q=attn_metadata.prefill.max_query_len,

957 + max_seqlen_k=attn_metadata.prefill.max_query_len,

958 softmax_scale=self.scale,

959 causal=True,

960 return_softmax_lse=has_context,

@@ -934,7 +987,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

987 q_nope: torch.Tensor,

988 q_pe: torch.Tensor,

989 kv_c_and_k_pe_cache: torch.Tensor,

937 - attn_metadata: T,

990 + attn_metadata: M,

991 ) -> torch.Tensor:

992 raise NotImplementedError

993

@@ -945,7 +998,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

998 k_c_normed: torch.Tensor, # key in unified attn

999 k_pe: torch.Tensor, # value in unified attn

1000 kv_cache: torch.Tensor,

948 - attn_metadata: T,

1001 + attn_metadata: M,

1002 output: Optional[torch.Tensor] = None,

1003 ) -> torch.Tensor:

1004

@@ -966,7 +1019,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1019

1020 # Restore head dim (for rotary embedding)

1021 k_pe = k_pe.unsqueeze(1)

969 - assert hasattr(attn_metadata, "input_positions")

1022

1023 assert attn_metadata.num_decodes is not None and \

1024 attn_metadata.num_prefills is not None and \

@@ -978,28 +1030,27 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1030

1031 decode_hs_or_q_c = hidden_states_or_q_c[:num_decode_tokens]

1032 decode_k_pe = k_pe[:num_decode_tokens]

981 - decode_input_positions = \

982 - attn_metadata.input_positions[:num_decode_tokens]

1033

1034 prefill_hs_or_q_c = hidden_states_or_q_c[num_decode_tokens:]

1035 prefill_k_pe = k_pe[num_decode_tokens:]

986 - prefill_input_positions = \

987 - attn_metadata.input_positions[num_decode_tokens:]

1036 prefill_k_c_normed = k_c_normed[num_decode_tokens:]

1037

1038 if has_decode:

1039 + assert attn_metadata.decode is not None

1040 decode_q_nope = self._q_proj_and_k_up_proj(decode_hs_or_q_c)

1041 decode_q_pe = torch.matmul(decode_hs_or_q_c, self.W_QR)\

1042 .view(-1, self.num_heads, self.qk_rope_head_dim)

1043 decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(

995 - decode_input_positions, decode_q_pe, decode_k_pe)

1044 + attn_metadata.decode.input_positions, decode_q_pe, decode_k_pe)

1045

1046 if has_prefill:

1047 + assert attn_metadata.prefill is not None

1048 prefill_q = self.q_proj(prefill_hs_or_q_c)[0]\

1049 .view(-1, self.num_heads, self.qk_head_dim)

1050 prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]

1051 prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(

1002 - prefill_input_positions, prefill_q_pe, prefill_k_pe)

1052 + attn_metadata.prefill.input_positions, prefill_q_pe,

1053 + prefill_k_pe)

1054

1055 # write the latent and rope to kv cache

1056 if kv_cache.numel() > 0:

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

+33 -25

@@ -11,6 +11,7 @@ from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,

11 is_flashmla_supported)

12 from vllm.logger import init_logger

13 from vllm.v1.attention.backends.mla.common import (MLACommonBackend,

14 + MLACommonDecodeMetadata,

15 MLACommonImpl,

16 MLACommonMetadata,

17 MLACommonMetadataBuilder)

@@ -38,34 +39,41 @@ class FlashMLABackend(MLACommonBackend):

41 @dataclass

41 -class FlashMLAMetadata(MLACommonMetadata):

42 - decode_tile_scheduler_metadata: Optional[tuple[torch.Tensor,

43 - torch.Tensor]] = None

44 - decode_num_splits: Optional[torch.Tensor] = None

42 +class FlashMLADecodeMetadata(MLACommonDecodeMetadata):

43 + tile_scheduler_metadata: tuple[torch.Tensor, torch.Tensor]

44 + num_splits: torch.Tensor

45 +

46 +

47 +@dataclass

48 +class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):

49 + pass

52 class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

54 def __init__(self, runner):

50 - super().__init__(runner, cls=FlashMLAMetadata)

55 + super().__init__(runner)

57 self.num_q_heads = self.runner.model_config.get_num_attention_heads(

58 self.runner.parallel_config)

55 - def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,

56 - common_prefix_len: int):

57 - m = super().build(num_reqs, num_actual_tokens, max_query_len,

58 - common_prefix_len)

59 -

60 - if m.num_decode_tokens is not None and m.num_decode_tokens > 0:

61 - m.decode_tile_scheduler_metadata, m.decode_num_splits = \

62 - get_mla_metadata(

63 - m.seq_lens[:m.num_decode_tokens],

64 - self.num_q_heads,

65 - 1, # MQA for the decode path

66 - )

60 + def _build_decode(self, input_positions: torch.Tensor,

61 + block_table: torch.Tensor,

62 + seq_lens: torch.Tensor) -> FlashMLADecodeMetadata:

63 + tile_scheduler_metadata, num_splits = \

64 + get_mla_metadata(

65 + seq_lens,

66 + self.num_q_heads,

67 + 1, # MQA for the decode path

68 + )

68 - return m

70 + return FlashMLADecodeMetadata(

71 + input_positions=input_positions,

72 + block_table=block_table,

73 + seq_lens=seq_lens,

74 + tile_scheduler_metadata=tile_scheduler_metadata,

75 + num_splits=num_splits,

76 + )

79 class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

@@ -115,6 +123,8 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

123 attn_metadata: FlashMLAMetadata,

124 ) -> torch.Tensor:

125 assert kv_c_and_k_pe_cache.numel() > 0

126 + assert attn_metadata.decode is not None

127 +

128 if self.kv_cache_dtype.startswith("fp8"):

129 raise NotImplementedError("FP8 FlashMLA not yet supported")

130

@@ -124,14 +134,12 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

134 o, _ = flash_mla_with_kvcache(

135 q=q,

136 k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1

127 - block_table=attn_metadata.block_table[:attn_metadata.num_decodes,

128 - ...],

129 - cache_seqlens=attn_metadata.seq_lens[:attn_metadata.

130 - num_decode_tokens],

137 + block_table=attn_metadata.decode.block_table,

138 + cache_seqlens=attn_metadata.decode.seq_lens,

139 head_dim_v=self.kv_lora_rank,

132 - tile_scheduler_metadata=attn_metadata.

133 - decode_tile_scheduler_metadata,

134 - num_splits=attn_metadata.decode_num_splits,

140 + tile_scheduler_metadata=attn_metadata.decode.

141 + tile_scheduler_metadata,

142 + num_splits=attn_metadata.decode.num_splits,

143 softmax_scale=self.scale,

144 causal=True,

145 )

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

+5 -2

@@ -69,6 +69,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

69 attn_metadata: MLACommonMetadata,

70 ) -> torch.Tensor:

71 assert kv_c_and_k_pe_cache.numel() > 0

72 + assert attn_metadata.decode is not None

73 +

74 if self.kv_cache_dtype.startswith("fp8"):

75 raise NotImplementedError("FP8 Triton MLA not yet supported")

@@ -104,7 +106,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

106

107 # Run MQA

108 decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,

107 - attn_metadata.block_table, attn_metadata.seq_lens,

108 - attn_logits, num_kv_splits, self.scale, PAGE_SIZE)

109 + attn_metadata.decode.block_table,

110 + attn_metadata.decode.seq_lens, attn_logits,

111 + num_kv_splits, self.scale, PAGE_SIZE)

112

113 return self._v_up_proj_and_o_proj(o)

vllm/v1/worker/gpu_input_batch.py

+23 -2

@@ -383,8 +383,6 @@ class InputBatch:

383 self.req_id_to_index[old_id_i2], self.req_id_to_index[old_id_i1]

384 self.num_tokens[i1], self.num_tokens[i2] =\

385 self.num_tokens[i2], self.num_tokens[i1]

386 - self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\

387 - self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]

386 self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] =\

387 self.num_tokens_no_spec[i2], self.num_tokens_no_spec[i1]

388 self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] =\

@@ -406,24 +404,47 @@ class InputBatch:

404 self.min_p_cpu[i1], self.min_p_cpu[i2] =\

405 self.min_p_cpu[i2], self.min_p_cpu[i1]

406

407 + # NOTE: the following is unsafe

408 + # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\

409 + # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]

410 + # instead, we need to temporiarily copy the data for one of the indices

411 + # TODO(lucas): optimize this by only copying valid indices

412 + tmp = self.token_ids_cpu[i1, ...].copy()

413 + self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]

414 + self.token_ids_cpu[i2, ...] = tmp

415 +

416 g1 = self.generators.get(i1)

417 g2 = self.generators.get(i2)

418 if g1 is not None:

419 self.generators[i2] = g1

420 + else:

421 + self.generators.pop(i2, None)

422 if g2 is not None:

423 self.generators[i1] = g2

424 + else:

425 + self.generators.pop(i1, None)

426

427 t1 = self.min_tokens.get(i1)

428 t2 = self.min_tokens.get(i2)

429 if t1 is not None:

430 self.min_tokens[i2] = t1

431 + else:

432 + self.min_tokens.pop(i2, None)

433 if t2 is not None:

434 self.min_tokens[i1] = t2

435 + else:

436 + self.min_tokens.pop(i1, None)

437

438 self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\

439 self.request_lora_mapping[i2], self.request_lora_mapping[i1]

440 self.logit_bias[i1], self.logit_bias[i2] =\

441 self.logit_bias[i2], self.logit_bias[i1]

442 +

443 + if self.allowed_token_ids_mask_cpu_tensor is not None:

444 + self.allowed_token_ids_mask_cpu_tensor[i1], \

445 + self.allowed_token_ids_mask_cpu_tensor[i2] =\

446 + self.allowed_token_ids_mask_cpu_tensor[i2], \

447 + self.allowed_token_ids_mask_cpu_tensor[i1]

448 self.block_table.swap_row(i1, i2)

449

450 def condense(self, empty_req_indices: list[int]) -> None:

vllm/v1/worker/gpu_model_runner.py

+4 -2

@@ -456,8 +456,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

456 # Some attention backends (namely MLA) may want to separate requests

457 # based on if the attention computation will be compute-bound or

458 # memory-bound. This gives them a hook to do that.

459 - self.attn_metadata_builder.reorder_batch(self.input_batch,

460 - scheduler_output)

459 + modified_batch = self.attn_metadata_builder.reorder_batch(

460 + self.input_batch, scheduler_output)

461 + if modified_batch:

462 + self.input_batch.refresh_sampling_metadata()

463

464 # OPTIMIZATION: Start copying the block table first.

465 # This way, we can overlap the copy with the following CPU operations.