Skip to content

Add residual top-up extraction to PPR forward push #702

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mkolodner-sc wants to merge 13 commits into
base: main
Choose a base branch
from
Draft

Add residual top-up extraction to PPR forward push #702

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
38be668
Reuse fetched PPR neighbor cache entries
mkolodner-sc Jul 13, 2026
d482f0c
Add residual top-up extraction to PPR forward push
mkolodner-sc Jul 10, 2026
531f290
Move regular PPR residual top-up wiring into sampler
mkolodner-sc Jul 10, 2026
ee6027f
Make PPR residual top-up configurable
mkolodner-sc Jul 13, 2026
8cb1565
Avoid over-extracting residual top-up candidates
mkolodner-sc Jul 13, 2026
a33f080
Remove redundant Python PPR top-up truncation
mkolodner-sc Jul 13, 2026
4a8d631
Sort residual top-up PPR results by score
mkolodner-sc Jul 13, 2026
5ace07f
Calibrate finalized PPR top-up scores
mkolodner-sc Jul 13, 2026
2bcb990
Fix PPR sampler test type narrowing
mkolodner-sc Jul 13, 2026
e06eaf3
Address PPR residual top-up review feedback
mkolodner-sc Jul 13, 2026
ae89d73
Address follow-up PPR review feedback
mkolodner-sc Jul 13, 2026
c21ace6
Simplify PPR top-up extraction API
mkolodner-sc Jul 14, 2026
788f813
Simplify residual top-up selection
mkolodner-sc Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
166 changes: 112 additions & 54 deletions gigl-core/core/sampling/ppr_forward_push.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#include <torch/torch.h>

#include <algorithm>

@kmontemayor2-sc kmontemayor2-sc Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

btw why do we include this?

@mkolodner-sc mkolodner-sc Jul 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

std::min, std::sort, and std::partial_sort require this. These functions worked before because we had the include in the .h file, but it's better for algorithm imports to live in the .cpp file, so we moved it from .h to here

#include <climits>
#include <cstdint>
#include <optional>
Expand Down Expand Up @@ -148,9 +149,10 @@ std::optional<std::unordered_map<int32_t, torch::Tensor>> PPRForwardPush::drainQ

void PPRForwardPush::pushResiduals(
const std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>& fetchedByEtypeId) {
// Step 1: Unpack the input map into a C++ map keyed by packKey(nodeId, edgeTypeId)
// for fast lookup during the residual-push loop below.
std::unordered_map<uint64_t, std::vector<int32_t>> fetched;
// Step 1: Persist fetched neighbor lists in the per-state cache. drainQueue()
// consults this cache before requesting future lookups, so storing every
// fetched row here avoids re-fetching a (node, edge type) pair if it re-enters
// the frontier later in the same PPR channel.
for (const auto& [edgeTypeId, neighborTensors] : fetchedByEtypeId) {
const auto& nodeIdsTensor = std::get<0>(neighborTensors);
const auto& flatNeighborIdsTensor = std::get<1>(neighborTensors);
Expand All @@ -172,7 +174,10 @@ void PPRForwardPush::pushResiduals(
for (int64_t neighborIdx = 0; neighborIdx < count; ++neighborIdx) {
neighborIds[neighborIdx] = static_cast<int32_t>(flatNeighborIdsAccessor[offset + neighborIdx]);
}
fetched[packKey(nodeId, edgeTypeId)] = std::move(neighborIds);
uint64_t cacheKey = packKey(nodeId, edgeTypeId);
if (_neighborCache.find(cacheKey) == _neighborCache.end()) {
_neighborCache.emplace(cacheKey, std::move(neighborIds));
}
offset += count;
}
}
Expand All @@ -197,21 +202,16 @@ void PPRForwardPush::pushResiduals(
srcNodeTypeState.pprScores[sourceNodeId] += sourceResidual;
srcNodeTypeState.residuals[sourceNodeId] = 0.0;

// b. Count total fetched/cached neighbors across all edge types for
// b. Count total cached neighbors across all edge types for
// this source node. We normalise by the number of neighbors we
// actually retrieved, not the true degree, so residual is fully
// distributed among known neighbors rather than leaking to unfetched
// ones (which matters when num_neighbors_per_hop < true_degree).
int32_t totalFetched = 0;
int32_t totalCachedNeighbors = 0;
for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) {
auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId));
if (fetchedEntry != fetched.end()) {
totalFetched += static_cast<int32_t>(fetchedEntry->second.size());
} else {
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
totalFetched += static_cast<int32_t>(cachedEntry->second.size());
}
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
totalCachedNeighbors += static_cast<int32_t>(cachedEntry->second.size());
}
}
// Two cases reach here:
Expand All @@ -221,32 +221,23 @@ void PPRForwardPush::pushResiduals(
// This overstates src and understates its neighbors. This is expected
// behavior when max_fetch_iterations is set, which intentionally trades
// theoretical PPR correctness for better throughput.
if (totalFetched == 0) {
if (totalCachedNeighbors == 0) {
continue;
}

double residualPerNeighbor = (1.0 - _alpha) * sourceResidual / static_cast<double>(totalFetched);
double residualPerNeighbor =
(1.0 - _alpha) * sourceResidual / static_cast<double>(totalCachedNeighbors);

for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) {
// Invariant: fetched and _neighborCache are mutually exclusive for
// any given (node, etype) key within one iteration. drainQueue()
// only requests a fetch for nodes absent from _neighborCache, so a
// key is in at most one of the two.
//
// Neighbor list for this (src, edgeTypeId) pair, borrowed from whichever
// map holds it. reference_wrapper is used because std::optional cannot
// hold a reference directly, and we want to avoid copying the vector —
// the data already exists in fetched or _neighborCache and both outlive
// this loop body. Access via neighborList->get().
// the data already exists in _neighborCache and outlives this loop body.
// Access via neighborList->get().
std::optional<std::reference_wrapper<const std::vector<int32_t>>> neighborList;
auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId));
if (fetchedEntry != fetched.end()) {
neighborList = std::cref(fetchedEntry->second);
} else {
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
neighborList = std::cref(cachedEntry->second);
}
auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId));
if (cachedEntry != _neighborCache.end()) {
neighborList = std::cref(cachedEntry->second);
}
if (!neighborList || neighborList->get().empty()) {
continue;
Expand All @@ -267,18 +258,6 @@ void PPRForwardPush::pushResiduals(
dstNodeTypeState.residuals[neighborNodeId] >= threshold) {
dstNodeTypeState.queue.insert(neighborNodeId);
++_numNodesInQueue;

// Promote neighbor lists to the persistent cache: this node will
// be processed next iteration, so caching avoids a re-fetch.
for (int32_t neighborEdgeTypeId : _nodeTypeToEdgeTypeIds[dstNodeTypeId]) {
uint64_t packedKey = packKey(neighborNodeId, neighborEdgeTypeId);
if (_neighborCache.find(packedKey) == _neighborCache.end()) {
auto fetchedNeighborEntry = fetched.find(packedKey);
if (fetchedNeighborEntry != fetched.end()) {
_neighborCache[packedKey] = fetchedNeighborEntry->second;
}
}
}
}
}
}
Expand All @@ -287,8 +266,10 @@ void PPRForwardPush::pushResiduals(
}
}

