fix pin memory device

[youkaichao]

right now flashinfer does not specify the device when creating pin-memory tensor. it will cause error when users changed the default pytorch device.

we should explicitly set the device to be "cpu" for pin-memory tensor.

[youkaichao]

example error trace:

kernels/test_flashinfer.py:229:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/usr/local/lib/python3.12/dist-packages/flashinfer/prefill.py:1099: in __init__
    self._pin_memory_int_workspace_buffer = torch.empty(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <torch.utils._device.DeviceContext object at 0x7ff7263fde50>, func = <built-in method empty of type object at 0x7ff8c8c5e280>
types = (), args = (torch.Size([8388608]),), kwargs = {'device': device(type='cuda'), 'dtype': torch.uint8, 'pin_memory': True}

    def __torch_function__(self, func, types, args=(), kwargs=None):
        kwargs = kwargs or {}
        if func in _device_constructors() and kwargs.get('device') is None:
            kwargs['device'] = self.device
>       return func(*args, **kwargs)
E       RuntimeError: Only dense CPU tensors can be pinned

[yzh119]

LGTM, thank you!