@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2396a61108

[Attention][MLA][DCP] Query replication for MLA decode (DeepSeek-V2/R1 + Kimi-K2.5) (#45964)

Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

sungsoo ha committed Jul 20, 2026 at 16:51 UTC 2396a611085d7bc78c0f7c41ab45eccbbc7ec0a7

6 files changed +157 -15

tests/v1/attention/test_mla_backends.py

@@ -119,6 +119,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):

119 layer.kv_b_proj.quant_method = None

120 layer.is_aiter_triton_fp4_bmm_enabled = False

121 layer.is_aiter_triton_fp8_bmm_enabled = False

122 + layer.dcp_q_replicate = False

123 layer.quant_config = None

124 layer.layer_name = "test"

125

vllm/envs.py

@@ -186,6 +186,7 @@ if TYPE_CHECKING:

186 VLLM_MOE_USE_DEEP_GEMM: bool = True

187 VLLM_USE_DEEP_GEMM_E8M0: bool = True

188 VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True

189 + VLLM_DCP_Q_REPLICATE: bool = False

190 VLLM_DEEP_GEMM_WARMUP: Literal[

191 "skip",

192 "full",

@@ -1490,6 +1491,8 @@ environment_variables: dict[str, Callable[[], Any]] = {

1491 "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool(

1492 int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1"))

1493 ),

1494 + # Opt-in MLA DCP query replication: skip the decode query all-gather.

1495 + "VLLM_DCP_Q_REPLICATE": lambda: bool(int(os.getenv("VLLM_DCP_Q_REPLICATE", "0"))),

1496 # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm

1497 # JIT all the required kernels before model execution so there is no

1498 # JIT'ing in the hot-path. However, this warmup increases the engine

vllm/model_executor/layers/attention/mla_attention.py

+46 -5

@@ -366,6 +366,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

366 q_lora_rank: int | None,

367 kv_lora_rank: int,

368 kv_b_proj: ColumnParallelLinear,

369 + dcp_q_replicate: bool = False,

370 cache_config: CacheConfig | None = None,

371 quant_config: QuantizationConfig | None = None,

372 prefix: str = "",

@@ -384,10 +385,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):

385 self.q_lora_rank = q_lora_rank

386 self.kv_lora_rank = kv_lora_rank

387 self.kv_b_proj = kv_b_proj

388 + self.dcp_q_replicate = dcp_q_replicate

389 + self.W_UK_T_dcp_qrep: torch.Tensor | None = None

390 self.head_size = kv_lora_rank + qk_rope_head_dim

391 self.layer_name = prefix

392 self.indexer = indexer

390 -

393 self.num_kv_heads = 1

394 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

395

@@ -591,6 +593,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

593 kv_c_normed: torch.Tensor,

594 k_pe: torch.Tensor,

595 output_shape: torch.Size | None = None,

596 + q_dcp_replicated: torch.Tensor | None = None,

597 ) -> torch.Tensor:

598 if self.calculate_kv_scales:

599 torch.ops.vllm.maybe_calc_kv_scales(

@@ -646,6 +649,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

649 self_kv_cache,

650 attn_metadata,

651 output=output,

652 + q_dcp_replicated=q_dcp_replicated,

653 )

654 return output

655 else:

@@ -665,6 +669,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

669 output,

670 encoded,

671 kv_cache_dummy_dep=kv_cache_dummy_dep,

672 + q_dcp_replicated=q_dcp_replicated,

673 )

674 return output

675

@@ -682,6 +687,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

687 quant_scale_ue8m0: bool | None = None,

688 quant_col_major: bool | None = None,

689 quant_tma_aligned: bool | None = None,

690 + q_dcp_replicated: torch.Tensor | None = None,

691 ) -> torch.Tensor:

692 assert output is not None, "Output tensor must be provided."

693

@@ -738,6 +744,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):

744 output_padded = output

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

746 q = q[:num_actual_toks, ...]

747 + if q_dcp_replicated is not None:

748 + q_dcp_replicated = q_dcp_replicated[:num_actual_toks, ...]

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

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

751

@@ -793,7 +801,12 @@ class MLAAttention(nn.Module, AttentionLayerBase):

801 )

802

803 if num_mqa_tokens > 0:

796 - mqa_q = q[:num_mqa_tokens]

804 + if q_dcp_replicated is not None:

805 + mqa_q = q_dcp_replicated[:num_mqa_tokens]

806 + qrep_decode = True

807 + else:

808 + mqa_q = q[:num_mqa_tokens]

809 + qrep_decode = False

810 mqa_output_slice = output[:num_mqa_tokens]

811

812 mqa_q_nope, mqa_q_pe = mqa_q.split(

@@ -833,7 +846,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):

846 else:

847 # Pads the head_dim if necessary (for the underlying kernel)

