Prefill PLE: STRATA_PLE_COMPACT=1 uploads the raw IQ4_NL rows and dequantizes them on the device

1,440 bytes a token over PCIe instead of 10,240 floats' bytes; the same floats as the host dequant (bitwise).
On a PCIe Gen2 x4 card (CMP 170HX) the cmp_bench3 prompts (4K and 16K) were +1.0..+1.7 % faster. Opt-in.

- the raw-row buffer's cudaFree and the error clear after a failed cudaMalloc are (void): HIP's hipError_t is
  [[nodiscard]]
- metal: ple_iq4nl_dequant is a stub (stubs_0145.cpp); the Metal prompt runs through its window and never reaches it

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

include/strata/kernels/ngram.hpp

@@ -187,6 +187,10 @@ public:
     /// as ONE reader request - page dedupe and sort across the whole batch, the reader's full queue depth.  Not
     /// while a single-token `issue` is pending.  The mapped mode gathers row by row.
     bool gather_batch(const uint32_t* rows, size_t n_tokens, float* out, std::string& err);
+    /// STRATA_PLE_COMPACT: the same rows as their raw IQ4_NL bytes (n_tokens x 16 rows of PLE_ROW_BYTES; a row past the
+    /// table: zeros, which decode to the zeros gather_batch writes) for ple_iq4nl_dequant on the device.  False (nothing
+    /// written) unless the table is in memory (--ple-io ram / mmap) and IQ4_NL.
+    bool gather_raw_iq4nl(const uint32_t* rows, size_t n_tokens, uint8_t* out) const;
 
     /// Fault injection (Direct mode): every row read completes no earlier than `us` after issue.
     void set_injected_delay_us(double us);

include/strata/kernels/ple.hpp

@@ -125,4 +125,8 @@ void ple_history_advance(float* hist, const float* normalized, void* stream);
 /// Whether the CUDA path is available.  The host hash and the table read work without it; the block does not.
 bool ple_block_available();
 
+/// STRATA_PLE_COMPACT: `rows` raw IQ4_NL PLE rows (PleTable::gather_raw_iq4nl, PLE_ROW_BYTES each) to 160 floats each,
+/// the floats iq4nl_dequant_row writes (d times the code's value, one multiply), on `stream`.
+void ple_iq4nl_dequant(const uint8_t* raw, int64_t rows, float* out, void* stream);
+
 }  // namespace strata::kernels

metal/src/engine/stubs_0145.cpp

@@ -149,3 +149,5 @@ namespace strata::kernels { bool native_q8_0_sm60_pack(void const* p0, int p1, i
 namespace strata::kernels { void native_q8_0_sm60_release(void const* p0) { } }
 // only reached with pf_wmma (off on Metal)
 namespace strata::prefill { void swiglu_split_f16(float const* p0, unsigned short* p1, long long p2, int p3, bool p4, void* p5) { } }
+// STRATA_PLE_COMPACT (opt-in, CUDA / HIP prompt path): prefill.cpp references it; the Metal prompt runs through its window and never reaches it.
+namespace strata::kernels { void ple_iq4nl_dequant(const uint8_t* p0, int64_t p1, float* p2, void* p3) { strata::metal_engine::missing("strata::kernels::ple_iq4nl_dequant"); } }

src/kernels/cuda/ple.cu

@@ -23,6 +23,7 @@
 #include "strata/kernels/native_ple_postops.hpp"
 
 #include <cuda_runtime.h>
+#include <cuda_fp16.h>
 
 #include <cmath>
 #include <cstdio>
@@ -387,4 +388,26 @@ void ple_block_projected(const float* projected_key, const float* projected_valu
     native_ple_postops(projected_key, hidden, projected_value, hist_rows, w, buffers, stream);
 }
 
+namespace {
+__constant__ int8_t kPleIq4Nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113};
+// one thread a value: row r, block b (18 bytes: fp16 d, 16 bytes of codes), element j of 32 (j < 16: low nibble of
+// qs[j], else the high nibble of qs[j - 16] - the split halves of iq4nl_dequant_row)
+__global__ void ple_iq4nl_dequant_kernel(const uint8_t* __restrict__ raw, int64_t n, float* __restrict__ out) {
+    const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    const int64_t r = i / PLE_HEAD_DIM;
+    const int e = (int) (i - r * PLE_HEAD_DIM), b = e >> 5, j = e & 31;
+    const uint8_t* blk = raw + r * PLE_ROW_BYTES + b * 18;
+    const float d = __half2float(__ushort_as_half((unsigned short) (blk[0] | (blk[1] << 8))));
+    const uint8_t q = blk[2 + (j & 15)];
+    out[i] = d * (float) kPleIq4Nl[j < 16 ? (q & 15) : (q >> 4)];
+}
+}  // namespace
+
+void ple_iq4nl_dequant(const uint8_t* raw, int64_t rows, float* out, void* stream) {
+    const int64_t n = rows * PLE_HEAD_DIM;
+    if (n <= 0) return;
+    ple_iq4nl_dequant_kernel<<<(unsigned) ((n + 255) / 256), 256, 0, (cudaStream_t) stream>>>(raw, n, out);
+}
+
 }  // namespace strata::kernels

