diff --git a/csrc/ops.h b/csrc/ops.h index 220031881..9bca26529 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -979,7 +979,8 @@ int64_t ple_fast_lane_start( std::vector offsets, int64_t eos, std::vector result_ptrs, std::vector flag_ptrs, int64_t cache_bytes, int64_t idle_spin_us, bool clamp_tokens, std::optional draft_ring, - int64_t ring_reqs, int64_t ring_spec); + int64_t ring_reqs, int64_t ring_spec, std::vector shard_ptrs, + int64_t shard_rows, bool iq4_rows); torch::Tensor ple_fast_lane_stats(int64_t handle); std::vector gguf_dmv13_decode_pack(torch::Tensor raw, diff --git a/csrc/ple_fast_lane.cpp b/csrc/ple_fast_lane.cpp index 8b7c6c645..518bdf6b9 100644 --- a/csrc/ple_fast_lane.cpp +++ b/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 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( - config_.table + row * config_.row_bytes); + uintptr_t address = reinterpret_cast(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(sysconf(_SC_PAGE_SIZE)); for (int64_t i : misses_) { - uintptr_t address = - reinterpret_cast(config_.table + rows_[i] * row_bytes); + uintptr_t address = reinterpret_cast(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 offsets, int64_t eos, std::vector result_ptrs, std::vector flag_ptrs, int64_t cache_bytes, int64_t idle_spin_us, bool clamp_tokens, std::optional draft_ring, - int64_t ring_reqs, int64_t ring_spec) { + int64_t ring_reqs, int64_t ring_spec, std::vector 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(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(); config.max_reqs = max_reqs; config.max_tokens = max_tokens; - config.table = reinterpret_cast(table_ptr); + for (int64_t pointer : shard_ptrs) { + config.shards.push_back(reinterpret_cast(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(multipliers[i]); } diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 5751c1eba..fe913222e 100644 --- a/csrc/torch_bindings.cpp +++ b/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); diff --git a/tests/v1/worker/test_ple_fast_lane.py b/tests/v1/worker/test_ple_fast_lane.py index 33af425ff..d6b45e2ad 100644 --- a/tests/v1/worker/test_ple_fast_lane.py +++ b/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() diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py index b21fec7a4..b41e6e91d 100644 --- a/vllm/v1/ple_offload/worker.py +++ b/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]):