848 N, B, P = mqa_q_nope.shape

836 - _, _, L = self.W_UK_T.shape

849 + W_UK_T = self.W_UK_T_dcp_qrep if qrep_decode else self.W_UK_T

850 + assert W_UK_T is not None

851 + _, _, L = W_UK_T.shape

852

853 if self.q_pad_num_heads is not None:

854 mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))

@@ -842,7 +857,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

857 mqa_ql_nope = mqa_q_nope.new_empty((N, B, L))

858

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

845 - torch.bmm(mqa_q_nope, self.W_UK_T, out=mqa_ql_nope)

860 + torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)

861

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

863 mqa_ql_nope = mqa_ql_nope.transpose(0, 1)

@@ -855,6 +870,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

870 )

871 else:

872 mqa_q = (mqa_ql_nope, mqa_q_pe)

873 + # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused)

874 if self.impl.dcp_world_size > 1:

875 if self.use_pcp:

876 if self.impl.dcp_world_size > self.impl.pcp_world_size:

@@ -863,8 +879,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):

879 mqa_q = get_tp_group().all_gather(mqa_q, dim=1)

880 else:

881 if isinstance(mqa_q, tuple):

882 + # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)

883 mqa_q = torch.cat(mqa_q, dim=-1)

867 - mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)

884 + if not qrep_decode:

885 + # mqa_q do allgather in head dim.

886 + mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)

887

888 # call decode attn

889 if not self.impl.is_sparse:

@@ -952,6 +971,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):

971 self.kv_b_proj, out_dtype=act_dtype

972 ).T

973

974 + if self.dcp_q_replicate:

975 + # qrep wired here: validate unsupported decode backends once.

976 + assert self.q_pad_num_heads in (None, self.num_heads), (

977 + "DCP query replication is unsupported on head-padding MLA "

978 + "backends (q_pad_num_heads)."

979 + )

980 + if (

981 + self.is_aiter_triton_fp4_bmm_enabled

982 + or self.is_aiter_triton_fp8_bmm_enabled

983 + ):

984 + raise NotImplementedError(

985 + "DCP query replication is not implemented for the aiter "

986 + "FP4/FP8 MLA BMM paths."

987 + )

988 +

989 assert kv_b_proj_weight.shape == (

990 self.kv_lora_rank,

991 self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),

@@ -1032,6 +1066,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):

1066 replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)

1067 # Convert from (L, N, P) to (N, P, L)

1068 replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)

1069 + if self.dcp_q_replicate:

1070 + self.W_UK_T_dcp_qrep = get_dcp_group().all_gather(

1071 + self.W_UK_T.contiguous(), dim=0

1072 + )

1073

1074 # If we should not load quant weights, we initialize the scales to 1.0

1075 # as the default value. See [Note: Register q/k/v/prob scales in state dict]

@@ -1176,6 +1214,7 @@ def unified_mla_attention_with_output(

1214 quant_scale_ue8m0: bool | None = None,

1215 quant_col_major: bool | None = None,

1216 quant_tma_aligned: bool | None = None,

1217 + q_dcp_replicated: torch.Tensor | None = None,

1218 ) -> None:

1219 # kv_cache_dummy_dep is not used but accepting it creates a data dependency

1220 # that ensures torch.compile preserves ordering between KV cache update and

@@ -1196,6 +1235,7 @@ def unified_mla_attention_with_output(

1235 quant_scale_ue8m0=quant_scale_ue8m0,

1236 quant_col_major=quant_col_major,

1237 quant_tma_aligned=quant_tma_aligned,

1238 + q_dcp_replicated=q_dcp_replicated,

1239 )

1240

1241

@@ -1212,6 +1252,7 @@ def unified_mla_attention_with_output_fake(

1252 quant_scale_ue8m0: bool | None = None,

1253 quant_col_major: bool | None = None,

1254 quant_tma_aligned: bool | None = None,

1255 + q_dcp_replicated: torch.Tensor | None = None,

1256 ) -> None:

1257 return

1258

vllm/model_executor/layers/linear.py

+75 -4

@@ -11,6 +11,7 @@ from torch.nn.parameter import Parameter

11 from typing_extensions import TypeIs

13 import vllm.envs as envs

14 +from vllm.config import get_current_vllm_config

15 from vllm.distributed import (

16 divide,

17 get_tensor_model_parallel_rank,

@@ -254,6 +255,8 @@ class LinearBase(PluggableLayer):

255 *,

256 return_bias: bool = True,

257 disable_tp: bool = False,

258 + tp_rank: int | None = None,

259 + tp_size: int | None = None,

260 ):

261 super().__init__()

262

@@ -277,8 +280,17 @@ class LinearBase(PluggableLayer):

280 raise ValueError("All linear layers should support quant method.")

