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

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

[Attention] MLA move rotary embedding to cuda-graph region (#17668)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

Lucas Wilkinson committed May 8, 2025 at 23:14 UTC 5e6f93948449b8095e8eef5c3d99a8726e216a44

6 files changed +35 -121

vllm/attention/backends/mla/common.py

+11 -60

@@ -211,8 +211,6 @@ from vllm.attention.utils.fa_utils import get_flash_attn_version

211 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

212 LinearBase,

213 UnquantizedLinearMethod)

214 -from vllm.model_executor.layers.rotary_embedding import (

215 - DeepseekScalingRotaryEmbedding, RotaryEmbedding)

214 from vllm.multimodal import MultiModalPlaceholderMap

215 from vllm.platforms import current_platform

216 from vllm.triton_utils import HAS_TRITON

@@ -377,7 +375,6 @@ class MLACommonState(AttentionState, Generic[T]):

375 seq_start_loc=None,

376 context_lens_tensor=None,

377 block_tables=self._graph_block_tables[:batch_size],

380 - input_positions=self._positions[:batch_size],

378 head_dim=self.runner.model_config.get_head_size())

379

380 if is_encoder_decoder_model:

@@ -393,7 +390,6 @@ class MLACommonState(AttentionState, Generic[T]):

390 "slot_mapping": attn_metadata.slot_mapping,

391 "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor,

392 "block_tables": attn_metadata.decode_metadata.block_tables,

396 - "input_positions": attn_metadata.decode_metadata.input_positions,

393 }

394 if is_encoder_decoder_model:

395 raise NotImplementedError(

@@ -405,16 +401,10 @@ class MLACommonState(AttentionState, Generic[T]):

401 input_buffers,

402 attn_metadata,

403 is_encoder_decoder_model: bool = False):

408 - input_positions = attn_metadata.input_positions

409 - num_positions = input_positions.shape[0]

404 input_buffers["seq_lens_tensor"].copy_(

405 attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True)

406 input_buffers["block_tables"].copy_(

407 attn_metadata.decode_metadata.block_tables, non_blocking=True)

414 - # CUDA graph buffer is padded so only perform a partial copy based on

415 - # num_positions

416 - input_buffers["input_positions"][:num_positions].copy_(

417 - input_positions, non_blocking=True)

408 if is_encoder_decoder_model:

409 raise NotImplementedError(

410 "TritonMLAState does not support encoder/decoder yet")

@@ -456,11 +446,6 @@ class MLACommonMetadata(AttentionMetadata):

446 # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.

447 use_cuda_graph: bool

448

459 - # New for MLA (compared to FlashAttention)

460 - # Input positions for rotrary embeddings since for MLA the rotary

461 - # position embeddings are applied inside the attention backend

462 - input_positions: torch.Tensor

463 -

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

450 # |---------- N-1 iteration --------|

451 # |---------------- N iteration ---------------------|

@@ -563,8 +548,6 @@ class MLACommonMetadata(AttentionMetadata):

548 self.context_lens_tensor[:self.num_prefills])

549 block_tables = (None if self.block_tables is None else

550 self.block_tables[:self.num_prefills])

566 - input_positions = (None if self.input_positions is None else

567 - self.input_positions[:self.num_prefill_tokens])

551

552 self._cached_prefill_metadata = self.__class__(

553 # Required by ModelRunner

@@ -578,7 +561,6 @@ class MLACommonMetadata(AttentionMetadata):

561 multi_modal_placeholder_index_maps=None,

562 enable_kv_scales_calculation=False,

563 # MLACommonMetadata

581 - input_positions=input_positions,

564 seq_lens=seq_lens,

565 seq_lens_tensor=seq_lens_tensor,

566 max_query_len=self.max_query_len,

@@ -615,8 +597,6 @@ class MLACommonMetadata(AttentionMetadata):

597 self.seq_lens_tensor[self.num_prefills:])

