[BUG] flash_attention_gpu (mo.mha.no_cache) fails kernel elaboration for float16 on H100

Bug description

max.nn.kernels.flash_attention_gpu (op mo.mha.no_cache, MHAMaskVariant.NULL_MASK) fails graph compilation for float16 inputs with a kernel-elaboration constraint failure in max/mojo/max/gpu/compute/mma.mojo (:878: function instantiation failed, :955: constraint failed). The identical graph compiles and runs correctly (cos = 1.000000 vs a PyTorch SDPA reference) in bfloat16.

Shapes tested: q [1, 2, 12, 128] with kv [1, 2, 12, 128] and [1, 6, 12, 128] (both fail in fp16, both pass in bf16). H100.

This matters in combination with modular/modular#6883: on the current nightly, bf16 plain ops.matmul is broken (M > 8) while fp16 is correct — so fp16 is the natural fallback dtype, but then the fused no-cache attention kernel is unavailable and one has to hand-compose attention out of matmul/softmax.

Repro

import torch
from max.driver import Accelerator, Buffer, CPU
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn.kernels import flash_attention_gpu
from max.nn.attention.mask_config import MHAMaskVariant

GPU = DeviceRef.GPU()
session = InferenceSession(devices=[Accelerator()])
for dt in (DType.float16, DType.bfloat16):
    try:
        with Graph(f"fa_{dt}", input_types=[
            TensorType(dt, (1, 2, 12, 128), device=GPU),
            TensorType(dt, (1, 6, 12, 128), device=GPU),
            TensorType(dt, (1, 6, 12, 128), device=GPU),
        ]) as g:
            q, k, v = (i.tensor for i in g.inputs)
            g.output(flash_attention_gpu(q, k, v, MHAMaskVariant.NULL_MASK, 128 ** -0.5))
        session.load(g)
        print(dt, "OK")
    except Exception as e:
        print(dt, "FAIL:", str(e).splitlines()[0])

Output:

DType.float16 FAIL: Failed to compile the model. ...
DType.bfloat16 OK

System information

• MAX 26.6.0.dev2026081105 (pip nightly), NVIDIA H100 80GB, Ubuntu 22.04.5
• driver 570.172.08 + cuda-compat-13-0 580.178.04 forward-compat libs (native CUDA 13 path)

Start with max/mojo/max/gpu/compute/mma.mojo around lines 878 and 955, then run the supplied Python repro on an H100 to compare float16 and bfloat16 elaboration. Trace the constraint failure and validate the fix against both reported query/key-value shapes; done means the float16 graph compiles and runs like the bfloat16 case.