File size: 9,456 Bytes
bbb6388 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | #include "strata/core/native_head.hpp"
#include "strata/artifact/gguf_reader.hpp"
#include "strata/kernels/iq_kernels.hpp"
#include "strata/kernels/native_mmvq.hpp"
#include <cuda_runtime.h>
#include <climits>
#include <cstdio>
#include <cstring>
#include <exception>
namespace strata::core {
NativeHead::~NativeHead() {
if (scratch_) cudaFree(scratch_);
if (weights_) cudaFree(weights_);
}
bool NativeHead::load(const std::vector<std::string>& shards, int64_t n_in, int64_t n_out, std::string& err) {
if (loaded()) { err = "native head is already loaded"; return false; }
if (n_in <= 0 || n_out <= 0 || n_in > INT_MAX || n_out > INT_MAX || n_in % 256) {
err = "native head requires positive int32 dimensions and whole 256-value rows";
return false;
}
try {
// The architecture is the metadata shard's; output.weight comes from whichever shard holds it (shard 2
// of Unsloth's UD-Q4_K_XL, whose shard 1 holds no tensor). GgufModel refuses a duplicate across shards.
const strata::GgufModel model(shards);
err = strata::check_architecture(model.meta());
if (!err.empty()) return false;
size_t at = 0;
const strata::TensorInfo* tensor = model.find("output.weight", &at);
const strata::GgufFile& gguf = model.shard(at);
if (!tensor || !strata::kernels::native_mmvq_supported((int) tensor->type) || tensor->shape.size() != 2 ||
tensor->shape[0] != (uint64_t) n_in || tensor->shape[1] != (uint64_t) n_out) {
err = "native head: expected a natively supported output.weight with the canonical head dimensions";
return false;
}
const uint64_t bytes = strata::kernels::native_mmvq_weight_bytes((int) tensor->type, (int) n_in, (int) n_out);
const uint64_t payload = gguf.file_size() - gguf.data_start();
if (tensor->offset > payload || bytes > payload - tensor->offset) {
err = "native head: truncated output.weight payload";
return false;
}
void* weights = nullptr;
void* scratch = nullptr;
cudaError_t status = cudaMalloc(&weights, bytes);
if (status == cudaSuccess)
status = cudaMalloc(&scratch, strata::kernels::native_q8_1_bytes((int) n_in, 1));
if (status == cudaSuccess)
status = cudaMemcpy(weights, gguf.tensor_data(*tensor), bytes, cudaMemcpyHostToDevice);
if (status != cudaSuccess) {
if (scratch) cudaFree(scratch);
if (weights) cudaFree(weights);
err = std::string("native head upload: ") + cudaGetErrorString(status);
return false;
}
weights_ = weights;
scratch_ = scratch;
bytes_ = bytes;
n_in_ = (int) n_in;
n_out_ = (int) n_out;
type_ = (int) tensor->type;
return true;
} catch (const std::exception& error) {
err = std::string("native head: ") + error.what();
return false;
}
}
bool NativeHead::run(const float* mixed, float* logits, void* stream, std::string& err) const {
if (!loaded() || !mixed || !logits || !stream) {
err = "native head requires loaded weights, device buffers and an explicit stream";
return false;
}
try {
if (type_ == 13) {
strata::kernels::native_q5_k_f32(weights_, mixed, scratch_, logits, n_in_, n_out_, 1, stream);
} else {
strata::kernels::native_quantize_q8_1(mixed, scratch_, n_in_, 1, stream);
strata::kernels::native_mmvq(type_, weights_, scratch_, logits, n_in_, n_out_, 1, stream);
}
} catch (const std::exception& error) {
err = std::string("native head launch: ") + error.what();
return false;
}
const cudaError_t status = cudaPeekAtLastError();
if (status != cudaSuccess) {
err = std::string("native head launch: ") + cudaGetErrorString(status);
return false;
}
return true;
}
// ================================ plan v0.3 P6: THE NATIVE EMBEDDING ================================
namespace {
const NativeEmbed* g_embed = nullptr;
}
void set_native_embed(const NativeEmbed* e) { g_embed = e; }
const NativeEmbed* native_embed() { return g_embed; }
NativeEmbed::~NativeEmbed() {
if (host_) cudaFreeHost(host_);
else if (dev_) cudaFree(const_cast<void*>(dev_)); // the VRAM fallback below
}
bool NativeEmbed::load(const std::vector<std::string>& shards, int64_t n_embd, int64_t n_vocab, std::string& err) {
try {
const strata::GgufModel model(shards);
// --embd-gguf's one-tensor file (tools/embd_bf16_pack.py) says "strata-embd": only its tensor is checked
const strata::MetaValue* arch = model.meta().get("general.architecture");
err = arch != nullptr && arch->s == "strata-embd" ? std::string() : strata::check_architecture(model.meta());
if (!err.empty()) { err = "native embedding: " + err; return false; }
size_t at = 0;
const strata::TensorInfo* t = model.find("token_embd.weight", &at);
const strata::GgufFile& gguf = model.shard(at);
if (!t || t->shape.size() != 2 || t->shape[0] != (uint64_t) n_embd || t->shape[1] != (uint64_t) n_vocab ||
!strata::kernels::embed_type_supported((int) t->type) || n_embd % 256) {
err = "native embedding: token_embd.weight is absent, of another shape, or of a type without a GPU "
"dequantizer";
return false;
}
row_ = strata::kernels::iq_row_bytes((int) t->type, n_embd);
bytes_ = (uint64_t) row_ * (uint64_t) n_vocab;
// the table is copied out of the mapping below: a truncated shard must be an error, not a read past EOF
if (!model.in_bounds(*t, at) || strata::tensor_payload_bytes(*t) != bytes_) {
err = "native embedding: token_embd.weight's payload is truncated or not " + std::to_string(bytes_) +
" B (" + gguf.path() + ")";
bytes_ = 0;
return false;
}
if (cudaHostAlloc(&host_, bytes_, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess) {
// Under WSL2 the driver's pinned/mapped host budget (~1 GiB) can be spent by the GPU contexts
// themselves (three cards). The table is only gathered from, so keep it in the current device's VRAM
// instead: it costs its size there and reads faster than over PCIe.
cudaGetLastError();
host_ = nullptr;
void* d = nullptr;
if (cudaMalloc(&d, bytes_) != cudaSuccess ||
cudaMemcpy(d, gguf.tensor_data(*t), bytes_, cudaMemcpyHostToDevice) != cudaSuccess) {
if (d) cudaFree(d);
cudaGetLastError();
err = "native embedding: cannot pin " + std::to_string(bytes_ >> 20) + " MiB, nor place it in VRAM";
return false;
}
std::fprintf(stderr, "strata: native embedding: cannot pin %llu MiB, kept in VRAM instead\n",
(unsigned long long) (bytes_ >> 20));
dev_ = d;
} else {
std::memcpy(host_, gguf.tensor_data(*t), bytes_);
void* d = nullptr;
#if defined(STRATA_USE_HIP)
// #325: the Windows HIP stack can refuse the device alias of a mapped allocation (and, when it gives
// one, it is the host address itself - unified addressing; kernels read it correctly there, a
// device-to-device copy into it does not land: tests/hip/mapped_alias). The table is only gathered from,
// so without an alias it goes into VRAM like the unpinnable case above, instead of failing the start.
if (cudaHostGetDevicePointer(&d, host_, 0) != cudaSuccess || d == nullptr) {
cudaGetLastError();
d = nullptr;
if (cudaMalloc(&d, bytes_) != cudaSuccess ||
cudaMemcpy(d, gguf.tensor_data(*t), bytes_, cudaMemcpyHostToDevice) != cudaSuccess) {
if (d) cudaFree(d);
cudaGetLastError();
err = "native embedding: no device alias for the mapped table, and no VRAM to copy it into";
return false;
}
cudaFreeHost(host_); // the destructor frees dev_ when host_ is null
host_ = nullptr;
std::fprintf(stderr, "strata: native embedding: no device alias for the mapped table, kept in VRAM\n");
}
#else
if (cudaHostGetDevicePointer(&d, host_, 0) != cudaSuccess) {
err = "native embedding: no device alias for the mapped table";
return false;
}
#endif
dev_ = d;
}
type_ = (int) t->type;
n_embd_ = n_embd;
n_vocab_ = n_vocab;
return true;
} catch (const std::exception& e) {
err = std::string("native embedding: ") + e.what();
return false;
}
}
void NativeEmbed::gather_dev(const int32_t* tokens, int64_t n_tok, float* out, void* stream) const {
strata::kernels::iq_embed_rows(type_, dev_, row_, tokens, n_tok, n_embd_, out, stream);
}
void NativeEmbed::gather_one(int64_t token, float* out, void* stream) const {
strata::kernels::iq_dequant_f32(type_, (const uint8_t*) dev_ + (size_t) token * row_, n_embd_, out, stream);
}
} // namespace strata::core
|