598 block_tables = (None if self.block_tables is None else

599 self.block_tables[self.num_prefills:])

618 - input_positions = (None if self.input_positions is None else

619 - self.input_positions[self.num_prefill_tokens:])

600

601 self._cached_decode_metadata = self.__class__(

602 # Required by ModelRunner

@@ -646,7 +626,6 @@ class MLACommonMetadata(AttentionMetadata):

626 if self.seq_start_loc is not None else None,

627 context_lens_tensor=None,

628 block_tables=block_tables,

649 - input_positions=input_positions,

629 head_dim=self.head_dim,

630 is_profile_run=self.is_profile_run)

631 return self._cached_decode_metadata

@@ -765,7 +744,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

744 self.context_lens: List[int] = []

745 self.block_tables: List[List[int]] = []

746 self.curr_seq_lens: List[int] = []

768 - self.input_positions: List[int] = []

747 self.multimodal_placeholder_maps: Dict[

748 str,

749 MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)

@@ -786,13 +764,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

764 block_tables = inter_data.block_tables

765

766 for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

789 - curr_sliding_window_block, input_positions) in zip(

767 + curr_sliding_window_block) in zip(

768 inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

769 inter_data.orig_seq_lens, inter_data.seq_lens,

770 inter_data.query_lens, inter_data.context_lens,

793 - inter_data.curr_sliding_window_blocks,

794 - inter_data.input_positions):

795 - self.input_positions.extend(input_positions)

771 + inter_data.curr_sliding_window_blocks):

772 self.context_lens.append(context_len)

773 if is_prompt:

774 self.num_prefills += 1

@@ -912,8 +888,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

888 device, self.runner.pin_memory)

889 seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,

890 self.runner.pin_memory)

915 - input_positions = async_tensor_h2d(self.input_positions, torch.long,

916 - device, self.runner.pin_memory)

891 slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,

892 device, self.runner.pin_memory)

893 query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,

@@ -987,7 +961,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

961 multi_modal_placeholder_index_maps=None, # Not Attention Related

962 enable_kv_scales_calculation=False,

963 # MLACommonMetadata

990 - input_positions=input_positions,

964 seq_lens=seq_lens,

965 seq_lens_tensor=seq_lens_tensor,

966 max_query_len=max_query_len,

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

1006 qk_rope_head_dim: int,

1007 qk_head_dim: int,

1008 v_head_dim: int,

1036 - rotary_emb: RotaryEmbedding,

1009 kv_b_proj: ColumnParallelLinear,

1010 ) -> None:

1011 self.num_heads = num_heads

@@ -1048,10 +1020,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1020 self.qk_rope_head_dim = qk_rope_head_dim

1021 self.qk_head_dim = qk_head_dim

1022 self.v_head_dim = v_head_dim

1051 -

1052 - self.rotary_emb = rotary_emb

1053 - self.use_yarn_rope = isinstance(rotary_emb,

1054 - DeepseekScalingRotaryEmbedding)

1023 self.kv_b_proj = kv_b_proj

1024

1025 self.triton_fa_func = triton_attention

@@ -1367,41 +1335,15 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1335 has_decode = attn_metadata.decode_metadata is not None

1336 has_prefill = attn_metadata.prefill_metadata is not None

1337

1370 - # Restore head dim (for rotary embedding)

1371 - k_pe = k_pe.unsqueeze(1)

1372 - assert hasattr(attn_metadata, "input_positions")

1373 -

1338 num_prefill_tokens: int = attn_metadata.num_prefill_tokens

1339 q = q.view(-1, self.num_heads, self.qk_head_dim)

1340

1341 decode_q = q[num_prefill_tokens:]

1378 - decode_k_pe = k_pe[num_prefill_tokens:]

1379 - decode_input_positions = \

1380 - attn_metadata.input_positions[num_prefill_tokens:]

1342

1343 prefill_q = q[:num_prefill_tokens]

1344 prefill_k_pe = k_pe[:num_prefill_tokens]

