Native head: STRATA_EMBED_VRAM=1 keeps the token-embedding table in VRAM

The table lived in mapped pinned host memory, so every prompt's row gather crossed the link. With the env, when
512 MiB stay free beside it, the table goes to VRAM (322 MiB for an IQ4_XS table) through the calls the existing
cannot-pin path uses; same bytes, same gather kernel. The start-up line says where the table is ("in VRAM" also for
the existing VRAM fallbacks, which it called mapped host memory). The failed copy's cudaFree and the error clear are
(void) (HIP's [[nodiscard]] hipError_t).

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>

include/strata/core/native_head.hpp

@@ -43,7 +43,8 @@ private:
 
 /// Plan v0.3 P6: `token_embd.weight` in its GGUF form (the IQ model files), in mapped pinned host memory: a row
 /// is read over PCIe per token, so the table costs no VRAM.  `embed_row`, the verify window and the MTP drafter
-/// use it instead of the canonical table when it is set (`set_native_embed`).
+/// use it instead of the canonical table when it is set (`set_native_embed`).  `STRATA_EMBED_VRAM=1` places the
+/// same bytes in VRAM instead (slow PCIe links); the gathers read `dev_` either way.
 class NativeEmbed {
 public:
     NativeEmbed() = default;
@@ -57,6 +58,8 @@ public:
     void gather_one(int64_t token, float* out, void* stream) const;
     uint64_t bytes() const { return bytes_; }
     int type() const { return type_; }
+    /// The table was placed in VRAM (STRATA_EMBED_VRAM=1, or the pinned / alias fallbacks) rather than mapped.
+    bool in_vram() const { return dev_ != nullptr && host_ == nullptr; }
 
 private:
     void* host_ = nullptr;

src/core/native_head.cpp

@@ -7,6 +7,7 @@
 #include <cuda_runtime.h>
 #include <climits>
 #include <cstdio>
+#include <cstdlib>
 #include <cstring>
 #include <exception>
 
@@ -146,6 +147,34 @@ bool NativeEmbed::load(const std::vector<std::string>& shards, int64_t n_embd, i
             return false;
         }
         strata::platform::advise_willneed(gguf.tensor_data(*t), bytes_);   // copied out of the mapping below
+        // STRATA_EMBED_VRAM=1: the table in the current device's VRAM instead of mapped host memory. The same
+        // bytes, read by the same gather kernels through dev_ (bitwise the same rows); what changes is where a
+        // prompt's scattered row reads go: over a PCIe Gen2 x4 link they cost more than the table's size in VRAM
+        // is worth. Taken only with the table plus 512 MiB free; otherwise, or when the upload fails, the mapped
+        // placement below is used as before (and the fallback is said).
+        if (const char* ev = std::getenv("STRATA_EMBED_VRAM"); ev != nullptr && ev[0] == '1') {
+            size_t free_b = 0, total_b = 0;
+            const bool have_free = cudaMemGetInfo(&free_b, &total_b) == cudaSuccess;
+            void* d = nullptr;
+            if (have_free && (uint64_t) free_b > bytes_ + ((uint64_t) 512 << 20) &&
+                cudaMalloc(&d, bytes_) == cudaSuccess) {
+                if (cudaMemcpy(d, gguf.tensor_data(*t), bytes_, cudaMemcpyHostToDevice) == cudaSuccess) {
+                    std::fprintf(stderr, "strata: native embedding: STRATA_EMBED_VRAM=1: %llu MiB placed in VRAM "
+                                         "(%llu MiB were free)\n",
+                                 (unsigned long long) (bytes_ >> 20), (unsigned long long) (free_b >> 20));
+                    dev_ = d;   // host_ stays null: the destructor frees dev_
+                    type_ = (int) t->type;
+                    n_embd_ = n_embd;
+                    n_vocab_ = n_vocab;
+                    return true;
+                }
+                (void) cudaFree(d);
+            }
+            (void) cudaGetLastError();
+            std::fprintf(stderr, "strata: native embedding: STRATA_EMBED_VRAM=1 not taken (%llu MiB free, the table "
+                                 "needs %llu MiB plus 512 MiB): kept in mapped host memory\n",
+                         (unsigned long long) (have_free ? free_b >> 20 : 0), (unsigned long long) (bytes_ >> 20));
+        }
         cudaError_t host_err = cudaHostAlloc(&host_, bytes_, cudaHostAllocMapped | cudaHostAllocPortable);
         if (host_err != cudaSuccess) {
             // Under WSL2 the driver's pinned/mapped host budget (~1 GiB) can be spent by the GPU contexts

src/program/generate.cpp

@@ -3717,9 +3717,11 @@ int main(int argc, char** argv) {
         }
         strata::core::set_native_embed(&native_embed);
         std::fprintf(stderr, "strata generate: native pack: %s experts (largest blob %.2f MB), token embedding "
-                             "%s in mapped host memory (%.0f MiB, %.1f s)\n",
+                             "%s %s (%.0f MiB, %.1f s)\n",
                      o.pack.c_str(), (double) strata::kernels::cpu::expert_layout().max_blob / 1e6,
-                     strata::ggml_type_name((uint32_t) native_embed.type()), (double) native_embed.bytes() / 1048576.0,
+                     strata::ggml_type_name((uint32_t) native_embed.type()),
+                     native_embed.in_vram() ? "in VRAM" : "in mapped host memory",
+                     (double) native_embed.bytes() / 1048576.0,
                      std::chrono::duration<double>(std::chrono::steady_clock::now() - embed_t0).count());
     }
     // Plan v0.3 P1: tensors served in native form are not also loaded in canonical form (~2.7 GB of VRAM back