[data][llm] Use numpy arrays for embeddings to avoid torch.Tensor serialization overhead

[jeffreywang-anyscale]

Description

Currently, Ray Data LLM uses torch.Tensor to store embeddings from pooling tasks (e.g., classify) in Ray dataset columns. This causes redundant copies during serialization and deserialization. Using NumPy arrays is better because they support zero-copy deserialization via Ray's shared memory object store.

When embeddings are stored in Ray Data columns:

• np.array: Data stored once in shared memory → multiple workers read via pointers (zero-copy)
• torch.Tensor: Data copied into pickle stream → copied again on deserialize (2x copies)

Benchmark Result

Using np.array, classification tasks with vLLMEngineStage show a ~9% throughput improvement.

Metric	torch.Tensor	np.array
Throughput (row/s)	953	1038

• Hardware: L4
• Repro command:

python benchmark_processor.py --mode classify --batch-size 2048 --concurrency 1 --num-prompts 102400 --model HuggingFaceTB/fineweb-edu-classifier

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request successfully improves performance by switching from torch.Tensor to numpy.ndarray for embeddings, which enables zero-copy deserialization in Ray. The core change in vllm_engine_stage.py is correct and well-targeted. The addition of a new classification mode to the benchmark script, along with a corresponding test case, effectively validates the performance gain and correctness of the new functionality.

I have one suggestion regarding an import statement in the benchmark script to improve its portability and maintainability. Overall, this is a solid contribution that delivers a measurable performance improvement.