1384 - prefill_input_positions = \

1385 - attn_metadata.input_positions[:num_prefill_tokens]

1345 prefill_k_c_normed = k_c_normed[:num_prefill_tokens]

1346

1388 - if has_decode:

1389 - decode_q_nope, decode_q_pe = decode_q.split(

1390 - [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

1391 - # Convert from (B, N, P) to (N, B, P)

1392 - decode_q_nope = decode_q_nope.transpose(0, 1)

1393 - # Multiply (N, B, P) x (N, P, L) -> (N, B, L)

1394 - decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)

1395 - # Convert from (N, B, L) to (B, N, L)

1396 - decode_ql_nope = decode_ql_nope.transpose(0, 1)

1397 - decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(

1398 - decode_input_positions, decode_q_pe, decode_k_pe)

1399 -

1400 - if has_prefill:

1401 - prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]

1402 - prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(

1403 - prefill_input_positions, prefill_q_pe, prefill_k_pe)

1404 -

1347 # write the latent and rope to kv cache

1348 if kv_cache.numel() > 0:

1349 ops.concat_and_cache_mla(

@@ -1424,6 +1366,15 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1366 attn_metadata)

1367

1368 if has_decode:

1369 + decode_q_nope, decode_q_pe = decode_q.split(

1370 + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

1371 + # Convert from (B, N, P) to (N, B, P)

1372 + decode_q_nope = decode_q_nope.transpose(0, 1)

1373 + # Multiply (N, B, P) x (N, P, L) -> (N, B, L)

1374 + decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)

1375 + # Convert from (N, B, L) to (B, N, L)

1376 + decode_ql_nope = decode_ql_nope.transpose(0, 1)

1377 +

1378 output[num_prefill_tokens:] = self._forward_decode(

1379 decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)

1380

vllm/attention/backends/rocm_aiter_mla.py

+2 -4

@@ -148,13 +148,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

148 block_tables = inter_data.block_tables

149

150 for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

151 - curr_sliding_window_block, input_positions) in zip(

151 + curr_sliding_window_block) in zip(

152 inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

153 inter_data.orig_seq_lens, inter_data.seq_lens,

154 inter_data.query_lens, inter_data.context_lens,

155 - inter_data.curr_sliding_window_blocks,

156 - inter_data.input_positions):

157 - self.input_positions.extend(input_positions)

155 + inter_data.curr_sliding_window_blocks):

156 self.context_lens.append(context_len)

157 if is_prompt:

158 self.num_prefills += 1

vllm/model_executor/layers/rotary_embedding.py

+3 -2

@@ -808,8 +808,9 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):

808 query_pass = query[..., self.rotary_dim:]

809 key_pass = key[..., self.rotary_dim:]

810

811 - self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(

812 - positions.device)

811 + if self.cos_sin_cache.device != positions.device:

812 + self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(

813 + positions.device)

814 cos_sin = self.cos_sin_cache[torch.add(positions, offsets)

815 if offsets is not None else positions]

816 cos, sin = cos_sin.chunk(2, dim=-1)

vllm/model_executor/models/deepseek_v2.py

+7 -1

@@ -453,7 +453,6 @@ class DeepseekV2MLAAttention(nn.Module):

453 qk_rope_head_dim=self.qk_rope_head_dim,

454 qk_head_dim=self.qk_head_dim,

455 v_head_dim=self.v_head_dim,

456 - rotary_emb=self.rotary_emb,

456 kv_b_proj=self.kv_b_proj,

457 )

458

@@ -475,6 +474,13 @@ class DeepseekV2MLAAttention(nn.Module):

474 [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

475 kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())

476

477 + q = q.view(-1, self.num_local_heads, self.qk_head_dim)

478 + # Add head dim of 1 to k_pe

479 + k_pe = k_pe.unsqueeze(1)

480 +

