FlashQLA GDN backend produces different tokens from FLA Triton on correctness-test prompts

[JessePinkman2019]

Setup

• H800 × 4, TP=4, bf16
• T=20468, B=1, chunk_size=64, varlen API

Observation

Comparing chunk_gated_delta_rule output with FLA Triton as reference (zero initial state):

┌───────┬──────────┬───────────┬─────────┐
│ T │ max_diff │ mean_diff │ cos_sim │
├───────┼──────────┼───────────┼─────────┤
│ 512 │ 0.026 │ 3.7e-5 │ 0.99806 │
├───────┼──────────┼───────────┼─────────┤
│ 20468 │ 0.047 │ 1.2e-5 │ 0.99947 │
└───────┴──────────┴───────────┴─────────┘

Mean error is small, but max_diff up to 0.047 is enough to flip argmax at specific positions. In an end-to-end generation test
with prompts padded to T=20468, FlashQLA produces wrong tokens on 2/3 prompts where Triton is correct (e.g. "The capital of the
UK is" → Triton: London, FlashQLA: .).

Question

Is this level of numerical deviation from FLA Triton expected? If so, is there a recommended tolerance or a reference correctness
test for non-padded long-context inputs?

[Starmys]

Thank you for your feedback!

Since FlashQLA is not bitwise identical to the FLA Triton implementation, it is expected that the output tokens may differ. However, we have not observed FlashQLA producing a large number of erroneous tokens on H200 – for example, it achieves consistent scores on benchmarks like MMLU. In our unit tests, we also haven't seen a max_diff as large as 0.047.

Here are our test results on H200 using QK normalized by L2‑norm and V drawn from a Gaussian distribution:

- attention output:
┌───────┬──────────┬─────────┐
│   T   | max_diff | max_ref |
├───────┼──────────┼─────────┤
│  512  |  0.0015  | 0.2500  |
├───────┼──────────┼─────────┤
│ 20468 |  0.0020  | 0.3105  |
└───────┴──────────┴─────────┘

- final_state:
┌───────┬──────────┬─────────┐
│   T   | max_diff | max_ref |
├───────┼──────────┼─────────┤
│  512  |  0.0158  | 2.7375  |
├───────┼──────────┼─────────┤
│ 20468 |  0.0175  | 3.0353  |
└───────┴──────────┴─────────┘

Since we do not have access to H800 devices, it is difficult for us to verify FlashQLA's correctness on H800. Could you please provide more details about your test setup, such as the input distribution and the way you integrated FlashQLA into your inference engine? That would help us better understand the discrepancy.

Thank you again for raising this issue!