over the last 6-8 months, ive been trying to move towards the ml systems and ai infra space. these are some of my favorite work ive done -

1. a small python inference engine which is mainly a testbed for me to implement the major inference techniques, its about 600 toks/s on an rtx 4060 ti with prefix caching - https://github.com/JINO-ROHIT/tachyon…
2. wrote a blog post that outperforms cublas on ada using cute dsl - https://jino-rohit.github.io/blogs/08_cute_dsl_matmul.html…
3. a detailed post on ncccl collective communication - https://jino-rohit.github.io/blogs/11_collective_communication.html…
4. a blog series on how torch compile works and internal mechanics - https://jino-rohit.github.io/blogs/
5. i maintain my notes and experiments of most of the work around ml systems here - https://github.com/JINO-ROHIT/ml-systems-notes…
6. open source work in sglang and llm-compressor(vllm).

im trying to become a stronger ml systems and inference engineer. what should i spend my next months getting better on?

JINO-ROHIT/tachyon

Source: https://github.com/JINO-ROHIT/tachyon

tachyon

a LLM inference engine to run on consumer hardware.

1. each request has its own state and you can choose to disable or enable kv cache on a request level.
2. supports continous batching where the decode step is batched to max out gpu, and requests get recycled as soon as it is completed from the active batch to allow new requests to enter the batch.
3. the prefill step is batched using a smaller batch size to avoid OOMs
4. vectorized operations to improve batch performance.
5. prefix caching to save on expensive prefill computation when there is shared system prompts or few shot examples etc.
6. 3 lines to invoke the engine and run inference!

Usage

1. first download the weights for llama 1B.

hf_hub_download(
    repo_id=f"meta-llama/Llama-3.2-1B-Instruct",
    filename="model.safetensors",
    local_dir=f"Llama-3.2-1B-Instruct"
)

2. invoke the engine.

from tachyon.engine.llm import Engine
engine = Engine("meta-llama/Llama-3.2-1B-Instruct")
print(engine.generate_text("Explain AGI")) # for single request

#for multiple requests
outputs = engine.generate_text([
    "Explain AGI",
    "What is vLLM?",
    "Tell me about SGLang"
])

for o in outputs:
    print(o)

OpenAI compatible server

You can spin up the server and then use the openai library to invoke the model as well -

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="anything"
)

resp = client.chat.completions.create(
    model="llama-3.2-1b-instruct",
    messages=[{"role": "user", "content": "Hello"}],
    stream=False,
)

3. benchmark script

python3 benchmark.py

current benchmarks(rtx 4060 ti 16GB)

to-do

• implement llama 3 family model.
• make it into a serving engine.
• write a benchmark script and check latency and throughput.
• add kv cache.
• continous batching
• prefix caching
• add an openai compatible api server
• test the effect of torch compile (open PR, come back to this)
• paged attention
• more techniques