From 935a6f4fb2def99fd70528b1535c95554ef88f2f Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 12:14:15 +0000 Subject: [PATCH 01/12] docs(rfc): plan to move model-topology logic from C++ to Python Scoped v2 extraction: Python computes is_sparse/corr_id, passed via a new set_topology_v2 binding; C++ keeps all memory/IO and the uint32_t TensorID runtime lookups. Momus-approved, design-only pre-implementation RFC. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- docs/rfcs/topology-logic-python-extraction.md | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 docs/rfcs/topology-logic-python-extraction.md diff --git a/docs/rfcs/topology-logic-python-extraction.md b/docs/rfcs/topology-logic-python-extraction.md new file mode 100644 index 0000000..8225f8e --- /dev/null +++ b/docs/rfcs/topology-logic-python-extraction.md @@ -0,0 +1,83 @@ +# Plan: Move non-performance-critical model-topology *logic* out of C++ into Python + +## 1. Goal + +Reduce the amount of **model-structure semantics** computed inside the C++ hot-path module `core/model/model_topology.cpp`, moving it into Python where the topology is already constructed, **without** moving any performance-critical runtime lookup or any memory/IO work. Keep the `uint32_t TensorID -> Node*` runtime path 100% in C++. + +Success = C++ `InitializeTopology` stops re-deriving `is_sparse` and `corr_id` (layer/expert identity) from loop indices; instead it consumes those values computed in Python. All existing behavior (device placement, caching, prefetch, dispatch, generation output) is byte-for-byte unchanged. + +## 2. Findings that constrain scope (from code analysis) + +Verified against the current tree: + +- Python already builds the topology grouping in `moe_infinity/runtime/model_offload.py::get_topology()` (~lines 1121-1245), returning `List[Tuple[str, List[List[TensorID]]]]`, and calls `self.archer_engine.set_topology(topo)` (line ~1271). +- pybind contract: `core/python/py_archer_prefetch.cpp` (~lines 62-66) binds `set_topology` to `std::vector>>>`; `TensorID = uint32_t` (`core/common/types.h:13`). +- `ArcherPrefetchHandle::SetTopology` (`core/prefetch/archer_prefetch_handle.cpp:348-353`) forwards to `ArcherTopologyHandle::InitializeTopology` (`core/model/model_topology.cpp:507-737`). +- `InitializeTopology` interleaves **movable** logic and **non-movable** work: + - MOVABLE (pure integer/label logic): `is_sparse = (num_groups_in_stage > 1)`; `corr_id = (layer_id & 0xFFFFFFFF) | ((expert_id & 0xFFFFFFFF) << 32)`; last-stage nodes get high bits `0xFFFFFFFF`; node ordering; `children_visit_cnt` parent-child sizing for sparse stages. + - NON-MOVABLE (must stay C++): `byte_size` from `kTensorIndex`; GPU round-robin `default_device` placement; `kHostMemoryPool` host allocation + partition-file reads via `kArcherTensorHandle`; `SetDevice(...)` data movement. +- Runtime lookups that MUST stay C++ and MUST remain unchanged: `GetNodeFromTensorID` (`model_topology.cpp:739-760`, ~12 hot-path callers in `archer_prefetch_handle.cpp` + `expert_dispatcher.cpp:268`), `GetNodeBodyFromCorrID` (762-780), `GetSparseNodes/GetDenseNodes/GetLFUNodes/GetSparseCacheLimit/GetNumLayersAndExperts`. + +**Honest scope caveat (for reviewer):** the movable surface is small (~30-60 lines of index math) and is coupled to the memory-allocation loop it lives in. The value is *separation of concerns / single source of truth for model semantics*, NOT performance. If the reviewer judges the risk/reward unfavorable, the correct outcome may be "do not refactor" — see Section 8. + +## 3. Non-goals (explicitly out of scope) + +- Do NOT move `GetNodeFromTensorID`, `GetNodeBodyFromCorrID`, or any `Get*Nodes`/cache function to Python. +- Do NOT move `byte_size`/`kTensorIndex`, device placement, `kHostMemoryPool`, partition reads, or `SetDevice` to Python. +- Do NOT change the on-disk format, `name_id_map.json`, or checkpoint layout. +- Do NOT change the numeric values of `corr_id`/`is_sparse` — only *where* they are computed. + +## 4. Design: enrich the Python->C++ topology contract + +Change the topology element from +`(stage_name, id_groups)` to +`(stage_name, is_sparse, id_groups, corr_ids)` +where `is_sparse: bool`, `id_groups: List[List[TensorID]]` (unchanged), `corr_ids: List[uint64]` (one packed `corr_id` per group, in group order, already including the `0xFFFFFFFF` high-bits marker for last-stage nodes). + +C++ `InitializeTopology` then: +- reads `is_sparse` and `corr_id` from the tuple instead of computing them from loop indices; +- keeps the byte_size / device-placement / host-alloc / SetDevice code paths exactly as-is. + +Backward-compat option (decision for reviewer): add a **new** binding `set_topology_v2` and keep `set_topology` intact, so a stale `name_id_map`/older caller path still works. Default recommendation: add v2, route Python through v2, leave v1 in place unused (lower blast radius, easy rollback). + +## 5. Step-by-step tasks (each independently verifiable) + +Ordering: C++ side first (compiles standalone), then Python switch, then verify. + +- [ ] **T1 (read/confirm):** Read `core/model/model_topology.cpp:507-737` and `model_topology.h:172-175` in full; confirm exact `corr_id` packing and the last-stage `0xFFFFFFFF` rule. Record the exact lines that compute `is_sparse` and `corr_id`. Evidence: quoted current lines in the PR description. +- [ ] **T2 (C++ signature):** Add `InitializeTopologyV2(const std::vector>, std::vector>>&)` in `model_topology.{h,cpp}` that consumes `is_sparse`+`corr_ids` and otherwise reuses the existing body (extract shared code into a helper to avoid duplication). Evidence: file compiles. +- [ ] **T3 (handle passthrough):** Add `ArcherPrefetchHandle::SetTopologyV2(...)` in `archer_prefetch_handle.{h,cpp}` forwarding to `InitializeTopologyV2`. Evidence: compiles. +- [ ] **T4 (pybind):** Bind `set_topology_v2` in `core/python/py_archer_prefetch.cpp`. Evidence: compiles; `import moe_infinity` exposes the method. +- [ ] **T5 (Python producer):** In `model_offload.py::get_topology()`, compute `is_sparse` and per-group `corr_id` (reusing existing name parsing / `parse_expert_id`) and return the enriched tuples; switch the call site (~line 1271) to `set_topology_v2`. Evidence: `python -c "import moe_infinity"` OK; unit assertion that produced `corr_id`/`is_sparse` equal the values C++ v1 would have produced (see T7). +- [ ] **T6 (build):** Rebuild the extension (`pip install --no-build-isolation -e .`). Evidence: exit code 0. **BLOCKED by disk — see Section 7.** +- [ ] **T7 (equivalence test):** Add a temporary parity check (or a small pytest under `tests/python/`) that runs a supported small model (e.g. `deepseek-ai/DeepSeek-V2-Lite-Chat` if weights are available, else a mocked topology) and asserts the resulting node `corr_id`/`is_sparse`/device placement are identical between v1 and v2 paths. Evidence: test passes. +- [ ] **T8 (smoke generation):** Run `examples/deepseek_v2_chat_example.py` (or the smallest available example) and confirm generation output is unchanged vs a pre-refactor run. Evidence: identical decoded output on a fixed prompt+seed. +- [ ] **T9 (diagnostics/cleanup):** `lsp_diagnostics` clean on changed files; remove temporary parity scaffolding; leave `set_topology` (v1) present but unused unless reviewer approves deletion. + +## 6. Verification matrix + +| Step | Command / check | Pass criteria | +|---|---|---| +| C++ edits (T2-T4) | build compiles | exit 0 | +| Python edits (T5) | `python -c "import moe_infinity"` | no error | +| Parity (T7) | pytest parity test | v1 == v2 corr_id/is_sparse/device | +| Smoke (T8) | example generation, fixed seed | output identical to baseline | +| Final (T9) | `lsp_diagnostics` on changed files | no new errors | + +## 7. Risks & BLOCKERS + +- **BLOCKER — disk full.** `/mnt/raid0nvme0` is at 100% (≈2.0 GB free of 14 TB). A from-source rebuild of the CUDA/C++ extension (T6) produces large object files and will very likely fail with `ENOSPC`, and generation (T8) writes an offload dir needing many GB. **T6/T7/T8 cannot be completed until disk space is freed.** This plan should not be marked done on assertion alone; build+run evidence is required. +- **ABI/rebuild risk:** any C++ signature change requires a successful rebuild before Python can use it. Mitigated by adding v2 alongside v1 (no removal) so a partial state still imports. +- **Semantic drift risk:** Python-computed `corr_id` must exactly match the C++ formula, including the last-stage `0xFFFFFFFF` marker and layer/expert index derivation for every supported model family (Mixtral/DeepSeek/Qwen3/GLM/GPT-OSS/NLLB naming differs). Mitigated by the T7 parity test across at least one dense + one sparse stage; ideally gated per model family. +- **Low reward:** net C++ reduction is small; primary benefit is clarity/single-source-of-truth, not speed. + +## 8. Decision gate for reviewer + +Given Section 2's caveat and Section 7's blocker, choose one: +- **(A) Proceed** with the scoped v2 extraction as specified. +- **(B) Narrow further** — only pass `is_sparse` (drop `corr_id` move) to minimize semantic-drift risk. +- **(C) Do not refactor** — conclude the C++ residue is inherently memory-coupled and the movable logic is too thin to justify a new ABI + parity burden. + +## 9. Rollback + +All changes are additive (v2 binding alongside v1). Rollback = revert the Python call site to `set_topology` (v1) and, if desired, remove the v2 symbols. No data/format migration involved, so rollback is a pure code revert with a rebuild. From edb83cb1df4c36905c376d51b3b4b0fd8f31a980 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 12:55:11 +0000 Subject: [PATCH 02/12] refactor(topology): consume is_sparse/corr_id via shared builder Extract InitializeTopology's pipeline construction into BuildTopologyFromSpecs and add InitializeTopologyV2 that reads is_sparse and corr_id from the caller instead of deriving them from loop indices. v1 path preserved. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- core/model/model_topology.cpp | 80 ++++++++++++++++++++++++++--------- core/model/model_topology.h | 12 ++++++ 2 files changed, 71 insertions(+), 21 deletions(-) diff --git a/core/model/model_topology.cpp b/core/model/model_topology.cpp index a84d533..ebab6ac 100644 --- a/core/model/model_topology.cpp +++ b/core/model/model_topology.cpp @@ -504,14 +504,11 @@ bool ArcherTopologyHandle::IsFirstNode(const NodePtr& node) { return false; } -void ArcherTopologyHandle::InitializeTopology( - const std::vector< - std::tuple>>>& - topology) { +void ArcherTopologyHandle::BuildTopologyFromSpecs( + const std::vector& specs) { std::lock_guard lock(mutex_); pipeline_.stages.clear(); std::size_t node_id = 0; - std::size_t layer_id = 0; std::size_t last_sparse_layer_id = UINT64_MAX; size_t num_sparse_layers = 0; @@ -519,12 +516,14 @@ void ArcherTopologyHandle::InitializeTopology( std::vector all_nodes; - for (auto& stage : topology) { - auto& stage_tensors = std::get<1>(stage); - auto stage_ptr = std::make_shared(stage_tensors.size() > 1); + for (std::size_t layer_id = 0; layer_id < specs.size(); ++layer_id) { + const auto& spec = specs[layer_id]; + const auto& stage_tensors = *spec.tensor_groups; + auto stage_ptr = std::make_shared(spec.is_sparse); - std::size_t expert_id = 0; - for (auto& tensor_ids : stage_tensors) { + for (std::size_t expert_id = 0; expert_id < stage_tensors.size(); + ++expert_id) { + const auto& tensor_ids = stage_tensors[expert_id]; auto node_ptr = std::make_shared(); node_ptr->tensor_ids = tensor_ids; int64_t byte_size = 0; @@ -540,8 +539,7 @@ void ArcherTopologyHandle::InitializeTopology( } node_ptr->byte_size = byte_size; node_ptr->id = node_id; - node_ptr->corr_id = - (layer_id & 0xFFFFFFFF) | ((expert_id & 0xFFFFFFFF) << 32); + node_ptr->corr_id = spec.corr_ids[expert_id]; node_ptr->is_sparse = stage_ptr->is_sparse; all_nodes.push_back(node_ptr); @@ -552,11 +550,9 @@ void ArcherTopologyHandle::InitializeTopology( stage_ptr->nodes.push_back(node_body_ptr); node_id++; - expert_id++; } pipeline_.stages.push_back(stage_ptr); auto current_layer_id = layer_id; - layer_id++; if (stage_ptr->is_sparse) { if (UINT64_MAX == last_sparse_layer_id) { @@ -577,13 +573,6 @@ void ArcherTopologyHandle::InitializeTopology( } } - // set last stage nodes corr_id higher 32 bits to be 0xFFFFFFFF - auto last_stage_ptr = pipeline_.stages.back(); - for (auto& node_body : last_stage_ptr->nodes) { - node_body->node->corr_id = - (node_body->node->corr_id & 0xFFFFFFFF) | (UINT64_MAX << 32); - } - // output every tensor id in node for (auto& stage : pipeline_.stages) { for (auto& node : stage->nodes) { @@ -736,6 +725,55 @@ void ArcherTopologyHandle::InitializeTopology( EnableTrace(); } +void ArcherTopologyHandle::InitializeTopology( + const std::vector< + std::tuple>>>& + topology) { + std::vector specs; + specs.reserve(topology.size()); + for (std::size_t layer_id = 0; layer_id < topology.size(); ++layer_id) { + const auto& stage_tensors = std::get<1>(topology[layer_id]); + StageSpec spec; + spec.is_sparse = stage_tensors.size() > 1; + spec.tensor_groups = &stage_tensors; + spec.corr_ids.reserve(stage_tensors.size()); + for (std::size_t expert_id = 0; expert_id < stage_tensors.size(); + ++expert_id) { + spec.corr_ids.push_back((layer_id & 0xFFFFFFFF) | + ((expert_id & 0xFFFFFFFF) << 32)); + } + specs.push_back(std::move(spec)); + } + // Last stage: corr_id high 32 bits = 0xFFFFFFFF end-of-pipeline marker. + if (!specs.empty()) { + for (auto& corr_id : specs.back().corr_ids) { + corr_id = (corr_id & 0xFFFFFFFF) | (UINT64_MAX << 32); + } + } + BuildTopologyFromSpecs(specs); +} + +void ArcherTopologyHandle::InitializeTopologyV2( + const std::vector< + std::tuple>, + std::vector>>& topology) { + std::vector specs; + specs.reserve(topology.size()); + for (const auto& stage : topology) { + StageSpec spec; + spec.is_sparse = std::get<1>(stage); + spec.tensor_groups = &std::get<2>(stage); + spec.corr_ids = std::get<3>(stage); + if (spec.corr_ids.size() != spec.tensor_groups->size()) { + DLOG_ERROR( + "InitializeTopologyV2: corr_ids count {} != tensor group count {}", + spec.corr_ids.size(), spec.tensor_groups->size()); + } + specs.push_back(std::move(spec)); + } + BuildTopologyFromSpecs(specs); +} + NodePtr ArcherTopologyHandle::GetNodeFromTensorID(const TensorID& tensor_id) { std::lock_guard lock(mutex_); diff --git a/core/model/model_topology.h b/core/model/model_topology.h index 27cd1d4..d5f8f5a 100644 --- a/core/model/model_topology.h +++ b/core/model/model_topology.h @@ -174,6 +174,11 @@ class ArcherTopologyHandle : public base::noncopyable { std::tuple>>>& topology); + void InitializeTopologyV2( + const std::vector< + std::tuple>, + std::vector>>& topology); + void EnableTrace() noexcept { trace_enabled_ = true; } void DisableTrace() noexcept { trace_enabled_ = false; } @@ -194,6 +199,13 @@ class ArcherTopologyHandle : public base::noncopyable { } private: + struct StageSpec { + bool is_sparse = false; + const std::vector>* tensor_groups = nullptr; + std::vector corr_ids; + }; + void BuildTopologyFromSpecs(const std::vector& specs); + Pipeline pipeline_; std::unordered_set visited_; std::unordered_map last_active_stage_; From e5943fbe38f098d2a88c689de2591e5c0d445ff0 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 12:55:11 +0000 Subject: [PATCH 03/12] feat(prefetch): expose set_topology_v2 binding Forward the enriched (name, is_sparse, groups, corr_ids) topology from Python to InitializeTopologyV2; set_topology (v1) kept intact. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- core/prefetch/archer_prefetch_handle.cpp | 7 +++++++ core/prefetch/archer_prefetch_handle.h | 4 ++++ core/python/py_archer_prefetch.cpp | 6 ++++++ 3 files changed, 17 insertions(+) diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index 533b288..239640c 100644 --- a/core/prefetch/archer_prefetch_handle.cpp +++ b/core/prefetch/archer_prefetch_handle.cpp @@ -352,6 +352,13 @@ void ArcherPrefetchHandle::SetTopology( kTopologyHandle->InitializeTopology(topology); } +void ArcherPrefetchHandle::SetTopologyV2( + const std::vector< + std::tuple>, + std::vector>>& topology) { + kTopologyHandle->InitializeTopologyV2(topology); +} + bool ArcherPrefetchHandle::IsTensorOffloaded(const std::uint32_t tensor_id) { std::unique_lock lock(mutex_); auto it = kTensorIndex->find(tensor_id); diff --git a/core/prefetch/archer_prefetch_handle.h b/core/prefetch/archer_prefetch_handle.h index d67acfd..e3146e4 100644 --- a/core/prefetch/archer_prefetch_handle.h +++ b/core/prefetch/archer_prefetch_handle.h @@ -44,6 +44,10 @@ class ArcherPrefetchHandle { void SetTopology(const std::vector< std::tuple>>>& topology); + void SetTopologyV2( + const std::vector< + std::tuple>, + std::vector>>& topology); void UpdateTensorMap(std::uint64_t old_ptr, std::uint64_t new_ptr); bool IsTensorIndexInitialized() const; bool IsTensorOnDevice(const torch::Tensor& tensor) const; diff --git a/core/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index 2f45990..6991003 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -64,6 +64,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { const std::vector>>>&)) & ArcherPrefetchHandle::SetTopology) + .def("set_topology_v2", + (void(ArcherPrefetchHandle::*)( + const std::vector>, + std::vector>>&)) & + ArcherPrefetchHandle::SetTopologyV2) .def("update_tensor_map", (void(ArcherPrefetchHandle::*)(std::uint64_t, std::uint64_t)) & ArcherPrefetchHandle::UpdateTensorMap) From 5f245fa8ac99f0ff59edcb341422bb070c9c854e Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 12:55:11 +0000 Subject: [PATCH 04/12] feat(offload): compute topology is_sparse/corr_id in Python get_topology results are lowered to the v2 tuple and sent via set_topology_v2, mirroring the C++ bit-packing exactly. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- moe_infinity/runtime/model_offload.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index 45d2d6b..a53211d 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -1079,6 +1079,24 @@ def get_topology(self, model): topology = list(ret_dict.items()) return topology + def _build_topology_v2(self, topo): + mask32 = 0xFFFFFFFF + num_stages = len(topo) + topo_v2 = [] + for stage_idx, (name, groups) in enumerate(topo): + is_sparse = len(groups) > 1 + is_last_stage = stage_idx == num_stages - 1 + # corr_id mirrors C++ InitializeTopology bit-packing: low 32 bits = + # stage index, high 32 bits = group index, except the last stage + # whose high 32 bits are the 0xFFFFFFFF end-of-pipeline marker. + corr_ids = [ + (stage_idx & mask32) + | ((mask32 if is_last_stage else group_idx & mask32) << 32) + for group_idx in range(len(groups)) + ] + topo_v2.append((name, is_sparse, groups, corr_ids)) + return topo_v2 + def setup_archer_hooks(self, model): for name, param in model.named_parameters(recurse=True): if name not in self.name_id_map: @@ -1103,7 +1121,7 @@ def setup_archer_hooks(self, model): f"TOPO: {len(topo)} stages, {sparse_count} sparse", flush=True, ) - self.archer_engine.set_topology(topo) + self.archer_engine.set_topology_v2(self._build_topology_v2(topo)) print("TOPO: set_topology done", flush=True) @torch.no_grad() From d0ffd65e572d185cb3b79bad3f6aa0393a85a851 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 13:10:01 +0000 Subject: [PATCH 05/12] refactor(topology): drop legacy set_topology, keep single spec API Remove the (name, groups) InitializeTopology/SetTopology path and its binding; the (name, is_sparse, groups, corr_ids) form is now the only set_topology. The signature change spans header+impl+binding so it lands atomically. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- core/model/model_topology.cpp | 30 +----------------------- core/model/model_topology.h | 5 ---- core/prefetch/archer_prefetch_handle.cpp | 9 +------ core/prefetch/archer_prefetch_handle.h | 5 +--- core/python/py_archer_prefetch.cpp | 7 +----- 5 files changed, 4 insertions(+), 52 deletions(-) diff --git a/core/model/model_topology.cpp b/core/model/model_topology.cpp index ebab6ac..7aeefe5 100644 --- a/core/model/model_topology.cpp +++ b/core/model/model_topology.cpp @@ -726,34 +726,6 @@ void ArcherTopologyHandle::BuildTopologyFromSpecs( } void ArcherTopologyHandle::InitializeTopology( - const std::vector< - std::tuple>>>& - topology) { - std::vector specs; - specs.reserve(topology.size()); - for (std::size_t layer_id = 0; layer_id < topology.size(); ++layer_id) { - const auto& stage_tensors = std::get<1>(topology[layer_id]); - StageSpec spec; - spec.is_sparse = stage_tensors.size() > 1; - spec.tensor_groups = &stage_tensors; - spec.corr_ids.reserve(stage_tensors.size()); - for (std::size_t expert_id = 0; expert_id < stage_tensors.size(); - ++expert_id) { - spec.corr_ids.push_back((layer_id & 0xFFFFFFFF) | - ((expert_id & 0xFFFFFFFF) << 32)); - } - specs.push_back(std::move(spec)); - } - // Last stage: corr_id high 32 bits = 0xFFFFFFFF end-of-pipeline marker. - if (!specs.empty()) { - for (auto& corr_id : specs.back().corr_ids) { - corr_id = (corr_id & 0xFFFFFFFF) | (UINT64_MAX << 32); - } - } - BuildTopologyFromSpecs(specs); -} - -void ArcherTopologyHandle::InitializeTopologyV2( const std::vector< std::tuple>, std::vector>>& topology) { @@ -766,7 +738,7 @@ void ArcherTopologyHandle::InitializeTopologyV2( spec.corr_ids = std::get<3>(stage); if (spec.corr_ids.size() != spec.tensor_groups->size()) { DLOG_ERROR( - "InitializeTopologyV2: corr_ids count {} != tensor group count {}", + "InitializeTopology: corr_ids count {} != tensor group count {}", spec.corr_ids.size(), spec.tensor_groups->size()); } specs.push_back(std::move(spec)); diff --git a/core/model/model_topology.h b/core/model/model_topology.h index d5f8f5a..881b4dc 100644 --- a/core/model/model_topology.h +++ b/core/model/model_topology.h @@ -170,11 +170,6 @@ class ArcherTopologyHandle : public base::noncopyable { std::uint64_t GetLastActivateStage(const HashID& hash_id); void InitializeTopology( - const std::vector< - std::tuple>>>& - topology); - - void InitializeTopologyV2( const std::vector< std::tuple>, std::vector>>& topology); diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index 239640c..2e18c4b 100644 --- a/core/prefetch/archer_prefetch_handle.cpp +++ b/core/prefetch/archer_prefetch_handle.cpp @@ -346,17 +346,10 @@ void ArcherPrefetchHandle::TraceRequest(const std::uint64_t request_id, } void ArcherPrefetchHandle::SetTopology( - const std::vector< - std::tuple>>>& - topology) { - kTopologyHandle->InitializeTopology(topology); -} - -void ArcherPrefetchHandle::SetTopologyV2( const std::vector< std::tuple>, std::vector>>& topology) { - kTopologyHandle->InitializeTopologyV2(topology); + kTopologyHandle->InitializeTopology(topology); } bool ArcherPrefetchHandle::IsTensorOffloaded(const std::uint32_t tensor_id) { diff --git a/core/prefetch/archer_prefetch_handle.h b/core/prefetch/archer_prefetch_handle.h index e3146e4..dff7ee8 100644 --- a/core/prefetch/archer_prefetch_handle.h +++ b/core/prefetch/archer_prefetch_handle.h @@ -41,10 +41,7 @@ class ArcherPrefetchHandle { torch::Tensor GetHitRate(); void SetTrace(const torch::Tensor& trace); void TraceRequest(const std::uint64_t request_id, const TensorID tensor_id); - void SetTopology(const std::vector< - std::tuple>>>& - topology); - void SetTopologyV2( + void SetTopology( const std::vector< std::tuple>, std::vector>>& topology); diff --git a/core/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index 6991003..8e20674 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -60,16 +60,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // std::uint32_t)) & // ArcherPrefetchHandle::TraceRequest) .def("set_topology", - (void(ArcherPrefetchHandle::*)( - const std::vector>>>&)) & - ArcherPrefetchHandle::SetTopology) - .def("set_topology_v2", (void(ArcherPrefetchHandle::*)( const std::vector>, std::vector>>&)) & - ArcherPrefetchHandle::SetTopologyV2) + ArcherPrefetchHandle::SetTopology) .def("update_tensor_map", (void(ArcherPrefetchHandle::*)(std::uint64_t, std::uint64_t)) & ArcherPrefetchHandle::UpdateTensorMap) From 84d12e271d2eda1769b3eaa0ef2b84a6f48cf87a Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 13:10:01 +0000 Subject: [PATCH 06/12] refactor(offload): lower topology via pure build_topology_specs Extract is_sparse/corr_id lowering into moe_infinity/utils/topology.build_topology_specs and call the single set_topology; drops the inline _build_topology_v2 method. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- moe_infinity/runtime/model_offload.py | 22 +++------------------- moe_infinity/utils/topology.py | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 19 deletions(-) create mode 100644 moe_infinity/utils/topology.py diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index a53211d..39040b8 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -1079,24 +1079,6 @@ def get_topology(self, model): topology = list(ret_dict.items()) return topology - def _build_topology_v2(self, topo): - mask32 = 0xFFFFFFFF - num_stages = len(topo) - topo_v2 = [] - for stage_idx, (name, groups) in enumerate(topo): - is_sparse = len(groups) > 1 - is_last_stage = stage_idx == num_stages - 1 - # corr_id mirrors C++ InitializeTopology bit-packing: low 32 bits = - # stage index, high 32 bits = group index, except the last stage - # whose high 32 bits are the 0xFFFFFFFF end-of-pipeline marker. - corr_ids = [ - (stage_idx & mask32) - | ((mask32 if is_last_stage else group_idx & mask32) << 32) - for group_idx in range(len(groups)) - ] - topo_v2.append((name, is_sparse, groups, corr_ids)) - return topo_v2 - def setup_archer_hooks(self, model): for name, param in model.named_parameters(recurse=True): if name not in self.name_id_map: @@ -1113,6 +1095,8 @@ def setup_archer_hooks(self, model): self.archer_engine.register(buffer.data, self.name_id_map[name]) self.offload_set.add(buffer.data.data_ptr()) + from moe_infinity.utils.topology import build_topology_specs + topo = self.get_topology(model) sparse_count = sum( 1 for _, t in topo if isinstance(t, list) and len(t) > 1 @@ -1121,7 +1105,7 @@ def setup_archer_hooks(self, model): f"TOPO: {len(topo)} stages, {sparse_count} sparse", flush=True, ) - self.archer_engine.set_topology_v2(self._build_topology_v2(topo)) + self.archer_engine.set_topology(build_topology_specs(topo)) print("TOPO: set_topology done", flush=True) @torch.no_grad() diff --git a/moe_infinity/utils/topology.py b/moe_infinity/utils/topology.py new file mode 100644 index 0000000..33dc587 --- /dev/null +++ b/moe_infinity/utils/topology.py @@ -0,0 +1,24 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +# EfficientMoE Team + +_MASK32 = 0xFFFFFFFF + + +def build_topology_specs(topology): + num_stages = len(topology) + specs = [] + for stage_idx, (name, groups) in enumerate(topology): + is_sparse = len(groups) > 1 + is_last_stage = stage_idx == num_stages - 1 + # corr_id mirrors C++ ArcherTopologyHandle::InitializeTopology: low 32 + # bits = stage index, high 32 bits = group index, except the last stage + # whose high 32 bits are the 0xFFFFFFFF end-of-pipeline marker. + corr_ids = [ + (stage_idx & _MASK32) + | ((_MASK32 if is_last_stage else group_idx & _MASK32) << 32) + for group_idx in range(len(groups)) + ] + specs.append((name, is_sparse, groups, corr_ids)) + return specs From 40f92b1bbc89be1df62685f029310b6af72abc70 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 13:10:01 +0000 Subject: [PATCH 07/12] test(topology): unit-test is_sparse/corr_id bit-packing Covers dense/sparse flag, structural corr_id packing, last-stage 0xFFFFFFFF marker, and equality with an independent reference. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- tests/python/unit/test_topology_specs.py | 62 ++++++++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 tests/python/unit/test_topology_specs.py diff --git a/tests/python/unit/test_topology_specs.py b/tests/python/unit/test_topology_specs.py new file mode 100644 index 0000000..7dd7ba8 --- /dev/null +++ b/tests/python/unit/test_topology_specs.py @@ -0,0 +1,62 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +# EfficientMoE Team + +from moe_infinity.utils.topology import build_topology_specs + +MASK32 = 0xFFFFFFFF + + +def _reference(topology): + n = len(topology) + out = [] + for layer_id, (name, groups) in enumerate(topology): + is_sparse = len(groups) > 1 + corr = [ + (layer_id & MASK32) | ((e & MASK32) << 32) + for e in range(len(groups)) + ] + if layer_id == n - 1: + corr = [(c & MASK32) | (MASK32 << 32) for c in corr] + out.append((name, is_sparse, groups, corr)) + return out + + +def test_dense_single_stage_is_last_and_marked(): + specs = build_topology_specs([("embed", [[0, 1]])]) + _, is_sparse, _, corr = specs[0] + assert is_sparse is False + assert corr == [MASK32 << 32] + + +def test_sparse_flag_by_group_count(): + topo = [("a", [[0]]), ("b", [[1], [2], [3]]), ("c", [[4]])] + specs = build_topology_specs(topo) + assert [s[1] for s in specs] == [False, True, False] + + +def test_middle_sparse_corr_packing(): + topo = [("a", [[0]]), ("b", [[1], [2], [3]]), ("c", [[4]])] + specs = build_topology_specs(topo) + assert specs[1][3] == [1, 1 | (1 << 32), 1 | (2 << 32)] + + +def test_last_stage_high_bits_marker(): + topo = [("a", [[0]]), ("b", [[1], [2]])] + specs = build_topology_specs(topo) + assert all((c >> 32) == MASK32 for c in specs[-1][3]) + + +def test_matches_reference_bitpacking(): + topo = [ + ("embed", [[0]]), + ("l0", [[1], [2], [3]]), + ("l1", [[4], [5], [6]]), + ("head", [[7]]), + ] + assert build_topology_specs(topo) == _reference(topo) + + +def test_empty_topology(): + assert build_topology_specs([]) == [] From a834fa49531b998c2ba0a69ca3f23bdc654df561 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Wed, 5 Aug 2026 14:08:20 +0000 Subject: [PATCH 08/12] style(prefetch): clang-format set_topology binding Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- core/python/py_archer_prefetch.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index 8e20674..d1bb5b9 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -61,9 +61,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // ArcherPrefetchHandle::TraceRequest) .def("set_topology", (void(ArcherPrefetchHandle::*)( - const std::vector>, - std::vector>>&)) & + const std::vector>, + std::vector>>&)) & ArcherPrefetchHandle::SetTopology) .def("update_tensor_map", (void(ArcherPrefetchHandle::*)(std::uint64_t, std::uint64_t)) & From c7abb9ce90765e9cb361398214270ef8d9628076 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sun, 9 Aug 2026 11:15:39 +0000 Subject: [PATCH 09/12] test(topology): verify legacy and Python-spec parity Signed-off-by: drunkcoding --- core/model/model_topology.cpp | 43 ++++++++- core/model/model_topology.h | 7 ++ core/prefetch/archer_prefetch_handle.cpp | 16 +++- core/prefetch/archer_prefetch_handle.h | 5 ++ core/python/py_archer_prefetch.cpp | 15 +++- moe_infinity/runtime/model_offload.py | 2 +- .../test_topology_runtime_parity.py | 88 +++++++++++++++++++ 7 files changed, 168 insertions(+), 8 deletions(-) create mode 100644 tests/python/integration/test_topology_runtime_parity.py diff --git a/core/model/model_topology.cpp b/core/model/model_topology.cpp index 7aeefe5..738ab13 100644 --- a/core/model/model_topology.cpp +++ b/core/model/model_topology.cpp @@ -726,6 +726,33 @@ void ArcherTopologyHandle::BuildTopologyFromSpecs( } void ArcherTopologyHandle::InitializeTopology( + const std::vector< + std::tuple>>>& + topology) { + std::vector specs; + specs.reserve(topology.size()); + for (std::size_t layer_id = 0; layer_id < topology.size(); ++layer_id) { + const auto& stage_tensors = std::get<1>(topology[layer_id]); + StageSpec spec; + spec.is_sparse = stage_tensors.size() > 1; + spec.tensor_groups = &stage_tensors; + spec.corr_ids.reserve(stage_tensors.size()); + for (std::size_t expert_id = 0; expert_id < stage_tensors.size(); + ++expert_id) { + spec.corr_ids.push_back((layer_id & 0xFFFFFFFF) | + ((expert_id & 0xFFFFFFFF) << 32)); + } + specs.push_back(std::move(spec)); + } + if (!specs.empty()) { + for (auto& corr_id : specs.back().corr_ids) { + corr_id = (corr_id & 0xFFFFFFFF) | (UINT64_MAX << 32); + } + } + BuildTopologyFromSpecs(specs); +} + +void ArcherTopologyHandle::InitializeTopologyV2( const std::vector< std::tuple>, std::vector>>& topology) { @@ -738,7 +765,7 @@ void ArcherTopologyHandle::InitializeTopology( spec.corr_ids = std::get<3>(stage); if (spec.corr_ids.size() != spec.tensor_groups->size()) { DLOG_ERROR( - "InitializeTopology: corr_ids count {} != tensor group count {}", + "InitializeTopologyV2: corr_ids count {} != tensor group count {}", spec.corr_ids.size(), spec.tensor_groups->size()); } specs.push_back(std::move(spec)); @@ -746,6 +773,20 @@ void ArcherTopologyHandle::InitializeTopology( BuildTopologyFromSpecs(specs); } +std::vector> +ArcherTopologyHandle::GetTopologySnapshot() { + std::lock_guard lock(mutex_); + std::vector> snapshot; + for (const auto& stage : pipeline_.stages) { + for (const auto& node_body : stage->nodes) { + const auto& node = node_body->node; + snapshot.emplace_back(static_cast(node->corr_id), + node->is_sparse, node->default_device.index()); + } + } + return snapshot; +} + NodePtr ArcherTopologyHandle::GetNodeFromTensorID(const TensorID& tensor_id) { std::lock_guard lock(mutex_); diff --git a/core/model/model_topology.h b/core/model/model_topology.h index 881b4dc..beaee25 100644 --- a/core/model/model_topology.h +++ b/core/model/model_topology.h @@ -170,10 +170,17 @@ class ArcherTopologyHandle : public base::noncopyable { std::uint64_t GetLastActivateStage(const HashID& hash_id); void InitializeTopology( + const std::vector< + std::tuple>>>& + topology); + + void InitializeTopologyV2( const std::vector< std::tuple>, std::vector>>& topology); + std::vector> GetTopologySnapshot(); + void EnableTrace() noexcept { trace_enabled_ = true; } void DisableTrace() noexcept { trace_enabled_ = false; } diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index 2e18c4b..edae130 100644 --- a/core/prefetch/archer_prefetch_handle.cpp +++ b/core/prefetch/archer_prefetch_handle.cpp @@ -347,11 +347,23 @@ void ArcherPrefetchHandle::TraceRequest(const std::uint64_t request_id, void ArcherPrefetchHandle::SetTopology( const std::vector< - std::tuple>, - std::vector>>& topology) { + std::tuple>>>& + topology) { kTopologyHandle->InitializeTopology(topology); } +void ArcherPrefetchHandle::SetTopologyV2( + const std::vector< + std::tuple>, + std::vector>>& topology) { + kTopologyHandle->InitializeTopologyV2(topology); +} + +std::vector> +ArcherPrefetchHandle::GetTopologySnapshot() { + return kTopologyHandle->GetTopologySnapshot(); +} + bool ArcherPrefetchHandle::IsTensorOffloaded(const std::uint32_t tensor_id) { std::unique_lock lock(mutex_); auto it = kTensorIndex->find(tensor_id); diff --git a/core/prefetch/archer_prefetch_handle.h b/core/prefetch/archer_prefetch_handle.h index dff7ee8..55a86ce 100644 --- a/core/prefetch/archer_prefetch_handle.h +++ b/core/prefetch/archer_prefetch_handle.h @@ -42,9 +42,14 @@ class ArcherPrefetchHandle { void SetTrace(const torch::Tensor& trace); void TraceRequest(const std::uint64_t request_id, const TensorID tensor_id); void SetTopology( + const std::vector< + std::tuple>>>& + topology); + void SetTopologyV2( const std::vector< std::tuple>, std::vector>>& topology); + std::vector> GetTopologySnapshot(); void UpdateTensorMap(std::uint64_t old_ptr, std::uint64_t new_ptr); bool IsTensorIndexInitialized() const; bool IsTensorOnDevice(const torch::Tensor& tensor) const; diff --git a/core/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index d1bb5b9..b3f05c8 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -60,11 +60,18 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // std::uint32_t)) & // ArcherPrefetchHandle::TraceRequest) .def("set_topology", + (void(ArcherPrefetchHandle::*)( + const std::vector>>>&)) & + ArcherPrefetchHandle::SetTopology) + .def("set_topology_v2", (void(ArcherPrefetchHandle::*)( - const std::vector>, - std::vector>>&)) & - ArcherPrefetchHandle::SetTopology) + const std::vector>, + std::vector>>&)) & + ArcherPrefetchHandle::SetTopologyV2) + .def("get_topology_snapshot", + &ArcherPrefetchHandle::GetTopologySnapshot) .def("update_tensor_map", (void(ArcherPrefetchHandle::*)(std::uint64_t, std::uint64_t)) & ArcherPrefetchHandle::UpdateTensorMap) diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index d6e1a23..e65b8db 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -1116,7 +1116,7 @@ def setup_archer_hooks(self, model): f"TOPO: {len(topo)} stages, {sparse_count} sparse", flush=True, ) - self.archer_engine.set_topology(build_topology_specs(topo)) + self.archer_engine.set_topology_v2(build_topology_specs(topo)) print("TOPO: set_topology done", flush=True) @torch.no_grad() diff --git a/tests/python/integration/test_topology_runtime_parity.py b/tests/python/integration/test_topology_runtime_parity.py new file mode 100644 index 0000000..618cb7e --- /dev/null +++ b/tests/python/integration/test_topology_runtime_parity.py @@ -0,0 +1,88 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +import moe_infinity._store as store +from moe_infinity.utils.topology import build_topology_specs + +MASK32 = 0xFFFFFFFF +MODEL = "deepseek-ai/DeepSeek-V2-Lite-Chat" +NUM_HIDDEN_LAYERS = 27 +FIRST_K_DENSE_REPLACE = 1 +NUM_ROUTED_EXPERTS = 64 + + +def _deepseek_v2_lite_topology(): + """Contract-level topology shaped from DeepSeek-V2-Lite config values.""" + next_tensor_id = 0 + topology = [("model.embed_tokens", [[next_tensor_id]])] + next_tensor_id += 1 + + for layer_id in range(NUM_HIDDEN_LAYERS): + topology.append((f"model.layers.{layer_id}", [[next_tensor_id]])) + next_tensor_id += 1 + if layer_id >= FIRST_K_DENSE_REPLACE: + groups = [ + [next_tensor_id + expert_id] + for expert_id in range(NUM_ROUTED_EXPERTS) + ] + topology.append((f"model.layers.{layer_id}.mlp.experts", groups)) + next_tensor_id += NUM_ROUTED_EXPERTS + + topology.append(("model.norm", [[next_tensor_id]])) + next_tensor_id += 1 + topology.append(("lm_head", [[next_tensor_id]])) + return topology, next_tensor_id + 1 + + +def _snapshot(tmp_path, api_name): + topology, tensor_count = _deepseek_v2_lite_topology() + store_path = tmp_path / api_name + store_path.mkdir() + handle = store.prefetch_handle(str(store_path), 0.01) + try: + for tensor_id in range(tensor_count): + tensor = torch.tensor([tensor_id], dtype=torch.float32) + handle.offload(tensor, tensor_id) + if api_name == "v1": + handle.set_topology(topology) + else: + handle.set_topology_v2(build_topology_specs(topology)) + return [tuple(item) for item in handle.get_topology_snapshot()] + finally: + handle.clean_up_resources() + + +@torch.no_grad() +def test_deepseek_v2_lite_v1_v2_topology_metadata_and_placement_match( + tmp_path, +): + assert torch.cuda.device_count() == 6 + topology, _ = _deepseek_v2_lite_topology() + + legacy = _snapshot(tmp_path, "v1") + enriched = _snapshot(tmp_path, "v2") + + assert enriched == legacy + + first_sparse_stage = 3 + first_sparse_node = sum(len(groups) for _, groups in topology[:3]) + assert legacy[first_sparse_node] == (first_sparse_stage, True, 3) + assert legacy[first_sparse_node + 63] == ( + first_sparse_stage | (63 << 32), + True, + 0, + ) + + last_stage = len(topology) - 1 + assert last_stage == 55 + assert legacy[-1][0] == last_stage | (MASK32 << 32) + assert legacy[-1][1] is False + assert legacy[-1][2] == 5 + + first_six_sparse_devices = [ + item[2] + for item in legacy[first_sparse_node : first_sparse_node + 6] + ] + assert first_six_sparse_devices == [3, 4, 5, 0, 1, 2] From 11277ad5b1cdca95fb5583d4c794ba66f4a89c29 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sun, 9 Aug 2026 11:15:59 +0000 Subject: [PATCH 10/12] test(topology): add deterministic generation probe Signed-off-by: drunkcoding --- .../integration/topology_generation_probe.py | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 tests/python/integration/topology_generation_probe.py diff --git a/tests/python/integration/topology_generation_probe.py b/tests/python/integration/topology_generation_probe.py new file mode 100644 index 0000000..c296425 --- /dev/null +++ b/tests/python/integration/topology_generation_probe.py @@ -0,0 +1,69 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import json +import random +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoTokenizer + +from moe_infinity import MoE + +CHECKPOINT = "deepseek-ai/DeepSeek-V2-Lite-Chat" +PROMPT = "What is 2+3? Answer with one short sentence." +SEED = 20260809 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--offload-dir", required=True) + parser.add_argument("--output", required=True) + args = parser.parse_args() + + random.seed(SEED) + np.random.seed(SEED) + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + torch.use_deterministic_algorithms(True, warn_only=True) + + tokenizer = AutoTokenizer.from_pretrained( + CHECKPOINT, trust_remote_code=True + ) + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": PROMPT}], + tokenize=False, + add_generation_prompt=True, + ) + input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda:0") + model = MoE( + CHECKPOINT, + {"offload_path": args.offload_dir, "device_memory_ratio": 0.5}, + ) + + with torch.no_grad(): + output_ids = model.generate( + input_ids, + max_new_tokens=16, + do_sample=False, + pad_token_id=tokenizer.eos_token_id, + ) + + payload = { + "checkpoint": CHECKPOINT, + "prompt": PROMPT, + "seed": SEED, + "input_ids": input_ids[0].cpu().tolist(), + "output_ids": output_ids[0].cpu().tolist(), + "decoded": tokenizer.decode(output_ids[0], skip_special_tokens=True), + } + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + print(output) + + +if __name__ == "__main__": + main() From e41bf9d8efb64427cdae0fce027b920bbfea104c Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sun, 9 Aug 2026 11:16:27 +0000 Subject: [PATCH 11/12] docs(topology): record compatibility decision Signed-off-by: drunkcoding --- docs/rfcs/topology-logic-python-extraction.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/rfcs/topology-logic-python-extraction.md b/docs/rfcs/topology-logic-python-extraction.md index 8225f8e..77e2229 100644 --- a/docs/rfcs/topology-logic-python-extraction.md +++ b/docs/rfcs/topology-logic-python-extraction.md @@ -40,6 +40,15 @@ C++ `InitializeTopology` then: Backward-compat option (decision for reviewer): add a **new** binding `set_topology_v2` and keep `set_topology` intact, so a stale `name_id_map`/older caller path still works. Default recommendation: add v2, route Python through v2, leave v1 in place unused (lower blast radius, easy rollback). +### Final compatibility decision for PR #133 + +PR #133 retains the original two-field `set_topology` binding as an unused, +tested rollback path and exposes the enriched four-field contract as +`set_topology_v2`. Production Python calls `set_topology_v2`. Both entry points +lower to the same C++ `BuildTopologyFromSpecs` implementation. This deliberately +reverses commit `d0ffd65`'s unreviewed removal of v1 and follows T9's lower-risk +compatibility recommendation. + ## 5. Step-by-step tasks (each independently verifiable) Ordering: C++ side first (compiles standalone), then Python switch, then verify. From f061297e149d00d98dbfd8a290b0ec45d77ae897 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sun, 9 Aug 2026 11:16:50 +0000 Subject: [PATCH 12/12] style(topology): apply repository formatting Signed-off-by: drunkcoding --- core/prefetch/archer_prefetch_handle.cpp | 2 +- core/prefetch/archer_prefetch_handle.h | 7 +++---- core/python/py_archer_prefetch.cpp | 17 ++++++++--------- .../integration/test_topology_runtime_parity.py | 3 +-- 4 files changed, 13 insertions(+), 16 deletions(-) diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index edae130..35acadc 100644 --- a/core/prefetch/archer_prefetch_handle.cpp +++ b/core/prefetch/archer_prefetch_handle.cpp @@ -355,7 +355,7 @@ void ArcherPrefetchHandle::SetTopology( void ArcherPrefetchHandle::SetTopologyV2( const std::vector< std::tuple>, - std::vector>>& topology) { + std::vector>>& topology) { kTopologyHandle->InitializeTopologyV2(topology); } diff --git a/core/prefetch/archer_prefetch_handle.h b/core/prefetch/archer_prefetch_handle.h index 55a86ce..b6ecf41 100644 --- a/core/prefetch/archer_prefetch_handle.h +++ b/core/prefetch/archer_prefetch_handle.h @@ -41,10 +41,9 @@ class ArcherPrefetchHandle { torch::Tensor GetHitRate(); void SetTrace(const torch::Tensor& trace); void TraceRequest(const std::uint64_t request_id, const TensorID tensor_id); - void SetTopology( - const std::vector< - std::tuple>>>& - topology); + void SetTopology(const std::vector< + std::tuple>>>& + topology); void SetTopologyV2( const std::vector< std::tuple>, diff --git a/core/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index b3f05c8..108de15 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -60,18 +60,17 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // std::uint32_t)) & // ArcherPrefetchHandle::TraceRequest) .def("set_topology", - (void(ArcherPrefetchHandle::*)( - const std::vector>>>&)) & - ArcherPrefetchHandle::SetTopology) - .def("set_topology_v2", (void(ArcherPrefetchHandle::*)( const std::vector>, - std::vector>>&)) & + std::string, std::vector>>>&)) & + ArcherPrefetchHandle::SetTopology) + .def("set_topology_v2", + (void(ArcherPrefetchHandle::*)( + const std::vector>, + std::vector>>&)) & ArcherPrefetchHandle::SetTopologyV2) - .def("get_topology_snapshot", - &ArcherPrefetchHandle::GetTopologySnapshot) + .def("get_topology_snapshot", &ArcherPrefetchHandle::GetTopologySnapshot) .def("update_tensor_map", (void(ArcherPrefetchHandle::*)(std::uint64_t, std::uint64_t)) & ArcherPrefetchHandle::UpdateTensorMap) diff --git a/tests/python/integration/test_topology_runtime_parity.py b/tests/python/integration/test_topology_runtime_parity.py index 618cb7e..2a3b098 100644 --- a/tests/python/integration/test_topology_runtime_parity.py +++ b/tests/python/integration/test_topology_runtime_parity.py @@ -82,7 +82,6 @@ def test_deepseek_v2_lite_v1_v2_topology_metadata_and_placement_match( assert legacy[-1][2] == 5 first_six_sparse_devices = [ - item[2] - for item in legacy[first_sparse_node : first_sparse_node + 6] + item[2] for item in legacy[first_sparse_node : first_sparse_node + 6] ] assert first_six_sparse_devices == [3, 4, 5, 0, 1, 2]