LLM Latency Optimization — How We Cut P99 from 12s to 1.8s Without Changing the Model

import asyncio
from typing import List, Dict
import time

class DynamicBatcher:
    def __init__(self, max_batch_size: int = 8, max_context_length: int = 4096, batch_timeout: float = 2.0):
        self.max_batch_size = max_batch_size
        self.max_context_length = max_context_length  # Truncate or reject longer contexts
        self.batch_timeout = batch_timeout
        self.queue: List[Dict] = []  # Each item: {'request': ..., 'context_length': int, 'future': asyncio.Future}

    async def submit(self, request: Dict) -> Dict:
        context_length = len(request['messages'])  # Simplified; actual token count
        if context_length > self.max_context_length:
            raise ValueError(f"Context length {context_length} exceeds max {self.max_context_length}")
        future = asyncio.get_event_loop().create_future()
        self.queue.append({'request': request, 'context_length': context_length, 'future': future})
        # If queue is full, trigger batch processing
        if len(self.queue) >= self.max_batch_size:
            asyncio.create_task(self._process_batch())
        return await future

    async def _process_batch(self):
        # Sort by context length to minimize straggler effect
        batch = sorted(self.queue, key=lambda x: x['context_length'])[:self.max_batch_size]
        self.queue = self.queue[len(batch):]  # Remove processed items
        try:
            # Simulate LLM call with timeout
            results = await asyncio.wait_for(
                self._llm_call([item['request'] for item in batch]),
                timeout=self.batch_timeout
            )
            for item, result in zip(batch, results):
                item['future'].set_result(result)
        except asyncio.TimeoutError:
            # Fallback: process each request individually
            for item in batch:
                try:
                    result = await self._llm_call([item['request']])
                    item['future'].set_result(result[0])
                except Exception as e:
                    item['future'].set_exception(e)

    async def _llm_call(self, requests: List[Dict]) -> List[Dict]:
        # Placeholder for actual LLM API call
        await asyncio.sleep(0.5)  # Simulate processing
        return [{'response': f'processed {len(requests)} requests'}]

# Usage
batcher = DynamicBatcher()
async def main():
    results = await asyncio.gather(*[batcher.submit({'messages': [{'role': 'user', 'content': 'hello'}]}) for _ in range(10)])
    print(results)

asyncio.run(main())

import asyncio
from typing import List, Optional

class SpeculativeDecoder:
    def __init__(self, draft_model, target_model, max_draft_tokens: int = 5):
        self.draft_model = draft_model  # Small, fast model
        self.target_model = target_model  # Large, accurate model
        self.max_draft_tokens = max_draft_tokens  # How many tokens to guess at once

    async def generate(self, prompt: str) -> str:
        # Step 1: Draft model generates candidate tokens
        draft_tokens = await self.draft_model.generate(prompt, max_tokens=self.max_draft_tokens)
        # Step 2: Target model verifies the draft tokens
        # It returns the logits for each position; we check if the draft token is in the top-k
        logits = await self.target_model.get_logits(prompt + draft_tokens)
        accepted_tokens = []
        for i, token in enumerate(draft_tokens):
            # Check if draft token is in the top 1 (greedy) or top-k (sampling)
            if self._is_token_accepted(logits[i], token):
                accepted_tokens.append(token)
            else:
                # Reject the rest and let target model generate from here
                break
        if len(accepted_tokens) == 0:
            # Fallback: target model generates from scratch
            return await self.target_model.generate(prompt, max_tokens=1)
        # Step 3: If all draft tokens accepted, generate one more token with target model
        if len(accepted_tokens) == self.max_draft_tokens:
            extra_token = await self.target_model.generate(prompt + draft_tokens, max_tokens=1)
            return draft_tokens + extra_token
        return ''.join(accepted_tokens)

    def _is_token_accepted(self, logits, token):
        # Simplified: check if token is the argmax
        import torch
        return token == torch.argmax(logits).item()

# Usage (pseudocode)
# decoder = SpeculativeDecoder(draft_model=SmallModel(), target_model=LargeModel())
# result = await decoder.generate("Write a Python function to sort a list")
# print(result)

import asyncio
from typing import List, Dict

