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 2, 23:15 UTC
[DSV4] More multi-stream enablement for c4a (#42925)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
May 22, 2026 at 21:22 UTC
367cb81966f98f6ba73e5e3f529573d017c446d0
2 files changed
+61
-29
vllm/models/deepseek_v4/nvidia/model.py
+7
@@ -935,6 +935,12 @@ class DeepseekV4Attention(nn.Module):
935
self.indexer = None
936
if self.compress_ratio == 4:
937
# Only C4A uses sparse attention and hence has indexer.
938
+ # aux_stream_list[0] runs indexer.forward() in the wrapper; [2] is
939
+ # free here (outer GEMMs joined) for the inner overlap of
940
+ # wq_b+fused_indexer_q_rope_quant vs compressor.
941
+ indexer_aux_stream = (
942
+ aux_stream_list[2] if aux_stream_list is not None else None
943
+ )
944
self.indexer = DeepseekV4Indexer(
945
vllm_config,
946
config=config,
@@ -945,6 +951,7 @@ class DeepseekV4Attention(nn.Module):
951
topk_indices_buffer=topk_indices_buffer,
952
compress_ratio=self.compress_ratio,
953
prefix=f"{prefix}.indexer",
954
+ aux_stream=indexer_aux_stream,
955
)
956
957
mla_modules = DeepseekV4MLAModules(
vllm/models/deepseek_v4/nvidia/ops/attention.py
+54
-29
@@ -442,33 +442,38 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
442
# downstream reads q on default). Indexer/compressor go on aux for
443
# overlap with default's GEMM + cache write.
444
if self.indexer is not None:
445
- aux_stream = (
446
- self.aux_stream_list[0] if self.aux_stream_list is not None else None
447
- )
445
+ aux_streams = self.aux_stream_list
446
indexer = self.indexer
447
# Local ref so the closure keeps a non-None type for mypy.
448
assert self.compressor is not None
449
compressor = self.compressor
450
453
- def wq_b_kv_insert_and_compress() -> torch.Tensor:
451
+ def wq_b_kv_insert() -> torch.Tensor:
452
q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
453
self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
456
- compressor(kv_score, positions, self.rotary_emb)
454
return q
455
459
- q, _ = maybe_execute_in_parallel(
460
- wq_b_kv_insert_and_compress,
461
- lambda: indexer(
462
- hidden_states,
463
- qr,
464
- indexer_kv_score,
465
- indexer_weights,
466
- positions,
467
- self.indexer_rotary_emb,
468
- ),
456
+ # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
457
+ # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the
458
+ # MLA compressor. Slot [2] is reserved for the indexer's inner
459
+ # overlap. ROCm (aux_streams is None) falls back to sequential.
460
+ q, _ = execute_in_parallel(
461
+ wq_b_kv_insert,
462
+ [
463
+ lambda: indexer(
464
+ hidden_states,
465
+ qr,
466
+ indexer_kv_score,
467
+ indexer_weights,
468
+ positions,
469
+ self.indexer_rotary_emb,
470
+ ),
471
+ lambda: compressor(kv_score, positions, self.rotary_emb),
472
+ ],
473
self.ln_events[0],
470
- self.ln_events[1],
471
- aux_stream,
474
+ [self.ln_events[1], self.ln_events[2]],
475
+ [aux_streams[0], aux_streams[1]] if aux_streams is not None else None,
476
+ enable=aux_streams is not None,
477
)
478
elif self.compressor is not None:
479
# wq_b + kv_insert on default, compressor on aux.
@@ -792,6 +797,7 @@ class DeepseekV4Indexer(nn.Module):
797
topk_indices_buffer: torch.Tensor | None,
798
compress_ratio: int = 1,
799
prefix: str = "",
800
+ aux_stream: torch.cuda.Stream | None = None,
801
):
802
super().__init__()
803
self.vllm_config = vllm_config
@@ -877,6 +883,13 @@ class DeepseekV4Indexer(nn.Module):
883
use_fp4_cache=self.use_fp4_kv,
884
)
885
886
+ # None on ROCm — maybe_execute_in_parallel falls back to sequential.
887
+ self.aux_stream = aux_stream
888
+ self.ln_events: list[torch.cuda.Event] = [
889
+ torch.cuda.Event(),
890
+ torch.cuda.Event(),
891
+ ]
892
+
893
def forward(
894
self,
895
hidden_states: torch.Tensor,
@@ -886,17 +899,29 @@ class DeepseekV4Indexer(nn.Module):
899
positions: torch.Tensor,
900
rotary_emb: nn.Module,
901
) -> torch.Tensor:
889
- # ReplicatedLinear returns (output, bias); bias is None.
890
- q, _ = self.wq_b(qr)
891
- q = q.view(-1, self.n_head, self.head_dim)
892
- k = self.compressor(compressed_kv_score, positions, rotary_emb)
893
- q_quant, weights = fused_indexer_q_rope_quant(
894
- positions,
895
- q,
896
- rotary_emb.cos_sin_cache,
897
- indexer_weights,
898
- self.softmax_scale,
899
- self.n_head**-0.5,
900
- use_fp4=self.use_fp4_kv,
902
+ compressor = self.compressor
903
+
904
+ def wq_b_and_q_quant():
905
+ # ReplicatedLinear returns (output, bias); bias is None.
906
+ q, _ = self.wq_b(qr)
907
+ q = q.view(-1, self.n_head, self.head_dim)
908
+ return fused_indexer_q_rope_quant(
909
+ positions,
910
+ q,
911
+ rotary_emb.cos_sin_cache,
912
+ indexer_weights,
913
+ self.softmax_scale,
914
+ self.n_head**-0.5,
915
+ use_fp4=self.use_fp4_kv,
916
+ )
917
+
918
+ # compressor returns None and writes K to the indexer KV cache; the
919
+ # join orders that write before indexer_op (skip_k_cache_insert=True).
920
+ (q_quant, weights), k = maybe_execute_in_parallel(
921
+ wq_b_and_q_quant,
922
+ lambda: compressor(compressed_kv_score, positions, rotary_emb),
923
+ self.ln_events[0],
924
+ self.ln_events[1],
925
+ self.aux_stream,
926
)
927
return self.indexer_op(hidden_states, q_quant, k, weights)