nvfp4-megamoe-kernel/dsv4/kernels/attention/production.py

"""DSV4 Blackwell Attention — Production kernel wrapper.

Wraps the CuTeDSL FMHA kernel with Python KV merge for multi-KV-tile.
Replaces vLLM's broken FlashMLA Blackwell implementation.

Supported attention types:
- CSA (Classical Self-Attention): 32 KV entries/block, full KV cache
- HCA (Hash-based Cross-Attention): 1 KV entry/block, compressed cache
- SWA (Sliding Window Attention): ring buffer, no KV cache

Limitations:
- head_dim > 256: MLIR compilation hang (known CuTeDSL issue)
- In-kernel multi-KV-tile: blocked on TMA layout matching (uses Python KV merge)
"""
import torch
import math
import cutlass.cute as cute
import cutlass.torch as ct
import cuda.bindings.driver as cuda
from dsv4.kernels.attention.fmha import FmhaKernel

# Kernel cache: compiled kernels keyed by config tuple
_kernel_cache: dict = {}


def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
                           normalize: bool = False, apply_swa_mask: bool = False,
                           is_causal: bool = False, n_comp: int = 0,
                           apply_sink_bias: bool = False) -> tuple:
    """Get or compile a kernel for the given configuration."""
    key = (head_dim, s_k, use_smem_p, normalize, apply_swa_mask, is_causal, n_comp, apply_sink_bias)
    if key in _kernel_cache:
        return _kernel_cache[key]

    kernel = FmhaKernel(
        head_dim=head_dim, s_k=s_k, use_smem_p=use_smem_p, normalize=normalize,
        apply_swa_mask=apply_swa_mask, is_causal=is_causal,
        n_comp=n_comp, apply_sink_bias=apply_sink_bias,
    )

    m = 128
    pv_n_tile = kernel.pv_n_tile
    q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, head_dim, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, pv_n_tile, dtype=torch.bfloat16, device='cuda')
    c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
    lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')

    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

    v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
    mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
    mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
    mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
    mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
    mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse))

    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE)
    _kernel_cache[key] = (compiled, kernel)
    return (compiled, kernel)


def dsv4_attention(
    q: torch.Tensor,      # (num_heads, seq_len, head_dim)
    k: torch.Tensor,      # (seq_len, head_dim) or (num_heads, seq_len, head_dim) for MQA
    v: torch.Tensor,      # (seq_len, head_dim) or (num_heads, seq_len, head_dim) for MQA
    scale: float = None,
    swa_len: int = None,
    is_causal: bool = False,
    n_comp: int = 0,
    sink_bias: torch.Tensor = None,
) -> torch.Tensor:
    """Production DSV4 attention using CuTeDSL FMHA kernel + Python KV merge.

    Args:
        q: Query tensor (num_heads, seq_len, head_dim) BF16
        k: Key tensor (seq_len, head_dim) or (num_heads, seq_len, head_dim) BF16
        v: Value tensor (seq_len, head_dim) or (num_heads, seq_len, head_dim) BF16
        scale: Softmax scale (default: 1/sqrt(head_dim))
        swa_len: Sliding window length (None = no mask)
        is_causal: Apply causal mask
        n_comp: Compressed KV length for D5c sink bias merge
        sink_bias: Per-head FP32 logit bias for sink attention

    Returns:
        Output tensor (num_heads, seq_len, head_dim) BF16
    """
    n_h, T, hd = q.shape
    N = k.shape[-2]
    scale = scale or (1.0 / math.sqrt(hd))
    use_smem_p = hd > 64
    apply_swa_mask = swa_len is not None
    apply_sink_bias = sink_bias is not None

    output = torch.zeros(n_h, T, hd, dtype=torch.bfloat16, device='cuda')

    for h in range(n_h):
        q_h = q[h:h+1]  # (1, T, hd)
        k_h = k.unsqueeze(0) if k.dim() == 2 else k[h:h+1]  # (1, N, hd)
        v_h = v.unsqueeze(0) if v.dim() == 2 else v[h:h+1]  # (1, N, hd)

        o_h = _attention_single_head(
            q_h, k_h, v_h, scale=scale,
            swa_len=swa_len, is_causal=is_causal,
            n_comp=n_comp, sink_bias=sink_bias[h] if sink_bias is not None else None,
            use_smem_p=use_smem_p,
        )
        output[h] = o_h[0]

    return output


