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 14, 05:30 UTC
[attention][DCP] use AttentionImpl.need_to_return_lse_for_decode (#24372)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Sep 7, 2025 at 09:18 UTC
558f0907dc67c804c5821c92f4d64ed43d10489b
4 files changed
+38
-9
vllm/attention/backends/abstract.py
+26
@@ -257,6 +257,32 @@ class AttentionLayer(Protocol):
257
258
class AttentionImpl(ABC, Generic[T]):
259
260
+ # Whether the attention impl can return the softmax lse for decode.
261
+ # Some features like decode context parallelism require the softmax lse.
262
+ can_return_lse_for_decode: bool = False
263
+
264
+ # some attention backends might not always want to return lse
265
+ # even if they can return lse (for efficiency reasons)
266
+ need_to_return_lse_for_decode: bool = False
267
+
268
+ dcp_world_size: int
269
+ dcp_rank: int
270
+
271
+ def __new__(cls, *args, **kwargs):
272
+ # use __new__ so that all subclasses will call this
273
+ self = super().__new__(cls)
274
+ try:
275
+ from vllm.distributed.parallel_state import get_dcp_group
276
+ self.dcp_world_size = get_dcp_group().world_size
277
+ self.dcp_rank = get_dcp_group().rank_in_group
278
+ except AssertionError:
279
+ # DCP might not be initialized in testing
280
+ self.dcp_world_size = 1
281
+ self.dcp_rank = 0
282
+ self.need_to_return_lse_for_decode = self.dcp_world_size > 1 \
283
+ and self.can_return_lse_for_decode
284
+ return self
285
+
286
@abstractmethod
287
def __init__(
288
self,
vllm/v1/attention/backends/mla/common.py
-4
@@ -1592,10 +1592,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1592
1593
# recorect dcp attn_out with lse.
1594
if self.dcp_world_size > 1:
1595
- assert lse is not None, (
1596
- "For a mla backend want to enable"
1597
- "DCP, it is mandatory that the corresponding decode attn"
1598
- "kernel return the softmax lse.")
1595
attn_out = cp_lse_ag_out_rs(attn_out, lse, get_dcp_group())
1596
1597
# v_up projection
vllm/v1/attention/backends/mla/flashmla.py
+2
@@ -133,6 +133,8 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
133
134
class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
135
136
+ can_return_lse_for_decode: bool = True
137
+
138
def __init__(
139
self,
140
num_heads: int,
vllm/v1/worker/gpu_model_runner.py
+10
-5
@@ -56,7 +56,6 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
56
GiB_bytes, LazyLoader, cdiv, check_use_alibi,
57
get_dtype_size, is_pin_memory_available, round_up,
58
supports_dynamo)
59
-from vllm.v1.attention.backends.mla.flashmla import FlashMLABackend
59
from vllm.v1.attention.backends.utils import (
60
AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,
61
create_fast_prefill_custom_backend,
@@ -3405,10 +3404,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3404
copy_kv_blocks)
3405
3406
if self.dcp_world_size > 1:
3408
- assert self.attn_groups[0][0].backend is FlashMLABackend, (
3409
- "DCP only support flashmla now."
3410
- "For a mla backend want to enable DCP, it is mandatory that the"
3411
- "corresponding decode attn kernel return the softmax lse.")
3407
+ layer_names = self.attn_groups[0][0].layer_names
3408
+ layers = get_layers_from_vllm_config(self.vllm_config,
3409
+ AttentionLayerBase,
3410
+ layer_names)
3411
+ for layer in layers.values():
3412
+ assert layer.impl.need_to_return_lse_for_decode, (
3413
+ "DCP requires attention impls to return"
3414
+ " the softmax lse for decode, but the impl "
3415
+ f"{layer.impl.__class__.__name__} "
3416
+ "does not return the softmax lse for decode.")
3417
3418
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
3419
"""