[Feature] rewrite rope kernel; remove flashinfer dependencies (#18844)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
@@ -11,6 +11,7 @@
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
#ifndef USE_ROCM
|
||||
using fp32_t = float;
|
||||
@@ -62,6 +63,31 @@ SGL_DEVICE void PDLTriggerSecondary() {
|
||||
#endif
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Load data with the specified type and offset from a void pointer.
|
||||
* \tparam T The type to load.
|
||||
* \param ptr The base pointer.
|
||||
* \param offset The offset in number of elements of type T.
|
||||
*/
|
||||
template <typename T>
|
||||
SGL_DEVICE T load_as(const void* ptr, int64_t offset = 0) {
|
||||
return static_cast<const T*>(ptr)[offset];
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Store data with the specified type and offset to a void pointer.
|
||||
* \tparam T The type to store.
|
||||
* \param ptr The base pointer.
|
||||
* \param val The value to store.
|
||||
* \param offset The offset in number of elements of type T.
|
||||
* \note we use type_identity_t to force the caller to explicitly specify
|
||||
* the template parameter `T`, which can avoid accidentally using the wrong type.
|
||||
*/
|
||||
template <typename T>
|
||||
SGL_DEVICE void store_as(void* ptr, std::type_identity_t<T> val, int64_t offset = 0) {
|
||||
static_cast<T*>(ptr)[offset] = val;
|
||||
}
|
||||
|
||||
namespace pointer {
|
||||
|
||||
// we only allow void * pointer arithmetic for safety
|
||||
|
||||
Reference in New Issue
Block a user