diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 5a7832d945..583378fdb4 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -120,6 +120,7 @@ steps: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - vllm/models/deepseek_v4/nvidia/ + - vllm/utils/deep_gemm.py - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - tests/models/test_deepseek_v4_mega_moe.py commands: diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index fafb37d97b..08b1a24c85 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -132,6 +132,7 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams) draft = SimpleNamespace( + use_mega_moe=False, use_sequence_parallel=False, hc_mult=hc_mult, layers=[make_layer(mix) for mix in mixes], @@ -452,7 +453,7 @@ def test_deepseek_v41_decoder_mixes_match_torch( device=DEVICE, dtype=torch.bfloat16 ) decoder.attn = lambda positions, x, _: x * 0.5 - decoder.ffn = lambda x, input_ids: x * 0.25 + decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25 with torch.device(DEVICE): decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 @@ -536,7 +537,7 @@ def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_conf device=DEVICE, dtype=torch.bfloat16 ) decoder.attn = lambda positions, x, _: x * 0.5 - decoder.ffn = lambda x, input_ids: x * 0.25 + decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25 with torch.device(DEVICE): decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 67029b50a9..ad029160c5 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -15,6 +15,7 @@ from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4MegaMoEExperts, DeepseekV4MoE, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs @@ -22,6 +23,7 @@ from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config +from vllm.utils.torch_utils import set_default_torch_dtype pytestmark = pytest.mark.skipif( not current_platform.is_cuda(), @@ -33,6 +35,7 @@ pytestmark = pytest.mark.skipif( def v41_moe_config(dist_init): return SimpleNamespace( model_config=SimpleNamespace( + dtype=torch.bfloat16, hf_config=DeepseekV41Config( text_config=dict( hidden_size=128, @@ -63,28 +66,52 @@ def v41_moe_config(dist_init): ) +@pytest.mark.parametrize("use_cudagraph", [False, True]) +@pytest.mark.parametrize("above_threshold", [False, True]) @pytest.mark.parametrize("vision", [False, True]) -@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 8, 2), (2, 4, 3)]) +@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 384, 6), (2, 128, 3)]) def test_deepseek_v41_moe_routes_without_hash_table( - v41_moe_config, monkeypatch, vision, layer_id, num_experts, top_k + v41_moe_config, + monkeypatch, + vision, + layer_id, + num_experts, + top_k, + above_threshold, + use_cudagraph, ): """Main and draft layers select experts by score and preserve image routing.""" + if not current_platform.is_device_capability_family(100): + pytest.skip("DeepGEMM Mega Gate requires SM100") + config = v41_moe_config.model_config.hf_config + config.hidden_size = 256 + config.n_routed_experts = 384 + config.num_experts_per_tok = 6 + config.dspark_n_routed_experts = 128 + num_tokens = (1 if num_experts == 128 else 16) + int(above_threshold) config.vision_n_layers = int(vision) - with torch.device("cuda"): + with ( + set_default_torch_dtype(v41_moe_config.model_config.dtype), + torch.device("cuda"), + ): moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn") - hidden_states = torch.randn(4, config.hidden_size) + hidden_states = torch.randn( + num_tokens, config.hidden_size, dtype=torch.bfloat16 + ) input_ids = ( - torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257]) + torch.tensor( + [42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257] + ).repeat((num_tokens + 3) // 4)[:num_tokens] if vision else None ) assert moe.gate.tid2eid is None assert moe.gate.weight.shape == (num_experts, config.hidden_size) assert isinstance(moe.experts, DeepseekV4MegaMoEExperts) - assert moe.experts.w13_weight.shape == (num_experts, 256, 64) + assert moe.experts.w13_weight.shape == (num_experts, 256, 128) assert moe.experts.top_k == top_k - assert (config.n_routed_experts, config.num_experts_per_tok) == (8, 2) + assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6) with torch.no_grad(): moe.gate.weight.normal_(std=0.01) @@ -93,7 +120,7 @@ def test_deepseek_v41_moe_routes_without_hash_table( moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias) scores = torch.nn.functional.softplus( - torch.nn.functional.linear(hidden_states, moe.gate.weight) + torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32) ).sqrt() bias = moe.gate.e_score_correction_bias if vision: @@ -105,15 +132,49 @@ def test_deepseek_v41_moe_routes_without_hash_table( dim=-1, keepdim=True ) + routed = {} + def check_routing(x, weights, ids, *, activation_clamp): - torch.testing.assert_close(ids, expected_ids) - torch.testing.assert_close(weights, expected_weights) + routed["ids"], routed["weights"] = ids, weights assert activation_clamp == config.swiglu_limit return x.clone() monkeypatch.setattr(moe.experts, "forward", check_routing) monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x) - torch.testing.assert_close(moe(hidden_states, input_ids), 3 * hidden_states) + routing_input_ids = ( + input_ids + if input_ids is not None + else torch.zeros(num_tokens, dtype=torch.int64, device="cuda") + ) + metadata = prepare_mega_gate_routing_metadata( + routing_input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID if vision else None, + ) + if use_cudagraph: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + moe(hidden_states, input_ids, metadata) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = moe(hidden_states, input_ids, metadata) + graph.replay() + torch.accelerator.synchronize() + else: + output = moe(hidden_states, input_ids, metadata) + + ids, weights = routed["ids"], routed["weights"] + torch.testing.assert_close(ids, expected_ids) + torch.testing.assert_close( + weights, + expected_weights, + rtol=1e-3, + atol=1e-4, + ) + torch.testing.assert_close(output, 3 * hidden_states) @pytest.mark.parametrize( @@ -156,6 +217,67 @@ def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config): moe(torch.zeros(1, 128)) +@pytest.mark.parametrize("num_tokens", [16, 17]) +def test_deepseek_v4_mega_gate_hash_routing_correctness( + v41_moe_config, monkeypatch, num_tokens +): + if not current_platform.is_device_capability_family(100): + pytest.skip("DeepGEMM Mega Gate requires SM100") + + config = v41_moe_config.model_config.hf_config + config.hidden_size = 256 + config.num_hash_layers = 1 + config.vocab_size = 32 + with ( + set_default_torch_dtype(v41_moe_config.model_config.dtype), + torch.device("cuda"), + ): + moe = DeepseekV4MoE( + v41_moe_config, + prefix="model.layers.0.ffn", + num_hash_layers=config.num_hash_layers, + ) + hidden_states = torch.randn( + num_tokens, config.hidden_size, dtype=torch.bfloat16 + ) + input_ids = torch.arange(num_tokens) + + token_ids = torch.arange(config.vocab_size, device="cuda") + fixed_ids = torch.stack( + ( + token_ids % config.n_routed_experts, + (token_ids + 3) % config.n_routed_experts, + ), + dim=1, + ) + with torch.no_grad(): + moe.gate.weight.normal_(std=0.01) + moe.gate.tid2eid.copy_(fixed_ids) + + expected_ids = fixed_ids[input_ids] + scores = torch.nn.functional.softplus( + torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32) + ).sqrt() + expected_weights = scores.gather(1, expected_ids) + expected_weights *= config.routed_scaling_factor / expected_weights.sum( + dim=-1, keepdim=True + ) + + def check_routing(x, weights, ids, *, activation_clamp): + torch.testing.assert_close(ids, expected_ids) + torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4) + return x.clone() + + monkeypatch.setattr(moe.experts, "forward", check_routing) + monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like) + metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=True, + image_sentinel_base_id=None, + ) + torch.testing.assert_close(moe(hidden_states, input_ids, metadata), hidden_states) + + def test_deepseek_v4_mega_moe_expert_mapping(): mapping = make_deepseek_v4_expert_params_mapping(2) @@ -513,6 +635,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( monkeypatch, fused ): class FakeGate(torch.nn.Module): + weight = torch.empty(2, 128) tid2eid = None e_score_correction_bias = None @@ -760,7 +883,8 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", ) -def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(): +@pytest.mark.parametrize("nonfinite_padding", [False, True]) +def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(nonfinite_padding): from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 device = torch.device("cuda") @@ -796,6 +920,8 @@ def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(): [False, True, False, False, True, False, True], device=device, ) + if nonfinite_padding: + topk_weights[is_padding] = float("nan") ref_x, ref_x_sf = per_token_cast_to_fp8( hidden_states, diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 52783889d6..987909702c 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -17,6 +17,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -51,12 +52,14 @@ from vllm.models.common.ops.sequence_parallel import ( sp_padding_mask, sp_shard, ) +from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from .model import ( DeepseekV4DecoderLayer, DeepseekV4Model, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) logger = init_logger(__name__) @@ -96,6 +99,7 @@ class DSparkDeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 @@ -235,6 +239,15 @@ class DSparkDeepseekV4Model(nn.Module): ) inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) @@ -247,6 +260,7 @@ class DSparkDeepseekV4Model(nn.Module): post_mix, res_mix, residual, + mega_gate_metadata=mega_gate_metadata, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) if self.use_sequence_parallel: diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 7fef691bd1..59156365c5 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -98,6 +98,37 @@ from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID logger = init_logger(__name__) +class MegaGateRoutingMetadata(typing.NamedTuple): + safe_hash_input_ids: torch.Tensor | None + image_token_mask: torch.Tensor | None + hash_token_mask: torch.Tensor | None + + +def prepare_mega_gate_routing_metadata( + input_ids: torch.Tensor, + *, + has_hash_routing: bool, + image_sentinel_base_id: int | None, +) -> MegaGateRoutingMetadata: + if image_sentinel_base_id is None: + return MegaGateRoutingMetadata( + input_ids if has_hash_routing else None, + None, + torch.ones_like(input_ids, dtype=torch.bool) if has_hash_routing else None, + ) + + image_token_mask = (input_ids >= image_sentinel_base_id) & ( + input_ids < image_sentinel_base_id + 5 + ) + if not has_hash_routing: + return MegaGateRoutingMetadata(None, image_token_mask, None) + return MegaGateRoutingMetadata( + torch.where(image_token_mask, 0, input_ids), + image_token_mask, + ~image_token_mask, + ) + + class DeepseekV4MLP(nn.Module): def __init__( self, @@ -803,6 +834,10 @@ class DeepseekV4MoE(nn.Module): raise NotImplementedError( "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." ) + if self.use_mega_moe and not self.renormalize: + raise NotImplementedError( + "DeepSeek V4 MegaMoE requires normalized top-k probabilities." + ) if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": raise NotImplementedError( "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" @@ -1000,7 +1035,10 @@ class DeepseekV4MoE(nn.Module): ) def forward( - self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None = None, + mega_gate_metadata: MegaGateRoutingMetadata | None = None, ) -> torch.Tensor: if self.gate.tid2eid is not None and input_ids is None: raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") @@ -1013,23 +1051,61 @@ class DeepseekV4MoE(nn.Module): return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape - router_logits, _ = self.gate(hidden_states) - topk_weights, topk_ids = fused_topk_bias( - hidden_states=hidden_states, - gating_output=router_logits, - scoring_func=self.scoring_func, - e_score_correction_bias=self.gate.e_score_correction_bias.data - if self.gate.e_score_correction_bias is not None - else None, - topk=self.n_activated_experts, - renormalize=self.renormalize, - indices_type=self.hash_indices_dtype, - input_tokens=input_ids, - hash_indices_table=self.gate.tid2eid, - routed_scaling_factor=self.routed_scaling_factor, - bias_vl=bias_vl.data if bias_vl is not None else None, - image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, - ) + # Small local padded batches favor GateLinear; 128-expert gates cross earlier. + gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16 + if hidden_states.shape[0] <= gate_threshold: + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None, + topk=self.n_activated_experts, + renormalize=self.renormalize, + indices_type=self.hash_indices_dtype, + input_tokens=input_ids, + hash_indices_table=self.gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + bias_vl=bias_vl.data if bias_vl is not None else None, + image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, + ) + else: + from vllm.utils.deep_gemm import bf16_mega_gate + + if mega_gate_metadata is None: + raise RuntimeError( + "MegaMoE routing metadata must be prepared once by the model." + ) + + unmapped_topk_idx = None + fix_routing_mask = None + if self.gate.tid2eid is not None: + if mega_gate_metadata.safe_hash_input_ids is None: + raise RuntimeError("Hash routing requires prepared input IDs.") + if mega_gate_metadata.hash_token_mask is None: + raise RuntimeError("Hash routing requires a prepared routing mask.") + unmapped_topk_idx = self.gate.tid2eid[ + mega_gate_metadata.safe_hash_input_ids + ] + fix_routing_mask = mega_gate_metadata.hash_token_mask + + topk_weights, topk_ids = bf16_mega_gate( + hidden_states, + self.gate.weight, + self.n_activated_experts, + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + ep_rank=self.ep_rank, + bias=self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None, + image_bias=bias_vl.data if bias_vl is not None else None, + image_token_mask=mega_gate_metadata.image_token_mask, + fix_routing_mask=fix_routing_mask, + unmapped_topk_idx=unmapped_topk_idx, + ) activation_clamp = ( float(self.swiglu_limit) if self.swiglu_limit is not None else None ) @@ -1248,6 +1324,7 @@ class DeepseekV4DecoderLayer(nn.Module): post_mix: torch.Tensor | None = None, res_mix: torch.Tensor | None = None, residual: torch.Tensor | None = None, + mega_gate_metadata: MegaGateRoutingMetadata | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: attn_norm_weight = self.attn_norm.weight.data attn_norm_eps = self.attn_norm.variance_epsilon @@ -1332,7 +1409,7 @@ class DeepseekV4DecoderLayer(nn.Module): norm_eps=ffn_norm_eps, ) - x = self.ffn(x, input_ids) + x = self.ffn(x, input_ids, mega_gate_metadata) return x, residual, post_mix, res_mix @@ -1394,6 +1471,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): ), prefix=f"{prefix}.layers", ) + self.has_local_hash_moe = self.start_layer < min( + self.end_layer, config.num_hash_layers + ) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) @@ -1496,6 +1576,16 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): hidden_states = sp_shard(hidden_states) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=self.has_local_hash_moe, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) + residual, post_mix, res_mix = None, None, None remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) aux_hidden_states: list[torch.Tensor] = [] @@ -1511,6 +1601,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): post_mix, res_mix, residual, + mega_gate_metadata=mega_gate_metadata, ) if idx + 1 in self.aux_hidden_state_layers: # Reconstruct the aux hidden state for draft models diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index a19cc707b2..50886b0afa 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -20,6 +20,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig +from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -51,6 +52,7 @@ from vllm.models.common.ops.sequence_parallel import ( sp_padding_mask, sp_shard, ) +from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import ( _FUSED_MTP_INPUT_RMSNORM_KERNEL, _MTP_SHARED_HEAD_RMSNORM_KERNEL, @@ -62,6 +64,7 @@ from .model import ( DeepseekV4Model, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) logger = init_logger(__name__) @@ -92,6 +95,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): self.config = config quant_config = vllm_config.quant_config self.rms_norm_eps = config.rms_norm_eps + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -181,8 +185,23 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( inputs_embeds ).unsqueeze(-2) + mega_gate_metadata = None + if self.use_mega_moe: + routing_input_ids = input_ids + if self.mtp_block.use_sequence_parallel: + routing_input_ids = sp_shard(routing_input_ids) + mega_gate_metadata = prepare_mega_gate_routing_metadata( + routing_input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=input_ids + positions=positions, + x=hidden_states, + input_ids=input_ids, + mega_gate_metadata=mega_gate_metadata, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) if self.mtp_block.use_sequence_parallel: diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 9a07189bc3..c08de12a12 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -18,6 +18,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -50,12 +51,14 @@ from vllm.models.common.ops.sequence_parallel import ( sp_shard, ) +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from .model import ( DeepseekV4DecoderLayer, DeepseekV4Model, _linear_scale_param_name, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) logger = init_logger(__name__) @@ -77,6 +80,7 @@ class DSparkDeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = ( @@ -202,6 +206,15 @@ class DSparkDeepseekV4Model(nn.Module): ) inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) @@ -215,6 +228,7 @@ class DSparkDeepseekV4Model(nn.Module): post_mix, res_mix, residual, + mega_gate_metadata=mega_gate_metadata, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) # Collapse the hc copies with the pre-mix from the last layer's FFN diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 21fa00e5c9..6d1faace65 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -61,7 +61,9 @@ from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) from vllm.models.deepseek_v4.nvidia.model import ( + MegaGateRoutingMetadata, make_deepseek_v4_expert_params_mapping, + prepare_mega_gate_routing_metadata, ) from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( @@ -339,6 +341,7 @@ class DeepseekV4DecoderLayer(nn.Module): engram_mask: torch.Tensor | None = None, *, capture_previous_aux: bool = False, + mega_gate_metadata: MegaGateRoutingMetadata | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -462,7 +465,7 @@ class DeepseekV4DecoderLayer(nn.Module): norm_weight=self.ffn_norm.weight, norm_eps=self.ffn_norm.variance_epsilon, ) - x = self.ffn(x, input_ids) + x = self.ffn(x, input_ids, mega_gate_metadata) return x, residual, post_mix, res_mix, ffn_pre, previous_aux @@ -705,6 +708,16 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): hidden_states = sp_shard(hidden_states) input_ids = sp_shard(input_ids) + mega_gate_metadata = None + if self.use_mega_moe: + mega_gate_metadata = prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) + residual, post_mix, res_mix = None, None, None pre_mix: torch.Tensor | None = None if not get_pp_group().is_first_rank: @@ -728,6 +741,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): engram_hashes, engram_mask, capture_previous_aux=idx in self.aux_hidden_state_layers, + mega_gate_metadata=mega_gate_metadata, ) if previous_aux is not None: # idx is the one-based id of the layer whose post this is. diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 3d96b496f1..2a0d4eee4a 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -164,6 +164,7 @@ _get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None _get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None _mega_mhc_impl: Callable[..., Any] | None = None +_bf16_mega_gate_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( @@ -240,6 +241,7 @@ def _lazy_init() -> None: global _get_sparse_mqa_logits_metadata_impl global _get_paged_sparse_mqa_logits_metadata_impl global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl + global _bf16_mega_gate_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl global _get_theoretical_mk_alignment_for_contiguous_layout_impl @@ -260,6 +262,7 @@ def _lazy_init() -> None: or _get_paged_mqa_logits_metadata_impl is not None or _tf32_hc_prenorm_gemm_impl is not None or _mega_mhc_impl is not None + or _bf16_mega_gate_impl is not None or _get_mk_alignment_for_contiguous_layout_impl is not None or _transform_sf_into_required_layout_impl is not None or _pack_ue8m0_to_int_impl is not None @@ -315,6 +318,7 @@ def _lazy_init() -> None: ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) _mega_mhc_impl = getattr(_dg, "mega_mhc", None) + _bf16_mega_gate_impl = getattr(_dg, "bf16_mega_gate", None) _get_mn_major_tma_aligned_tensor_impl = getattr( _dg, "get_mn_major_tma_aligned_tensor", None ) @@ -482,6 +486,55 @@ def cublaslt_gemm_nt(*args, **kwargs): return _cublaslt_gemm_nt_impl(*args, **kwargs) +def bf16_mega_gate( + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + *, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run DeepGEMM Mega Gate. + + The output order follows vLLM's router contract: weights, then indices. + Passing PyTorch-owned outputs to DeepGEMM makes their storage part of CUDA + graph capture. + """ + _lazy_init() + if _bf16_mega_gate_impl is None: + raise RuntimeError( + "DeepGEMM MegaMoE requires a DeepGEMM build with bf16_mega_gate." + ) + + topk_idx = torch.empty((x.shape[0], num_topk), dtype=torch.int64, device=x.device) + topk_weights = torch.empty( + (x.shape[0], num_topk), dtype=torch.float32, device=x.device + ) + _bf16_mega_gate_impl( + x, + weight, + num_topk, + use_shared_as_routed=False, + num_shared_experts=0, + routed_scaling_factor=routed_scaling_factor, + ep_rank=ep_rank, + scoring_func=scoring_func, + bias=bias, + image_bias=image_bias, + image_token_mask=image_token_mask, + fix_routing_mask=fix_routing_mask, + unmapped_topk_idx=unmapped_topk_idx, + out=(topk_idx, topk_weights), + ) + return topk_weights, topk_idx + + def fp8_gemm_nt(*args, **kwargs): _lazy_init() if _fp8_gemm_nt_impl is None: @@ -988,6 +1041,7 @@ def should_use_deepgemm_for_fp8_linear( __all__ = [ "calc_diff", + "bf16_mega_gate", "DeepGemmQuantScaleFMT", "fp8_gemm_nt", "fp8_einsum",