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

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

[Attention] MLA move o_proj q_proj into cuda-graph region (#17484)

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

Lucas Wilkinson committed May 1, 2025 at 23:16 UTC afcb3f8863ee060f8464901a1cf0dfdf57465a90

9 files changed +55 -85

vllm/attention/backends/cpu_mla.py

+2 -3

@@ -281,8 +281,7 @@ class CPUMLAImpl(MLACommonImpl[CPUMLAMetadata]):

281 # remove padding

282 output = output.view(-1, self.num_heads,

283 q.shape[-1])[..., :v.shape[-1]]

284 - output = output.reshape(-1, self.num_heads * v.shape[-1])

285 - return self.o_proj(output)[0]

284 + return output.reshape(-1, self.num_heads * v.shape[-1])

285

286 def _forward_decode(

287 self,

@@ -303,4 +302,4 @@ class CPUMLAImpl(MLACommonImpl[CPUMLAMetadata]):

302 ops.mla_decode_kvcache_cpu(o, q, kv_c_and_k_pe_cache, self.scale,

303 decode_meta.block_tables,

304 decode_meta.seq_lens_tensor)

306 - return self._v_up_proj_and_o_proj(o)

305 + return self._v_up_proj(o)

vllm/attention/backends/flashmla.py

+1 -1

@@ -239,4 +239,4 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

239 causal=True,

240 )

241

242 - return self._v_up_proj_and_o_proj(o)

242 + return self._v_up_proj(o)

vllm/attention/backends/mla/common.py

+19 -35

@@ -207,7 +207,7 @@ from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,

207 from vllm.attention.ops.merge_attn_states import merge_attn_states

208 from vllm.attention.utils.fa_utils import get_flash_attn_version

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

210 - LinearBase, RowParallelLinear,

210 + LinearBase,

211 UnquantizedLinearMethod)

212 from vllm.model_executor.layers.rotary_embedding import (

213 DeepseekScalingRotaryEmbedding, RotaryEmbedding)

@@ -1032,12 +1032,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1032 qk_head_dim: int,

1033 v_head_dim: int,

1034 rotary_emb: RotaryEmbedding,

1035 - # q_proj should be q_b_proj if q_lora_rank is not None, but from an

1036 - # attention backend perspective we rely on the layer to pass in the

1037 - # correct matrix

1038 - q_proj: ColumnParallelLinear,

1035 kv_b_proj: ColumnParallelLinear,

1040 - o_proj: RowParallelLinear,

1036 ) -> None:

1037 self.num_heads = num_heads

1038 self.head_size = head_size

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

1050 self.rotary_emb = rotary_emb

1051 self.use_yarn_rope = isinstance(rotary_emb,

1052 DeepseekScalingRotaryEmbedding)

1058 - self.q_proj = q_proj

1053 self.kv_b_proj = kv_b_proj

1060 - self.o_proj = o_proj

1054

1055 self.triton_fa_func = triton_attention

1056 # Handle the differences between the flash_attn_varlen from flash_attn

@@ -1141,27 +1134,13 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1134 return attn_out, rest[0]

1135 return attn_out

1136

1144 - def _v_up_proj_and_o_proj(self, x):

1137 + def _v_up_proj(self, x):

1138 # Convert from (B, N, L) to (N, B, L)

1139 x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)

1140 # Multiply (N, B, L) x (N, L, V) -> (N, B, V)

1141 x = torch.bmm(x, self.W_UV)

1142 # Convert from (N, B, V) to (B, N * V)

1150 - x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)

1151 - return self.o_proj(x)[0]

1152 -

1153 - # Return `ql_nope`, `q_pe`

1154 - def _q_proj_and_k_up_proj(self, x):

1155 - q_nope, q_pe = self.q_proj(x)[0]\

1156 - .view(-1, self.num_heads, self.qk_head_dim)\

1157 - .split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

1158 -

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

1160 - q_nope = q_nope.transpose(0, 1)

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

1162 - ql_nope = torch.bmm(q_nope, self.W_UK_T)

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

1164 - return ql_nope.transpose(0, 1), q_pe

1143 + return x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)

1144

1145 def process_weights_after_loading(self, act_dtype: torch.dtype):

1146

@@ -1345,7 +1324,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1324 suffix_lse=suffix_lse,

1325 )

1326

1348 - return self.o_proj(output.flatten(start_dim=-2))[0]

1327 + return output.flatten(start_dim=-2)

1328

1329 @abstractmethod