def _attention_single_head(
    q: torch.Tensor,      # (1, T, hd)
    k: torch.Tensor,      # (1, N, hd)
    v: torch.Tensor,      # (1, N, hd)
    scale: float,
    swa_len: int = None,
    is_causal: bool = False,
    n_comp: int = 0,
    sink_bias: torch.Tensor = None,
    use_smem_p: bool = False,
) -> torch.Tensor:
    """Run FMHA for a single head, with Python KV merge for s_k > 128."""
    _, T, hd = q.shape
    N = k.shape[1]
    apply_swa_mask = swa_len is not None
    apply_sink_bias = sink_bias is not None

    # Segment the KV sequence into 128-entry tiles
    s_k_per_seg = 128
    n_segments = (N + s_k_per_seg - 1) // s_k_per_seg

    compiled, kernel = _get_or_compile_kernel(
        head_dim=hd, s_k=s_k_per_seg, use_smem_p=use_smem_p,
        normalize=False, apply_swa_mask=apply_swa_mask,
        is_causal=is_causal, n_comp=n_comp if n_comp > 0 else None,
        apply_sink_bias=apply_sink_bias,
    )
    pv_n_tile = kernel.pv_n_tile
    n_pv_tiles = kernel.n_pv_tiles

    # Accumulate across KV segments using log-sum-exp merge
    o_accum = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
    lse_accum = torch.full((T, 1), float('-inf'), dtype=torch.float32, device='cuda')

    for seg in range(n_segments):
        k_start = seg * s_k_per_seg
        k_end = min(k_start + s_k_per_seg, N)
        k_seg = k[:, k_start:k_end]
        v_seg = v[:, k_start:k_end]

        # Pad last segment if shorter than s_k_per_seg
        if k_end - k_start < s_k_per_seg:
            pad_len = s_k_per_seg - (k_end - k_start)
            k_seg = torch.cat([k_seg, torch.zeros(1, pad_len, hd, dtype=k.dtype, device='cuda')], dim=1)
            v_seg = torch.cat([v_seg, torch.zeros(1, pad_len, hd, dtype=v.dtype, device='cuda')], dim=1)

        seg_o = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
        seg_lse = torch.zeros(T, 1, dtype=torch.float32, device='cuda')

        for nt in range(n_pv_tiles):
            v_start = nt * pv_n_tile
            v_end = v_start + pv_n_tile
            v_tile = v_seg[0, :, v_start:v_end].contiguous()
            v_kernel = v_tile.unsqueeze(-1)
            c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
            lse_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')

            q_input = q[0].contiguous().unsqueeze(-1)
            k_input = k_seg[0].contiguous().unsqueeze(-1)
            stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

            mQ = ct.from_dlpack(q_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_input))
            mK = ct.from_dlpack(k_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_input))
            mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
            mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
            mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))

            compiled(mQ, mK, mV, mC, stream, lse=mLSE)
            torch.cuda.synchronize()

            seg_o[:, v_start:v_end] = c_tile[:, :, 0].float()
            if nt == 0:
                seg_lse[:, 0] = lse_tensor[:, 0, 0].float()

        # Merge with accumulator using log-sum-exp
        e_old = torch.exp(lse_accum)
        e_new = torch.exp(seg_lse)
        e_sum = e_old + e_new

        o_accum = (e_old * o_accum + e_new * seg_o) / e_sum
        lse_accum = torch.log(e_sum)

    output = o_accum.to(torch.bfloat16).unsqueeze(0)
    return output