Add Habana Gaudi (HPU) Support & Performance Benchmarks for Khoj

[BartoszBLL]

This PR introduces support for Habana Gaudi accelerators (HPUs) to the project, enabling the application to run on HPU devices in addition to the existing support for CUDA, MPS, and CPU. The changes include:

🚀 Key Updates

1. HPU Dockerfile (Dockerfile.hpu)
   • Added a new Dockerfile for running Khoj on Habana Gaudi devices.
   • Installs necessary dependencies (optimum-habana).
   • Configures environment variables for Habana optimizations.
2. Device Selection (helpers.py)
   • Enhanced get_device() to detect HPU if available.
   • Supports cuda, hpu, mps, or cpu based on availability or user preference.
3. Memory Management (helpers.py)
   • get_device_memory() now supports Habana HPU memory queries.
4. Dependency Updates (pyproject.toml)
   • Added optimum-habana, torch-geometric, and numba.
5. Documentation
   • README.md & src/khoj/app/README.md: Instructions for building and running Khoj with HPU.

💎 Why This Matters:

HPU Support

This PR enables the application to leverage Habana Gaudi accelerators, which can provide significant performance improvements for deep learning workloads.

Flexibility

Users can now choose their preferred device (CUDA, HPU, MPS, or CPU) for running the application, making it more versatile across different hardware setups.

Optimization

The addition of optimum-habana ensures that models are optimized for HPU and other hardware, improving efficiency and performance.

⚡ Performance Benchmarks

HPU: ~0.2703s average runtime (10 runs)
CPU: ~76.3144s average runtime (10 runs)
Result: ~282× speedup using HPU compared to CPU.

🛠 How to Test

Use the new Dockerfile.hpu to build and run the application on a system with Habana Gaudi accelerators.

# Build the HPU Docker image
docker build -t khoj-hpu -f Dockerfile.hpu .

# Run with Habana runtime
docker run --runtime=habana -e HABANA_VISIBLE_DEVICES=all -p <PORT>:<PORT> khoj-hpu

Check logs to confirm that HPU is recognized and in use.

✅ Checklist

• Tested on Habana Gaudi accelerator
• Verified compatibility with CPU and other devices
• Updated documentation
• Added required dependencies

📝 Notes

This PR is part of the effort to expand hardware support for the application, ensuring it can run efficiently on a wide range of devices.

[debanjum]

Hi, thanks for creating a PR to add support for Habana HPU to Khoj. Apologize for the delayed review. Not sure if this change should be included in Khoj (yet). Some questions below:

1. Are you using Khoj on Intel Gaudi machines? What's the use-case? Gaudi support seems more relevant for production scenarios. For such setups with Khoj, you should offload both the LLM heavy components (embedding generation and chat model interactions) to appropriate llm inference servers (like vllm, sglang, tensort etc.).
2. You mention performance benchmarks for Khoj with HPU support. Can you clarify what kind of workloads you tested? Was it the rag indexing, interacting with a local/offline chat model or something else? More details on the perf benchmarks would be useful for context

[debanjum]

The preferred_device arg seems unused. Should we remove it?

[debanjum]

Seems like the only dependency explicitly added is optimum-habana (in pyproject.toml)?

[debanjum]

This isn't used and is relevant to Khoj users or deployers. So should be removed from the documention

[debanjum]

These seem like optional runtime variables to configure habana support based on the Habana docs?

If these runtime env vars are the only change this Dockerfile.hpu adds, we can drop this Dockerfile.hpu file and just mention in our setup documentation that folks wanting to run khoj on habana hpu can setup these (and other required) environment variables for their setup by referring to the Habana documentation before starting khoj?

[debanjum]

These seems to be a required runtime environment variable to enable habana hpu? If so, it just be mentioned in the Khoj setup docs under the HPU tab. See /documentation/docs/get-started/setup.mdx for reference

[debanjum]

Updates to this should be instead moved to a new HPU tab under the Khoj setup docs at /documentation/docs/get-started/setup.mdx (which maps to https://docs.khoj.dev/get-started/setup/)

[debanjum]

This may not be required if the previous comments on the Dockerfile.hpu are valid. Folks can just use the default Khoj dockerfile or image