Add residual top-up extraction to PPR forward push by mkolodner-sc · Pull Request #702 · Snapchat/GiGL

mkolodner-sc · 2026-07-10T19:54:18Z

Summary

Reasoning for Change

Previously, the PPR forward push would return all neighbors re-enqueued during forward push. This is how the base PPR forward push algorithm works, but results in small sequences volumes that could be more meaningful if they were larger. Decreasing epsilon previously provided a knob to get these larger sequence lengths, but came at heavy throughput costs. Instead of relying on epsilon for this behavior, this change exposes a way to allow us to "top-off" the existing PPR sequence with nodes that were found during forward-push but didn't meet the threshold for re-enqueue. The logic here is that these are the nodes that would have been re-enqueued next by decreasing epsilon, so this gives us a clean way to get these nodes and their ppr scores without having to pay heavy throughput costs.

On internal benchmarking, we have seen sequences toped off this way perform as good as the sequences with lower epsilon at a much faster runtime.

Changes Summary

Adds an configurable but opt-in residual top-up extraction API to PPRForwardPush.

The existing extractTopK behavior is unchanged. The new API returns the normal top-k finalized PPR nodes first, then appends additional candidates from residual mass that was discovered during forward push but did not necessarily cross the requeue threshold. This gives downstream samplers a way to fill more sequence slots without issuing extra neighbor fetches.

Changes

Adds extractTopKWithResidualTopUp(maxPprNodes, maxResidualNodes) in C++.
Exposes the new method through pybind and the Python stub.
Scores residual top-up candidates on the same scale as finalized candidates via ppr_score + residual.

Validation

Ran focused ruff check on the updated Python stub.
Full uv run validation in the temp worktree was blocked by local environment setup because gigl-core attempted to rebuild and could not find Torch CMake config.

mkolodner-sc · 2026-07-13T19:37:36Z

/unit_test

github-actions · 2026-07-13T19:37:49Z

GiGL Automation

@ 19:37:48UTC : 🔄 C++ Unit Test started.

@ 19:39:40UTC : ✅ Workflow completed successfully.

github-actions · 2026-07-13T19:37:49Z

GiGL Automation

@ 19:37:49UTC : 🔄 Python Unit Test started.

@ 19:42:19UTC : ❌ Workflow failed.
Please check the logs for more details.

github-actions · 2026-07-13T19:38:02Z

GiGL Automation

@ 19:38:02UTC : 🔄 Scala Unit Test started.

@ 19:49:31UTC : ✅ Workflow completed successfully.

mkolodner-sc · 2026-07-13T19:58:18Z

/unit_test

github-actions · 2026-07-13T19:58:29Z

GiGL Automation

@ 19:58:29UTC : 🔄 Scala Unit Test started.

@ 20:06:43UTC : ✅ Workflow completed successfully.

github-actions · 2026-07-13T19:58:29Z

GiGL Automation

@ 19:58:29UTC : 🔄 C++ Unit Test started.

@ 20:00:20UTC : ✅ Workflow completed successfully.

github-actions · 2026-07-13T19:58:30Z

GiGL Automation

@ 19:58:30UTC : 🔄 Python Unit Test started.

@ 21:08:08UTC : ✅ Workflow completed successfully.

kmontemayor2-sc

Thanks Matt! In generally I think we have a good amount of algorithmic changes here and it'd be nice if we could comment what / why the new code is doing.

kmontemayor2-sc · 2026-07-13T19:48:15Z


