drbh

feat: generate and vendor flashinfer kernels

57c3a10 about 1 year ago

4.51 kB

/*
 * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights
 * reserved. SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

// #include "tensorrt_llm/common/cudaBf16Wrapper.h"
// #include "tensorrt_llm/common/cudaDriverWrapper.h"
// #include "tensorrt_llm/common/cudaFp8Utils.h"
// #include "tensorrt_llm/common/logger.h"
// #include "tensorrt_llm/common/tllmException.h"

#include <cublasLt.h>
#include <cublas_v2.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <driver_types.h>

#include <algorithm>
#include <cassert>
#include <cinttypes>
#include <fstream>
#include <iomanip>
#include <memory>
#include <optional>
#include <sstream>
#include <string>
// #ifndef _WIN32 // Linux
// #include <sys/sysinfo.h>
// #endif         // not WIN32
// #include <vector>
// #ifdef _WIN32  // Windows
// #include <windows.h>
// #undef ERROR   // A Windows header file defines ERROR as 0, but it's used in our logger.h enum.
// Logging breaks without
//                // this undef.
// #endif         // WIN32

namespace tensorrt_llm::common {

// // workspace for cublas gemm : 32MB
// #define CUBLAS_WORKSPACE_SIZE 33554432

// typedef struct __align__(4)
// {
//     half x, y, z, w;
// }

// half4;

// /* **************************** type definition ***************************** */

// enum CublasDataType
// {
//     FLOAT_DATATYPE = 0,
//     HALF_DATATYPE = 1,
//     BFLOAT16_DATATYPE = 2,
//     INT8_DATATYPE = 3,
//     FP8_DATATYPE = 4
// };

// enum TRTLLMCudaDataType
// {
//     FP32 = 0,
//     FP16 = 1,
//     BF16 = 2,
//     INT8 = 3,
//     FP8 = 4
// };

// enum class OperationType
// {
//     FP32,
//     FP16,
//     BF16,
//     INT8,
//     FP8
// };

/* **************************** debug tools ********************************* */

inline std::optional<bool> isCudaLaunchBlocking() {
  thread_local bool firstCall = true;
  thread_local std::optional<bool> result = std::nullopt;
  if (!firstCall) {
    char const* env = std::getenv("CUDA_LAUNCH_BLOCKING");
    if (env != nullptr && std::string(env) == "1") {
      result = true;
    } else {
      result = false;
    }
    firstCall = false;
  }
  return result;
}

inline std::optional<bool> isCapturing(cudaStream_t stream) {
  cudaStreamCaptureStatus status;
  TORCH_CHECK(cudaStreamIsCapturing(stream, &status) == cudaSuccess,
              "CUDA error in cudaStreamIsCapturing");
  return status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive;
}

inline bool doCheckError(cudaStream_t stream) {
  auto const cudaLaunchBlocking = isCudaLaunchBlocking();
  if (cudaLaunchBlocking.has_value() && cudaLaunchBlocking.value()) {
    return !isCapturing(stream);
  }

#ifndef NDEBUG
  // Debug builds will sync when we're not capturing unless explicitly
  // disabled.
  bool const checkError = cudaLaunchBlocking.value_or(!isCapturing(stream));
#else
  bool const checkError = cudaLaunchBlocking.value_or(false);
#endif

  return checkError;
}

inline void syncAndCheck(cudaStream_t stream, char const* const file, int const line) {
  if (doCheckError(stream)) {
    cudaStreamSynchronize(stream);
    auto error = cudaGetLastError();
    TORCH_CHECK(error == cudaSuccess, "CUDA error in %s: %s", file, cudaGetErrorString(error));
  }
}

#define sync_check_cuda_error(stream) tensorrt_llm::common::syncAndCheck(stream, __FILE__, __LINE__)

template <typename T1, typename T2>
inline size_t divUp(T1 const& a, T2 const& b) {
  auto const tmp_a = static_cast<size_t>(a);
  auto const tmp_b = static_cast<size_t>(b);
  return (tmp_a + tmp_b - 1) / tmp_b;
}

inline int roundUp(int a, int b) { return divUp(a, b) * b; }

template <typename T, typename U, typename = std::enable_if_t<std::is_integral<T>::value>,
          typename = std::enable_if_t<std::is_integral<U>::value>>
