[XPU] Ensure RNG offset alignment with PyTorch requirements in XPU sampler (#43028)
Signed-off-by: chaojun-zhang <chaojun.zhang@intel.com> Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
d4004455d2
commit
71d810bbf4
@@ -291,7 +291,8 @@ class TopKTopPSampler(nn.Module):
|
||||
)
|
||||
# The custom XPU sampler kernel consumes RNG values internally, so advance
|
||||
# the default generator's offset to keep future draws deterministic.
|
||||
offset += logits.numel()
|
||||
# pytorch: offset must be multiple of 4
|
||||
offset = (offset + logits.numel() + 3) // 4 * 4
|
||||
state.view(torch.int64)[1] = offset
|
||||
generator.set_state(state)
|
||||
return random_sampled, logits_to_return
|
||||
|
||||
Reference in New Issue
Block a user