1330 def _forward_decode(

@@ -1360,7 +1339,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1339 def forward(

1340 self,

1341 layer: AttentionLayer,

1363 - hidden_states_or_q_c: torch.Tensor, # query in unified attn

1342 + q: torch.Tensor, # query in unified attn

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

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

1345 kv_cache: torch.Tensor,

@@ -1391,27 +1370,32 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1370 assert hasattr(attn_metadata, "input_positions")

1371

1372 num_prefill_tokens: int = attn_metadata.num_prefill_tokens

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

1374

1395 - decode_hs_or_q_c = hidden_states_or_q_c[num_prefill_tokens:]

1375 + decode_q = q[num_prefill_tokens:]

1376 decode_k_pe = k_pe[num_prefill_tokens:]

1377 decode_input_positions = \

1378 attn_metadata.input_positions[num_prefill_tokens:]

1379

1400 - prefill_hs_or_q_c = hidden_states_or_q_c[:num_prefill_tokens]

1380 + prefill_q = q[:num_prefill_tokens]

1381 prefill_k_pe = k_pe[:num_prefill_tokens]

1382 prefill_input_positions = \

1383 attn_metadata.input_positions[:num_prefill_tokens]

1384 prefill_k_c_normed = k_c_normed[:num_prefill_tokens]

1385

1386 if has_decode:

1407 - decode_ql_nope, decode_q_pe = \

1408 - self._q_proj_and_k_up_proj(decode_hs_or_q_c)

1387 + decode_q_nope, decode_q_pe = decode_q.split(

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

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

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

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

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

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

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

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

1396 decode_input_positions, decode_q_pe, decode_k_pe)

1397

1398 if has_prefill:

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

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

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

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

1401 prefill_input_positions, prefill_q_pe, prefill_k_pe)

@@ -1429,9 +1413,9 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

1413

1414 output = torch.empty(attn_metadata.num_prefill_tokens +

1415 attn_metadata.num_decode_tokens,

1432 - self.o_proj.output_size,

1433 - device=hidden_states_or_q_c.device,

1434 - dtype=hidden_states_or_q_c.dtype)

1416 + self.v_head_dim * self.num_heads,

1417 + device=q.device,

1418 + dtype=q.dtype)

1419 if has_prefill:

1420 output[:num_prefill_tokens] = self._forward_prefill(

1421 prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,

vllm/attention/backends/rocm_aiter_mla.py

+1 -1

@@ -409,4 +409,4 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

409 attn_metadata.paged_kv_indices,

410 attn_metadata.paged_kv_last_page_lens)

411

412 - return self._v_up_proj_and_o_proj(o)

412 + return self._v_up_proj(o)

vllm/attention/backends/triton_mla.py

+1 -1

@@ -110,4 +110,4 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

110 decode_meta.seq_lens_tensor, attn_logits,

111 num_kv_splits, self.scale, PAGE_SIZE)

112

113 - return self._v_up_proj_and_o_proj(o)

113 + return self._v_up_proj(o)

vllm/model_executor/models/deepseek_v2.py

+12 -9

@@ -454,9 +454,7 @@ class DeepseekV2MLAAttention(nn.Module):

454 qk_head_dim=self.qk_head_dim,

455 v_head_dim=self.v_head_dim,

456 rotary_emb=self.rotary_emb,

457 - q_proj=self.q_proj if self.q_lora_rank is None else self.q_b_proj,

457 kv_b_proj=self.kv_b_proj,

459 - o_proj=self.o_proj,

458 )

459

460 self.prefix = prefix

@@ -468,17 +466,22 @@ class DeepseekV2MLAAttention(nn.Module):

466 hidden_states: torch.Tensor,

467 ) -> torch.Tensor:

468 if self.q_lora_rank is not None:

471 - ckq = self.q_a_proj(hidden_states)[0]

472 - hidden_states_or_q_c = self.q_a_layernorm(ckq)

469 + q_c = self.q_a_proj(hidden_states)[0]

470 + q_c = self.q_a_layernorm(q_c)

471 + q = self.q_b_proj(q_c)[0]

472 else:

474 - hidden_states_or_q_c = hidden_states

473 + q = self.q_proj(hidden_states)[0]

474 kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(

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

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

478 - return self.mla_attn(hidden_states_or_q_c,

479 - kv_c_normed,

480 - k_pe,

481 - output_shape=hidden_states.shape)

477 +

478 + attn_out = self.mla_attn(

479 + q,

480 + kv_c_normed,

481 + k_pe,

482 + output_shape=(hidden_states.shape[0],

483 + self.num_local_heads * self.v_head_dim))

484 + return self.o_proj(attn_out)[0]

485

486

487 class DeepseekV2DecoderLayer(nn.Module):

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

+17 -33

@@ -200,7 +200,7 @@ from vllm.attention.ops.merge_attn_states import merge_attn_states

200 from vllm.attention.utils.fa_utils import get_flash_attn_version

201 from vllm.logger import init_logger

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

203 - LinearBase, RowParallelLinear,

203 + LinearBase,

204 UnquantizedLinearMethod)

