From cee23bf44a86e48d70f8fda3ca3b6d4d9890072e Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 23 Sep 2026 22:22:48 -0700 Subject: [PATCH 01/20] fix(runtime): correct overlap token limits and cache safety (#546) * fix(scheduler): count delivered tokens under overlap * fix(moe): initialize slots used by inactive hybrid routes * fix(kvcache): report usable window capacity floors * refactor(moe): inline the zeroed slot cache allocation * test: drop the new regression tests --------- Co-authored-by: andy-yang-1 --- python/freetoken/kvcache/cache_status.py | 11 ++++------- python/freetoken/moe/offload_cache.py | 4 ++-- python/freetoken/scheduler/scheduler.py | 3 ++- tests/kvcache/test_cache_unit_bytes.py | 2 +- 4 files changed, 9 insertions(+), 11 deletions(-) diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 10169d6515..6aa13e7f75 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/python/freetoken/kvcache/cache_status.py @@ -120,21 +120,18 @@ def _mamba() -> int: return int(_linear_pool_min_slots(config) - 1) def _swa() -> int: - # Window-pool floor in tokens (matches the pool's own page unit x count): - # DSV4 -- the working-set floor pages _dsv4_pool_sizes clamps n_win_pages to, x P. - # radix-SWA -- the concurrent working set (+1 slot-0 sentinel) _swa_paged_num_tokens - # floors to. - # 0 for models without a window pool. + # Report usable tokens, matching num_swa_pages in rebuild requests. + # DSV4's physical floor includes a dummy page; the SWA floor already excludes slot 0. from .dsv4_cost_model import _dsv4_window_floor_pages from .hybrid_swa_pool import _swa_pool_floor mc = config.model_config if mc.dsv4_args is not None: P = mc.dsv4_args.window_size - return int(_dsv4_window_floor_pages(config, P) * P) + return int((_dsv4_window_floor_pages(config, P) - 1) * P) if not (mc.has_swa_attention and config.cache_type == "swa_radix"): return 0 - return int(_swa_pool_floor(config) + 1) + return int(_swa_pool_floor(config)) for key, fn in (("kv_tokens", _kv), ("moe_experts", _moe), ("mamba_slots", _mamba), ("swa_tokens", _swa)): diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 7abb9c8003..32e0f3cc1e 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -354,7 +354,7 @@ def set_bank_sources( name, layer_id, source.shape, source.dtype, ) self.bank_sources[name] = list(per_layer) - self.bank_caches[name] = torch.empty( + self.bank_caches[name] = torch.zeros( (self.cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device, @@ -492,7 +492,7 @@ def rebuild(self, cache_size: int) -> None: # 3. Reallocate the slot cache from the retained host sources. for name in self.bank_schema: head = self.bank_sources[name][0] - self.bank_caches[name] = torch.empty( + self.bank_caches[name] = torch.zeros( (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device ) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index ac1bf322e3..fa0d96a714 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -349,7 +349,8 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: next_token = int(next_token.item()) # EOS / stop-string -> "stop", output budget exhausted -> "length"; # EOS and stop strings win over length. - hit_length = not req.can_decode + # Overlap can advance device_len ahead of the token delivered to the host. + hit_length = req.input_ids.numel() >= req.max_device_len hit_eos = ( not req.sampling_params.ignore_eos and next_token in self.eos_token_ids ) diff --git a/tests/kvcache/test_cache_unit_bytes.py b/tests/kvcache/test_cache_unit_bytes.py index c0eec34eff..496622f985 100644 --- a/tests/kvcache/test_cache_unit_bytes.py +++ b/tests/kvcache/test_cache_unit_bytes.py @@ -213,7 +213,7 @@ def test_floors_dsv4_reports_real_window_floor(): assert floors["kv_tokens"] > 0 # The window (swa) floor is reported too, in tokens (window pages x P) -- not the bogus 0 the # radix-only branch used to give for DSV4. - assert floors["swa_tokens"] == _dsv4_window_floor_pages(cfg, P) * P + assert floors["swa_tokens"] == (_dsv4_window_floor_pages(cfg, P) - 1) * P def test_floors_missing_config_all_zero(): From f5b9700cc10312d52169d9195af043c040f66d70 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Thu, 24 Sep 2026 20:16:03 -0700 Subject: [PATCH 02/20] refactor(ple): extract a shared disk row store (#548) --- .../row_store_ext.cpp} | 170 ++++++++++++------ python/freetoken/kernel/row_store.py | 54 ++++++ python/freetoken/models/qwen4_exp/ple_disk.py | 37 +--- setup.py | 6 +- tests/kernels/test_row_store.py | 87 +++++++++ tests/models/qwen4_exp/test_ple_disk.py | 8 +- 6 files changed, 270 insertions(+), 92 deletions(-) rename python/freetoken/kernel/csrc/{ple_store/ple_store_ext.cpp => row_store/row_store_ext.cpp} (83%) create mode 100644 python/freetoken/kernel/row_store.py create mode 100644 tests/kernels/test_row_store.py diff --git a/python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp b/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp similarity index 83% rename from python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp rename to python/freetoken/kernel/csrc/row_store/row_store_ext.cpp index 2d54d692d6..e33b18a8e0 100644 --- a/python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp +++ b/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp @@ -1,7 +1,12 @@ -// Disk-backed PLE row store: rows read straight from the checkpoint's fp8 shard tensors -// through an extent table (PleRowSource in ple_ssd.py). Engine-thread only, no locks. -// Duplicate rows in one fill dedup into ONE batched read round; no RAM cache and no -// per-sequence state. Hash reference: tests/models/qwen4_exp/test_ple_disk.py. +// Disk-backed row store: fixed-width rows read straight from checkpoint shard tensors +// through an extent table (row i of extent e at extent_base[e] + i * row_stride). +// Engine-thread only, no locks. Duplicate rows in one fill dedup into ONE batched read +// round; no RAM cache and no per-sequence state. +// +// RowStore -- hash-agnostic: stage_rows(row_ids) + flush(); callers supply row ids. +// PleStore -- RowStore plus the Qwen3.8-Flash-Next PLE n-gram hash: stage(tokens). +// Hash reference: tests/models/qwen4_exp/test_ple_disk.py. +// // Platform seams: TableFile (O_DIRECT+pread; Win: NO_BUFFERING), BatchReader (io_uring, // pread-pool fallback = the portable shape), cumemop_* (dlopen libcuda; Win: nvcuda). @@ -183,7 +188,9 @@ class ThreadPoolBatchReader final : public BatchReader { ThreadPoolBatchReader() { // these threads block on I/O, not compute, so the core count is only a default unsigned n = std::max(1u, std::min(kReaderThreads, std::thread::hardware_concurrency())); - if (const char *env = std::getenv("FREETOKEN_PLE_READER_THREADS")) { + const char *env = std::getenv("FREETOKEN_ROW_STORE_READER_THREADS"); + if (env == nullptr) env = std::getenv("FREETOKEN_PLE_READER_THREADS"); // the PLE-era name + if (env != nullptr) { // kBatchEntries is the submit depth, so threads past it never get a read const int v = std::atoi(env); if (v > 0) n = std::min((unsigned)v, kBatchEntries); @@ -432,29 +439,24 @@ int64_t pos_mod(int64_t v, int64_t m) { return r < 0 ? r + m : r; } -class PleStore { +class RowStore { struct Extent { const TableFile *file; int64_t base; }; public: - PleStore(std::vector paths, std::vector extent_file, + RowStore(std::vector paths, std::vector extent_file, std::vector extent_base, int64_t rows_per_extent, int64_t row_bytes, - int64_t row_stride, std::vector multipliers, std::vector head_vocab_sizes, - std::vector head_offsets, int64_t eos_token_id, bool use_io_uring) - : row_bytes_(row_bytes), - row_stride_(row_stride), - rows_per_extent_(rows_per_extent), - mult_(std::move(multipliers)), - sizes_(std::move(head_vocab_sizes)), - offsets_(std::move(head_offsets)), - eos_(eos_token_id) { - if (mult_.size() != 3 || sizes_.size() != offsets_.size() || sizes_.empty()) - throw std::runtime_error("PLE hash geometry: want 3 multipliers and equal-length head tables"); + int64_t row_stride, bool use_io_uring) + : row_bytes_(row_bytes), row_stride_(row_stride), rows_per_extent_(rows_per_extent) { + if (row_bytes_ <= 0 || row_stride_ < row_bytes_ || rows_per_extent_ <= 0) + throw std::runtime_error("row store geometry: want row_bytes > 0, row_stride >= row_bytes, rows_per_extent > 0"); if (row_bytes_ > kPage) - throw std::runtime_error("PLE row_bytes " + std::to_string(row_bytes_) + + throw std::runtime_error("row_bytes " + std::to_string(row_bytes_) + " exceeds a page; bounce slots assume one-page rows"); + if (extent_file.size() != extent_base.size() || extent_file.empty()) + throw std::runtime_error("row store geometry: extent_file and extent_base must be equal-length and non-empty"); for (const std::string &p : paths) files_.push_back(std::make_unique(p)); const int64_t extent_bytes = (rows_per_extent_ - 1) * row_stride_ + row_bytes_; @@ -472,36 +474,33 @@ class PleStore { } // reader first: a still-running read must not land in freed bounce memory - ~PleStore() { + virtual ~RowStore() { reader_.reset(); free(bounce_); } - PleStore(const PleStore &) = delete; - PleStore &operator=(const PleStore &) = delete; + RowStore(const RowStore &) = delete; + RowStore &operator=(const RowStore &) = delete; - // Row ids for the token at w[2] with context (w[0], w[1]); mirrors NGramEmbedding.row_ids incl. the eos barrier. - void hash_rows(const int64_t *w, int64_t *rows) { - const int64_t prev1 = w[1]; - const int64_t prev2 = prev1 == eos_ ? eos_ : w[0]; - const int64_t bigram = wrap_mul(w[2], mult_[0]) ^ wrap_mul(prev1, mult_[1]); - const int64_t trigram = bigram ^ wrap_mul(prev2, mult_[2]); - const size_t half = sizes_.size() / 2; - for (size_t h = 0; h < sizes_.size(); h++) - rows[h] = pos_mod(h < half ? bigram : trigram, sizes_[h]) + offsets_[h]; - } + int64_t row_bytes() const { return row_bytes_; } + int64_t total_rows() const { return (int64_t)extents_.size() * rows_per_extent_; } - // Hash and queue one run: tokens_addr holds n+2 ids, the leading two are context. No I/O until flush(). - void stage(uintptr_t tokens_addr, int64_t n, uintptr_t staging_addr) { - const int64_t *tokens = reinterpret_cast(tokens_addr); + // Queue n row ids (int64 at rows_addr); row i lands at staging_addr + i * dst_stride + // (dst_stride = 0 packs rows back to back). Out-of-range ids throw before any I/O is + // issued. No I/O until flush(). + void stage_rows(uintptr_t rows_addr, int64_t n, uintptr_t staging_addr, int64_t dst_stride) { + const int64_t *rows = reinterpret_cast(rows_addr); uint8_t *staging = reinterpret_cast(staging_addr); - const size_t heads = sizes_.size(); - std::vector rows(heads); + if (dst_stride == 0) dst_stride = row_bytes_; + if (dst_stride < row_bytes_) + throw std::runtime_error("stage_rows: dst_stride below row_bytes would overlap rows"); + const int64_t total = total_rows(); for (int64_t i = 0; i < n; i++) { - hash_rows(tokens + i, rows.data()); - for (size_t h = 0; h < heads; h++) - request_row(rows[h], staging + ((size_t)i * heads + h) * row_bytes_); + if (rows[i] < 0 || rows[i] >= total) + throw std::out_of_range("stage_rows: row id " + std::to_string(rows[i]) + " outside [0, " + + std::to_string(total) + ")"); } + for (int64_t i = 0; i < n; i++) request_row(rows[i], staging + (size_t)i * (size_t)dst_stride); } // One batched disk round for everything staged; signals even when nothing was. @@ -519,15 +518,7 @@ class PleStore { std::to_string(files_.size()) + " files"; } - private: - struct Pending { - const TableFile *file; - int64_t read_off; - int64_t read_len; - int64_t row_off; // row payload start inside the read buffer - std::vector dsts; - }; - + protected: // Queue dst on this fill's pending batch; duplicate rows fan out from one read. void request_row(int64_t row_id, uint8_t *dst) { auto pit = pending_index_.find(row_id); @@ -548,11 +539,20 @@ class PleStore { pending_.push_back(std::move(p)); } + private: + struct Pending { + const TableFile *file; + int64_t read_off; + int64_t read_len; + int64_t row_off; // row payload start inside the read buffer + std::vector dsts; + }; + // Read every pending row in reader-capacity batches and fan out the copies. void flush_pending() { if (pending_.empty()) return; struct Cleanup { - PleStore *s; + RowStore *s; ~Cleanup() { s->pending_.clear(); s->pending_index_.clear(); @@ -586,8 +586,6 @@ class PleStore { } int64_t row_bytes_, row_stride_, rows_per_extent_; - std::vector mult_, sizes_, offsets_; - int64_t eos_; std::vector> files_; std::vector extents_; std::unique_ptr reader_; @@ -597,10 +595,71 @@ class PleStore { std::unordered_map pending_index_; }; +// The Qwen3.8-Flash-Next PLE table: RowStore plus the checkpoint's n-gram hash, so a fill +// takes raw token runs and never crosses the Python boundary per row. +class PleStore final : public RowStore { + public: + PleStore(std::vector paths, std::vector extent_file, + std::vector extent_base, int64_t rows_per_extent, int64_t row_bytes, + int64_t row_stride, std::vector multipliers, std::vector head_vocab_sizes, + std::vector head_offsets, int64_t eos_token_id, bool use_io_uring) + : RowStore(std::move(paths), std::move(extent_file), std::move(extent_base), rows_per_extent, + row_bytes, row_stride, use_io_uring), + mult_(std::move(multipliers)), + sizes_(std::move(head_vocab_sizes)), + offsets_(std::move(head_offsets)), + eos_(eos_token_id) { + if (mult_.size() != 3 || sizes_.size() != offsets_.size() || sizes_.empty()) + throw std::runtime_error("PLE hash geometry: want 3 multipliers and equal-length head tables"); + } + + // Row ids for the token at w[2] with context (w[0], w[1]); mirrors NGramEmbedding.row_ids incl. the eos barrier. + void hash_rows(const int64_t *w, int64_t *rows) { + const int64_t prev1 = w[1]; + const int64_t prev2 = prev1 == eos_ ? eos_ : w[0]; + const int64_t bigram = wrap_mul(w[2], mult_[0]) ^ wrap_mul(prev1, mult_[1]); + const int64_t trigram = bigram ^ wrap_mul(prev2, mult_[2]); + const size_t half = sizes_.size() / 2; + for (size_t h = 0; h < sizes_.size(); h++) + rows[h] = pos_mod(h < half ? bigram : trigram, sizes_[h]) + offsets_[h]; + } + + // Hash and queue one run: tokens_addr holds n+2 ids, the leading two are context. No I/O until flush(). + void stage(uintptr_t tokens_addr, int64_t n, uintptr_t staging_addr) { + const int64_t *tokens = reinterpret_cast(tokens_addr); + uint8_t *staging = reinterpret_cast(staging_addr); + const size_t heads = sizes_.size(); + std::vector rows(heads); + for (int64_t i = 0; i < n; i++) { + hash_rows(tokens + i, rows.data()); + for (size_t h = 0; h < heads; h++) + request_row(rows[h], staging + ((size_t)i * heads + h) * (size_t)row_bytes()); + } + } + + private: + std::vector mult_, sizes_, offsets_; + int64_t eos_; +}; + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - py::class_(m, "PleStore") + py::class_(m, "RowStore") + .def(py::init, std::vector, std::vector, + int64_t, int64_t, int64_t, bool>(), + py::arg("paths"), py::arg("extent_file"), py::arg("extent_base"), + py::arg("rows_per_extent"), py::arg("row_bytes"), py::arg("row_stride"), + py::arg("use_io_uring") = true) + .def("stage_rows", &RowStore::stage_rows, py::arg("rows_addr"), py::arg("n"), + py::arg("staging_addr"), py::arg("dst_stride") = 0, + py::call_guard()) + .def("flush", &RowStore::flush, py::arg("signal_addr") = 0, + py::call_guard()) + .def("io_backend", &RowStore::io_backend) + .def_property_readonly("row_bytes", &RowStore::row_bytes) + .def_property_readonly("total_rows", &RowStore::total_rows); + py::class_(m, "PleStore") .def(py::init, std::vector, std::vector, int64_t, int64_t, int64_t, std::vector, std::vector, std::vector, int64_t, bool>(), @@ -610,10 +669,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("head_vocab_sizes"), py::arg("head_offsets"), py::arg("eos_token_id"), py::arg("use_io_uring") = true) .def("stage", &PleStore::stage, py::arg("tokens_addr"), py::arg("n"), - py::arg("staging_addr"), py::call_guard()) - .def("flush", &PleStore::flush, py::arg("signal_addr") = 0, - py::call_guard()) - .def("io_backend", &PleStore::io_backend); + py::arg("staging_addr"), py::call_guard()); m.def("memop_write", &memop_write, py::arg("stream"), py::arg("addr"), py::arg("value")); m.def("memop_wait_geq", &memop_wait_geq, py::arg("stream"), py::arg("addr"), py::arg("value")); m.def("memop_wait_reset", &memop_wait_reset, py::arg("stream"), py::arg("flag_addr")); diff --git a/python/freetoken/kernel/row_store.py b/python/freetoken/kernel/row_store.py new file mode 100644 index 0000000000..5b42479ebc --- /dev/null +++ b/python/freetoken/kernel/row_store.py @@ -0,0 +1,54 @@ +"""Python face of the disk row store extension (``kernel/csrc/row_store``) and its CUDA-graph sync. + +``RowStore`` reads fixed-width rows straight out of checkpoint shard files (io_uring or a pread pool, +O_DIRECT) into pinned staging by row id; ``PleStore`` adds PLE's n-gram hashing. + +Graph sync: a captured graph cannot wait for host I/O, so a captured ``lookup`` WAITs on a pinned +flag through a stream memop (``wait_reset``) that the host ``signal``s once the fill landed; where +the driver rejects memops in capture, the caller falls back to launch-gating (fill before replay). +``probe_wait_sync`` decides which, once, at startup. +""" + +from __future__ import annotations + +import torch + +from freetoken.kernel import _row_store +from freetoken.kernel.pinned import alloc_pinned_tensor + +RowStore = _row_store.RowStore +PleStore = _row_store.PleStore + + +def probe_wait_sync(mode: str, device: torch.device) -> bool: + """``mode``: ``auto`` (probe), ``wait`` (require memops) or ``gate`` (never use them).""" + if mode not in ("auto", "wait", "gate"): + raise ValueError(f"unknown row-store sync mode {mode!r}; expected auto, wait or gate") + if mode == "gate": + return False + scratch = alloc_pinned_tensor(1, dtype=torch.int64) + scratch.zero_() + stream = torch.cuda.current_stream(device) + ok = ( + _row_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0 + and _row_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0 + ) + if ok: + stream.synchronize() + ok = int(scratch[0]) == 7 + if mode == "wait" and not ok: + raise RuntimeError("wait-sync requested but stream memops are unavailable") + return ok + + +def wait_reset(stream: torch.cuda.Stream, flag: torch.Tensor) -> None: + """Enqueue WAIT(flag >= 1) then flag := 0 on ``stream`` (capturable).""" + _row_store.memop_wait_reset(stream.cuda_stream, flag.data_ptr()) + + +def signal(flag: torch.Tensor) -> None: + """Host side: release a pending WAIT (a release-store of 1).""" + _row_store.signal_flag(flag.data_ptr()) + + +__all__ = ["RowStore", "PleStore", "probe_wait_sync", "wait_reset", "signal"] diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index cea6269357..ec7da80d56 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -112,7 +112,7 @@ def __init__( max_extend_tokens: int = 8192, dtype: torch.dtype = torch.bfloat16, ) -> None: - from freetoken.kernel import _ple_store + from freetoken.kernel.row_store import PleStore self.num_rows = source.total_rows self.head_dim = source.row_bytes # fp8: one byte per element @@ -128,7 +128,7 @@ def __init__( raise ValueError( f"PLE row source holds {source.total_rows} rows but the hash addresses {need}; incomplete checkpoint?" ) - self._store = _ple_store.PleStore( + self._store = PleStore( paths=list(source.paths), extent_file=list(source.extent_file), extent_base=list(source.extent_base), @@ -155,7 +155,9 @@ def __init__( self._eager_pinned.zero_() # the warmup prefill stages nothing and reads whatever sits here self._eager_dev = torch.empty(eager_bytes, dtype=torch.uint8, device=self._device) # probe picks flag-sync (graph WAITs at the consume, host fills then signals) or launch-gating - self._wait_sync = self._probe_wait_sync(os.getenv(_SYNC_ENV, "auto")) + from freetoken.kernel.row_store import probe_wait_sync + + self._wait_sync = probe_wait_sync(os.getenv(_SYNC_ENV, "auto"), self._device) # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap self._flag = alloc_pinned_tensor(1, dtype=torch.int64) self._flag.zero_() @@ -164,25 +166,6 @@ def __init__( sync = "wait-sync" if self._wait_sync else "launch-gating" logger.info_rank0(f"PLE disk backend: {self._store.io_backend()}, {sync}") - def _probe_wait_sync(self, mode: str) -> bool: - from freetoken.kernel import _ple_store - - if mode == "gate": - return False - scratch = alloc_pinned_tensor(1, dtype=torch.int64) - scratch.zero_() - stream = torch.cuda.current_stream(self._device) - ok = ( - _ple_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0 - and _ple_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0 - ) - if ok: - stream.synchronize() - ok = int(scratch[0]) == 7 - if mode == "wait" and not ok: - raise RuntimeError("FREETOKEN_PLE_SYNC=wait but stream memops are unavailable") - return ok - # ---------------- host side (engine thread, before the forward launches) ---------------- def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: @@ -222,10 +205,10 @@ def _complete() -> None: for r, t in zip(reqs, tokens)] self.fill(runs, graph=True) except BaseException: - from freetoken.kernel import _ple_store + from freetoken.kernel.row_store import signal # unblock the stream before surfacing; the step's output is discarded - _ple_store.signal_flag(self._flag.data_ptr()) + signal(self._flag) raise return _complete @@ -260,11 +243,9 @@ def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torc rows = row_ids.shape[0] capturing = torch.cuda.is_current_stream_capturing() if capturing and self._wait_sync: - from freetoken.kernel import _ple_store + from freetoken.kernel.row_store import wait_reset - _ple_store.memop_wait_reset( - torch.cuda.current_stream(self._device).cuda_stream, self._flag.data_ptr() - ) + wait_reset(torch.cuda.current_stream(self._device), self._flag) pinned, dev = ( (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned, self._eager_dev) ) diff --git a/setup.py b/setup.py index d8e705ab69..7132eea088 100644 --- a/setup.py +++ b/setup.py @@ -64,12 +64,12 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: libraries=["cudart"], extra_compile_args=["-O3", "-std=c++17", "-pthread"], ), - # --ple-backend disk row store; Linux-only until the TableFile/BatchReader seams grow Windows bodies + # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies *([ CppExtension( - name="freetoken.kernel._ple_store", + name="freetoken.kernel._row_store", sources=[ - "python/freetoken/kernel/csrc/ple_store/ple_store_ext.cpp", + "python/freetoken/kernel/csrc/row_store/row_store_ext.cpp", ], extra_compile_args=["-O3", "-std=c++17"], ) diff --git a/tests/kernels/test_row_store.py b/tests/kernels/test_row_store.py new file mode 100644 index 0000000000..d4a195026a --- /dev/null +++ b/tests/kernels/test_row_store.py @@ -0,0 +1,87 @@ +"""The hash-agnostic ``RowStore`` (``kernel/csrc/row_store``): byte-exact staged rows across extents and files, +dedup of repeated ids, a destination stride, and bounds checks before any I/O. The PLE hash wrapper on top of +it is covered by ``tests/models/qwen4_exp/test_ple_disk.py``.""" + +from __future__ import annotations + +import pytest +import torch + +_row_store = pytest.importorskip("freetoken.kernel._row_store") + +ROW_BYTES = 264 # deliberately unaligned rows exercise reads across page boundaries + + +def _table(rows: int, seed: int) -> torch.Tensor: + gen = torch.Generator().manual_seed(seed) + return torch.randint(0, 256, (rows, ROW_BYTES), dtype=torch.uint8, generator=gen) + + +def _stage(store, ids: torch.Tensor, dst_stride: int = 0) -> torch.Tensor: + stride = dst_stride or store.row_bytes + staging = torch.zeros(ids.numel() * stride, dtype=torch.uint8) + store.stage_rows(ids.data_ptr(), ids.numel(), staging.data_ptr(), dst_stride) + store.flush(0) + return staging.view(ids.numel(), stride) + + +@pytest.mark.parametrize("use_io_uring", [True, False]) +def test_rows_come_back_bitwise_across_extents_and_files(tmp_path, use_io_uring): + # two files; file 0 carries extents 0 and 2 behind an unrelated header, file 1 carries extent 1 + per_extent = 1000 + table = _table(3 * per_extent, seed=1) + header = b"safetensors-like header, not a row" + bytes(4096 - 34) + f0 = tmp_path / "shard0.bin" + f1 = tmp_path / "shard1.bin" + f0.write_bytes(header + table[:per_extent].numpy().tobytes() + table[2 * per_extent :].numpy().tobytes()) + f1.write_bytes(table[per_extent : 2 * per_extent].numpy().tobytes()) + store = _row_store.RowStore( + paths=[str(f0), str(f1)], + extent_file=[0, 1, 0], + extent_base=[len(header), 0, len(header) + per_extent * ROW_BYTES], + rows_per_extent=per_extent, + row_bytes=ROW_BYTES, + row_stride=ROW_BYTES, + use_io_uring=use_io_uring, + ) + assert store.total_rows == 3 * per_extent and store.row_bytes == ROW_BYTES + + ids = torch.tensor([0, 999, 1000, 1999, 2000, 2999, 1234, 1234, 42, 0], dtype=torch.int64) + got = _stage(store, ids) + assert torch.equal(got, table[ids]) + + # a wider destination stride leaves the gap untouched + got = _stage(store, ids, dst_stride=ROW_BYTES + 40) + assert torch.equal(got[:, :ROW_BYTES], table[ids]) and int(got[:, ROW_BYTES:].sum()) == 0 + + # a fill larger than the reader's in-flight window pipelines through + many = torch.randint(0, 3 * per_extent, (5000,), dtype=torch.int64, generator=torch.Generator().manual_seed(2)) + assert torch.equal(_stage(store, many), table[many]) + + # flushing nothing still signals the flag + flag = torch.zeros(1, dtype=torch.int64) + store.flush(flag.data_ptr()) + assert int(flag[0]) == 1 + + +def test_geometry_and_bounds_are_checked_up_front(tmp_path): + table = _table(64, seed=3) + path = tmp_path / "t.bin" + path.write_bytes(table.numpy().tobytes()) + make = lambda **kw: _row_store.RowStore( + paths=[str(path)], extent_file=[0], extent_base=[0], rows_per_extent=64, + row_bytes=ROW_BYTES, row_stride=ROW_BYTES, use_io_uring=False, **kw, + ) + store = make() + ids = torch.tensor([3, 64], dtype=torch.int64) + staging = torch.zeros(2 * ROW_BYTES, dtype=torch.uint8) + with pytest.raises(IndexError): + store.stage_rows(ids.data_ptr(), 2, staging.data_ptr(), 0) + with pytest.raises(RuntimeError, match="overlap"): + store.stage_rows(ids[:1].data_ptr(), 1, staging.data_ptr(), ROW_BYTES - 1) + with pytest.raises(RuntimeError, match="extent needs"): + _row_store.RowStore(paths=[str(path)], extent_file=[0], extent_base=[0], rows_per_extent=65, + row_bytes=ROW_BYTES, row_stride=ROW_BYTES, use_io_uring=False) + with pytest.raises(RuntimeError, match="exceeds a page"): + _row_store.RowStore(paths=[str(path)], extent_file=[0], extent_base=[0], rows_per_extent=1, + row_bytes=4097, row_stride=4097, use_io_uring=False) diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index fa11807390..1112210b84 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -13,7 +13,7 @@ from .common import EOS, hash_constants, requires_cuda, toy_hf_config from .test_ple import _meta -_ple_store = pytest.importorskip("freetoken.kernel._ple_store") +_row_store = pytest.importorskip("freetoken.kernel._row_store") _KEY_PREFIX = "model.layers.1.ple.ple_embedding.ngram_embedding" @@ -43,7 +43,7 @@ def _make_store(tmp_path, *, write=True, use_io_uring=True): path = tmp_path / "ple-table.bin" if write: path.write_bytes(table.numpy().tobytes()) - store = _ple_store.PleStore( + store = _row_store.PleStore( paths=[str(path)], extent_file=[0], extent_base=[0], @@ -175,10 +175,10 @@ def test_layouts_readers_and_errors(tmp_path): multipliers=[3, 5, 7], head_vocab_sizes=sizes, head_offsets=offsets, eos_token_id=eos, ) - ref = _ple_store.PleStore( + ref = _row_store.PleStore( paths=[str(flat)], extent_file=[0, 0, 0, 0], extent_base=[0, nb, 2 * nb, 3 * nb], **kwargs ) - multi = _ple_store.PleStore( + multi = _row_store.PleStore( paths=[str(fa), str(fb)], extent_file=[0, 1, 0, 1], extent_base=[1231, 0, 1231 + nb + 77, nb + 4095], **kwargs, ) From 0d652e73a452d014ac5441a15baa75348e9fcb0a Mon Sep 17 00:00:00 2001 From: ZihaoMu Date: Sat, 26 Sep 2026 08:49:45 +0800 Subject: [PATCH 03/20] feat(rocm): add RDNA3 and RDNA4 runtime foundation (#132) * feat: add simple ROCm GPU support for RDNA3 (gfx1100-1103) - Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents - Update pinned_tensor.cpp to compile under both nvcc and hipcc - Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family) - Guard NVIDIA arch checks to return None on ROCm - Skip nvcc version check in _toolchain.py when on ROCm - Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch) - Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm - Add is_rocm() and driver_hip_version() in backend.py - Add rocm-smi fallback in __main__.py for clangd generation - Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds, Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache - Add AMD ROCm classifier in pyproject.toml * fix(rocm): fall back for single-bank fast index copy * test(rocm): cover fast index copy fallback * test(rocm): use explicit fast copy module import * fix(rocm): complete RDNA3 runtime path * fix(rocm): preserve CPU MoE implementation * build(rocm): support RDNA4 discovery and modular SDKs * fix(rocm): complete modular SDK runtime support * test(rocm): harden foundation platform contracts * fix(rocm): pin jit target and split install guide * perf(kernel): restore ROCm non-temporal index copy --------- Co-authored-by: bouclem Co-authored-by: Yuu <206304251+nekomario28@users.noreply.github.com> Co-authored-by: LZ-QWQ --- docs/install.md | 2 + docs/install_amd.md | 39 +++++ pyproject.toml | 6 +- python/freetoken/kernel/__main__.py | 35 ++-- python/freetoken/kernel/_toolchain.py | 9 +- python/freetoken/kernel/backend.py | 18 ++ .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 2 +- .../csrc/include/freetoken/hip_compat.h | 162 ++++++++++++++++++ .../kernel/csrc/include/freetoken/utils.cuh | 6 + .../kernel/csrc/jit/fast_index_copy.cuh | 55 ++++-- .../freetoken/kernel/csrc/pinned_tensor.cpp | 2 +- python/freetoken/kernel/pynccl.py | 1 + python/freetoken/kernel/triton/activation.py | 10 +- python/freetoken/kernel/triton/e4m3_compat.py | 4 + python/freetoken/kernel/triton/norm.py | 8 +- python/freetoken/kernel/utils.py | 145 ++++++++++++++-- python/freetoken/utils/__init__.py | 4 + python/freetoken/utils/arch.py | 52 ++++++ setup.py | 77 +++++++-- tests/kernels/test_e4m3_compat.py | 23 +++ tests/kernels/test_pinned_tensor.py | 7 +- tests/kernels/test_rocm_launch_kwargs.py | 51 ++++++ tests/utils/test_rocm_arch.py | 127 ++++++++++++++ 23 files changed, 780 insertions(+), 65 deletions(-) create mode 100644 docs/install_amd.md create mode 100644 python/freetoken/kernel/csrc/include/freetoken/hip_compat.h create mode 100644 tests/kernels/test_rocm_launch_kwargs.py create mode 100644 tests/utils/test_rocm_arch.py diff --git a/docs/install.md b/docs/install.md index 9438c21a67..f474216bf0 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1,5 +1,7 @@ # Install +AMD users: see the [AMD ROCm installation guide (WIP)](install_amd.md). + ## Requirements - Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13) diff --git a/docs/install_amd.md b/docs/install_amd.md new file mode 100644 index 0000000000..d8b0c1bcb1 --- /dev/null +++ b/docs/install_amd.md @@ -0,0 +1,39 @@ +# AMD ROCm installation (WIP) + +AMD support is experimental and remains a work in progress. See the +[AMD Support roadmap](https://github.com/FlashML-org/FreeToken/issues/541) for +the current integration and qualification status. + +## Requirements + +- Linux x86_64 +- AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`) +- ROCm 7.14 +- Python >= 3.10 + +## Install from source + +Use an official ROCm PyTorch image whose PyTorch version satisfies the project's +`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is: + +```bash +VIDEO_GID="$(getent group video | cut -d: -f3)" +RENDER_GID="$(getent group render | cut -d: -f3)" +docker run --rm -it \ + --device=/dev/kfd --device=/dev/dri \ + --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \ + -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \ + rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash +``` + +Inside the container, preserve the ROCm-enabled PyTorch already supplied by the +image and disable build isolation so it is also used to compile the extensions: + +```bash +python -m pip install --no-build-isolation -e . +``` + +Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the +actual target reported by `rocminfo`. diff --git a/pyproject.toml b/pyproject.toml index 25cb204bb4..8d00d2bbc6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Environment :: GPU :: NVIDIA CUDA", + "Environment :: GPU :: AMD ROCm", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", @@ -59,7 +60,10 @@ dependencies = [ "torchvision>=0.26,<0.27", "tqdm>=4.66,<5", "transformers>=5.16,<5.17", - "triton==3.6.0; platform_system == 'Linux'", + # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its + # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the + # runtime-specific wheel selected by the PyTorch distribution. + "triton>=3.6,<3.8; platform_system == 'Linux'", "uvicorn>=0.30,<1", ] diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py index 7be541a67c..5b66d484df 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -6,27 +6,32 @@ def generate_clangd(): import subprocess from freetoken.kernel.utils import DEFAULT_INCLUDE - from freetoken.utils import init_logger + from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path logger = init_logger(__name__) logger.info("Generating .clangd file...") include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE - status = subprocess.run( - args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], - capture_output=True, - check=True, - ) - compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] - major, minor = compute_cap.split(".") + + # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries. + if is_rocm(): + arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"] + else: + try: + status = subprocess.run( + args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], + capture_output=True, + check=True, + ) + compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] + major, minor = compute_cap.split(".") + except (subprocess.CalledProcessError, FileNotFoundError, ValueError): + import torch + + major, minor = torch.cuda.get_device_capability() + arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"] compile_flags = ",\n ".join( - [ - "-xcuda", - f"--cuda-gpu-arch=sm_{major}{minor}", - "-std=c++20", - "-Wall", - "-Wextra", - ] + arch_flags + ["-std=c++20", "-Wall", "-Wextra"] + [f"-isystem{path}" for path in include_paths] ) clangd_content = f""" diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb43..93d0322151 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -1,4 +1,4 @@ -"""CUDA toolchain/torch consistency checks. +"""CUDA/HIP toolchain/torch consistency checks. Standalone on purpose: setup.py and the kernel-cache build backend load this file by path, so it must not import the freetoken package. @@ -16,6 +16,11 @@ _TRUE_VALUES = {"1", "true", "yes", "on"} +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + def _nvcc_path() -> str | None: from torch.utils.cpp_extension import CUDA_HOME @@ -49,6 +54,8 @@ def check_nvcc_matches_torch() -> None: nvcc-built binaries link libcudart.so.; at runtime only the torch wheel's own CUDA runtime is guaranteed to be loadable. """ + if _is_rocm(): + return # ROCm uses hipcc, not nvcc if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: return torch_major = torch_cuda_major() diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 2a886cf77d..6d1b4a9f8e 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -47,6 +47,24 @@ def device_capability() -> tuple[int, int]: return (int(major), int(minor)) +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU).""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def driver_hip_version() -> int | None: + """ROCm driver version, or None if undetermined.""" + # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used. + try: + from freetoken.kernel.pinned import _load_pinned_extension + return int(_load_pinned_extension().driver_cuda_version()) or None + except Exception: + return None + + @functools.cache def driver_cuda_version() -> int | None: """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0), diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 8349123bd6..e2cb0d9c33 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -29,7 +29,7 @@ #include #include -#include +#include #include #if defined(__linux__) diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h new file mode 100644 index 0000000000..1898b41189 --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -0,0 +1,162 @@ +#pragma once + +// HIP compatibility shim: maps CUDA runtime API names to HIP equivalents so +// the same C++ source compiles under both nvcc and hipcc. Include this instead +// of directly when the file needs the runtime API. +// +// On NVIDIA platforms the CUDA headers are included as-is and every macro below +// resolves to the original CUDA symbol, so there is zero overhead. +// +// Supported ROCm targets: +// gfx1100 — RX 7900 XTX / XT +// gfx1101 — RX 7900 GRE +// gfx1102 — RX 7700 / XT +// gfx1103 — RX 7600 / XT +// gfx1200 — RX 9060 family +// gfx1201 — RX 9070 family / Radeon AI PRO R9700 + +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) + +#define FREETOKEN_USE_ROCM 1 + +// --- HIP runtime headers --- +#include +#include + +// --- API name mapping (CUDA -> HIP) --- +// HIP already defines most cuda* names as macros that expand to hip* equivalents +// via hip_runtime.h, but a few are missing or differ in signature. Define them +// here so call-sites stay unchanged. + +#ifndef cudaSuccess +#define cudaSuccess hipSuccess +#endif + +#ifndef cudaError_t +#define cudaError_t hipError_t +#endif + +#ifndef cudaGetErrorString +#define cudaGetErrorString hipGetErrorString +#endif + +#ifndef cudaGetLastError +#define cudaGetLastError hipGetLastError +#endif + +#ifndef cudaMallocHost +#define cudaMallocHost hipHostMalloc +#endif + +#ifndef cudaFreeHost +#define cudaFreeHost hipHostFree +#endif + +#ifndef cudaHostAlloc +#define cudaHostAlloc hipHostMalloc +#endif + +#ifndef cudaHostRegister +#define cudaHostRegister hipHostRegister +#endif + +#ifndef cudaHostRegisterPortable +#define cudaHostRegisterPortable hipHostRegisterPortable +#endif + +#ifndef cudaHostRegisterMapped +#define cudaHostRegisterMapped hipHostRegisterMapped +#endif + +#ifndef cudaHostAllocPortable +#define cudaHostAllocPortable hipHostMallocPortable +#endif + +#ifndef cudaHostAllocMapped +#define cudaHostAllocMapped hipHostMallocMapped +#endif + +#ifndef cudaHostGetDevicePointer +#define cudaHostGetDevicePointer hipHostGetDevicePointer +#endif + +#ifndef cudaGetDevice +#define cudaGetDevice hipGetDevice +#endif + +#ifndef cudaDriverGetVersion +#define cudaDriverGetVersion hipDriverGetVersion +#endif + +#ifndef cudaDeviceGetAttribute +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#endif + +#ifndef cudaDevAttrUnifiedAddressing +#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaDevAttrCanUseHostPointerForRegisteredMem +// HIP does not expose this attribute; assume UVA identity on ROCm (true on Linux). +// TODO(ROCm): re-enable proper UVA query if HIP adds this attribute. +#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaFuncSetAttribute +#define cudaFuncSetAttribute hipFuncSetAttribute +#endif + +#ifndef cudaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#endif + +#ifndef cudaLaunchKernelEx +// ROCm 7 exposes the CUDA-compatible extended launch configuration through HIP. +#define cudaLaunchKernelEx hipLaunchKernelEx +#endif + +#ifndef cudaLaunchConfig_t +#define cudaLaunchConfig_t hipLaunchConfig_t +#endif + +#ifndef cudaLaunchAttribute +#define cudaLaunchAttribute hipLaunchAttribute +#endif + +#ifndef cudaLaunchAttributeProgrammaticStreamSerialization +// PDL (Programmatic Dependent Launch) is NVIDIA-specific. +// TODO(ROCm): PDL has no ROCm equivalent — disabled, may affect overlap scheduling latency. +#define cudaLaunchAttributeProgrammaticStreamSerialization 0 +#endif + +#ifndef cudaStream_t +#define cudaStream_t hipStream_t +#endif + +#ifndef cudaStreamSynchronize +#define cudaStreamSynchronize hipStreamSynchronize +#endif + +#ifndef cudaLaunchHostFunc +#define cudaLaunchHostFunc hipLaunchHostFunc +#endif + +#ifndef CUDART_CB +#define CUDART_CB +#endif + +#ifndef __grid_constant__ +#define __grid_constant__ +#endif + +#ifndef dim3 +// HIP already provides dim3; this is a no-op guard. +#endif + +#else // NVIDIA CUDA path + +#define FREETOKEN_USE_ROCM 0 + +#include + +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832c1..72a4959698 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -115,6 +116,10 @@ public: } auto with_attr(bool use_pdl) -> LaunchKernel & { +#if FREETOKEN_USE_ROCM + (void)use_pdl; + m_config.numAttrs = 0; +#else if (use_pdl) { m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization; m_attr_cache.val.programmaticStreamSerializationAllowed = 1; @@ -123,6 +128,7 @@ public: } else { m_config.numAttrs = 0; } +#endif return *this; } diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23ed2..43f3fc1ba5 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -20,6 +20,13 @@ struct device_vec { namespace details { +#if FREETOKEN_USE_ROCM +// Use native Clang vector types so each __builtin_nontemporal_load/store +// maps to one vectorized 8B/16B global memory instruction. +using native_uint2 = uint32_t __attribute__((ext_vector_type(2))); +using native_uint4 = uint32_t __attribute__((ext_vector_type(4))); +#endif + template inline constexpr auto get_mem_package() { if constexpr (kUnit == 16) { @@ -34,40 +41,68 @@ inline constexpr auto get_mem_package() { } __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { +#if FREETOKEN_USE_ROCM + return uint1{__builtin_nontemporal_load(&src->x)}; +#else uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; +#endif } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { +#if FREETOKEN_USE_ROCM + const auto value = __builtin_nontemporal_load(reinterpret_cast(src)); + return __builtin_bit_cast(uint2, value); +#else uint32_t tmp0, tmp1; asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; +#endif } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { +#if FREETOKEN_USE_ROCM + const auto value = __builtin_nontemporal_load(reinterpret_cast(src)); + return __builtin_bit_cast(uint4, value); +#else uint32_t tmp0, tmp1, tmp2, tmp3; asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; +#endif } __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { +#if FREETOKEN_USE_ROCM + __builtin_nontemporal_store(value.x, &dst->x); +#else uint32_t tmp = value.x; asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +#endif } __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { +#if FREETOKEN_USE_ROCM + __builtin_nontemporal_store( + __builtin_bit_cast(native_uint2, value), reinterpret_cast(dst)); +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +#endif } __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { +#if FREETOKEN_USE_ROCM + __builtin_nontemporal_store( + __builtin_bit_cast(native_uint4, value), reinterpret_cast(dst)); +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w; asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +#endif } __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) { @@ -75,7 +110,7 @@ __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag auto* flag = reinterpret_cast(const_cast(flag_ptr)); uint32_t sleep_ns = 128; while (atomicAdd(flag, 0) > 0) { -#if __CUDA_ARCH__ >= 700 +#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 700 __nanosleep(sleep_ns); #endif sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048; @@ -147,7 +182,7 @@ inline bool host_ptr_identity() { } inline void* device_alias(void* ptr, DLDevice dev) { - if (dev.device_type == kDLCUDA || host_ptr_identity()) { + if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) { return ptr; } void* mapped = nullptr; @@ -269,7 +304,7 @@ inline auto get_sync_flag_ptr( auto flag_dtype = host::SymbolicDType{}; host::TensorMatcher({1}) .with_dtype(flag_dtype) - .with_device(device) + .with_device(device) .verify(sync_flag); return static_cast(sync_flag.data_ptr()); } @@ -344,17 +379,17 @@ struct FastIndexCopyKernel { TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(src); TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(dst); TensorMatcher({L}) .with_dtype(indices_dtype) - .with_device(device) + .with_device(device) .verify(src_indices) .verify(dst_indices); @@ -363,7 +398,7 @@ struct FastIndexCopyKernel { const auto num_indices_tensor = num_indices.value(); TensorMatcher({1}) .with_dtype(num_indices_dtype) - .with_device(device) + .with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -529,14 +564,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); - TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) .verify(dst_indices).verify(src_indices); const int64_t* valid_length = nullptr; if (num_indices.has_value()) { - TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adfa3..9355f57aa5 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,5 +1,5 @@ #include -#include +#include #include namespace { diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea57351d..71bde7342a 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -27,6 +27,7 @@ def get_buffer(self) -> int: ... @functools.cache def _load_nccl_module() -> Module: + # TODO(ROCm): NCCL -> RCCL migration for multi-GPU tensor parallelism on AMD. return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"]) diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 0e4d28c8fc..9645960d11 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -20,8 +20,9 @@ import triton.language as tl from triton.language.extra import libdevice from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents +from triton.language import target_info -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported SILU = 0 GELU = 1 @@ -50,6 +51,8 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): + if target_info.is_hip(): + return libdevice.tanh(x) # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. return tl.inline_asm_elementwise( "tanh.approx.f32 $0, $1;", "=f,f", [x], @@ -59,6 +62,8 @@ def _fast_tanh(x): @triton.jit def _fast_ex2(x): + if target_info.is_hip(): + return libdevice.exp2(x) # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu. return tl.inline_asm_elementwise( "ex2.approx.f32 $0, $1;", "=f,f", [x], @@ -141,8 +146,9 @@ def _act_and_mul( block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512) num_stages = 2 if block_d == 1024 else 3 _act_and_mul_kernel[grid]( - o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl, + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + **({} if is_rocm() else {"launch_pdl": pdl}), ) return out diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 1d9f744ce1..6ea0c7537d 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -59,6 +59,10 @@ def e4m3_native() -> bool: if _native is None: if FORCE_EMU: _native = False + elif torch.version.hip is not None: + # ROCm reports gfx1101 as capability (11, 0), which is not a CUDA + # compute capability and must not select the native fp8e4nv path. + _native = False else: from freetoken.gpu_select import assigned_visible_gpu diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f0d..9071e1dfff 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -30,7 +30,7 @@ import triton.language as tl from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])} @@ -144,7 +144,8 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): pdl = contig and is_sm90_supported() _rmsnorm_kernel[(A, B)]( out, input, weight, eps, H, sxa, sxb, soa, sob, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) return out @@ -172,7 +173,8 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): pdl = contig and is_sm90_supported() _fused_add_rmsnorm_kernel[(A, B)]( input, residual, weight, eps, H, sxa, sxb, sra, srb, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index cdd765c7d6..9af61526ae 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -5,6 +5,7 @@ import os import pathlib import re +from functools import cache from typing import TYPE_CHECKING, Iterator, List, NamedTuple, Tuple, TypeAlias, Union if TYPE_CHECKING: @@ -20,15 +21,23 @@ DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] -ARCH_LIST_ENV = "TVM_FFI_CUDA_ARCH_LIST" +CUDA_ARCH_LIST_ENV = "TVM_FFI_CUDA_ARCH_LIST" +ROCM_ARCH_LIST_ENV = "TVM_FFI_ROCM_ARCH_LIST" + + +def _is_rocm() -> bool: + import torch + + return getattr(torch.version, "hip", None) is not None def _cuda_arch_list() -> List[str]: """Archs a CUDA build targets: the AOT build's TVM_FFI_CUDA_ARCH_LIST, else the GPU this process is bound to.""" - arch_list = os.getenv(ARCH_LIST_ENV, "").split() + arch_list = os.getenv(CUDA_ARCH_LIST_ENV, "").split() if arch_list: return arch_list import torch @@ -39,17 +48,34 @@ def _cuda_arch_list() -> List[str]: return [f"{major}.{minor}"] +def _rocm_arch_list() -> List[str]: + """The single ROCm GPU this process is bound to, or an explicit cross-compile target.""" + from freetoken.utils.arch import get_rocm_gfx_arch + + arch = get_rocm_gfx_arch() + if arch is None: + raise RuntimeError( + "Could not determine the bound ROCm GPU architecture; select a visible GPU " + "or set FREETOKEN_ROCM_ARCH for cross-compilation" + ) + return [arch] + + @contextlib.contextmanager -def _pin_tvm_ffi_arch_ctx(arch_list: List[str]) -> Iterator[None]: - """Hand tvm-ffi the arch list through its env var for one build. Left unset, tvm-ffi asks nvidia-smi and takes the first GPU listed, which under --gpu or CUDA_VISIBLE_DEVICES on a mixed box is not the bound one.""" - if os.getenv(ARCH_LIST_ENV) or not arch_list: +def _pin_tvm_ffi_arch_ctx(arch_list: List[str], env_var: str) -> Iterator[None]: + """Pin tvm-ffi to the resolved targets for one build, then restore the caller's environment.""" + if not arch_list: yield return - os.environ[ARCH_LIST_ENV] = " ".join(arch_list) + old_arch_list = os.environ.get(env_var) + os.environ[env_var] = " ".join(arch_list) try: yield finally: - os.environ.pop(ARCH_LIST_ENV, None) + if old_arch_list is None: + os.environ.pop(env_var, None) + else: + os.environ[env_var] = old_arch_list def _cuda_cflags(extra: List[str], arch_list: List[str]) -> List[str]: @@ -63,6 +89,61 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _hip_cflags(extra: List[str], arch_list: List[str]) -> List[str]: + """HIP flags for a kernel build on ROCm.""" + # TODO(ROCm): Triton autotune configs need RDNA-specific tuning (wave count, LDS size). + return DEFAULT_HIP_CFLAGS + extra + [f"--offload-arch={arch}" for arch in arch_list] + + +@cache +def _rocm_link_flags() -> List[str]: + """Make ROCm's runtime library discoverable to JIT link commands. + + Traditional ROCm installs provide ``libamdhip64.so`` under ``$ROCM_HOME/lib``. + ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI + still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via + an explicit linker search path without modifying the Python environment. + """ + candidates: list[pathlib.Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(pathlib.Path(os.environ["ROCM_HOME"])) + try: + from torch.utils.cpp_extension import ROCM_HOME + + if ROCM_HOME: + candidates.append(pathlib.Path(ROCM_HOME)) + except ImportError: + pass + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(pathlib.Path(next(iter(spec.submodule_search_locations)))) + candidates.append(pathlib.Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + library_dir = rocm_home / "lib" + unversioned = library_dir / "libamdhip64.so" + link_dir = library_dir + if not unversioned.exists(): + versioned = sorted(library_dir.glob("libamdhip64.so.*")) + if not versioned: + continue + link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib" + link_dir.mkdir(parents=True, exist_ok=True) + compat_link = link_dir / "libamdhip64.so" + if not compat_link.exists() and not compat_link.is_symlink(): + try: + compat_link.symlink_to(versioned[-1]) + except FileExistsError: + # Multiple tensor-parallel ranks may prepare the same cache. + pass + + return [f"-L{link_dir}", f"-Wl,-rpath,{library_dir}"] + + raise RuntimeError("Unable to locate libamdhip64 for ROCm JIT linking") + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -223,12 +304,16 @@ def load_aot( if prebuilt is not None: return prebuilt + is_rocm = _is_rocm() arch_list: List[str] = [] if cuda_files: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + if is_rocm: + arch_list = _rocm_arch_list() + else: + from freetoken.kernel._toolchain import check_nvcc_matches_torch - check_nvcc_matches_torch() - arch_list = _cuda_arch_list() + check_nvcc_matches_torch() + arch_list = _cuda_arch_list() from tvm_ffi.cpp import load @@ -242,14 +327,23 @@ def load_aot( cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files] cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files] - with _pin_tvm_ffi_arch_ctx(arch_list): + if is_rocm: + cuda_cflags = _hip_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = _rocm_link_flags() + arch_list_env = ROCM_ARCH_LIST_ENV + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = [] + arch_list_env = CUDA_ARCH_LIST_ENV + + with _pin_tvm_ffi_arch_ctx(arch_list, arch_list_env): return load( name, cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags, arch_list), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) @@ -272,12 +366,16 @@ def load_jit( if prebuilt is not None: return prebuilt + is_rocm = _is_rocm() arch_list: List[str] = [] if cuda_files or cuda_wrappers: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + if is_rocm: + arch_list = _rocm_arch_list() + else: + from freetoken.kernel._toolchain import check_nvcc_matches_torch - check_nvcc_matches_torch() - arch_list = _cuda_arch_list() + check_nvcc_matches_torch() + arch_list = _cuda_arch_list() from tvm_ffi.cpp import load_inline @@ -300,14 +398,23 @@ def load_jit( cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] - with _pin_tvm_ffi_arch_ctx(arch_list): + if is_rocm: + cuda_cflags = _hip_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = _rocm_link_flags() + arch_list_env = ROCM_ARCH_LIST_ENV + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags, arch_list) + runtime_ldflags = [] + arch_list_env = CUDA_ARCH_LIST_ENV + + with _pin_tvm_ffi_arch_ctx(arch_list, arch_list_env): return load_inline( name, cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags, arch_list), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 2e4ad15f2f..af68bf5816 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,7 @@ from .arch import ( is_arch_supported, + is_rocm, + get_rocm_gfx_arch, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -35,6 +37,8 @@ "load_toolcall_anchor_id", "init_logger", "is_arch_supported", + "is_rocm", + "get_rocm_gfx_arch", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d56..304505966c 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -1,14 +1,66 @@ from __future__ import annotations import functools +import os +import re from typing import Tuple +_GFX_ARCH_RE = re.compile(r"gfx\d+[a-z]?") + + +def _gfx_arch_from(value: object) -> str | None: + match = _GFX_ARCH_RE.search(str(value).lower()) + return match.group(0) if match else None + + +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU) instead of CUDA.""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def get_rocm_gfx_arch() -> str | None: + """Return the current AMD GPU target (for example ``gfx1201``). + + Prefer the runtime device because build variables may contain multiple + semicolon-separated targets. Environment variables remain useful for + cross-compilation and systems where no GPU is currently visible; in that + fallback mode the first target is returned. The result is process-cached + for FreeToken's one-process-per-GPU execution model, so callers must select + the intended device before the first call. + """ + if not is_rocm(): + return None + + import torch + + if torch.cuda.is_available(): + try: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + for attr in ("gcnArchName", "arch"): + arch = _gfx_arch_from(getattr(props, attr, "")) + if arch: + return arch + except (AttributeError, RuntimeError): + pass + + for env_var in ("FREETOKEN_ROCM_ARCH", "PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): + arch = _gfx_arch_from(os.getenv(env_var, "")) + if arch: + return arch + return None + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch import torch.version + if is_rocm(): + return None if not torch.cuda.is_available() or not torch.version.cuda: return None return torch.cuda.get_device_capability() diff --git a/setup.py b/setup.py index 7132eea088..997745f1a8 100644 --- a/setup.py +++ b/setup.py @@ -1,15 +1,17 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path import sys from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension +from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension, ROCM_HOME ROOT = Path(__file__).parent +KERNEL_INCLUDE = str(ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include") def _check_toolchain() -> None: @@ -20,6 +22,43 @@ def _check_toolchain() -> None: module.check_nvcc_matches_torch() +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + +def _rocm_paths() -> tuple[list[str], list[str], str]: + candidates: list[Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(Path(os.environ["ROCM_HOME"])) + if ROCM_HOME: + candidates.append(Path(ROCM_HOME)) + + # ROCm 7.14 PyTorch images ship the SDK as a Python package instead of + # installing it at /opt/rocm. + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(Path(next(iter(spec.submodule_search_locations)))) + candidates.append(Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + include_dir = rocm_home / "include" + library_dir = rocm_home / "lib" + if not (include_dir / "hip" / "hip_runtime.h").exists(): + continue + if (library_dir / "libamdhip64.so").exists(): + return [str(include_dir)], [str(library_dir)], "amdhip64" + versioned = sorted(library_dir.glob("libamdhip64.so.*")) + if versioned: + return [str(include_dir)], [str(library_dir)], f":{versioned[-1].name}" + + searched = ", ".join(str(path) for path in dict.fromkeys(candidates)) + raise RuntimeError( + "A ROCm SDK with HIP headers and libamdhip64 is required to build on ROCm; " + f"searched: {searched}. Set ROCM_HOME to override." + ) + + def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if CUDA_HOME is None: raise RuntimeError( @@ -33,7 +72,21 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: return [str(cuda_home / "include")], library_dirs -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +IS_ROCM = _is_rocm() + +if IS_ROCM: + runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() + runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"] + # These extensions contain host code only. BuildExtension supplies the ROCm + # platform defines to the C++ compiler; offload architecture flags belong on + # HIP device sources and would be rejected by the host compiler here. + extra_compile = ["-O3", "-std=c++17"] +else: + runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() + runtime_lib = "cudart" + runtime_link_args = [] + extra_compile = ["-O3", "-std=c++17"] + _check_toolchain() @@ -44,12 +97,13 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/pinned_tensor.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17"], + include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile, + extra_link_args=runtime_link_args, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the + # CPU-compute MoE executor for --moe-backend cpu. Links cudart/amdhip64 for the # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels # use per-function target attributes (avx512bf16/avx512f) + a runtime # __builtin_cpu_supports dispatch, so the single binary stays portable @@ -59,10 +113,11 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17", "-pthread"], + include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile + ["-pthread"], + extra_link_args=runtime_link_args, ), # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies *([ diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index 8bcb8e5810..fa8a1f1d32 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/tests/kernels/test_e4m3_compat.py @@ -41,6 +41,29 @@ def _native_cc() -> bool: return torch.cuda.get_device_capability() >= (8, 9) +@pytest.mark.skipif(torch.version.hip is None, reason="needs ROCm") +def test_rocm_host_and_compile_time_native_gates_match(): + """The host wrapper and Triton constexpr must choose the same buffer ABI.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import e4m3_native, e4m3_native_cx + + @triton.jit + def gate_kernel(output): + if e4m3_native_cx(): + value = 1 + else: + value = 0 + tl.store(output, value) + + output = torch.empty(1, dtype=torch.int32, device="cuda") + gate_kernel[(1,)](output) + + assert e4m3_native() is False + assert output.item() == int(e4m3_native()) + + # ====================================================================================== # 1. Primitives vs the native fp8 unit (needs sm_89+ hardware for the reference). # ====================================================================================== diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py index e61108fd53..2fee4f25c8 100644 --- a/tests/kernels/test_pinned_tensor.py +++ b/tests/kernels/test_pinned_tensor.py @@ -122,7 +122,12 @@ def test_host_device_ptr_is_identity_under_uva(): pytest.skip("non-UVA platform: host_device_ptr rejects unregistered memory instead") # Under UVA cudaHostGetDevicePointer degenerates to identity for any host pointer # (no registration validation); rejection of pageable memory only exists on - # non-identity platforms (Windows/WDDM), where the translation is real. + # non-identity CUDA platforms (Windows/WDDM), where the translation is real. + # HIP validates registration even though registered/pinned memory uses the + # identity address on Linux. Calling it with pageable memory also leaves a + # sticky HIP error, so the pinned identity case above is the relevant check. + if torch.version.hip is not None: + return pageable = torch.empty(64, dtype=torch.uint8) ext = _load_pinned_extension() assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr() diff --git a/tests/kernels/test_rocm_launch_kwargs.py b/tests/kernels/test_rocm_launch_kwargs.py new file mode 100644 index 0000000000..52fac54519 --- /dev/null +++ b/tests/kernels/test_rocm_launch_kwargs.py @@ -0,0 +1,51 @@ +import pytest +import torch + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is None, + reason="a ROCm GPU is required", +) + + +class _KernelRecorder: + def __init__(self): + self.kwargs = None + + def __getitem__(self, _grid): + def launch(*_args, **kwargs): + self.kwargs = kwargs + + return launch + + +def test_rocm_activation_launch_omits_cuda_pdl_attribute(monkeypatch): + import freetoken.kernel.triton.activation as activation + + recorder = _KernelRecorder() + monkeypatch.setattr(activation, "_act_and_mul_kernel", recorder) + monkeypatch.setattr(activation, "_pdl_supported", lambda: True) + monkeypatch.setattr(activation, "is_rocm", lambda: True) + + activation.silu_and_mul(torch.ones((2, 128), device="cuda")) + + assert recorder.kwargs is not None + assert recorder.kwargs["ENABLE_PDL"] is True + assert "launch_pdl" not in recorder.kwargs + + +def test_rocm_norm_launch_omits_cuda_pdl_attribute(monkeypatch): + import freetoken.kernel.triton.norm as norm + + recorder = _KernelRecorder() + monkeypatch.setattr(norm, "_rmsnorm_kernel", recorder) + monkeypatch.setattr(norm, "is_sm90_supported", lambda: True) + monkeypatch.setattr(norm, "is_rocm", lambda: True) + + x = torch.ones((2, 64), device="cuda") + weight = torch.ones(64, device="cuda") + norm.rmsnorm(x, weight) + + assert recorder.kwargs is not None + assert recorder.kwargs["ENABLE_PDL"] is True + assert "launch_pdl" not in recorder.kwargs diff --git a/tests/utils/test_rocm_arch.py b/tests/utils/test_rocm_arch.py new file mode 100644 index 0000000000..fcfab75812 --- /dev/null +++ b/tests/utils/test_rocm_arch.py @@ -0,0 +1,127 @@ +import importlib +import os +import pathlib +import sys +import types +from types import SimpleNamespace + +import torch + +from freetoken.utils import arch + + +def _clear_arch_caches() -> None: + arch.get_rocm_gfx_arch.cache_clear() + + +def test_rocm_arch_prefers_visible_device_over_multi_arch_build_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(gcnArchName="gfx1201:sramecc-:xnack-"), + ) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1200") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1201" + + _clear_arch_caches() + + +def test_rocm_arch_falls_back_to_cross_compile_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1200" + + _clear_arch_caches() + + +def test_hip_cflags_target_only_resolved_arch(monkeypatch): + from freetoken.kernel import utils + + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + monkeypatch.setattr(arch, "get_rocm_gfx_arch", lambda: "gfx1201") + + arch_list = utils._rocm_arch_list() + flags = utils._hip_cflags(["-Wno-unused-command-line-argument"], arch_list) + + assert arch_list == ["gfx1201"] + assert "--offload-arch=gfx1201" in flags + assert "--offload-arch=gfx1200" not in flags + + +def test_rocm_loaders_pin_tvm_ffi_to_resolved_arch(monkeypatch): + from freetoken.kernel import utils + + builds = [] + + def record_build(*_args, **kwargs): + builds.append( + ( + os.environ[utils.ROCM_ARCH_LIST_ENV], + kwargs["extra_cuda_cflags"], + ) + ) + return object() + + tvm_ffi = types.ModuleType("tvm_ffi") + tvm_ffi.__path__ = [] + tvm_ffi_cpp = types.ModuleType("tvm_ffi.cpp") + tvm_ffi_cpp.load = record_build + tvm_ffi_cpp.load_inline = record_build + monkeypatch.setitem(sys.modules, "tvm_ffi", tvm_ffi) + monkeypatch.setitem(sys.modules, "tvm_ffi.cpp", tvm_ffi_cpp) + monkeypatch.setenv(utils.DISABLE_KERNEL_CACHE_ENV, "1") + monkeypatch.setenv(utils.ROCM_ARCH_LIST_ENV, "gfx1100 gfx1200") + monkeypatch.setattr(utils, "_is_rocm", lambda: True) + monkeypatch.setattr(utils, "_rocm_arch_list", lambda: ["gfx1201"]) + monkeypatch.setattr(utils, "_rocm_link_flags", lambda: []) + + utils.load_aot("test_rocm_aot_arch", cuda_files=["unused.cu"]) + utils.load_jit("test_rocm_jit_arch", cuda_files=["unused.cu"]) + + assert builds == [ + ("gfx1201", [*utils.DEFAULT_HIP_CFLAGS, "--offload-arch=gfx1201"]), + ("gfx1201", [*utils.DEFAULT_HIP_CFLAGS, "--offload-arch=gfx1201"]), + ] + assert os.environ[utils.ROCM_ARCH_LIST_ENV] == "gfx1100 gfx1200" + + +def test_rocm_link_flags_support_versioned_modular_sdk(monkeypatch, tmp_path): + import torch.utils.cpp_extension as cpp_extension + + from freetoken.kernel import utils + + sdk = tmp_path / "sdk" + library_dir = sdk / "lib" + library_dir.mkdir(parents=True) + versioned_runtime = library_dir / "libamdhip64.so.7" + versioned_runtime.write_bytes(b"") + real_find_spec = importlib.util.find_spec + + def find_spec(name: str): + if name == "_rocm_sdk_core": + return SimpleNamespace(submodule_search_locations=[str(sdk)]) + return real_find_spec(name) + + monkeypatch.delenv("ROCM_HOME", raising=False) + monkeypatch.setattr(cpp_extension, "ROCM_HOME", None) + monkeypatch.setattr(importlib.util, "find_spec", find_spec) + monkeypatch.setattr(pathlib.Path, "home", lambda: tmp_path) + utils._rocm_link_flags.cache_clear() + + flags = utils._rocm_link_flags() + + compat_dir = tmp_path / ".cache" / "freetoken" / "rocm-lib" + compat_link = compat_dir / "libamdhip64.so" + assert f"-L{compat_dir}" in flags + assert f"-Wl,-rpath,{library_dir}" in flags + assert compat_link.resolve() == versioned_runtime.resolve() + + utils._rocm_link_flags.cache_clear() From 5031e3d578a3b49572776c3536b3ff4b69769c9b Mon Sep 17 00:00:00 2001 From: YevheniiKotyrlo <44449990+YevheniiKotyrlo@users.noreply.github.com> Date: Wed, 30 Sep 2026 02:29:17 +0300 Subject: [PATCH 04/20] perf(qsa): reduce block top-k candidates in levels past 16 splits (#571) --- python/freetoken/kernel/triton/qsa/topk.py | 226 ++++++++++++++++++--- tests/models/qwen4_exp/test_qsa_kernels.py | 39 +++- 2 files changed, 225 insertions(+), 40 deletions(-) diff --git a/python/freetoken/kernel/triton/qsa/topk.py b/python/freetoken/kernel/triton/qsa/topk.py index ae57566cab..0f37a0f9d9 100644 --- a/python/freetoken/kernel/triton/qsa/topk.py +++ b/python/freetoken/kernel/triton/qsa/topk.py @@ -17,10 +17,18 @@ top-k is a subset of that union -- a global winner beats at most k-1 columns, so it beats all but at most k-1 of its own slice -- so the answer stays exact. ``_split_plan`` derives the geometry from the buffer width alone, which is fixed when a graph is captured. + +A row wide enough to leave more candidates than one program holds resident (262,144 columns, +a 1M-token window, leaves 32 x k) gets reduction levels between the two phases: each program +of a level takes ``_MAX_RESIDENT`` consecutive candidates of the level before and writes their +own top-k, until the survivors fit the merge. The same subset argument holds per level, and +candidates stay in ascending column order throughout, so ties still go to the lowest column. """ from __future__ import annotations +from dataclasses import dataclass + import torch import triton import triton.language as tl @@ -125,6 +133,44 @@ def _compact( return above, equal +@triton.jit +def _live_candidates( + visible, + CHUNK: tl.constexpr, + N_SPLITS: tl.constexpr, + TOP_K: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, +): + """Candidate slots level ``LEVEL`` filled for a row: its splits past the row's visible tail + are one skipped suffix, so every later stage keeps costing what the live part costs.""" + live = tl.minimum(tl.cdiv(visible, CHUNK), N_SPLITS) * TOP_K + for _ in tl.static_range(LEVEL): + live = tl.cdiv(live, RESIDENT) * TOP_K + return live + + +@triton.jit +def _emit_candidates( + key, + column, + k_eff, + key_out, + col_out, + TOP_K: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +): + """Write a resident tile's own top-k as ranked (key, column) candidates; returns how many.""" + prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES) + rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) + write = take & (rank < TOP_K) + tl.store(key_out + rank, key.to(tl.int32, bitcast=True), mask=write) + tl.store(col_out + rank, column, mask=write) + return above + tl.minimum(equal, ties) + + @triton.jit def _qsa_block_topk_kernel( logits_ptr, @@ -225,17 +271,95 @@ def _qsa_topk_split_kernel( if k_eff > 0: offsets = tl.arange(0, CHUNK) key = _load_keys(logits_ptr + row.to(tl.int64) * stride_logits_row + base, offsets, limit) - prefix, ties = _resident_prefix(key, k_eff, BINS, RADIX, PASSES) - rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) - write = take & (rank < TOP_K) - tl.store(key_ptr + slot + rank, key.to(tl.int32, bitcast=True), mask=write) - tl.store(col_ptr + slot + rank, (base + offsets).to(tl.int32), mask=write) - emitted = above + tl.minimum(equal, ties) + emitted = _emit_candidates( + key, (base + offsets).to(tl.int32), k_eff, key_ptr + slot, col_ptr + slot, + TOP_K, BINS, RADIX, PASSES, + ) # A key of 0 is the dead-column sentinel, so the merge needs no separate count per slot. pad = tl.arange(0, PAD_K) tl.store(key_ptr + slot + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) +@triton.jit +def _reduce_tile( + in_key, + in_col, + out_key, + out_col, + count, + k_eff, + TOP_K: tl.constexpr, + BLOCK: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +): + """Top-k of ``count`` candidates held as one resident tile; returns how many it wrote.""" + tl.static_assert(BLOCK <= 0xFFFF, "packed cumsum keeps 16 bits per half") + offsets = tl.arange(0, BLOCK) + inside = offsets < count + key = tl.load(in_key + offsets, mask=inside, other=0).to(tl.uint32, bitcast=True) + column = tl.load(in_col + offsets, mask=inside, other=-1) + return _emit_candidates(key, column, k_eff, out_key, out_col, TOP_K, BINS, RADIX, PASSES) + + +@triton.jit +def _qsa_topk_reduce_kernel( + visible_ptr, + in_key_ptr, + in_col_ptr, + out_key_ptr, + out_col_ptr, + stride_scratch_row, + num_columns, + TOP_K: tl.constexpr, + PAD_K: tl.constexpr, + CHUNK: tl.constexpr, + N_SPLITS: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, + BLOCK_SMALL: tl.constexpr, + BLOCK_MID: tl.constexpr, + BINS: tl.constexpr, + RADIX: tl.constexpr, + PASSES: tl.constexpr, +) -> None: + """A reduction level: one program per (row, ``RESIDENT`` candidates of level ``LEVEL - 1``), + writing their own top-k as this level's candidates.""" + row = tl.program_id(0) + split = tl.program_id(1) + visible = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0) + k_eff = tl.minimum(visible, TOP_K) + live = _live_candidates(visible, CHUNK, N_SPLITS, TOP_K, LEVEL - 1, RESIDENT) + base = split * RESIDENT + count = tl.minimum(tl.maximum(live - base, 0), RESIDENT) + source = row.to(tl.int64) * stride_scratch_row + base + slot = row.to(tl.int64) * stride_scratch_row + split * TOP_K + in_key, in_col = in_key_ptr + source, in_col_ptr + source + out_key, out_col = out_key_ptr + slot, out_col_ptr + slot + + emitted = 0 + # As in the merge: residency costs the whole tile, so a short range takes a narrower one. + if count > 0: + if count <= BLOCK_SMALL: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, BLOCK_SMALL, BINS, RADIX, PASSES, + ) + elif count <= BLOCK_MID: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, BLOCK_MID, BINS, RADIX, PASSES, + ) + else: + emitted = _reduce_tile( + in_key, in_col, out_key, out_col, count, k_eff, + TOP_K, RESIDENT, BINS, RADIX, PASSES, + ) + pad = tl.arange(0, PAD_K) + tl.store(out_key + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) + + @triton.jit def _merge_tile( key_row, @@ -274,6 +398,8 @@ def _qsa_topk_merge_kernel( PAD_K: tl.constexpr, CHUNK: tl.constexpr, N_SPLITS: tl.constexpr, + LEVEL: tl.constexpr, + RESIDENT: tl.constexpr, BLOCK_SMALL: tl.constexpr, BLOCK_MID: tl.constexpr, BLOCK_FULL: tl.constexpr, @@ -281,13 +407,11 @@ def _qsa_topk_merge_kernel( RADIX: tl.constexpr, PASSES: tl.constexpr, ) -> None: - """Phase 2: one program per row over the candidates phase 1 left behind.""" + """Phase 2: one program per row over the candidates the last level ``LEVEL`` left behind.""" row = tl.program_id(0) limit = tl.maximum(tl.minimum(tl.load(visible_ptr + row), num_columns), 0) k_eff = tl.minimum(limit, TOP_K) - # Candidates sit chunk-major, so the splits past the visible tail are one skipped suffix - # and the merge keeps costing what the live part of the row costs. - candidates = tl.minimum(tl.cdiv(limit, CHUNK), N_SPLITS) * TOP_K + candidates = _live_candidates(limit, CHUNK, N_SPLITS, TOP_K, LEVEL, RESIDENT) key_row = key_ptr + row.to(tl.int64) * stride_scratch_row col_row = col_ptr + row.to(tl.int64) * stride_scratch_row out_row = out_ptr + row.to(tl.int64) * stride_out_row @@ -325,27 +449,37 @@ def _qsa_topk_merge_kernel( tl.store(out_row + pad, -1, mask=(pad >= emitted) & (pad < TOP_K)) -def _split_plan(columns: int, top_k: int) -> tuple[int, int] | None: - """``(chunk, n_splits)`` for the split+merge path, or None to keep the one-program path.""" +@dataclass(frozen=True) +class _SplitPlan: + """``chunk`` columns per first-level program, and the program count of every level: the + first splits the row, each later one reduces the candidates of the level before it.""" + + chunk: int + levels: tuple[int, ...] + + +def _split_plan(columns: int, top_k: int) -> _SplitPlan | None: + """The split geometry of a ``columns``-wide buffer, or None to keep the one-program path.""" if top_k <= 0 or columns <= _MIN_CHUNK: return None max_splits = _MAX_RESIDENT // triton.next_power_of_2(top_k) if max_splits < 2: return None - chunk = max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits))) - if chunk > _MAX_RESIDENT: - return None + chunk = min(_MAX_RESIDENT, max(_MIN_CHUNK, triton.next_power_of_2(-(-columns // max_splits)))) n_splits = -(-columns // chunk) # Merging n_splits*top_k candidates has to be cheaper than scanning the row once. if n_splits < 2 or n_splits * top_k >= columns: return None - return chunk, n_splits + levels = [n_splits] + while levels[-1] * top_k > _MAX_RESIDENT: + levels.append(-(-levels[-1] * top_k // _MAX_RESIDENT)) + return _SplitPlan(chunk, tuple(levels)) def qsa_block_topk_scratch_width(columns: int, top_k: int) -> int: """int32 columns of scratch ``qsa_block_topk`` wants per row; 0 when it needs none.""" plan = _split_plan(columns, top_k) - return 0 if plan is None else 2 * plan[1] * top_k + return 0 if plan is None else 2 * top_k * sum(plan.levels) def qsa_block_topk( @@ -396,24 +530,31 @@ def qsa_block_topk( ) return out - chunk, n_splits = plan - half = n_splits * top_k + width = 2 * top_k * sum(plan.levels) if scratch is None: - scratch = torch.empty((rows, 2 * half), dtype=torch.int32, device=logits.device) + scratch = torch.empty((rows, width), dtype=torch.int32, device=logits.device) elif ( scratch.ndim != 2 or scratch.shape[0] < rows - or scratch.shape[1] < 2 * half + or scratch.shape[1] < width or scratch.dtype != torch.int32 or scratch.stride(1) != 1 ): raise ValueError( f"QSA block top-k needs a row-contiguous int32 scratch of at least " - f"[{rows}, {2 * half}], got {tuple(scratch.shape)} {scratch.dtype}" + f"[{rows}, {width}], got {tuple(scratch.shape)} {scratch.dtype}" ) - keys = scratch[:rows, :half] - cols = scratch[:rows, half : 2 * half] - _qsa_topk_split_kernel[(rows, n_splits)]( + # Each level's candidates: its keys, then their columns, level after level along the row. + candidates: list[tuple[torch.Tensor, torch.Tensor]] = [] + offset = 0 + for programs in plan.levels: + half = programs * top_k + candidates.append( + (scratch[:rows, offset : offset + half], scratch[:rows, offset + half : offset + 2 * half]) + ) + offset += 2 * half + keys, cols = candidates[0] + _qsa_topk_split_kernel[(rows, plan.levels[0])]( logits, visible, keys, @@ -423,14 +564,39 @@ def qsa_block_topk( columns, TOP_K=top_k, PAD_K=pad_k, - CHUNK=chunk, + CHUNK=plan.chunk, BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1, ) - merge_block = triton.next_power_of_2(half) + for level in range(1, len(plan.levels)): + (in_keys, in_cols), (out_keys, out_cols) = candidates[level - 1], candidates[level] + _qsa_topk_reduce_kernel[(rows, plan.levels[level])]( + visible, + in_keys, + in_cols, + out_keys, + out_cols, + keys.stride(0), + columns, + TOP_K=top_k, + PAD_K=pad_k, + CHUNK=plan.chunk, + N_SPLITS=plan.levels[0], + LEVEL=level, + RESIDENT=_MAX_RESIDENT, + BLOCK_SMALL=min(1024, _MAX_RESIDENT), + BLOCK_MID=min(4096, _MAX_RESIDENT), + BINS=_BINS, + RADIX=_RADIX, + PASSES=_PASSES, + num_warps=8, + num_stages=1, + ) + keys, cols = candidates[-1] + merge_block = triton.next_power_of_2(keys.shape[1]) _qsa_topk_merge_kernel[(rows,)]( visible, keys, @@ -441,8 +607,10 @@ def qsa_block_topk( columns, TOP_K=top_k, PAD_K=pad_k, - CHUNK=chunk, - N_SPLITS=n_splits, + CHUNK=plan.chunk, + N_SPLITS=plan.levels[0], + LEVEL=len(plan.levels) - 1, + RESIDENT=_MAX_RESIDENT, BLOCK_SMALL=min(1024, merge_block), BLOCK_MID=min(4096, merge_block), BLOCK_FULL=merge_block, diff --git a/tests/models/qwen4_exp/test_qsa_kernels.py b/tests/models/qwen4_exp/test_qsa_kernels.py index 4b678b3bdc..6715519f05 100644 --- a/tests/models/qwen4_exp/test_qsa_kernels.py +++ b/tests/models/qwen4_exp/test_qsa_kernels.py @@ -331,7 +331,17 @@ def test_torch_topk_env_picks_the_fallback(monkeypatch): # Block top-k, split + merge path (wide buffers) # -------------------------------------------------------------------------------------- -SPLIT_CHUNK = 4096 # _split_plan's chunk for every buffer these tests use +# (columns, _split_plan's chunk). A 1M-token window scores 262,144 compressed rows, and its 32 +# splits leave more candidates than one program holds, so a reduction level runs before the merge. +SPLIT_SHAPES = [(65536, 4096), (262144, 8192)] + + +def _split_chunk(columns: int, width: int) -> int: + from freetoken.kernel.triton.qsa.topk import _split_plan + + plan = _split_plan(columns, width) + assert plan is not None, f"{columns} columns must take the split path" + return plan.chunk def _policy_topk_blocks(logits, visible, width): @@ -348,29 +358,34 @@ def _policy_topk_blocks(logits, visible, width): return out -def _split_topk_case(n_blocks: int, width: int, bs: int, mode: str, seed: int): +def _split_topk_case(n_blocks: int, chunk: int, width: int, bs: int, mode: str, seed: int): if mode != "boundary": return _topk_case(n_blocks, bs, mode, seed) # width - 212 columns beat the tie, so the 212 remaining winners start 100 columns below - # a chunk boundary and run past it into a chunk that is all tie. + # a chunk boundary and run past it into a chunk that is all tie. The last row cuts at the + # middle split boundary, which is where a reduction level divides the candidates. logits = torch.zeros(bs, n_blocks, device="cuda") + splits = n_blocks // chunk for row in range(bs): - cut = SPLIT_CHUNK * (1 + row % (n_blocks // SPLIT_CHUNK - 1)) + cut = chunk * (splits // 2 if row == bs - 1 else 1 + row % (splits - 1)) logits[row, : width - 212] = 2.0 logits[row, cut - 100 :] = 1.0 return logits, torch.full((bs,), n_blocks, dtype=torch.int32, device="cuda") @requires_cuda -@pytest.mark.parametrize("n_blocks", [65536]) +@pytest.mark.parametrize("n_blocks,chunk", SPLIT_SHAPES) @pytest.mark.parametrize("bs", [4]) @pytest.mark.parametrize("mode", ["random", "boundary", "ragged", "dead"]) -def test_block_topk_split_path_matches_torch_topk(n_blocks: int, bs: int, mode: str): +def test_block_topk_split_path_matches_torch_topk(n_blocks: int, chunk: int, bs: int, mode: str): from freetoken.kernel.triton.qsa import qsa_block_topk, qsa_block_topk_scratch_width width = 512 assert qsa_block_topk_scratch_width(n_blocks, width) > 0, "case must take the split path" - logits, visible = _split_topk_case(n_blocks, width, bs, mode, seed=n_blocks + bs + len(mode)) + assert _split_chunk(n_blocks, width) == chunk + logits, visible = _split_topk_case( + n_blocks, chunk, width, bs, mode, seed=n_blocks + bs + len(mode) + ) blocks = torch.empty(bs, width, dtype=torch.int32, device=logits.device) qsa_block_topk(logits, visible, blocks) @@ -431,8 +446,10 @@ def test_block_topk_split_path_cost_tracks_live_blocks(): """A wide buffer with a short row must not pay for the splits past its visible tail.""" from freetoken.kernel.triton.qsa import qsa_block_topk, qsa_block_topk_scratch_width - # wide enough that the split work dwarfs the 20-launch floor even at boosted clocks - rows, columns, width = 1, 262144, 512 + # enough rows that the live splits fill the card, so their work outweighs the launch floor + rows, columns, width = 8, 262144, 512 + assert qsa_block_topk_scratch_width(columns, width) > 0, "case must take the split path" + chunk = _split_chunk(columns, width) device = torch.device("cuda") logits = torch.randn(rows, columns, device=device) visible = torch.full((rows,), columns, dtype=torch.int32, device=device) @@ -463,8 +480,8 @@ def replay_us(live: int) -> float: best = min(best, start.elapsed_time(stop) * 1000.0 / 20) return best - full, short = replay_us(columns), replay_us(SPLIT_CHUNK) - assert full > 2.0 * short, f"{columns} live {full:.1f}us vs {SPLIT_CHUNK} live {short:.1f}us" + full, short = replay_us(columns), replay_us(chunk) + assert full > 2.0 * short, f"{columns} live {full:.1f}us vs {chunk} live {short:.1f}us" @requires_cuda From eba7120eabde70c7414e0e14c551e1caa3b0f174 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 29 Sep 2026 19:32:36 -0700 Subject: [PATCH 05/20] feat(install): record the install root as FREETOKEN_HOME in environment.d (#576) --- install.sh | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/install.sh b/install.sh index 2a9a3dec30..f623ffece7 100755 --- a/install.sh +++ b/install.sh @@ -20,7 +20,8 @@ # Defaults to $DEFAULT_KERNEL_CACHE_WHEEL_URL; if # unset and FREETOKEN_WHEEL is local, the script # auto-detects a sibling freetoken_kernel_cache-*.whl. -# FREETOKEN_HOME install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv +# FREETOKEN_HOME install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv; +# a successful install writes it to environment.d # FREETOKEN_PY_VERSION python for the venv (default: 3.12 — must match the wheel tag) # FREETOKEN_BIN_DIR where to symlink `ft` (default: ~/.local/bin) # FREETOKEN_ENV_DIR environment.d dir (default: ~/.config/environment.d) @@ -248,8 +249,8 @@ ln -sf "$FT_BIN" "$BIN_DIR/ft" say "symlinked $BIN_DIR/ft -> $FT_BIN" mkdir -p "$ENV_DIR" -printf 'FREETOKEN_FT_BIN=%s\n' "$FT_BIN" > "$ENV_DIR/50-freetoken.conf" -say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_FT_BIN) — GUI picks it up after next login" +printf 'FREETOKEN_HOME=%s\n' "$(cd "$FT_HOME" && pwd)" >"$ENV_DIR/50-freetoken.conf" +say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_HOME) - FreeToken Desktop reads the file, other programs see the variable after next login" # --- 5. Self-check --------------------------------------------------------- if "$FT_BIN" --help >/dev/null 2>&1; then @@ -264,10 +265,9 @@ ${C_GREEN}FreeToken engine installed.${C_RESET} ft binary $FT_BIN on PATH as $BIN_DIR/ft (ensure $BIN_DIR is on PATH) - Desktop env FREETOKEN_FT_BIN via environment.d (re-login to apply) + Desktop env FREETOKEN_HOME in $ENV_DIR/50-freetoken.conf -Run in this shell without re-login: - export FREETOKEN_FT_BIN="$FT_BIN" +Run: ft serve --model --port 1919 EOF From 235e2016ddb64c8eb1f0d877237062f087fb2ab2 Mon Sep 17 00:00:00 2001 From: YevheniiKotyrlo <44449990+YevheniiKotyrlo@users.noreply.github.com> Date: Thu, 1 Oct 2026 10:40:27 +0300 Subject: [PATCH 06/20] feat(rotary): size the rope table for YaRN and add --hf-overrides (#573) * feat(rotary): size the rope table for YaRN and add --hf-overrides * fix(rotary): guard mrope's table width and copy --hf-overrides values - MRotaryEmbedding asserts its cos/sin table is rotary_dim wide. A proportional rope's table spans the head, and the mrope kernel indexes the rotary_dim // 2 section table with the table's width. - _with_overrides deep-copies every value it stores: transformers fills rope defaults into the rope_parameters it is handed, which was the caller's dict. - The rope-table tests the engine test covers fold into it, the two shared-table tests are one, and the --hf-overrides test moves to tests/utils/test_hf.py, where it loads the override against the same values written into config.json, for a model type transformers knows and one it does not. - docs/cli.md lists --hf-overrides. * fix(server): give --hf-overrides a fresh dict per parse argparse hands its default object to every parse that omits the flag, so default={} put one shared dict in every ServerArgs. The flag defaults to None and becomes a new dict where the arguments are assembled, as --mm-processor-kwargs does. * docs(rotary): an indexer's heads differ in width from the attention's, not wider --- docs/cli.md | 1 + python/freetoken/attention/qsa_sparse.py | 17 +- python/freetoken/engine/config.py | 8 +- python/freetoken/engine/engine.py | 9 +- python/freetoken/layers/rotary.py | 193 +++++++++--------- python/freetoken/models/config.py | 7 + python/freetoken/server/args.py | 12 ++ python/freetoken/utils/hf.py | 31 ++- tests/engine/test_attention_backend_matrix.py | 4 +- tests/engine/test_cache_budget.py | 21 +- tests/kernels/test_mrope.py | 13 ++ tests/kernels/test_rotary.py | 61 +++++- tests/models/qwen4_exp/test_config.py | 51 +++++ tests/models/test_qwen3_vl.py | 2 +- tests/utils/test_hf.py | 34 +++ 15 files changed, 328 insertions(+), 136 deletions(-) create mode 100644 tests/utils/test_hf.py diff --git a/docs/cli.md b/docs/cli.md index 0c33943987..fb5e613f7b 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -33,6 +33,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--model-path`, `--model` | required | Local dir, HF repo id, or an FTW dir (auto-detected) | | `--served-model-name` | basename of `--model` | Model id reported by `/v1/models` | +| `--hf-overrides` | — | JSON object applied to the checkpoint's config as vLLM's `--hf-overrides`: a nested config section updates key by key, any other value is replaced whole. A YaRN `rope_parameters` override serves `original_max_position_embeddings * factor` positions | ### Server & runtime diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index b07c952323..8ce63c4242 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -167,23 +167,22 @@ def _cmp_pages(self, slot: int) -> torch.Tensor: def _index_rope_cache(self) -> torch.Tensor: """cos/sin table of the indexer rope: same rotary_dim and frequencies as the main - attention, ``head_size`` 128 instead of 256, so it is a separate get_rope instance. + attention with ``head_size`` 128 instead of 256, so the attention's own table. The table itself (not RotaryEmbedding.forward) because the indexer's norm+rope is one fused kernel and the compressed keys rope at their group's position, not the query's.""" if self._index_cos_sin is None: - from freetoken.layers.rotary import get_rope + from freetoken.layers.rotary import rope_cos_sin_table rotary = self.rotary_config with torch.device(self.device): - rope = get_rope( - head_dim=self.index_head_dim, - rotary_dim=rotary.rotary_dim, - max_position=rotary.max_position, - base=rotary.base, - rope_scaling=tuple(rotary.scaling.items()) if rotary.scaling else None, + self._index_cos_sin = rope_cos_sin_table( + self.index_head_dim, + rotary.rotary_dim, + rotary.max_position, + rotary.base, + tuple(rotary.scaling.items()) if rotary.scaling else None, ) - self._index_cos_sin = rope._cos_sin_cache.to(self.device) return self._index_cos_sin def _index_rope_rows(self, positions: torch.Tensor) -> torch.Tensor: diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 44c5a38e73..3e38cb2165 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -3,7 +3,7 @@ import copy from dataclasses import dataclass, field, replace from functools import cached_property -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, Any, List, Mapping import torch from freetoken.distributed import DistributedInfo @@ -23,6 +23,8 @@ class EngineConfig: model_path: str tp_info: DistributedInfo dtype: torch.dtype + # --hf-overrides: applied to the checkpoint config the model is built from (cached_load_hf_config) + hf_overrides: Mapping[str, Any] = field(default_factory=dict) max_running_req: int = 4 attention_backend: str = "auto" moe_strategy: str = "auto" @@ -103,7 +105,7 @@ def __post_init__(self): @cached_property def hf_config(self): - return cached_load_hf_config(self.model_path) + return cached_load_hf_config(self.model_path, self.hf_overrides) @cached_property def model_spec(self) -> ModelSpec: @@ -142,7 +144,7 @@ def model_config(self) -> ModelConfig: def max_seq_len(self) -> int: if self.max_seq_len_override is not None: return self.max_seq_len_override - return self.model_config.rotary_config.max_position + return self.model_config.rotary_config.table_positions @property def max_forward_len(self) -> int: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cc14cf5339..e855bcd32f 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1770,18 +1770,17 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) override("num_page_override", config.num_token_override // config.page_size) - # The rope cos/sin table is baked to rotary_config.max_position, and neither rope kernel + # The rope cos/sin table covers rotary_config.table_positions, and neither rope kernel # bounds-checks the position it gathers with -- a longer ceiling reads past the table. # DSV4 is exempt: it sizes its own table from the resolved max_seq_len (_adjust_dsv4_config). rotary = getattr(model_config, "rotary_config", None) seq_override = getattr(config, "max_seq_len_override", None) if seq_override is not None and rotary is not None and not is_dsv4: - if seq_override > rotary.max_position: + if seq_override > rotary.table_positions: raise ValueError( f"--max-seq-len-override {seq_override} exceeds the model's " - f"rope table ({rotary.max_position} positions). Serving past it would read " - "out of bounds; extend the checkpoint's rope_scaling / " - "max_position_embeddings in config.json instead." + f"rope table ({rotary.table_positions} positions). Serving past it would read " + "out of bounds; extend the rope with --hf-overrides (YaRN rope_parameters) instead." ) # The startup ServerArgs dump is the *requested* config, printed in the frontend process diff --git a/python/freetoken/layers/rotary.py b/python/freetoken/layers/rotary.py index 0d04645d99..d0178f3bc1 100644 --- a/python/freetoken/layers/rotary.py +++ b/python/freetoken/layers/rotary.py @@ -2,7 +2,7 @@ import functools import math -from typing import Any, Callable, Dict, Tuple +from typing import Any, Mapping, Tuple import torch @@ -14,50 +14,21 @@ def __init__( self, head_size: int, rotary_dim: int, - max_position_embeddings: int, - base: float, - post_process: None | Callable[[torch.Tensor], torch.Tensor] = None, - proportional: bool = False, - attention_factor: float = 1.0, + cos_sin_cache: torch.Tensor, is_neox: bool = True, ) -> None: super().__init__() + assert 0 < rotary_dim <= head_size + assert rotary_dim % 2 == 0 + assert head_size in [64, 128, 256, 512] self.head_size = head_size self.rotary_dim = rotary_dim # NeoX (half-rotation, HF default) vs GPT-J interleaved (adjacent pairs, # ``rope_interleave`` models: GLM MLA lineage). Both underlying kernels # accept the flag; the cos/sin cache layout is identical. self.is_neox = is_neox - if proportional: - assert 0 < rotary_dim <= head_size - assert rotary_dim % 2 == 0 - inv_freq = 1.0 / ( - base ** (torch.arange(0, head_size, 2, dtype=torch.float) / head_size) - ) - if rotary_dim < head_size: - inv_freq[rotary_dim // 2 :] = 0.0 - else: - # Standard (NeoX) rope. Supports partial rotary (rotary_dim < head_size): - # rope is applied to the first ``rotary_dim`` dims of each head, the rest pass - # through. Frequencies are spaced over ``rotary_dim`` (matches HF default - # partial rope, e.g. Qwen3.5 partial_rotary_factor, MiniMax-M2's - # ``apply_rotary_pos_emb``). Full rope is rotary_dim == head_size and is - # unaffected. ``head_size`` is passed to flashinfer separately so it rotates - # only the first ``rotary_dim`` dims. - assert 0 < rotary_dim <= head_size - assert rotary_dim % 2 == 0 - inv_freq = 1.0 / ( - base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim) - ) - if post_process is not None: - inv_freq = post_process(inv_freq) - t = torch.arange(max_position_embeddings, dtype=torch.float) - freqs = torch.einsum("i,j -> ij", t, inv_freq) - cos = freqs.cos() * attention_factor - sin = freqs.sin() * attention_factor - # buffer, so don't load/save - self._cos_sin_cache = torch.cat((cos, sin), dim=-1) - assert self.head_size in [64, 128, 256, 512] + # shared with every rope of the same frequencies (rope_cos_sin_table), so never written + self._cos_sin_cache = cos_sin_cache from freetoken.kernel.backend import is_flashinfer_installed @@ -160,6 +131,7 @@ def __init__( ) -> None: super().__init__(*args, **kwargs) assert self.is_neox, "mrope is defined on the NeoX half-rotation layout" + assert self._cos_sin_cache.shape[1] == self.rotary_dim, "mrope reads a rotary_dim-wide table, not a proportional rope's" half = self.rotary_dim // 2 assert sum(mrope_section) == half, (mrope_section, half) self._section_table = build_section_table(tuple(mrope_section), layout) @@ -193,56 +165,50 @@ def forward( return query, key -def _get_rope( - head_dim: int, - rotary_dim: int, - max_position: int, - base: float, - rope_scaling: Dict[str, Any] | None = None, - is_neox: bool = True, -) -> RotaryEmbedding: - if rope_scaling is None: - return RotaryEmbedding(head_dim, rotary_dim, max_position, base, is_neox=is_neox) - # need to test some cases: - match rope_scaling["rope_type"]: +def rope_table_positions(max_position: int, rope_scaling: Mapping[str, Any] | None) -> int: + """Rows the cos/sin table needs: YaRN serves ``original_max_position_embeddings * factor``, + above ``max_position_embeddings`` when only the rope parameters are overridden.""" + if rope_scaling is None or rope_scaling.get("rope_type") != "yarn": + return max_position + extended = int(rope_scaling["original_max_position_embeddings"] * rope_scaling["factor"]) + return max(max_position, extended) + + +def _rope_frequencies( + head_size: int, rotary_dim: int, base: float, rope_scaling: Mapping[str, Any] | None +) -> Tuple[torch.Tensor, float]: + """``inv_freq`` over the table's half width, and the factor cos and sin are both scaled by.""" + rope_type = "default" if rope_scaling is None else rope_scaling["rope_type"] + if rope_type == "proportional": + # Frequencies spaced over the whole head; the pairs past rotary_dim do not rotate. + inv_freq = 1.0 / (base ** (torch.arange(0, head_size, 2, dtype=torch.float) / head_size)) + inv_freq[rotary_dim // 2 :] = 0.0 + return inv_freq, 1.0 + # Standard (NeoX) rope. Supports partial rotary (rotary_dim < head_size): rope is applied + # to the first ``rotary_dim`` dims of each head, the rest pass through, and frequencies + # are spaced over ``rotary_dim`` (HF default partial rope, e.g. Qwen3.5 + # partial_rotary_factor, MiniMax-M2's ``apply_rotary_pos_emb``). + inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)) + match rope_type: case "default": - return RotaryEmbedding(head_dim, rotary_dim, max_position, base, is_neox=is_neox) - - case "proportional": - return RotaryEmbedding( - head_dim, - rotary_dim, - max_position, - base, - proportional=True, - is_neox=is_neox, - ) + return inv_freq, 1.0 case "llama3": scaling_factor: float = rope_scaling["factor"] low_freq_factor: float = rope_scaling["low_freq_factor"] high_freq_factor: float = rope_scaling["high_freq_factor"] original_max_position: int = rope_scaling["original_max_position_embeddings"] - - def post_process(inv_freq: torch.Tensor) -> torch.Tensor: - # no smooth if low_freq_factor == high_freq_factor - wave_len = 2 * math.pi / inv_freq - if low_freq_factor == high_freq_factor: - return torch.where( - wave_len < original_max_position / high_freq_factor, - inv_freq, - inv_freq / scaling_factor, - ) - - delta = high_freq_factor - low_freq_factor - smooth = (original_max_position / wave_len - low_freq_factor) / delta - smooth = torch.clamp(smooth, 0, 1) - factor = (1 - smooth) / scaling_factor + smooth - return factor * inv_freq - - return RotaryEmbedding( - head_dim, rotary_dim, max_position, base, post_process, is_neox=is_neox - ) + wave_len = 2 * math.pi / inv_freq + # no smooth if low_freq_factor == high_freq_factor + if low_freq_factor == high_freq_factor: + return torch.where( + wave_len < original_max_position / high_freq_factor, + inv_freq, + inv_freq / scaling_factor, + ), 1.0 + delta = high_freq_factor - low_freq_factor + smooth = torch.clamp((original_max_position / wave_len - low_freq_factor) / delta, 0, 1) + return ((1 - smooth) / scaling_factor + smooth) * inv_freq, 1.0 case "yarn": factor: float = rope_scaling["factor"] @@ -288,27 +254,54 @@ def _find_correction_dim(num_rotations: float) -> float: high = min(high, rotary_dim - 1) if low == high: # HF nudges instead of flooring the gap at 1 ("truncate": false) high += 0.001 - - def post_process(inv_freq: torch.Tensor) -> torch.Tensor: - ramp = torch.clamp( - (torch.arange(rotary_dim // 2, dtype=torch.float32) - low) / (high - low), - 0, 1, - ) - return (inv_freq / factor) * ramp + inv_freq * (1 - ramp) - - return RotaryEmbedding( - head_dim, - rotary_dim, - max_position, - base, - post_process, - attention_factor=float(attention_factor), - is_neox=is_neox, + ramp = torch.clamp( + (torch.arange(rotary_dim // 2, dtype=torch.float32) - low) / (high - low), 0, 1 ) + return (inv_freq / factor) * ramp + inv_freq * (1 - ramp), float(attention_factor) raise ValueError(f"Unsupported {rope_scaling = }") +@functools.cache +def _cos_sin_table( + head_size: int, + rotary_dim: int, + positions: int, + base: float, + rope_scaling: Tuple[Tuple[str, Any], ...] | None, + device: torch.device, +) -> torch.Tensor: + inv_freq, attention_factor = _rope_frequencies( + head_size, rotary_dim, base, dict(rope_scaling) if rope_scaling is not None else None + ) + freqs = torch.arange(positions, dtype=torch.float)[:, None] * inv_freq[None, :] + return torch.cat((freqs.cos() * attention_factor, freqs.sin() * attention_factor), dim=-1) + + +def rope_cos_sin_table( + head_size: int, + rotary_dim: int, + max_position: int, + base: float, + rope_scaling: Tuple[Tuple[str, Any], ...] | None = None, +) -> torch.Tensor: + """The ``[positions, width]`` cos|sin table, one per distinct rope: only proportional rope + depends on the head size, so an indexer whose heads differ in width reads the attention's table.""" + rope_map = dict(rope_scaling) if rope_scaling is not None else None + proportional = rope_map is not None and rope_map["rope_type"] == "proportional" + device = torch.get_default_device() + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return _cos_sin_table( + head_size if proportional else rotary_dim, + rotary_dim, + rope_table_positions(max_position, rope_map), + base, + rope_scaling, + device, + ) + + _ROPE_DEVICE: torch.device | None = None @@ -328,17 +321,15 @@ def get_rope( mrope_section: Tuple[int, ...] | None = None, mrope_layout: str = "interleaved", ) -> RotaryEmbedding: - rope_map = dict(rope_scaling) if rope_scaling is not None else None - def build() -> RotaryEmbedding: + table = rope_cos_sin_table(head_dim, rotary_dim, max_position, base, rope_scaling) if mrope_section is not None: - assert rope_map is None or rope_map.get("rope_type", "default") == "default" return MRotaryEmbedding( - head_dim, rotary_dim, max_position, base, + head_dim, rotary_dim, table, is_neox=is_neox, mrope_section=tuple(mrope_section), layout=mrope_layout, ) - return _get_rope(head_dim, rotary_dim, max_position, base, rope_map, is_neox) + return RotaryEmbedding(head_dim, rotary_dim, table, is_neox=is_neox) t = torch.tensor([]) if t.device == torch.device("meta"): @@ -358,5 +349,7 @@ def build() -> RotaryEmbedding: "RotaryEmbedding", "build_section_table", "get_rope", + "rope_cos_sin_table", + "rope_table_positions", "set_rope_device", ] diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index d7b504a769..c5943dce1b 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -92,6 +92,13 @@ class RotaryConfig: mrope_section: list | None = None mrope_layout: str = "contiguous" # see freetoken.layers.rotary.build_section_table + @property + def table_positions(self) -> int: + """Positions the cos/sin table covers: the longest sequence this rope can serve.""" + from freetoken.layers.rotary import rope_table_positions + + return rope_table_positions(self.max_position, self.scaling) + def mrope_layout_from_rope_params(rope_params: Any) -> str: """rope_parameters flags -> layout name: mrope_interleaved_glm, else mrope_interleaved, else contiguous.""" diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index ab2fb9b748..23a62f8419 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -272,6 +272,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Data type for model weights and activations. 'auto' will use FP16 for FP32/FP16 models and BF16 for BF16 models.", ) + parser.add_argument( + "--hf-overrides", + type=_json_object, + default=None, + metavar="JSON", + help="JSON object applied to the checkpoint config the model is built from, as vLLM's " + "--hf-overrides: a nested config section is updated key by key, any other value is " + "replaced whole. A YaRN rope_parameters override extends the servable context to " + "original_max_position_embeddings * factor.", + ) + parser.add_argument( "--tensor-parallel-size", "--tp-size", @@ -892,6 +903,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: image_max_tokens=image_max_tokens, processor_kwargs=kwargs.pop("mm_processor_kwargs") or {}, ) + kwargs["hf_overrides"] = kwargs["hf_overrides"] or {} result = ServerArgs(**kwargs) logger.info(f"Parsed arguments:\n{result}") return result, run_shell diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py index b93d3b1203..489585e8c2 100644 --- a/python/freetoken/utils/hf.py +++ b/python/freetoken/utils/hf.py @@ -1,9 +1,8 @@ +import copy import functools import json import os -from typing import Any - -from typing import FrozenSet +from typing import Any, FrozenSet, Mapping from huggingface_hub import hf_hub_download, snapshot_download from huggingface_hub.utils import EntryNotFoundError @@ -235,20 +234,40 @@ def _load_hf_config(model_path: str) -> Any: return config -def cached_load_hf_config(model_path: str) -> PretrainedConfig: +def _with_overrides(config: Any, data: dict, overrides: Mapping[str, Any]) -> dict: + merged = dict(data) + for key, value in overrides.items(): + section = getattr(config, key, None) + if isinstance(value, Mapping) and isinstance(section, (PretrainedConfig, RawConfigShim)): + merged[key] = _with_overrides(section, merged.get(key) or {}, value) + else: + merged[key] = copy.deepcopy(value) + return merged + + +def cached_load_hf_config( + model_path: str, overrides: Mapping[str, Any] | None = None +) -> PretrainedConfig: + """A fresh copy of the checkpoint's config with ``overrides`` applied as vLLM's --hf-overrides + does: nested config sections update key by key, every other value is replaced whole.""" # A .gguf file (or an FTW dir converted from one) carries its own metadata (no HF # config.json); return a shim the model registry dispatches on instead of a # PretrainedConfig. from freetoken.models.gguf.reader import gguf_config_source if (gguf_src := gguf_config_source(model_path)) is not None: + if overrides: + raise ValueError("--hf-overrides applies to a HF config.json; a GGUF carries its own metadata") from freetoken.models.gguf.config import build_gguf_shim return build_gguf_shim(gguf_src) config = _load_hf_config(model_path) + data = config.to_dict() + if overrides: + data = _with_overrides(config, data, overrides) if isinstance(config, RawConfigShim): - return RawConfigShim(config.to_dict()) - return type(config)(**config.to_dict()) + return RawConfigShim(data) + return type(config)(**data) def _weight_allow_patterns(repo_id: str) -> list[str]: diff --git a/tests/engine/test_attention_backend_matrix.py b/tests/engine/test_attention_backend_matrix.py index ef7dc7ff85..ad07a9dd09 100644 --- a/tests/engine/test_attention_backend_matrix.py +++ b/tests/engine/test_attention_backend_matrix.py @@ -40,6 +40,8 @@ def _spec(name, attn_type, *, mla=False, sliding_window=None, index_head_dim=0, def _model_config(kind): + from freetoken.models.config import RotaryConfig + mc = SimpleNamespace( model_type=kind, single_stream_only=False, @@ -48,7 +50,7 @@ def _model_config(kind): has_swa_attention=False, has_linear_attention=False, num_layers=4, - rotary_config=SimpleNamespace(max_position=1024), + rotary_config=RotaryConfig(head_dim=128, rotary_dim=128, max_position=1024, base=1e4, scaling=None), ) if kind == "full": specs = (_spec("full", AttnType.FULL),) diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a4e47aa02f..073112bb5a 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -430,13 +430,17 @@ def test_page_table_width_covers_whole_trailing_pages(): assert w > last_col and w % 32 == 0 -def _generic_rotary_cfg(max_position, override): +def _generic_rotary_cfg(max_position, override, scaling=None): from types import SimpleNamespace + from freetoken.models.config import RotaryConfig + model_config = SimpleNamespace( single_stream_only=False, is_moe=False, expert_quant="none", has_swa_attention=False, has_linear_attention=False, - rotary_config=SimpleNamespace(max_position=max_position), + rotary_config=RotaryConfig( + head_dim=128, rotary_dim=128, max_position=max_position, base=1e4, scaling=scaling + ), ) class Cfg: @@ -476,6 +480,19 @@ def test_adjust_config_allows_override_at_rope_table_boundary(): _adjust_config(_generic_rotary_cfg(max_position=1024, override=1024)) # must not raise +YARN_4X = {"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 1024} + + +def test_adjust_config_allows_override_to_the_yarn_extended_table(): + # A YaRN override that leaves max_position_embeddings at the trained length (Qwen3.8's 1M + # recipe) serves original * factor positions; the gate must follow the table, not the config. + from freetoken.engine.engine import _adjust_config + + _adjust_config(_generic_rotary_cfg(max_position=1024, override=4096, scaling=YARN_4X)) + with pytest.raises(ValueError, match=r"rope table \(4096 positions\)"): + _adjust_config(_generic_rotary_cfg(max_position=1024, override=4097, scaling=YARN_4X)) + + def test_adjust_config_rope_gate_exempts_dsv4(): # DSV4 sizes its own rope table from the resolved max_seq_len (_adjust_dsv4_config), # so the generic gate must not fire even when the override dwarfs max_position. diff --git a/tests/kernels/test_mrope.py b/tests/kernels/test_mrope.py index f8402f8754..08dff5cb57 100644 --- a/tests/kernels/test_mrope.py +++ b/tests/kernels/test_mrope.py @@ -32,6 +32,19 @@ def test_section_tables(): build_section_table(SECTION, "diagonal") +def test_refuses_a_table_wider_than_the_rotary_dims(): + # A proportional rope's table spans the whole head; mrope takes its width from the table + # and indexes the rotary_dim // 2 section table with it, reading past the end. + from freetoken.layers.rotary import get_rope + + get_rope.cache_clear() + with pytest.raises(AssertionError, match="rotary_dim"): + get_rope( + head_dim=HEAD, rotary_dim=ROT, max_position=8, base=1e7, + rope_scaling=(("rope_type", "proportional"),), mrope_section=SECTION, + ) + + def _make(mrope: bool, layout: str = "interleaved"): from freetoken.layers.rotary import get_rope diff --git a/tests/kernels/test_rotary.py b/tests/kernels/test_rotary.py index f3e8654fc8..82ee01537d 100644 --- a/tests/kernels/test_rotary.py +++ b/tests/kernels/test_rotary.py @@ -101,7 +101,7 @@ def test_yarn_correction_range_matches_hf_where_the_clamp_and_gap_bind(): longest-wavelength dims; flooring the gap at 1 flattens the ramp for a sub-1 gap.""" from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS - from freetoken.layers.rotary import get_rope + from freetoken.layers.rotary import _rope_frequencies class _Shim: # duck-typed config for HF's real _compute_yarn_parameters def __init__(self, head_dim, max_position, rope_parameters): @@ -122,16 +122,18 @@ def standardize_rope_params(self): (128, 1e6, {"factor": 16.0, "beta_fast": 1.15, "beta_slow": 1.0, "original_max_position_embeddings": 32768, "truncate": False}), ] + def table(inv_freq: torch.Tensor, attention_factor: float) -> torch.Tensor: + freqs = torch.outer(torch.arange(max_position, dtype=torch.float), inv_freq.float()) + return torch.cat((freqs.cos() * attention_factor, freqs.sin() * attention_factor), dim=-1) + + # The frequencies, not get_rope's table: YaRN sizes that to original * factor (up to 2M + # rows here), and the correction range is a property of the frequencies alone. for rotary_dim, base, scaling in cases: - get_rope.cache_clear() - rope = get_rope(head_dim=rotary_dim, rotary_dim=rotary_dim, max_position=max_position, - base=base, rope_scaling=(("rope_type", "yarn"), *scaling.items())) + ours = table(*_rope_frequencies(rotary_dim, rotary_dim, base, {"rope_type": "yarn", **scaling})) params = {"rope_type": "yarn", "rope_theta": base, "partial_rotary_factor": 1.0, **scaling} - inv_freq, attn_factor = ROPE_INIT_FUNCTIONS["yarn"]( - _Shim(rotary_dim, max_position, params), device=torch.device("cpu")) - freqs = torch.outer(torch.arange(max_position, dtype=torch.float), inv_freq.float()) - expected = torch.cat((freqs.cos() * attn_factor, freqs.sin() * attn_factor), dim=-1) - torch.testing.assert_close(rope._cos_sin_cache, expected, rtol=0, atol=1e-6) + expected = table(*ROPE_INIT_FUNCTIONS["yarn"]( + _Shim(rotary_dim, max_position, params), device=torch.device("cpu"))) + torch.testing.assert_close(ours, expected, rtol=0, atol=1e-6) def test_yarn_mscale_all_dim_zero_falls_back_like_hf(): @@ -160,3 +162,44 @@ def test_yarn_mscale_all_dim_zero_falls_back_like_hf(): expected = 0.1 * math.log(factor) + 1.0 # 1.277, not 1.196 (= 1 + 0.1*0.707*ln 16) cos = rope._cos_sin_cache[:, :32] torch.testing.assert_close(cos[0], torch.full_like(cos[0], expected)) + + +def _yarn(factor: float, original: int) -> tuple: + return (("rope_type", "yarn"), ("factor", factor), ("original_max_position_embeddings", original)) + + +def test_yarn_table_covers_the_extended_context(): + # A checkpoint extended by overriding only its rope parameters (Qwen3.8's 1M recipe) keeps + # max_position_embeddings at the trained length; YaRN serves original * factor positions. + from freetoken.layers.rotary import get_rope + + get_rope.cache_clear() + rope = get_rope(head_dim=64, rotary_dim=64, max_position=64, base=10000.0, rope_scaling=_yarn(4.0, 64)) + assert rope._cos_sin_cache.shape == (256, 64) + + longer = get_rope(head_dim=64, rotary_dim=64, max_position=512, base=10000.0, rope_scaling=_yarn(4.0, 64)) + assert longer._cos_sin_cache.shape == (512, 64) + torch.testing.assert_close(longer._cos_sin_cache[:256], rope._cos_sin_cache, rtol=0, atol=0) + + +def test_one_table_serves_every_rope_of_the_same_frequencies(): + # The QSA indexer ropes 128-wide heads with the attention's 64 rotary dims, and mrope reads + # the same rows per axis: one table, not a second copy of a 1M-row fp32 tensor. + from freetoken.layers.rotary import MRotaryEmbedding, get_rope + + get_rope.cache_clear() + attention = get_rope(head_dim=256, rotary_dim=64, max_position=64, base=1e7, rope_scaling=_yarn(4.0, 64)) + indexer = get_rope(head_dim=128, rotary_dim=64, max_position=64, base=1e7, rope_scaling=_yarn(4.0, 64)) + mrope = get_rope( + head_dim=256, rotary_dim=64, max_position=64, base=1e7, rope_scaling=_yarn(4.0, 64), + mrope_section=(11, 11, 10), mrope_layout="interleaved", + ) + assert indexer.head_size == 128 and isinstance(mrope, MRotaryEmbedding) + assert attention._cos_sin_cache.shape == (256, 64) + table = attention._cos_sin_cache.data_ptr() + assert indexer._cos_sin_cache.data_ptr() == table and mrope._cos_sin_cache.data_ptr() == table + + proportional = (("rope_type", "proportional"),) + wide = get_rope(head_dim=128, rotary_dim=64, max_position=4, base=1e4, rope_scaling=proportional) + narrow = get_rope(head_dim=64, rotary_dim=64, max_position=4, base=1e4, rope_scaling=proportional) + assert wide._cos_sin_cache.shape == (4, 128) and narrow._cos_sin_cache.shape == (4, 64) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index be532eeca9..f850927b99 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -215,3 +215,54 @@ def test_released_checkpoints_keep_the_dense_projections_bf16(quantization_confi quant = _quant(_hf_config(quantization_config), tmp_path) for prefix in DENSE_PREFIXES + BF16_PREFIXES: assert quant.scheme_for(prefix) is None, prefix + + +# Qwen3.8-Flash-Next model card, "Processing Ultra-Long Texts": the 1M-token override. +MODEL_CARD_YARN = { + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + "rope_type": "yarn", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + "factor": 4.0, + "original_max_position_embeddings": 262144, +} + + +def _write_checkpoint_config(path, **fields) -> None: + import json + + text = {k: v for k, v in vars(_text_config()).items()} + config = { + "model_type": "qwen4_exp", + "architectures": ["Qwen4ExpForConditionalGeneration"], + "image_token_id": 248056, + "text_config": text, + "quantization_config": RADIXARK_NVFP4, + **fields, + } + (path / "config.json").write_text(json.dumps(config)) + + +def test_the_engine_serves_the_model_card_1m_recipe(tmp_path): + import torch + + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + from freetoken.mm.config import ENCODER_KINDS, MultimodalConfig + + _write_checkpoint_config(tmp_path) + config = EngineConfig( + model_path=str(tmp_path), + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.bfloat16, + hf_overrides={"text_config": {"rope_parameters": MODEL_CARD_YARN}}, + mm=MultimodalConfig(disabled_encoders=frozenset(ENCODER_KINDS)), + ) + rotary = config.model_config.rotary_config + assert rotary.scaling["rope_type"] == "yarn" and rotary.scaling["factor"] == 4.0 and rotary.mrope_section is None + assert rotary.max_position == 262144 and rotary.rotary_dim == 64 and rotary.base == 10000000 + assert config.max_seq_len == 4 * 262144 + plain = EngineConfig(model_path=str(tmp_path), tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16, + mm=MultimodalConfig(disabled_encoders=frozenset(ENCODER_KINDS))) + assert plain.model_config.rotary_config.scaling is None and plain.max_seq_len == 262144 diff --git a/tests/models/test_qwen3_vl.py b/tests/models/test_qwen3_vl.py index c45a9f150e..6c773d3deb 100644 --- a/tests/models/test_qwen3_vl.py +++ b/tests/models/test_qwen3_vl.py @@ -39,7 +39,7 @@ def _engine_config(monkeypatch, hf, processor, **overrides): import freetoken.engine.config as engine_config import freetoken.mm.processor as mm_processor - monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf) + monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path, overrides: hf) monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None) monkeypatch.setattr(mm_processor, "get_mm_processor", lambda path, mm=None: processor) return EngineConfig( diff --git a/tests/utils/test_hf.py b/tests/utils/test_hf.py new file mode 100644 index 0000000000..28630b1431 --- /dev/null +++ b/tests/utils/test_hf.py @@ -0,0 +1,34 @@ +"""cached_load_hf_config with --hf-overrides, against the same config written into config.json.""" + +import copy +import json + +import pytest + +from freetoken.utils.hf import cached_load_hf_config + + +def _checkpoint(path, model_type: str, dtype: str, rope_parameters: dict) -> str: + path.mkdir() + text_config = {"hidden_size": 64, "max_position_embeddings": 4096, "rope_parameters": rope_parameters} + (path / "config.json").write_text(json.dumps({"model_type": model_type, "dtype": dtype, "text_config": text_config})) + return str(path) + + +@pytest.mark.parametrize("model_type", ["qwen3_5", "a_model_type_transformers_does_not_know"]) +def test_hf_overrides_load_as_if_the_checkpoint_config_said_so(tmp_path, model_type): + # vLLM's merge: a nested config section updates key by key, any other value is replaced whole + yarn = {"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 4096} + checkpoint = _checkpoint(tmp_path / "checkpoint", model_type, "bfloat16", {"rope_type": "default", "rope_theta": 1e6}) + edited = _checkpoint(tmp_path / "edited", model_type, "float16", yarn) + overrides = {"dtype": "float16", "text_config": {"rope_parameters": yarn}} + requested = copy.deepcopy(overrides) + + hf = cached_load_hf_config(checkpoint, overrides) + expected = cached_load_hf_config(edited) + assert hf.dtype == expected.dtype + assert hf.text_config.to_dict() == expected.text_config.to_dict() + # transformers fills rope defaults into the rope_parameters it is handed + assert overrides == requested + # a fresh copy each time: the override never reaches the cached checkpoint config + assert cached_load_hf_config(checkpoint).text_config.rope_parameters["rope_type"] == "default" From d3512b43affe981465e03ee28cbd88f49c39b9aa Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Thu, 1 Oct 2026 16:59:45 -0700 Subject: [PATCH 07/20] fix(ple): fence disk staging against in-flight lookups (#588) Co-authored-by: andy-yang-1 --- python/freetoken/models/qwen4_exp/ple_disk.py | 31 ++++++++++++++++--- tests/models/qwen4_exp/test_ple_disk.py | 18 +++++++++++ 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index ec7da80d56..69ee9fa93a 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -150,9 +150,18 @@ def __init__( self._graph_dev = torch.empty( max_graph_rows * self._token_bytes, dtype=torch.uint8, device=self._device ) + # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the + # dispatch context after the step's fill: a record inside the capture would not be observable) + self._graph_consumed = torch.cuda.Event() eager_bytes = max_extend_tokens * self._token_bytes - self._eager_pinned = alloc_pinned_tensor(eager_bytes, dtype=torch.uint8) - self._eager_pinned.zero_() # the warmup prefill stages nothing and reads whatever sits here + # the overlap scheduler fills batch k+1 while batch k's lookup copy may still be queued: the + # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy + # that last read it has run (``_eager_read[i]``, recorded right after that copy) + self._eager_pinned = [alloc_pinned_tensor(eager_bytes, dtype=torch.uint8) for _ in range(2)] + for buf in self._eager_pinned: + buf.zero_() # the warmup prefill stages nothing and reads whatever sits here + self._eager_read = [torch.cuda.Event(), torch.cuda.Event()] + self._eager_slot = 0 self._eager_dev = torch.empty(eager_bytes, dtype=torch.uint8, device=self._device) # probe picks flag-sync (graph WAITs at the consume, host fills then signals) or launch-gating from freetoken.kernel.row_store import probe_wait_sync @@ -170,7 +179,13 @@ def __init__( def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: """Stage per-request token runs (two context ids, then the new tokens) in batch order.""" - pinned = self._graph_pinned if graph else self._eager_pinned + if graph: + self._graph_consumed.synchronize() + pinned = self._graph_pinned + else: + self._eager_slot ^= 1 + self._eager_read[self._eager_slot].synchronize() + pinned = self._eager_pinned[self._eager_slot] offset = 0 for run in runs: self._store.stage(run.data_ptr(), run.numel() - 2, pinned.data_ptr() + offset * self._token_bytes) @@ -236,21 +251,27 @@ def forward_host_ctx(self, batch: Batch, use_graph: bool): # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run if deferred is not None: deferred() + if use_graph: + # only after this step's fill: the next fill waits on this launch, not the fill on it + self._graph_consumed.record(torch.cuda.current_stream(self._device)) # ---------------- device side (PLETableBackend protocol) ---------------- def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: rows = row_ids.shape[0] + stream = torch.cuda.current_stream(self._device) capturing = torch.cuda.is_current_stream_capturing() if capturing and self._wait_sync: from freetoken.kernel.row_store import wait_reset - wait_reset(torch.cuda.current_stream(self._device), self._flag) + wait_reset(stream, self._flag) pinned, dev = ( - (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned, self._eager_dev) + (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned[self._eager_slot], self._eager_dev) ) nbytes = rows * self._token_bytes dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True) + if not capturing: + self._eager_read[self._eager_slot].record(stream) values = dev[:nbytes].view(torch.float8_e4m3fn).to(self.dtype) if self.scale != 1.0: values = values * self.scale diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index 1112210b84..ba9cd98c63 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -287,6 +287,24 @@ def test_disk_table_matches_oracle(tmp_path): assert _bitwise_equal(disk.lookup(ids), oracle.lookup(ids)), "hook prefill" +@requires_cuda +def test_eager_fill_does_not_overwrite_an_in_flight_lookup(tmp_path): + """The overlap scheduler fills batch k+1 while batch k's lookup copy is still queued behind the + layers before it (a sleep kernel here): batch k must read its own rows.""" + disk, oracle, args = _make_table(tmp_path) + emb = _embedding() + steps = [(3, 4, 7), (4, 7, 9), (7, 9, 13), (9, 13, 5)] + ids = [emb.row_ids(_meta([[token]], [[older, newer]], decode=True)).cuda() for older, newer, token in steps] + got = [] + for (older, newer, token), step_ids in zip(steps, ids): + disk.fill([torch.tensor([older, newer, token])], graph=False) + torch.cuda._sleep(50_000_000) + got.append(disk.lookup(step_ids)) + torch.cuda.synchronize() + for i, (out, step_ids) in enumerate(zip(got, ids)): + assert _bitwise_equal(out, oracle.lookup(step_ids)), f"step {i} read another step's rows" + + @requires_cuda def test_graph_sync_protocol(tmp_path, monkeypatch): disk, oracle, args = _make_table(tmp_path) From 08c223ea55c720f4a3c97a9a9ee759e835a06a5f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 18:23:04 -0400 Subject: [PATCH 08/20] test(kvcache): give the qsa fp8 pool fakes the mrope flag main now reads Assisted-by: Claude Opus 5 (cherry picked from commit 150c95d6f39b83e0516f5c9c42fef54c679eae09) (cherry picked from commit 4db8b858375b01ff7edf37031dc81671b0181853) (cherry picked from commit 495770e33dcb5c3aa6979c6515b61dfade5e29b3) --- tests/kvcache/test_qsa_pool_fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py index c0a91f862a..02b525f679 100644 --- a/tests/kvcache/test_qsa_pool_fp8.py +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -74,7 +74,7 @@ def _spec(): def _config(kv_quant, *, page_size=PAGE_SIZE, max_running_req=3): - mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True) + mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True, model_is_mrope=False) mc.kv_cache_group_specs = lambda: (_spec(),) return SimpleNamespace( model_config=mc, @@ -200,7 +200,7 @@ def test_factory_threads_kv_quant_into_the_qsa_pool(): mc = SimpleNamespace( num_layers=8, has_swa_attention=False, has_linear_attention=True, - num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, + num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, model_is_mrope=False, ) mc.kv_cache_group_specs = lambda: (_spec(),) pool = create_kvcache_pool( From 8ff7f4ad0a7e0c8c437c041fdb8005e8fc02ae38 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:02:28 -0400 Subject: [PATCH 09/20] fix(models): restore the os import and truthy set the fp8 env helpers read Assisted-by: Claude Opus 5 (cherry picked from commit 82a21ae2432fb891175c8c881f52cb4b710241bc) (cherry picked from commit 0369ed94d7d085ece40ffdcc0adf93a5c26a6b4e) (cherry picked from commit b5c7812fceb68b4b7835a80ab7c533e0385d9577) --- python/freetoken/models/config.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 117d43010e..afe8000c2b 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -1,4 +1,6 @@ from __future__ import annotations + +import os from dataclasses import dataclass from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias @@ -7,12 +9,14 @@ # State-dict key prefixes of the vision stack; load_weight drops them when the engine serves text-only. VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.") +_ENV_TRUE = {"1", "true", "yes", "on"} + def fp8_dense_enabled() -> bool: """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF): per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``.""" - return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE + return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _ENV_TRUE def fp8_lmhead_enabled() -> bool: @@ -20,7 +24,7 @@ def fp8_lmhead_enabled() -> bool: for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits: a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own quality gate. ``FREETOKEN_FP8_LMHEAD=1``.""" - return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _VISION_TRUE + return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _ENV_TRUE def detect_expert_quant(hf_config: Any) -> str: From a1398c8ae662c0d7a8e32003da992a7b51324696 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:00:44 -0400 Subject: [PATCH 10/20] test(qwen4_exp): restore the rope device and clear get_rope after the cpu build test Assisted-by: Claude Opus 5 (cherry picked from commit 27ad92e0a554be016d6e533683c40a2eb67c6164) (cherry picked from commit e152f665cc1f35f073fbb82daedece77d02f91e1) (cherry picked from commit 333f7f9f0501116694ede067202070353faf341b) --- tests/models/qwen4_exp/test_build_cpu.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py index 0cf2027d14..58ad0c4dc3 100644 --- a/tests/models/qwen4_exp/test_build_cpu.py +++ b/tests/models/qwen4_exp/test_build_cpu.py @@ -10,7 +10,7 @@ import torch from freetoken.distributed import set_tp_info, try_get_tp_info -from freetoken.layers import set_rope_device +from freetoken.layers import rotary, set_rope_device from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM @@ -20,9 +20,15 @@ def _build(): if try_get_tp_info() is None: set_tp_info(rank=0, size=1) + previous = rotary._ROPE_DEVICE set_rope_device(torch.device("cpu")) - with torch.device("meta"): - return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + try: + with torch.device("meta"): + return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + finally: + # get_rope is cached per geometry: a CPU cos/sin table left behind feeds later CUDA tests + rotary.get_rope.cache_clear() + set_rope_device(previous) def test_model_constructs_on_meta_without_cuda(): From 5dee633f54ccc994be2ea6660e69d2edebdb0e6f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:19:07 -0400 Subject: [PATCH 11/20] test(kernels): give the block_ends backend fake the kv scale accessors the quantized-kv backend reads Assisted-by: Claude Opus 5 (cherry picked from commit 600819fd7ce18c813cd4b8400cf534bff38f6f7f) (cherry picked from commit dd2ebdfda3e1c53e5269bd9adfbee7291cc8ffc2) (cherry picked from commit 12754d41c5fe646ccc6081ce26f05fdf7b0a57b9) --- tests/kernels/test_triton_attention.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index a5abb3d7b3..3176ad71ce 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -804,6 +804,12 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + def k_scale(self, layer_id): # 16-bit pool: no fp8 scales + return None + + def v_scale(self, layer_id): + return None + device = torch.device("cuda") head_dim = 256 kv_cache = FakeKVCache(device, head_dim) From b5c3c4f399becdbe5e532d3383515c9691cdd4fd Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:53:22 -0400 Subject: [PATCH 12/20] fix(models): skip the ftw layout check for a bare ftw without an hf config Fork #19 asks the model family to check the stored fp8 split against the env flags; a synthetic FTW written straight from tensors (upstream's test_ftw_weights) carries no config.json, so there is no family to ask and the read proceeds as before. Assisted-by: Claude Fable 5.1 (cherry picked from commit cab9be35b8e56fa3cf08a109017183691e2e69c5) (cherry picked from commit e59e5441e44dac96d530f4e2cac34d2c6273411f) --- python/freetoken/models/weight.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 9016b9477e..29cce4eefa 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -224,9 +224,13 @@ def load_weight( # a text-only engine never built the tower, so its tensors are not even read keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES)) if is_ftw_checkpoint(model_path): - spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) - # a family whose stored tensors depend on load-time settings checks them against the index before the read - check = _model_override(spec, "check_ftw_weights") + # a family whose stored tensors depend on load-time settings checks them against the index before the + # read; a bare FTW without a readable HF config (synthetic test checkpoints) has no family to ask + try: + spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) + except (OSError, ValueError): + spec = None + check = _model_override(spec, "check_ftw_weights") if spec is not None else None if check is not None: reader = FTWReader(model_path) names = {e["name"] for e in reader.entries("weight")} From 230059071a2e56143a51d64269b3404cb014d329 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:53:25 -0400 Subject: [PATCH 13/20] test(scheduler): give the rejection test's engine fake the pool size the admission clamp reads Assisted-by: Claude Fable 5.1 (cherry picked from commit 4584957f3d63449a153097f380ffdac035a245dd) (cherry picked from commit 0ab1764343228845db270505185dd2be7090086d) --- tests/scheduler/test_cost_accounting_core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py index 045010c6b4..116d4b88c7 100644 --- a/tests/scheduler/test_cost_accounting_core.py +++ b/tests/scheduler/test_cost_accounting_core.py @@ -132,7 +132,8 @@ def fail_prepare(_batch): def test_scheduler_rejection_emits_error_but_no_admission(): scheduler = Scheduler.__new__(Scheduler) - scheduler.engine = SimpleNamespace(max_seq_len=4) + scheduler.engine = SimpleNamespace(max_seq_len=4, num_pages=4) + scheduler.config = SimpleNamespace(page_size=1) added = [] scheduler.prefill_manager = SimpleNamespace(add_one_req=added.append) sent = [] From 8430f004e60923c5c57353d6080de5ceb8514595 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:41:19 -0400 Subject: [PATCH 14/20] test(server): abort every sample of an n>1 non-streaming request on disconnect Guards the #222 resolution against #393's fan-out: the disconnect watcher takes the whole uid list, not the first sample's uid. Assisted-by: Claude Fable 5.1 (cherry picked from commit 93ff96349235ea3ca1504733a2abbd6e6bcccf94) (cherry picked from commit 69a5efa3300e93bd331b4394e54eea8cef82ad3f) (cherry picked from commit 40e994121fb79a21ba0d7ffff70e4394a4498899) --- tests/server/test_stream_cancellation.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py index e38fa56bfb..ba7da7c17b 100644 --- a/tests/server/test_stream_cancellation.py +++ b/tests/server/test_stream_cancellation.py @@ -294,3 +294,25 @@ def test_non_stream_connected_client_gets_result_without_abort(monkeypatch): assert result["choices"][0]["message"]["content"] == "Hi" assert state.aborted == [] + + +def test_chat_non_stream_disconnect_aborts_every_sample(monkeypatch): + # n > 1 fans out into one engine request per sample; a disconnect must stop all of them + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + class _State(_ApiState): + def __init__(self): + super().__init__(acks=None) + self._next = 7 + + def new_user(self): + self._next += 1 + return self._next + + state = _State() + resp = asyncio.run( + openai_api.handle_chat_completion(_chat_req(n=2), _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert sorted(state.aborted) == [8, 9] From 8ece6379730d5d2f35ee0d169434892eb0c4da18 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 23:04:45 -0400 Subject: [PATCH 15/20] fix(server): record requests with a pure asgi middleware so handlers can see a client disconnect Starlette's @app.middleware("http") (BaseHTTPMiddleware) wraps the handler's receive channel and never yields the client's http.disconnect while a non-streaming request is still computing, so request.is_disconnected() stays False and the disconnect watcher from FlashML-org/FreeToken#222 could not stop an abandoned request. Streaming requests only worked because the socket write failed. The recorder is the one such middleware; as pure ASGI it records the same row at response start. Assisted-by: Claude Fable 5.1 (cherry picked from commit 8adde913c94e2f76b3962b8d10f10ded4b1cccc6) (cherry picked from commit 202d7fec30b29a9a4320370b0c62426422422339) --- python/freetoken/server/api_server.py | 79 ++++++++++++------- .../test_request_recorder_disconnect.py | 75 ++++++++++++++++++ 2 files changed, 124 insertions(+), 30 deletions(-) create mode 100644 tests/server/test_request_recorder_disconnect.py diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index b097ea9c6f..1a039402f1 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -470,38 +470,57 @@ def _served_model_name() -> str | None: return getattr(cfg, "served_model_name", None) -@app.middleware("http") -async def _record_request_middleware(request: Request, call_next): +class _RecordRequestMiddleware: """Time every generation request into the ring for /v1/requests + /v1/stats p95. Single- model server, so model = served_model_name; stream is inferred from the response media - type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None.""" - path = request.url.path - if path.startswith(_UNTRACKED_REQUEST_PREFIXES) or not path.startswith( - _TRACKED_REQUEST_PREFIXES - ): - return await call_next(request) - import time as _time - - start = _time.monotonic() - response = await call_next(request) - duration_ms = int((_time.monotonic() - start) * 1000) - ctype = response.headers.get("content-type", "") - request_ring.record_request( - request_ring.RequestRecord( - ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), - method=request.method, - path=path, - status=response.status_code, - model=_served_model_name(), - duration_ms=duration_ms, - ttft_ms=None, - prompt_tokens=None, - completion_tokens=None, - stream=ctype.startswith("text/event-stream"), - error=None, - ) - ) - return response + type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None. + + Pure ASGI on purpose: Starlette's ``@app.middleware("http")`` (BaseHTTPMiddleware) hands the + handler a wrapped receive channel that never yields the client's ``http.disconnect`` while a + non-streaming request is still computing, so ``request.is_disconnected()`` stays False and an + abandoned request runs to max_tokens. The recorder is the only such middleware here.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + path = scope.get("path", "") if scope["type"] == "http" else "" + if not path.startswith(_TRACKED_REQUEST_PREFIXES) or path.startswith(_UNTRACKED_REQUEST_PREFIXES): + await self.app(scope, receive, send) + return + import time as _time + + start = _time.monotonic() + + async def send_recording(message): + # the base-class version recorded once the handler returned its response object, + # i.e. at response start, before a streaming body -- keep that timing + if message["type"] == "http.response.start": + ctype = "" + for name, value in message.get("headers", ()): + if name.lower() == b"content-type": + ctype = value.decode("latin-1") + request_ring.record_request( + request_ring.RequestRecord( + ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), + method=scope["method"], + path=path, + status=message["status"], + model=_served_model_name(), + duration_ms=int((_time.monotonic() - start) * 1000), + ttft_ms=None, + prompt_tokens=None, + completion_tokens=None, + stream=ctype.startswith("text/event-stream"), + error=None, + ) + ) + await send(message) + + await self.app(scope, receive, send_recording) + + +app.add_middleware(_RecordRequestMiddleware) class CacheRebuildRequest(BaseModel): diff --git a/tests/server/test_request_recorder_disconnect.py b/tests/server/test_request_recorder_disconnect.py new file mode 100644 index 0000000000..06ff17f99b --- /dev/null +++ b/tests/server/test_request_recorder_disconnect.py @@ -0,0 +1,75 @@ +"""The request recorder must not hide a client disconnect from a non-streaming handler. + +Real sockets on purpose: the failure only shows through uvicorn's receive channel (a +BaseHTTPMiddleware wrapper never yields ``http.disconnect`` while the handler computes), which +TestClient does not model. +""" + +from __future__ import annotations + +import asyncio +import socket +import threading +import time + +from fastapi import FastAPI, Request + +from freetoken.server.api_server import _RecordRequestMiddleware + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _serve(app: FastAPI, port: int): + import uvicorn + + server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + deadline = time.monotonic() + 10 + while not server.started: + assert time.monotonic() < deadline, "uvicorn did not start" + time.sleep(0.02) + return server + + +def _post_and_hang_up(port: int) -> None: + body = b'{"prompt": "x"}' + with socket.create_connection(("127.0.0.1", port)) as sock: + sock.sendall( + b"POST /v1/completions HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n" + b"Content-Length: " + str(len(body)).encode() + b"\r\n\r\n" + body + ) + time.sleep(0.2) # the handler is now polling + + +def test_non_streaming_handler_sees_the_client_leave(): + app = FastAPI() + app.add_middleware(_RecordRequestMiddleware) + seen: dict[str, float | None] = {} + + @app.post("/v1/completions") + async def slow(request: Request): + t0 = time.monotonic() + for _ in range(60): + await asyncio.sleep(0.05) + if await request.is_disconnected(): + seen["after"] = time.monotonic() - t0 + return {"detected": True} + seen["after"] = None + return {"detected": False} + + port = _free_port() + server = _serve(app, port) + try: + _post_and_hang_up(port) + deadline = time.monotonic() + 5 + while "after" not in seen and time.monotonic() < deadline: + time.sleep(0.05) + finally: + server.should_exit = True + assert seen.get("after") is not None, "handler never saw the disconnect" + assert seen["after"] < 2.0 From f02c3bbcce2f8b120c0720e293f1ff6edc33c5ce Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 19 Sep 2026 17:21:12 -0400 Subject: [PATCH 16/20] fix(kvcache): refuse the host kv tiers on an nvfp4 pool whose block scales they cannot mirror Upstream #499/#525 mirror KV codes plus the fp8 pool's per-token _scale_buffer; the nvfp4 pool (#408) keeps per-block scales in _block_scale_buffer, which neither the offloader nor the prefix tier copies, so --kv-host-pages or FT_PREFIX_HOST would rehydrate codes without scales. Both constructors now raise a clear error naming the knob; the GDN-snapshot-only tier (FT_GDN_HOST_TIER) never reads the KV pool and stays usable. Assisted-by: Claude Fable 5.1 (cherry picked from commit 4bb35fc34b0cb65b41584c65d7390b6c5f4f854a) --- python/freetoken/kvcache/host_prefix_tier.py | 3 ++ python/freetoken/kvcache/kv_host_offload.py | 14 +++++++ tests/kvcache/test_host_tier_nvfp4_guard.py | 44 ++++++++++++++++++++ 3 files changed, 61 insertions(+) create mode 100644 tests/kvcache/test_host_tier_nvfp4_guard.py diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 0265b126ef..3b2d2d5dc3 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -79,6 +79,9 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, self._kv_scales: HostBank | None = None self._kv_free: list[int] = [] if kv_pool is not None and kv_budget_bytes > 0: + from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool + + _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST") codes = kv_pool._kv_buffer # [2, L, P, page, H, D] scales = kv_pool._scale_buffer # [2, L, P*page, H] | None (fp8) _, layers, _, page_size, heads, dim = codes.shape diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py index dbb49f4814..1f8c358426 100644 --- a/python/freetoken/kvcache/kv_host_offload.py +++ b/python/freetoken/kvcache/kv_host_offload.py @@ -33,10 +33,24 @@ logger = init_logger(__name__) +def _reject_block_scaled_pool(pool, knob: str) -> None: + """The host mirrors carry codes plus the fp8 pool's per-token scales; an nvfp4 pool + (--kv-cache-dtype nvfp4) keeps per-block scales in _block_scale_buffer, which no mirror + copies, so a rehydrated page would be codes without scales. Refuse instead of decoding + garbage.""" + if getattr(pool, "_block_scale_buffer", None) is not None: + raise ValueError( + f"{knob} needs a KV pool with per-token scales (bf16 or fp8); the nvfp4 pool keeps " + "block scales the host mirror does not carry -- leave the tier off or serve with " + "--kv-cache-dtype fp8" + ) + + class KVHostOffloader: """GPU-slot LRU cache over the pinned host mirror of the logical KV page space.""" def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None: + _reject_block_scaled_pool(pool, "--kv-host-pages") buf = pool._kv_buffer # [2, L, P_gpu + 1, page_size, kv_heads, head_dim] assert buf.shape[0] == 2 and buf.is_contiguous() _, num_layers, num_slots, page_size, kv_heads, head_dim = buf.shape diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py new file mode 100644 index 0000000000..a793fb294a --- /dev/null +++ b/tests/kvcache/test_host_tier_nvfp4_guard.py @@ -0,0 +1,44 @@ +"""The host KV tiers (upstream #499 / #525) mirror codes plus the fp8 pool's per-token scales; +an nvfp4 pool keeps block scales they do not copy, so both must refuse it up front.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.kv_host_offload import KVHostOffloader + + +def _kv_pool(block_scaled: bool): + return SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 3, 4, 1, 8, dtype=torch.uint8), # [2, L, P+1, page, H, D] + _scale_buffer=None, + _block_scale_buffer=torch.zeros(1) if block_scaled else None, + ) + + +def _state_pool(): + return SimpleNamespace( + conv_states=torch.zeros(2, 2, 4, 3), + recurrent_states=torch.zeros(2, 2, 2, 4), + slot_states={}, + ) + + +def test_kv_host_offloader_refuses_an_nvfp4_pool(): + with pytest.raises(ValueError, match="--kv-host-pages .* nvfp4"): + KVHostOffloader(_kv_pool(block_scaled=True), num_logical_pages=4, device=torch.device("cpu")) + + +def test_prefix_tier_refuses_kv_pages_on_an_nvfp4_pool(): + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* nvfp4"): + HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=1 << 16, device="cpu") + + +def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool(): + # the snapshot-only tier (no KV budget) never reads the pool, nvfp4 or not + tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=0, device="cpu") + assert tier.num_slots == 2 and not tier.has_kv From b658deeb7c998a3e639f8acb34058956146f5b2e Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:57:03 -0400 Subject: [PATCH 17/20] test(scheduler): give the scheduler stubs the prefill-start map the metrics timing reads Upstream #504 times the prefill span in Scheduler._prefill_start; the stubs from #222 and #505's tests build the scheduler without it. Assisted-by: Claude Opus 5.5 --- tests/scheduler/test_abort_inflight_prefill.py | 1 + tests/scheduler/test_hybrid_cache_manager.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py index 85c217859f..f2cbd6e0d2 100644 --- a/tests/scheduler/test_abort_inflight_prefill.py +++ b/tests/scheduler/test_abort_inflight_prefill.py @@ -63,6 +63,7 @@ def _setup(): config=SimpleNamespace(page_size=1), status_reporter=SimpleNamespace(report_batch=lambda *_, **__: None), send_result=sent.extend, + _prefill_start={}, _kv_usage_pages=cm.page_usage, _mamba_slot_usage=lambda: None, _swa_token_usage=lambda: None, diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index 282ab2ca43..07385d7046 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/tests/scheduler/test_hybrid_cache_manager.py @@ -213,7 +213,7 @@ def test_chunked_prefill_retains_resumable_snapshot(monkeypatch, full_chunks, ta assert torch.all(pool.conv_states[:, match.mamba_value] == states[expected]) # Finish/abort may occur before the prefill commit; use the real idempotent cleanup path. - stub = SimpleNamespace(cache_manager=cm, table_manager=tm) + stub = SimpleNamespace(cache_manager=cm, table_manager=tm, _prefill_start={}) Scheduler._free_req_resources(stub, final) Scheduler._free_req_resources(stub, final) match = cm.match_req(_pend(prompt.tolist() + [999])) From fb6b9a5d67e1a8f110e62dd93bb8db842adb7782 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:59:17 -0400 Subject: [PATCH 18/20] test(engine): give the logprobs sampler fakes the logits-processing flag the sampler plan reads Fork #47's sampler plan asks each request whether it needs logits processing; the upstream #224 test fakes predate it. Assisted-by: Claude Opus 5.5 --- tests/engine/test_sample_logprobs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index aabbdc5fe7..6bd204604d 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -137,8 +137,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)), ] ) args = sampler.prepare(batch) From 4985d83c39e65aaebf87a8417082f7cc000d265d Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 24 Sep 2026 00:10:43 -0400 Subject: [PATCH 19/20] revert the post-load FP8 shared expert / hyper-connection merge A/B on the 5080 (offload, 1650 slots, 32K and 2K contexts, 4 reps each): no decode gain (~38.3 vs ~38.8 tok/s), idle VRAM not reduced, and the quality probe lost one arithmetic item (13/15 vs 14/15). At batch size 1 the bf16 GEMVs already run near memory bandwidth; the FP8 path pads a single row into a _scaled_mm tile and adds a quantize launch per call. The branch perf/fp8-shared-expert-hc stays for reference. Assisted-by: Claude Opus 5.5 --- .../models/qwen4_exp/fp8_post_load.py | 94 ------------------ python/freetoken/models/qwen4_exp/model.py | 7 -- tests/models/qwen4_exp/test_fp8_post_load.py | 95 ------------------- 3 files changed, 196 deletions(-) delete mode 100644 python/freetoken/models/qwen4_exp/fp8_post_load.py delete mode 100644 tests/models/qwen4_exp/test_fp8_post_load.py diff --git a/python/freetoken/models/qwen4_exp/fp8_post_load.py b/python/freetoken/models/qwen4_exp/fp8_post_load.py deleted file mode 100644 index dedd6783bf..0000000000 --- a/python/freetoken/models/qwen4_exp/fp8_post_load.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Opt-in post-load FP8 for the qwen4_exp projections the dense FP8 path leaves in bf16. - -At batch size 1 these skinny GEMMs run at ~600 GB/s on the 5080, so halving their weight -bytes halves their time: the shared expert (``FREETOKEN_FP8_SHARED_EXPERT=1``, ~0.4 ms/token) -and the hyper-connection mixes (``FREETOKEN_FP8_HC=1``, ~1 ms/token, two per layer). They are -quantized after ``load_state_dict`` from the loaded bf16 weights, with the same per-tensor -quantizer as the dense path, so HF and FTW checkpoints alike need no layout change. TP=1 only. -The router stays bf16: routing decisions are the precision-sensitive part. -""" - -from __future__ import annotations - -import os - -import torch -from freetoken.distributed import get_tp_info -from freetoken.layers.base import BaseOP -from freetoken.layers.fp8_dynamic import Fp8DynamicLinear -from freetoken.models.config import _ENV_TRUE -from freetoken.utils import init_logger - -logger = init_logger(__name__) - - -def fp8_shared_expert_enabled() -> bool: - return os.getenv("FREETOKEN_FP8_SHARED_EXPERT", "0").strip().lower() in _ENV_TRUE - - -def fp8_hc_enabled() -> bool: - return os.getenv("FREETOKEN_FP8_HC", "0").strip().lower() in _ENV_TRUE - - -def _fp8_from(linear: BaseOP) -> Fp8DynamicLinear: - from .weight import _quantize_per_tensor - - weight = linear.weight - out_features, in_features = weight.shape - op = Fp8DynamicLinear(in_features, out_features) - op.weight, op.weight_scale = _quantize_per_tensor(weight) - return op - - -def _walk(root: BaseOP): - seen: set[int] = set() - stack = [root] - while stack: - op = stack.pop() - if id(op) in seen: - continue - seen.add(id(op)) - yield op - for value in vars(op).values(): - if isinstance(value, BaseOP): - stack.append(value) - elif isinstance(value, (list, tuple)): - stack.extend(v for v in value if isinstance(v, BaseOP)) - - -def quantize_after_load(root: BaseOP) -> int: - """Swap the opted-in bf16 projections under ``root`` for FP8 ones; returns how many.""" - shared, hc = fp8_shared_expert_enabled(), fp8_hc_enabled() - if not (shared or hc): - return 0 - if get_tp_info().size > 1: - logger.warning("FREETOKEN_FP8_SHARED_EXPERT / FREETOKEN_FP8_HC are TP=1 only; ignored") - return 0 - from freetoken.models.qwen3_5_moe.moe import _SharedExpert - - from .hc import GatedResidual - - swapped = 0 - for op in list(_walk(root)): - if shared and isinstance(op, _SharedExpert): - names = ("gate_up_proj", "down_proj") - elif hc and isinstance(op, GatedResidual): - names = ("input_mix_weight_down_block_inject", "input_mix_weight_down", "input_mix_weight_up") - else: - continue - for name in names: - linear = getattr(op, name, None) - if linear is None or getattr(linear, "weight", None) is None or linear.weight.dtype != torch.bfloat16: - continue - setattr(op, name, _fp8_from(linear)) - swapped += 1 - if swapped: - if torch.cuda.is_available(): - torch.cuda.empty_cache() # the dropped bf16 weights are what the cache planner sees - logger.info_rank0( - f"post-load FP8: {swapped} projections (shared expert={shared}, hyper-connections={hc})" - ) - return swapped - - -__all__ = ["fp8_hc_enabled", "fp8_shared_expert_enabled", "quantize_after_load"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index f97df5012f..5d4c6f9d2f 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -152,13 +152,6 @@ def __init__(self, config: ModelConfig) -> None: ) super().__init__() - def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = False) -> None: - super().load_state_dict(state_dict, prefix=prefix, _internal=_internal) - if not _internal: - from .fp8_post_load import quantize_after_load - - quantize_after_load(self) - def load_host_tables(self, engine_config) -> int: """Attach the PLE n-gram table (pinned checkpoint bank, or zeros for dummy weights); returns the pinned host bytes the engine reserves from its pin budget.""" ple_layers = self.model.ple_layers diff --git a/tests/models/qwen4_exp/test_fp8_post_load.py b/tests/models/qwen4_exp/test_fp8_post_load.py deleted file mode 100644 index 09de72b0f2..0000000000 --- a/tests/models/qwen4_exp/test_fp8_post_load.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Opt-in post-load FP8 for the qwen4_exp shared expert and hyper-connection projections.""" - -from __future__ import annotations - -import copy -from types import SimpleNamespace - -import pytest -import torch - -from freetoken.distributed import set_tp_info, try_get_tp_info -from freetoken.layers.base import BaseOP -from freetoken.layers.fp8_dynamic import Fp8DynamicLinear - - -class _Root(BaseOP): - def __init__(self, **ops): - self.__dict__.update(ops) - - -def _config(): - args = SimpleNamespace(hc_count=4, hidden_size=64, hc_lowrank=32, ple_state_width=4 * 64) - return SimpleNamespace(qwen4_args=args, rms_norm_eps=1e-6, quant=None) - - -def _linears(root): - return { - "se.gate_up_proj": root.se.gate_up_proj, "se.down_proj": root.se.down_proj, - "hc.down": root.hc.input_mix_weight_down_block_inject, "hc.up": root.hc.input_mix_weight_up, - "mixer.down": root.mixer.input_mix_weight_down, "mixer.up": root.mixer.input_mix_weight_up, - } - - -def _root(device="cpu"): - if try_get_tp_info() is None: - set_tp_info(rank=0, size=1) - from freetoken.models.qwen3_5_moe.moe import _SharedExpert - from freetoken.models.qwen4_exp.hc import GatedResidual - - root = _Root( - se=_SharedExpert(_config(), 64, 32, prefix="se"), - hc=GatedResidual(_config(), prefix="hc"), - mixer=GatedResidual(_config(), use_combine=False, prefix="mixer"), - ) - g = torch.Generator().manual_seed(0) - for lin in _linears(root).values(): - lin.weight = (torch.randn(lin.weight.shape, generator=g) * 0.05).to(torch.bfloat16).to(device) - root.hc.hc_norm.weight = torch.zeros_like(root.hc.hc_norm.weight, dtype=torch.bfloat16).to(device) - root.mixer.hc_norm.weight = torch.zeros_like(root.mixer.hc_norm.weight, dtype=torch.bfloat16).to(device) - return root - - -def test_flags_off_leaves_everything_bf16(monkeypatch): - from freetoken.models.qwen4_exp.fp8_post_load import quantize_after_load - - monkeypatch.delenv("FREETOKEN_FP8_SHARED_EXPERT", raising=False) - monkeypatch.delenv("FREETOKEN_FP8_HC", raising=False) - root = _root() - assert quantize_after_load(root) == 0 - assert all(lin.weight.dtype == torch.bfloat16 for lin in _linears(root).values()) - - -@pytest.mark.parametrize("shared, hc, expect", [("1", "0", 2), ("0", "1", 4), ("1", "1", 6)]) -def test_swaps_exactly_the_opted_in_projections(monkeypatch, shared, hc, expect): - from freetoken.models.qwen4_exp.fp8_post_load import quantize_after_load - - monkeypatch.setenv("FREETOKEN_FP8_SHARED_EXPERT", shared) - monkeypatch.setenv("FREETOKEN_FP8_HC", hc) - root = _root() - before = {k: v.weight.float().clone() for k, v in _linears(root).items()} - assert quantize_after_load(root) == expect - for name, lin in _linears(root).items(): - want_fp8 = (name.startswith("se.") and shared == "1") or (not name.startswith("se.") and hc == "1") - assert isinstance(lin, Fp8DynamicLinear) == want_fp8, name - if want_fp8: - deq = lin.weight.float() * lin.weight_scale - rel = (deq - before[name]).norm() / before[name].norm() - assert rel < 0.05, (name, float(rel)) - assert root.hc.input_mix_weight_down_block_inject.weight.shape[0] % 16 == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA (_scaled_mm)") -def test_fp8_forward_tracks_the_bf16_forward(monkeypatch): - from freetoken.models.qwen4_exp.fp8_post_load import quantize_after_load - - monkeypatch.setenv("FREETOKEN_FP8_SHARED_EXPERT", "1") - monkeypatch.setenv("FREETOKEN_FP8_HC", "1") - ref = _root("cuda") - fp8 = copy.deepcopy(ref) - quantize_after_load(fp8) - x = torch.randn(3, 64, dtype=torch.bfloat16, device="cuda") - R = torch.randn(3, 256, dtype=torch.bfloat16, device="cuda") - torch.testing.assert_close(fp8.se.forward(x).float(), ref.se.forward(x).float(), rtol=0.1, atol=0.02) - for a, b in zip(fp8.hc.mix(R), ref.hc.mix(R)): - torch.testing.assert_close(a.float(), b.float(), rtol=0.1, atol=0.05) From 43e0cadab2e727821b2bfccb054161769ec272dc Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:18:22 -0400 Subject: [PATCH 20/20] test(engine): give the kv quant config fake the rope table size #573's max_seq_len reads Assisted-by: Claude Opus 5.5 --- tests/engine/test_kv_quant_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py index ba4196be09..1f460166d1 100644 --- a/tests/engine/test_kv_quant_config.py +++ b/tests/engine/test_kv_quant_config.py @@ -39,7 +39,7 @@ def _model_config(kind): has_swa_attention=False, has_linear_attention=False, num_layers=4, - rotary_config=SimpleNamespace(max_position=1024), + rotary_config=SimpleNamespace(max_position=1024, table_positions=1024), ) specs = { "full": (_spec("full", AttnType.FULL),),