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

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

[DSV4] Remove unncessary classes & functions (#44246)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

Woosuk Kwon committed Jun 1, 2026 at 14:43 UTC 8c3cc98cffd31b910c41b11076e8c175fc6dabe9

3 files changed +62 -124

vllm/models/deepseek_v4/amd/model.py

+14 -18

@@ -48,8 +48,7 @@ from vllm.model_executor.models.utils import (

48 )

49 from vllm.models.deepseek_v4.attention import (

50 DeepseekV4Indexer,

51 - DeepseekV4MLAModules,

52 - DeepseekV4MultiHeadLatentAttentionWrapper,

51 + DeepseekV4MLA,

52 )

53 from vllm.platforms import current_platform

54 from vllm.sequence import IntermediateTensors

@@ -314,7 +313,7 @@ class DeepseekV4Attention(nn.Module):

313

314 self.rope_parameters = config.rope_scaling

315

317 - # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it)

316 + # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it)

317 rope_parameters = config.rope_parameters

318 rope_parameters["rope_theta"] = (

319 config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta

@@ -351,7 +350,17 @@ class DeepseekV4Attention(nn.Module):

350 prefix=f"{prefix}.indexer",

351 )

352

354 - mla_modules = DeepseekV4MLAModules(

353 + self.mla_attn = DeepseekV4MLA(

354 + hidden_size=self.hidden_size,

355 + num_heads=self.n_local_heads,

356 + head_dim=self.head_dim,

357 + scale=self.softmax_scale,

358 + qk_nope_head_dim=self.nope_head_dim,

359 + qk_rope_head_dim=self.rope_head_dim,

360 + v_head_dim=self.head_dim,

361 + q_lora_rank=self.q_lora_rank,

362 + kv_lora_rank=self.head_dim,

363 + o_lora_rank=self.o_lora_rank,

364 vllm_config=vllm_config,

365 fused_wqa_wkv=self.fused_wqa_wkv,

366 q_norm=self.q_norm,

@@ -365,19 +374,6 @@ class DeepseekV4Attention(nn.Module):

374 indexer_rotary_emb=self.rotary_emb,

375 topk_indices_buffer=topk_indices_buffer,

376 aux_stream_list=aux_stream_list,

368 - )

369 - self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper(

370 - hidden_size=self.hidden_size,

371 - num_heads=self.n_local_heads,

372 - head_dim=self.head_dim,

373 - scale=self.softmax_scale,

374 - qk_nope_head_dim=self.nope_head_dim,

375 - qk_rope_head_dim=self.rope_head_dim,

376 - v_head_dim=self.head_dim,

377 - q_lora_rank=self.q_lora_rank,

378 - kv_lora_rank=self.head_dim,

379 - o_lora_rank=self.o_lora_rank,

380 - mla_modules=mla_modules,

377 window_size=self.window_size,

378 compress_ratio=self.compress_ratio,

379 cache_config=vllm_config.cache_config,

@@ -618,7 +614,7 @@ class DeepseekV4Model(nn.Module):

614 self.rms_norm_eps = config.rms_norm_eps

615

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

621 - # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute

617 + # DeepseekV4MLA.attn_gemm_parallel_execute

618 # (compressor kv_score, indexer.weights_proj, indexer.compressor

619 # kv_score). fused_wqa_wkv stays on the default stream.

620 # Disable them on ROCm because of hang issues.

vllm/models/deepseek_v4/attention.py

+34 -88

@@ -5,7 +5,6 @@ DeepseekV4 MLA Attention Layer

5 """

7 from collections.abc import Callable

8 -from dataclasses import dataclass

8 from typing import TYPE_CHECKING, Any, cast

10 import torch

@@ -38,9 +37,8 @@ from vllm.config import (

37 get_current_vllm_config,

38 )

39 from vllm.distributed import get_tensor_model_parallel_world_size

41 -from vllm.forward_context import ForwardContext, get_forward_context

40 +from vllm.forward_context import get_forward_context

41 from vllm.logger import init_logger

43 -from vllm.model_executor.custom_op import PluggableLayer

42 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

43 from vllm.model_executor.layers.layernorm import RMSNorm

44 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -90,46 +88,7 @@ def _select_v4_sparse_impl() -> "type[DeepseekV4SparseMLAAttentionImpl]":

88 return DeepseekV4FlashMLASparseImpl

93 -@dataclass

94 -class DeepseekV4MLAModules:

95 - """Modules used in DeepseekV4 MLA."""

96 -

97 - vllm_config: VllmConfig

98 - fused_wqa_wkv: torch.nn.Module

99 - q_norm: torch.nn.Module

100 - wq_b: torch.nn.Module

101 - kv_norm: torch.nn.Module

102 - wo_a: torch.nn.Module

103 - wo_b: torch.nn.Module

104 - attn_sink: torch.nn.Module

105 - rotary_emb: torch.nn.Module

106 - indexer: torch.nn.Module | None

107 - indexer_rotary_emb: torch.nn.Module

108 - topk_indices_buffer: torch.Tensor | None

109 - aux_stream_list: list[torch.cuda.Stream] | None = None

110 -

111 -

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

113 -@PluggableLayer.register("deepseek_v4_multi_head_latent_attention")

114 -class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

115 - """Pluggable MLA layer which allows OOT backends to add

116 - custom implementations of the outer MLA layer (including rope & o_proj).

117 - Note that currently oot platforms can still use CustomOp.register_oot to

118 - replace MLA layer entirely, although we use PluggableLayer to register

119 - this layer now.

120 -

121 - This class takes positions and hidden_states as input.

122 - The input tensors can either contain prefill tokens or decode tokens.

123 - The class does the following:

124 -

125 - 1. MLA Preprocess.

126 - 2. Perform multi-head attention to prefill tokens and

127 - multi-query attention to decode tokens separately.

128 - 3. Return the output tensor.

129 - """

130 -

131 - # --8<-- [end:multi_head_latent_attention]

132 -

91 +class DeepseekV4MLA(nn.Module):

92 def __init__(

93 self,

94 hidden_size: int,

@@ -142,7 +101,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

101 q_lora_rank: int | None,

102 kv_lora_rank: int,

103 o_lora_rank: int | None,

145 - mla_modules: DeepseekV4MLAModules,

104 + vllm_config: VllmConfig,

105 + fused_wqa_wkv: torch.nn.Module,

106 + q_norm: torch.nn.Module,

107 + wq_b: torch.nn.Module,

108 + kv_norm: torch.nn.Module,

109 + wo_a: torch.nn.Module,

110 + wo_b: torch.nn.Module,

111 + attn_sink: torch.nn.Module,

112 + rotary_emb: torch.nn.Module,

113 + indexer: torch.nn.Module | None,

114 + indexer_rotary_emb: torch.nn.Module,

115 + topk_indices_buffer: torch.Tensor | None,

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

117 window_size: int,

118 compress_ratio: int | None,

119 cache_config: CacheConfig | None = None,

@@ -162,7 +133,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

133 self.prefix = prefix

134

135 # Extract config from vllm_config

165 - config = mla_modules.vllm_config.model_config.hf_config

136 + config = vllm_config.model_config.hf_config

137 tp_size = get_tensor_model_parallel_world_size()

138

139 # DeepseekV4-specific attributes (num_heads is already TP-adjusted)

@@ -173,12 +144,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

144 self.o_lora_rank = config.o_lora_rank

145

146 # Store projection modules

176 - self.fused_wqa_wkv = mla_modules.fused_wqa_wkv

177 - self.q_norm = mla_modules.q_norm

178 - self.wq_b = mla_modules.wq_b

147 + self.fused_wqa_wkv = fused_wqa_wkv

148 + self.q_norm = q_norm

149 + self.wq_b = wq_b

150

180 - self.kv_norm = mla_modules.kv_norm

181 - self.wo_a = mla_modules.wo_a

151 + self.kv_norm = kv_norm

152 + self.wo_a = wo_a

153

154 self._wo_a_act_quant = QuantFP8(

155 static=False,

@@ -188,7 +159,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

159 # Bypass packed-for-deepgemm path — we need FP32 scales (not packed

160 # INT32) so fp8_einsum can handle layout transform internally.

161 self._wo_a_act_quant.use_deep_gemm_supported = False

191 - self.wo_b = mla_modules.wo_b

162 + self.wo_b = wo_b

163

164 # Pick fp8_einsum recipe based on GPU arch:

165 # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128

@@ -198,11 +169,11 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

169 self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)

170 self._tma_aligned_scales = cap.major >= 10

171

201 - self.rotary_emb = mla_modules.rotary_emb

202 - self.indexer_rotary_emb = mla_modules.indexer_rotary_emb

203 - self.topk_indices_buffer = mla_modules.topk_indices_buffer

172 + self.rotary_emb = rotary_emb

173 + self.indexer_rotary_emb = indexer_rotary_emb

174 + self.topk_indices_buffer = topk_indices_buffer

175

205 - self.indexer = mla_modules.indexer

176 + self.indexer = indexer

177

178 # Per-head RMS normalization for Q (no learnable weights)

179 self.q_head_norm = RMSNorm(head_dim, eps=self.eps, has_weight=False)

@@ -216,7 +187,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

187 )

188

189 # Will be None on ROCm for now.

219 - self.aux_stream_list = mla_modules.aux_stream_list

190 + self.aux_stream_list = aux_stream_list

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

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

193 # before post-GEMM starts.

@@ -243,7 +214,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

214 window_size=self.window_size,

215 head_bytes=head_bytes,

216 swa_cache_layer=self.swa_cache_layer,

246 - attn_sink=mla_modules.attn_sink, # already padded with -inf

217 + attn_sink=attn_sink, # already padded with -inf

218 cache_config=cache_config,

219 quant_config=quant_config,

220 prefix=prefix,

@@ -253,21 +224,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

224 # Mirror the inner layer's padded head count (single source of truth).

225 self.padded_heads = self.mla_attn.padded_heads

226

256 - # Register this layer in the compilation config's static forward context

257 - # This allows the custom op to retrieve the layer during execution

258 - compilation_config = mla_modules.vllm_config.compilation_config

259 - # HACK

260 - self.layer_name = prefix + ".deepseek_v4_multi_head_latent_attention"

261 - if self.layer_name in compilation_config.static_forward_context:

262 - raise ValueError(f"Duplicate layer name: {self.layer_name}")

263 - compilation_config.static_forward_context[self.layer_name] = self

264 -

227 # Create the compressor for layers with compress_ratio > 1; after

228 # creating the DeepseekV4MLAAttention layer to get its cache.

229 self.compressor = None

230 if self.compress_ratio > 1:

231 self.compressor = DeepseekCompressor(

270 - vllm_config=mla_modules.vllm_config,

232 + vllm_config=vllm_config,

233 compress_ratio=self.compress_ratio,

234 hidden_size=self.hidden_size,

235 head_dim=self.head_dim,

@@ -291,15 +253,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

253 device=hidden_states.device,

254 )

255

294 - # @eager_break_during_capture: this is where the breakable

295 - # cudagraph capture breaks (the attention op runs eagerly between

296 - # captured graph segments).

297 - deepseek_v4_attention(

298 - hidden_states,

299 - positions,

300 - o_padded,

301 - self.layer_name,

302 - )

256 + # attention_impl is wrapped with @eager_break_during_capture: this is

257 + # where the breakable cudagraph capture breaks (the attention op runs

258 + # eagerly between captured graph segments).

259 + self.attention_impl(hidden_states, positions, o_padded)

260 o = o_padded[:, : self.n_local_heads, :]

261

262 # Keep ROCm on the BF16 reference wo_a path util kernel ready.

@@ -405,6 +362,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

362

363 return qr_kv, kv_score, indexer_kv_score, indexer_weights

364

365 + @eager_break_during_capture

366 def attention_impl(

367 self,

368 hidden_states: torch.Tensor,

@@ -541,18 +499,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

499 )

500

501

544 -@eager_break_during_capture

545 -def deepseek_v4_attention(

546 - hidden_states: torch.Tensor,

547 - positions: torch.Tensor,

548 - out: torch.Tensor,

549 - layer_name: str,

550 -) -> None:

551 - forward_context: ForwardContext = get_forward_context()

552 - self = forward_context.no_compile_layers[layer_name]

553 - self.attention_impl(hidden_states, positions, out)

554 -

555 -

502 class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):

503 def __init__(

504 self,

vllm/models/deepseek_v4/nvidia/model.py

+14 -18

@@ -54,8 +54,7 @@ from vllm.model_executor.models.utils import (

54 from vllm.model_executor.utils import set_weight_attrs

55 from vllm.models.deepseek_v4.attention import (

56 DeepseekV4Indexer,

57 - DeepseekV4MLAModules,

58 - DeepseekV4MultiHeadLatentAttentionWrapper,

57 + DeepseekV4MLA,

58 )

59 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs

60 from vllm.sequence import IntermediateTensors

@@ -697,7 +696,7 @@ class DeepseekV4Attention(nn.Module):

696

697 self.rope_parameters = config.rope_scaling

698

700 - # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it)

699 + # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it)

700 rope_parameters = config.rope_parameters

701 rope_parameters["rope_theta"] = (

702 config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta

@@ -741,7 +740,17 @@ class DeepseekV4Attention(nn.Module):

740 aux_stream=indexer_aux_stream,

741 )

742

744 - mla_modules = DeepseekV4MLAModules(

743 + self.mla_attn = DeepseekV4MLA(

744 + hidden_size=self.hidden_size,

745 + num_heads=self.n_local_heads,

746 + head_dim=self.head_dim,

747 + scale=self.softmax_scale,

748 + qk_nope_head_dim=self.nope_head_dim,

749 + qk_rope_head_dim=self.rope_head_dim,

750 + v_head_dim=self.head_dim,

751 + q_lora_rank=self.q_lora_rank,

752 + kv_lora_rank=self.head_dim,

753 + o_lora_rank=self.o_lora_rank,

754 vllm_config=vllm_config,

755 fused_wqa_wkv=self.fused_wqa_wkv,

756 q_norm=self.q_norm,

@@ -755,19 +764,6 @@ class DeepseekV4Attention(nn.Module):

764 indexer_rotary_emb=self.rotary_emb,

765 topk_indices_buffer=topk_indices_buffer,

766 aux_stream_list=aux_stream_list,

758 - )

759 - self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper(

760 - hidden_size=self.hidden_size,

761 - num_heads=self.n_local_heads,

762 - head_dim=self.head_dim,

763 - scale=self.softmax_scale,

764 - qk_nope_head_dim=self.nope_head_dim,

765 - qk_rope_head_dim=self.rope_head_dim,

766 - v_head_dim=self.head_dim,

767 - q_lora_rank=self.q_lora_rank,

768 - kv_lora_rank=self.head_dim,

769 - o_lora_rank=self.o_lora_rank,

770 - mla_modules=mla_modules,

767 window_size=self.window_size,

768 compress_ratio=self.compress_ratio,

769 cache_config=vllm_config.cache_config,

@@ -955,7 +951,7 @@ class DeepseekV4Model(nn.Module):

951 self.rms_norm_eps = config.rms_norm_eps

952

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

958 - # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute

954 + # DeepseekV4MLA.attn_gemm_parallel_execute

955 # (compressor kv_score, indexer.weights_proj, indexer.compressor

956 # kv_score). fused_wqa_wkv stays on the default stream.

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