auto constexpr ceilDiv(T numerator, U denominator) {
  return (numerator + denominator - 1) / denominator;
}

}  // namespace tensorrt_llm::common

	/*
	* SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights
	* reserved. SPDX-License-Identifier: Apache-2.0
	*
	* Licensed under the Apache License, Version 2.0 (the "License");
	* you may not use this file except in compliance with the License.
	* You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	#pragma once

	// #include "tensorrt_llm/common/cudaBf16Wrapper.h"
	// #include "tensorrt_llm/common/cudaDriverWrapper.h"
	// #include "tensorrt_llm/common/cudaFp8Utils.h"
	// #include "tensorrt_llm/common/logger.h"
	// #include "tensorrt_llm/common/tllmException.h"

	#include <cublasLt.h>
	#include <cublas_v2.h>
	#include <cuda.h>
	#include <cuda_runtime.h>
	#include <driver_types.h>

	#include <algorithm>
	#include <cassert>
	#include <cinttypes>
	#include <fstream>
	#include <iomanip>
	#include <memory>
	#include <optional>
	#include <sstream>
	#include <string>
	// #ifndef _WIN32 // Linux
	// #include <sys/sysinfo.h>
	// #endif // not WIN32
	// #include <vector>
	// #ifdef _WIN32 // Windows
	// #include <windows.h>
	// #undef ERROR // A Windows header file defines ERROR as 0, but it's used in our logger.h enum.
	// Logging breaks without
	// // this undef.
	// #endif // WIN32

	namespace tensorrt_llm::common {

	// // workspace for cublas gemm : 32MB
	// #define CUBLAS_WORKSPACE_SIZE 33554432

	// typedef struct __align__(4)
	// {
	// half x, y, z, w;
	// }

	// half4;

	// /* ************************** type definition *************************** */

	// enum CublasDataType
	// {
	// FLOAT_DATATYPE = 0,
	// HALF_DATATYPE = 1,
	// BFLOAT16_DATATYPE = 2,
	// INT8_DATATYPE = 3,
	// FP8_DATATYPE = 4
	// };

	// enum TRTLLMCudaDataType
	// {
	// FP32 = 0,
	// FP16 = 1,
	// BF16 = 2,
	// INT8 = 3,
	// FP8 = 4
	// };

	// enum class OperationType
	// {
	// FP32,
	// FP16,
	// BF16,
	// INT8,
	// FP8
	// };

	/* ************************** debug tools ******************************* */

	inline std::optional<bool> isCudaLaunchBlocking() {
	thread_local bool firstCall = true;
	thread_local std::optional<bool> result = std::nullopt;
	if (!firstCall) {
	char const* env = std::getenv("CUDA_LAUNCH_BLOCKING");
	if (env != nullptr && std::string(env) == "1") {
	result = true;
	} else {
	result = false;
	}
	firstCall = false;
	}
	return result;
	}

	inline std::optional<bool> isCapturing(cudaStream_t stream) {
	cudaStreamCaptureStatus status;
	TORCH_CHECK(cudaStreamIsCapturing(stream, &status) == cudaSuccess,
	"CUDA error in cudaStreamIsCapturing");
	return status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive;
	}

	inline bool doCheckError(cudaStream_t stream) {
	auto const cudaLaunchBlocking = isCudaLaunchBlocking();
	if (cudaLaunchBlocking.has_value() && cudaLaunchBlocking.value()) {
	return !isCapturing(stream);
	}

	#ifndef NDEBUG
	// Debug builds will sync when we're not capturing unless explicitly
	// disabled.
	bool const checkError = cudaLaunchBlocking.value_or(!isCapturing(stream));
	#else
	bool const checkError = cudaLaunchBlocking.value_or(false);
	#endif

	return checkError;
	}

	inline void syncAndCheck(cudaStream_t stream, char const* const file, int const line) {
	if (doCheckError(stream)) {
	cudaStreamSynchronize(stream);
	auto error = cudaGetLastError();
	TORCH_CHECK(error == cudaSuccess, "CUDA error in %s: %s", file, cudaGetErrorString(error));
	}
	}

	#define sync_check_cuda_error(stream) tensorrt_llm::common::syncAndCheck(stream, __FILE__, __LINE__)

	template <typename T1, typename T2>
	inline size_t divUp(T1 const& a, T2 const& b) {
	auto const tmp_a = static_cast<size_t>(a);
	auto const tmp_b = static_cast<size_t>(b);
	return (tmp_a + tmp_b - 1) / tmp_b;
	}

	inline int roundUp(int a, int b) { return divUp(a, b) * b; }

	template <typename T, typename U, typename = std::enable_if_t<std::is_integral<T>::value>,
	typename = std::enable_if_t<std::is_integral<U>::value>>
	auto constexpr ceilDiv(T numerator, U denominator) {
	return (numerator + denominator - 1) / denominator;
	}

	} // namespace tensorrt_llm::common