Text Generation, Decoding & Serving at Scale

# Demonstrates how temperature, top-k, and top-p are set at the API edge.
# These are product controls: lowering temperature makes output deterministic,
# while top-p dynamically adjusts the sampling nucleus.
from transformers import pipeline

# Initialize a text-generation pipeline with a base model
generator = pipeline("text-generation", model="gpt2")

# Generate with explicit decoding controls
result = generator(
    "Explain retrieval-augmented generation in simple terms:",
    max_new_tokens=120,       # Hard cap on output length (cost + latency control)
    temperature=0.7,           # Moderate randomness - good for explanations
    top_k=40,                  # Hard ceiling: only top 40 tokens considered
    top_p=0.9,                 # Nucleus: include tokens until 90% cumulative prob
    do_sample=True             # Enable sampling (False = greedy decoding)
)

# Output the generated text
print(result[0]["generated_text"])

# Compare greedy decoding with beam search on the same prompt.
# Greedy: fast but can get stuck.
# Beam: explores multiple candidates but costs B times more.
from transformers import pipeline

generator = pipeline("text-generation", model="gpt2")

# Greedy decoding: always pick the top token
greedy_out = generator(
    "The key advantage of transformers is",
    max_new_tokens=50,
    do_sample=False,         # Disable sampling = greedy
    num_beams=1              # Single beam = greedy search
)

# Beam search: keep 4 candidate sequences alive
beam_out = generator(
    "The key advantage of transformers is",
    max_new_tokens=50,
    do_sample=False,         # Beam search is deterministic
    num_beams=4,             # 4 beams = 4 parallel candidates
    early_stopping=True      # Stop when all beams reach EOS
)

print("Greedy:", greedy_out[0]["generated_text"])
print("Beam:", beam_out[0]["generated_text"])

# Streaming generation using an OpenAI-compatible client.
# Each chunk arrives as a server-sent event; the client
# processes tokens incrementally instead of waiting.
from openai import OpenAI

client = OpenAI()

# Create a streaming completion request
stream = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "Explain KV caching."}],
    stream=True,               # Enable SSE streaming
    max_tokens=200             # Budget cap for cost control
)

# Process each token as it arrives
for chunk in stream:
    # Each chunk contains a delta with the new token
    delta = chunk.choices[0].delta.content
    if delta:
        print(delta, end="", flush=True)  # Flush for real-time display

# Demonstrates the throughput advantage of batching.
# In real serving, continuous batching handles this dynamically,
# but this shows the core principle: shared GPU work.
import time

def simulate_inference(n_requests, batch_size, per_token_ms=10, tokens=100):
    """Simulate sequential vs batched inference timing."""
    # Sequential: each request waits for the previous one
    sequential_ms = n_requests * tokens * per_token_ms

    # Batched: GPU processes batch_size requests in parallel
    # Batching adds ~20% overhead per step but shares across requests
    n_batches = (n_requests + batch_size - 1) // batch_size
    batched_ms = n_batches * tokens * per_token_ms * 1.2

    print(f"Sequential: {sequential_ms:.0f}ms")
    print(f"Batched (bs={batch_size}): {batched_ms:.0f}ms")
    print(f"Speedup: {sequential_ms/batched_ms:.1f}x")

simulate_inference(n_requests=8, batch_size=8)

# Calculate KV cache memory for a given model and sequence configuration.
# This helps engineers understand why memory (not compute) is the binding
# constraint for LLM serving concurrency.

def kv_cache_memory_gb(
    n_layers: int,        # Number of transformer layers
    n_heads: int,         # Number of attention heads (for KV)
    head_dim: int,        # Dimension per head
    seq_len: int,         # Total sequence length (prompt + output)
    batch_size: int = 1,  # Concurrent requests
    precision: int = 2   # Bytes per value (2=FP16, 1=INT8)
) -> float:
    # Factor of 2 for keys AND values
    total_bytes = 2 * n_layers * n_heads * head_dim * seq_len * batch_size * precision
    return total_bytes / (1024 ** 3)  # Convert to GB

# Example: Llama-70B style model
mem = kv_cache_memory_gb(
    n_layers=80, n_heads=8,   # GQA: 8 KV heads (not 64 query heads)
    head_dim=128, seq_len=4096,
    batch_size=32, precision=2  # FP16
)
print(f"KV cache for 32 concurrent 4K requests: {mem:.1f} GB")

# Load a model in 4-bit quantization using bitsandbytes.
# This reduces a 70B model from ~140 GB (FP16) to ~35 GB,
# making it fit on a single A100 80GB GPU.
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# Configure 4-bit quantization
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,                 # Enable 4-bit weight quantization
    bnb_4bit_quant_type="nf4",         # NormalFloat4: optimized for normal distributions
    bnb_4bit_compute_dtype="bfloat16", # Compute in BF16 for stability
    bnb_4bit_use_double_quant=True    # Quantize the quantization constants too
)

