SIGN IN SIGN UP

Fix DFlash candidate token device mismatch with device_map="auto" (#47877)

* Fix DFlash candidate token device mismatch with device_map="auto"

This PR fixes a device mismatch in DFlashTokenCandidateGenerator when models are split across devices with device_map="auto".

candidate_logits can be produced on the device of main_model_output_embeddings, while input_ids / candidate_ids may live on another device. In the logits-processor path, this could pass logits from a different device than candidate_ids into self.logits_processor. In the vectorized path, sampled or argmax candidate tokens could then be concatenated with input_ids across devices.

The fix keeps token-id construction device-consistent by:

moving each logits slice to candidate_ids.device before applying logits_processor
moving vectorized generated candidate ids to input_ids.device before concatenation

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* update

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* fmt

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* update test

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* update

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* update case

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* let logits on output device, similar to models

* fix device in general and fix tests

---------

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
Co-authored-by: Cyril Vallez <cyril.vallez@gmail.com>
W
Wang, Yi committed
18388f3b684f6e80ba21e25be5b27bf7b342b7ec
Parent: 242f5df
Committed by GitHub <noreply@github.com> on 8/17/2026, 4:53:39 AM