[XPU] Fix Eagle3 initialization on XPU

[chaojun-zhang]

Purpose:

Fix Eagle3 draft model embedding sharing when the target and draft models have different embedding widths.

Why

tests/models/test_initialization.py uses a patched _initialize_kv_caches_v1 to skip the warmup path. That patch works under fork with create_process_for_each_test because the child inherits the test process state, but it does not work under spawn because the child starts fresh. On XPU, that means the real warmup path still runs and exposes the invalid embedding sharing.

Changes

Add an embedding-dimension check before sharing embed_tokens.
Keep target and draft embeddings separate when their widths differ.
Log the mismatch so the fallback is visible.

Test plan (CUDA):

VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]

Test output

EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 341, in _fn
(EngineCore pid=4091)     return fn(*args, **kwargs)
(EngineCore pid=4091)            ^^^^^^^^^^^^^^^^^^^
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 863, in trace_frame
(EngineCore pid=4091)     run_tracer()
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 844, in run_tracer
(EngineCore pid=4091)     tracer.run()
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1794, in run
(EngineCore pid=4091)     while self.step():
(EngineCore pid=4091)           ^^^^^^^^^^^
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1459, in step
(EngineCore pid=4091)     self.dispatch_table[inst.opcode](self, inst)
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1001, in wrapper
(EngineCore pid=4091)     return inner_fn(self, inst)
(EngineCore pid=4091)            ^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 4093, in CALL
(EngineCore pid=4091)     self._call(inst)
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 4084, in _call
(EngineCore pid=4091)     self.call_function(fn, args, kwargs)
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1360, in call_function
(EngineCore pid=4091)     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
(EngineCore pid=4091)               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/variables/lazy.py", line 294, in realize_and_forward
(EngineCore pid=4091)     return getattr(self.realize(), name)(*args, **kwargs)
(EngineCore pid=4091)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/_dynamo/variables/torch.py", line 2169, in call_function
(EngineCore pid=4091)     self.as_python_constant()(
(EngineCore pid=4091)   File "/.venv/lib/python3.12/site-packages/torch/__init__.py", line 2254, in _assert
(EngineCore pid=4091)     raise AssertionError(message)
(EngineCore pid=4091) AssertionError: hidden_states and input_embeds must have the same last dimension
[rank0]:[W624 10:18:48.695258091 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

[chaojun-zhang]

@jikunshang @benchislett please help review

[yma11]

Please take a look the failed UTs.

[eugr]

@jikunshang - This PR breaks Gemma 4 MTP. It fails with "Target embedding dim (2816) differs from draft embedding dim (1024). Keeping separate embedding weights."

Gemma4MultiTokenPredictor.pre_projection expects 2 * backbone_hidden_size, so for this model it expects 2816 + 2816 = 5632. After #43957 prevents embedding sharing, the draft path feeds 1024 + 2816 = 3840, producing the observed error:
[s47, 3840] X [5632, 1024].

The easiest fix would be to change the guard from. I tested it on DGX Spark and it seems to work:

if share_embeddings:

to:

if share_embeddings and hasattr(self.model, "has_own_embed_tokens"):

@mgoin - FYI

[chaojun-zhang]

@jikunshang - This PR breaks Gemma 4 MTP. It fails with "Target embedding dim (2816) differs from draft embedding dim (1024). Keeping separate embedding weights."

Gemma4MultiTokenPredictor.pre_projection expects 2 * backbone_hidden_size, so for this model it expects 2816 + 2816 = 5632. After #43957 prevents embedding sharing, the draft path feeds 1024 + 2816 = 3840, producing the observed error: [s47, 3840] X [5632, 1024].

The easiest fix would be to change the guard from. I tested it on DGX Spark and it seems to work:

if share_embeddings:

to:

if share_embeddings and hasattr(self.model, "has_own_embed_tokens"):

@mgoin - FYI

@eugr ,can you paste your test command ?

[eugr]

@chaojun-zhang - here we go:

vllm serve nvidia/Gemma-4-26B-A4B-NVFP4 \
  --max-model-len 262144 \
  --gpu-memory-utilization 0.8 \
  --port 8000 \
  --host 0.0.0.0 \
  --max-num-seqs 8 \
  --load-format fastsafetensors \
  --enable-prefix-caching \
  --enable-auto-tool-choice \
  --tool-call-parser gemma4 \
  --reasoning-parser gemma4 \
  --max-num-batched-tokens 8192 \
  --speculative-config '{"method":"mtp","model":"google/gemma-4-26B-A4B-it-assistant","num_speculative_tokens":4, "moe_backend": "triton"}'

This is for DGX Spark, so you'd need to adjust parameters accordingly for your test platform.

[chaojun-zhang]

@eugr thank you, seems #47833 try to fix this one.