Unify llama generate utils

[guangyey]

Motivation

Unify llama generate utils via torch.accelerator APIs.

cc @albanD

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3318

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Some B200 runners are down due to network issues

✅ You can merge normally! (1 Unrelated Failure)

As of commit b7156ec with merge base 415e0e8 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/dtypes/test_nf4.py::TestComm::test_comm

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[albanD]

Sounds good to me.
But I'll let an ao maintainer accept the PR

[jerryzh168]

looks good, as long as it works

[jerryzh168]

maybe spell out the arg of current_accelerator to be clearer https://docs.pytorch.org/docs/stable/generated/torch.accelerator.current_accelerator.html

[guangyey]

@jerryzh168 Done.

[guangyey]

@pytorchbot label "topic: improvement"

[guangyey]

The CI workflow needs to be triggered.

[guangyey]

@jerryzh168 I rebase this PR to the latest main already. Could you help re-trigger the CI.

[guangyey]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) are pending/not yet run. The first few are:

• Facebook CLA Check

Dig deeper by viewing the pending checks on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: superuser

[jerryzh168]

@guangyey we can't use pytorchbot here to merge the PR. we just click the button to merge

[pytorch-bot]

To add the ciflow label ciflow/xpu please first approve the workflows that are awaiting approval (scroll to the bottom of this page).

This helps ensure we don't trigger CI on this PR until it is actually authorized to do so. Please ping one of the reviewers if you do not have access to approve and run workflows.

[pytorch-bot]

To add the ciflow label ciflow/xpu please first approve the workflows that are awaiting approval (scroll to the bottom of this page).

This helps ensure we don't trigger CI on this PR until it is actually authorized to do so. Please ping one of the reviewers if you do not have access to approve and run workflows.

[guangyey]

Thanks, @jerryzh168. I’ll reach out to the Intel team to re-trigger the CI and land the PR once all checks are green.

[guangyey]

The failure is irrelevant to this PR. The same failure can be found in https://github.com/pytorch/ao/actions/runs/20529844203/job/58979473220 that already landed in #3548

File "/opt/conda/envs/venv/lib/python3.10/site-packages/torchao/dtypes/nf4tensor.py", line 940, in __torch_dispatch__
    raise NotImplementedError(
NotImplementedError: NF4Tensor dispatch: attempting to run _c10d_functional._wrap_tensor_autograd.default, this is not supported

To execute this test, run the following from the base repo dir:
    python test/dtypes/test_nf4.py TestComm.test_comm

[liangan1]

Double confirmed, the UT fail is not related to this PR. Merged.

[guangyey]

Thanks for your help!