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