diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 614d903ca1..b7f5ca29ec 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -8,6 +8,7 @@ import pytest import torch from vllm.config import CompilationConfig, CUDAGraphMode +from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.forward_context import ( BatchDescriptor, ForwardContext, @@ -144,8 +145,18 @@ def test_piecewise_graphs_break_out_by_size(): write_batch_kv.assert_called_once_with(*states) -def test_replay_graph_matches_eager(monkeypatch): - """The replay graph of the next captured size pads and runs the rows.""" +@pytest.mark.parametrize( + ("moe_backend", "ids_dtype"), + [ + ("auto", torch.int32), + *((b, torch.int64) for b in sorted(NATIVE_MEGA_MOE_BACKENDS)), + *((b, torch.int32) for b in sorted(FLASHINFER_MOE_EP_BACKENDS)), + ], + ids=str, +) +def test_replay_graph_matches_eager(monkeypatch, moe_backend, ids_dtype): + """The replay graph of the next captured size pads and runs the rows. Input ids + arrive as the model passes them: int32, cast to int64 for native MegaMoE only.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) capture = torch.cuda.stream(torch.cuda.Stream()) monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture) @@ -157,6 +168,7 @@ def test_replay_graph_matches_eager(monkeypatch): cfg.cache_config.use_kda_recoverssm = False cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN) cfg.model_config.dtype = torch.float32 + cfg.kernel_config.moe_backend = moe_backend def run_layers(hidden, positions, ids, pre, post, mix, residual): return residual + hidden[:, None] + positions[:, None, None], pre + ids[:, None] @@ -173,6 +185,7 @@ def test_replay_graph_matches_eager(monkeypatch): torch.randn(12, *b.shape[1:], device=DEVICE).to(b.dtype) for b in manager.input_buffers ] + states[2] = torch.arange(12, device=DEVICE, dtype=ids_dtype) # input_ids rows = torch.tensor([11, 2, 7], device=DEVICE) desc = manager.dispatch(2, 3, None, 0) with override_forward_context(_graph_context(desc.num_tokens)): diff --git a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py index 871f535584..c985f973bb 100644 --- a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py +++ b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py @@ -10,7 +10,7 @@ import torch from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.config import CUDAGraphMode, VllmConfig -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.forward_context import get_forward_context, override_forward_context from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -56,8 +56,9 @@ class DecoderReplayCudaGraphManager(CudaGraphManager): config, act = vllm_config.model_config.hf_config, vllm_config.model_config.dtype assert isinstance(act, torch.dtype) hc, h, f32 = config.hc_mult, config.hidden_size, torch.float32 - mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS - ids = torch.int64 if mega_moe else torch.int32 + # DeepseekV4Model.forward casts input_ids to int64 for native MegaMoE only. + native = vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ids = torch.int64 if native else torch.int32 # hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual shapes = [(h,), (), (), (hc,), (hc, 1), (hc, hc), (hc, h)] dtypes = [act, torch.int64, ids, f32, f32, f32, act]