iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[DSv4.1] Integrate Mega-Gate from DeepGEMM (#56266)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
This commit is contained in:
Thien Tran
2026-09-17 12:06:36 -07:00
committed by GitHub
co-authored by OpenAI Codex zjy0516
parent a9a7e45f31
commit 41f9104fa6
9 changed files with 369 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
+1
View File
@@ -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:
tests/kernels/test_mhc_kernels.py
+3 -2
View File
@@ -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
tests/models/test_deepseek_v4_mega_moe.py
+138 -12
View File
@@ -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,
vllm/models/deepseek_v4/nvidia/dspark.py
+14
View File
@@ -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:
vllm/models/deepseek_v4/nvidia/model.py
+110 -19
View File
@@ -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
vllm/models/deepseek_v4/nvidia/mtp.py
+20 -1
View File
@@ -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:
vllm/models/deepseek_v41/nvidia/dspark.py
+14
View File
@@ -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
vllm/models/deepseek_v41/nvidia/model.py
+15 -1
View File
@@ -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.
vllm/utils/deep_gemm.py
+54
View File
@@ -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",