FP16 forward pass blows up with NaNs past 32k context length

[kds1123001]

Ran into an issue while testing long sequences using FP16. Everything works great up to around 32k tokens, but as soon as I push past seq_len = 32769, the forward pass starts dumping NaNs in the attention output.

to reproduce:

import torch
from flash_qla import flash_qla_forward

Anything > 32768 causes issues in FP16

q = torch.randn(1, 32769, 16, 128, dtype=torch.float16, device="cuda")
k = torch.randn(1, 32769, 16, 128, dtype=torch.float16, device="cuda")
v = torch.randn(1, 32769, 16, 128, dtype=torch.float16, device="cuda")

out = flash_qla_forward(q, k, v)
print(torch.isnan(out).any()) # True

Switching over to bfloat16 completely resolves it, so it looks like a numerical overflow in the Triton kernel accumulator when processing long sequence blocks.

Tested on an H100 with CUDA 12.4 / PyTorch 2.4.

[kds1123001]

Can you resolve this issue quickly, and also Has anyone else hit this wall on long contexts with FP16, or is the recommendation to strictly stick to BF16 for 32k+?

[Erix025]

Thanks for the report! However, based on the information currently provided we're unable to reproduce this, so we can't yet tell whether it's a bug or a usage issue.

Following our unit-test procedure, we tested both bf16 and fp16 at seq_len = 32769 and at a longer 40k sequence. In all cases the outputs match our float64 reference within tolerance — no NaNs in either dtype. So we haven't been able to reproduce the "fp16-only, past-32k" behavior on our end.

A couple of things we noticed while digging in:

1. q/k in the snippet aren't L2-normalized. FlashQLA implements the Gated Delta Rule, and L2-normalized q/k is a numerical requirement for GDN (the within-chunk delta-rule solve assumes unit-norm keys). When we feed raw randn q/k without normalization, we do get NaNs — but importantly this happens in both bf16 and fp16, and at any sequence length (we see it as low as T=256), so it's unrelated to fp16 or to the 32k boundary. Normalizing q/k (or passing use_qk_l2norm_in_kernel=True) makes the NaNs go away in our tests.
2. flash_qla_forward isn't part of our public API. Our entry point is chunk_gated_delta_rule (and the lower-level chunk_gated_delta_rule_fwd / chunk_gated_delta_rule_bwd), which also require g (gating, in log space) and beta in addition to q/k/v. The custom flash_qla_forward in the snippet may be hiding some input-construction details — for instance, we can't see how g and beta are generated.

Could you share a more complete, runnable reproduction? Specifically:

• the definition of flash_qla_forward (or the actual public function you're calling),
• how g and beta are constructed,
• whether q/k are normalized anywhere before the call.

With that we should be able to pin down what's happening. Thanks again!