Prepare the runtime
Load the tokenizer and model
Prepare the model input
TODO 1Format a user prompt with the model chat template and move the resulting tensors to the model device.
Check: System and user messages with a generation prompt, thinking disabled, and PyTorch tensors returned as a dictionary.
Show hint
Use apply_chat_template with add_generation_prompt, return_tensors, and return_dict enabled.
Add temperature and top-k sampling
TODO 2Implement greedy decoding, temperature scaling, top-k filtering, and multinomial sampling.
Check: Temperature zero uses argmax. Sampled tokens come only from the retained top-k candidates.
Show hint
Divide logits by temperature, mask values below the kth logit, apply softmax, then call torch.multinomial.
Build generation without KV cache
TODO 3Implement the baseline autoregressive loop that sends the complete sequence through the model for every new token.
Check: Generated tokens and attention mask grow together, generation stops on EOS, and only completion text is returned.
Show hint
Concatenate next_token to generated_ids and append a column of ones to attention_mask.
Add KV-cached generation
TODO 4Reuse the model KV cache so every decode step processes only the newest token.
Check: The first call processes the prompt, later calls receive one token plus past_key_values, and greedy output matches the uncached loop.
Show hint
Select from the current logits, then call model with next_token, the extended mask, past_key_values, and use_cache enabled.
Calculate serving metrics
TODO 5Turn raw generation timings into the serving metrics used to compare both implementations.
Check: Output token count and output tokens per second calculated from the generated token IDs and total latency.
Show hint
Divide the number of generated token IDs by latency in seconds.
Compare generation with and without KV cache
Stream generated text
TODO 6Stream the completion as it is decoded without repeating text that was already sent.
Check: Each yielded chunk contains only new text, and joining all chunks reproduces the cached completion.
Show hint
Decode completion tokens, subtract previous_text from current_text, then update previous_text.
Define the request and response
TODO 7Define a safe API contract for generation requests and measured responses.
Check: Bounded generation settings and response fields for text, TTFT, latency, input and output token counts, throughput, and cache usage.
Show hint
Limit max_new_tokens to 512, temperature to 0 through 2, and top_k to 1 through 100.
Create the server endpoints
TODO 8Connect cached generation and streaming to the HTTP endpoints.
Check: /generate returns GenerateResponse with metrics, while /stream returns a plain-text StreamingResponse.
Show hint
Summarize generate_cached for the regular endpoint and pass generate_stream to StreamingResponse.
Test the inference server
TODO 9Send requests through the FastAPI test client and verify the health and generation responses.
Check: Both endpoints return HTTP 200 and the generation response contains at least one output token.
Show hint
Use client.get for /health and client.post with a short prompt and max_new_tokens in the JSON body.