ML Systems Design: Feature Stores, Model Serving & Inference Pipelines

Production machine learning architectures face a fundamental tension between training time requirements and inference time constraints. During offline training, models require massive historical datasets spanning months or years to discover predictive signals. During real-time inference, models require the latest feature values computed and retrieved within a strict sub-5 millisecond budget.

When data engineering teams write separate pipelines for training (SQL/Spark batch scripts) and inference (Python/Java microservices), subtle mathematical discrepancies emerge—a phenomenon known as training-serving skew. For example, calculating a user's 30-day rolling purchase sum using UTC midnight boundaries in Spark vs. rolling sliding windows in Python causes production models to behave erratically on live traffic.

The Unified Feature Store Pattern

A Feature Store (such as Feast, Hopsworks, or Tecton) provides a centralized registry and dual-sink architecture:

1. Offline Columnar Store (Parquet, Delta Lake, Snowflake, BigQuery): Stores immutable historical event logs partitioned by timestamp. It enables Point-in-Time (AS-OF) Joins, ensuring that when generating training records for an event that occurred on May 12th at 14:02:11, the feature store joins only features that existed prior to that millisecond, completely preventing target leakage.
2. Online Low-Latency Store (Redis Cluster, Aerospike, DynamoDB): Holds only the latest feature vector per entity key (e.g., user_id:98472, merchant_id:441). Key-value lookups are executed in O(1) time over persistent connection pools in sub-3ms.
3. Streaming Ingestion Engines (Apache Flink): Flink processes incoming Kafka event streams, updates tumbling/sliding stateful window aggregations, and writes simultaneously to the online Redis store and the offline Delta Lake sink.

Exposing deep neural networks or gradient-boosted decision trees (XGBoost/LightGBM) over standard synchronous REST endpoints (such as naive Flask/Gunicorn workers) leads to disastrous GPU underutilization. Modern GPUs (such as NVIDIA H100 with 3.35 TFLOPS of FP32 or 2,000 TFLOPS of FP8 Tensor Core compute) are massive SIMD (Single Instruction, Multiple Data) parallel matrix engines. Processing requests with a batch size of 1 leaves >85% of GPU SMs (Streaming Multiprocessors) idle while incurring the full memory transfer overhead.

Dynamic Batching Mechanics

Inference engines like NVIDIA Triton Inference Server and TorchServe implement Dynamic Batching at the server gateway:

• When an inference request arrives, the server places the input tensor into a concurrent request queue.
• The server waits up to a configurable max delay window (e.g., max_queue_delay_microseconds = 2000 or 2ms) or until the batch reaches max_batch_size = 64.
• All queued requests are merged into a single contiguous batch tensor, transferred over PCIe Gen5 / NVLink to GPU VRAM, executed in parallel through the compiled model engine, and split back to individual client response sockets.
• Under peak traffic (10,000 QPS), GPU utilization surges from 12% to >90%, increasing system throughput by 8× to 15× with only a deterministic 2ms latency penalty.

import asyncio
import time
import numpy as np
import tritonclient.grpc.aio as grpcclient
from feast import FeatureStore

class RealTimeInferenceGateway:
    def __init__(self, feast_repo_path: str, triton_url: str = "localhost:8001"):
        self.store = FeatureStore(repo_path=feast_repo_path)
        self.triton_client = grpcclient.InferenceServerClient(url=triton_url)
        self.model_name = "ensemble_fraud_ranking_v3"

    async def predict_fraud_score(self, user_id: str, transaction_amount: float, merchant_id: str) -> dict:
        t0 = time.perf_counter()
        
        # Step 1: Sub-3ms parallel retrieval of entity features from Redis via Feast
        entity_keys = [{"user_id": user_id}, {"merchant_id": merchant_id}]
        features_to_fetch = [
            "user_features:user_avg_tx_30d",
            "user_features:failed_auth_count_1h",
            "merchant_features:merchant_risk_score",
            "merchant_features:merchant_chargeback_rate"
        ]
        
        # Non-blocking async fetch from online feature store
        feature_response = await asyncio.to_thread(
            self.store.get_online_features,
            features=features_to_fetch,
            entity_rows=entity_keys
        )
        online_dict = feature_response.to_dict()
        
        # Step 2: Assemble input feature tensor [1, 5]
        user_avg = online_dict["user_avg_tx_30d"][0] or 0.0
        failed_auth = float(online_dict["failed_auth_count_1h"][0] or 0)
        m_risk = online_dict["merchant_risk_score"][1] or 0.5
        m_cb = online_dict["merchant_chargeback_rate"][1] or 0.01
        
        input_data = np.array([[transaction_amount, user_avg, failed_auth, m_risk, m_cb]], dtype=np.float32)
        
        # Step 3: Prepare gRPC Triton Tensor payload
        inputs = [grpcclient.InferInput("FLOAT_INPUTS", input_data.shape, "FP32")]
        inputs[0].set_data_from_numpy(input_data)
        outputs = [grpcclient.InferRequestedOutput("PROBABILITIES")]
        
        # Step 4: Execute inference (Triton automatically batches this with concurrent requests)
        infer_response = await self.triton_client.infer(
            model_name=self.model_name,
            inputs=inputs,
            outputs=outputs,
            timeout=15.0 # 15ms hard latency SLA
        )
        
        probabilities = infer_response.as_numpy("PROBABILITIES")
        fraud_probability = float(probabilities[0][1])
        
        total_latency_ms = (time.perf_counter() - t0) * 1000.0
        return {
            "user_id": user_id,
            "fraud_probability": fraud_probability,
            "decision": "BLOCK" if fraud_probability > 0.85 else "ALLOW",
            "latency_ms": round(total_latency_ms, 2)
        }

