7f2a703ed5
* Merge with private repo * Update README * Update README * Update README * Add PyTorch requirements * Fix sync scopes for MQA logits (#256) * Update README
42 lines
1.2 KiB
Plaintext
42 lines
1.2 KiB
Plaintext
#pragma once
|
|
|
|
#include <deep_gemm/common/exception.cuh>
|
|
|
|
namespace deep_gemm::layout {
|
|
|
|
constexpr static uint32_t kNumMaxRanks = 72;
|
|
|
|
template <uint32_t kNumRanks = kNumMaxRanks>
|
|
struct SymBuffer {
|
|
int64_t base;
|
|
int64_t offsets[kNumMaxRanks];
|
|
uint32_t rank_idx;
|
|
|
|
DG_STATIC_ASSERT(kNumRanks <= kNumMaxRanks, "Too many ranks");
|
|
|
|
SymBuffer() = default;
|
|
|
|
template <typename Container>
|
|
explicit SymBuffer(const Container& c, const uint32_t& rank_idx): rank_idx(rank_idx) {
|
|
const auto size = static_cast<uint32_t>(c.size());
|
|
base = c[rank_idx];
|
|
for (uint32_t i = 0; i < kNumMaxRanks; ++ i)
|
|
offsets[i] = i < size ? (c[i] - base) : 0;
|
|
}
|
|
|
|
#if defined(__CUDA_ARCH__) or defined(__CLION_IDE__)
|
|
template <typename ptr_t = void*>
|
|
CUTLASS_DEVICE ptr_t get_base_ptr() const {
|
|
return reinterpret_cast<ptr_t>(base);
|
|
}
|
|
|
|
template <typename ptr_t>
|
|
CUTLASS_DEVICE ptr_t map(const ptr_t& ptr, const uint32_t& dst_rank_idx) const {
|
|
int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast<int64_t>(ptr);
|
|
return *reinterpret_cast<ptr_t*>(&mapped_ptr);
|
|
}
|
|
#endif
|
|
};
|
|
|
|
} // namespace deep_gemm::layout
|