minimax_h3_fusions ¶

def fused_qknorm_rope(
    x: torch.Tensor,
    weight: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    eps: float,
) -> torch.Tensor:
    """Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel.

    RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the
    final store. Triton's reduction order and the absence of eager's BF16
    intermediate materializations can produce small, expected rounding drift.

    Row offsets are computed in int64, so inputs beyond 2**31 total elements
    (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address
    correctly.
    """
    if torch.is_grad_enabled() and any(
            isinstance(tensor, torch.Tensor) and tensor.requires_grad for tensor in (x, weight, cos, sin)):
        raise RuntimeError("fused_qknorm_rope is inference-only and does not implement autograd")
    if torch.compiler.is_compiling():
        return torch.ops.fastvideo._minimax_h3_qknorm_rope(x, weight, cos, sin, eps)
    return _fused_qknorm_rope_impl(x, weight, cos, sin, eps)

def fused_residual_gate_rmsnorm_modulate(
    residual: torch.Tensor,
    branch: torch.Tensor,
    gate: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    shift: torch.Tensor,
    index: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fuse residual update, row-indexed gate, RMSNorm, and modulation.

    ``index`` values must lie in ``[0, table_rows)``; see
    :func:`fused_rmsnorm_modulate` for why the wrapper does not check them.
    """
    _require_forward_only(residual, branch, gate, weight, scale, shift)
    if torch.compiler.is_compiling():
        return torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate(
            residual,
            branch,
            gate,
            weight,
            scale,
            shift,
            index,
            eps,
        )
    return _fused_residual_gate_rmsnorm_modulate_impl(
        residual,
        branch,
        gate,
        weight,
        scale,
        shift,
        index,
        eps,
    )

def fused_rmsnorm_modulate(
    x: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    shift: torch.Tensor,
    index: torch.Tensor,
    eps: float,
) -> torch.Tensor:
    """Run RMSNorm and row-indexed modulation in one strict Triton kernel.

    ``index`` values must lie in ``[0, table_rows)``. Unlike eager
    ``index_select``, the kernel does not raise on out-of-range values (a
    device-side bounds check would synchronize); callers are safe by
    construction (``timestep_indices * 3 + token_tags``, SP pads with 0).
    """
    _require_forward_only(x, weight, scale, shift)
    if torch.compiler.is_compiling():
        return torch.ops.fastvideo._minimax_h3_rmsnorm_modulate(x, weight, scale, shift, index, eps)
    return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps)

def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
    """Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input.

    This is intentionally a strict kernel wrapper: callers own fallback policy and
    must only invoke it for a supported CUDA inference path.
    """
    if torch.is_grad_enabled() and x.requires_grad:
        raise RuntimeError("MiniMax H3 fused SwiGLU is forward-only and does not implement autograd")
    if torch.compiler.is_compiling():
        return torch.ops.fastvideo._minimax_h3_swiglu(x)
    return _minimax_h3_swiglu_impl(x)