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

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 4, 13:45 UTC

[DSV4] Enable Multi-stream for Pre-Attn GEMM (#41061)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed Apr 28, 2026 at 12:08 UTC 5aa371dc8e38e053754d89b444abca0a1d63f676

4 files changed +187 -57

vllm/model_executor/layers/deepseek_compressor.py

+2 -7

@@ -14,7 +14,6 @@ from vllm.model_executor.layers.layernorm import RMSNorm

14 from vllm.model_executor.layers.linear import (

15 MergedColumnParallelLinear,

16 )

17 -from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32

17 from vllm.platforms import current_platform

18 from vllm.triton_utils import tl, triton

19 from vllm.v1.attention.backend import (

@@ -271,16 +270,12 @@ class DeepseekCompressor(nn.Module):

270

271 def forward(

272 self,

274 - # [num_tokens, hidden_size]

275 - x: torch.Tensor,

273 + # [num_tokens, 2 * self.coff * self.head_dim]

274 + kv_score: torch.Tensor,

275 # [num_tokens]

276 positions: torch.Tensor,

277 rotary_emb,

278 ) -> None:

280 - num_tokens, _ = x.shape

281 - # bf16 weights/activations but fp32 output for numerical stability of

282 - # the downstream compressor math.

283 - kv_score = cublas_gemm_bf16_bf16_fp32(x, self.fused_wkv_wgate.weight)

279 # Each of shape [num_tokens, coff * self.head_dim]

280 # input bf16, output are fp32

281 kv, score = kv_score.split(

vllm/model_executor/layers/deepseek_v4_attention.py

+111 -38

@@ -4,8 +4,9 @@

4 DeepseekV4 MLA Attention Layer

5 """

7 +from collections.abc import Callable

8 from dataclasses import dataclass

8 -from typing import TYPE_CHECKING, cast

9 +from typing import TYPE_CHECKING, Any, cast

11 import torch

12 import torch.nn as nn

@@ -16,6 +17,7 @@ from vllm.model_executor.layers.linear import (

17 ReplicatedLinear,

18 )

19 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer

20 +from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32

21 from vllm.utils.deep_gemm import fp8_einsum

22 from vllm.utils.torch_utils import direct_register_custom_op

23 from vllm.v1.attention.ops.deepseek_v4_ops import (

@@ -51,7 +53,10 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import (

53 from vllm.model_executor.layers.quantization.utils.quant_utils import (

54 GroupShape,

55 )

54 -from vllm.utils.multi_stream_utils import maybe_execute_in_parallel

56 +from vllm.utils.multi_stream_utils import (

57 + execute_in_parallel,

58 + maybe_execute_in_parallel,

59 +)

60 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata

61 from vllm.v1.attention.backends.mla.flashmla_sparse import (

62 DeepseekV4FlashMLASparseBackend,

@@ -94,7 +99,7 @@ class DeepseekV4MLAModules:

99 indexer: torch.nn.Module | None

100 indexer_rotary_emb: torch.nn.Module

101 topk_indices_buffer: torch.Tensor | None

97 - aux_stream: torch.cuda.Stream | None = None

102 + aux_stream_list: list[torch.cuda.Stream] | None = None

103

104

105 # --8<-- [start:multi_head_latent_attention]

@@ -217,8 +222,11 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

222 + 1 # 1B pad

223 )

224

220 - self.aux_stream = mla_modules.aux_stream

221 - self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]

225 + self.aux_stream_list = mla_modules.aux_stream_list

226 + # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;

227 + # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins

228 + # before post-GEMM starts.

229 + self.ln_events = [torch.cuda.Event() for _ in range(4)]

230

231 assert cache_config is not None, "DeepseekV4 attention requires cache_config"

232 self.swa_cache_layer = DeepseekV4SWACache(

@@ -277,9 +285,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

285 hidden_states: torch.Tensor,

286 llama_4_scaling: torch.Tensor | None = None,

287 ) -> torch.Tensor:

280 - qr_kv, _ = self.fused_wqa_wkv(hidden_states)

281 - qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)

282 -

288 # Pre-allocate attention output with FlashMLA-padded head count.

289 # The op writes into `o_padded`; we slice to n_local_heads after.

290 num_tokens = hidden_states.shape[0]

@@ -292,8 +297,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

297 # Attention (inside custom op for torch.compile boundary)

298 torch.ops.vllm.deepseek_v4_attention(

299 hidden_states,

295 - qr,

296 - kv,

300 positions,

301 o_padded,

302 self.layer_name,

@@ -332,17 +335,71 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

335

336 return self.wo_b(z.flatten(1))

337

338 + def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:

339 + assert self.aux_stream_list is not None

340 + assert len(self.aux_stream_list) >= 3

341 +

342 + # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs

343 + # on aux streams 0..2 when their owning module exists. ln_events[0]

344 + # is the fan-out start event; ln_events[1..3] are per-aux done events.

345 + aux_fns: list[Callable[[], Any] | None] = [None, None, None]

346 +

347 + if self.compressor is not None:

348 + # Local ref so the closure keeps a non-None type for mypy.

349 + compressor = self.compressor

350 +

351 + def compressor_kv_score() -> torch.Tensor:

352 + return cublas_gemm_bf16_bf16_fp32(

353 + hidden_states, compressor.fused_wkv_wgate.weight

354 + )

355 +

356 + aux_fns[0] = compressor_kv_score

357 +

358 + if self.indexer is not None:

359 + indexer = self.indexer

360 +

361 + def indexer_weights_proj() -> torch.Tensor:

362 + # ReplicatedLinear returns (output, bias); bias is None.

363 + weights, _ = indexer.weights_proj(hidden_states)

364 + return weights

365 +

366 + def indexer_compressor_kv_score() -> torch.Tensor:

367 + return cublas_gemm_bf16_bf16_fp32(

368 + hidden_states, indexer.compressor.fused_wkv_wgate.weight

369 + )

370 +

371 + aux_fns[1] = indexer_weights_proj

372 + aux_fns[2] = indexer_compressor_kv_score

373 +

374 + def fused_wqa_wkv() -> torch.Tensor:

375 + # MergedColumnParallelLinear returns (output, bias); bias is None.

376 + qr_kv, _ = self.fused_wqa_wkv(hidden_states)

377 + return qr_kv

378 +

379 + qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel(

380 + fused_wqa_wkv,

381 + aux_fns,

382 + self.ln_events[0],

383 + self.ln_events[1:4],

384 + self.aux_stream_list[:3],

385 + )

386 +

387 + return qr_kv, kv_score, indexer_kv_score, indexer_weights

388 +

389 def attention_impl(

390 self,

391 hidden_states: torch.Tensor,

338 - qr: torch.Tensor,

339 - kv: torch.Tensor,

392 positions: torch.Tensor,

393 out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place

394 ) -> None:

395 forward_context = get_forward_context()

396 attn_metadata = forward_context.attn_metadata

397

398 + qr_kv, kv_score, indexer_kv_score, indexer_weights = (

399 + self.attn_gemm_parallel_execute(hidden_states)

400 + )

401 +

402 + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)

403 qr, kv = fused_q_kv_rmsnorm(

404 qr,

405 kv,

@@ -350,42 +407,60 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

407 self.kv_norm.weight.data,

408 self.eps,

409 )

353 - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)

410

355 - # Overlap kv_insert with whichever of indexer/compressor is present.

356 - # Indexer implies compressor; when both exist, compressor rides on the

357 - # aux stream alongside kv_insert so the heavy indexer owns default.

411 + # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride

412 + # on the default stream so q stays on its consumer stream (mla_attn

413 + # downstream reads q on default). Indexer/compressor go on aux for

414 + # overlap with default's GEMM + cache write.

415 if self.indexer is not None:

416 + assert self.aux_stream_list is not None

417 + aux_stream = self.aux_stream_list[0]

418 indexer = self.indexer

419 # Local ref so the closure keeps a non-None type for mypy.

420 assert self.compressor is not None

421 compressor = self.compressor

422

364 - def kv_insert_and_compress() -> None:

423 + def wq_b_kv_insert_and_compress() -> torch.Tensor:

424 + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)

425 self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)

366 - compressor(hidden_states, positions, self.rotary_emb)

367 -

368 - maybe_execute_in_parallel(

369 - lambda: indexer(hidden_states, qr, positions, self.indexer_rotary_emb),

370 - kv_insert_and_compress,

426 + compressor(kv_score, positions, self.rotary_emb)

427 + return q

428 +

429 + q, _ = maybe_execute_in_parallel(

430 + wq_b_kv_insert_and_compress,

431 + lambda: indexer(

432 + hidden_states,

433 + qr,

434 + indexer_kv_score,

435 + indexer_weights,

436 + positions,

437 + self.indexer_rotary_emb,

438 + ),

439 self.ln_events[0],

440 self.ln_events[1],

373 - self.aux_stream,

441 + aux_stream,

442 )

443 elif self.compressor is not None:

376 - # Compressor on default, kv_insert on aux.

444 + # wq_b + kv_insert on default, compressor on aux.

445 + assert self.aux_stream_list is not None

446 + aux_stream = self.aux_stream_list[0]

447 compressor = self.compressor

378 - maybe_execute_in_parallel(

379 - lambda: compressor(hidden_states, positions, self.rotary_emb),

380 - lambda: self._fused_qnorm_rope_kv_insert(

381 - q, kv, positions, attn_metadata

382 - ),

448 +

449 + def wq_b_kv_insert() -> torch.Tensor:

450 + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)

451 + self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)

452 + return q

453 +

454 + q, _ = maybe_execute_in_parallel(

455 + wq_b_kv_insert,

456 + lambda: compressor(kv_score, positions, self.rotary_emb),

457 self.ln_events[0],

458 self.ln_events[1],

385 - self.aux_stream,

459 + aux_stream,

460 )

461 else:

462 # SWA-only layer: no compressor, no overlap.

463 + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)

464 self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)

465

466 # Handle dummy run (no metadata).

@@ -455,21 +530,17 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

530

531 def deepseek_v4_attention(

532 hidden_states: torch.Tensor,

458 - qr: torch.Tensor,

459 - kv: torch.Tensor,

533 positions: torch.Tensor,

534 out: torch.Tensor,

535 layer_name: str,

536 ) -> None:

537 forward_context: ForwardContext = get_forward_context()

538 self = forward_context.no_compile_layers[layer_name]

466 - self.attention_impl(hidden_states, qr, kv, positions, out)

539 + self.attention_impl(hidden_states, positions, out)

540

541

542 def deepseek_v4_attention_fake(

543 hidden_states: torch.Tensor,

471 - qr: torch.Tensor,

472 - kv: torch.Tensor,

544 positions: torch.Tensor,

545 out: torch.Tensor,

546 layer_name: str,

@@ -1057,18 +1128,20 @@ class DeepseekV4Indexer(nn.Module):

1128 self,

1129 hidden_states: torch.Tensor,

1130 qr: torch.Tensor,

1131 + compressed_kv_score: torch.Tensor,

1132 + indexer_weights: torch.Tensor,

1133 positions: torch.Tensor,

1134 rotary_emb: nn.Module,

1135 ) -> torch.Tensor:

1136 + # ReplicatedLinear returns (output, bias); bias is None.

1137 q, _ = self.wq_b(qr)

1138 q = q.view(-1, self.n_head, self.head_dim)

1065 - k = self.compressor(hidden_states, positions, rotary_emb)

1066 - weights, _ = self.weights_proj(hidden_states)

1139 + k = self.compressor(compressed_kv_score, positions, rotary_emb)

1140 q_quant, weights = fused_indexer_q_rope_quant(

1141 positions,

1142 q,

1143 rotary_emb.cos_sin_cache,

1071 - weights,

1144 + indexer_weights,

1145 self.softmax_scale,

1146 self.n_head**-0.5,

1147 use_fp4=self.use_fp4_kv,

vllm/model_executor/models/deepseek_v4.py

+10 -12

@@ -54,7 +54,6 @@ from vllm.model_executor.utils import set_weight_attrs

54 from vllm.platforms import current_platform

55 from vllm.sequence import IntermediateTensors

56 from vllm.triton_utils import tl, triton

57 -from vllm.utils.multi_stream_utils import AuxStreamType

57 from vllm.utils.torch_utils import direct_register_custom_op

59 from .utils import (

@@ -926,7 +925,7 @@ class DeepseekV4Attention(nn.Module):

925 vllm_config: VllmConfig,

926 prefix: str,

927 topk_indices_buffer: torch.Tensor | None = None,

929 - aux_stream: torch.cuda.Stream | None = None,

928 + aux_stream_list: list[torch.cuda.Stream] | None = None,

929 ):

930 super().__init__()

931 config = vllm_config.model_config.hf_config

@@ -1059,7 +1058,7 @@ class DeepseekV4Attention(nn.Module):

1058 indexer=self.indexer,

1059 indexer_rotary_emb=self.rotary_emb,

1060 topk_indices_buffer=topk_indices_buffer,

1062 - aux_stream=aux_stream,

1061 + aux_stream_list=aux_stream_list,

1062 )

1063 self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper(

1064 hidden_size=self.hidden_size,

@@ -1095,7 +1094,7 @@ class DeepseekV4DecoderLayer(nn.Module):

1094 vllm_config,

1095 prefix,

1096 topk_indices_buffer: torch.Tensor | None = None,

1098 - aux_stream_dict: dict[AuxStreamType, torch.cuda.Stream] | None = None,

1097 + aux_stream_list: list[torch.cuda.Stream] | None = None,

1098 ):

1099 super().__init__()

1100 config = vllm_config.model_config.hf_config

@@ -1106,9 +1105,7 @@ class DeepseekV4DecoderLayer(nn.Module):

1105 vllm_config,

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

1107 topk_indices_buffer=topk_indices_buffer,

1109 - aux_stream=aux_stream_dict.get(AuxStreamType.Attention)

1110 - if aux_stream_dict is not None

1111 - else None,

1108 + aux_stream_list=aux_stream_list,

1109 )

1110 self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")

1111

@@ -1236,10 +1233,11 @@ class DeepseekV4Model(nn.Module):

1233 self.hc_dim = self.hc_mult * config.hidden_size

1234 self.rms_norm_eps = config.rms_norm_eps

1235

1239 - aux_stream_list = [torch.cuda.Stream() for _ in range(1)]

1240 - self.aux_stream_dict = {

1241 - AuxStreamType.Attention: aux_stream_list[0],

1242 - }

1236 + # Three aux streams: one per non-default input GEMM in

1237 + # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute

1238 + # (compressor kv_score, indexer.weights_proj, indexer.compressor

1239 + # kv_score). fused_wqa_wkv stays on the default stream.

1240 + aux_stream_list = [torch.cuda.Stream() for _ in range(3)]

1241

1242 self.device = current_platform.device_type

1243 # Reserved topk indices buffer for all Indexer layers to reuse.

@@ -1263,7 +1261,7 @@ class DeepseekV4Model(nn.Module):

1261 vllm_config,

1262 prefix=prefix,

1263 topk_indices_buffer=self.topk_indices_buffer,

1266 - aux_stream_dict=self.aux_stream_dict,

1264 + aux_stream_list=aux_stream_list,

1265 ),

1266 prefix=f"{prefix}.layers",

1267 )

vllm/utils/multi_stream_utils.py

+64

@@ -56,3 +56,67 @@ def maybe_execute_in_parallel(

56 result0 = fn0()

57 result1 = fn1()

58 return (result0, result1)

59 +

60 +

61 +def execute_in_parallel(

62 + default_fn: Callable[[], Any],

63 + aux_fns: list[Callable[[], Any] | None],

64 + start_event: torch.cuda.Event,

65 + done_events: list[torch.cuda.Event],

66 + aux_streams: list[torch.cuda.Stream] | None = None,

67 +) -> tuple[Any, list[Any]]:

68 + """Run default_fn on the current stream and aux_fns concurrently on

69 + aux_streams.

70 +

71 + Generalizes maybe_execute_in_parallel to N aux callables. Slots where

72 + aux_fns[i] is None are skipped (no stream switch, no event record); their

73 + corresponding entry in the returned aux_results list is None.

74 +

75 + start_event fans out from the current stream to every launched aux stream;

76 + done_events[i] is recorded after aux_fns[i] so the current stream joins

77 + before returning. When aux_streams is None, all aux_fns run sequentially

78 + on the current stream.

79 +

80 + Args:

81 + default_fn: Callable for the default (current) stream.

82 + aux_fns: Per-aux callables; entries may be None to skip.

83 + start_event: CUDA event recorded on the current stream before

84 + default_fn so each launched aux stream can wait on it.

85 + done_events: One CUDA event per aux slot, recorded after the

86 + corresponding aux_fn. Length must match aux_fns.

87 + aux_streams: Per-aux CUDA streams. Length must match aux_fns.

88 + Multi-stream is disabled when None.

89 +

90 + Returns:

91 + Tuple of (default_result, aux_results) where aux_results[i] is the

92 + result of aux_fns[i] (or None when skipped).

93 + """

94 + aux_results: list[Any]

95 + if aux_streams is None:

96 + default_result = default_fn()

97 + aux_results = [fn() if fn is not None else None for fn in aux_fns]

98 + return default_result, aux_results

99 +

100 + assert len(aux_fns) == len(aux_streams) == len(done_events), (

101 + "aux_fns, aux_streams, and done_events must be the same length"

102 + )

103 +

104 + aux_results = [None] * len(aux_fns)

105 + pending: list[torch.cuda.Event] = []

106 +

107 + start_event.record()

108 + for i, fn in enumerate(aux_fns):

109 + if fn is None:

110 + continue

111 + with torch.cuda.stream(aux_streams[i]):

112 + start_event.wait()

113 + aux_results[i] = fn()

114 + done_events[i].record()

115 + pending.append(done_events[i])

116 +

117 + default_result = default_fn()

118 +

119 + for ev in pending:

120 + ev.wait()

121 +

122 + return default_result, aux_results