fix: Mark FP8 scale to have the same batch size as input

[ani300]

Description of the change

This PR adds an assertion that the FP8 scales have the same batch dimension as the inputs for paged attention. This is important for torch compile to not create extra dynamic dimensions, which is needed for AIU compilation and also accelerates tracing.

Related issues or PRs

Internal AIU issue tracker.

How to verify the PR

Was the PR tested

• I have added >=1 unit test(s) for every new method I have added (if that coverage is difficult, please briefly explain the reason)
• I have ensured all unit tests pass

Checklist for passing CI/CD:

• All commits are signed showing "Signed-off-by: Name <email@domain.com>" with git commit -signoff or equivalent
• PR title and commit messages adhere to Conventional Commits
• Contribution is formatted with tox -e fix
• Contribution passes linting with tox -e lint
• Contribution passes spellcheck with tox -e spellcheck
• Contribution passes all unit tests with tox -e unit

Note: CI/CD performs unit tests on multiple versions of Python from a fresh install. There may be differences with your local environment and the test environment.