optimize multi gpu support

[Presburger]

/kind improvement
issue: #147

[mergify]

@Presburger 🔍 Important: PR Classification Needed!

For efficient project management and a seamless review process, it's essential to classify your PR correctly. Here's how:

1. If you're fixing a bug, label it as kind/bug.
2. For small tweaks (less than 20 lines without altering any functionality), please use kind/improvement.
3. Significant changes that don't modify existing functionalities should be tagged as kind/enhancement.
4. Adjusting APIs or changing functionality? Go with kind/feature.

For any PR outside the kind/improvement category, ensure you link to the associated issue using the format: “issue: #”.

Thanks for your efforts and contribution to the community!.

[foxspy]

/lgtm

[Presburger]

kind/improvement

[liliu-z]

/approve

[sre-ci-robot]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: foxspy, liliu-z, Presburger

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

• OWNERS [Presburger,foxspy,liliu-z]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment