From 6b277351c614eb3befda3d0495a351850992645c Mon Sep 17 00:00:00 2001
From: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com>
Date: Mon, 13 Jul 2026 22:48:19 +0000
Subject: [PATCH] fix(scheduler): count active pages across groups

Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com>
---
 .../csrc/fsm/forward_states.h                 | 22 ++++++++++++++++++-
 tokenspeed-scheduler/csrc/scheduler/request.h | 16 ++++++++++++++
 .../csrc/scheduler/scheduler.cpp              |  6 ++++-
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/tokenspeed-scheduler/csrc/fsm/forward_states.h b/tokenspeed-scheduler/csrc/fsm/forward_states.h
index 1cc5386d7..abdd2e548 100644
--- a/tokenspeed-scheduler/csrc/fsm/forward_states.h
+++ b/tokenspeed-scheduler/csrc/fsm/forward_states.h
@@ -175,7 +175,8 @@ struct ForwardState : public BaseState {
     std::int32_t GetReqPoolIndex() const { return req_pool_index_ ? req_pool_index_->slot_ : -1; }
 
     // Flat: group-0 SAMPLE (ids -- and with per-group block sizes, counts -- differ per group);
-    // feeds one-group stats and the radix-era req_to_page fallback, never cross-group accounting.
+    // feeds one-group stats and the radix-era req_to_page fallback, never cross-group accounting
+    // (that is GetOccupiedPagesAllGroups).
     std::vector<std::int32_t> GetOccupiedPages() const {
         if (!block_tables_.empty()) {
             return BlockTablePageIds(block_tables_[0]);  // flat: first-group sample (see above)
@@ -183,6 +184,25 @@ struct ForwardState : public BaseState {
         return GetPageContainer().Pages();  // radix
     }
 
+    // Flat: real pages of EVERY group's block table (null holes excluded — block 0 is the
+    // never-allocated punch placeholder, not occupancy). Block ids are pool-wide and each id
+    // is owned by one group, so concatenation has no cross-group duplicates. This is the
+    // pool-wide occupancy counterpart of the group-0 sample above.
+    std::vector<std::int32_t> GetOccupiedPagesAllGroups() const {
+        if (block_tables_.empty()) {
+            return GetPageContainer().Pages();  // radix
+        }
+        std::vector<std::int32_t> ids;
+        for (const BlockTable& table : block_tables_) {
+            for (CacheBlock* b : table.Blocks()) {
+                if (!b->IsNull()) {
+                    ids.push_back(b->BlockId());
+                }
+            }
+        }
+        return ids;
+    }
+
     std::vector<std::int32_t> GetLocalAllocatorPages() const {
         if (!block_tables_.empty()) {
             return BlockTablePageIds(block_tables_[0]);  // flat: first-group sample (see GetOccupiedPages)
diff --git a/tokenspeed-scheduler/csrc/scheduler/request.h b/tokenspeed-scheduler/csrc/scheduler/request.h
index 2157be683..951b8a0f2 100644
--- a/tokenspeed-scheduler/csrc/scheduler/request.h
+++ b/tokenspeed-scheduler/csrc/scheduler/request.h
@@ -152,6 +152,22 @@ class Request {
             state_);
     }
 
+    // Pool-wide occupancy: every group's real pages on the flat path (radix: same as
+    // GetOccupiedPages). Forward states only — the occupancy stats that need this never
+    // query queued or terminal requests.
+    std::vector<std::int32_t> GetOccupiedPagesAllGroups() const {
+        return std::visit(Overloaded{
+            []<typename T>(const T& s) -> std::vector<std::int32_t>
+                requires std::derived_from<T, fsm::ForwardState>
+            { return s.GetOccupiedPagesAllGroups(); },
+            [this](const auto&) -> std::vector<std::int32_t> {
+                throw std::logic_error("Request::GetOccupiedPagesAllGroups: expected a forward state; got state=" +
+                                       StateName());
+            },
+            },
+            state_);
+    }
+
     // Flat KV-cache: true when the current state carries no per-group block
     // tables (radix path, or a non-forward state). Forward states on the flat
     // path return false once allocation has populated block_tables_.
diff --git a/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp b/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp
index 0f54243b9..14b1b64fe 100644
--- a/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp
+++ b/tokenspeed-scheduler/csrc/scheduler/scheduler.cpp
@@ -260,10 +260,14 @@ std::size_t Scheduler::AvailableKvPages() const {
 }
 
 std::size_t Scheduler::ActiveKvPages() const {
+    // Distinct pages pinned by running requests, in the same units as AvailableKvPages():
+    // flat pool ids across ALL groups (a group-0 sample here understated the ratio that
+    // Python monitoring derives against the whole pool), radix device pages otherwise.
+    // The set dedups pages shared between requests via prefix hits.
     std::unordered_set<std::int32_t> active_pages;
     for (const auto& [_, req] : requests_) {
         if (req->Is<fsm::Prefilling>() || req->Is<fsm::PrefillDone>() || req->Is<fsm::Decoding>()) {
-            for (std::int32_t page : req->GetOccupiedPages()) {
+            for (std::int32_t page : req->GetOccupiedPagesAllGroups()) {
                 active_pages.insert(page);
             }
         }