281 self.return_bias = return_bias

282 self.disable_tp = disable_tp

280 - self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0

281 - self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1

283 + if disable_tp:

284 + self.tp_rank, self.tp_size = 0, 1

285 + else:

286 + self.tp_rank = (

287 + tp_rank if tp_rank is not None else get_tensor_model_parallel_rank()

288 + )

289 + self.tp_size = (

290 + tp_size

291 + if tp_size is not None

292 + else get_tensor_model_parallel_world_size()

293 + )

294

295 def update_param_tp_status(self):

296 # Single source of truth for a parameter's TP state. BasevLLMParameter

@@ -425,6 +437,11 @@ class ColumnParallelLinear(LinearBase):

437 (e.g. model.layers.0.qkv_proj)

438 return_bias: If true, return bias together with outputs in forward pass.

439 disable_tp: If true, weights matrix won't be sharded through tp rank.

440 + tp_rank: Override the tensor-parallel rank used for sharding. Defaults to

441 + the global TP rank. Used to shard at a coarser granularity than one

442 + shard per rank (see ``DCPGroupColumnParallelLinear``).

443 + tp_size: Override the tensor-parallel world size used for sharding.

444 + Defaults to the global TP world size.

445 """

446

447 # --8<-- [end:column_parallel_linear]

@@ -442,10 +459,21 @@ class ColumnParallelLinear(LinearBase):

459 *,

460 return_bias: bool = True,

461 disable_tp: bool = False,

462 + tp_rank: int | None = None,

463 + tp_size: int | None = None,

464 ):

465 # Divide the weight matrix along the last dimension.

447 - self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0

448 - self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1

466 + if disable_tp:

467 + self.tp_rank, self.tp_size = 0, 1

468 + else:

469 + self.tp_rank = (

470 + tp_rank if tp_rank is not None else get_tensor_model_parallel_rank()

471 + )

472 + self.tp_size = (

473 + tp_size

474 + if tp_size is not None

475 + else get_tensor_model_parallel_world_size()

476 + )

477 self.input_size_per_partition = input_size

478 self.output_size_per_partition = divide(output_size, self.tp_size)

479 self.output_partition_sizes = [self.output_size_per_partition]

@@ -465,6 +493,8 @@ class ColumnParallelLinear(LinearBase):

493 prefix,

494 return_bias=return_bias,

495 disable_tp=disable_tp,

496 + tp_rank=self.tp_rank,

497 + tp_size=self.tp_size,

498 )

499

500 self._maybe_allow_fp8_block_shape_mismatch()

@@ -586,6 +616,47 @@ class ColumnParallelLinear(LinearBase):

616 return s

617

618

619 +class DCPGroupColumnParallelLinear(ColumnParallelLinear):

620 + """Column-parallel linear whose weight is sharded across DCP groups.

621 +

622 + With Decode Context Parallelism (DCP) the KV cache is sharded across a DCP

623 + group, so MLA decode must attend the group's full head set. This layer shards

624 + its output across DCP *groups* (effective tp size ``tp_size //

625 + dcp_world_size``) rather than across every rank, so each rank in a group

626 + holds the whole group's heads, letting decode skip the query all-gather.

627 +

628 + :meth:`forward` returns the group's full head set. :meth:`_local_view`

629 + extracts this rank's TP shard for prefill.

630 + """

631 +

632 + def __init__(self, *args, **kwargs):

633 + dcp_world_size = (

634 + get_current_vllm_config().parallel_config.decode_context_parallel_size

635 + )

636 + rank = get_tensor_model_parallel_rank()

637 + world_size = get_tensor_model_parallel_world_size()

638 + self.group_size = max(dcp_world_size, 1)

639 + self.qrep_active = self.group_size > 1

640 + self.rank_in_group = rank % self.group_size

641 + super().__init__(

642 + *args,

643 + **kwargs,

644 + tp_rank=rank // self.group_size,

645 + tp_size=world_size // self.group_size,

646 + )

647 +

648 + def _local_view(self, out: torch.Tensor) -> torch.Tensor:

649 + """Slice this rank's tp head shard from a group-heads output.

650 +

651 + ``out`` is head-shaped, i.e. ``(..., group_heads, head_dim)``.

652 + """

653 + if self.group_size == 1:

654 + return out

655 + n = out.shape[-2] // self.group_size

656 + start = self.rank_in_group * n

657 + return out[..., start : start + n, :].contiguous()

658 +

659 +

660 class MergedColumnParallelLinear(ColumnParallelLinear):

661 """Packed linear layers with column parallelism.

662

vllm/model_executor/layers/mla.py

+20 -4

@@ -93,6 +93,10 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):

93 # the topk_tokens buffer written by a previous layer in the same pass.

94 # Refer: https://arxiv.org/abs/2603.12201 for more details.

