GPU: batch query pipeline with CUDA stream overlap

[TinDang97]

Context

The current search path processes one query at a time through the HNSW beam search. For single-query latency this is optimal (1.5-2.5 ms). However, under high concurrent load (e.g., 100+ parallel FT.SEARCH), each query independently traverses the graph.

A GPU batch query pipeline would amortize device setup and memory transfer costs across multiple queries, significantly improving throughput (queries/second) while trading slightly higher per-query latency.

What needs to happen

1. Query batching infrastructure

pub struct QueryBatch {
    queries: Vec<AlignedBuffer<f32>>,  // rotated query vectors
    filters: Vec<Option<RoaringBitmap>>,
    k_values: Vec<usize>,
    ef_values: Vec<usize>,
    /// Response channels (one per query)
    senders: Vec<oneshot::Sender<Vec<SearchResult>>>,
}

• Accumulate incoming FT.SEARCH queries in a bounded queue per shard
• Flush when: batch full (e.g., 32 queries) OR timeout (e.g., 1ms)
• Configurable via FT.CREATE ... BATCH_SIZE 32 BATCH_TIMEOUT_MS 1

2. GPU batch search kernel

For mutable segment (brute-force):

• Upload batch of Q queries + N database vectors
• Compute Q × N distance matrix via cuBLAS GEMM
• Top-k selection on GPU (radix select or bitonic sort)
• Download only top-k results per query

For immutable segment:

• HNSW beam search is inherently sequential per query (data-dependent traversal)
• GPU helps only with distance computation within each beam step
• Alternative: run Q independent beam searches in parallel GPU threads (one warp per query)

3. CUDA stream pipelining

Stream 1: Upload batch[i] → Compute batch[i] → Download batch[i]
Stream 2:                    Upload batch[i+1] → Compute batch[i+1] → Download batch[i+1]

• Double-buffered: while one batch computes, next batch uploads
• Overlap PCIe transfer with GPU computation
• Requires pinned host memory (from GPU memory pool issue)

When this matters

Scenario	Single-query	Batch (32 queries)
10K mutable, brute-force	300 μs × 32 = 9.6 ms	~1 ms total
100K mutable, brute-force	3 ms × 32 = 96 ms	~5 ms total
HNSW search (not batchable)	1.5 ms × 32 = 48 ms	~48 ms (no gain)

Key insight: GPU batch helps most for mutable segment brute-force at scale. HNSW search benefits minimally since traversal is sequential.

Acceptance criteria

• Query batching queue in VectorStore or SegmentHolder
• Configurable batch size and timeout
• GPU batch brute-force for mutable segment
• CUDA stream double-buffering for transfer/compute overlap
• Latency measurement: p50/p99 with batching enabled vs disabled
• Throughput measurement: QPS at 32/64/128 concurrent clients
• Fallback: batch processed on CPU when GPU unavailable

Dependencies

• GPU memory pool (issue: device memory pool)
• Batch f32 L2 distance kernel (issue: batch distance computation)
• PTX compilation pipeline (issue: FWHT CUDA kernel)

Files to create/modify

• src/vector/batch_query.rs — new query batching module
• src/vector/gpu/stream.rs — CUDA stream management
• src/vector/store.rs — integrate batch queue into VectorStore
• src/vector/segment/mutable.rs — GPU-accelerated brute-force path
• src/command/vector_search.rs — opt-in batching in FT.SEARCH handler