fix triton kernel with the correct device

[jiqing-feng]

The triton kernel didn't consider model on multi-card. To reproduce the error, run ZE_AFFINITY_MASK=2,3 python test_triton_bnb.py on the following script:

from transformers import pipeline

pipe = pipeline("text-generation", model="hugging-quants/Meta-Llama-3.1-8B-Instruct-BNB-NF4", device_map="auto")
print(pipe.model.hf_device_map)
output = pipe("Hello, my name is", max_new_tokens=10)
print(output)

The error is:

AssertHandler::printMessage
/pytorch/third_party/torch-xpu-ops/src/ATen/native/xpu/sycl/TensorCompareKernels.cpp:163: operator(): global id: [0,0,0], local id: [0,0,0] Assertion `input_[0] != 0` failed.
Aborted (core dumped)

After this PR, the output be like:

[{'generated_text': 'Hello, my name is Dr. Smith and I am a professor of mathematics'}]

[jiqing-feng]

Hi @matthewdouglas . Do you mind reviewing this PR? Thanks!

[matthewdouglas]

LGTM, thanks for the fix!

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.