[Bug]: Attention does not validate element type correctness, leading to CUDA invalid memory reference

[caelunshun]

Before submitting

• I searched existing issues and did not find a duplicate.

Affected area

None

Bug report

Executing an SDPA operation using the cudnn-frontend graph API, when requesting float32 precision, leads to a CUDA exception from an invalid memory reference. Inspecting this, it looks like cudnn dispatches a kernel with fp16 in the name that ends up writing out of bounds inside shared memory.

To reproduce: run the following script.

script

import math

import cudnn
import torch


def main():
    b, h, s, d = 1, 8, 512, 64
    shape = (b, h, s, d)

    q = torch.randn(shape, device="cuda", dtype=torch.float32)
    k = torch.randn(shape, device="cuda", dtype=torch.float32)
    v = torch.randn(shape, device="cuda", dtype=torch.float32)
    o = torch.empty(shape, device="cuda", dtype=torch.float32)

    graph = cudnn.pygraph(
        io_data_type=cudnn.data_type.FLOAT,
        intermediate_data_type=cudnn.data_type.FLOAT,
        compute_data_type=cudnn.data_type.FLOAT,
    )

    q_t = graph.tensor(name="Q", dim=list(q.shape), stride=list(q.stride()), data_type=cudnn.data_type.FLOAT)
    k_t = graph.tensor(name="K", dim=list(k.shape), stride=list(k.stride()), data_type=cudnn.data_type.FLOAT)
    v_t = graph.tensor(name="V", dim=list(v.shape), stride=list(v.stride()), data_type=cudnn.data_type.FLOAT)

    o_t, _ = graph.sdpa(
        q=q_t,
        k=k_t,
        v=v_t,
        is_inference=True,
        attn_scale=1.0 / math.sqrt(d),
        use_causal_mask=True,
        name="prefill_sdpa_f32",
    )
    o_t.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_dim(list(o.shape)).set_stride(list(o.stride()))

    graph.validate()
    graph.build_operation_graph()
    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
    graph.check_support()
    graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE)

    workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)
    graph.execute({q_t: q, k_t: k, v_t: v, o_t: o}, workspace)
    torch.cuda.synchronize()
if __name__ == "__main__":
    main()

Expected behavior

Since FP32 seems to be unsupported for SDPA, cudnn should throw a validation error on the graph rather than executing an invalid kernel.