205 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

206 from vllm.platforms import current_platform

@@ -597,12 +597,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

597 qk_head_dim: int,

598 v_head_dim: int,

599 rotary_emb: RotaryEmbedding,

600 - # q_proj should be q_b_proj if q_lora_rank is not None, but from an

601 - # attention backend perspective we rely on the layer to pass in the

602 - # correct matrix

603 - q_proj: ColumnParallelLinear,

600 kv_b_proj: ColumnParallelLinear,

605 - o_proj: RowParallelLinear,

601 ) -> None:

602 self.num_heads = num_heads

603 self.head_size = head_size

@@ -625,9 +620,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

620 if current_platform.is_cuda():

621 self.rotary_emb = rotary_emb.forward_cuda

622

628 - self.q_proj = q_proj

623 self.kv_b_proj = kv_b_proj

630 - self.o_proj = o_proj

624 self.vllm_flash_attn_version = get_flash_attn_version()

625

626 # Handle the differences between the flash_attn_varlen from flash_attn

@@ -684,27 +677,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

677 return attn_out, lse

678 return attn_out

679

687 - def _v_up_proj_and_o_proj(self, x):

680 + def _v_up_proj(self, x):

681 # Convert from (B, N, L) to (N, B, L)

682 x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)

683 # Multiply (N, B, L) x (N, L, V) -> (N, B, V)

684 x = torch.bmm(x, self.W_UV)

685 # Convert from (N, B, V) to (B, N * V)

693 - x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)

694 - return self.o_proj(x)[0]

695 -

696 - # Return `ql_nope`, `q_pe`

697 - def _q_proj_and_k_up_proj(self, x):

698 - q_nope, q_pe = self.q_proj(x)[0]\

699 - .view(-1, self.num_heads, self.qk_head_dim)\

700 - .split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

701 -

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

703 - q_nope = q_nope.transpose(0, 1)

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

705 - ql_nope = torch.bmm(q_nope, self.W_UK_T)

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

707 - return ql_nope.transpose(0, 1), q_pe

686 + return x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)

687

688 def process_weights_after_loading(self, act_dtype: torch.dtype):

689

@@ -874,7 +853,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

853 suffix_lse=suffix_lse,

854 )

855

877 - return self.o_proj(output.flatten(start_dim=-2))[0]

856 + return output.flatten(start_dim=-2)

857

858 @abstractmethod

859 def _forward_decode(

@@ -889,7 +868,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

868 def forward(

869 self,

870 layer: AttentionLayer,

892 - hidden_states_or_q_c: torch.Tensor, # query in unified attn

871 + q: torch.Tensor,

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

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

874 kv_cache: torch.Tensor,

@@ -908,7 +887,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

887 # Inputs and outputs may be padded for CUDA graphs

888 output_padded = output

889 output = output[:num_actual_toks, ...]

911 - hidden_states_or_q_c = hidden_states_or_q_c[:num_actual_toks, ...]

890 + q = q[:num_actual_toks, ...]

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

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

893

@@ -923,24 +902,29 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

902 has_prefill = attn_metadata.num_prefills > 0

903 num_decode_tokens = attn_metadata.num_decode_tokens

904

926 - decode_hs_or_q_c = hidden_states_or_q_c[:num_decode_tokens]

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

906 + decode_q = q[:num_decode_tokens]

907 decode_k_pe = k_pe[:num_decode_tokens]

908

929 - prefill_hs_or_q_c = hidden_states_or_q_c[num_decode_tokens:]

909 + prefill_q = q[num_decode_tokens:]

910 prefill_k_pe = k_pe[num_decode_tokens:]

911 prefill_k_c_normed = k_c_normed[num_decode_tokens:]

912

913 if has_decode:

914 assert attn_metadata.decode is not None

935 - decode_ql_nope, decode_q_pe = \

936 - self._q_proj_and_k_up_proj(decode_hs_or_q_c)

915 + decode_q_nope, decode_q_pe = decode_q.split(

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

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

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

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

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

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

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

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

924 attn_metadata.decode.input_positions, decode_q_pe, decode_k_pe)

925

926 if has_prefill:

927 assert attn_metadata.prefill is not None

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

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

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

929

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

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

+1 -1

@@ -146,4 +146,4 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

146 causal=True,

147 )

148

149 - return self._v_up_proj_and_o_proj(o)

149 + return self._v_up_proj(o)

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

+1 -1

@@ -115,4 +115,4 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

115 attn_metadata.decode.seq_lens, attn_logits,

116 num_kv_splits, self.scale, PAGE_SIZE)

117

118 - return self._v_up_proj_and_o_proj(o)

118 + return self._v_up_proj(o)