# Load model with quantization applied
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b-hf",
    quantization_config=quant_config,
    device_map="auto"                  # Auto-distribute across available GPUs
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-hf")
print(f"Model loaded. Memory: {model.get_memory_footprint() / 1e9:.1f} GB")

# Measure Time to First Token (TTFT) and throughput for a streaming request.
# In production, these metrics feed SLO dashboards and autoscaling rules.
import time
from openai import OpenAI

client = OpenAI()

start = time.perf_counter()
first_token_time = None
token_count = 0

# Stream a completion and measure timing
stream = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "List 10 LLM serving optimizations."}],
    stream=True, max_tokens=300
)

for chunk in stream:
    if chunk.choices[0].delta.content:
        if first_token_time is None:
            first_token_time = time.perf_counter()  # Capture TTFT
        token_count += 1

end = time.perf_counter()

# Report metrics
ttft = (first_token_time - start) * 1000  # ms
total = end - start                          # seconds
throughput = token_count / total              # tokens/sec
print(f"TTFT: {ttft:.0f}ms | Tokens: {token_count} | Throughput: {throughput:.1f} tok/s")

# Show how serving cost and memory scale with context length.
# This helps engineers justify retrieval over context-stuffing.

def context_cost_scaling(base_len=4096, target_lengths=None):
    """Compare relative cost at different context lengths."""
    if target_lengths is None:
        target_lengths = [4096, 8192, 16384, 32768, 65536, 131072]

    print(f"{'Context':>10} {'KV Memory':>12} {'Prefill':>10} {'Max Batch':>10}")
    print("-" * 46)

    for length in target_lengths:
        ratio = length / base_len
        # KV cache scales linearly
        kv_ratio = ratio
        # Prefill compute scales quadratically (standard attention)
        prefill_ratio = ratio ** 2
        # Max concurrent batch inversely proportional to KV size
        batch_ratio = 1.0 / ratio

        print(f"{length:>10,} {kv_ratio:>11.1f}x {prefill_ratio:>9.1f}x {batch_ratio:>9.2f}x")

context_cost_scaling()

# A simplified layered safety pipeline showing pre/post generation checks.
# In production, each layer is a separate service with its own SLO.

class SafetyPipeline:
    def __init__(self, input_classifier, output_classifier, pii_detector):
        self.input_clf = input_classifier    # Pre-generation safety check
        self.output_clf = output_classifier  # Post-generation content filter
        self.pii_det = pii_detector          # PII redaction layer

    def check_input(self, prompt: str) -> dict:
        """Screen input before generation begins."""
        result = self.input_clf.classify(prompt)
        if result.risk_score > 0.85:
            return {"action": "block", "reason": result.category}
        if result.risk_score > 0.6:
            return {"action": "escalate", "reason": result.category}
        return {"action": "allow"}

    def check_output(self, response: str) -> str:
        """Post-process output: moderate content and redact PII."""
        # Step 1: Content safety
        if self.output_clf.is_unsafe(response):
            return "I cannot provide that information."  # Refusal
        # Step 2: PII redaction
        response = self.pii_det.redact(response)
        return response

# Skeleton of a tiered LLM service with routing, safety, and observability.
# Each component would be a separate service in production.

class LLMService:
    def __init__(self, router, safety, models, logger):
        self.router = router       # Classifies request complexity
        self.safety = safety       # Layered safety pipeline
        self.models = models       # Dict of model tiers: {"fast": ..., "frontier": ...}
        self.logger = logger       # Structured logging + metrics

    async def handle_request(self, request):
        """Full request lifecycle: route -> assemble -> generate -> moderate."""

        # 1. Input safety check
        safety_result = self.safety.check_input(request.prompt)
        if safety_result["action"] == "block":
            return RefusalResponse(safety_result["reason"])

        # 2. Route to appropriate model tier
        tier = self.router.classify(request)  # "fast", "balanced", "frontier"
        model = self.models[tier]

        # 3. Assemble prompt (system prompt + retrieved context + user input)
        full_prompt = assemble_prompt(request, tier)

        # 4. Generate with streaming
        async for token in model.generate_stream(full_prompt):
            yield token  # Stream tokens to client

        # 5. Post-generation: log metrics, check output safety
        self.logger.log(tier=tier, tokens=token_count, latency_ms=elapsed)

        # 6. Output moderation (in production, runs in parallel with streaming)
        final_output = self.safety.check_output(full_response)