fix: Fixes for paged fp8 attention with chunked prefill

[ani300]

Description of the change

This PR ports the fixes to paged attention coming from chunked prefill implementations to the FP8 variant, and restores our custom CPU implementation for scaled_bmm, as the Pytorch CPU does not support per-sentence kv-scales, which we need for our current fp8 implementation.

Related issues or PRs

See https://ibm.enterprise.slack.com/archives/C04FKKQ11LH/p1762528936408799?thread_ts=1762528936.408799&cid=C04FKKQ11LH for details on the fix.

How to verify the PR

Running DPP (https://github.com/foundation-model-stack/aiu-fms-testing-utils/blob/main/aiu_fms_testing_utils/scripts/drive_paged_programs.py) like this:

export DTLOG_LEVEL=warning
export DT_DEEPRT_VERBOSE=-1
export HF_HUB_CACHE=/tmp/models/hf_cache/hub
export VLLM_DT_MAX_CONTEXT_LEN=8192
export VLLM_DT_MAX_BATCH_SIZE=2
export VLLM_DT_MAX_BATCH_TKV_LIMIT=524288
export DEM_COMPILE_VERSION=0
export DTCOMPILER_KEEP_EXPORT=1
export DEE_DUMP_GRAPHS=granite8b
export TORCH_LOGS=recompiles,dynamo

python scripts/drive_paged_programs.py --model_variant=/tmp/models/granite-3.3-8b-instruct-FP8  --program_criteria_json_path=/mnt/home/program_criteria.json --timing per-token --dataset_type sharegpt --dataset_path=/tmp/models/ShareGPT_V3_unfiltered_cleaned_split.json --test_type tokens --skip_validation --max_new_tokens=128 --validation_info_outputs_dir=/mnt/home/tmp_validation_info_dir --save_validation_info_outputs --dist_timeout 120 --prioritize_large_batch_sizes --enforce_homogeneous_prompt_programs --prefill_chunk_size=4096 --attention_type=paged_fp8 |& tee drive_out.txt

produces the expected graph and results now.

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.

[andrea-fasoli]

@ani300 by registering a custom _scaled_mm_cpu_out in torch >= 2.8, we go back to the seeing the warning
Overriding a previously registered kernel for the same operator and the same dispatch key
which triggers an error in CI/CD (exact same error we tried to fix with #187).

We need to either suppress it or declare a kernel with a name distinct from the official Aten operation.

[andrea-fasoli]

what functionality do we need that is not provided by the conventional _scaled_mm_cpu_out in torch?

[ani300]

Scales that aren't per tensor, we can refactor it I guess.

[ani300]

Fixed the warning, now trying if it works with DPP before asking for merge

[ani300]

@andrea-fasoli The code compiled successfully for CPU using DPP tests in AIU, so the fix works.

[andrea-fasoli]

looks good to me