std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> PPRForwardPush::extractTopK(
int32_t maxPprNodes) {
std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> PPRForwardPush::
extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp) {
TORCH_CHECK(maxPPRNodes >= 0, "maxPPRNodes must be non-negative, got ", maxPPRNodes, ".");

std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> result;
// Emit an entry for every node type, even if unreachable in this batch (empty tensors,
// all-zero valid_counts). This keeps the output shape consistent across batches so
Expand All @@ -299,21 +280,98 @@ std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tens
std::vector<int64_t> validCounts;

for (int32_t seedIdx = 0; seedIdx < _batchSize; ++seedIdx) {
const auto& scores = _state[seedIdx][nodeTypeId].pprScores;
int32_t topK = std::min(maxPprNodes, static_cast<int32_t>(scores.size()));
const auto& nodeTypeState = _state[seedIdx][nodeTypeId];
const auto& scores = nodeTypeState.pprScores;
const auto higherScore = [](const auto& a, const auto& b) {
return a.second > b.second;
};

const int32_t topK = std::min(maxPPRNodes, static_cast<int32_t>(scores.size()));
const int32_t residualTopUpBudget = enableResidualTopUp ? maxPPRNodes - topK : 0;
std::vector<std::pair<int32_t, double>> selectedPairs;
selectedPairs.reserve(static_cast<size_t>(topK + residualTopUpBudget));
std::unordered_set<int32_t> selectedPPRNodeIds;
if (topK > 0) {
if (residualTopUpBudget > 0) {
selectedPPRNodeIds.reserve(static_cast<size_t>(topK));
}
std::vector<std::pair<int32_t, double>> scorePairs(scores.begin(), scores.end());
std::partial_sort(scorePairs.begin(),
scorePairs.begin() + topK,
scorePairs.end(),
[](const auto& a, const auto& b) { return a.second > b.second; });
// Selection is intentionally two-phase: finalized nodes are selected
// first by raw PPR score, and residual candidates only compete for
// the remaining budget.
if (enableResidualTopUp) {
// The final emitted order is sorted by ppr_score + residual
// after top-up candidates are selected, so this pass only
// needs to partition out the raw-PPR top K.
if (topK < static_cast<int32_t>(scorePairs.size())) {
std::nth_element(scorePairs.begin(),
scorePairs.begin() + topK,
scorePairs.end(),
higherScore);
}
} else {
std::partial_sort(scorePairs.begin(),
scorePairs.begin() + topK,
scorePairs.end(),
higherScore);
}

for (int32_t rankIdx = 0; rankIdx < topK; ++rankIdx) {
Comment on lines +297 to 319

@kmontemayor2-sc kmontemayor2-sc Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Robot review:

  ▎ Selection metric and emitted metric disagree — the result isn't top-k under the score it's sorted by.
  ▎
  ▎ The primary top-k is selected by raw pprScore (the partial_sort on scorePairs at line 310), but each selected node is then emitted with the
  ▎ calibrated score pprScore + residual (lines 318-322) and the combined set is re-sorted by that calibrated score later (lines 331+). So the
  ▎ returned set is top-k by raw ppr, merely re-ordered by calibrated score — it is not the top-k under the emitted metric.
  ▎
  ▎ Concretely, a residual candidate whose calibrated score would outrank a selected finalized node can still be excluded, because it competes
  ▎ only for residualBudget = totalNodeLimit - topK slots rather than against the finalized set. Reproduced against the built extension:
  ▎ topup(2,2,2):  ids=[0,1], weights=[0.5, 0.0625]     # node 3 (calibrated 0.1875) excluded
  ▎ topup(2,2,-1): ids=[0,3,1], weights=[0.5, 0.1875, 0.0625]
  ▎ This matches the "do not displace finalized-PPR nodes" note in ppr_forward_push.h:309-320, so it may be intended. If so, please add a
  ▎ one-line note here that the emitted set is top-k by raw ppr score and only re-sorted by the calibrated score. If displacement is actually
  ▎ desired, build a single candidate pool scored by pprScore + residual and take top-maxTotalNodes instead of the two-phase select.

@mkolodner-sc mkolodner-sc Jul 13, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is intended. Residual top-up is a fill-only mechanism: finalized PPR nodes are selected first by raw PPR score, and residual candidates only compete for the remaining budget. After that, the selected rows are emitted/sorted by ppr_score + residual when top-up is enabled. I added source comments to make explicit that this is a two-phase policy and not a global top-k over the emitted score when maxTotalNodes is tight.

flatIds.push_back(static_cast<int64_t>(scorePairs[rankIdx].first));
flatWeights.push_back(scorePairs[rankIdx].second);
int32_t nodeId = scorePairs[rankIdx].first;
double outputScore = scorePairs[rankIdx].second;
if (enableResidualTopUp) {
auto residualIter = nodeTypeState.residuals.find(nodeId);
if (residualIter != nodeTypeState.residuals.end()) {
outputScore += residualIter->second;
}
}
selectedPairs.emplace_back(nodeId, outputScore);
if (residualTopUpBudget > 0) {
selectedPPRNodeIds.insert(nodeId);
}
}
}
validCounts.push_back(static_cast<int64_t>(topK));

if (residualTopUpBudget > 0) {
std::vector<std::pair<int32_t, double>> residualPairs;
residualPairs.reserve(nodeTypeState.residuals.size());
for (const auto& [nodeId, residual] : nodeTypeState.residuals) {
if (residual <= 0.0 || selectedPPRNodeIds.find(nodeId) != selectedPPRNodeIds.end()) {
continue;
}

auto scoreIter = scores.find(nodeId);
double pprScore = (scoreIter != scores.end()) ? scoreIter->second : 0.0;
double outputScore = pprScore + residual;
residualPairs.emplace_back(nodeId, outputScore);
}

const int32_t residualTopK =
std::min(residualTopUpBudget, static_cast<int32_t>(residualPairs.size()));
if (residualTopK > 0) {
// Residual candidates only need selection here; selected
// finalized and residual rows are sorted together below.
if (residualTopK < static_cast<int32_t>(residualPairs.size())) {
std::nth_element(residualPairs.begin(),
residualPairs.begin() + residualTopK,
residualPairs.end(),
higherScore);
}

for (int32_t rankIdx = 0; rankIdx < residualTopK; ++rankIdx) {
selectedPairs.emplace_back(residualPairs[rankIdx].first, residualPairs[rankIdx].second);
}
}
}

if (enableResidualTopUp && selectedPairs.size() > 1) {
std::sort(selectedPairs.begin(), selectedPairs.end(), higherScore);
}
for (const auto& [nodeId, score] : selectedPairs) {
flatIds.push_back(static_cast<int64_t>(nodeId));
flatWeights.push_back(score);
}
validCounts.push_back(static_cast<int64_t>(selectedPairs.size()));
}

result[nodeTypeId] = {torch::tensor(flatIds, torch::kLong),
Expand Down
Loading