[BUG] CAGRA search recall inconsistent with query batch size.

[singhmanas1]

Describe the bug
Batching CAGRA search with query_batch_size >= 512 has lower recall than batch size < 512

Steps/Code to reproduce bug

def calc_recall(found_indices, ground_truth):
    found_indices = cp.asarray(found_indices)
    bs, k = found_indices.shape
    if bs != ground_truth.shape[0]:
        raise RuntimeError(
            "Batch sizes do not match {} vs {}".format(
                bs, ground_truth.shape[0]
            )
        )
    if k > ground_truth.shape[1]:
        raise RuntimeError(
            "Not enough indices in the ground truth ({} > {})".format(
                k, ground_truth.shape[1]
            )
        )
    n = 0
    # Go over the batch
    for i in range(bs):
        # Note, ivf-pq does not guarantee the ordered input, hence the use of intersect1d
        n += cp.intersect1d(found_indices[i, :k], ground_truth[i, :k]).size
        # To-do: Change to account for equidistant indices that are not captured.
    
    #recall = n / found_indices.size
    recall = n / (bs * ground_truth.shape[1])
    return recall


def test_cagra_batch_recall():
    """Test CAGRA batch consistency using only cuVS functions."""
    import cupy as cp
    from cuvs.neighbors import cagra, brute_force
    
    # Setup test data
    np.random.seed(42)
    cp.random.seed(42)
    
    n_samples, n_queries, dim, k = 5000000, 10000, 64, 10
    vectors = np.random.randn(n_samples, dim).astype(np.float32)  # cuVS prefers float32
    queries = np.random.randn(n_queries, dim).astype(np.float32)
    
    # Convert to GPU
    vectors_gpu = cp.asarray(vectors)
    queries_gpu = cp.asarray(queries)
    
    # Generate ground truth using cuVS brute force
    print("Generating ground truth with cuVS brute force...")
    bf_index = brute_force.build(vectors_gpu, metric="sqeuclidean")
    gt_distances, gt_indices = brute_force.search(bf_index, queries_gpu, k)
    gt_indices = cp.asnumpy(gt_indices)  # Convert to CPU for comparison
    
    # Build CAGRA index
    print("Building CAGRA index...")
    cagra_index_params = cagra.IndexParams(graph_degree=32, intermediate_graph_degree=64)
    cagra_index = cagra.build(cagra_index_params, vectors_gpu)
    search_params = cagra.SearchParams(itopk_size=64, search_width=8)
    
    # Test different batch sizes
    batch_sizes = [1, 8, 64, 256, 512, 1024]
    recalls = []
    
    for batch_size in batch_sizes:
        print(f"Testing batch_size={batch_size}...")
        
        # Manual batching with pure cuVS
        all_indices = np.zeros((n_queries, k), dtype=np.int32)
        
        for start_idx in range(0, n_queries, batch_size):
            end_idx = min(start_idx + batch_size, n_queries)
            batch_queries = queries_gpu[start_idx:end_idx]
            
            # Pure cuVS search
            distances, indices = cagra.search(search_params, cagra_index, batch_queries, k=k)
            
            # Store results
            indices_cpu = cp.asnumpy(indices)
            all_indices[start_idx:end_idx] = indices_cpu
        
        # Calculate recall manually (no helper functions)
        correct = 0
        total = 0
        
        for i in range(n_queries):
            # Count how many CAGRA results are in the ground truth for this query
            cagra_neighbors = set(all_indices[i])
            gt_neighbors = set(gt_indices[i])
            correct += len(cagra_neighbors.intersection(gt_neighbors))
            total += k
        
        recall = correct / total
        recalls.append(recall)
        print(f"  batch_size={batch_size}, recall={recall:.4f}")
    
    # Check consistency
    recall_std = np.std(recalls)
    min_recall = min(recalls)
    max_recall = max(recalls)
    
    print(f"\nCUVS CAGRA Batch Consistency Results:")
    print(f"  Recalls: {[f'{r:.4f}' for r in recalls]}")
    print(f"  Min: {min_recall:.4f}, Max: {max_recall:.4f}")
    print(f"  Std Dev: {recall_std:.8f}")
    
    # The test - CAGRA should be consistent across batch sizes
    if recall_std < 1e-4:
        print("PASS: CAGRA shows good batch consistency")
        return True
    else:
        print(f"FAIL: CAGRA shows poor batch consistency (std={recall_std:.8f})")
        return False
        
        recall = calc_recall(search_indices, gt_indices)
        recalls.append(recall)
    
    # All recalls should be identical
    recall_std = np.std(recalls)
    assert recall_std < 1e-6, f"Recall varies with batch size! Values: {recalls}, Std: {recall_std:.8f}"
    
    print(f"PASS: Recall batch independence test (recall: {recalls[0]:.4f}, std: {recall_std:.8f})")
    return True

Output-

Testing batch_size=1...
  batch_size=1, recall=0.2986
Testing batch_size=8...
  batch_size=8, recall=0.2976
Testing batch_size=64...
  batch_size=64, recall=0.2972
Testing batch_size=256...
  batch_size=256, recall=0.2965
Testing batch_size=512...
  batch_size=512, recall=0.1231
Testing batch_size=1024...
  batch_size=1024, recall=0.1178

Expected behavior
Recall shouldn't be dependent on query batch size.

Environment details (please complete the following information):

• AWS g6e.16xlarge (1 L40S)
• cuvs 25.06

[zbennett10]

I've opened a PR to fix this: #1841

Root cause: The AUTO algorithm selector switches to SINGLE_CTA at large batch sizes, but SINGLE_CTA computes max_iterations = itopk_size / search_width. When search_width > 1 (e.g., 8), this results in significantly fewer iterations than MULTI_CTA (which uses its own internal mc_search_width=1), causing the recall degradation.

Fix: Added search_width <= 1 guard to the SINGLE_CTA selection path, so AUTO only picks SINGLE_CTA when the search_width won't reduce iteration count.

Includes a Python regression test that verifies recall consistency across batch sizes [64, 256, 512, 1024] with search_width values of 1, 4, and 8.

[zbennett10]

I've submitted a PR to fix this: #1841

The fix prevents recall degradation at batch sizes >= 512 by clamping the internal topk parameter when search_width > 1.

Root cause analysis: When search_width > 1, the algorithm internally multiplies k × search_width, but this can exceed thresholds that degrade search quality at large batch sizes.

This was causing cache hit rate degradation in our production semantic caching system (Synapse) when processing large batches of embeddings.