Keep masked candidates out of underfilled retrieval responses

phoenix/crates/serving/xai-recsys-engine/src/python.rs

@@ -51,6 +51,24 @@ fn bloom_may_contain(bit_array: &[u64], num_bits: usize, post_id: i64) -> bool {
         idx < num_bits && (bit_array[idx >> 6] & (1u64 << (idx & 63))) != 0
     })
 }
+
+#[inline]
+fn should_emit_retrieval_candidate(is_valid: bool, post_id: i64) -> bool {
+    is_valid && post_id != 0
+}
+
+#[cfg(test)]
+mod retrieval_reply_tests {
+    use super::should_emit_retrieval_candidate;
+
+    #[test]
+    fn masked_candidates_are_not_emitted() {
+        assert!(should_emit_retrieval_candidate(true, 42));
+        assert!(!should_emit_retrieval_candidate(false, 42));
+        assert!(!should_emit_retrieval_candidate(true, 0));
+    }
+}
+
 use axum::routing::get;
 use bytes::Bytes;
 use chrono::TimeDelta;
@@ -768,6 +786,16 @@ impl RetrieveRequestBatch {
         Ok(())
     }
 
+    #[pyo3(signature = (
+        dataset_types,
+        all_top_k_indices,
+        all_top_k_scores,
+        all_post_ids,
+        all_author_ids,
+        large_k,
+        batch_size,
+        all_top_k_validity=None
+    ))]
     pub fn reply(
         &mut self,
         py: Python<'_>,
@@ -778,6 +806,7 @@ impl RetrieveRequestBatch {
         all_author_ids: Bound<'_, PyArray1<i64>>,
         large_k: usize,
         batch_size: usize,
+        all_top_k_validity: Option<Vec<Bound<'_, PyArray2<bool>>>>,
     ) -> PyResult<()> {
         let donated_indices: Vec<DonatedArray2<i32>> = all_top_k_indices
             .into_iter()
@@ -787,16 +816,33 @@ impl RetrieveRequestBatch {
             .into_iter()
             .map(|a| DonatedArray2::new(a))
             .collect::<PyResult<Vec<_>>>()?;
+        let donated_validity: Option<Vec<DonatedArray2<bool>>> = all_top_k_validity
+            .map(|arrays| {
+                arrays
+                    .into_iter()
+                    .map(DonatedArray2::new)
+                    .collect::<PyResult<Vec<_>>>()
+            })
+            .transpose()?;
         let post_ids = DonatedArray1::new(all_post_ids)?;
         let author_ids = DonatedArray1::new(all_author_ids)?;
 
         let num_datasets = dataset_types.len();
-        if donated_indices.len() != num_datasets || donated_scores.len() != num_datasets {
+        if donated_indices.len() != num_datasets
+            || donated_scores.len() != num_datasets
+            || donated_validity
+                .as_ref()
+                .is_some_and(|validity| validity.len() != num_datasets)
+        {
+            let validity_len = donated_validity
+                .as_ref()
+                .map_or(num_datasets, std::vec::Vec::len);
             return Err(pyo3::exceptions::PyValueError::new_err(format!(
-                "Mismatched array lengths: dataset_types={}, indices={}, scores={}",
+                "Mismatched array lengths: dataset_types={}, indices={}, scores={}, validity={}",
                 num_datasets,
                 donated_indices.len(),
-                donated_scores.len()
+                donated_scores.len(),
+                validity_len
             )));
         }
 
@@ -827,12 +873,25 @@ impl RetrieveRequestBatch {
                         let scores = &donated_scores[ds_idx];
                         let indices_view = unsafe { indices.as_array_view() };
                         let scores_view = unsafe { scores.as_array_view() };
-
-                        let k = large_k.min(indices_view.ncols());
+                        let validity_view = donated_validity
+                            .as_ref()
+                            .map(|validity| unsafe { validity[ds_idx].as_array_view() });
+
+                        let k = large_k
+                            .min(indices_view.ncols())
+                            .min(scores_view.ncols())
+                            .min(
+                                validity_view
+                                    .as_ref()
+                                    .map_or(usize::MAX, |validity| validity.ncols()),
+                            );
                         for j in 0..k {
                             let idx = indices_view[[user_idx, j]] as usize;
                             let pid = unsafe { post_ids.get(idx) };
-                            if pid == 0 {
+                            let is_valid = validity_view
+                                .as_ref()
+                                .is_none_or(|validity| validity[[user_idx, j]]);
+                            if !should_emit_retrieval_candidate(is_valid, pid) {
                                 continue;
                             }
                             let aid = unsafe { author_ids.get(idx) };

phoenix/xrex/cuda/top_k_by_key/__init__.py

@@ -3,6 +3,11 @@
 import jax
 import jax.numpy as jnp
 
+
+def gather_selected_validity(eligibility: jax.Array, selected_indices: jax.Array) -> jax.Array:
+    return jnp.take_along_axis(eligibility, selected_indices, axis=1)
+
+
 try:
     from xrex.cuda.top_k_by_key.src import top_k_by_key_api
 except ImportError:

phoenix/xrex/inference/model_runner.py

@@ -4530,7 +4530,7 @@ def gather_embeddings_and_forward(
         request: xai_recsys_engine.RetrieveRequestBatch | None = None,
         eligible_mask: jax.Array | None = None,
         bucket_size: int | None = None,
-    ) -> dict[int, tuple[jax.Array, jax.Array]]:
+    ) -> dict[int, tuple[jax.Array, jax.Array, jax.Array]]:
         if self._live_swap_enabled:
             state = self.state
         forward_jit = self._forward_jit_for_bucket(bucket_size)
@@ -4624,7 +4624,7 @@ def _sds(x: Any) -> Any:
     def reply_request(
         self,
         request: xai_recsys_engine.RetrieveRequestBatch,
-        output_dict: dict[int, tuple[np.ndarray, np.ndarray]],
+        output_dict: dict[int, tuple[np.ndarray, np.ndarray, np.ndarray]],
         orig_batch_size: int,
         bucket_size: int = 0,
     ) -> None:
@@ -4640,6 +4640,10 @@ def reply_request(
             np.array(output_dict[ds][1][:, : self.large_k], dtype=np.float32, copy=True)
             for ds in ds_types
         ]
+        all_validity = [
+            np.array(output_dict[ds][2][:, : self.large_k], dtype=np.bool_, copy=True)
+            for ds in ds_types
+        ]
 
         request.reply(
             ds_types,
@@ -4649,6 +4653,7 @@ def reply_request(
             self.all_author_ids,
             self.large_k,
             orig_batch_size,
+            all_validity,
         )
 
     def create_server(
@@ -4859,6 +4864,7 @@ def two_tower_forward_fn(
                 dataset_ranges=dataset_ranges,
                 use_async_topk=self.enable_async_topk,
                 use_radix_select_topk=self.enable_radix_select_topk,
+                return_validity=True,
             )
 
         return JittedOrCompiled(

phoenix/xrex/inference/serving_filters_runner.py

@@ -293,7 +293,7 @@ def gather_embeddings_and_forward(
         request: xai_recsys_engine.RetrieveRequestBatch | None = None,
         eligible_mask: jax.Array | None = None,
         bucket_size: int | None = None,
-    ) -> dict[int, tuple[jax.Array, jax.Array]]:
+    ) -> dict[int, tuple[jax.Array, jax.Array, jax.Array]]:
         forward_jit = self._forward_jit_for_bucket(bucket_size)
         bs = bucket_size if bucket_size is not None else self.inference_batch_size
         assert state.post_embeddings is not None

phoenix/xrex/models/recsys_two_tower_model.py

@@ -1317,7 +1317,8 @@ def forward(
         use_async_topk: bool = False,
         use_radix_select_topk: bool = False,
         post_scales: jax.Array | None = None,
-    ) -> tuple[tuple[jax.Array, jax.Array], ...]:
+        return_validity: bool = False,
+    ) -> tuple[tuple[jax.Array, jax.Array] | tuple[jax.Array, jax.Array, jax.Array], ...]:
         saxis = "expert"
 
         user_representation, _, _ = self(
@@ -1437,7 +1438,13 @@ def slice_and_top_k(all_scores, _start=start, _end=end):
                     return top_k_scores, top_k_indices
 
                 top_k_scores, top_k_indices = slice_and_top_k(all_scores)
-                results.append((top_k_indices, top_k_scores.astype(jnp.float32)))
+                if return_validity:
+                    top_k_validity = jnp.ones_like(top_k_indices, dtype=jnp.bool_)
+                    results.append(
+                        (top_k_indices, top_k_scores.astype(jnp.float32), top_k_validity)
+                    )
+                else:
+                    results.append((top_k_indices, top_k_scores.astype(jnp.float32)))
             return tuple(results)
 
         results = []
@@ -1448,7 +1455,11 @@ def slice_and_top_k(all_scores, _start=start, _end=end):
                 type_mask = jnp.ones(post_embeddings.shape[0], dtype=jnp.bool_)
 
             top_k_scores, top_k_indices = mask_and_top_k(all_scores, type_mask)
-            results.append((top_k_indices, top_k_scores.astype(jnp.float32)))
+            if return_validity:
+                top_k_validity = jnp.take(type_mask, top_k_indices)
+                results.append((top_k_indices, top_k_scores.astype(jnp.float32), top_k_validity))
+            else:
+                results.append((top_k_indices, top_k_scores.astype(jnp.float32)))
 
         return tuple(results)
 

phoenix/xrex/models/recsys_two_tower_serving_filters.py

@@ -2,13 +2,12 @@
 # Copyright 2026 X.AI Corp.
 import os
 
-import haiku as hk
 import jax
 import jax.numpy as jnp
 from jax import shard_map
 from jax.sharding import PartitionSpec as P
 
-from xrex.cuda.top_k_by_key import top_k_by_key
+from xrex.cuda.top_k_by_key import gather_selected_validity, top_k_by_key
 from xrex.data.recsys.recsys_batch import RecsysFeaturesBatch
 from xrex.models.recsys_embedding import RecsysEmbeddings
 from xrex.models.topic_categories import NUM_TOPIC_INT32S
@@ -28,7 +27,7 @@ def forward_with_filters(
     eval_bs_per_device: int = 0,
     dataset_ranges: tuple[tuple[int, int], ...] | None = None,
     use_async_topk: bool = False,
-) -> tuple[tuple[jax.Array, jax.Array], ...]:
+) -> tuple[tuple[jax.Array, jax.Array, jax.Array], ...]:
     saxis = "expert"
 
     user_representation, _, _ = model(
@@ -78,7 +77,7 @@ def compute_top_k(post_embeddings: jax.Array, user_embedding: jax.Array) -> jax.
     @shard_map(
         mesh=mesh,
         in_specs=(P(), P(), mask_in_spec, topic_bitmaps_in_spec, topic_user_bitmasks_in_spec),
-        out_specs=(P(), P()),
+        out_specs=(P(), P(), P()),
         check_vma=False,
     )
     def mask_and_top_k(
@@ -87,15 +86,14 @@ def mask_and_top_k(
         user_eligible_mask: jax.Array,
         topic_bitmaps_shard: jax.Array,
         topic_user_bitmasks_full: jax.Array,
-    ) -> tuple[jax.Array, jax.Array]:
+    ) -> tuple[jax.Array, jax.Array, jax.Array]:
         combined = type_mask & user_eligible_mask
-        masked_scores = jnp.where(combined, all_scores, jnp.finfo(jnp.bfloat16).min)
 
         if use_topic_filter:
             topic_bitmaps_full = jax.lax.all_gather(
                 topic_bitmaps_shard, axis_name=saxis, axis=0, tiled=True
             )
-            B_local = masked_scores.shape[0]
+            B_local = combined.shape[0]
             shard_idx = jax.lax.axis_index(saxis)
             start_idx = shard_idx * B_local
             user_bitmasks_local = jax.lax.dynamic_slice(
@@ -107,12 +105,15 @@ def mask_and_top_k(
             )
             no_filter_mask = jnp.all(user_bitmasks_local == 0, axis=-1)[:, None]
             topic_mask = jnp.where(no_filter_mask, True, topic_mask)
-            masked_scores = jnp.where(topic_mask, masked_scores, jnp.finfo(jnp.bfloat16).min)
+            combined = combined & topic_mask
 
+        masked_scores = jnp.where(combined, all_scores, jnp.finfo(jnp.bfloat16).min)
         sorted_scores, sorted_indices = local_top_k(masked_scores, top_k)
+        sorted_validity = gather_selected_validity(combined, sorted_indices)
         top_k_scores = jax.lax.all_gather(sorted_scores, axis_name=saxis, axis=0, tiled=True)
         top_k_indices = jax.lax.all_gather(sorted_indices, axis_name=saxis, axis=0, tiled=True)
-        return top_k_scores, top_k_indices
+        top_k_validity = jax.lax.all_gather(sorted_validity, axis_name=saxis, axis=0, tiled=True)
+        return top_k_scores, top_k_indices, top_k_validity
 
     all_scores = compute_top_k(post_embeddings, user_representation)
 
@@ -144,7 +145,8 @@ def slice_and_top_k(all_scores, _start=start, _end=end):
                 return top_k_scores, top_k_indices
 
             top_k_scores, top_k_indices = slice_and_top_k(all_scores)
-            results.append((top_k_indices, top_k_scores.astype(jnp.float32)))
+            top_k_validity = jnp.ones_like(top_k_indices, dtype=jnp.bool_)
+            results.append((top_k_indices, top_k_scores.astype(jnp.float32), top_k_validity))
         return tuple(results)
 
     B = all_scores.shape[0]
@@ -165,9 +167,9 @@ def slice_and_top_k(all_scores, _start=start, _end=end):
         _topic_user_bitmasks = (
             topic_user_bitmasks if use_topic_filter else jnp.zeros((), dtype=jnp.int32)
         )
-        top_k_scores, top_k_indices = mask_and_top_k(
+        top_k_scores, top_k_indices, top_k_validity = mask_and_top_k(
             all_scores, type_mask, user_eligible_mask, _topic_bitmaps, _topic_user_bitmasks
         )
-        results.append((top_k_indices, top_k_scores.astype(jnp.float32)))
+        results.append((top_k_indices, top_k_scores.astype(jnp.float32), top_k_validity))
 
     return tuple(results)

phoenix/xrex/models/recsys_two_tower_serving_filters_test.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 X.AI Corp.
+import unittest
+
+import jax.numpy as jnp
+import numpy as np
+
+from xrex.cuda.top_k_by_key import gather_selected_validity, top_k_by_key
+
+
+class SelectedValidityTest(unittest.TestCase):
+    def test_underfilled_top_k_marks_masked_candidates(self):
+        eligibility = jnp.array([[False, True, False, False]])
+        scores = jnp.array([[3.0, 8.0, 7.0, 6.0]], dtype=jnp.bfloat16)
+        masked_scores = jnp.where(eligibility, scores, jnp.finfo(jnp.bfloat16).min)
+
+        _, selected_indices = top_k_by_key(masked_scores, 3, heuristic_pivot_ratio=0.1)
+        selected_validity = np.asarray(gather_selected_validity(eligibility, selected_indices))
+
+        self.assertEqual(selected_validity.shape, (1, 3))
+        self.assertEqual(int(selected_validity.sum()), 1)
+        self.assertTrue(selected_validity[0, 0])
+
+
+if __name__ == "__main__":
+    unittest.main()