[WIP] add lora codes for SglangJaxRollout#844

aolemila wants to merge 1 commit intomaingoogle/tunix:mainfrom

fix/replace-device-put-with-reshardgoogle/tunix:fix/replace-device-put-with-reshardCopy head branch name to clipboard