src/kernels/ngram.cpp

@@ -617,6 +617,19 @@ bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, s
     return true;
 }
 
+bool PleTable::gather_raw_iq4nl(const uint32_t* rows, size_t n_tokens, uint8_t* out) const {
+    std::lock_guard<std::recursive_mutex> lk(impl_->mu);
+    if (impl_->data == nullptr || impl_->fmt->id != PleFormat::IQ4_NL || impl_->rb != PLE_ROW_BYTES) return false;
+    const size_t n = n_tokens * (size_t) PLE_N_HEADS;
+    for (size_t i = 0; i < n; ++i) {
+        uint8_t* o = out + i * PLE_ROW_BYTES;
+        if (rows[i] >= impl_->n_rows) std::memset(o, 0, PLE_ROW_BYTES);
+        else std::memcpy(o, impl_->data + (size_t) rows[i] * PLE_ROW_BYTES, PLE_ROW_BYTES);
+    }
+    impl_->bytes_read += (uint64_t) n * PLE_ROW_BYTES;
+    return true;
+}
+
 void PleTable::set_injected_delay_us(double us) { impl_->reader.set_injected_delay_us(us); }
 
 std::string PleTable::io_report() const {

src/prefill/prefill.cpp

@@ -127,6 +127,14 @@ bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* con
 }  // namespace strata::metal_engine
 #endif
 