class PromptCompressor:
    def __init__(self, summarizer_model, max_context_length: int = 2048):
        self.summarizer_model = summarizer_model  # Small, fast model for summarization
        self.max_context_length = max_context_length

    async def compress(self, conversation: List[Dict]) -> str:
        # Convert conversation to text
        full_text = '\n'.join([f"{msg['role']}: {msg['content']}" for msg in conversation])
        # If it's short enough, return as-is
        if len(full_text.split()) < self.max_context_length:
            return full_text
        # Summarize the conversation, focusing on key facts
        summary_prompt = f"Summarize the following conversation in under {self.max_context_length // 2} words, keeping all important facts and the user's original request:\n\n{full_text}"
        summary = await self.summarizer_model.generate(summary_prompt, max_tokens=self.max_context_length // 2)
        # Validate: check that key entities from the original are in the summary
        # Simplified: just return the summary
        return summary

# Usage
# compressor = PromptCompressor(summarizer_model=SmallModel())
# compressed = await compressor.compress(conversation)
# print(compressed)

from collections import OrderedDict
from typing import Dict, Any
import torch

class KVCacheManager:
    def __init__(self, max_cache_size_gb: float = 2.0, eviction_policy: str = 'LRU'):
        self.max_cache_size_bytes = int(max_cache_size_gb * 1024**3)
        self.cache: OrderedDict[str, Dict[str, torch.Tensor]] = OrderedDict()  # session_id -> KV cache
        self.current_size = 0

    def get(self, session_id: str) -> Dict[str, torch.Tensor]:
        if session_id in self.cache:
            # Move to end (most recently used)
            self.cache.move_to_end(session_id)
            return self.cache[session_id]
        return None

    def set(self, session_id: str, kv_cache: Dict[str, torch.Tensor]):
        # Estimate size of KV cache (simplified)
        size = sum(tensor.element_size() * tensor.numel() for tensor in kv_cache.values())
        # Evict if needed
        while self.current_size + size > self.max_cache_size_bytes and len(self.cache) > 0:
            # Evict least recently used (first item in OrderedDict)
            evicted_id, evicted_cache = self.cache.popitem(last=False)
            evicted_size = sum(tensor.element_size() * tensor.numel() for tensor in evicted_cache.values())
            self.current_size -= evicted_size
        self.cache[session_id] = kv_cache
        self.current_size += size

    def clear_session(self, session_id: str):
        if session_id in self.cache:
            kv_cache = self.cache.pop(session_id)
            size = sum(tensor.element_size() * tensor.numel() for tensor in kv_cache.values())
            self.current_size -= size

# Usage
# manager = KVCacheManager(max_cache_size_gb=2.0)
# cache = manager.get('session_123')
# if cache is None:
#     cache = compute_kv_cache(...)
#     manager.set('session_123', cache)

def should_optimize_software(daily_requests: int, current_p99: float, target_p99: float) -> str:
    """
    Simple heuristic: if traffic is low, buy more hardware. If traffic is high, optimize software.
    """
    if daily_requests < 1000:
        return "Upgrade hardware (e.g., A100 -> H100). Engineering time is better spent elsewhere."
    elif daily_requests < 10000:
        return "Consider a hybrid approach: upgrade hardware for immediate gains, then optimize software for the long term."
    else:
        return "Optimize software. Hardware costs will dominate at this scale."

# Example
print(should_optimize_software(100, 5.0, 2.0))  # "Upgrade hardware..."
print(should_optimize_software(50000, 5.0, 2.0))  # "Optimize software..."

import asyncio
from typing import AsyncGenerator

class StreamingLLM:
    async def generate_stream(self, prompt: str, cancel_token: asyncio.Event) -> AsyncGenerator[str, None]:
        # Simulate streaming generation
        for i in range(10):
            if cancel_token.is_set():
                break  # Stop generating if cancelled
            yield f"token_{i}"
            await asyncio.sleep(0.1)  # Simulate generation time

async def main():
    cancel_token = asyncio.Event()
    llm = StreamingLLM()
    # Start streaming in background
    async def consume():
        async for token in llm.generate_stream("hello", cancel_token):
            print(token)
    task = asyncio.create_task(consume())
    # Simulate user cancellation after 0.5 seconds
    await asyncio.sleep(0.5)
    cancel_token.set()
    await task

asyncio.run(main())

import asyncio
from typing import List, AsyncGenerator

class BatchedStreamingLLM:
    def __init__(self, max_batch_size: int = 8):
        self.max_batch_size = max_batch_size
        self.queue: List[dict] = []

    async def submit(self, prompt: str) -> AsyncGenerator[str, None]:
        # Create a queue for this request's tokens
        token_queue = asyncio.Queue()
        self.queue.append({'prompt': prompt, 'token_queue': token_queue})
        if len(self.queue) >= self.max_batch_size:
            asyncio.create_task(self._process_batch())
        # Yield tokens as they arrive
        while True:
            token = await token_queue.get()
            if token is None:
                break
            yield token

    async def _process_batch(self):
        batch = self.queue[:self.max_batch_size]
        self.queue = self.queue[self.max_batch_size:]
        # Simulate batched generation with streaming
        # In reality, you'd call the LLM with a batch of prompts and stream tokens
        for i in range(10):  # Simulate 10 tokens
            for item in batch:
                await item['token_queue'].put(f"token_{i}")
            await asyncio.sleep(0.1)
        # Signal end of stream
        for item in batch:
            await item['token_queue'].put(None)

# Usage
# llm = BatchedStreamingLLM()
# async def consume(prompt):
#     async for token in llm.submit(prompt):
#         print(token)
# asyncio.run(asyncio.gather(consume("hello"), consume("world")))

import time
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

# Set up OpenTelemetry
tracer_provider = TracerProvider()
span_exporter = OTLPSpanExporter(endpoint="http://localhost:4318/v1/traces")
span_processor = BatchSpanProcessor(span_exporter)
tracer_provider.add_span_processor(span_processor)
trace.set_tracer_provider(tracer_provider)
tracer = trace.get_tracer(__name__)

class MonitoredLLM:
    def __init__(self, model_name: str):
        self.model_name = model_name

    def generate(self, prompt: str) -> str:
        with tracer.start_as_current_span("llm_generate") as span:
            span.set_attribute("model_name", self.model_name)
            span.set_attribute("prompt_length", len(prompt))
            start_time = time.time()
            # Simulate LLM call
            response = self._llm_call(prompt)
            latency = time.time() - start_time
            span.set_attribute("latency_seconds", latency)
            span.set_attribute("response_length", len(response))
            return response

    def _llm_call(self, prompt: str) -> str:
        # Placeholder
        time.sleep(0.5)
        return "response"

# Usage
# llm = MonitoredLLM(model_name="gpt-4")
# response = llm.generate("hello")
# print(response)

# io.thecodeforge.com/examples/phase_profile
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import time

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
prompt = "Explain quantum computing to a 5-year-old" * 100
inputs = tokenizer(prompt, return_tensors="pt")

# Measure prefill phase individually
start = time.perf_counter()
with torch.no_grad():
    outputs = model(**inputs, use_cache=True)
prefill_time = time.perf_counter() - start
print(f"Prefill TTFT: {prefill_time*1000:.1f}ms for {len(inputs[0]):,} tokens")

# Measure decode phase (generate 50 tokens)
input_ids = inputs["input_ids"]
start = time.perf_counter()
with torch.no_grad():
    for _ in range(50):
        outputs = model(input_ids, use_cache=True)
        next_token = outputs.logits[:, -1, :].argmax(dim=-1).unsqueeze(0)
        input_ids = torch.cat([input_ids, next_token], dim=-1)
decode_time = time.perf_counter() - start
print(f"Decode TPOT: {decode_time/50*1000:.1f}ms per token")

# io.thecodeforge.com/examples/flash_attention
import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load model once with standard attention
model_standard = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.float16,
    attn_implementation="eager"  # Standard attention
).cuda()

# Load with FlashAttention (same model, different backend)
model_flash = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2"  # FlashAttention
).cuda()

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
prompt = "Write a 2000 word essay on" + " deep learning" * 500
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Benchmark throughput
for name, model in [("Standard", model_standard), ("FlashAttention", model_flash)]:
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    start = time.perf_counter()
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=True
        )
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    tokens = outputs.shape[1] - inputs["input_ids"].shape[1]
    print(f"{name}: {elapsed:.2f}s for {tokens} tokens ({tokens/elapsed:.1f} tok/s)")