95 self.skip_topk = skip_topk

96 + # qrep is active when the query projection is a DCP-group-sharded layer

97 + # that materializes the full group head set locally.

98 + q_proj_layer = self.q_b_proj if self.q_lora_rank is not None else self.q_proj

99 + self.dcp_q_replicate = getattr(q_proj_layer, "qrep_active", False)

100 if self.indexer is not None:

101 assert hasattr(self.indexer, "topk_tokens")

102 self.topk_tokens = self.indexer.topk_tokens

@@ -110,6 +114,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):

114 quant_config=quant_config,

115 prefix=f"{prefix}.attn",

116 kv_b_proj=self.kv_b_proj,

117 + dcp_q_replicate=self.dcp_q_replicate,

118 use_sparse=self.is_sparse,

119 indexer=self.indexer,

120 topk_indices_buffer=mla_modules.topk_indices_buffer,

@@ -143,7 +148,8 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):

148 dim=-1,

149 )

150 q_c = self.q_a_layernorm(q_c)

146 - q = self.q_b_proj(q_c)[0]

151 + q_proj_layer = self.q_b_proj

152 + q_proj_input = q_c

153 else:

154 assert self.kv_a_proj_with_mqa is not None, (

155 "kv_a_proj_with_mqa is required when q_lora_rank is None"

@@ -152,15 +158,20 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):

158 "q_proj is required when q_lora_rank is None"

159 )

160 kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]

155 - q = self.q_proj(hidden_states)[0]

161 + q_proj_layer = self.q_proj

162 + q_proj_input = hidden_states

163

164 kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

165 kv_c_normed = self.kv_a_layernorm(kv_c)

159 -

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

166 # Add head dim of 1 to k_pe

167 k_pe = k_pe.unsqueeze(1)

168

169 + q = q_proj_layer(q_proj_input)[0]

170 + heads = self.num_heads

171 + if self.dcp_q_replicate:

172 + heads *= q_proj_layer.group_size

173 + q = q.view(-1, heads, self.qk_head_dim)

174 +

175 if self.rotary_emb is not None:

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

177 positions, q[..., self.qk_nope_head_dim :], k_pe

@@ -172,11 +183,16 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):

183 if llama_4_scaling is not None:

184 q *= llama_4_scaling

185

186 + q_dcp_replicated = None

187 + if self.dcp_q_replicate:

188 + q_dcp_replicated, q = q, q_proj_layer._local_view(q)

189 +

190 attn_out = self.mla_attn(

191 q,

192 kv_c_normed,

193 k_pe,

194 output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),

195 + q_dcp_replicated=q_dcp_replicated,

196 )

197

198 return self.o_proj(attn_out)[0]

vllm/model_executor/models/deepseek_v2.py

+12 -2

@@ -33,6 +33,7 @@ from torch import nn

33 from transformers import DeepseekV2Config, DeepseekV3Config

35 import vllm._custom_ops as ops

36 +import vllm.envs as envs

37 from vllm._aiter_ops import rocm_aiter_ops

38 from vllm.compilation.decorators import support_torch_compile

39 from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config

@@ -56,6 +57,7 @@ from vllm.model_executor.layers.fused_moe import (

57 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm

58 from vllm.model_executor.layers.linear import (

59 ColumnParallelLinear,

60 + DCPGroupColumnParallelLinear,

61 MergedColumnParallelLinear,

62 QKVParallelLinear,

63 ReplicatedLinear,

@@ -1015,9 +1017,17 @@ class DeepseekV2MLAAttention(nn.Module):

1017 prefix=f"{prefix}.kv_a_proj_with_mqa",

1018 )

1019

1020 + qrep_enabled = (

1021 + envs.VLLM_DCP_Q_REPLICATE

1022 + and vllm_config.parallel_config.decode_context_parallel_size > 1

1023 + and vllm_config.parallel_config.prefill_context_parallel_size <= 1

1024 + )

1025 + q_proj_cls = (

1026 + DCPGroupColumnParallelLinear if qrep_enabled else ColumnParallelLinear

1027 + )

1028 if self.q_lora_rank is not None:

1029 self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)

1020 - self.q_b_proj = ColumnParallelLinear(

1030 + self.q_b_proj = q_proj_cls(

1031 self.q_lora_rank,

1032 self.num_heads * self.qk_head_dim,

1033 bias=False,

@@ -1025,7 +1035,7 @@ class DeepseekV2MLAAttention(nn.Module):

1035 prefix=f"{prefix}.q_b_proj",

1036 )

1037 else:

1028 - self.q_proj = ColumnParallelLinear(

1038 + self.q_proj = q_proj_cls(

1039 proj_input_size,

1040 self.num_heads * self.qk_head_dim,

1041 bias=False,