[llama-mm] Add torch.cond to replace if condition in MHA

[larryliu0820]

Summary:

In torchtune's MultiHeadAttention we have this logic:

If y is not None, calculate the values of k and v from y and update the KVCache.

Otherwise (if y is None), retrieve the value of k and v from KVCache.

This logic is not able to be handled by export world. Here I'm proposing a rewrite:

If y does not have all values equal to nan (not a number), calculate the values of k and v from y and update the KVCache.

Otherwise (if all of the values of y are nan), retrieve the value of k and v from KVCache.

This rewrite allows the module to satisfy the requirement of torch.cond and avoid specialization:

• The operands to torch.cond should have the same shape for the true branch and the false branch.

This means we will have to change this logic in torchtune:

        if encoder_input is not None:
            encoder_embed = self.encoder(**encoder_input)

        output = self.decoder(
            tokens=tokens,
            mask=mask,
            encoder_input=encoder_embed,
            encoder_mask=encoder_mask,
            input_pos=input_pos,
        )

To be:

        if encoder_input is not None:
            encoder_embed = self.encoder(**encoder_input)
        else:
            encoder_embed = torch.full_like(encoder_input, torch.nan)
        output = self.decoder(
            tokens=tokens,
            mask=mask,
            encoder_input=encoder_embed,
            encoder_mask=encoder_mask,
            input_pos=input_pos,
        )

Test Plan: Rely on unit tests

Reviewers:

Subscribers:

Tasks:

Tags:

Summary

[PLEASE REMOVE] See CONTRIBUTING.md's Pull Requests for ExecuTorch PR guidelines.

[PLEASE REMOVE] If this PR closes an issue, please add a Fixes #<issue-id> line.

[PLEASE REMOVE] If this PR introduces a fix or feature that should be the upcoming release notes, please add a "Release notes: " label. For a list of available release notes labels, check out CONTRIBUTING.md's Pull Requests.

Test plan

[PLEASE REMOVE] How did you test this PR? Please write down any manual commands you used and note down tests that you have written if applicable.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/6869

• 📄 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:

• [DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit b7763e9 with merge base ad15852 ():
💚 Looks good so far! There are no failures yet. 💚

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

[jackzhxng]

Seems fine - so instead of passing in y is None we pass in a tensor full of zeroes