It would be great in a multi-device system, the device id that will run a KA kernel could be set through a function call. cc: @vchuravy