Move Cutlass kernels to FBGEMM Gen-AI

Summary: As title

Differential Revision: D81189637

Author: sryap

fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py

@@ -0,0 +1,259 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional, Tuple
+
+import torch
+from torch.library import register_fake
+
+
+torch.library.define(
+    "blackwell_fmha::fmha_fwd",
+    "(Tensor q, Tensor k, Tensor v, Tensor? cu_seqlens_q, Tensor? cu_seqlens_k, int? max_seq_len_q, int? max_seq_len_k, float? softmax_scale, bool? causal, Tensor? seqlen_kv) -> (Tensor, Tensor)",
+    tags=[torch.Tag.pt2_compliant_tag],
+)
+
+torch.library.define(
+    "blackwell_fmha::fmha_bwd",
+    "(Tensor dout, Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, Tensor? cu_seqlens_q, Tensor? cu_seqlens_k, int? max_seq_len_q, int? max_seq_len_k, bool? causal) -> (Tensor, Tensor, Tensor)",
+    tags=[torch.Tag.pt2_compliant_tag],
+)
+
+
+@torch.library.impl("blackwell_fmha::fmha_fwd", "cuda")
+def custom_op_fmha(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor = None,
+    cu_seqlens_k: torch.Tensor = None,
+    max_seq_len_q: int = None,
+    max_seq_len_k: int = None,
+    softmax_scale: float = None,
+    causal: bool = False,
+    seqlen_kv: torch.Tensor = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert q.is_contiguous(), "q is not contiguous"
+    assert k.is_contiguous(), "k is not contiguous"
+    assert v.is_contiguous(), "v is not contiguous"
+    assert q.is_cuda, "q must be on GPU"
+    assert k.is_cuda, "k must be on GPU"
+    assert v.is_cuda, "v must be on GPU"
+    return torch.ops.fbgemm.fmha_fwd(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+        max_seq_len_k=max_seq_len_k,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        seqlen_kv=seqlen_kv,
+    )
+
+
+@register_fake("blackwell_fmha::fmha_fwd")
+def fmha_fwd_meta(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor = None,
+    cu_seqlens_k: torch.Tensor = None,
+    max_seq_len_q: int = None,
+    max_seq_len_k: int = None,
+    softmax_scale: float = None,
+    causal: bool = False,
+    seqlen_kv: torch.Tensor = None,
+):
+    if q.dtype == torch.float16:
+        out_dtype = torch.float16
+    elif q.dtype == torch.bfloat16:
+        out_dtype = torch.bfloat16
+    elif q.dtype == torch.float8_e4m3fn:
+        # Output is BF16 when input is FP8
+        out_dtype = torch.bfloat16
+    else:
+        raise RuntimeError(f"Unsupported dtype for q: {q.dtype}")
+
+    kIsVarlen = max_seq_len_q is not None
+    if kIsVarlen:
+        SQ = q.shape[0]
+        H_Q = q.shape[1]
+        B = cu_seqlens_q.shape[0] - 1
+    else:
+        SQ = q.shape[1]
+        H_Q = q.shape[2]
+        B = q.shape[0]
+    device = q.device
+    options2 = {"dtype": torch.float32, "device": device}
+    if kIsVarlen:
+        out = torch.empty_like(q, dtype=out_dtype)
+        size = out.size()
+        stride = out.stride()
+        storage_offset = q.shape[-1] * max_seq_len_q * H_Q  # example scalar offset
+        out1 = torch.as_strided(
+            out, size=size, stride=stride, storage_offset=storage_offset
+        )
+    else:
+        out1 = torch.empty_like(q, dtype=out_dtype)
+
+    if kIsVarlen:
+        out2 = torch.empty((1, H_Q, SQ), **options2)
+    else:
+        out2 = torch.empty((B, H_Q, SQ), **options2)
+    return out1, out2
+
+
+@torch.library.impl("blackwell_fmha::fmha_bwd", "cuda")
+def custom_op_fmha_bwd(
+    dOutput: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor = None,
+    cu_seqlens_k: torch.Tensor = None,
+    max_seq_len_q: int = None,
+    max_seq_len_k: int = None,
+    causal: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return torch.ops.fbgemm.fmha_bwd(
+        dOutput,
+        query,
+        key,
+        value,
+        output,
+        softmax_lse,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+        max_seq_len_k=max_seq_len_k,
+        causal=causal,
+    )
+
+
+@register_fake("blackwell_fmha::fmha_bwd")
+def fmha_bwd_meta(
+    dOutput: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor = None,
+    cu_seqlens_k: torch.Tensor = None,
+    max_seq_len_q: int = None,
+    max_seq_len_k: int = None,
+    causal: bool = False,
+):
+    return (
+        torch.empty_like(query),
+        torch.empty_like(key),
+        torch.empty_like(value),
+    )
+
+
+def _backward(ctx, *grad):
+    if ctx.is_gen:
+        # For gen case, no backward pass is needed (generation is inference only)
+        raise RuntimeError("Backward pass is not supported for generation phase (sq=1)")
+    q, k, v, out, softmax_lse = ctx.saved_tensors
+    if not grad[0].is_contiguous():
+        grad0 = grad[0].contiguous()
+    else:
+        grad0 = grad[0]
+    if not softmax_lse.is_contiguous:
+        softmax_lse = softmax_lse.contiguous()
+    if not out.is_contiguous:
+        out = out.contiguous()
+    if not q.is_contiguous:
+        q = q.contiguous()
+    if not k.is_contiguous:
+        k = k.contiguous()
+
+    if not softmax_lse.is_contiguous:
+        softmax_lse = softmax_lse.contiguous()
+    if not out.is_contiguous:
+        out = out.contiguous()
+    if not q.is_contiguous:
+        q = q.contiguous()
+    if not k.is_contiguous:
+        k = k.contiguous()
+
+    dq, dk, dv = torch.ops.blackwell_fmha.fmha_bwd(
+        grad0,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        ctx.cu_seqlens_q,
+        ctx.cu_seqlens_k,
+        ctx.max_seq_len_q,
+        ctx.max_seq_len_k,
+        ctx.causal,
+    )
+    return dq, dk, dv, None, None, None, None, None, None, None
+
+
+def _setup_context(ctx, inputs, output):
+    (
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seq_len_q,
+        max_seq_len_k,
+        softmax_scale,
+        causal,
+        seqlen_kv,
+    ) = inputs
+    (out, softmax_lse) = output
+    ctx.save_for_backward(q, k, v, out, softmax_lse)
+    ctx.softmax_scale = softmax_scale
+    ctx.causal = causal
+    ctx.max_seq_len_q = max_seq_len_q
+    ctx.max_seq_len_k = max_seq_len_k
+    ctx.cu_seqlens_q = cu_seqlens_q
+    ctx.cu_seqlens_k = cu_seqlens_k
+    ctx.is_gen = False
+
+
+# This code adds training support for the operator. You must provide us
+# the backward formula for the operator and a `setup_context` function
+# to save values to be used in the backward.
+torch.library.register_autograd(
+    "blackwell_fmha::fmha_fwd", _backward, setup_context=_setup_context
+)
+
+
+def cutlass_blackwell_fmha_custom_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seq_len_q: Optional[int] = None,
+    max_seq_len_k: Optional[int] = None,
+    seqlen_kv: Optional[torch.Tensor] = None,
+):
+    return torch.ops.blackwell_fmha.fmha_fwd(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+        max_seq_len_k=max_seq_len_k,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        seqlen_kv=seqlen_kv,
+    )[0]