diff --git a/phoenix/crates/serving/xai-recsys-engine/src/python.rs b/phoenix/crates/serving/xai-recsys-engine/src/python.rs index fa3e6651..69be657f 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/python.rs +++ b/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>, large_k: usize, batch_size: usize, + all_top_k_validity: Option>>>, ) -> PyResult<()> { let donated_indices: Vec> = all_top_k_indices .into_iter() @@ -787,16 +816,33 @@ impl RetrieveRequestBatch { .into_iter() .map(|a| DonatedArray2::new(a)) .collect::>>()?; + let donated_validity: Option>> = all_top_k_validity + .map(|arrays| { + arrays + .into_iter() + .map(DonatedArray2::new) + .collect::>>() + }) + .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) }; diff --git a/phoenix/xrex/cuda/top_k_by_key/__init__.py b/phoenix/xrex/cuda/top_k_by_key/__init__.py index 79b9e2b0..4507bb37 100644 --- a/phoenix/xrex/cuda/top_k_by_key/__init__.py +++ b/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: diff --git a/phoenix/xrex/inference/model_runner.py b/phoenix/xrex/inference/model_runner.py index 3455b865..8d444ee1 100644 --- a/phoenix/xrex/inference/model_runner.py +++ b/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( diff --git a/phoenix/xrex/inference/serving_filters_runner.py b/phoenix/xrex/inference/serving_filters_runner.py index dfb2bda5..c675417c 100644 --- a/phoenix/xrex/inference/serving_filters_runner.py +++ b/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 diff --git a/phoenix/xrex/models/recsys_two_tower_model.py b/phoenix/xrex/models/recsys_two_tower_model.py index e142e9c2..5aeec0d8 100644 --- a/phoenix/xrex/models/recsys_two_tower_model.py +++ b/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) diff --git a/phoenix/xrex/models/recsys_two_tower_serving_filters.py b/phoenix/xrex/models/recsys_two_tower_serving_filters.py index c20e1a7a..d365c04e 100644 --- a/phoenix/xrex/models/recsys_two_tower_serving_filters.py +++ b/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) diff --git a/phoenix/xrex/models/recsys_two_tower_serving_filters_test.py b/phoenix/xrex/models/recsys_two_tower_serving_filters_test.py new file mode 100644 index 00000000..596d9075 --- /dev/null +++ b/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()