[Perf][PLE] Serve FP8 checkpoint PLE rows from the native fast lane (#1343)

The native PLE request lane only accepted a CPU-owned GGUF IQ4_NL table, so
the Qwen3.8-Flash-Next NVFP4 checkpoint (raw E4M3 n-gram rows in 128 mapped
safetensors shards) went through the Python owner: socket, msgpack and Python
scheduling on every decode request, about 1.7 ms per MTP4 round.

The lane now addresses rows through a list of equally sized mapped shards
(one shard is the old contiguous table) and validates rows only for IQ4_NL.
The worker admits whole-table FP8 disk layers and hands the lane their shard
pointers; results stay raw rows in n-gram id order, as the Python owner
writes them. The existing hashing, row cache, parallel cold-page reads and
draft-token page prefetch apply unchanged. The op gains three defaulted
arguments; GGUF callers are unchanged.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

csrc/ops.h

@@ -979,7 +979,8 @@ int64_t ple_fast_lane_start(
     std::vector<int64_t> offsets, int64_t eos, std::vector<int64_t> result_ptrs,
     std::vector<int64_t> flag_ptrs, int64_t cache_bytes, int64_t idle_spin_us,
     bool clamp_tokens, std::optional<torch::Tensor> draft_ring,
-    int64_t ring_reqs, int64_t ring_spec);
+    int64_t ring_reqs, int64_t ring_spec, std::vector<int64_t> shard_ptrs,
+    int64_t shard_rows, bool iq4_rows);
 torch::Tensor ple_fast_lane_stats(int64_t handle);
 
 std::vector<torch::Tensor> gguf_dmv13_decode_pack(torch::Tensor raw,

csrc/ple_fast_lane.cpp

@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 //
-// Native PLE request lane for CPU-owned packed IQ4_NL n-gram tables.
+// Native PLE request lane for CPU-owned n-gram tables: packed IQ4_NL rows of a
+// GGUF table, or raw FP8 rows of a checkpoint split into equal mapped shards.
 //
 // TP0's model stream publishes each request (token IDs, query offsets and
 // n-gram context) into a host-mapped buffer and releases a sequence number.
@@ -185,9 +186,13 @@ struct LaneConfig {
   int64_t* request;  // host-mapped request buffer (int64)
   int64_t max_reqs;
   int64_t max_tokens;
-  const uint8_t* table;
+  // Row r lives at shards[r / shard_rows] + (r % shard_rows) * row_bytes. A
+  // single shard holding every row describes one contiguous table.
+  std::vector<const uint8_t*> shards;
+  int64_t shard_rows;
   int64_t row_bytes;
   int64_t num_rows;
+  bool iq4_rows;  // validate packed IQ4_NL rows; raw FP8 rows are copied as is
   uint64_t multipliers[3];
   int64_t sizes[kHeads];
   int64_t offsets[kHeads];
@@ -227,6 +232,11 @@ class Lane {
     return __atomic_load_n(address, __ATOMIC_ACQUIRE);
   }
 
+  const uint8_t* row_ptr(int64_t row) const {
+    return config_.shards[row / config_.shard_rows] +
+           (row % config_.shard_rows) * config_.row_bytes;
+  }
+
   void run() {
 #ifdef __linux__
     prctl(PR_SET_TIMERSLACK, 1UL, 0, 0, 0);
@@ -370,8 +380,7 @@ class Lane {
           if (row < 0 || row >= config_.num_rows) continue;
           if (cache_.enabled() && cache_.contains(row)) continue;
 #ifdef __linux__
-          uintptr_t address = reinterpret_cast<uintptr_t>(
-              config_.table + row * config_.row_bytes);
+          uintptr_t address = reinterpret_cast<uintptr_t>(row_ptr(row));
           uintptr_t first = address / page * page;
           uintptr_t last = (address + config_.row_bytes - 1) / page * page;
           for (uintptr_t q = first; q <= last; q += page) {
@@ -472,8 +481,7 @@ class Lane {
     static const uintptr_t page =
         static_cast<uintptr_t>(sysconf(_SC_PAGE_SIZE));
     for (int64_t i : misses_) {
-      uintptr_t address =
-          reinterpret_cast<uintptr_t>(config_.table + rows_[i] * row_bytes);
+      uintptr_t address = reinterpret_cast<uintptr_t>(row_ptr(rows_[i]));
       uintptr_t first = address / page * page;
       uintptr_t last = (address + row_bytes - 1) / page * page;
       for (uintptr_t p = first; p <= last; p += page) {
@@ -487,11 +495,12 @@ class Lane {
     }
 #endif
     for (int64_t i : misses_) {
-      const uint8_t* source = config_.table + rows_[i] * row_bytes;
       uint8_t* target = out + i * row_bytes;
-      std::memcpy(target, source, row_bytes);
-      int64_t error = validate_row(target, row_bytes);
-      if (error != kOk) return error;
+      std::memcpy(target, row_ptr(rows_[i]), row_bytes);
+      if (config_.iq4_rows) {
+        int64_t error = validate_row(target, row_bytes);
+        if (error != kOk) return error;
+      }
       if (cache_.enabled()) cache_.insert(rows_[i], target);
     }
     return kOk;
@@ -522,16 +531,27 @@ int64_t ple_fast_lane_start(
     std::vector<int64_t> offsets, int64_t eos, std::vector<int64_t> result_ptrs,
     std::vector<int64_t> flag_ptrs, int64_t cache_bytes, int64_t idle_spin_us,
     bool clamp_tokens, std::optional<torch::Tensor> draft_ring,
-    int64_t ring_reqs, int64_t ring_spec) {
+    int64_t ring_reqs, int64_t ring_spec, std::vector<int64_t> shard_ptrs,
+    int64_t shard_rows, bool iq4_rows) {
   TORCH_CHECK(request.device().is_cpu() && request.scalar_type() == at::kLong &&
                   request.is_contiguous(),
               "PLE fast lane request must be a contiguous CPU int64 tensor");
   TORCH_CHECK(max_reqs > 0 && max_tokens > 0, "PLE fast lane needs capacity");
   TORCH_CHECK(request.numel() >= kHeader + max_reqs * 3 + 1 + max_tokens,
               "PLE fast lane request buffer is too small");
-  TORCH_CHECK(table_ptr != 0 && num_rows > 0 && row_bytes > 0 &&
-                  row_bytes % kBlockBytes == 0,
-              "PLE fast lane requires a packed IQ4_NL table");
+  if (shard_ptrs.empty()) {
+    shard_ptrs.push_back(table_ptr);
+    shard_rows = num_rows;
+  }
+  TORCH_CHECK(num_rows > 0 && row_bytes > 0 && shard_rows > 0 &&
+                  static_cast<int64_t>(shard_ptrs.size()) ==
+                      (num_rows + shard_rows - 1) / shard_rows,
+              "PLE fast lane shards must cover every row");
+  for (int64_t pointer : shard_ptrs) {
+    TORCH_CHECK(pointer != 0, "PLE fast lane shard pointers must be valid");
+  }
+  TORCH_CHECK(!iq4_rows || row_bytes % kBlockBytes == 0,
+              "PLE fast lane requires whole packed IQ4_NL blocks");
   TORCH_CHECK(multipliers.size() == 3 && sizes.size() == kHeads &&
                   offsets.size() == kHeads,
               "PLE fast lane supports trigram tables with 16 heads");
@@ -542,9 +562,13 @@ int64_t ple_fast_lane_start(
   config.request = request.data_ptr<int64_t>();
   config.max_reqs = max_reqs;
   config.max_tokens = max_tokens;
-  config.table = reinterpret_cast<const uint8_t*>(table_ptr);
+  for (int64_t pointer : shard_ptrs) {
+    config.shards.push_back(reinterpret_cast<const uint8_t*>(pointer));
+  }
+  config.shard_rows = shard_rows;
   config.row_bytes = row_bytes;
   config.num_rows = num_rows;
+  config.iq4_rows = iq4_rows;
   for (int i = 0; i < 3; ++i) {
     config.multipliers[i] = static_cast<uint64_t>(multipliers[i]);
   }

csrc/torch_bindings.cpp

@@ -232,7 +232,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int[] sizes, int[] offsets, int eos, int[] result_ptrs, "
       "int[] flag_ptrs, int cache_bytes, int idle_spin_us, "
       "bool clamp_tokens, Tensor? draft_ring, int ring_reqs, "
-      "int ring_spec) -> int");
+      "int ring_spec, int[] shard_ptrs=[], int shard_rows=0, "
+      "bool iq4_rows=True) -> int");
   ops.impl("ple_fast_lane_start", torch::kCPU, &ple_fast_lane_start);
   ops.def("ple_fast_lane_stats(int handle) -> Tensor");
   ops.impl("ple_fast_lane_stats", &ple_fast_lane_stats);

tests/v1/worker/test_ple_fast_lane.py

@@ -164,6 +164,77 @@ def test_native_lane_matches_python_owner(monkeypatch, cache_rows, clamp):
         lane.close()
 
 
+def _fp8_lane(shards, shard_rows, cache_rows, targets=2):
+    """A Lane over raw 160-byte FP8 rows split into separately allocated shards."""
+    lane = Lane.__new__(Lane)
+    lane.request = torch.zeros(request_shape(MAX_REQS, MAX_TOKENS), dtype=torch.long)
+    lane.results = [
+        torch.full((MAX_TOKENS, 16 * 160), 0xEE, dtype=torch.uint8)
+        for _ in range(targets)
+    ]
+    lane.flags = [torch.zeros(16, dtype=torch.int32) for _ in range(targets)]
+    lane.data = shards
+    lane.seq = 0
+    pointers = [shard.ctypes.data for shard in shards]
+    lane.handle = torch.ops._C.ple_fast_lane_start(
+        lane.request,
+        MAX_REQS,
+        MAX_TOKENS,
+        pointers[0],
+        160,
+        TABLE_ROWS,
+        MULTIPLIERS,
+        SIZES,
+        OFFSETS,
+        EOS,
+        [r.data_ptr() for r in lane.results],
+        [f.data_ptr() for f in lane.flags],
+        cache_rows * 160,
+        1000,
+        True,
+        None,
+        0,
+        0,
+        pointers,
+        shard_rows,
+        False,
+    )
+    return lane
+
+
+@pytest.mark.parametrize("cache_rows", [0, 1 << 16])
+def test_native_lane_serves_sharded_fp8_rows(monkeypatch, cache_rows):
+    monkeypatch.setattr(ple_module, "is_offload_process", lambda: True)
+    rng = np.random.default_rng(1339)
+    full = rng.integers(0, 256, (TABLE_ROWS, 160), dtype=np.uint8)
+    full[::7, :2] = 0x7F  # FP8 NaN bytes: copied as is, never validated as IQ4
+    shard_rows = 997  # the last shard is shorter
+    shards = [
+        np.array(full[i : i + shard_rows])  # separate allocations
+        for i in range(0, TABLE_ROWS, shard_rows)
+    ]
+    module = ngram_module()
+    lane = _fp8_lane(shards, shard_rows, cache_rows)
+    try:
+        for _ in range(16):
+            reqs = int(rng.integers(1, MAX_REQS + 1))
+            lengths = rng.integers(1, 6, reqs)
+            starts = [0, *np.cumsum(lengths).tolist()]
+            tokens = rng.integers(0, 200000, starts[-1]).tolist()
+            context = rng.integers(0, 200000, (reqs, 2)).tolist()
+            rows = module.compute_ngram_ids(
+                torch.tensor(tokens, dtype=torch.int32),
+                torch.tensor(starts, dtype=torch.int32),
+                torch.tensor(context, dtype=torch.int32),
+            )
+            want = torch.from_numpy(full[rows.numpy()].reshape(len(tokens), -1))
+            for out in lane.submit(tokens, starts, context):
+                assert torch.equal(out, want)
+        assert lane.stats()[6] == 0  # no error
+    finally:
+        lane.close()
+
+
 def test_native_lane_waits_for_consumers_and_reports_bad_rows(monkeypatch):
     monkeypatch.setattr(ple_module, "is_offload_process", lambda: True)
     data = table()

vllm/v1/ple_offload/worker.py

@@ -935,11 +935,14 @@ class PleOffloadRunner:
         if set(self._worker_targets) != {0} or len(self._layers) != 1:
             return "requires_one_dp_rank_and_one_ple_layer"
         layer = next(iter(self._layers.values()))
-        reader = getattr(getattr(layer, "ngram_embedding", None), "_cpu_reader", None)
-        if reader is None or reader.source_type != 20:
-            return "requires_cpu_owned_iq4nl_rows"
-        if getattr(layer, "_packed_result_layout", None) is None:
-            return "requires_packed_results"
+        if self._fast_lane_fp8_shards(layer) is None:
+            reader = getattr(
+                getattr(layer, "ngram_embedding", None), "_cpu_reader", None
+            )
+            if reader is None or reader.source_type != 20:
+                return "requires_cpu_owned_iq4nl_rows"
+            if getattr(layer, "_packed_result_layout", None) is None:
+                return "requires_packed_results"
         if getattr(layer, "_cascade", False):
             # Packed rows split across device/host tiers merge in Python.
             return "tiered_cascade_uses_the_python_owner"
@@ -950,6 +953,24 @@ class PleOffloadRunner:
             return "requires_mapped_result_transport"
         return None
 
+    @staticmethod
+    def _fast_lane_fp8_shards(layer) -> tuple[list[int], int, int, int] | None:
+        """(shard pointers, rows per shard, row bytes, rows) of a whole-table
+        FP8 disk layer whose results are raw rows in n-gram id order."""
+        pointers = list(getattr(layer, "_disk_shard_pointers", None) or [])
+        if (
+            not pointers
+            or getattr(layer, "_cascade", False)
+            or getattr(layer, "_packed_result_layout", None) is not None
+        ):
+            return None
+        return (
+            pointers,
+            int(layer._disk_shard_size),
+            int(layer.head_dim),
+            int(layer.ngram_embedding.org_vocab_size),
+        )
+
     @staticmethod
     def _native_draft_ring(registrations: list[PleOffloadRegistration]):
         """Draft ring arguments for the native lane's page prefetch."""
@@ -973,17 +994,25 @@ class PleOffloadRunner:
         assert registration.fast_lane_geometry is not None
         max_reqs, max_tokens = registration.fast_lane_geometry
         layer = next(iter(self._layers.values()))
-        reader = layer.ngram_embedding._cpu_reader
         targets = sorted(
             next(iter(self._worker_targets[0].values())), key=lambda t: t.tp_rank
         )
+        fp8 = self._fast_lane_fp8_shards(layer)
+        if fp8 is not None:
+            shard_ptrs, shard_rows, row_bytes, num_rows = fp8
+            table_ptr = shard_ptrs[0]
+        else:
+            reader = layer.ngram_embedding._cpu_reader
+            shard_ptrs, shard_rows = [], 0
+            table_ptr = reader.data.ctypes.data
+            row_bytes, num_rows = reader.data.shape[1], reader.logical_rows
         handle = torch.ops._C.ple_fast_lane_start(
             request,
             max_reqs,
             max_tokens,
-            reader.data.ctypes.data,
-            reader.data.shape[1],
-            reader.logical_rows,
+            table_ptr,
+            row_bytes,
+            num_rows,
             layer.layer_multipliers.tolist(),
             layer.ngram_heads_vocab_sizes.tolist(),
             layer.ngram_heads_offsets.tolist(),
@@ -994,10 +1023,15 @@ class PleOffloadRunner:
             50_000,
             self._clamp_input_ids,
             *self._native_draft_ring(registrations),
+            shard_ptrs,
+            shard_rows,
+            fp8 is None,
         )
         self._fast_lane = FastLaneOwner(request, handle)
         logger.info(
-            "PLE fast lane serving %d TP result buffer(s) natively.", len(targets)
+            "PLE fast lane serving %d TP result buffer(s) natively (%s rows).",
+            len(targets),
+            "IQ4_NL" if fp8 is None else f"FP8 x{len(shard_ptrs)} shards",
         )
 
     def _start_draft_prefetch(self, registrations: list[PleOffloadRegistration]):