-void PPRForwardPush::pushResiduals(
-    const std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>& fetchedByEtypeId) {
+void PPRForwardPush::pushResiduals(const NeighborFetchMap& fetchedByEtypeId) {


btw I'd prefer if these changes could be made in a separate PR :) makes it easier to distinguish the business logic changes vs refactors

Makes sense, I reverted the type-alias/refactor-only changes from this PR and left a TODO to move the repeated tensor tuple/map types into a follow-up refactor-only PR, so this PR stays focused on the residual top-up logic

kmontemayor2-sc · 2026-07-13T19:49:10Z

-    int32_t maxPprNodes) {
-    std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> result;
+PPRExtractResult PPRForwardPush::extractTopK(int32_t maxPprNodes) {
+    return extractTopKWithResidualTopUp(maxPprNodes, /*maxResidualNodes=*/0);


Why do we always call with 0 here?

We've made the residual top-off a configurable knob (which defaults to True) and this is the previous path, so 0 means “do not include residual top-up nodes” and preserves the previous finalized-PPR-only behavior. The nonzero residual budget is only used by extractTopKWithResidualTopUp; this wrapper delegates to the shared implementation to avoid duplicating extraction code.

I've gone ahead and removed this since we can just use extractTopKWithResidualTopUp in both cases later on, thanks!

kmontemayor2-sc · 2026-07-13T19:55:48Z

+            std::vector<std::pair<int32_t, double>> selectedPairs;
+            selectedPairs.reserve(static_cast<size_t>(topK + residualBudget));


any reason we do it this way vs just give it the right size at init?

I found https://stackoverflow.com/questions/8928547/vector-initialization-or-reserve which implies it doesn't really matter but it does read "cleaner" to me to do the sizing at init?

topK + residualBudget is only an upper bound, not the final size. We always add topK finalized PPR nodes, but we may find fewer residual candidates than residualBudget. If we initialized the vector with that size, it would contain placeholder pairs and we’d need extra indexing/trimming logic. So reserve gives us the allocation benefit without changing the logical size.

kmontemayor2-sc · 2026-07-13T20:06:35Z

+            max_total_nodes_for_extraction = (
+                max_total_nodes if max_total_nodes is not None else -1
+            )
+            extracted_results = ppr_state.extract_top_k_with_residual_top_up(


qq why not always call this? it still works with the other path right?

kmontemayor2-sc · 2026-07-13T20:17:09Z

+        sampler = object.__new__(DistPPRNeighborSampler)
+        sampler._alpha = _TEST_ALPHA
+        sampler._requeue_threshold_factor = _TEST_ALPHA * _TEST_EPS
+        sampler._max_ppr_nodes = max_ppr_nodes
+        sampler._enable_residual_topup = enable_residual_topup
+        sampler._max_fetch_iterations = None
+        sampler._is_homogeneous = True
+        sampler._node_type_to_id = {DEFAULT_HOMOGENEOUS_NODE_TYPE: 0}
+        sampler._ntype_id_to_ntype = [DEFAULT_HOMOGENEOUS_NODE_TYPE]
+        sampler._node_type_id_to_edge_type_ids = [[]]
+        sampler._edge_type_id_to_dst_ntype_id = []
+        sampler._degree_tensors_for_cpp = [torch.zeros(0, dtype=torch.int32)]
+        return sampler


Can we avoid monkey patching like this?

pulled the C++ extraction-method selection into a small helper and updated the tests to call it directly with fake PPR states, so we no longer need object.new or manually assigned private sampler fields.

Hmm, the tests still seem very mock heavy - what are they testing here?

It seems like we coverage of the new behavior in the cpp tests right? What are the python tests adding?

Hmm, that's fair. I'll remove these tests since they aren't adding much over the current cpp coverage

kmontemayor2-sc · 2026-07-13T20:18:03Z

+                }
                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; });

                for (int32_t rankIdx = 0; rankIdx < topK; ++rankIdx) {


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.

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.

kmontemayor2-sc · 2026-07-13T20:18:42Z

+        residual_topup_nodes = (
+            self._max_ppr_nodes if self._enable_residual_topup else 0
+        )
+        return self._extract_ppr_state_top_k(
+            ppr_state,
+            device,
+            max_ppr_nodes=self._max_ppr_nodes,
+            residual_topup_nodes=residual_topup_nodes,
+            max_total_nodes=self._max_ppr_nodes,
+        )


Robot review:

▎ Because the regular path passes residual_topup_nodes = max_ppr_nodes and max_total_nodes = max_ppr_nodes, the C++ side computes ▎ residualBudget = max_ppr_nodes - topK. Whenever a seed already has >= max_ppr_nodes finalized PPR scores, residualBudget == 0 and residual ▎ top-up never fires — which is precisely the well-explored/hub-seed case in the truncated (max_fetch_iterations) regime that motivated top-up, ▎ where large sub-threshold residuals are most likely to exist. Worth confirming this is the intended envelope (see Comment A), and ▎ documenting it.

Yep, that is the intended envelope for the regular sampler right now. We cap total output at max_ppr_nodes, so residual candidates only fill slots when finalized PPR produces fewer than that cap; they do not displace finalized PPR nodes. I added documentation at the sampler call site and in the constructor docs to make this clearer.

mkolodner-sc · 2026-07-13T21:08:22Z

Thanks Matt! In generally I think we have a good amount of algorithmic changes here and it'd be nice if we could comment what / why the new code is doing.

@kmontemayor2-sc Done, added more context in the PR description and in the PR for the reasoning/change made here

kmontemayor2-sc · 2026-07-13T23:09:46Z


 #include <torch/torch.h>

+#include <algorithm>


btw why do we include this?

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

kmontemayor2-sc · 2026-07-13T23:25:24Z

+            device,
+            max_ppr_nodes=self._max_ppr_nodes,
+            residual_topup_nodes=residual_topup_nodes,
+            max_total_nodes=self._max_ppr_nodes,


why do we add max_total_nodes if it's always max_ppr_nodes?

kmontemayor2-sc · 2026-07-13T23:26:20Z

@@ -44,6 +44,12 @@ class PPRSamplerOptions:

    For homogeneous graphs these live directly on ``data.edge_index`` / ``data.edge_attr``.



Can we try to make this comment more "why" and not "what" e.g. "Enable residual top-up to get longer sequence lengths ... "

kmontemayor2-sc · 2026-07-13T23:28:10Z

+        sampler = object.__new__(DistPPRNeighborSampler)
+        sampler._alpha = _TEST_ALPHA
+        sampler._requeue_threshold_factor = _TEST_ALPHA * _TEST_EPS
+        sampler._max_ppr_nodes = max_ppr_nodes
+        sampler._enable_residual_topup = enable_residual_topup
+        sampler._max_fetch_iterations = None
+        sampler._is_homogeneous = True
+        sampler._node_type_to_id = {DEFAULT_HOMOGENEOUS_NODE_TYPE: 0}
+        sampler._ntype_id_to_ntype = [DEFAULT_HOMOGENEOUS_NODE_TYPE]
+        sampler._node_type_id_to_edge_type_ids = [[]]
+        sampler._edge_type_id_to_dst_ntype_id = []
+        sampler._degree_tensors_for_cpp = [torch.zeros(0, dtype=torch.int32)]
+        return sampler


Hmm, the tests still seem very mock heavy - what are they testing here?

It seems like we coverage of the new behavior in the cpp tests right? What are the python tests adding?

kmontemayor2-sc · 2026-07-14T18:01:08Z

-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, int32_t maxResidualNodes, int32_t maxTotalNodes) {


We never pass in maxTotalNodes via python right?

I guess what I was getting at earlier is why do we have this API separate from maxNodes: int, enableTopUp: bool?

Since we do just pass in maxPPRNodes for both if top up is enabled?

mkolodner-sc requested review from kmontemayor2-sc, nshah-sc, svij-sc, xgao4-sc, yliu2-sc and zfan3-sc as code owners July 10, 2026 19:54

mkolodner-sc marked this pull request as draft July 10, 2026 19:56

mkolodner-sc mentioned this pull request Jul 10, 2026

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

Closed

mkolodner-sc added 6 commits July 13, 2026 18:29

Reuse fetched PPR neighbor cache entries

38be668

Add residual top-up extraction to PPR forward push

d482f0c

Move regular PPR residual top-up wiring into sampler

531f290

Make PPR residual top-up configurable

ee6027f

Avoid over-extracting residual top-up candidates

8cb1565

Remove redundant Python PPR top-up truncation

a33f080

mkolodner-sc force-pushed the mkolodner/ppr-residual-topup branch from 98d95fa to a33f080 Compare July 13, 2026 18:35

mkolodner-sc changed the base branch from main to mkolodner/ppr-neighbor-cache-reuse July 13, 2026 18:36

mkolodner-sc mentioned this pull request Jul 13, 2026

Reuse fetched PPR neighbor cache entries #703

Merged

mkolodner-sc added 2 commits July 13, 2026 19:00

Sort residual top-up PPR results by score

4a8d631

Calibrate finalized PPR top-up scores

5ace07f

Fix PPR sampler test type narrowing

2bcb990

kmontemayor2-sc reviewed Jul 13, 2026

View reviewed changes

Address PPR residual top-up review feedback

e06eaf3

Base automatically changed from mkolodner/ppr-neighbor-cache-reuse to main July 13, 2026 23:03

kmontemayor2-sc reviewed Jul 13, 2026

View reviewed changes

Address follow-up PPR review feedback

ae89d73

kmontemayor2-sc reviewed Jul 14, 2026

View reviewed changes

		std::vector<std::pair<int32_t, double>> selectedPairs;
		selectedPairs.reserve(static_cast<size_t>(topK + residualBudget));

		@@ -44,6 +44,12 @@ class PPRSamplerOptions:

		For homogeneous graphs these live directly on ``data.edge_index`` / ``data.edge_attr``.

Uh oh!

Conversation

mkolodner-sc commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Reasoning for Change

Changes Summary

Changes

Validation

Uh oh!

mkolodner-sc commented Jul 13, 2026

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GiGL Automation

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GiGL Automation

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GiGL Automation

Uh oh!

mkolodner-sc commented Jul 13, 2026

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GiGL Automation

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GiGL Automation

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

GiGL Automation

Uh oh!

kmontemayor2-sc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mkolodner-sc commented Jul 13, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

mkolodner-sc commented Jul 10, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading

github-actions Bot commented Jul 13, 2026 •

edited

Loading

mkolodner-sc Jul 13, 2026 •

edited

Loading