481 + q[..., self.qk_nope_head_dim:], k_pe = self.rotary_emb(

482 + positions, q[..., self.qk_nope_head_dim:], k_pe)

483 +

484 attn_out = self.mla_attn(

485 q,

486 kv_c_normed,

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

+11 -51

@@ -204,7 +204,6 @@ from vllm.logger import init_logger

204 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

205 LinearBase,

206 UnquantizedLinearMethod)

207 -from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

207 from vllm.platforms import current_platform

208 from vllm.utils import cdiv, round_down

209 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

@@ -269,9 +268,6 @@ class MLACommonPrefillMetadata:

268 max_seq_lens: list[int]

269 workspace: torch.Tensor

270

272 - # Input positions for rotrary embeddings since for MLA the rotary

273 - # position embeddings are applied inside the attention backend

274 - input_positions: torch.Tensor

271 block_table: torch.Tensor

272 query_start_loc: torch.Tensor

273 max_query_len: int

@@ -280,9 +276,6 @@ class MLACommonPrefillMetadata:

276

277 @dataclass

278 class MLACommonDecodeMetadata:

283 - # Input positions for rotrary embeddings since for MLA the rotary

284 - # position embeddings are applied inside the attention backend

285 - input_positions: torch.Tensor

279 block_table: torch.Tensor

280 seq_lens: torch.Tensor

281

@@ -443,10 +436,8 @@ class MLACommonMetadataBuilder(Generic[M]):

436

437 return modified_batch

438

446 - def _build_decode(self, input_positions: torch.Tensor,

447 - block_table: torch.Tensor, seq_lens: torch.Tensor):

439 + def _build_decode(self, block_table: torch.Tensor, seq_lens: torch.Tensor):

440 return MLACommonDecodeMetadata(

449 - input_positions=input_positions,

441 block_table=block_table,

442 seq_lens=seq_lens,

443 )

@@ -464,8 +455,6 @@ class MLACommonMetadataBuilder(Generic[M]):

455 self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])

456 slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(

457 device, non_blocking=True).long()

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

468 - device, non_blocking=True).long()

458

459 query_start_loc = common_attn_metadata.query_start_loc

460 seq_lens = common_attn_metadata.seq_lens

@@ -473,7 +462,6 @@ class MLACommonMetadataBuilder(Generic[M]):

462 prefill_metadata = None

463 if self._num_prefills > 0:

464 reqs_start = self._num_decodes # prefill_start

476 - tokens_start = self._num_decode_tokens

465

466 context_lens_cpu = self.runner.input_batch.\

467 num_computed_tokens_cpu_tensor[reqs_start:num_reqs]

@@ -542,7 +530,6 @@ class MLACommonMetadataBuilder(Generic[M]):

530 self.chunked_prefill_workspace_size

531

532 prefill_metadata = MLACommonPrefillMetadata(

545 - input_positions=input_positions[tokens_start:],

533 block_table=block_table[reqs_start:, ...],

534 query_start_loc=prefill_query_start_loc,

535 max_query_len=max_query_len,

@@ -552,7 +539,6 @@ class MLACommonMetadataBuilder(Generic[M]):

539 decode_metadata = None

540 if self._num_decodes > 0:

541 decode_metadata = self._build_decode(

555 - input_positions=input_positions[:self._num_decode_tokens],

542 block_table=block_table[:self._num_decodes, ...],

543 seq_lens=seq_lens[:self._num_decodes],

544 )

@@ -599,7 +585,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

585 qk_rope_head_dim: int,

586 qk_head_dim: int,

587 v_head_dim: int,

602 - rotary_emb: RotaryEmbedding,

588 kv_b_proj: ColumnParallelLinear,

589 ) -> None:

590 self.num_heads = num_heads

@@ -614,15 +599,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

599 self.qk_rope_head_dim = qk_rope_head_dim

600 self.qk_head_dim = qk_head_dim

601 self.v_head_dim = v_head_dim

617 -

