fix(models): revisions triton gt kernel to improve stability

[ssmmnn11]

Improve numerical stability of the Triton GT kernel. However, this results in a ~ 10% memory increase and a similar % in terms of performance degradation (pure kernel - less in a real training run). Tests.

blue: current triton kernel bf16
red: current triton kernel bf16, second try
green: resume blue curve with pyg compiled bf16
orange: pyg compiled bf16
violet: pyg compiled bf16, cast to float32
brown: revised triton kernel bf16

test:

• make the attention logits very large -> softmax becomes very sharp
• every neighbor sends the same message -> output should no longer depend on the attention weights
• the gradients dq and dk should be zero or very close to zero
• check whether the backward pass still creates nonzero gradients

|dq| / |dk|

• float32
  • GT before change: 0.000078 / 0.000703
  • GT now: 0.000091 / 0.000454
• bfloat16
  • GT before change: 1.468 / 18.760
  • GT now: 0.321 / 3.615