All projects

Forward Deployed Engineer - Project Case Study · Jan – May 2025

TITANS: Long-Context Dialogue Memory

01

Engagement Snapshot

What I built, who it is for, and where my ownership starts and stops.

Memory-Augmented Transformer for Long Conversations

A memory-augmented transformer that lets attention read across unbounded history through a persistent memory module, without paying quadratic attention cost. A fused Triton kernel merges attention and memory read into a single pass.

+0%
Multi-turn recall
0%
Inference latency
0×
Throughput
0
Benchmarks (BoolQ, GSM8K)

Role
Research engineer: architecture, kernel, evaluation

Scope
Model architecture, Triton kernel, quantisation, serving integration, benchmarking

PyTorchTrTritonvLvLLMTransformersFastAPIPythonNVIDIA
02

Customer & Problem

Who this is for, the messy reality they work in, and why the system needed to exist.

Who it is for

Teams building conversational systems that need to remember long multi-turn exchanges without paying quadratic attention cost.

What was wrong before

A standard KV cache discards context beyond the window, which is the failure point in long multi-turn conversations. Extending the window makes attention cost quadratic.

Who needed it

Conversational AI systems that need to remember what was said many turns ago, especially for multi-hop reasoning across turns.

If nothing changed

Assistants forget earlier facts, or pay memory and latency costs that make long conversations unaffordable to serve.

The environment I had to work in

  • A standard KV cache loses context beyond the window.
  • Longer windows make attention cost grow quadratically.
  • Results only count if they hold up in a real serving stack.
03

Discovery & Requirements

The questions I ask first, and the requirements every decision below traces back to.

Discovery questions

  1. 01How long are real conversations?
  2. 02Which tasks need recall across turns?
  3. 03What latency and throughput targets does serving have?

Requirements

RecallRetain facts across long multi-turn exchanges
Attention costNo quadratic growth with history
LatencyLower than the baseline, not higher
ThroughputHigher than the baseline
ValidationMulti-hop reasoning tasks (BoolQ, GSM8K)
ServingRuns behind vLLM / FastAPI
04

System Architecture

The centerpiece. The diagram is the easy part - the value is in why each boundary exists.

Why this boundary exists · Tokens

Standard tokenised input keeps the model drop-in compatible with existing tooling.

Click any component to see why it is a separate boundary.

05

Integrations & Data

Where each category of information lives, what it connects to, and why.

06

Critical Architecture Decisions

Architecture decision records: the problem, the options, the call, and when to revisit it.

ADR-001 - How to carry long context

Options considered

Bigger windowRAG over historyPersistent memory moduleSummarisation

Decision

Persistent memory module.

Problem
Context beyond the window is discarded by a KV cache; enlarging the window is quadratic.
Reasoning
Lets attention read across unbounded history at fixed cost and stays inside the model, avoiding retrieval hops.
Trade-offs
Custom architecture and kernel to maintain; not a drop-in checkpoint.
Revisit when
If off-the-shelf long-context models match recall at comparable cost.
07

Deployment & Handoff

How it gets installed and shipped, what it runs on, and how it is handed over so the customer can run it.

step 1 / 5

Git.

FastAPI
vLvLLM
GPU

Rollout and handoff

Served through vLLM and FastAPI

The model is evaluated in a realistic serving stack, not only in a notebook.

Quantised for deployment

Lower-precision weights reduce memory and latency at serving time.

08

Reliability Engineering

Failure thinking: what breaks, what catches it, and what actually went wrong.

Click a stage to knock it offline, then send a request.

Served by Fused kernel

numerical parity checks vs referencefallback to unfused pathseeded evaluation
09

Security Architecture

Every request passes these gates in order.

step 1 / 4

Authenticate.

Research scope

Serving is behind FastAPI with standard input validation; the project is model research, not a multi-tenant platform.

10

Results

Only outcomes that can be substantiated. Adoption figures are added when they are real.

Performance

+0%
Multi-turn recall
0×
Throughput

Latency

0%
Inference latency
11

Field Lessons

What this project taught me about shipping AI into someone else's environment.

L1

Measure where it will run

Latency and throughput are only meaningful in the serving engine you will deploy.

L2

Fuse what shares data

Merging attention and memory read into one pass removed duplicated I/O.

12

Request / Execution Flow

One real request traced through the system - dynamic behaviour, not just boxes.

step 1 / 6

Turn arrives.New user message.

13

Deep-Dive Architecture

The technically difficult pieces, step by step.

step 1 / 5

Current window.Recent tokens attend locally.

14

AI Architecture

Model strategy, retrieval, and evaluation - the parts specific to an AI system.

Model strategy

Baseline

Standard transformer with KV cache: the comparison point.

TITANS variant

Memory-augmented model with fused kernel.

Tr

Quantised serving

Lower-precision weights for latency and throughput.

vL

Evaluation

Baseline run
TITANS run
BoolQ + GSM8K
Multi-turn recall
Latency + throughput
Compare
15

Observability

The observability model, not a tool name.

Generation request
Tokenise
TrFused attention + memory
Sample + stream

Illustrative span layout - shows nesting, not measured durations.

multi-turn recalllatencythroughputmemory footprint
16

Performance Engineering

Where the bottlenecks were, and what moved the numbers.

Multi-turn recall (index)baseline
Inference latency (index)baseline
Throughput (index)baseline · 1×

What caused the improvement

  • Fusing attention and memory read into one Triton pass avoids double I/O.
  • Persistent memory replaces ever-growing KV context for old turns.
  • Quantisation reduces weight bandwidth.

Bars are relative to the baseline (=100) using the published percentages.

17

Architecture Principles

The architecture followed from these, not from a shopping list of technologies.

P1

Memory as a module

Long-term recall is an explicit component, not an ever-larger window.

P2

Fuse what shares data

If two ops read the same tensors, do them in one kernel.

P3

Validate on reasoning, not just perplexity

BoolQ and GSM8K require using remembered information across steps.

P4

Measure end to end

Latency and throughput are measured in a serving stack, not a notebook.

18

Alternatives Rejected

Understanding the solution space, not just the final implementation.

Attention cost grows quadratically with window length; memory cost here stays fixed.

19

What I Would Change Today

A self-critique of the architecture.

  • Evaluate on more purpose-built long-dialogue benchmarks than BoolQ and GSM8K.
  • Study memory saturation and eviction behaviour over very long conversations.
  • Compare against newer long-context and state-space models on identical hardware.
20

My Role

What I owned, co-designed, influenced, and implemented personally.

Owned

  • · Architecture
  • · Evaluation design

Co-designed

  • -

Influenced

  • -

Implemented personally

  • · Fused Triton kernel
  • · Quantisation
  • · Serving integration
  • · Benchmarks