Kernel issues, with docker image:main

[gane5hvarma]

System Info

spec - aws g6e.12xLarge

Hi, I'm trying out lorax. I ran a docker container with image tag as main(ghcr.io/predibase/lorax:main) and was facing some kernel issues. Attaching logs. After changing the image tag to latest, the server has started. Reporting this issue, so it can help you debug and fix. Thank you

indicies, layer_idx, 1.0)\nRuntimeError: No suitable kernel. h_in=256 h_out=1024 dtype=BFloat16\n"},"target":"lorax_launcher"}
{"timestamp":"2025-04-02T05:56:29.502899Z","level":"ERROR","message":"Server error: No suitable kernel. h_in=256 h_out=1024 dtype=BFloat16","target":"lorax_client","filename":"router/client/src/lib.rs","line_number":38,"span":{"name":"warmup"},"spans":[{"max_input_length":4095,"max_prefill_tokens":4145,"max_total_tokens":4096,"name":"warmup"},{"name":"warmup"}]}
{"timestamp":"2025-04-02T05:56:29.504029Z","level":"ERROR","message":"Server error: No suitable kernel. h_in=256 h_out=1024 dtype=BFloat16","target":"lorax_client","filename":"router/client/src/lib.rs","line_number":38,"span":{"name":"warmup"},"spans":[{"max_input_length":4095,"max_prefill_tokens":4145,"max_total_tokens":4096,"name":"warmup"},{"name":"warmup"}]}
{"timestamp":"2025-04-02T05:56:29.504301Z","level":"ERROR","message":"Server error: No suitable kernel. h_in=256 h_out=1024 dtype=BFloat16","target":"lorax_client","filename":"router/client/src/lib.rs","line_number":38,"span":{"name":"warmup"},"spans":[{"max_input_length":4095,"max_prefill_tokens":4145,"max_total_tokens":4096,"name":"warmup"},{"name":"warmup"}]}
{"timestamp":"2025-04-02T05:56:29.506594Z","level":"ERROR","message":"Server error: No suitable kernel. h_in=256 h_out=1024 dtype=BFloat16","target":"lorax_client","filename":"router/client/src/lib.rs","line_number":38,"span":{"name":"warmup"},"spans":[{"max_input_length":4095,"max_prefill_tokens":4145,"max_total_tokens":4096,"name":"warmup"},{"name":"warmup"}]}
Error: Warmup(Generation("No suitable kernel. h_in=256 h_out=1024 dtype=BFloat16"))

Information

• [x] Docker
• [ ] The CLI directly

Tasks

• [x] An officially supported command
• [ ] My own modifications

Reproduction

FROM ghcr.io/predibase/lorax:main
ENV HUGGINGFACE_HUB_CACHE=/data 
ENV HF_HUB_ENABLE_HF_TRANSFER=1

ENTRYPOINT ["lorax-launcher", "--json-output", "--model-id", "meta-llama/Llama-3.1-70B-Instruct", "--num-shard", "4", "--port", "80"]

docker build -f Dockerfile . -t lorax volume=$PWD/data docker run --gpus all --env-file .env --shm-size 1g -p 8080:80 -v $volume:/data lorax

Expected behavior

The webserver should start