Large Language Models (LLMs based on the Transformer decoder architecture, e.g., LLaMA-3, Mistral, GPT-4) exhibit distinct computational characteristics that break standard ML serving paradigms. LLM generation operates in two distinct phases:

1. Prefill Phase (Prompt Processing): Highly compute-bound. The input prompt is processed in parallel to generate the initial Key and Value matrices across all attention layers.
2. Decode Phase (Autoregressive Token Generation): Highly memory-bandwidth bound. The model generates one token at a time. For each new token, it must re-attend to all preceding tokens by reading their Key and Value vectors from GPU High Bandwidth Memory (HBM).

The KV-Cache Bottleneck & Memory Formula

To avoid recomputing keys and values for past tokens on every decoding step, inference engines store them in a KV-Cache. The memory consumption for a single request scales linearly with context length:

KV-Cache Memory (Bytes) = 2 × 2 × n_{layers} × n_{heads} × d_{head} × context\_length × precision\_bytes

For a LLaMA-3 70B model (80 layers, 64 attention heads with 8 KV heads in Grouped-Query Attention, d_{head} = 128, FP16 precision = 2 bytes):

Per-Token KV-Cache = 2 × 2 × 80 × 8 × 128 × 2 = 655,360 Bytes ≈ 0.655 MB per token

For a context length of 4,096 tokens, a single user session consumes 2.68 GB of GPU VRAM solely for the KV-Cache!

PagedAttention & Continuous Batching (vLLM Engine)

Legacy serving runtimes allocated contiguous virtual memory buffers for the worst-case context length (e.g., 8,192 tokens), causing 60-80% memory fragmentation and out-of-memory crashes.

• PagedAttention (vLLM): Adopts operating system virtual memory paging concepts. The KV-cache is partitioned into fixed-size physical blocks (e.g., 16 tokens per block). Blocks are dynamically allocated on-demand in non-contiguous VRAM locations and mapped via a page table.
• Continuous / Iteration-Level Batching: Instead of waiting for an entire batch to complete all generation steps (which wastes compute when short requests finish early while long requests keep running), continuous batching evicts completed requests after every decoding step and immediately inserts newly arrived prompts into the active iteration batch.

Unlike traditional software services that crash with HTTP 500 errors or panic logs upon code regression, machine learning models fail silently. The model continues responding with HTTP 200 OK and valid JSON floats, but its underlying real-world accuracy degrades as world dynamics change.

Types of ML Degradation:

1. Data Drift (Covariate Shift): The input distribution P(X) changes while the underlying conditional probability P(Y|X) remains constant. For example, during Black Friday, e-commerce transaction volumes and amounts spike drastically compared to baseline training data.
2. Concept Drift: The fundamental relationship between inputs and targets P(Y|X) changes. For example, a new cybercrime fraud technique uses legitimate-looking login patterns that bypass existing feature weights.

Quantifying Drift Mathematically: Population Stability Index (PSI)

Systems continuously bucket live feature streams and compare their probability distributions P_{live} against reference baseline distributions P_{training}:

PSI = \sum_{i=1}^{B} ≤ft( P_{live}(i) - P_{train}(i) \right) × \ln≤ft( \frac{P_{live}(i)}{P_{train}(i)} \right)

• PSI < 0.1: Negligible shift; no intervention required.
• 0.1 ≤ PSI < 0.25: Moderate drift; schedule retraining job.
• PSI ≥ 0.25: Severe distribution divergence; raise P1 automated alert and trigger fallback heuristic or active learning pipeline.

Dark Launches (Shadowing) vs. Multi-Armed Bandits

To safely deploy model updates without risking catastrophic revenue loss:

• Shadow Deployments: The API gateway routes 100% of live traffic to Production Model A, but asynchronously clones input payloads to Candidate Model B in the background. Model B's output is logged to Kafka for offline statistical parity checks without returning its results to the user.
• Multi-Armed Bandits (Thompson Sampling / Upper Confidence Bound): Dynamically allocates 5% to 50% of real user traffic to candidate variants, automatically routing more traffic to variants that demonstrate statistically superior conversion or click-through rates.