618 - # Hack for V1 for now to avoid torch library overhead (since we are

619 - # already inside an attention custom op), pull out the forward

620 - # method from the rotary embedding and call it directly

621 - # TODO(lucas): we should probably find a cleaner way to do this

622 - self.rotary_emb = rotary_emb.forward_native

623 - if current_platform.is_cuda():

624 - self.rotary_emb = rotary_emb.forward_cuda

625 -

602 self.kv_b_proj = kv_b_proj

603 self.vllm_flash_attn_version = get_flash_attn_version()

604

@@ -894,9 +870,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

870 k_c_normed = k_c_normed[:num_actual_toks, ...]

871 k_pe = k_pe[:num_actual_toks, ...]

872

897 - # Restore head dim (for rotary embedding)

898 - k_pe = k_pe.unsqueeze(1)

899 -

873 assert attn_metadata.num_decodes is not None and \

874 attn_metadata.num_prefills is not None and \

875 attn_metadata.num_decode_tokens is not None

@@ -905,35 +878,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

878 has_prefill = attn_metadata.num_prefills > 0

879 num_decode_tokens = attn_metadata.num_decode_tokens

880

908 - q = q.view(-1, self.num_heads, self.qk_head_dim)

881 decode_q = q[:num_decode_tokens]

910 - decode_k_pe = k_pe[:num_decode_tokens]

882

883 prefill_q = q[num_decode_tokens:]

884 prefill_k_pe = k_pe[num_decode_tokens:]

885 prefill_k_c_normed = k_c_normed[num_decode_tokens:]

886

916 - if has_decode:

917 - assert attn_metadata.decode is not None

918 - decode_q_nope, decode_q_pe = decode_q.split(

919 - [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

920 - # Convert from (B, N, P) to (N, B, P)

921 - decode_q_nope = decode_q_nope.transpose(0, 1)

922 - # Multiply (N, B, P) x (N, P, L) -> (N, B, L)

923 - decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)

924 - # Convert from (N, B, L) to (B, N, L)

925 - decode_ql_nope = decode_ql_nope.transpose(0, 1)

926 - decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(

927 - attn_metadata.decode.input_positions, decode_q_pe, decode_k_pe)

928 -

929 - if has_prefill:

930 - assert attn_metadata.prefill is not None

931 - prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]

932 -

933 - prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(

934 - attn_metadata.prefill.input_positions, prefill_q_pe,

935 - prefill_k_pe)

936 -

887 # write the latent and rope to kv cache

888 if kv_cache.numel() > 0:

889 ops.concat_and_cache_mla(

@@ -951,6 +901,16 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

901 attn_metadata)

902

903 if has_decode:

904 + assert attn_metadata.decode is not None

905 + decode_q_nope, decode_q_pe = decode_q.split(

906 + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

907 + # Convert from (B, N, P) to (N, B, P)

908 + decode_q_nope = decode_q_nope.transpose(0, 1)

909 + # Multiply (N, B, P) x (N, P, L) -> (N, B, L)

910 + decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)

911 + # Convert from (N, B, L) to (B, N, L)

912 + decode_ql_nope = decode_ql_nope.transpose(0, 1)

913 +

914 output[:num_decode_tokens] = self._forward_decode(

915 decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)

916

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

+1 -3

@@ -58,8 +58,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

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

59 self.runner.parallel_config)

61 - def _build_decode(self, input_positions: torch.Tensor,

62 - block_table: torch.Tensor,

61 + def _build_decode(self, block_table: torch.Tensor,

62 seq_lens: torch.Tensor) -> FlashMLADecodeMetadata:

63 tile_scheduler_metadata, num_splits = \

64 get_mla_metadata(

@@ -69,7 +68,6 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

68 )

70 return FlashMLADecodeMetadata(

72 - input_positions=input_positions,

71 block_table=block_table,

72 seq_lens=seq_lens,

73 tile_scheduler_metadata=tile_scheduler_metadata,