[rocm-libraries] ROCm/rocm-libraries#6498 (commit 5961a2e)

[CK_TILE] Fix conditional rescale numerical instability in
 FMHA forward (#6498)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

[CK_TILE] Fix conditional rescale numerical instability in FMHA forward

## Motivation

Fix numerical instability in the conditional O-accumulator rescaling
optimization
for CK-Tile FMHA forward (FlashAttention-4, Algorithm 6, Eq. 6).

The conditional rescale optimization skips the expensive O-accumulator
rescale when
the running row-max shift is within a threshold (tau = log2(256) = 8.0).
The original
implementation had a bug: attention weights P were computed in the
`m_new` reference
frame before the skip/rescale decision. In the skip branch, `m` was
reverted to
`m_old`, but P remained in the `m_new` frame, causing incorrect softmax
normalization.

This fix introduces a `p_row_correction` factor: in the skip branch, P
is multiplied
by `exp2(m_new - m_old)` to bring it back to the `m_old` reference
frame.

- **Correctness:** Fixes broken inference on long sequences where
running-max drift
causes exp2 overflow (observed as degraded image quality on MI350X Flux2
generation)
- **Performance:** Neutral to +4% depending on workload shape

## Technical Details

6 pipeline header files (same pattern in each):
- `block_fmha_pipeline_qr_ks_vs.hpp`
- `block_fmha_pipeline_qr_ks_vs_async.hpp`
- `block_fmha_pipeline_qr_ks_vs_async_trload.hpp`
- `block_fmha_pipeline_qr_ks_vs_fp8.hpp`
- `block_fmha_pipeline_qr_ks_vs_whole_k_prefetch.hpp`
- `block_fmha_pipeline_qs_ks_vs.hpp`

In each file:
- Lower threshold from 10.0 to 8.0 (tau = log2(256))
- Add `p_row_correction` distributed tensor initialized to 1.0
- Rescale branch: standard rescale of O_acc and l; correction = 1.0
- Skip branch: compute correction = exp2(-acc_scale_log2), update l,
revert m, store correction
- New `p_spans` sweep applies per-row correction to `p_compute` before
P*V GEMM
- Move P-to-PDataType cast to after correction sweep

## Dependencies

None — this PR is standalone.

## Test Plan

- GPU validation on MI300X (gfx942, ROCm 6.4.1):
- Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128
-prec=bf16 -v=1 -warmup=1 -repeat=3`
- GPU validation on MI350X (gfx950, ROCm 7.0):
- Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128
-prec=bf16 -v=1 -warmup=1 -repeat=3`
- Command: `./build/bin/tile_example_fmha_fwd -b=2 -h=8 -s=4096 -d=128
-prec=fp16 -v=1 -warmup=1 -repeat=3`

## Test Result

Accuracy vs FP32 reference (MI350X, gfx950):

| Shape | max_diff | mean_diff |
|-------|----------|-----------|
| B=1 H=24 M=4096 K=128 bf16 | 9.1e-4 | 4.6e-5 |
| B=4 H=32 M=4096 K=128 bf16 | 9.9e-4 | 4.6e-5 |
| B=1 H=24 M=4096 K=128 fp16 | 1.2e-4 | 9.0e-6 |

Performance (MI350X, gfx950, ROCm 7.0):

| Shape | FA4 (TFlops) | Always-rescale (TFlops) | Delta |
|-------|-------------|------------------------|-------|
| B=1 H=24 M=4096 K=128 bf16 | 425.9 | 428.5 | neutral |
| B=2 H=8 M=2048 K=256 bf16 | 513.9 | 509.0 | +1.0% |
| B=1 H=64 M=2048 K=64 bf16 | 481.7 | 464.3 | +3.7% |

Benchmark results (MI300X, gfx942, ROCm 6.4.1):

No regression on MI300X. This correctness fix is performance-neutral.

| Config | TFlops / GB/s | Time (ms) |
|--------|-------------|-----------|
| MHA bf16 b=2 h=8 s=4096 d=128 | 342.49 TFlops | 0.401 |
| MHA fp16 b=2 h=8 s=4096 d=128 | 391.70 TFlops | 0.351 |
| Causal MHA bf16 b=2 h=8 s=4096 d=128 | 227.07 TFlops | 0.303 |
| GQA 4:1 bf16 b=2 h=32 hk=8 s=2048 d=128 | 324.69 TFlops | 0.423 |
| GQA 8:1 bf16 b=2 h=64 hk=8 s=2048 d=128 | 348.09 TFlops | 0.790 |
| LLaMA-70B prefill b=1 h=64 hk=8 s=4096 d=128 bf16 | 376.71 TFlops |
1.459 |
| Long-seq bf16 b=1 h=16 s=16384 d=128 | 383.42 TFlops | 5.735 |
| Decode b=64 h=32 hk=8 s_k=4096 d=128 bf16 | 691.64 GB/s | 1.554 |

All validation tests pass (`valid:y`) on both MI300X and MI350X.

Additional validation:
- Uniform scores: softmax output matches FP32 reference (max_diff <
1e-3)
- Large seqlen (4096+): no overflow or NaN in O-accumulator
- Spike pattern: correct handling of sudden row-max jumps
- Multiple spikes: correction applied correctly across multiple
skip/rescale transitions
- Deterministic: identical outputs across repeated runs
- No performance regression on standard workloads