+namespace {
+// STRATA_PLE_COMPACT=1 (opt-in): the prompt chunk's PLE rows go up raw and are dequantized on the device (IQ4_NL tables
+// held in memory); the same floats as the host dequant, 7x fewer bytes over PCIe
+bool ple_compact_on() {
+    static const bool on = [] { const char* v = std::getenv("STRATA_PLE_COMPACT"); return v != nullptr && v[0] == '1'; }();
+    return on;
+}
+}  // namespace
 namespace strata::prefill {
 namespace {
 
@@ -982,6 +990,8 @@ struct Prefill::Impl {
     float* ple_emb_host[2] = {};             // pinned, double-buffered: the next chunk's rows are read while this
     cudaEvent_t ple_copied[2] = {};          // one runs; the event marks that buffer's upload done
     std::vector<uint32_t> ple_rows[2];
+    bool ple_compact[2] = {false, false};    // STRATA_PLE_COMPACT: the pinned buffer holds raw IQ4_NL rows
+    uint8_t* ple_raw_dev = nullptr;          // their device copy (cudaMalloc'd, the chunk's 16 rows a token)
     float* ple_norm = nullptr;
     uint8_t* region = nullptr;               // the attention/MoE scratch region (idle while the PLE block runs)
     uint64_t region_bytes = 0;
@@ -1080,6 +1090,7 @@ void Prefill::release() {
             if (e) cudaEventDestroy(e);
         if (impl_->ple_copied[b]) cudaEventDestroy(impl_->ple_copied[b]);
         if (impl_->ple_emb_host[b] && impl_->ple_pageable[b].empty()) cudaFreeHost(impl_->ple_emb_host[b]);
+        if (b == 0 && impl_->ple_raw_dev) { (void) cudaFree(impl_->ple_raw_dev); impl_->ple_raw_dev = nullptr; }
     }
     if (impl_->copy) cudaStreamDestroy(impl_->copy);
     if (impl_->kv_copy) cudaStreamDestroy(impl_->kv_copy);
@@ -1508,6 +1519,11 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co
         if (!m.ple_copied[b] && cudaEventCreateWithFlags(&m.ple_copied[b], cudaEventDisableTiming) != cudaSuccess)
             ok = false;
         m.ple_rows[b].resize(T * strata::kernels::PLE_N_HEADS);
+        if (ple_compact_on() && m.ple_raw_dev == nullptr &&
+            cudaMalloc((void**) &m.ple_raw_dev, (size_t) T * strata::kernels::PLE_N_HEADS * strata::kernels::PLE_ROW_BYTES) != cudaSuccess) {
+            m.ple_raw_dev = nullptr;
+            (void) cudaGetLastError();
+        }
     }
     if (ss.qsa_states[ss.qsa_primary()].kv_mode == 1) {   // KV streaming: the staging pool's identity page table
         const int64_t pages = ss.qsa_states[ss.qsa_primary()].n_pages;
@@ -2677,6 +2693,14 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
             pv[0] = pv[1];
             pv[1] = tok;
         }
+        // STRATA_PLE_COMPACT=1: the raw rows (1,440 bytes a token instead of 10,240 floats' bytes) into the same pinned
+        // buffer; ple_land uploads them and dequantizes on the device
+        if (ple_compact_on() && m.ple_pageable[buf].empty() && m.ple_raw_dev != nullptr &&
+            ss.ple.table->gather_raw_iq4nl(m.ple_rows[buf].data(), (size_t) T, (uint8_t*) m.ple_emb_host[buf])) {
+            m.ple_compact[buf] = true;
+            return true;
+        }
+        m.ple_compact[buf] = false;
         return ss.ple.table->gather_batch(m.ple_rows[buf].data(), (size_t) T, m.ple_emb_host[buf], e);
     };
     std::string ple_next_err;
@@ -2810,7 +2834,10 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
                 return false;
             }
             // #1996: the pageable fallback (no RAM to pin) goes up through a pinned slot; the pinned buffers as before
-            if ((m.ple_pageable[ple_buf].empty()
+            const size_t raw_bytes = (size_t) T * strata::kernels::PLE_N_HEADS * strata::kernels::PLE_ROW_BYTES;
+            if ((m.ple_compact[ple_buf]
+                     ? cudaMemcpyAsync(m.ple_raw_dev, m.ple_emb_host[ple_buf], raw_bytes, cudaMemcpyHostToDevice, m.cs)
+                     : m.ple_pageable[ple_buf].empty()
                      ? cudaMemcpyAsync(m.ple_emb, m.ple_emb_host[ple_buf], (size_t) T * N * 4, cudaMemcpyHostToDevice, m.cs)
                      : m.up_ring.h2d(m.ple_emb, m.ple_emb_host[ple_buf], (size_t) T * N * 4, m.cs)) !=
                     cudaSuccess ||
@@ -2818,6 +2845,8 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
                 err = std::string("prefill: the PLE rows' upload failed: ") + cudaGetErrorString(cudaGetLastError());
                 return false;
             }
+            if (m.ple_compact[ple_buf])
+                strata::kernels::ple_iq4nl_dequant(m.ple_raw_dev, T * strata::kernels::PLE_N_HEADS, m.ple_emb, m.cs);
             if (c0 + C < n) {
                 // the other buffer's upload (a chunk ago) is done before the SSD thread refills it
                 if (cudaEventSynchronize(m.ple_copied[ple_buf ^ 1]) != cudaSuccess) {