Production ML, responsible AI and alignment blog

Whose Memory Is It? Building Multi-Tenant, Multi-Tier Memory for AI Agents (Part 2)

Mon, 10 Aug 2026 00:00:00 GMT

Alice and bob are talking to the same agents. Alice interacts with various agents for infra management. Bob interacts with various agents for app development. Both are able to build on learnings from the last month. But how far should this go? Should Alice be able to recall memories from Bob’s interactions? Should a user allow a single agent to recall memories from across their agents?

This captures the design choices required in multi-tenancy for agentic memory management

Recently I spent some time extending the Kubernetes Agent Orchestration System (KAOS) to support multi-tiered memory persistence (aka short-, medium- and long-term memory). Along the way I hit most of the same issues that anyone would whilst building or integrating multi-tiered memory into a multi-tenant system, so I thought it would be useful to compile the learnings, design choices and examples into this series.

This is Part 2 of the series, and here I go through some of the design choices made for 3-tier multi-tenant memory. This follows Part 1, where we surveyed ~30 memory engines, built a working taxonomy, and landed on adopting Mem0 as a library behind our own interface, together with the list of gaps (observability, tenant isolation, kubernetes packaging, framework bridging) that become our integration work.

The objective throughout the series is:

Let’s make the memory layer BORING, so that the agents can continue to be the fun part.

This part consists of two sections:

Three memory tiers: Defining the memory adopted, which includes a short-term window memory, a medium-term rolling summary, and long-term semantic “facts”.
Scope model: A hierarchical multi-tenant read model scope, that spans across session > agent > user, defined by a verified identity and a maxReadScope ceiling

Finally we wrap up with five hard lessons we learned about tier and scope design that carry beyond KAOS.

As with my previous posts on observability for agentic systems and autonomous always-on agentic patterns, I use KAOS as the concrete implementation example, but the goal is to provide practical intuition for the primitives (tiers, scopes, folding, degradation), so that it applies whether you use KAOS, Mem0 directly, LangGraph, CrewAI, or a memory layer you wrote yourself.

Here’s a refresher on this 4-part series on Multi-Tiered / Multi-Tenant Agent Memory:

Part 1: What agent memory is and what to build on. The taxonomy, the baseline implementations everyone starts with, and the engine landscape from surveying ~30 tools.
Part 2 (this post): Tiers and scopes for multi-tenant agents. The three-tier design and the answer to whose memory it is.
Part 3: Memory as infrastructure. The Kubernetes MemoryStore resource, its deployment topology, and how to integrate it in your own agent (coming soon…).
Part 4: Agent memory in action. A worked example that runs end to end on a secured cluster, with real outputs (coming soon…).

Let’s get started.

Designing our Memory Architecture: The Three Tiers

As a reminder, the taxonomy defined in Part 1 consisted of five memory types: short-term (working), episodic, semantic, procedural and temporal.

However when aligning with our requirements, I settled with a simplified three tier model: a short-term window, a medium-term summary, and long-term “facts”. These are intuitively defined as follows:

Tier	What it holds	When it updates	Backing
Short-term	The context window of the live session, bounded by a token budget	Every turn (cheap append)	Relational rows
Medium-term	Rolling summary per session, versioned so past summaries stay accessible	On compaction, when the window hits its token budget	Relational rows, append-only
Long-term	Atomic facts extracted from context window, keyed by scope, recalled semantically	In the background, after compaction	Mem0 into the vector store

Now that these tiers are defined, it was possible to also formalise the following design decisions:

Long-term memory functionality is enabled via Mem0; short- and medium-term memory are built custom.
These three tiers should cohesively integrate as a single interoperable unit.
Medium- and long-term extraction is lossy; we could enable provenance, however this adds significant complexity so I decided to keep this out of scope for now.
Medium- and long-term extraction are always off the write path; it triggers when compaction threshold is crossed as opposed to in every insert, which is also how Mem0’s own platform behaves.
The medium-term summary stays out of the vector store: Mem0 wants atomic, individually revisable facts, whereas a summary is a narrative whose whole value is its continuity.
Underneath all three tiers, the raw turns are the source of truth and everything else (summaries, facts, embeddings) is a recomputable projection, which is also what makes lossy extraction and fire-and-forget background processing acceptable.
Temporal (bi-temporal validity) and procedural (aka skill persistence) memory are deliberately deferred in their explicit form, but achievable through the long-term memory.

These definitions also allow us to design the single coherent service that offers the short-, medium- and long-term memory tiers; the “MemoryStore Service”.

I will cover more on the MemoryStore service in Part 3 where we actually design the infrastructure components for Kubernetes. Before we get there however, we need to talk about another important (+ tricky) topic:

Access Scopes: or who should be able to remember what?

Access Scopes: Whose Memory Is It Anyway?

Every memory operation in a multi-tenant system needs an answer to “whose memory is it?”. And the answer has to come from the design of the system components.

We first have to start on the write path before we can define a solution for the read path access.

For the write, the constraint introduced was the following:

A single conversation is authored through an agent, on behalf of a user (or autonomous agent), inside of a session, on one memory store.

This means the service records all of these as metadata provenance for every memory input stored. Writes are therefore compound and invariant, while a read resolves to a single scope level and is a matter of policy.

That separation is what lets a write be recalled at several levels later without being duplicated. This is possible because the same memory data that an agent stored for “Alice” carries: 1) her user_id, 2) the agent identity, and 3) the session identifier.

This allows recalls at different levels to each find it through a different owner key.

For reads, we decided to use a hierarchical relationship across these three levels, where each wider level contains the previous one:

The access scope is what restricts read-level access, as it can be understood intuitively in the graph above.

This gives each level one documented meaning under each security posture:

Level	Meaning	With user auth on	With user auth off
`session`	the current conversation	current agent x user x session	current agent x session
`agent`	this agent’s memory of the verified context	current agent x user	this agent’s whole pool
`user`	the verified user across all agents on the store	current user, across agents	rejected at deploy time
`store`	the entire memory store for the agents using it	all users x agents x sessions in that store	all agents x sessions

The identity always comes from the gateway-verified request headers, never from the request body or the model, which means that when auth is enabled, memory access is enforced by design at the control and data plane level.

The only way to access the store level (or the agent-wide pool across users) is through cluster-admin permissions: the memory service refuses these levels on any request arriving through an agent, so no grant, tool, or prompt can reach them, and the only remaining path is kubectl port-forward, which Kubernetes RBAC gates. If auth is disabled, the level meanings still hold logically, however there is no verified identity to enforce them against.

Because the read levels are totally ordered as a hierarchy, we are also able to define them with a maxReadScope that allows for a threshold definition (for example “max agent level and below”).

This also allows the MemoryStore to carry its own maxReadScope ceiling (default is agent), and an agent may not claim above its store’s, so cross-agent user reads exist only where the store owner deliberately raised the ceiling.

As covered in my previous post on autonomous always-on agents, KAOS supports agents that run without user input. I was able to also cover these instances by ensuring that a self-initiated iteration runs with the agent’s own identity as its user identifier. This means that all requirements for memory scopes are still satisfied, and an autonomous loop’s memory stays private to the loop, even if a user also queries that same autonomous agent.

This scope model is probably the obvious choice; the trickier question is how do we enforce the sharing restrictions at a project-like scope.

There were a few design options for this:

Many groups inside one MemoryStore. One store holds the memories of several groups at once. This sounds efficient, however it means building and operating a whole group-management layer. This would involve an API to create and delete groups and to add and remove members, per-group quotas, and a single store whose failure affects every group in it.
One group per MemoryStore. The store itself is the group: whichever agents are bound to the same store share it, so membership is just the existing binding and no new API is needed. The cost is that every group needs its own store deployment, and sharing across two groups means binding to a second store.
Hierarchical scope paths. A richer model where scopes are nested paths (for example org:team:agent) and agents share memory up to the point where their paths diverge. Every version of this I drafted ended up re-creating an authorization system that the two simpler options already covered.

Interestingly enough, when looking at how the managed platforms handle this, they expose a two-level version of the same tradeoff. Each one has a hard container that their control plane creates and manages, and lighter logical partitions inside it. A few examples:

Mem0 platform: A project is the container that memories cannot cross, and the user and agent partitions live within it, with API keys scoped to the project.
Vertex Memory Bank: Provisions one Memory Bank per Agent Engine instance, and within it memories are partitioned by scope, with retrieval only returning memories whose scope exactly matches the request.
Zep Cloud: Each subject (a user, or a group via their group-graph API) gets its own isolated context graph, and the cloud platform is the control plane that manages millions of them.

Based on these tradeoffs, I went for option (2), one group per MemoryStore as it enforces this at the control plane. The store itself is the sharing boundary, which is exactly why the whole-store read scope is named store. This meant that I don’t have to build a full intra-store group management layer, and the data layer simply records the store’s group key as internal metadata on each record.

The way it’s designed to is set up to support finer grouping at the MemoryStore level by design, as we basically are storing everything under one global group per store.

Now finally once I adopted these design choices, I realised that there were a few caveats that came up, which I had to accept / address:

Security Attack Surfaces: Interesting research such as AgentPoison show the impact of poisoning memory (ie 0.1% poisoned memory yields over 80% attack success), as well as MINJA which shows that an attacker needs no write access at all, because if the agent writes its own memory from conversations then every user is a write path. To mitigate this, it was decided for KAOS to derive the scope server-side from the authenticated agent identity, fail-closed, and never from model- or tool-supplied arguments.
Depth of intra-store isolation: Within one store, the boundaries between agents, users and sessions are enforced by application-level filtering. Application-level predicates carry a classic risk, where one forgotten WHERE clause silently returns another tenant’s rows. To mitigate this, the filtering is centralised in a single storage module so there is one place to audit, and on Postgres the relational tiers can be hardened further with row-level security, where a FORCE ROW LEVEL SECURITY policy pins every query to the scope set on the transaction, so a missing predicate returns nothing instead of everything. Worth noting this RLS cannot be enforced at the vector level - however isolation between stores relies on none of this, since each store is its own deployment with its own database connection, which we cover in Part 3.
Right to Erasure: Compliance requirements such as GDPR mean you must be able to answer “delete everything you know about this user” reliably, and in a multi-tier design the same information lives in several derived forms at once (raw turns, summaries, extracted facts, and their embeddings), so deleting from one tier is not enough. To mitigate this, KAOS implements forget as a single operation that fans out across all three tiers in one pass, deleting the short-term rows, the summaries, and the scope-filtered long-term facts. Note this is destruction, which is different from supersession, where facts are merely marked invalid but kept for history.

Now that we have sorted the tiers and the access scopes, let’s distil the lessons from this part before we make it all run as infrastructure in part 3.

Lessons for Production Agentic Memory

Here are the patterns from this part that I would carry into any agentic memory system.

1. Separate conversational continuity from learned knowledge

Same-session verbatim windows and cross-session distilled facts are different memory tiers with different stores, lifecycles, and failure modes. Conflating them for any reason would add more complexity than simplification.

2. Raw conversations are the source of truth

Summaries, facts, and embeddings are lossy, but recomputable. Keep the verbatim record durable and you can survive both a lost extraction and a change of mind about your extraction strategy.

3. Keep rolling summaries out of the vector store

We use long-term memory stores for atomic facts, whereas the medium-term memory is built with a rolling summary that provides continuity. Summaries and windows can be stored relationally and deterministically.

4. The control plane should enforce the memory scope

Derive scope server-side from authenticated identity. When the model is allowed to search, bound the levels it can reach with a maxReadScope ceiling. Treat what comes back as untrusted data with provenance, since memory poisoning and cross-session injection are demonstrated attacks with published success rates.

5. The store is the group (and vice-versa)

Sharing topology can be a deployment choice instead of an authorization system, with scope filtering within a store and physical isolation by deploying a store per tenant. This may seem restrictive, but it’s the model that most production platforms (+ cloud providers) follow.

Closing Thoughts for Part 2

We opened Part 2 with Alice and Bob talking to the same agents, and with the question of how should their memories be available and accessable. After this initial design, we can now answer it precisely.

Should Alice recall memories from Bob’s interactions? Never through an agent, because every recall is bound to the verified identity on the request.

Should a single agent reach across a user’s other agents? Only when the user level sits within its maxReadScope ceiling, which both the agent and the store owner have to allow.

Everything wider than that belongs to the cluster admin, behind Kubernetes RBAC.

That is the conceptual core of the series: three tiers that decide what an agent remembers, and a scope model that decides who it remembers it for.

In Part 3 we turn this design into running infrastructure with the MemoryStore Kubernetes resource, the topology decision behind it, and the degradation contract that keeps a memory outage from taking an agent down, together with how you can integrate the same pattern in your own agent from scratch.

Stay tuned for next week!

The series:

Part 1: What agent memory is and what to build on. The taxonomy, the baseline implementations everyone starts with, and the engine landscape from surveying ~30 tools.
Part 2 (this post): Tiers and scopes for multi-tenant agents. The three-tier design and the answer to whose memory it is.
Part 3: Memory as infrastructure. The Kubernetes MemoryStore resource, its deployment topology, and how to integrate it in your own agent (coming soon…).
Part 4: Agent memory in action. A worked example that runs end to end on a secured cluster, with real outputs (coming soon…).

Whose Memory Is It? Building Multi-Tenant, Multi-Tier Memory for AI Agents (Part 1)

Sat, 08 Aug 2026 00:00:00 GMT

LLMs are stateless by design, and without added memory logic every session starts from zero.

A number of dedicated memory layers have emerged (and continue emerging almost daily) to tackle this, each with different approaches and tradeoffs. Which one should you adopt?

Recently I spent some time extending the Kubernetes Agent Orchestration System (KAOS) to support multi-tiered memory persistence (aka short-, medium- and long-term memory). Along the way I hit most of the same issues that anyone would whilst building or integrating multi-tiered memory into an agentic system, so I thought it would be useful to compile all the learnings, design choices and examples into this 4-part series.

Hopefully this post is useful for anyone looking to do this on their own project. My objective:

Let’s make the memory layer BORING, so that the agents can continue to be the fun part.

This first part includes the research findings from exploring ~30 tools, together with the taxonomy and the baseline implementations that motivate needing them, including tools like Mem0, Zep/Graphiti, Letta (MemGPT), Cognee, Memobase, Redis Agent Memory Server, as well as native implementations in OpenAI’s products, Claude, LangGraph, CrewAI, and Google ADK, among many others.

Throughout the series, I also share the learnings and best practices that came out of navigating through a large number of architecture tradeoffs, and getting my hands dirty on the implementation that now ships as a distributed, highly available, and scalable MemoryStore resource that any agent can bind to.

As with my previous posts on observability for agentic systems and autonomous always-on agentic patterns, I will use KAOS as the concrete implementation example (concretely from Part 2, hands-on in Part 4), but the goal is to provide practical intuition for the primitives (tiers, scopes, folding, degradation), so that it applies whether you use KAOS, Mem0 directly, LangGraph, CrewAI, or a memory layer you wrote yourself.

This post is the first of a 4-part series, which we plan to release weekly over the next couple of weeks:

Part 1 (this post): What agent memory is and what to build on. The taxonomy, the baseline implementations everyone starts with, and the engine landscape from surveying ~30 tools.
Part 2: Tiers and scopes for multi-tenant agents. The three-tier design and the answer to whose memory it is (coming soon…).
Part 3: Memory as infrastructure. The Kubernetes MemoryStore resource, its deployment topology, and how to integrate it in your own agent (coming soon…).
Part 4: Agent memory in action. A worked example that runs end to end on a secured cluster, with real outputs (coming soon…).

A Working Taxonomy of Agent Memory

“Memory” is one of the most overloaded words in agentic systems, so it is worth separating it from the concepts it gets conflated with, such as:

The context window, which holds working state for a single model call.
Session history, which holds an auditable transcript of what was said.
Prompt telemetry, which holds the specific prompts relative to events in the system.

To be more precise we can look at Princeton University’s paper on Cognitive Architectures for Language Agents (CoALA) to provide a more precise definition for “Memory” in agentic systems. We can define “Memory” as the component that holds the short-, medium- and long-term information an agent carries across turns and sessions to inform its reasoning.

This research paper quoted also provides a useful taxonomy for “memory types” that we will use to reason throughout the series, and especially in the tier design of part 2. This includes the memory types for episodic, semantic, procedural and temporal memory.

These memory types are also mentioned in the Berkeley paper that released MemGPT, as well as how the Stanford paper on large-scale LLM simulations Generative Agents: Interactive Simulacra of Human Behavior structured their memory event stream.

The formal definition of these memory types (+ a few examples) is outlined as follows:

Memory type	What it holds	Example
Short-term (working)	Verbatim recent turns of the live conversation	“The user just said port 8080”
Episodic	Records of specific past events	“On Tuesday the deploy failed twice”
Semantic	Distilled, durable facts	“The user prefers blue-green deploys”
Procedural	Learned skills and how-tos	“Here is how we roll back this service”
Temporal	Facts with validity intervals	“Joe was in a relationship until March, but not anymore.”

In practice what I found out however is that most frameworks only implement a small number of these, namely short-term is always present, episodic and semantic are bundled (the only difference is whether time is preserved), procedural tends to be present mainly in coding agents (eg creating skills, commands, extensions), and temporal tends to be replaced with “forgetting memory” functionality instead, or embedded with episodic/semantic.

These appear more informally defined as:

Conversational continuity: The agent remembers what was said three turns ago; a same-session problem.
Learned knowledge: The agent remembers what it figured out last week; a cross-session problem.

For example, frameworks like LangGraph separate thread-scoped checkpointers from a cross-thread store. Another example is Letta, which separates always-in-context memory blocks from an archival tier.

Most of the design mistakes I made early came from either trying to tackle all of these “memory-types” separately, by bundling sub-optimally, or by oversimplifying too much.

But before we dive into the implementation, let’s cover the basics.

Memory 101: The Version Everyone Starts With

Almost every agent system starts with the same memory implementation:

memory = []

async def handle_message(user_message):
    memory.append({"role": "user", "content": user_message})
    response = await run_agent(memory[-20:], tools)
    memory.append({"role": "assistant", "content": response})
    return response

And to be honest, the original KAOS memory was exactly this. It was an in-process queue with a max length, which ensured it was replaying the last N events into the next prompt.

The second version everyone builds is “just embed everything”:

async def handle_message(user_message):
    hits = await vector_store.search(embed(user_message), top_k=5)
    context = "\n".join(h.text for h in hits)
    response = await run_agent([context, user_message], tools)
    await vector_store.add(embed(user_message), user_message)
    return response

This is better, but this is not memory in the form that we introduced eariler, it is just a better search mechanism across the prompt history.

Another tempting alternative as the next step is “context windows are huge now, just replay everything”.

However this is not a great approach, and there are some benchmarks like UCLA’s Bench on Long-Term Interactive Memory, which showed that models reasoning over full ~115K-token interaction histories lose 30-60% accuracy versus the same models given oracle retrieval.

If we look at it from a feature / functionality standpoint, we can summarise the gaps between the base and the production implementation as follows:

Naive memory	Production memory
Last-N turns, unbounded token growth	Token-budgeted window with principled eviction
Verbatim replay of everything	Distilled facts, separated from the transcript (eg long- / short-term)
One user, one process	Many tenants, many agents, many replicas
Memory lives inside the agent pod	Memory survives restarts and is shared across the fleet
Writes block the response	Extraction runs off the hot path
Nothing is ever forgotten	Decay, retention, and right-to-erasure
Memory failure crashes the turn	Memory failure degrades the turn

In this case we can position “production memory” a tiered, scoped, context-specific and dynamic store, as opposed to purely a vector database connected to an agent.

Achieving this in a way that scales does get complex, as we need to decide who can see each memory tier, when we store facts, and how the agent behaves when memory fails.

However now that we have the conceptual foundation in place, we can start looking at these functionalities relative to the frameworks available.

Choosing an Engine: Build, Adopt, or Wrap

Before designing anything, I surveyed the landscape thoroughly, assessing dozens of tools across three tiers, and we will cover the scope, approach and learnings in this section, starting with an overview of all the tiers as follows.

Tier 1: Dedicated memory frameworks. This tier encompasses purpose-built frameworks whose whole job is agent memory. From the longer list, we reduced it to the actively maintained ones:

Candidate	Approach	Store	Strength	Weakness
Mem0	extracts facts from conversations into a vector store and recalls them by similarity	Qdrant, pgvector, others	most adopted, cleanest library integration	no OTel, tenant isolation only at application level
Zep / Graphiti	builds a temporal knowledge graph where facts carry validity intervals and provenance	Neo4j or FalkorDB	richest memory model, time-aware fact invalidation	heaviest to operate, costliest writes
Cognee	combines a knowledge graph with vector search, populated by an extract-and-load pipeline	LanceDB by default, Postgres or Neo4j optional	multi-tenancy and OTel built in	early stage, heavy dependencies, changing API
Memobase	maintains structured user profiles and event timelines, with no embeddings on the hot path	Postgres + Redis	cheapest write path	profile-only recall, weak self-hosted multi-tenancy
Redis Agent Memory Server	serves two memory tiers (working and long-term) behind one REST API	Redis	the two-tier model mirrors what agents actually need	young project, no OTel

What this tier taught me is that the architectural differences are really differences in recall pattern and write cost. Vector-first designs answer “what do we know about X” cheaply, graph-first designs answer “how did this fact change over time” at the price of an LLM-heavy ingestion pipeline plus a graph database. The profile-first designs answer “who is this user” with no embeddings on the hot path at all, and the two-tier designs bake in the working-versus-long-term split directly.

There were also clear shared gaps, mainly at the infrastructure level; none of them enforces tenant isolation below the application level, and almost none ships OpenTelemetry, so whichever one you pick, scope enforcement and observability become your integration work. That shared gap shaped the KAOS design more than any individual feature did.

It’s also worth noting that several of these libraries also offer an enterprise tier, so it was important to validate that basic features are not gated behind a paywall (similar to what we previously experienced with Google ADK and Vertex). More specifically Mem0’s own platform-versus-OSS documentation gates temporal reasoning, memory decay, webhooks, export, analytics, and auto-scaling behind the managed platform, and Zep draws the line at governed multi-tenancy and compliance, with the OSS Graphiti engine giving you a single context graph to run yourself. The pattern across vendors is that the memory algorithms are open while the operational maturity is the commercial product, which previews the exact layer a platform adopting one of these engines has to build.

Tier 2. Agent frameworks with native memory. This tier encompasses the embedded memory functionality across end-to-end agentic frameworks, and included the usual suspect / popular frameworks like LangGraph’s Store and LangMem, CrewAI memory, LlamaIndex memory, Google ADK’s MemoryService, and the Microsoft Agent Framework. These were reviewed for their design choices, but adopting one for its memory means importing a second agent runtime next to your own, so they served as references and not as candidates.

The learning from this tier is actually what they all had in common. Every framework independently separates session-scoped state and cross-session knowledge, such as how LangGraph has thread-scoped checkpointers versus its cross-thread Store.

There was also a clear separation between “local playground” and “production grade” when it comes to memory for all frameworks.

LangGraph: In-memory store is for development, but it’s recommended to use a database-backed checkpointer and store for production.
Google ADK: Heavier paywall, as it only offers the InMemoryMemoryService as open source, but anything serious would need to use Vertex AI.
Microsoft Agent Framework: Defaults to an in-memory context provider, with durable state via first-party Cosmos checkpoint storage, and ships a first-party Mem0ContextProvider.
CrewAI: Community documented replacing its native store with Mem0 after hitting redeploy and user-isolation gaps.

Also interesting learnings from agent runtimes and coding agents:

OpenClaw: layers markdown memory files (MEMORY.md, dated notes, per-skill SKILL.md) and runs a “Skill Workshop” where the agent proposes new skills from successful conversations for human approval.
Hermes agent: Uses skills explicitly as procedural memory, which are auto-proposed after repeated successful tool-call patterns, carry semver versions bumped on each self-improvement, and follow an anti-sediment principle where a skill should get shorter and sharper over time.
Claude Code skills: Uses progressive disclosure, loading only ~100 tokens of skill metadata until a skill is triggered.

Native memory is increasingly an extensible interface where the shipped default is a placeholder, which means memory is being externalized by design across the ecosystem, and the dominant production pattern is framework plus engine. This is something that we had to take into consideration as well.

Tier 3. Managed and commercial services. This tier included commercial services with managed memory platforms, which provided insights on the broader design of the system and the interactions with the memory, as opposed to just the design of the memory capability itself. These included the Mem0 Platform, Zep Cloud, Letta Cloud, OpenAI memory, and Google’s Vertex Memory Bank.

The learnings from this tier were also quite helpful to understand some of the architectural and feature tradeoffs that were done at the platform level. Every managed platform has the same two-layer model: namely 1) an explicit, user-curated layer (eg. OpenAI’s saved memories, Claude’s editable memory summary) that is visible/available at the surface, and that is build on top of; 2) an automatically inferred and consolidated layer (eg. OpenAI’s chat-history reference, Vertex Memory Bank’s LLM extraction with per-scope deduplication and contradiction checks) where the memory store/retrieval algorithms live.

There was however a clear distinction on the scope in which memory is available across each platform: For Claude, memory scope is per project, in Vertex scope is per identity with configurable memory “topics”, and Zep scopes per subject graph. None of them defaults to one global memory per account, which makes it clear that there is a design decision required on the isolation boundary.

Conclusions from Surveying the Ecosystem

Given the review was done in context of KAOS, the lens / considerations through which these were reviewed included the following non-exhaustive list:

Long-term capability coverage
Retrieval quality
Embeddability as a library
Pluggable storage backends
Infrastructure delta / overhead
Multi-tenancy hooks
Observability
Licensing
Maturity
Write-path cost

Based on these, the library that clearly stood out was Mem0. At least at the time of writing, Mem0 maximized the features and capabilities with the lowest integration friction. Mem0 also has the strongest ecosystem maturity, and pluggable stores.

It’s however worth noting that despite Mem0 being the right choice for this context, one learning that may seem obvious in retrospect was that there is no “Perfect Candidate”. The graph-first leaders (Graphiti, Cognee) have the most features but at the highest cost. Low-delta options (Redis AMS) buy fit at maturity cost. Building it yourself directly on the raw vector or graph stores, which we also weighed as the baseline option, allows you to have all the features and fit, but at the cost of rebuilding mature extraction and retrieval that already exists under permissive licenses.

This also applies to the numbers the frameworks publish about themselves. Interestingly enough, Mem0’s own research supports that extraction-based memory improves latency and cost, however it does not improve raw accuracy: in Mem0’s own evaluation on LoCoMo, a full-context baseline beats Mem0 on raw accuracy (72.9% vs 66.9%), while memory buys a 91% cut in p95 latency (1.44s vs 17.1s) and over 90% fewer tokens per conversation. At fleet scale that trade is exactly right, since you cannot ship 17-second turns and 26K-token replays, but it is a trade you should make knowingly.

As part of this, despite Mem0 being the strongest choice, it became clear that adopting a memory engine means choosing which 60% of the system you do not have to build, and committing to build the remaining 40% around it.

For Mem0, this meant working on the bridge to close some of the gaps, particularly at the infrastructure and interoperability layer. These included:

Enabling telemetry by instrumenting every operation and ensure correlation+consistency with the broader KAOS telemetry.
Introduce tenant isolation, as this is enforced at the Mem0 application level, so enforce scope through the memory service.
Bundle up the kubernetes packaging to ensure high availability and scalability as a distributed service.
Bridge the short- and medium-term memory with a native integration with the Pydantic AI server that we have built as part of KAOS.

Each of these gaps becomes a design decision in parts 2 and 3.

Based on these initial decisions we were able to proceed to the architecture of the memory system itself, which is exactly where part 2 picks up.

Closing Thoughts for Part 1

This first part covered the ground you need before writing any memory code with a working taxonomy that separates memory from the context window and the session transcript. We also reviewed the baseline implementations and the understood some of the limitations. This included the survey of the memory engine landscape, and the thinking process that went into selecting the memory framework to build upon.

However this is only the beginning as choosing the memory framework is only 40% of the work, we still need to build the remaining 60% to ensure we can integrate it in a coherent and scalable way for our distributed agent system use-case.

In part 2 we take that position and design the memory system itself. This includes the three tiers that separate “memory tiers”, and the scope model that answers the title’s question of “whose memory is it?”. Stay tuned, part 2

The series:

Part 1 (this post): What agent memory is and what to build on.
Part 2: Tiers and scopes for multi-tenant agents (coming soon…).
Part 3: Memory as infrastructure (coming soon…).
Part 4: Agent memory in action (coming soon…).

Autonomous Agentic Systems: A Practical Guide to Always-On Agents

Mon, 01 Jun 2026 00:00:00 GMT

The pattern of “always-on” autonomous agents that work proactively in the background has been taking over the agentic systems space.

This pattern was launched to massive popularity by OpenClaw, as it basically moved the field away from purely Chat-based interactions (aka request-response), and towards agents that proactively take over tasks autonomously in the background (aka an infinite loop).

This is an interesting new architectural innovation which has set a clear path where agentic systems are heading towards:

24/7 agents with a heartbeat-style loop that keeps checking the world even when you are not prompting it.

And although OpenClaw was one of the first movers, the rest of the frameworks have also followed (+ catching up): LangGraph, CrewAI, OpenAI Agents SDK, Google ADK…

…and as part of this trend, we also ventured into the world of autonomous agentic systems at scale, learning from implementing autonomous agentic loops in Kubernetes in KAOS - this led us through the good, the bad and the ugly of autnomous agent patterns.

In this post we want to share the learnings designing, developing and scaling autonomous agentic patterns, including what actually changes when an agentic loop becomes an autonomous workload, and why Kubernetes starts to become relevant once you need to orchestrate many of them.

The Useful Part of the Hype

The phrase “autonomous agent” is overloaded enough to be almost useless.

Anthropic’s guidance on building effective agents makes a useful distinction between workflows with predefined code paths and agents where the LLM dynamically directs its own process and tool usage.

For a broader research view, recent surveys of LLM-based autonomous agents and large language model based agents cover the recurring architecture pieces: memory, planning, action, tools, and evaluation.

Sometimes it may mean “a chatbot that keeps working in the background”. Sometimes it may mean “a stateful graph workflow”. And sometimes it just means “we put while True around the model call”.

But there is a useful distinction here - a normal tool-using agent is usually request/response:

User asks a question.
Agent calls the model.
Model decides whether to use tools.
Tools return data.
Agent returns an answer.

An autonomous agent in the context of this post, ignores whether a human is waiting on the other side of the HTTP response.

The work may continue; the environment may change; the agent may run again; it may call tools repeatedly; it may need to remember what happened last time.

The hard part is not making the model call itself again (aka the loop).

The hard part is making that agentic loop safe to run when nobody is staring at the chat.

AI Agents 101: The Loop Everyone Starts With

Most agent systems begin with a deceptively simple loop:

async def run_agent(messages, tools, max_steps=5):
    for step in range(max_steps):
        response = await model.chat(messages, tools=tools)

        if response.tool_calls:
            for call in response.tool_calls:
                result = await tools[call.name](**call.arguments)
                messages.append({
                    "role": "tool",
                    "tool_call_id": call.id,
                    "content": result,
                })
            continue

        return response.text

    raise RuntimeError("agent exceeded max_steps")

This is the core pattern behind a lot of the current wave of agentic software, and it maps closely to the reasoning/action loop described in ReAct and the broader tool-use direction represented by Toolformer.

This loop is powerful, but it is still usually bounded by a synchronous request/response paradigm.

What Changes When the Loop Keeps Running?

The naive autonomous version looks like this:

async def run_autonomous(goal, tools, interval_seconds=60):
    memory = []

    while True:
        messages = [
            {"role": "system", "content": "You are an autonomous worker."},
            {"role": "user", "content": goal},
            *memory[-20:],
        ]

        response = await run_agent(messages, tools)
        memory.append({"role": "assistant", "content": response})

        await sleep(interval_seconds)

This is useful as a mental model (but not something I would want to deploy yet).

Once the loop runs without a synchronous caller, the engineering problem changes:

Request/response agent	Autonomous agent
User waits for an answer	Work continues after the caller leaves
Loop ends with a response	Loop may run periodically or indefinitely
Failure is a request error	Failure becomes an operational incident
Context can be request-local	Needs task state, memory, and persistence boundaries
Tool calls happen inside one request	Tool calls may become ongoing side effects
Debugging starts with one trace	Debugging starts with task history, memory, state, and logs

This is the main thesis:

Autonomy is not the loop. Autonomy is the operating model around the loop.

The Missing Primitive: A Unit of Agent Work

If the agent can keep working after the caller leaves, this is where we start the need to introduce an ability to reason around the tasks being performed.

The simplest form is something as simple as:

class TaskState(str, Enum):
    SUBMITTED = "submitted"
    WORKING = "working"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELED = "canceled"

@dataclass
class Task:
    id: str
    goal: str
    state: TaskState
    output: str = ""
    history: list[dict] = field(default_factory=list)
    events: list[dict] = field(default_factory=list)

This is indeed not an AI breakthrough… it’s ordinary distributed-systems plumbing.

Production autonomous agents inherit all the boring concerns that make systems operable:

submission,
lifecycle state,
output capture,
error reporting,
cancellation,
retention,
auditing,
ownership.

If a user starts a long-running research task and closes the browser, they need a task ID. If an agent monitors a Kubernetes namespace, an operator needs to know whether it is working, stuck, failed, or canceled. If a tool starts returning bad data, you need to know which tasks used it.

The A2A protocol specification similarly treats a task as the fundamental unit of work for long-running agent interaction - and this is a good place to start.

Budgets Are Not Just About Cost

The first reason people add budgets is usually cost. But budgets are also safety controls, especially given the risks OWASP groups under excessive agency in LLM applications:

Budget	What it bounds
max iterations	runaway reasoning loops
max runtime	stuck or excessively long work
max tool calls	API pressure and side effects
token/cost budget	spend and context growth
per-iteration timeout	one blocked tool/model call

A minimal check can be as simple as:

def budget_exhausted(budgets, started_at, iteration, tool_calls):
    if budgets.max_iterations and iteration >= budgets.max_iterations:
        return "max_iterations"
    if budgets.max_runtime_seconds and time.monotonic() - started_at >= budgets.max_runtime_seconds:
        return "max_runtime_seconds"
    if budgets.max_tool_calls and tool_calls >= budgets.max_tool_calls:
        return "max_tool_calls"
    return None

Why OpenClaw, LangGraph, CrewAI, ADK and the Others Point in the Same Direction

OpenClaw’s design pattern is meant to keep a private, self-hosted agent running with memory, integrations, skills, and a heartbeat. Let’s have a look at how other frameworks approach it:

Framework	Approach
LangGraph	Durable, stateful agent orchestration.
CrewAI	Crews, flows, memory, guardrails, and human-in-the-loop controls
OpenAI Agents SDK	Gives developers a managed loop with sessions.
Google ADK	Frames the problem around production agents, graph workflows, evaluation, debugging, context management, and deployment
Semantic Kernel	Brings the same agent/tool orchestration pattern into Microsoft’s enterprise application ecosystem
AutoGen	Extends Microsoft Research’s multi-agent conversation work into an official framework

The abstractions and approaches are different in some cases, but the same primitives keep reappearing in the source code implementation, including:

state
tools
memory
guardrails
tracing
human intervention
deployment
resumability
task control

Once agents stop being one-off prompt handlers, frameworks have to become work managers. They need to manage units of agent work, not just model calls. And when we have to take it to the next level of scale, this is where things get more complicated…

Kubernetes Enters the Picture

A platform cannot guarantee that the model will reason correctly. But if we have a way in which we can abstract some of these complex agentic concepts into architectural abstractions, we can then answer practical questions that become unavoidable when you run many agents:

Where does each agent run?
What identity does it have?
Which tools can it reach?
Which secrets can it read?
What network is it allowed to access?
How do we restart it?
How do we observe it?
How do we isolate it?
How do we scale it?

AI workloads are moving away from short-lived stateless requests, and more toward coordinated agents that run constantly, maintain context, use tools, execute code, and communicate over longer periods.

If OpenClaw is one vision of the always-on agent, Kubernetes is one answer to the fleet question:

What happens when every team, service, workflow, or tenant wants its own autonomous agents?

At that point you need scheduling, identity, isolation, policy, rollouts, configuration, and observability. Basically we need the same things we already learned to need for microservices, except the workload is now non-deterministic, tool-using, and stateful. Kubernetes already has much of this substrate through the Operator pattern, RBAC, NetworkPolicy, HPA, and event-driven scaling systems like KEDA.

We made infrastructure much harder again. Let’s now make it simpler.

Scaling Your Agentic KAOS

To make this less abstract, let’s dive into it with a Kubernetes example using Pydantic AI and KAOS.

KAOS is a Kubernetes-native agent orchestration framework. It defines agents, MCP tool servers, and model APIs as Kubernetes resources. The most recent learnings are shared from the implementation of the new autonomous/A2A milestone, which supports asynchronous task lifecycle, JSON-RPC task methods, autonomous self-looping execution, budgets, cancellation, task history, CLI/UI debugging, and examples.

We’ll cover it with a fun example, where we’ll have a Production Operations Agent which will monitor a Kubernetes cluster - namely this will include:

an agent with a monitoring goal,
a read-only Kubernetes service account,
an MCP server exposing Kubernetes tools,
a reporting tool,
a model API,
budgets and task controls,
an endpoint for async task interaction.

The key part of the Agent configuration looks like this:

apiVersion: kaos.tools/v1alpha1
kind: Agent
metadata:
  name: cluster-monitor
spec:
  modelAPI: monitor-modelapi
  model: 'smollm2:135m'
  mcpServers:
    - monitor-k8s-mcp
    - monitor-report-mcp
  config:
    description: 'Autonomous cluster monitoring agent'
    instructions: |
      You are a Kubernetes cluster monitoring agent.
      List pods, check status, and generate a health report.
    autonomous:
      goal: 'Monitor the Kubernetes cluster health. List pods, check their status, and generate a health report.'
      intervalSeconds: 60
      maxIterRuntimeSeconds: 120
    taskConfig:
      maxIterations: 5
      maxRuntimeSeconds: 300
      maxToolCalls: 20

These are some of the key configurations:

autonomous.goal defines the persistent objective.
intervalSeconds controls the loop cadence.
maxIterRuntimeSeconds bounds one iteration.
taskConfig gives bounded defaults for async tasks.
MCP servers define the tools.
Kubernetes RBAC defines what those tools can actually access.

Monitoring is a good use-case for autonomous agents as the goal persists over time, the environment changes, and the agent needs tools but should be heavily constrained.

For the full end-to-end example, including the MCP Servers and Cluster RBAC, you can try it yourself in the hands on KAOS example section.

Continuous Mode vs Async Task Mode

Another design decision that we had to come across which was interesting was the distinction between “continuous autonomous execution” and “async task execution”.

“Continuous mode” does what it suggests, runs in the background with a goal, and infinitely iterates towards that goal: i.e. when the pod starts, it begins working toward that goal. In this case it is daemon-like: monitoring, checking, reporting, maintaining, watching.

“Async task mode” is the superset capability that allows the framework to execute without user / API interaction once the task is submitted; it would just not loop once the goal is achieved or the budgets are depleted.

For an example of an Async task, a caller sends a task, gets a task ID, and the agent continues working in the background.

The caller can later inspect or cancel it, which maps to the async task model in the A2A protocol.

{
  "jsonrpc": "2.0",
  "method": "SendMessage",
  "id": 1,
  "params": {
    "message": {
      "role": "user",
      "parts": [
        {
          "type": "text",
          "text": "Research recent autonomous agent frameworks and summarize findings."
        }
      ]
    },
    "configuration": {
      "mode": "autonomous"
    }
  }
}

Through the KAOS CLI, that becomes:

kaos agent a2a send researcher \
  --message "Research recent autonomous agent frameworks and summarize findings." \
  --async

Then the caller can poll:

{
  "jsonrpc": "2.0",
  "method": "GetTask",
  "id": 2,
  "params": { "id": "task_abc123" }
}

Or cancel:

{
  "jsonrpc": "2.0",
  "method": "CancelTask",
  "id": 3,
  "params": { "id": "task_abc123" }
}

Task State Is Not Memory

Another small but important lesson that we had to reason about was that task state and memory are not the same thing; this is obvious when saying it out-loud, but it was important to figure out what we need to have available where for the agent framework to have the right context. This distinction lines up with agent research that treats memory and reflection as part of the agent’s internal execution context, such as Reflexion and Generative Agents.

We found that task state should be small and stable:

submitted,
working,
completed,
failed,
canceled,
budget exhausted.

Memory can be richer:

user messages,
agent responses,
tool calls,
tool results,
delegations,
observations,
session history.

If you mix them, your task API becomes noisy and your memory system becomes responsible for lifecycle control (which is not what anyone would want).

Debugging Many Autonomous Agents

The debugging story changes as soon as the work is no longer attached to one waiting user - everything gets exponentially harder.

You need to answer:

What task did the agent start?
Is the agent still running?
What did the agent do?
Which tools did the agent call?
What did those tools return?
Did the agent hit a budget?
Can I stop the agent?

We had to think about those questions and expose abstractions to support on those workflows; with KAOS these are through CLI and UI:

kaos agent a2a send <agent> --message "..." --async
kaos agent a2a get <agent> --task-id <id>
kaos agent a2a cancel <agent> --task-id <id>

The UI adds also agent-card inspection, SendMessage, task viewer, auto-polling, cancellation, task history, and memory conversation views. At the telemetry layer, the OpenTelemetry GenAI semantic conventions now define model and agent spans that are directly relevant for debugging these systems.

How You Could Build the Basics Yourself

You do not need a full framework to understand the minimal shape, we can start with a single-iteration primitive:

async def run_agent_once(goal, history) -> tuple[str, int]:
    messages = build_messages(goal, history)
    response = await run_agent(messages, tools)
    return response.text, response.tool_call_count

Then wrap it in the smallest useful control loop:

async def run_autonomous_task(task, budgets, cancel_event):
    started_at = time.monotonic()
    iteration = 0
    tool_calls = 0

    task.state = TaskState.WORKING

    while True:
        if cancel_event.is_set():
            task.state = TaskState.CANCELED
            return task

        reason = budget_exhausted(budgets, started_at, iteration, tool_calls)
        if reason:
            task.events.append({"type": "budget.exhausted", "reason": reason})
            task.state = TaskState.COMPLETED
            return task

        output, calls = await run_agent_once(task.goal, task.history)
        task.output = output
        task.history.append({"iteration": iteration, "output": output})
        tool_calls += calls
        iteration += 1

        if calls == 0:
            task.state = TaskState.COMPLETED
            return task

This is not production-ready: i.e it does not handle persistence, retries, distributed workers, authentication, policy, observability, or recovery after process restart. However it shows the skeleton:

task identity,
state transitions,
budgets,
cancellation,
output,
history,
completion detection.

If you use an existing framework, look for where these pieces live.

When Not to Make Your agents Autonomous

There is a temptation to turn every useful agent into a background process. Anthropic’s advice is a useful counterweight here: start with the simplest solution possible, and only increase agentic complexity when it is needed.

Autonomy is helpful when:

the environment changes over time
the goal persists beyond one request
the work is too long for a synchronous response
the agent can safely observe or act with scoped tools
the user benefits from periodic or event-driven progress

Autonomy is not helpful when:

the action is high-risk and lacks approval controls
tool permissions are broad or unclear
success criteria are vague
cancellation is missing
progress cannot be inspected
cost or side effects are unbounded

Those risks also show up in more formal governance and security guidance, including the NIST AI Risk Management Framework and the OWASP LLM Top 10 category of excessive agency.

Lessons for Production Autonomous Agents

To wrap up, I want to share a few set of learnings / patterns I would carry into any autonomous-agent system.

1. Start with the loop, but design the task contract early

The agentic loop is the easy part to prototype, however the task contract is what makes it operable.

2. Separate continuous autonomy from bounded background work

An agent that monitors forever and an agent that writes a report in the background need different controls.

3. Treat budgets as safety controls

Budgets bound cost, time, tool side effects, API pressure, and runaway reasoning.

4. Keep task state separate from memory

Task state is the external lifecycle, memory is the execution context, and mixing them makes APIs noisy and debugging harder.

5. Scope tools with permissions

Autonomy becomes risky through tools, ensure you support this with read-only service accounts, scoped roles, network policy, and secret boundaries matter more than the prompt. This is where Kubernetes RBAC, NetworkPolicy, and OWASP’s LLM security guidance become more important than prompt wording.

6. Build cancellation into the first version

Cancellation is not an advanced feature, it is a foundational feature that should be integrated by design.

7. Use Kubernetes for workload concerns, not reasoning quality

Kubernetes can help with lifecycle, identity, permissions, isolation, networking, rollouts, and observability - but it will not make a bad model / agent less bad. The Kubernetes Agent Sandbox work is an example of the same platform question being addressed with agent-specific workload primitives.

8. Instrument everything

Agent loops are variable, non-deterministic, and tool-heavy: Traces, logs, metrics, task IDs, and memory events are how you understand them later. The OpenTelemetry GenAI agent spans and model spans are useful standards to track here. This is something that we skimmed through in this post, if you are interested on a more in-depth post on this you should check out: Monitoring KAOS: Observability for Multi-Agent Systems.

Production Observability for Multi-Agent AI (with KAOS + OTel + SigNoz)

Thu, 05 Mar 2026 00:00:00 GMT

You’ve built an AI agent that works in your laptop. It automatically chains tools together, delegates tasks to specialist sub-agents, and produces sound results.

Then you deploy it to production:

A user reports that a request “took forever”.
Another says they got a strange response.
Your logs show the agent ran

But, what happened in those 45 seconds between request and response?

Welcome to the observability challenge of agentic systems.

“Traditional logs tell you that it ran. Observability tells you why it ran like that.”

In this article, we’ll walk through a complete end-to-end example of observability for multi-agent systems.

Below is a preview of the multi-agent system that we’ll analyse, monitor and diagnose in this post:

The multi-agent system that we will monitoring

Let’s start with the main question of…

…Why do Multi-Agent Systems Need (Different) Observability?

Traditional microservices have predictable patterns: a request comes in, some processing happens, a response goes out. Latency is relatively consistent, code paths are deterministic, and debugging usually involves tracing a single thread of execution.

It’s Not Just a Request-Response

Agentic systems break traditional assumptions:

Traditional API	Agentic System
Synchronous request-response	Iterative reasoning loops
Predictable latency (50-500ms)	Variable: 100ms to 60+ seconds
Deterministic code paths	Non-deterministic LLM decisions
Single service per request	Model calls + tool calls + delegations
Fixed cost per request	Cost varies by token usage

AI Agents 101: The Agentic Loop

Consider the core loop of a multi-AI agent system. This is the deceptively simple pattern that has led to the current wave of innovation in AI systems.

Below is a simplified agentic loop as we skip a lot of the nuances, but the idea is that we send a first call for the LLM to respond with tool calls or delegation calls, and if there’s none or we run into the max steps, then we send the final response.

async def process_message(self, messages):
    for step in range(self.max_steps):
        # 1. Call the LLM specifically for tool/delegation requests
        response = await self.model.process_tool_request(messages)

        # 2. If the model wants to use a tool, execute it
        if response.has_tool_call:
            result = await self.execute_tool(response.tool_call)
            messages.append({"role": "tool", "content": result})
            continue

        # 3. If the model wants to delegate, call another agent
        elif response.needs_delegation:
            result = await self.delegate_to_agent(response.delegation)
            messages.append({"role": "assistant", "content": result})
            continue

        else:
            # If no tools/delegations we move to final answer
            break

    # 4. We return our final answer with any context gathered
    return await self.model.process_final_answer(messages)

As we can see, each iteration of this loop may take a different path. The model might need one tool call or five.

It might delegate to one sub-agent or chain through three. Traditional logging (e.g. “request started” … “request completed”) tells you almost nothing about what actually happened.

And once we take this into a distributed system it gets even more complex to understand what is going on - as we will see in this post.

In the last section of this post we will show you also how we will instrument this agentic loop specifically as well.

We’ve replaced ‘request-response’ with ‘request—panic—tool—panic—delegate—panic—response’.

The Three Pillars of Agent Observability

OpenTelemetry provides three types of telemetry data, each serving a distinct purpose for agentic systems.

Traces

The trace hierarchy maps directly to what the agent did, capturing and connecting every hop across the journey.

HTTP POST /v1/chat/completions (15.2s total)
-> agent.agentic_loop
    -> agent.step.1 (3.1s)
        -> model.inference (3.0s)
    -> agent.step.2 (8.5s)
        ->model.inference (2.1s)
        -> tool.web_search (6.3s)   <- Here may be your bottleneck
    -> agent.step.3 (3.4s)
        -> model.inference (3.3s)

Traces answer: “What path did this request take through my agents?”

Logs

Traditionally in software, logging provides an inside view into the behaviour and flow of the application.

These provide a way to understand what happened throughout a particular request or session, and catch also critical information such as exceptions.

In OpenTelemetry, the logs also are captured with the respective traces so they can be connected respectively.

2024-01-15 10:30:45 INFO [trace_id=abc123] Starting message processing
2024-01-15 10:30:47 DEBUG [trace_id=abc123] Model response: calling tool 'web_search'
2024-01-15 10:30:53 ERROR [trace_id=abc123] Tool execution failed: API rate limited

Logs answer: “What did the agent ‘think’ at each step?”

Metrics

The metrics provide granular time-based KPIs that evolve over time, and can be aggregated and windowed to tell a particular historical story.

<- Store number of tokens per request

<- Store success / failure rate of requests

<- Store latency for model calls

Metrics answer: “How is my system performing overall?”

Bringing it all together

The magic happens when these three are correlated (aka connected).

This allows us to do things like “click on that ERROR log in your observability backend” and diagnose the exact span in the trace where the failure occurred.

Multi-Agent Context Propagation

The real challenge comes with multi-agent systems. When Agent A delegates to Agent B, which delegates to Agent C, you want a single unified trace - not three disconnected ones.

This requires context propagation.

This involves passing trace context through HTTP headers using the W3C Trace Context standard.

The result is a unified trace across all agents - here’s an example that shows how a trace context spans across the coordinator agent, the researcher agent and the analyst agent.

coordinator.agent.agentic_loop (trace_id: abc123)
	-> coordinator.model.inference
	-> coordinator.delegate.researcher
	    -> researcher.agent.agentic_loop (trace_id: abc123)
	        -> researcher.model.inference
	        -> researcher.tool.web_search
	-> coordinator.model.inference
	-> coordinator.delegate.analyst
	    -> analyst.agent.agentic_loop (trace_id: abc123)
	        -> analyst.model.inference
	        -> analyst.tool.calculator

Without context propagation, multi-agent debugging is just distributed guessing.

The Practical Use-Case: A Multi-Agent Research System

Let’s now start building something concrete.

We’ll use KAOS(Kubernetes Agent Orchestration System), an open-source framework to deploy, manage and scale multi-agent systems in Kubernetes.

The Multi-Agent System to Monitor

Our use-case consists of a coordinator agent that delegates research and analysis tasks to specialist sub-agents:

Simplified architecture of the multi-agent system to monitor

Each component—agents, tools, and model APIs—sends traces, metrics, and logs to an OpenTelemetry collector, which forwards everything to your chosen backend for visualization and analysis.

Prerequisites: Commands & Tools

Before we start, you’ll need:

KAOS CLI installed: pip install kaos-cli==0.2.7
An LLM API key (Any provider like OpenAI, Nebius, etc)
kubectl, helm and a Kubernetes cluster (KIND, minikube, or a cloud cluster)

Installing KAOS

First, let’s install the KAOS operator with OpenTelemetry enabled and an observability backend. We’ll use SigNoz as an open-source, OpenTelemetry-native option.

Using the KAOS CLI (recommended):

# Install the KAOS operator
kaos system install \
    --set logLevel=DEBUG \
    --wait \
    --monitoring-enabled # Enables monitoring setup (supports signoz and jaeger)

# Verify the installation
kaos system status

# Change context to use (+create) this current namespace for convenience
kaos system working-namespace kaos-hierarchy

Once it’s running we can create our multi-agent system using one of the samples provided with the CLI.

You can also see the ANNEX at the end of the blog post to deploy each of the components one by one.

kaos samples deploy \
    3-hierarchical-agents \
    --provider openai \
    --wait \
    --api-secret # This will prompt your api-key secret

You can see the —provider flag which specifies the backend, such as nebius, gemini, bedrock and 100s of other providers [KAOS docs], as well as your —api-secret which will be prompted interactively.

Once installed you open the UI with the following command.

kaos ui --monitoring-enabled

This allows us to see the deployed multi-agent system:

KAOS Multi-Agent System - Hierarchical Sample

Putting It All Together: Monitoring KAOS

Now let’s generate some traffic and start monitoring KAOS.

Interacting with Agents

You can interact with agents in multiple ways:

# Invoke the coordinator agent directly
kaos agent invoke supervisor \
  --message "Research the current AI chip market and calculate the market share of the top 3 companies."

Activating the Chat through the User Interface:

Querying the Supervisor through the User Interface

Behind the scenes, this triggers a complex chain of operations:

Supervisor receives the request
Supervisor calls the LLM, which decides to delegate
Researcher agent is invoked for market research (and calls two researcher sub-agents)
Analyst agent calculates market shares (and calls two analyst sub-agents)
Supervisor synthesizes the final response

All of this is caputred in auditable traces: every LLM call, every tool execution, every delegation.

Viewing Traces: Understanding Request Flow

Why traces matter for agentic systems: Unlike traditional request-response services, agents make multiple decisions per request. Traces let you see each decision point, how long it took, and what path the agent chose.

The trace list shows all requests flowing through your agents. Each trace represents a complete user interaction.

Click on a trace to see the full request flow:

View the request traces and spans.

A single trace showing the coordinator delegating to researcher and analyst agents, with each span representing a distinct operation.

This trace visualization answers questions that would otherwise require hours of log spelunking:

Why did this request take 15 seconds? The web_search tool took 8 seconds.
Which agents were involved? Coordinator → Researcher → Analyst → Coordinator.
How many LLM calls were made? 6 calls across the three agents.
Did any tools fail? All tools completed successfully (green spans).

Log Correlation: Understanding Agent Reasoning

Traces tell you what happened. Logs tell you why. OpenTelemetry correlates them automatically through trace_id and span_id attributes.

Every log entry includes these identifiers, enabling you to:

Click on a span in your trace
View all logs emitted during that span
Understand the agent’s reasoning at each step

It is possible to see the view of the logs themselves as well as further details.

Double click on specific trace to see correlated logs

We can then drill deeper into individual log entries.

Here for example we can view the log “Delegation task: …”.

This shows us the entire prompt that was delegated, which is quite useful for diagnostics and debugging.

Clicking into a log for viewing the metadata and content

This includes full log context including all attributes, resource labels, and the complete message.

Exception Tracking: Finding Production Issues

In production, things fail. OpenTelemetry captures exceptions as first-class citizens, and attaches them to the span where they occurred.

We can try this by asking the supervisor to delegate to a non-existing sub-agent:

# Invoke the coordinator agent directly
kaos agent invoke supervisor \
  --message "We are now testing valid exception functionality. Try to delegate to a non-existing agent to validate that it works correctly."

Or directly via the UI:

Sending an incorrect delegation request to test the system.

If we now navigate to the Exceptions tab, we can now see that there is an entry. This is the Exceptions list view, where all exceptions for the time period are listed.

This list captures errors across the system and correlates them also with the respective request traces, and logs.

We are also able to filter by different (agent) services, as well as other attributes.

Showing the list of exceptions captured.

We can see clearly that the exception was due to an attempted incorrect delegation.

And we can also visualise the error in the trace view.

Trace view for exceptions

Metrics: Operational Overview

While traces show individual requests, metrics show trends over time.

Here are some example key metrics in agentic systems:

Metric	What It Tells You
`kaos.requests`	Request volume by agent
`kaos.request.duration`	Latency distribution (P50, P95, P99)
`kaos.model.calls`	LLM API usage (cost indicator)
`kaos.tool.calls`	Tool execution frequency
`kaos.delegations`	Multi-agent coordination patterns

These metrics enable alerting on production issues:

Request latency > 30s
Error rate > 5%
Model call failures > 1%

We can also visualise them:

Metrics dashboard showing request rates, latency percentiles, error rates, and token usage across all agents.

Under the Hood: How It Works

Now that you’ve seen observability in action, let’s dive into how it’s implemented. The challenges here aren’t obvious until you start building—and the solutions are broadly applicable to any agentic system.

The Architecture

KAOS separates control plane (Go) from data plane (Python); inside the Python application we have a OpenTelemetry manager (KaosOtelManager) to provide utilities to simplify the workflows.

KAOS architecture

The key insight: telemetry configuration flows from Operator -> Data Plane -> OTEL Collector. Users configure telemetry once in values.yaml, and the operator propagates it to all components.

Instrumenting the Core Logic

When it comes to agentic systems, it is not just about instrumenting the request-response, but it’s also about capturing the flow across the agent iterations.

To start with we ensure that we capture the overarching span across the top level request/response. We then should also instrument other key components such as the Agentic Loop (example below), as well as extra calls such as MCP calls, and agent delegations.

# This is the same function pseudo-code that we showed initially but instrumented
async def process_message(self, session_id: str, messages: List[Dict]) -> str:
    """Process message through agentic loop with full tracing."""

    # Start root span for entire message processing
    span = otel.span_begin("agent.agentic_loop", SpanKind.INTERNAL)
    span.set_attribute("agent.name", self.name)
    span.set_attribute("session.id", session_id)
    span.set_attribute("agent.max_steps", self.max_steps)
    span_failure = False

    try:
        # Agentic loop logic (see below)

    except Exception as e:
        # If exception mark as failure
        span_failure = True
        otel.span_failure(span, e)
        raise
    finally:
        if not span_failure:
            otel.span_success(span, e)

Instrumenting the Agentic Loop

We can now use the same pattern in the agentic loop, where we can capture the iterations with the respective context. This will be important as we can also capture the correlated logs and metrics that are connected to this particular request.

# Previous logic outlined above...

        for step in range(self.max_steps):
            # Span for each iteration
            step_span = otel.span_begin(f"agent.step.{step + 1}")
            step_span.set_attribute("step", step + 1)
            agent_span_failed = False

            try:

                # MCP Calls (with OTEL span)...

                # Delegation Calls (with OTEL span)...

            except Exception as e:
                otel.span_failure(step_span, e)
                agent_span_failed = True
                raise
            finally:
                if not agent_span_failed:
                    otel.span_success(span)

    # Preview logic outlined above...

Having the spans defined in placed is what allows us to trace the request across hops.

Some useful patterns to note:

Hierarchical spans: There is a parent span for the loop, child spans for each step, grandchild spans for operations, etc.
Log before span close: Logs are emitted while trace context is active, which allows correlating the logs
Explicit span management: We are using explicit try/finally pattern to ensure spans are always closed, but we can also use context managers (i.e. with clause)

This is also what allows us to then visualise the breakdown of the request traces and spans.

Zoomed-in the traces captured.

Context Propagation for Multi-Agent Systems

When delegating to sub-agents (running in separate pods), we must propagate trace context:

# Inject context into outgoing request
from opentelemetry.propagate import inject

async def delegate(self, target_agent: str, task: str) -> str:
    headers = {"Content-Type": "application/json"}

    # Inject current trace context into headers
    inject(headers)  # Adds 'traceparent' and 'tracestate' headers

    async with httpx.AsyncClient() as client:
        response = await client.post(
            f"http://{target_agent}/v1/chat/completions",
            headers=headers,
            json={"messages": [{"role": "user", "content": task}]}
        )
    return response.json()["choices"][0]["message"]["content"]

Ensuring Trace Context Propagates

We also need to make sure that the context is received and embedded.

# Extract context from incoming request
from opentelemetry.propagate import extract

@app.post("/v1/chat/completions")
async def chat(request: Request, body: ChatRequest):
    # Extract trace context from incoming headers
    context = extract(request.headers)

    # Attach to current context so new spans are children
    token = otel_context.attach(context)
    try:
        return await agent.process_message(body.messages)
    finally:
        otel_context.detach(token)

Log Export and Correlation

For log-trace correlation, we connect Python’s logging to OpenTelemetry.

If you have used previously something like the ELK stack, with OpenTelemetry now you have a setup where the logs are pushed as opposed to pulled, using the GRPC OTEL connection.

The logging instrumentation automatically injects trace_id and span_id into log records when there’s an active span context.

To set up this connection, you can configure it with a few lines as outlined below.

from opentelemetry.sdk._logs import LoggerProvider
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter
from opentelemetry.instrumentation.logging import LoggingInstrumentor

# Set up OTLP log export
logger_provider = LoggerProvider(resource=resource)
logger_provider.add_log_record_processor(
    BatchLogRecordProcessor(OTLPLogExporter(endpoint=endpoint))
)

# Attach handler to Python root logger
handler = LoggingHandler(level=logging.DEBUG, logger_provider=logger_provider)
logging.getLogger().addHandler(handler)

Metrics for Agent Operations

We are also able to track metrics using the OpenTelemetry SDK.

If you have used Prometheus in the past, the metrics in this case are not exposed through an endpoint that would be queried through a prometheus collector. Instead this also enables a push architecture where it’s sent to the OTEL collector through the GRPC OLTP connection to the OTel collector.

When selecting metrics we aim to track metrics with low-cardinality labels to avoid cardinality explosions, as outlined in the sample below.

from opentelemetry import metrics

meter = metrics.get_meter("kaos-agent")

# Counters
request_counter = meter.create_counter(
    "kaos.requests",
    description="Number of requests processed",
    unit="1"
)

model_call_counter = meter.create_counter(
    "kaos.model.calls",
    description="Number of model inference calls",
    unit="1"
)

# Histograms for latency
request_duration = meter.create_histogram(
    "kaos.request.duration",
    description="Request processing duration in milliseconds",
    unit="ms"
)

# Usage example
request_counter.add(1, {"agent.name": self.name, "status": "success"})
request_duration.record(duration_ms, {"agent.name": self.name})

Avoid high-cardinality labels: Never use session IDs, user IDs, prompt content, or other unbounded values as metric labels. Put those in logs or trace attributes instead.

The Tip of The Iceberg

This is a gentle introduction to observability in multi-agent systems, and provides a high level view of what are some of the main components involved.

This should also provide you with enough intuition to instrument your agentic systems conscientiously, however you will need to identify which patterns work best in your particular contexts.

Instrumenting agentic AI systems with OpenTelemetry requires understanding the unique challenges these systems present:

Iterative loops need span hierarchies that map to logical operations
Multi-agent delegation requires explicit context propagation using W3C Trace Context
Tool execution benefits from dedicated spans with clear naming
Log-trace correlation requires emitting logs before span close
Metrics need low-cardinality labels to avoid storage explosions

The patterns we’ve covered apply to any agentic system, not just KAOS. Start instrumenting now.

The agents of tomorrow will be as ubiquitous as microservices are today, and OpenTelemetry gives you the visibility to operate them with confidence.

Resources

KAOS Framework: github.com/axsaucedo/kaos - The open-source framework used in this article
KAOS Documentation: axsaucedo.github.io/kaos - Full CLI and CRD documentation
OpenTelemetry Python: opentelemetry.io/docs/languages/python - Official Python SDK documentation
OpenTelemetry GenAI Conventions: github.com/open-telemetry/semantic-conventions - Emerging standards for AI observability
SigNoz: signoz.io - Open-source APM with native OpenTelemetry support

ANNEX: Creating the Multi-Agent System Manually

In this section we create the agentic system components manually instead of using the utilities for the curious ones!

We’ll show you how you can do this with CLI but you can do this also with the UI as well as with kubectl directly.

Step 1: We first connect to LLMs with a ModelAPI

The ModelAPI resource in KAOS provides a unified interface for LLM access. It supports two modes:

Proxy Mode: Routes requests through LiteLLM to external providers (OpenAI, Anthropic, Nebius, etc.)
Hosted Mode: Pulls models into your cluster (via side-car) and runs it on the server for inference

The ModelAPI can be deployed easily via CLI.

Note that in order for our agents to use the model APIs we need to provide our authentication API Key. For this example we will be using Nebius as it’s easy to set up, but you can also set up OpenAI, Gemini and dozen others.

kaos modelapi deploy llm-proxy \
    --namespace kaos-hierarchy \
    --provider nebius \
    --api-key # When provided without value this prompts the key securely

From an observability perspective, the ModelAPI gives us visibility into model call latency, token usage, and error rates—critical metrics for understanding agent performance and controlling costs.

Step 2: Deploy the MCP Tool Servers

MCP (Model Context Protocol) Servers in KAOS provide tools that agents can use.

KAOS enables FastMCP native servers with ability to create and deploy your own images.

KAOS also supports multiple native MCP runtimes via a registry. The most commonly used are:

Runtime	Description
`python-string`	Define tools as inline Python functions for testing
`kubernetes`	Kubernetes CRUD operations
`slack`	Slack messaging integration
`custom`	Your own container image

For our demo, we’ll create a calculator server. In production, you’d connect to real APIs, databases, or external services.

Note: In this section we use the python-string runtime for quick testing, however for production-ready deployment use the custom-image deployment.

First we create the calculator mcp, which will have a simple add tool that will add two numbers and return the result.

export ADD_TOOL='
def add(a: float, b: float) -> float:
    """Add two numbers together."""
    return a + b
'

kaos mcp deploy calculator \
    --runtime python-string \
    --params $ADD_TOOL \
    --wait

And we then create an echomcp, which is also a simple tool that receives a string and returns the same value as the string.

export ECHO_TOOL='
def echo(message: str) -> str:
    """Echo back the message for testing."""
    return f"Echo: {message}"
'

kaos mcp deploy echo-search \
    --runtime python-string \
    --params $ECHO_TOOL \
    --wait

And we can send a request to test the mcp.

# Run 2 + 2 on mcp calculator
kaos mcp invoke calculator \
	--tool add \
	-a 2 \
	-a 2

# Run echo hello on mcp echo
kaos mcp invoke echo-search \
	--tool echo \
	-a "Hello"

From an observability standpoint we will show how it is important to understand the calls that are sent by any agent (or any external service) and processed by the MCP servers themselves

3a. The Researcher Agent

The Agent resource in KAOS represents an AI entity that can process requests, call models, execute tools, and delegate to other agents. Each agent:

Exposes an OpenAI-compatible /v1/chat/completions endpoint
Implements the agentic loop (model > tools > model > …)
Supports configurable memory for session state

The Researcher agent will specialise in gathering information:

kaos agent deploy researcher \
    --model "openai/gpt-4o" \
    --mcp echo-search \
    --description "Research specialist" \
    --instructions "You research topics and provide ." \
    --expose=true \
    --wait

For single-agent observability, we care about:

Model call latency: How long does inference take?
Tool execution time: Are tools responding quickly?
Step count: How many iterations does the agent need?

3b. The Analyst Agent

The Analyst agent focuses on data analysis and calculations:

kaos agent deploy analyst \
    --model "openai/gpt-4o" \
    --mcp calculator \
    --description "Data analyst with calculation capabilities" \
    --instructions "You analyze data and perform calculations." \
    --expose=true \
    --wait

3c. The Supervisor Agent

Finally, the Coordinator orchestrates the other agents. For multi-agent observability, we gain additional concerns:

Delegation patterns: Which agents are called and how often?
Cross-agent latency: How much time is spent in delegation vs. local processing?
Trace correlation: Can we see the full request flow across agents?

kaos agent deploy subervisor \
    --model "openai/gpt-4o" \
    --description "Coordinator that delegates to specialist agents" \
    --instructions "You are a coordinator. Analyze user requests and delegate to your analyst and researcher." \
    --sub-agent researcher \
    --sub-agent analyst \
    --wait

And we’re done!

You can now go back to the top and run through the observability walkthrough with your deployed setup!

Claude Code: It's not replacing devs. It's moving them to a higher altitude.

Mon, 26 Jan 2026 00:00:00 GMT

Join me in this post as I explore new paradigms brought by agentic development tools, lessons learned from building systems with AI tooling, and what that could mean for the software engineering practice (not “just coding”) within the next few years.

“AI is making developers 10x” [link], “AI is making developers less productive” [link], “AI is coming for our jobs” [link], “AI enables coders” [link].

Well, which one is it?

Let’s start this exploration with this meme:

I’ve been thinking about this meme for a while. The punchline is feels like a great take to the question on how agentic development is impacting the software engineering practice.

I’ve been able to further process this thought particularly as throughout the winter break I got a Copilot Pro subscription and went on a building spree; initially building one project per day, but then getting stuck on a longer project which provided me with a more clear view of what the future may look like for the tech industry.

Building a team-sized projects in days instead of weeks

If someone was to tell me they single-handedly built a project that spans across custom Kubernetes operators, backend, frontend and SDKs, with robust CI & release process within a couple of days, I would have struggled to believe them.

That is about how shocked I was when I took on the challenge, and successfully built K.A.O.S.in just a couple of days**, a Kubernetes Agent Orchestration System** that deploys and manges multi-agent systems, which consists of a non-trivial set of components:

Golang control-plane that manages the kubernetes native resources
Python data-plane that implements Agents, MCPs and ModelAPI client/servers
React based user interface with interactive CRUD, Agent chat, MCP debugging, etc
CLI for managing resources and easily performing workflows
CI/CD with KIND based e2e testing and automated release management

Traditionally building a project like KAOS would’ve taken weeks, and would’ve required a handful of skilled individuals with different skill-sets (eg frontend, backend, infra, etc).

But oh how things have changed.

I was blazing through development milestones and the code was produced at real time speed. It got to the point where iteration cycles were faster than my capability to review changes and design further extensions. This brought two particular themes that really surprised me.

Surprise #1: Time Constant — Watching tech debt form in real time

The first surprise was the time constant.

We all know tech debt is a real problem; in practice it creeps in slowly & silently, and we only realise it when it’s too late. The interest payments eventually show up: small changes require larger investments, developer productivity stalls, interesting new bugs arise from everywhere, code areas become fragile, etc.

What I was particularly surprised about was not about tech debt itself, but about the time constant in which was being created**.** If at some point I came across an implementation that was slightly off (eg. a quick interface hack, a decision in the wrong layer, etc) and didn’t address it right there and then, within hours I could feel the codebase start bending around it.

I have to say it was fascinating seeing tech debt form in a few hours in real time, when normally this spans across weeks, or even months.

This also coupled with the realisation that now I had the luxury to backtrack on particular implementation directions when seeing something was off. This was only possible due to low cost/effort for development.

This enabled for an exploration-exploitation branching-likedevelopment approach, as opposed to a purely linear one where otherwise retracting may be too expensive compared to accepting the short term debt.

Surprise #2: Cognitive Level — Operating at a higher level of abstraction whilst staying tactical

The second surprise was at the cognitive level.

Historically when coding the traditional cycle is along the lines of: pick up a task, dive into a module, load that context into your brain’s RAM, and descend through levels of abstractions into the internals to fix what you can as further rabbit-holes arise (aka “being in the zone”). Basically this was the status quo, where we had accepted that our brains (+ typing speed) are the bottleneck.

With agentic workflows, I noticed a different mode:

Load the entire picture at once and reason about at the system level(backend, infrastructure, frontend, SDK, CI, etc) without losing the tactical thread.
Delegate chunks of implementation whilst being able to steer the direction in real time whilst having the big picture in mind.
Review the resulting artifacts as discrete, comprehensible deltas noting down extensions and improvements.
Match against overarching direction otherwise backtrack and iterate; often with various threads in parallel.

This didn’t feel like I was “no longer coding”. It just felt like I was coding at a different altitude.

The way I like to think about this is the following analogy: how a tech lead or staff engineer thinks when they’re guiding the team(s) and shaping the system across months of development cycles, instead of just editing files.

The real unlock is abstraction: specs become complex systems

Now bringing these two “surprises” together is what makes me think: we have been over-emphasising the improved time constant and under-discussing the higher cognitive level that is being unlocked.

If we do now focus on the higher cognitive level unlocked, a more complex question arises:

If we have unlocked this higher level of abstraction, what could this enable in practice?

In organisations we may have a functional setups (teams of backend vs teams of frontend, etc) or cross functional setups (teams with mixes of frontend, backend, etc). Each of these team members would have a manager; teams of managers would report to a head of department; heads would report to a director; directors would report to… you get the point.

As of today, we build these layers because cognition and coordination don’t scale linearly. At some point, you need a layer whose job is not writing the code, but steering the system; i.e. aligning architecture, interfaces, guardrails, and operational quality. Further layers then arise to steer strategy, business priorities, etc — which may or may not be aligned with the execution (and vice-versa).

Why are we talking about organisational setups? Here is the twist: with this higher level of abstraction, a single developer may be able to move upwards to that next layer of cognitive abstraction without giving up the ability to ship tactical changes. The metric is no longer “how fast do I code”; it becomes:

How many layers of complexity can one person hold and steer effectively?

And once you look at it through that lens, the most interesting question becomes not “how many tickets can I close”, but:

What is an example of what we could build at this higher level of abstraction?

The first potential answer to this question came up as I was exploring it through various conversations. Namely on the topic of a common organizational reality: departmental systems optimizing in silos.

In a typical e-commerce setup you might have separate systems for:

Pricing — optimising discounting, demand shaping, stock clearance
Replenishment — optimising inventory coverage to enable re-stocking
Marketing — optimising campaigns, conversion, and demand spikes

Each system can be locally smart, but globally dumb.

Pricing may discount inventory down because it sees a stock issue. Replenishment may ramps stock up because it sees demand. Marketing may run a campaign that shifts demand curves. If these systems aren’t aligned, you get cannibalisation: adversarial optimisation that makes the overall customer lifetime value outcome worse.

Today, organizations try to solve this with process and coordination: roadmaps, planning, architecture governance, quarterly alignment, and a lot of “let’s align”. It works… sometimes. But it’s expensive and slow.

If individuals can operate one abstraction layer higher, it enforces that more of this coordination gets embodied in systems, not meetings:

shared constraints,
shared feedback loops,
shared intent,
and shared operational reality.

However, this still doesn’t answer my question; as it only provides the “means”, and not the “end”.

Namely it doesn’t answer the question: What is an _actual_ example of what could be built at this level of abstraction?

The next attempt was explored through the concept of programming languages:

From objects and functions to architectural primitives (K8s CRDs as a mental model)

When thinking of this, the mental model I kept returning to was Kubernetes: specifically the Operator Pattern — also known as Custom Resource Definitions (CRD).

A CRD is not just a schema. It’s a way to define a higher-level concept as architectural components that can be instantiated, reconciled, and operationalized into sophisticated and complex applications.

In the example of KAOS we abstract the concept of an Agent, or an MCP or a ModelAPI into an architectural component with a set of attributes, that reconciles into a broader set of system and application components that interact across a complex system.

So the thought experiment then becomes:

What if “programming language constructs” were CRD-level like architectural components; namely if we didn’t operate in objects but instead with dynamic architectural components via a data plane of concepts which reconcile into services, policies, interactions, dependencies and so on.

This is where another project I built (bit more half baked this one) gave me a different angle.

Basically, I wondered:

What would a programming language look like if it natively operated at the next level of abstraction?

And before someone comes here to just shout “erlang!”, I want to explore a bit of blue sky thinking with a slightly more naive exploration. Starting with the idea of a programming language where an agent can be a first-class architectural component in a specification; then what happens if we design it where:

modules are agents,
control flow is agentic,
and even the interpreter loop is an agent?

Not an interpreter in the classic sense (“parse -> execute”), but something closer to a runtime that plans, delegates, reconciles, validates, and iterates.

That line of thinking is what led me to experiment with AgenticScript: another weekend project — this time more half-baked — but that may help answer the bigger question.

AgenticScript was an idea to a distributed agent programming language explicitly designed around coordination: agent spawning, inter-agent communication (ask / tell), tool management, a message bus, and rich debugging via a REPL that surfaces agent status, message bus performance, and communication flows.

It was really fun (and impressive) to work through a lexer, parser and interpreter with a code copilot (+ learned there’s a whole sub-field of Agent-Oriented Programming). However although it wasn’t a success in regards to pursuing the project end-to-end, it was a success in a different sense as it helped progress this question.

And the answer seemed to lead to the potential of a more uncomfortable conclusion:

A language at this next abstraction layer might not look like a programming language at all

If we actually want a layer above today’s languages, it may not be a rigid syntax like we’re used to.

This line of thinking is closer to Andrej Karpathy’s famous 2023 phrase of “the hottest new programming language is English”. Not sure if it’s right to call this a “superset of human language”, but analogous to semi-structured Technical Design Documents (TDDs) that describe concepts, constraints, and components in ways that can compile into living artifacts that interact with broader systems.

Unsurprisingly, this is where the industry is already at the moment; large changes already require TDDs to be written before a solution is implemented.

The difference is that here “spec-driven development” would be framed as treating the specification as the source of truth and regenerating derived artifacts when it changes. And at the more radical end, there’s an explicit argument that code becomes a byproduct between requirements and system outcomes.

If you zoom out, you can see this as the latest iteration of a long lineage of code generation that has been exploring “spec/models -> generated artifacts” for years. The difference now is the “runtime” that would be required to execute such a new paradigm: the interpreter isn’t just executing instructions, it’s continuously reconciling intent against reality.

Which implies something important:

If we really want that future, the execution layer will likely be more complex than most of the software we build today. Namely because it needs to absorb the semantic and operational burden that currently lives in teams, process, and organizational structure.

A pragmatic cop-out (that might actually be the answer)

I can use this as a cop-out and say: maybe this already answers the question of what could be an example of what could be built at this higher level of abstraction.

One plausible answer is:

an ecosystem that itself exists at a higher abstraction layer (primitives, protocols, declarative intent),
enabling individuals (and teams) to build at that same layer,
with higher leverage in speed, integration, and reliability — iff the specs and guardrails are treated as first-class.

Which brings us right back to the meme.

A “comprehensive and precise spec” is called code.

The shift is that we’re increasingly moving the definition of “code” upward: from functions and classes toward specifications, architectural primitives, and executable intent.

Three final reflections from this philosophical rant

1) We’re heading toward an “agentic SDLC” standardization moment

I am getting flashbacks of the early 2010s, where many of us may remember those were the days where Scrum coaches were being hired everywhere to teach us “the way” of agile (and move away from the dark side of waterfall).

Whether or not you love Scrum, the pattern is the interesting part: practice standardizes after fragmentation.

I expect something similar here:

A small minority becomes visibly faster.
Organizations try to replicate the gains.
We get a proliferation of incompatible “agentic workflows.”
Then we converge into a more standardized agentic SDLC — maybe even with something like “agentic coaches” (a phrase that sounds ridiculous until it doesn’t).

So it may not be too far until we start seeing agentic coaches as part of teams, pushing for day-long cycles, renaming spikes to something else (agentic runs?), and inventing another set of agentic poker cards (oh my).

2) Developer profiles are diverging (fast)

And we’re already seeing the developer profiles shift. At least I have noticed three broad categories:

High performers becoming disproportionately more productive.
Low-discipline workflows becoming disproportionately more dangerous (generating code, not reading it, opening PRs before running it, creating more work downstream).
The middle majority watching from the sidelines, waiting for “the way we do it here” to crystallize, instead of venturing into the unknown.

This last group is important. Most industries don’t shift because early adopters are excited. They shift when norms, expectations, and playbooks form.

3) The bottleneck moves upstream: from “writing code” to judgment, constraints, and verification

The new 10x isn’t typing faster; it’s steering better. When you can generate and integrate at a higher level, the differentiator becomes (a) what you choose to build, (b) how precisely you specify it, and © how aggressively you verify it.

This is also where the tech debt surprise lands me: accelerated iteration doesn’t just accelerate output, it accelerates misalignment. People have started naming the new type failure “epistemic debt”: shipping systems you can’t explain, defend, or reliably change (even if tests pass today).

So the agentic SDLC endgame isn’t “everyone becomes 10x”. It is: speed becomes commoditised; rigour becomes the moat. Teams that turn specs, guardrails, evals, and operational feedback loops into first-class artifacts get compounding leverage. Teams that don’t… compound confusion at machine speed.

It is now our responsibility as practitioners and leaders figure out this shift upwards in the cognitive stack; individuals can now do what teams could; teams can do what departments; and what follows should be able to invent the future. This analogy may go beyond the field of software engineering.

As they say: the future is here, it’s just not evenly distributed.

Beyond CUDA: GPU Accelerated Python for Machine Learning on Cross-Vendor Graphics Cards Made Simple

Fri, 13 Nov 2020 00:00:00 GMT

Machine learning algorithms — together with many other advanced data processing paradigms — fit incredibly well to the parallel-architecture that GPU computing offers. This has driven massive growth in the advancement and adoption of graphics cards for accelerated computing in recent years. This has also driven exciting research around techniques that optimize towards concurrency, such as model parallelism and data parallelism.

In this article you’ll learn how to write your own GPU accelerated algorithms in Python, which you will be able to run on virtually any GPU hardware — including non-NVIDIA GPUs. We’ll introduce core concepts and show how you can get started with the Kompute Python frameworkwithonly a handful of lines of code.

Ben-Nun, Tal, and Torsten Hoefler. “Demystifying parallel and distributed deep learning: An in-depth concurrency analysis.” ACM Computing Surveys (CSUR) 52.4 (2019): 1–43.

First we will be building a simple GPU Accelerated Python script that will multiply two arrays in parallel which this will introduce the fundamentals of GPU processing. We will then write a Logistic Regression algorithm from scratch on the GPU. Below are the core topics that we will cover, together with the respective resource links:

Kompute python package Installation
GPU Accelerated Array Multiplication Example in Python
GPU Accelerated Logistic Regression Example in Python

Only basic programming experience is required for anyone reading this article, no knowledge of GPU computing is required. You can find the full code in the main repository, and we also created an online Google Colab Notebook where you’ll be able to run the example with a GPU for free — you can find it in this link.

Google Colab Notebook with Examples for GPU

Enter Kompute & the Vulkan SDK

There are two parts to the Python framework we will be using today, both which are in the name itself —the Vulkan SDK, and Kompute.

Playing “where’s waldo” with Khronos Membership (Image by Vincent Hindriksen via StreamHPC)

The VulkanSDK is an Open Source project led by the Khronos Group, a consortium consisting of numerous tech companies that have come together to work towards defining and advancing the open standards for mobile and desktop media (and compute) technologies.

A large number of high profile (and new) machine learning frameworks such as Google’s Tensorflow, Facebook’s Pytorch, Tencent’s NCNN, Alibaba’s MNN—between others — have been adopting Vulkan as their core cross-vendor GPU computing SDK. This is primarily to enable the frameworks for cross platform and cross vendor graphics card support.

As you can imagine, the Vulkan SDK provides very low-level C / C++ access to GPUs, which allows for very specialized optimizations. This is a great asset for GPU computing— the main disadvantage is the verbosity involved, requiring 500–2000+ lines of C++ code to only get the base boilerplate required to even start writing the application logic. This can result in expensive developer cycles and errors that can lead to larger problems. This was one of the main motivations for us to start the Komputeproject.

The Kompute Python package is built on top of the Vulkan SDK through optimized C++ bindings, which exposes Vulkan’s core computing capabilities. Kompute is the Python GPGPU framework that we will be using in this tutorial to build the GPU Accelerated machine learning algorithms.

Kompute Documentation (Image by Author)

Installing the Python Kompute Package

In order for us to start using the Kompute Python Package we will need to install its required dependencies. The package is available in Pypi, which means we can install it with pip install. You will however require the following key components installed on your machine before being able to use it:

CMAKE v3.41+ (install in Windows, Linux (Ubuntu), Mac)
Vulkan SDK installed via official website
C++ compiler (eg. gcc for linux / mac, MSVC for Windows)

Once you have these dependencies installed, you can simply run:

pip install kp==0.5.1

You should now see a success message confirming that the Kompute Python package has been installed. You can try it out yourself in the Google Colab Notebook provided in the repository, which you can set up with a GPU.

Writing your first Kompute: GPU Multiplication

To build our first simple array-multiplication GPU computing application using Kompute, we will write a simple python program that will do the following:

Create a Kompute Manager (selects device 0 by default)
Create Kompute Tensors to hold data (two input one output)
Initialise the Kompute Tensors in the GPU
Define the code to run on the GPU
Dispatch GPU shader execution against Kompute Tensors
Use Kompute Operation to map GPU output data into local Tensors
Print your results

The full Python code required is quite minimal, so we are able to show the full script below. We’ll break down each of the sections in more detail.

import kp
import pyshader as ps

# 1. Create Kompute Manager (selects device 0 by default)
mgr = kp.Manager()

# 2. Create Kompute Tensors to hold data
tensor_in_a = kp.Tensor([2, 2, 2])
tensor_in_b = kp.Tensor([1, 2, 3])
tensor_out = kp.Tensor([0, 0, 0])

# 3. Initialise the Kompute Tensors in the GPU
mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out])

# 4. Define the multiplication shader code to run on the GPU
@ps.python2shader
def compute_shader_multiply(index=("input", "GlobalInvocationId", ps.ivec3),
                            data1=("buffer", 0, ps.Array(ps.f32)),
                            data2=("buffer", 1, ps.Array(ps.f32)),
                            data3=("buffer", 2, ps.Array(ps.f32))):
    i = index.x # Fetch the current run index being processed
    data3[i] = data1[i] * data2[i] # Perform multiplication

# 5. Dispatch algorithm execution against Kompute Tensors
mgr.eval_algo_data_def(
  [tensor_in_a, tensor_in_b, tensor_out],
  compute_shader_multiply.to_spirv())

# 6. Sync tensor data from GPU back to local
mgr.eval_tensor_sync_local_def([tensor_out])

# 7. Print results
print(tensor_out.data()) # prints [2.0, 4.0, 6.0]

1. Create a Kompute Manager (selects device 0 by default)

First, we’ll create our Kompute Manager, which is in charge of creating and managing all the underlying Vulkan resources.

# ...previous code blocks

mgr = kp.Manager()

# ...latter code blocks

As you can see, here we are initializing our Kompute Manager, which by default creates all the base Vulkan resources on Device 0 (in my case it’s an NVIDIA card, and Device 1 is my integrated graphics card). For more advanced use-cases it’s also possible to provide the underlying GPU queues that you’d like to load — in this other tutorial we show how this can lead to significant speedups, but this is outside of scope of this article.

2. Create Kompute Tensors to hold data (two input one output)

We will now create the Kompute Tensors that will be used for input and output. These will hold the data required which will be mapped into the GPU to perform this simple multiplication.

# ...previous code blocks

tensor_in_a = kp.Tensor([2, 2, 2])
tensor_in_b = kp.Tensor([1, 2, 3])
tensor_out = kp.Tensor([0, 0, 0])

# ...latter code blocks

When the tensors are created, the data is only initialized in the local CPU memory (aka RAM), but in order to use it in the GPU we’ll have to map the data into the GPU memory.

3. Initialise the Kompute Tensors in the GPU

Now that we have our Tensors created with local data, we will map the data into the GPU. For this we will use the eval_tensor_create_def, which will initialize the underlying Vulkan buffer and GPU memory, and perform the respective mapping into the GPU.

# ...previous code blocks

mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out])

# ...latter code blocks

4. Define the code to run on the GPU

Now that we’ve initialized the necessary Kompute Tensor components and they are mapped in GPU memory, we can add the Kompute Algorithm that will be executed in the GPU. This is referred to as the “shader” code, which we build using the pyshaderlibrary. You can see the full shader code below, and we’ll break down each of the section below.

# ...previous code blocks

@ps.python2shader
def compute_shader_multiply(index=("input", "GlobalInvocationId", ps.ivec3),
                            data1=("buffer", 0, ps.Array(ps.f32)),
                            data2=("buffer", 1, ps.Array(ps.f32)),
                            data3=("buffer", 2, ps.Array(ps.f32))):
    i = index.x # Fetch the current run index being processed
    data3[i] = data1[i] * data2[i]

# ...latter code blocks

The GPU shader code can be defined as a Python function with the decorator @ps.python2shader , and the parameters in this case include the variables that we’ll be using. This includes the Tensor inputs and outputs that we’ll be processing — the parameter format is the following:

<param≥=(“<memory>”, <binding>, <type>, ...)

In this case we are using Tensors with float values, which inherently would be equivalent to the ps.Array value, with ps.f32 float values as elements.

The first parameter index is of type GlobalInvocationId , and provides the shader with the current index location in the execution GPU dispatch structure. This is what allows us to know what index in the parallel execution loop we are currently running, which is what we extract from the component i = index.x — the reason why here we select x is because the execution index can be defined as a vec3 component, where there would be execution indices for inedx.x , index.y and index.z .

The final component is the actual equation used, which in this case is a simple multiplication of the first and second parameter, and stored in the output (third) parameter.

5. Dispatch GPU shader execution against Kompute Tensors

In order to run the shader above we will use the eval_algo_data_deffunction. The parameters required for this Kompute Operation includes the Tensors to bind into the GPU instructions, as well as the GPU shader code that we defined in the Python function above.

# ...previous code blocks

mgr.eval_algo_data_def(
  [tensor_in_a, tensor_in_b, tensor_out],
  compute_shader_multiply.to_spirv())

# ...latter code blocks

It’s worth mentioning that Kompute allows the user to also pass the shader as a raw glsl string, or alternatively a file path to a SPIR-V binary or raw glsl/hlsl file. For context, SPIR-V is the intermediate representation that GPUs can use to process relevant operations.

6. Use Kompute Operation to map GPU output data into local Tensors

Once the algorithm runs successfully, the result data will now be we held in the GPU memory of our output tensor. We can now use the function eval_tensor_sync_local_def to sync the Tensor GPU memory into the local tensor.

# ...previous code blocks

mgr.eval_tensor_sync_local_def([tensor_out])

# ...latter code blocks

7. Print your results

Finally, we can print the output data of our tensor.

# ...previous code blocks

print(tensor_out.data()) # prints [2.0, 4.0, 6.0]

# ...latter code blocks

When you run this, you will see the values of your output tensor printed. That’s it, you’ve written your first Kompute!

Although it may not seem obvious, the above introduced some intuition around core concepts and design thinking in GPU computing, whilst still abstracting a couple of the more in-depth concepts. In the following sections we will be providing more concrete terminology and at the end we’ll also outline a set of articles to dive into if you’re interested to learn more.

Diving into the Machine Learning intuition

Now we’ll look into the more advanced GPU compute use-case, specifically implementing the “hello world of machine learning”: logistic regression. Before we cover the implementation we will provide some intuition on the concepts and the terminology that we’ll be using throughout the following sections.

In machine learning we always have two stages, training and inference. In the diagram below you can see the two simplified flows. At the top is the training flow, where you identify some data, extract some features, and train a model until you are happy with the accuracy. Once you have a trained model, you persist the model “weights” and deploy the model into the second workflow, where the model would perform inference on unseen data.

Data Science Process (Image by Author)

In this case we will have an input dataset X , where each element is a pair xi and xj . Our input data will be the following:

xi = { 0, 1, 1, 1, 1 }
xj = { 0, 0, 0, 1, 1 }

With this input data, the expected target value Y to be predicted will be the following:

Y = {0, 0, 0, 1, 1}

Logistic Regression Example from DS Central

Our core objective in machine learning is to learn using this training data to find the function (and parameters) that will allow us to predict values Y from new “previously unseen” inputs.

It’s worth noting that the predicted values will be defined as ŷ , which are specifically the values computed with our “prediction” function, distinct to the “true” or “actual” values of Y that we defined above.

The functions that we will be using for logistic regression will be the following:

Let’s break down this function:

z — is our linear mapping function
ŷ —is the resulting predicted outputs
Xᵀ —Transpose of the matrix of vectors we’ll represent as x_i and x_j
σ — The sigmoid function which is covered in more detail below

And the parameters that we’ll be looking to learn with our machine learning algorithm are:

w— The weights that will be applied to the inputs
b — The bias that will be added

There is also the surrounding function σwhich is the sigmoid function. This function forces our input to be closer to 0 or 1, which could be intuitively seen as the probability of our prediction to be “true” or “false”, and is defined as following:

This is now the prediction/inference function that will allow us to process predictions from new data points. If we say for example that we have a new unseen set of inputs X = { (0, 1) }, and we assume that the learned parameters were W = (1, 1), b = 0after running our machine learning algorithm through our training data (which we’ll do later on), then we’ll be able to run this through our prediction function by substituting the values as follows:

In this case the prediction is 0.73..., which would be a positive prediction. This of course is just to demonstrate what our inference function will look like once we learn the parameters W and b.

Gradient descent visualized from ML Academy

The way that we will be learning the parameters is by performing a prediction, calculating the error, and then re-adjusting the weights accordingly. The method used to “re-adjust” the weights based on the “prediction error” will be done by leveraging gradient descent. This will be repeated multiple times to find more accurate parameters.

For this we will need to use the derivatives of each of the formulas. The first one, which is the derivative of our linear mapping function z is using the partial derivatives of the variables w, zand b.First, the partial derivative ∂z:

∂z = z(X) — y

Where the variables are defined as follows:

∂z — The partial derivative of the linear mapping function z(x)
z(X) — the result of the linear mapping function applied to input x
y — the actual value label expected for that input x

Similarly the derivatives for w and b respectively are the following:

∂w = (X — ∂z)/m
∂b = ∂z/m

In this case m is the total number of input elements.

We will now be able to re-adjust the parameters using the above as follows:

w = w — θ · ∂w
b = b — θ · ∂b

In this case θ is the learning rate, which as the name suggests controls the ratio by which the parameters will be modified on each iteration. Intuitively, the smaller, the more iterations it will be required for the algorithm to converge, however if the learning rate is too big, it will overshoot, leading to never being able to converge (from the image above you can imagine it will keep bouncing from side to side never reaching the bottom).

In order for us to calculate loss, we will be using the log loss function, known also as cross-entropy loss function. This function is defined as follows:

Log loss (cross entropy loss) function

Intuitive diagram to visualize cost function from ML Mastery

The function itself is set up such that the larger the difference between the predicted class and the expected class, the larger the error (you can see how much it punishes if the predicted class is on the complete different label).

The loss function will provide us an idea of the improvement of our algorithm across iterations.

Finally, one of the most important points here will be the intuition behind how we can leverage the parallel architecture of the GPU to optimize computation. In this case, we’ll be able to do it by processing multiple input parameters at the same time, referred to as a micro-batch, and then re-adjusting the parameters in batch. This is known as data-parallelization, and is one of many techniques available. In the next section we will see how this is implemented, namely passing a mini-batch of inputs, storing the weights, and then re-adjusting them before the next iteration.

Note: In this post we won’t delve into much detail, nor best practices on machine learning, however at the end of the article we will be listing a broad range of sources for people interested to take their machine learning (or GPU compute) knowledge to the next level.

Now that we have covered some of the core concepts, we will be able to learn about the implementation.

Machine Learning GPU Shader Implementation

First we will start with the GPU compute shader, which is the code that will be executed in the GPU. The full shader is outlined below, and we’ll be breaking down each section in detail to explain what each part is doing.


@ps.python2shader
def compute_shader(
        index   = ("input", "GlobalInvocationId", ps.ivec3),
        x_i     = ("buffer", 0, ps.Array(ps.f32)),
        x_j     = ("buffer", 1, ps.Array(ps.f32)),
        y       = ("buffer", 2, ps.Array(ps.f32)),
        w_in    = ("buffer", 3, ps.Array(ps.f32)),
        w_out_i = ("buffer", 4, ps.Array(ps.f32)),
        w_out_j = ("buffer", 5, ps.Array(ps.f32)),
        b_in    = ("buffer", 6, ps.Array(ps.f32)),
        b_out   = ("buffer", 7, ps.Array(ps.f32)),
        l_out   = ("buffer", 8, ps.Array(ps.f32)),
        M       = ("buffer", 9, ps.Array(ps.f32))):

    i = index.x # Fetch the current run index being processed

    m = M[0]

    w_curr = vec2(w_in[0], w_in[1])
    b_curr = b_in[0]

    x_curr = vec2(x_i[i], x_j[i])
    y_curr = y[i]

    z_dot = w_curr @ x_curr
    z = z_dot + b_curr
    y_hat = 1.0 / (1.0 + exp(-z))

    d_z = y_hat - y_curr
    d_w = (1.0 / m) * x_curr * d_z
    d_b = (1.0 / m) * d_z

    loss = -((y_curr * log(y_hat)) + ((1.0 + y_curr) * log(1.0 - y_hat)))

    w_out_i[i] = d_w.x
    w_out_j[i] = d_w.y
    b_out[i] = d_b
    l_out[i] = loss

1. Define input and output parameters

First we define all input parameters that are analogous to the input and output components we mentioned in the previous sections.

@python2shader
def compute_shader(
        index   = ("input", "GlobalInvocationId", ivec3),
        x_i     = ("buffer", 0, Array(f32)),
        x_j     = ("buffer", 1, Array(f32)),
        y       = ("buffer", 2, Array(f32)),
        w_in    = ("buffer", 3, Array(f32)),
        w_out_i = ("buffer", 4, Array(f32)),
        w_out_j = ("buffer", 5, Array(f32)),
        b_in    = ("buffer", 6, Array(f32)),
        b_out   = ("buffer", 7, Array(f32)),
        l_out   = ("buffer", 8, Array(f32)),
        M       = ("buffer", 9, Array(f32))):

    # ... latter code blocks

If you remember, at the end of the last section we mentioned how we will be leveraging the concept of micro-batches in order to use the parallel architecture of GPU processing. What this means in practice, is that we will be passing multiple instances of X to the GPU to process at a time, instead of expecting the GPU to process it one by one. This is why we see that above we have an array for xi, xj, y, wOuti, wOutj,andbOut respectively.

In more detail:

The input Xas arrays x_i and x_j will hold the micro-batch of inputs
The array ywill hold all the expected labels for micro-batch inputs
The two input weight parameters w_in_i and w_out_j will be used for calculating predictions
The input parameter b which will be used for calculating the predictions
The output weights w_out_i and w_out_jcontains weights and will store the derivative of W for all micro-batches that should be subtracted
Similarly the output bias array contains the derivatives of bfor all micro-batches that should be subtracted in batch
Finally l_out contains the output array where losses will be returned

2. Define the size of the input buffers as M

We also receive the constant M, which will be the total number of elements — if you remember this parameter will be used for the calculation of the derivatives. We will also see how these parameters are actually passed into the shader from the Python Kompute side.

        # ... previous code blocks
        M       = ("buffer", 9, Array(f32))):

    m = M[0]

    # ... latter code blocks

Now that we have all the input and output parameters defined, we can start defining the core logic, which will contain the implementation of our machine learning training algorithm.

3. Keep track of the execution index

We will need to keep track of the current index of the global invocation. Since the GPU executes in parallel, each of these runs will be running directly in parallel, so this allows the current execution to consistently keep track of what iteration index is currently being executed.

    # ... previous code blocks

    i = index.x

    # ... latter code blocks

4. Define the variables from the input parameters

We now can start preparing all the variables that we’ll be using throughout the algorithms. All our inputs are buffer arrays, so we’ll want to store them in vec2and float32 variables.

    # ... previous code blocks

    w_curr = vec2(w_in[0], w_in[1])
    b_curr = b_in[0]

    x_curr = vec2(x_i[i], x_j[i])
    y_curr = y[i]

    # ... latter code blocks

In this case we’re basically making explicit the variables that are being used for the current “thread run”. The GPU architecture consists of slightly more nuanced execution structures that involve thread blocks, memory access limitations, etc — however we won’t be covering these in this article.

Now we get into the more fun part — implementing the inference / predict logic. Below we will implement the inference logic to calculate ŷ, which involves both the linear mapping function, as well as the sigmoid function which we defined above.

    # ... previous code blocks

    # Inference and sigmoid logic
    z_dot = w_curr @ x_curr
    z = z_dot + b_curr
    y_hat = 1.0 / (1.0 + exp(-z))

    # ... latter code blocks

5. Calculate derivatives to “re-adjust” parameters

Now that we have y_hat, we can now use it to calculate the derivatives (∂z, ∂w and ∂b), which in this case are the derivative of the currently-executed index input element.

    # ... previous code blocks

    d_z = y_hat - y_curr
    d_w = (1.0 / m) * x_curr * d_z
    d_b = (1.0 / m) * d_z

    # ... latter code blocks

6. Calculate the loss from the current iteration

Using the expected prediction output and the calculated prediction output we are now able to compute the loss for the current iteration. As covered above, we are using the log loss (cross entropy) function to calculate the loss.

    # ... previous code blocks

    loss = -((y_curr * log(y_hat)) + ((1.0 + y_curr) * log(1.0 - y_hat)))

    # ... latter code blocks

7. Store the data on the output parameters

Finally we are able to pass all respective calculated metrics to our output buffers. This will allow us to re-adjust for the next iteration.

    # ... previous code blocks

    w_out_i[i] = d_w.x
    w_out_j[i] = d_w.y
    b_out[i] = d_b
    l_out[i] = loss

We’ve now finished the shader that will enable us to train a Logistic Regression algorithm in the GPU —we will now cover the rest of the logic that will call this shader and orchestrate the machine learning training and inference. The full script is outlined below, and you can also try it in the Google Colab notebook with a GPU.

Machine Learning Orchestration from Kompute

We will be using a few more advanced components from Kompute, which can be more intuitively visualised in the diagram below.

Kompute Architecture Design (Image by Author)

At the core of Kompute are Kompute “Sequences” and “Operations”, which are used for GPU actions. A Kompute Section can record and execute a batch of Kompute Operations for more efficient processing. In this example we will be leveraging Sequences to manage more efficient execution of the machine learning processing.

Similar to the example above, we will be will setting up the following steps:

Create Kompute Manager with the device explicitly defined
Create all the Kompute Tensors required
Execute the Kompute Tensor GPU initialization via Kompute Manager
Create Kompute Sequence and record operations for execution
Iterate 100 times: Run micro-batch execution & update weights
Print resulting parameters to use for future inference

As you can see this is more involved than the simpler example we used above. In this case we will use the Kompute Sequence instead of the Kompute Manager directly, as we want to have deeper control on the commands that can be recorded to send in batch to the GPU. We will discuss this in more detail as we cover each of the steps. Let’s get started.

1. Create Kompute Manager with the device explicitly defined

We will be creating the Kompute Manager with the device 0 explicitly defined — you can define another device as required.

# ... previous code blocks

mgr = kp.Manager(0)

# ...latter code blocks

2. Create all the Kompute Tensors required

Now we’ll be creating all the tensors required. In this sub-section you will notice that we will be referencing all the buffers/arrays that are being used in the shader. We’ll also cover how the order in the parameters passed relates to the way data is bound into the shaders so it’s accessible.

# ... previous code blocks

tensor_x_i = kp.Tensor([0.0, 1.0, 1.0, 1.0, 1.0])
tensor_x_j = kp.Tensor([0.0, 0.0, 0.0, 1.0, 1.0])

tensor_y = kp.Tensor([0.0, 0.0, 0.0, 1.0, 1.0])

tensor_w_in = kp.Tensor([0.001, 0.001])
tensor_w_out_i = kp.Tensor([0.0, 0.0, 0.0, 0.0, 0.0])
tensor_w_out_j = kp.Tensor([0.0, 0.0, 0.0, 0.0, 0.0])

tensor_b_in = kp.Tensor([0.0])
tensor_b_out = kp.Tensor([0.0, 0.0, 0.0, 0.0, 0.0])

tensor_l_out = kp.Tensor([0.0, 0.0, 0.0, 0.0, 0.0])

tensor_m = kp.Tensor([ tensor_y.size() ])

# ...latter code blocks

We also store them in a list params for easier access:

# ... previous code blocks

params = [tensor_x_i, tensor_x_j, tensor_y, tensor_w_in, tensor_w_out_i,
    tensor_w_out_j, tensor_b_in, tensor_b_out, tensor_l_out, tensor_m]

# ...latter code blocks

3. Execute the Kompute Tensor GPU initialization via Kompute Manager

The Kompute Tensor initialisation is quite standard so we’ll be able to do this step directly through the manager as we did in the simple array multiplication example previously.

# ... previous code blocks

mgr.eval_tensor_create_def(params)

# ...latter code blocks

4. Create Kompute Sequence and record operations for execution

In this section we will want to clear the previous recordings of the Kompute Sequence and begin recording a set of sequences. You will notice that unlike the previous section, in this case we won’t be running the eval() straight away as we’ll have to first record the operations.

You will also notice that we will be recording three types of Kompute Operations through separate functions:

# ... previous code blocks

# Clear previous operations and begin recording for new operations
sq.begin()

# Record operation to sync memory from local to GPU memory
sq.record_tensor_sync_device([tensor_w_in, tensor_b_in])

# Record operation to execute GPU shader against all our parameters
sq.record_algo_data(params, compute_shader.to_spirv())

# Record operation to sync memory from GPU to local memory
sq.record_tensor_sync_local([tensor_w_out_i, tensor_w_out_j, tensor_b_out, tensor_l_out])

# Stop recording operations
sq.end()

# ... latter code blocks

record_tensor_sync_device(...) — This operation ensures that the Tensors are synchronized with their GPU memory by mapping their local data into the GPU data. In this case, these Tensors use Device-only memory for processing efficiency, so the mapping is performed with a staging Tensor inside the operation (which is re-used throughout the operations for efficiency). Here we’re only wanting to sync the input weights, as these will be updated locally with the respective derivatives.
record_algo_base_data(...) — This is the Kompute Operation that binds the shader that we wrote above with all the local CPU/host resources. This includes making available the Tensors. It’s worth mentioning that the index of the tensors provided as parameters is the order in which they are mapped in the shaders via their respective bindings.
record_tensor_sync_local(...)— This Kompute Operation performs a similar set of instructions as the sync operation above, but instead of copying the data to the GPU memory, it does the converse. This Kompute Operation maps the data in the GPU memory into the local Tensor vector so it’s accessible from the GPU/host. As you can see we’re only running this operation in the output tensors.

5. Iterate 100 times: Run micro-batch execution & update weights

Now that we have the command recorded, we can start running executions of these pre-loaded commands. In this case, we will be running the execution of a micro-batch iteration, followed by updating the parameters locally, so they are used in the following iteration.

# ... previous code blocks

# Perform machine learning training and inference across all input X and Y
for i_iter in range(ITERATIONS):

    # Execute an iteration of the algorithm
    sq.eval()

    # Calculate the parameters based on the respective derivatives calculated
    for j_iter in range(tensor_b_out.size()):
        tensor_w_in[0] -= learning_rate * tensor_w_out_i.data()[j_iter]
        tensor_w_in[1] -= learning_rate * tensor_w_out_j.data()[j_iter]
        tensor_b_in[0] -= learning_rate * tensor_b_out.data()[j_iter]

# ... latter code blocks

7. Print resulting parameters to use for future inference

We now have a trained logistic regression model, or at least we’ve been able to optimize its respective function to identify suitable parameters. We are now able to print these parameters and use the parameters for inference in unseen datasets.

# ... previous code blocks

print(tensor_w_in.data()) # prints [0.00032, 1.58746]
print(tensor_b_in.data()) # prints [-0.794009]

# ... latter code blocks

And we’re done!

You are able to find this entire example in the example repository, which you’ll be able to run and extend.

What next?

Congratulations, you’ve made it all the way to the end! Although there was a broad range of topics covered in this post, there is a massive amount of concepts that were skimmed through. These include the underlying Vulkan concepts, GPU computing fundamentals, machine learning best practices, and more advanced Kompute concepts. Luckily, there are a broad range of resources online to expand your knowledge on each of these. Some links I recommend as further reading include the following:

Kompute Documentation for more details and further examples
The Machine Learning Engineer Newsletter if you want to keep updated on articles around Machine Learning
Awesome Production Machine Learning list for open source tools to deploy, monitor, version and scale your machine learning
Introduction to ML for Coders course by FastAI to learn further machine learning concepts

Image by Author

Parallelizing GPU-intensive Workloads via Multi-Queue Operations

Sun, 18 Oct 2020 00:00:00 GMT

GPUs have proven extremely useful for highly parallelizable data processing use-cases. The computational paradigms found in machine learning & deep learning for example fit extremely well to the processing architecture graphics cards provide.

However, when it comes to multiple GPU workloads, one would assume that these would be processed concurrently, but this is not the case. Whilst a single GPU compute workload is parallelized across the numerous GPU cores, multiple workloads are run one by one sequentially. That is of course until recent improvements in graphics card architectureswhich are now enabling for hardware parallelization across multiple workloads. This can be achieved by submitting the workloads to different underlying physical GPU “queue families” that support concurrency. Practical tecniques in machine learning that would benefit from this include model parallelism and data parallelism.

In this example we will show how we can achieve a 2x performance improvement on a synchronous example by simply submitting multiple workloads across two queue families, resulting in these workloads running in parallel.

This is an important optimization technique, as recent announcements outlined in NVIDIA’s Ampere GA10x architecture specifications page 19 in this document will enable for 3x performance improvements(i.e. concurrency across one graphics queue and two compute queues), making it clear that this trend will only continue to bring further optimization improvements in this area.

We will be implementing this using Vulkan and the Kompute framework. More specifically we will cover:

Disambiguation of “asynchronous” and “parallel” in GPU processing
A base synchronous example that we will build upon
Steps to extend the example for asynchronous workload submission
Steps to extend the example for parallel multi-queue GPU processing

You can find the full code in this file — instructions on how to run the full suite using CMAKE can be found in the main Kompute repository build section.

About Vulkan and Kompute

Khronos Membership (Image by Vincent Hindriksen via StreamHPC)

The Vulkan SDKis an Open Source project led by the Khronos Group, which enables for highly optimized cross-vendor/cross-platform GPU processing.

Kompute is a framework built on top of the Vulkan SDK which abstracts the thousands of lines of boilerplate code required, introducing best practices that expose Vulkan’s core computing capabilities. Kompute is the GPGPU computing framework that we will be using in this tutorial to build the core asynchronous and parallel code implementation.

The “Komputer” from Kompute Repo (Image by Author)

Asynchronous vs Parallel Processing

Before diving into the code, it is important to disambiguate two concepts — asynchronous workload submission and parallel workload processing.

Simplified Vulkan Architecture (Image by Author)

The way parallel workloads are submitted for processing when using the Vulkan SDKthrough GPU Queues. This can be visualised in the simplified Vulkan Architecture diagram (pipeline and descriptor components were left out for simplicity).

Asynchronous Workload Submission

Asynchronous processing encompasses the ability for the CPU host side to be able to do other work whilst the GPU is processing the workload. “Other work” can include calling other C++ functions, or even submitting further workloads to the same or other GPU queues. When the CPU wants to check whether the GPU workload is finished, it can use a **Vulkan “Fence”**which is basically a semaphore resource that allows the CPU to be notified when a GPU workload finishes.

It is important to note that when multiple workloads are submitted to the same queue, even if these are done from multiple C++ threads, the expected execution ordering will still be sequential — at least as of today’s GPU architectures.

Parallel Workload Processing

Parallel workload processing consists of the concurrent execution of two or more workloads by the GPU. More specifically, if you had two GPU tasks that would take 10 seconds each to process, the theoretical parallel execution would still take 10 seconds for both as they would be carried out at the same time.

In order for parallel workload processing to be achieved, this is something that first and foremost has to be supported by the underlying GPU. The reason why this is important is because even if you were to submit workloads across different GPU queues, the processing may still be done sequentially by the underlying hardware based on its limitations.

Base Sequential Processing Example

We will now take a look at the code that we will be using throughout this article. This first version of the code will be the sequential flow — we will be able to then convert it into asynchronous code, and finally into parallel code. We will basically be running a workload where we will be doing the following:

Creating a Kompute Manager to orchestrate all GPU work
Create Kompute Tensors in the CPU host that will be used to process data
Map the Kompute Tensors into GPU Device memory
Define compute shader which will keep the GPU busy for a few 100s ms
Run compute shader in the GPU using the Tensors for data processing
Map results of the Kompute Tensors back into CPU Host memory
Verify that the operation was successful

For measuring time we will be using <chrono> from the standard library. We will be mainly using it to calculate the difference across a start and end time retrieved with std::chrono::high_resolution_clock::now() as follows:

int main() {
    auto startSync = std::chrono::high_resolution_clock::now();

    // ... code implementation

    auto end = std::chrono::high_resolution_clock::now();
    auto durationSync =
      std::chrono::duration_cast<std::chrono::microseconds>(
        end - start).count();
}

You can find the runnable code in this file, which is part of the Kompute test suite.

1. Creating a Kompute Manager to orchestrate all GPU work

First we have to create the Kompute Manager, which performs all the required memory management and creates all required Vulkan resources. By default the Kompute Manager will pick GPU Device 0, but you are able to pass the specific device index you would prefer to initialise with, and if preferred you can pass your Vulkan resources if you already have a Vulkan application.


int main()
{
    // ... previous code blocks

    kp::Manager mgr; // Selects device 0 and first compute queue unless explicitly requested

    // ... latter code blocks

2. Create the Kompute Tensors in CPU host that will be used to process data

We will now be able to create a set of Kompute Tensors. We first initialise the data in the CPU Host, consisting of an array of zeros with length of 10. We will be using two tensors as we’ll be running two algorithm executions. We will be able to check these Kompute Tensors at the end to confirm that the execution has been successful.


int main()
{
    // ... previous code blocks

    std::vector<float> zeros(10, 0);

    // 1. Create a set of data tensors in host memory for processing
    auto tensorA = std::make_shared<kp::Tensor>(kp::Tensor(zeros));
    auto tensorB = std::make_shared<kp::Tensor>(kp::Tensor(zeros));

    // ... latter code blocks

3. Map the Kompute Tensors into GPU Device memory

Stanford CS149 Course2019 Slides

We are now able to copy the host data of the Kompute Tensors into the GPU Device memory.


int main()
{
   // ... previous code blocks

    mgr.evalOpDefault<kp::OpTensorCreate>({ tensorA, tensorB });

    // ... latter code blocks

This is an important step as by default the Kompute Tensors use device-only-visible memory which means that a GPU operation will need to copy it with a staging tensor.

Kompute allows us to create the buffer and GPU memory block, as well as performing a copy with a staging buffer through the kp::OpTensorCreate operation.

4. Define compute shader which will keep the GPU busy for a few 100s ms

The compute shader that we create has a relatively large loop to simulate an “expensive computation”. It basically performs a unit addition for 100000000 iterations and adds the result to the input Tensor.


int main()
{
    // ... previous code blocks

    std::string shader(R"(
        #version 450
        layout (local_size_x = 1) in;
        layout(set = 0, binding = 0) buffer b { float pb[]; };

        shared uint sharedTotal[1];

        void main() {
            uint index = gl_GlobalInvocationID.x;

            sharedTotal[0] = 0;
            for (int i = 0; i < 100000000; i++)
            {
                atomicAdd(sharedTotal[0], 1);
            }

            pb[index] = sharedTotal[0];
        }
    )");

    // ... latter code blocks

5. Run compute shader in the GPU using the Tensors for data processing

Now we are able to submit the compute shader for execution through the kp::OpAlgoBase operation. This basically allows us to perform a submission of the shader with the respective tensor. This initial implementation runs the execution synchronously, so it will first run the execution of the shader with tensorA, and then the execution of the same shader with tensorB.


int main()
{
    // ... previous code blocks

    mgr.evalOpDefault<kp::OpAlgoBase<>>(
        { tensorA },
        std::vector<char>(shader.begin(), shader.end()));

    mgr.evalOpDefault<kp::OpAlgoBase<>>(
        { tensorB },
        std::vector<char>(shader.begin(), shader.end()));

    // ... latter code blocks

6. Map results of the Kompute Tensors back into CPU Host memory

Finally we want to retrieve the results from the GPU device memory into the CPU host memory so we can access it from C++. For this we can use the kp::OpTensorSync operation.


int main()
{
  // ... previous code blocks

    mgr.evalOpDefault<kp::OpTensorSyncLocal>({ tensorA, tensorB });

    // ... latter code blocks

7. Verify that the operation was successful

Finally we can just check that both resulting kp::Tensor contain the expected value of 100000000.


int main()
{
    // ... previous code blocks

    std::vector<float> expected(10, 100000000)
    EXPECT_EQ(inputsSyncB[i]->data(), expected);

}

Extending for Asynchronous Workload Submission

The steps that we will need to extend for asynchronous submission in this case are quite minimal. The only thing we need to do is to substitute the evalOpDefault function for the evalOpAsyncDefault function, and then using the evalOpAwaitDefault(<timeInNanoSecs>) to wait until the job is finished. This basically would look as follows:


int main()
{
    // ... previous code blocks

    mgr.evalOpAsyncDefault<kp::OpAlgoBase<>>(
      { tensorB },
      std::vector<char>(shader.begin(), shader.end()));

    mgr.evalOpAsyncDefault<kp::OpAlgoBase<>>(
      { tensorA },
      std::vector<char>(shader.begin(), shader.end()));

    mgr.evalOpAwaitDefault();

    // ... latter code blocks

As you can see we are able to submit two tasks for processing asynchronously, and then wait until they are finished with the Await function.

It’s worth pointing out that every time we call evalOpAsyncDefault it creates a new managed sequence, and evalOpAwaitDefault only waits for the most recent default sequence. This means that in the snippet above, we are only waiting for the second asynchronous operation. This isn’t a problem for our example, but this could introduce bugs if we’re now aware. The proper way to do this is with explicitly created “named sequences” — we will do this in the next section.

Extending for Parallel Workload Processing

Now that we know we are able to execute multiple workloads asynchronously, we are able to extend this to leverage the multiple queues in the GPU to achieve parallel execution of workloads.

Running on an NVIDIA 1650 Video Card

In order to show a useful example, we will dive into how this would be achieved in an NVIDIA 1650 video card. You are able to try this yourself by checking the device report of your video card — namely on the queue families and parallel processing capabilities available.

Conceptual Overview of Queues in NVIDIA 1650 (Image by Author)

The NVIDIA 1650 GPU has 3 queue families. Using G for GRAPHICS, T for TRANSFER and C for COMPUTE capabilities, the NVIDIA 1650 has a G+T+C family in familyIndex 0with 16 queues, a T family on familyIndex 1 with 2 queues, and a T+C family on familyIndex 2 with 8 queues.

As of today (October 2020), NVIDIA does not support parallel processing of workloads when work is submitted across multiple queues within the same family. However it supports parallelizing when workloads are submitted across queue families. This means that workloads between graphics and compute family queues can be parallelized — we will be using this knowledge in our implementation.

Implementation of Parallel Workflow Execution

So far we have been submitting all GPU workloads to a single queue, namely the GRAPHICS familyIndex 0 using the underlying queue index 0. In our case using the GPU 1650, we will be able to achieve parallel processing if we submit workloads across the GRAPHICS family and the COMPUTE family. The diagram below should provide an intuition on what we will be doing.

Operation Execution in Parallel through Multiple Family Queues (Image by Author)

In order for us to do this, we will need to modify three key things:

We initialise Kompute Manager with the respective queues available
We create two Kompute Sequences with each respective queue allocated
We run the operations on each respective queue

We will dive into each of these three points.

1. We initialise Kompute Manager with the respective queues available

When initialising a manager we are able to pass an array containing the queues that we would like to fetch. In this case, we only fetch one graphics queue and one compute queue, however, based on the hardware specs of the NVIDIA 1650, we would be able to request up to 16 graphics queues (familyIndex 0), 2 transfer queues (familyIndex 1), and 8 compute queues (familyIndex 2).


int main()
{
    // ... previous code blocks

    uint32_t deviceId(0);
    std::vector<uint32_t> queues({ 0, 2 });

    kp::Manager mgr(deviceId, queues);

    // ... latter code blocks

2. We create two Kompute Sequences with each respective queue allocated

Now we are able to explicitly initialise two managed sequences, each allocated to a different queue, referencing the index of the array we passed in the previous step.


int main()
{
    // ... previous code blocks

    // The index is relative to the array we used to initialise
    mgr.createManagedSequence("graphicsQueueSequence", 0);
    mgr.createManagedSequence("computeQueueSequence", 1);

    // ... latter code blocks

3. We run the operations on each respective queue

Now we are able to run operations submitting to each respective queue. In this case both of the GPU workloads are submitted in parallel.


int main()
{
    // ... previous code blocks

    mgrAsync.evalOpAsync<kp::OpAlgoBase<>>(
      { tensorA },
      "graphicsQueueSequence",
      std::vector<char>(shader.begin(), shader.end()));

    mgrAsync.evalOpAsync<kp::OpAlgoBase<>>(
      { tensorB },
      "computeQueueSequence",
      std::vector<char>(shader.begin(), shader.end()));

    mgrAsync.evalOpAwait("graphicsQueueSequence");
    mgrAsync.evalOpAwait("computeQueueSequence");

    // ... latter code blocks

Parallel Workload Execution Results

When running the code provided above, we can see a 2x speed improvement in execution time thanks to the parallel family queue submission of workload. You can also see that if we were to submit to extra queues from the GRAPHICS or COMPUTE queues, we would not see any further speed improvements as intra-queue parallelization is not supported in this NVIDIA 1650 card.

You can find the full code and run it in this file — instructions on how to run the full suite using CMAKE can be found in the main Kompute repository.

This is a particularly important result, as based on the recent announcement from NVIDIA coming together with the release of their 300x video cards, there are improvements via the Ampere GA10x architecture that allows for two compute workloads simultaneously. Relative to the example above, this means that we could see a 3x improvement if we were to use one GRAPHICS queue and two COMPUTE queue (together with the extra performance using the TRANSFER queue for transfer operations).

Next Steps

Congratulations, you’ve made it all the way to the end! Although there was a broad range of topics covered in this post, there is a massive amount of concepts that were skimmed through. These include the underlying Vulkan concepts, GPU computing fundamentals, and more advanced Kompute concepts. Luckily, there are resources online to expand your knowledge on each of these. Here are some links I recommend for further reading:

“Machine Learning in Mobile & Cross-Vendor GPUs Made Simple With Kompute & Vulkan” article with a deeper dive in theory and concepts
Kompute Documentation for more details and further examples
Machine Learning in Mobile & Cross-Vendor GPUs Made Simple With Kompute & Vulkan
Supercharging your Mobile Apps with GPU Accelerated Machine Learning using the Android NDK & Kompute
GPU Accelerated ML using the Godot Engine and Kompute
Vulkan SDK Tutorial for a deep dive into the underlying Vulkan components

Supercharging your Mobile Apps: GPU Accelerated Machine Learning using Android NDK & Vulkan Kompute

Mon, 05 Oct 2020 00:00:00 GMT

Kompute in Android NDK (Image by Author)

Some smartphones nowadays pack laptop-level hardware — carrying up to 16GB RAM, high-speed multi-core CPUs, and GPUs that can render high-performance complex graphical applications on 4k displays.

Image by Author

Tapping into that power — especially the GPU processing power — for on-device data processing capabilities becomes growingly important as mobile hardware only continues to improve. Recently, this has been opening exciting opportunities around edge computing, federated architectures, mobile deep learning, and more.

This article provides a technical deep dive that shows you how to tap into the power of mobile cross-vendor GPUs. You will learn how to use the Android Native Development Kitand the Kompute framework to write GPU optimized code for Android devices. The end result will be a mobile app created in Android Studio that is able to use a GPU accelerated machine learning model which we will write from scratch, together with a user interface that will allow the user to send the input to the model.

Android Studio Running Project in Emulator (Image by Author)

No background knowledge beyond programming experience is required, however if you are curious about the underlying AI / GPU compute concepts referenced, we suggest checking out our previous article, “Machine Learning in Mobile & Cross-Vendor GPUs Made Simple With Kompute & Vulkan”.

You can find the full code in the example folder in the repository.

Android Native Development Kit (NDK)

Android NDK Diagram (Image by Author)

The Native Development Kit (NDK) is one of Android’s solution to address the increasing computational demands from mobile apps. The NDK framework enables developers to write low level, highly efficient C and C++ code that can interoperate with the Java/Kotlin application code through the popular Java Native Interfacebindings.

This tooling enables mobile application developers not only to write highly efficient code, but also leverage existing optimized frameworks written in C++ for advanced data processing or machine learning.

Enter Kompute & the Vulkan SDK

Playing “where’s waldo” with Khronos Membership (Image by Vincent Hindriksen via StreamHPC)

Vulkanis an Open Source project led by the Khronos Group, a consortium consisting of several tech companies who have come together to work towards defining and advancing the open standards for mobile and desktop media (and compute) technologies.

A large number of high profile (and new) frameworks have been adopting Vulkan as their core GPU processing SDK. The Android NDK’s main documentation page has a full section dedicated to Vulkan, together with hands on examples showing how it can be used in Android mobile devices.

As you can imagine, the Vulkan SDK provides very low-level access to GPUs, which allows for very specialized optimizations. This is a great asset for data processing and GPU developers — the main disadvantage is the verbosity involved, requiring 500–2000+ lines of code to only get the base boilerplate required to even start writing the application logic. This can result in expensive developer cycles and errors that can lead to larger problems. This was one of the main motivations for us to start the Komputeproject.

Kompute is a framework built on top of the Vulkan SDK which abstracts a lot of boilerplate code required, introducing best practices that expose Vulkan’s computing capabilities. Kompute is the GPU computing framework that we will be using in this tutorial to build the machine learning module in our mobile Android app.

Kompute Documentation (Image by Author)

Machine Learning in Mobile Development

In this post we will be building upon the Machine Learning use-case we created in the “Machine Learning in Mobile & Cross-Vendor GPUs Made Simple With Kompute & Vulkan” article. We will not be covering the underlying concepts in as much detail as in that article, but we’ll still introduce the high level intuition required in this section.

To start with, we will need an interface that allows us to expose our Machine Learning logic, which will require primarily two functions:

train(…)— function which will allow the machine learning model to learn to predict outputs from the inputs provided
predict(...)— function that will predict the output of an unknown instance.

This can be visualised in the two workflows outlined in the image below.

Data Science Process (Image by Author)

Particularly in app development, this would also be a common pattern for machine learning workflows, for both predictive and explanatory modelling use cases. This often consists of leveraging data generated by your users as they interact directly (or indirectly) with the app itself. This data can then serve as training features for machine learning models. Training of new models can be performed through manual “offline” workflows that data scientists would carry out, or alternatively through automated triggers retraining models.

Android Studio Project Overview

Project File Structure (Image by Author)

We will start by providing a high level overview of the core components in the Android Studio project, including the Kompute C++ bindings, the Android User Interface, the App logic build in Kotlin, and the build files required. If you are interested in a particular area you can jump to its respective section below.

You will need to make sure you install Android Studio, and also install the Android NDK— the rest of the dependencies will be installed and configured automatically when opening the project in the IDE.

Now that you have everything installed, you are able to import the project. For this, you first have to clone the full Kompute repository and import the Android Studio project under examples/android/android-simple/. You should now be able to see the project load and configure the build. Once it opens you are able to run it in an emulator or in your own physical Android phone. This project was tested in the Pixel 2 emulator, and in a physical Samsung Galaxy phone.

Final GPU Accelerated Kompute App (Image by Author)

When you load the project you will notice the following key components in the file structure, which we will be breaking down further in detail in the following sections:

Android SDK Application— The Android UI, asset resources, build files, and Kotlin/Java components that provide the relevant application logic that interacts with the UI and C++ Kompute ML bindings.
Android NDK Kompute ML Module— The Kompute ML C++ code and bindings configured through Android NDK for GPU accelerated processing.

Android SDK Application

This section covers the Android SDK Java/Kotlin and User Interface components, which should provide an intuition on how the high level business logic interacts with the native C++ JNI bindings.

The user interface consists primarily of input text boxes and display text labels to enable users to interact with the GPU Accelerated ML processing C++ core (as shown in the screenshot below). If you are curious on exactly the views used, you can inspect it in your Android Studio project, or directly open the activity_kompute_jni.xmlfile.

The core of our mobile app can be found in the app/src/main/java/com/ethicalml/kompute folder, inside the KomputeJni.ktfile. This Kotlin file contains the main business logic for our Android app.

If we look at the shortened version of the class in the code block below we will notice the following key points:

// imports

class KomputeJni : AppCompatActivity() {

    override fun onCreate(savedInstanceState: Bundle?) {
        // ...
    }

    fun KomputeButtonOnClick(v: View) {
        // ...
    }

    external fun initVulkan(): Boolean

    external fun kompute(xi: FloatArray, xj: FloatArray, y: FloatArray): FloatArray

    external fun komputeParams(xi: FloatArray, xj: FloatArray, y: FloatArray): FloatArray

    companion object {
        init {
            System.loadLibrary("kompute-jni")
        }
    }
}

fun onCreate(…) — This function is called on initialisation of the Android Activity (when the app is loaded)
fun KomputeButtonOnClick(…)— This function is triggered when the main “KOMPUTE” button gets pressed, and triggers the C++ Jni binding functions using the data from the user interface text boxes.
external fun initVulkan(): Boolean — This function is one of the C++ JNI functions that will be bound to the Vulkan initialisation C++ function.
external fun kompute(...): FloatArray — The is the C++ JNI function that will train the ML model and run inference on the inputs provided, returned the inference results.
external fun komputeParams(...): FloatArray — The C++ JNI function that trains the model and returns the learned parameters weight 1 , weight 2 and bias .
companion object { ...("kompute-jni") } — This is the name you will give to your C++ output shared library, which will contain the compiled C++ source with all relevant binding functions.

As you will also notice, the external fun functions do not have any definition — this is because the definition is provided in the C++ JNI function bindings, which will be covered in the C++ JNI bindings section.

Now to cover each of the functions in more detail, we start with the onCreate function. This function is in charge of initialising all relevant components in our application. This includes:

val binding — This is the main object that will allow us to access all the text boxes and elements in the UI.
val successVulkanInit = initVulkan() — This is our first call to a C++ JNI function, which is primarily in charge of initialising Vulkan. If it’s successful it returns true, and we display the respective success message in a android.widget.Toast popup — an error is displayed otherwise.


// ... prior code blocks

class KomputeJni : AppCompatActivity() {

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)

        val binding = ActivityKomputeJniBinding.inflate(layoutInflater)
        setContentView(binding.root)

        binding.komputeGifView.loadUrl("file:///android_asset/komputer-2.gif")

        binding.komputeGifView.getSettings().setUseWideViewPort(true)
        binding.komputeGifView.getSettings().setLoadWithOverviewMode(true)

        val successVulkanInit = initVulkan()
        if (successVulkanInit) {
            Toast.makeText(applicationContext, "Vulkan Loaded SUCCESS", Toast.LENGTH_SHORT).show()
        } else {
            binding.KomputeButton.isEnabled = false
            Toast.makeText(applicationContext, "Vulkan Load FAILED", Toast.LENGTH_SHORT).show()
        }
        Log.i("KomputeJni", "Vulkan Result: " + successVulkanInit)

        binding.predictionTextView.text = "N/A"
    }

// ... latter code blocks

Next up we have the KomputeButtonOnClick(...) function. This function gets triggered when the user presses the “KOMPUTE” button in the app. The main purpose of this function is to retrieve the inputs from the input text boxes in the UI, then use the input data to perform an ML train/inference step through the JNI C++ bindings, and finally display the resulting outputs back in the UI text labels. In further detail:

// ... prior code blocks

    fun KomputeButtonOnClick(v: View) {

        val xiEditText = findViewById<EditText>(R.id.XIEditText)
        val xjEditText = findViewById<EditText>(R.id.XJEditText)
        val yEditText = findViewById<EditText>(R.id.YEditText)

        val wOneEditText = findViewById<TextView>(R.id.wOneTextView)
        val wTwoEditText = findViewById<TextView>(R.id.wTwoTextView)
        val biasEditText = findViewById<TextView>(R.id.biasTextView)

        val komputeJniTextview = findViewById<TextView>(R.id.predictionTextView)

        val xi = xiEditText.text.removeSurrounding("[", "]").split(",").map { it.toFloat() }.toFloatArray()
        val xj = xjEditText.text.removeSurrounding("[", "]").split(",").map { it.toFloat() }.toFloatArray()
        val y = yEditText.text.removeSurrounding("[", "]").split(",").map { it.toFloat() }.toFloatArray()

        val out = kompute(xi, xj, y)

        Log.i("KomputeJni", "RESULT:")
        Log.i("KomputeJni", out.contentToString())

        komputeJniTextview.text = out.contentToString()

        val params = komputeParams(xi, xj, y)

        Log.i("KomputeJni", "Params:")
        Log.i("KomputeJni", params.contentToString())

        wOneEditText.text = params[0].toString()
        wTwoEditText.text = params[1].toString()
        biasEditText.text = params[2].toString()
    }

// ... latter code blocks

val <elementname> = findViewById<EditText>(R.id.<elementname>)— This is the format of the steps that create the variable that holds the respective text box with inputs. In this case <elementname> is the name of the element we are interacting with.
xi, xj and y — The FloatArray elements created from the text in the respective input boxes, which are then used for the ML model processing.
val out = kompute(xi, xj, y) — Here we run the C++ JNI Binding function kompute which trains and processes the data through the KomputeModelML class we create in C++.
val params = komputeParams(xi, xj, y) — Here we run the C++ JNI Binding function which trains and returns the learned parameters of the Kompute machine learning model.
<elementname>.text = <value> — The lines that follow this format basically override the text labels in the UI to display the outputs.

The last few functions are only explicitly set as external functions to be bound to the C++ JNI bindings. Furthermore the companion object section provides the name of the shared library that will contain the respective bindings referenced in this activity.

// ... prior code blocks

    external fun initVulkan(): Boolean

    external fun kompute(xi: FloatArray, xj: FloatArray, y: FloatArray): FloatArray

    external fun komputeParams(xi: FloatArray, xj: FloatArray, y: FloatArray): FloatArray

    companion object {
        init {
            System.loadLibrary("kompute-jni")
        }
    }
}

You can find the full file in the KomputeJni.ktfile in the repository.

Android NDK Kompute ML Module

This section covers the Android NDK Kompute ML Module files, which includes the build-system, and the C++ source code using the Kompute framework.

Kompute Architecture Design (Image by Author)

We will be using the core components of Kompute which are outlined in this accompanying diagram. Namely, we will be loading the relevant data in the GPU using Kompute Tensors, processing it with the respective Kompute Operations, and orchestrating this with a Kompute Sequence and a Kompute Manager. We won’t be covering the Kompute architecture in detail but if you want to learn more about the underlying concepts, you can check out the more detailed article on the underlying implementation.

The core components in the Android NDK bindings module consist of the following:

JNI Binding Functions — The native functions that can be called from the Java/Kotlin Android SDK application code.
KomputeModelML Class — The class that exposes the Kompute GPU Accelerated ML model logic.
CMake build file— The C++ build file responsible for compiling and linking all relevant libraries.

JNI Binding Functions

The JNI bindings in this case are provided via the KomputeJniNative.cppfile. The skeleton of the class is below — the function code logic has been redacted for simplicity, and will be explained in more detail below.

// Imports & util functions

extern "C" {

JNIEXPORT jboolean JNICALL
Java_com_ethicalml_kompute_KomputeJni_initVulkan(JNIEnv *env, jobject thiz) {
    // ...
}


JNIEXPORT jfloatArray JNICALL
Java_com_ethicalml_kompute_KomputeJni_kompute(
        JNIEnv *env,
        jobject thiz,
        jfloatArray xiJFloatArr,
        jfloatArray xjJFloatArr,
        jfloatArray yJFloatArr) {
    // ...
}

JNIEXPORT jfloatArray JNICALL
Java_com_ethicalml_kompute_KomputeJni_komputeParams(
        JNIEnv *env,
        jobject thiz,
        jfloatArray xiJFloatArr,
        jfloatArray xjJFloatArr,
        jfloatArray yJFloatArr) {
    // ...
}

}

The JNI binding functions have to match the class functions defined in the Java/Kotlin code. The format for the function is:

Java_<modulepath>_<class>_<functionname>(env, thiz, ...params)

In our case the class is in the com.ethicalml.kompute module, in the class KomputeJni and its respective function — below the name of the functions will reflect this structure.

Diving one level deeper, we can now go through each section of the file. Starting with the imports, we can see below the imports together with comments outlining their core functionality.

// Includes the Jni utilities for Android to be able to create the
// relevant bindings for java, including JNIEXPORT, JNICALL , and
// other "j-variables".
#include <jni.h>

// The ML class exposing the Kompute ML workflow for training and
// prediction of inference data.
#include "KomputeModelML.hpp"

// Allows us to use the C++ sleep function to wait when loading the
// Vulkan library in android
#include <unistd.h>

In Android applications, we actually need to initialize the Vulkan dynamic library manually (which is something that you normally wouldn’t do outside of Android). The reason why this is required, is because the Vulkan library is not actually linked in Android phones. The reason why Android avoids doing any linking is for backwards compatibility, mainly to ensure the app doesn’t crash if the Vulkan library is not found in older phones.

This means we need to manually find the library in the C++ code and if found, link each function to its respective memory address pointer so our C++ framework can use it. Fortunately, this is something that Kompute does automatically, and we won’t be covering the details in this article as it probably would require an article in itself, but if you’re interested you can read more about it in this post, and you can see how Kompute imports Vulkan dynamically in the Core.hpp header file using the vk_ndk_wrapper_includefiles.

Below you can see the implementation of the function that exposes the initVulkan logic—Java_com_ethicalml_kompute_KomputeJni_initVulkan(...). You can see inside this function we run InitVulkan()until the Vulkan library is successfully initialised, or alternatively fails if the maximum number of retries is reached.

// ... previous code blocks

JNIEXPORT jboolean JNICALL
Java_com_ethicalml_kompute_KomputeJni_initVulkan(JNIEnv *env, jobject thiz) {

    SPDLOG_INFO("Initialising vulkan");

    uint32_t totalRetries = 0;

    while (totalRetries < KOMPUTE_VK_INIT_RETRIES) {
        SPDLOG_INFO("VULKAN LOAD TRY NUMBER: %u", totalRetries);
        if(InitVulkan()) {
            break;
        }
        sleep(1);
        totalRetries++;
    }

    return totalRetries < KOMPUTE_VK_INIT_RETRIES;
}

// ... latter code blocks

Once Vulkan has been initialised, it is possible to call the remaining functions. The first one is the kompute function, which is in charge of training the model and running an inference request. The function receives the input Xi and Xj values, together with the expected predictions that the model will learn from. It will then return the prediction treating Xi and Xj as unseen data. The function will basically call the KomputeModelML class train function and predict function.

// ... previous code blocks

JNIEXPORT jfloatArray JNICALL
Java_com_ethicalml_kompute_KomputeJni_kompute(
        JNIEnv *env,
        jobject thiz,
        jfloatArray xiJFloatArr,
        jfloatArray xjJFloatArr,
        jfloatArray yJFloatArr) {

    SPDLOG_INFO("Creating manager");

    std::vector<float> xiVector = jfloatArrayToVector(env, xiJFloatArr);
    std::vector<float> xjVector = jfloatArrayToVector(env, xjJFloatArr);
    std::vector<float> yVector = jfloatArrayToVector(env, yJFloatArr);

    KomputeModelML kml;
    kml.train(yVector, xiVector, xjVector);

    std::vector<float> pred = kml.predict(xiVector, xjVector);

    return vectorToJFloatArray(env, pred);
}

// ... latter code blocks

The last remaining JNI function that will be exposed to the Java/Kotlin code is the komputeParams function, which is in charge of returning the parameters that the machine learning model learns, namely the weight 1 , weight 2 and the bias parameters.

// ... previous code blocks

JNIEXPORT jfloatArray JNICALL
Java_com_ethicalml_kompute_KomputeJni_komputeParams(
        JNIEnv *env,
        jobject thiz,
        jfloatArray xiJFloatArr,
        jfloatArray xjJFloatArr,
        jfloatArray yJFloatArr) {

    SPDLOG_INFO("Creating manager");

    std::vector<float> xiVector = jfloatArrayToVector(env, xiJFloatArr);
    std::vector<float> xjVector = jfloatArrayToVector(env, xjJFloatArr);
    std::vector<float> yVector = jfloatArrayToVector(env, yJFloatArr);

    KomputeModelML kml;
    kml.train(yVector, xiVector, xjVector);

    std::vector<float> params = kml.get_params();

    return vectorToJFloatArray(env, params);
}

// ... latter code blocks

The only remaining functions are the utility functions that we used in the JNI logic — namely jfloatArrayToVector and vectorToJFloatArray — these functions are self explanatory, so we’ll leave it to the reader to explore further in the source if interested.

KomputeModelML Class

Now that we’ve covered the key functions that are bound to the Kotlin / Java class, we can cover the KomputeModelML C++ class that contains the Kompute GPU Accelerated logic.

The header file for the KomputeModelML class is outlined below, and contains the following key components:

#pragma once

#include "kompute/Kompute.hpp"

class KomputeModelML {

public:
    KomputeModelML();
    virtual ~KomputeModelML();

    void train(std::vector<float> yData, std::vector<float> xIData, std::vector<float> xJData);

    std::vector<float> predict(std::vector<float> xI, std::vector<float> xJ);

    std::vector<float> get_params();

private:
    kp::Tensor mWeights;
    kp::Tensor mBias;

};

static std::string LR_SHADER = R"(
//... rest of header file

#include "kompute/Kompute.hpp” — header containing all the Kompute dependencies that we’ll use in this project
void train(...) —Trains the machine learning model using the GPU native code for the logistic regression model. It takes the input array(s) X, and the array ycontaining the expected outputs.
std::vector<float> predict(...) —Perform the inference request. In this implementation it is not using GPU code as generally there tends to be less performance gains through parallelization on the inference side. However there are still expected performance gains if multiple inputs are processed in parallel (which this function allows for).
std::vector<float> get_params() —Returns an array containing the learned parameters in the format of [ <weight_1>, <weight_2>, <bias> ].
static std::string LR_SHADER— The shader code that will be executed as machine code inside of the GPU. Kompute allows us to pass a string containing the code, however for production deployments it is possible to convert the shaders to binary, and also use the utilities available to convert into header files.

If you are interested in the full implementation you can find all the files in the repository. Furthermore if you are interested in the theoretical and underlying foundational concepts of these techniques, this is covered fully in our previous post.

CMake Build File

The CMakeLists.txtbuild file is a very important component in the Android NDK workflow. This section becomes particularly important if you wish to add Kompute into your own project. The cmake file is quite small so we’ll be covering each of the lines separately.

First we need to make sure the Kompute library is available. Usually you would run the INSTALLtarget of the Kompute build to be able to use/import the respective library. However in this case we need to make sure Kompute is built for the right Android CPU architecture —our simplest option is adding the main repository as part of the build, which means that Kompute will also be built for the right mobile architectures. If you want to include this in your project, you just need to make sure the path is relative to the Kompute cloned folder.

add_subdirectory(
    ../../../../../../../
    ${CMAKE_CURRENT_BINARY_DIR}/kompute_build)

We now set the variable VK_ANDROID_INCLUDE_DIR to the vulkan include directory. This contains all the include files we need for Vulkan — for completeness, Kompute uses the vulkan.h header as well as the vulkan.hpp C++ headers.

set(
  VK_ANDROID_INCLUDE_DIR
  ${ANDROID_NDK}/sources/third_party/vulkan/src/include)

We are now able to add the library that will be used by the Java/Kotlin Android Studio project, which in this case is the shared library kompute-jni.

add_library(kompute-jni SHARED
        KomputeJniNative.cpp
        KomputeModelML.cpp)

We now are able to add all relevant include directories. This includes the VK_ANDROID_INCLUDE_DIRwhich we defined above, as well as the VK_ANDROID_COMMON_DIR which contains Android log.h . The single_include is what contains the kompute/Kompute.hpp header from Kompute. Finally, we need to import the Kompute Dynamic library wrapper vk_ndk_wrapper_include which is necessary as the Vulkan library is imported dynamically. The logic behind this could become a series of articles in itself, so we won’t go down this rabbit-hole, but if you’re interested you can read more in this post, and you can see how Kompute imports Vulkan dynamically.

include_directories(
        ${VK_ANDROID_COMMON_DIR}
        ${VK_ANDROID_INCLUDE_DIR}
        ../../../../../../../single_include/
        ../../../../../../../vk_ndk_wrapper_include/)

To compile the project we’ll want to make sure the VK_USE_PLATFORM_ANDROID_KHR is set, as this is what enables the Android configuration. For this project we also disable the Vulkan debug layers with KOMPUTE_DISABLE_VK_DEBUG_LAYERS.

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 \
                     -DVK_USE_PLATFORM_ANDROID_KHR=1 \
                     -DKOMPUTE_DISABLE_VK_DEBUG_LAYERS=1")

Finally, we are able to link the relevant libraries to our kompute-jni shared library target. This includes:

target_link_libraries(kompute-jni
        # Libraries from kompute build
        kompute
        kompute_vk_ndk_wrapper
        # Libraries from android build
        log
        android)

kompute— This is the library created in the Kompute build.
kompute_vk_ndk_wrapper — This library also gets created by the Kompute build and contains the code to dynamically load and wrap the Vulkan library.
log — This is the Android log library, which is required by Kompute to override logging.
android — This is the Android library which is used by the Android project.

That’s it — you are now able to run the application, which will execute the full build. You should then be able to see the Kompute app in your Android Studio emulator, or in you physical phone, where you’ll be able to trigger the processing on your on-device GPU.

Android Studio Running Project in Emulator (Image by Author)

What’s next?

Congratulations, you’ve made it all the way to the end! Although there was a broad range of topics covered in this post, there is a massive amount of concepts that were skimmed through. These include the underlying Android Development workflows, Vulkan concepts, GPU computing fundamentals, machine learning best practices, and more advanced Kompute concepts. Luckily, there are resources online to expand your knowledge on each of these. Here are some links I recommend for further reading:

“Machine Learning in Mobile & Cross-Vendor GPUs Made Simple With Kompute & Vulkan” article with a deeper dive in theory and concepts
“Parallelizing GPU-intensive Workloads via Multi-Queue Concurrency using Kompute” with more advanced Kompute concepts
Kompute Documentation for more details and further examples
The Machine Learning Engineer Newsletter if you want to keep updated on articles around Machine Learning
Android NDK Getting Started Documentation if you want to dive further into the Native Development Kit
Awesome Production Machine Learning list for open source tools to deploy, monitor, version and scale your machine learning
Introduction to ML for Coders course by FastAI to learn further machine learning concepts
Vulkan SDK Tutorial for a deep dive into the underlying Vulkan components

Beyond CUDA: GPU Accelerated C++ for Machine Learning on Cross-Vendor Graphics Cards Made Simple

Sat, 12 Sep 2020 00:00:00 GMT

Machine learning, together with many other advanced data processing paradigms, fits incredibly well to the parallel-processing architecture that GPU computing offers.

Image by Author

In this article you’ll learn how to write your own ML algorithm from scratch in GPU optimized code, which will be able to run in virtually any hardware — including your mobile phone. We’ll introduce core GPU & ML concepts and show how you can use the Kompute framework to implement it in only a handful lines of code.

We will be building first a simple algorithm that will multiply two arrays in parallel, which will introduce the fundamentals of GPU processing. We will then write a Logistic Regression algorithm from scratch in the GPU. You can find the repo with the full code in the following links:

Kompute Repository
GPU Array Multiplication Repository and Kompute Code
GPU Logistic Regression Repository, Kompute Code and Shader Code

Motivation

The potential and adoption of GPU computing has been exploding in recent years — you can get a glimpse of the increasing speed in adoption from the charts in the image below. In deep learning there has been a massive increase in adoption of GPUs for processing, together with paradigms that have enabled massively parallelizable distribution of compute tasks across increasing number of GPU nodes. There is a lot of exciting research around techniques that propose new approaches towards model parallelism and data parallelism— both which allow algorithms and data respectively to be sub-divided in a broad range of approaches to maximize processing efficiency.

Ben-Nun, Tal, and Torsten Hoefler. “Demystifying parallel and distributed deep learning: An in-depth concurrency analysis.” ACM Computing Surveys (CSUR) 52.4 (2019): 1–43.

In this article we outline the theory, and hands on tools that will enable both, beginners and seasoned GPU compute practitioners, to make use of and contribute to the current development and discussions across these fascinating high-performance computing areas.

The Vulkan Framework

Before diving right in, it is worth introducing the core framework that is making it possible to build hyper-optimized, cross platform and scalable GPU algorithms — and that is the Vulkan Framework.

Playing “where’s waldo” with Khronos Membership (Image by Vincent Hindriksen via StreamHPC)

Vulkan is an Open Source project led by the Khronos Group, a consortium of a very large number of tech companies who have come together to work towards defining and advancing the open standards for mobile and desktop media (and compute) technologies. On the left you can see the broad range of Khronos Members.

You may be wondering, why do we need yet another new GPU framework where there are already many options available for writing parallelizable GPU code? The main reason is that unlike some of its closed source counterparts (e.g. NVIDIA’s CUDA, or Apple’s Metal) Vulkan is fully Open Source, and unlike some of the older options (e.g. OpenGL), Vulkan is built with the modern GPU architecture in mind, providing very granular access to GPU optimizations. Finally, whilst some alternatives provide vendor-specific support for GPUs, Vulkan provides cross-platform, and cross-vendor support, which means that it opens the doors to opportunities in mobile processing, edge computing, and more.

The Vulkan SDK provides very low-level access to GPUs, which allows for very specialized optimizations. This is a great asset for GPU developers — the main disadvantage is the verbosity involved, requiring 500–2000+ lines of code to only get the base boilerplate required to even start writing the application logic. This can result not only in expensive developer cycles but also prone to small errors that can lead to larger problems.

This can actually be seen across many new and renowned machine learning & deep learning projects like Pytorch, Tensorflow, and Alibaba DNN — between others — which have either integrated or are looking to integrate the Vulkan GPU SDK to add mobile GPU (and cross-vendor GPU) support. All of these frameworks end up with very similar and extremely verbose boilerplate code, which means they would have benefited (and could still benefit) from using a unified baseline. This was one of the main motivations for us to start the Komputeproject.

Enter Kompute

Kompute is a framework built on top of the Vulkan SDK, specifically designed to extend its compute capabilities as a simple to use, highly optimized, and mobile friendly General Purpose GPU computing framework.

Kompute Documentation (Image by Author)

Kompute was not built to hide any of the core Vulkan concepts — the Vulkan API is very well designed —instead it augments Vulkan’s Computing capabilities with a BYOV (bring your own Vulkan) design, enabling developers by reducing boilerplate code required and automating some of the more common workflows involved in writing Vulkan applications.

For new developers curious to learn more, it provides a solid base to get started into GPU computing. For more advanced Vulkan developers, Kompute allows them to integrate it into their existing Vulkan applications, and perform very granular optimizations by getting access to all of the Vulkan internals when required. The project is fully open source, and we welcome bug reports, documentation extensions, new examples or suggestions — please feel free to open an issue in the repo.

Writing your first Kompute

To build our first simple array-multiplication GPU computing application using Kompute, we will create the following:

TwoKompute Tensorsto store the input data
One Kompute Tensor to store the output data
AKompute Operation to create and copy the tensors to the GPU
A Kompute Operation with an Kompute Algorithm that will hold the code to be executed in the GPU (called a “shader”)
A Kompute Operation to sync the GPU data back to the local tensors
A Kompute Sequence to record the operations to send to the GPU in batches (we’ll use theKompute Managerto simplify the workflow)

Kompute Architecture Design (Image by Author)

At the core of Kompute are “Kompute Operations”, which are used for GPU actions, as well as “Kompute Tensor” operations to handle the GPU data and memory. More specifically, this diagram shows the relationship between Kompute components (including explicit memory ownership).

When interacting with the GPU, you have to send the instructions to the GPU to execute, and you need to make sure that the GPU has all of the relevant data available in GPU memory to begin processing. With Vulkan you send these instructions to the GPU via a queue, so to simplify things intuitively you can think of your GPU as a remote server, where data serialization, resource creation and memory allocation is expensive, and instructions are submitted via a queue — there is still GPU-CPU shared memory but you tend to only use this for data transfer to the GPU.

Let’s jump into the code. Typically, in a Kompute application we’ll follow the following steps:

Create a Kompute Manager to manage resources
Create Kompute Tensors to hold data
Initialise the Kompute Tensors in the GPU with a Kompute Operation
Define the code to run on the GPU as a “compute shader”
Use Kompute Operation to run shader against Kompute Tensors
Use Kompute Operation to map GPU output data into local Tensors
Print your results

1. Create a Kompute Manager to manage resources

First, we’ll create our Kompute Manager, which is in charge of creating and managing all the underlying Vulkan resources.

// Single header include for Kompute
#include "kompute/Kompute.hpp"

int main() {

    // Vulkan resources get created unless passed
    kp::Manager mgr(0); // Selects GPU device at index 0

    //... continued in next section
}

As you can see, here we are initializing our Kompute Manager, expecting it to create all the base Vulkan resources on Device 0 (in my case Device 0 is my NVIDIA card, and Device 1 is my integrated graphics card). For more advanced use-cases it’s also possible to initialize the Kompute Manager with your own Vulkan resources (Device, Queue, etc) but this is out of scope of this article.

2. Create Kompute Tensors to hold data

We will now create the Kompute Tensors that will be used for input and output. These will hold the data required which will be mapped into the GPU to perform this simple multiplication.

int main()
{
    //... continued from previous section

    std::shared_ptr<kp::Tensor> tensorInA{ new kp::Tensor({ 2.0, 4.0, 6.0 }) };
    std::shared_ptr<kp::Tensor> tensorInB{ new kp::Tensor({ 0.0, 1.0, 2.0 }) };
    std::shared_ptr<kp::Tensor> tensorOut{ new kp::Tensor({ 0.0, 0.0, 0.0 }) };

    //... continued in next section
}

int main()
{
    //... continued from previous section

    std::vector<std::shared_ptr<kp::Tensor>> tensorParams({
        std::shared_ptr<kp::Tensor>{&tensorInA},
        std::shared_ptr<kp::Tensor>{&tensorInB},
        std::shared_ptr<kp::Tensor>{&tensorOut}
    });

    //... continued in next section
}

The reason why Kompute uses std::shared_ptr by design to avoid passing the objects by value, and instead passing them using smart pointers.

3. Initialise the Kompute Tensors in the GPU with a Kompute Operation

Now that we have our Tensors created with local data, we will map the data into the GPU. For this we will use the kp::OpTensorCreate Kompute Operation, which will initialize the underlying Vulkan buffer and GPU memory, and perform the respective mapping into the GPU.

int main()
{
    //... continued from previous section

    // Create the tensors by passing them as parameters
    mgr.evalOpDefault<kp::OpTensorCreate>({ tensorInA, tensorInB, tensorOut });

    //... continued in next section
}

It’s also worth mentioning that it’s possible to shorten the tensor creation steps by leveraging the Kompute Manager buildTensor helper function. This would allow you to skip the need to create the shared_ptr explicitly as well as the kp::OpTensorCreate Operation as outlined below (you can also find the full code implementation of this variation here).

int main()
{
    //... continued from previous section

    // Using this instead of Operations
    auto tensorInA = mgr.buildTensor({ 2.0, 4.0, 6.0 });
    auto tensorInB = mgr.buildTensor({ 0.0, 1.0, 2.0 });
    auto tensorOut = mgr.buildTensor({ 0.0, 0.0, 0.0 });

    //... continued in next section
}

4. Define the code to run on the GPU as a “compute shader”

Now that we’ve initialized the necessary Kompute Tensor components and they are mapped in GPU memory, we can add the Kompute Algorithm that will be executed in the GPU. This is referred to as the “shader” code, which follows a C-like syntax. You can see the full shader code below, and we’ll break down each of the sections below.

int main()
{
    //... continued from previous section

    // Define your shader as a string (using string literals for simplicity)
    // (You can also pass the raw compiled bytes, or even path to file)
    std::string shader(R"(
        // The version to use
        #version 450

        // The execution structure
        layout (local_size_x = 1) in;

        // The buffers are provided via the tensors
        layout(binding = 0) buffer bufA { float a[]; };
        layout(binding = 1) buffer bufB { float b[]; };
        layout(binding = 2) buffer bufOut { float o[]; };

        void main() {
            uint index = gl_GlobalInvocationID.x;

            o[index] = a[index] * b[index];
        }
    )");


    //... continued in next section
}

The #version 450 and layout(local_size_x = 1) in; sections specify the version and parallel thread execution structure (which we’ll look at further down the article). We then can see the GPU data inputs and outputs defined in the format:

layout(binding = <INDEX>) buffer <UNIQUENAME> {float <VARNAME>[]};

— index that maps Tensors to the GPU input
— This must be a unique name for the buffer
— This is the variable name to use in the shader code

These are the parameters that can be used throughout the shader code for processing. Namely in this case, the processing is done inside the mainfunction. The first variable uint index = gl_GlobalInvocationID.x; is the currently parallel execution index which will allow us to process each data input.

We then come into the core of this algorithm which is the multiplication o[index] = a[index] * b[index].This part is quite self-explanatory — we multiply the elements of the GPU arrays a[] and b[] , then store the output on the array o[].

5. Use Kompute Operation to run shader against Kompute Tensors

In order to run the shader above we will create the Kompute Operation kp::OpAlgoBase. The parameters required for this Kompute Operation includes the Tensors to bind into the GPU instructions, as well as the shader code itself.

int main()
{
    //... continued from previous section

    // Run Kompute operation on the Tensor parameters provided
    mgr.evalOpDefault<kp::OpAlgoBase<>>(
        tensorParams,
        std::vector<char>(shader.begin(), shader.end()));

    //... continued in next section
}

It’s worth mentioning that Kompute allows the user to also pass the shader through a file path, or alternatively there are also Kompute tools that will allow you to convert the shader binaries into C++ header files.

6. Use Kompute Operation to map GPU output data into local Tensors

Once the algorithm gets triggered, the result data will now be we held in the GPU memory of our output tensor. We can now use the kp::OpTensorSyncLocalKompute Operation to sync the Tensor GPU memory as per the code block below.

int main()
{
    //... continued from previous section

    // Run Kompute operation on the parameters provided with dispatch layout
    mgr.evalOpDefault<kp::OpTensorSyncLocal>({ tensorOut });

    //... continued in next section
}

7. Print your results

Finally, we can print the output data of our tensor.

int main()
{
    //... continued from previous section

    // prints "Output {  0  4  12  }"
    std::cout<< "Output: {  ";
    for (const float& elem : tensorOut->data()) {
      std::cout << elem << "  ";
    }
    std::cout << "}" << std::endl;
}

When you run this, you will see the values of your output tensor printed. That’s it, you’ve written your first Kompute!

You can also find the full example in the repo so you can run it and extend it as desired. You can find the full standalone example in this repository which includes the instructions on how to build it as well as the Kompute C++ code.

Although it may not seem obvious, the above introduced some intuition around core concepts and design thinking in GPU computing, whilst still abstracting a couple of the more in-depth concepts. In the following sections we will be providing more concrete terminology and we’ll be scratching the surface of some of the more advanced concepts such as threads, blocks, memory strides and shared memory (although a lot will be provided as further reading).

Diving into the Machine Learning intuition

Let’s look at a more advanced GPU compute use-case, specifically implementing the hello world of machine learning, logistic regression. Before we cover the implementation we will provide some intuition on the theory, and the terminology that we’ll be using throughout.

In machine learning we always have two stages, training and inference. In the diagram below you can see the two simplified flows. At the top is the training flow, where you identify some training data, extract some features, and train a model until you are happy with the accuracy. Once you have a trained model, you persist the model “weights” and deploy the model into the second workflow, where the model would perform inference on unseen data.

Data Science Process (Image by Author)

In this case we will have an input dataset X , where each element is a pair xi and xj . Our input data will be the following:

xi = { 0, 1, 1, 1, 1 }
xj = { 0, 0, 0, 1, 1 }

With this input data, the expected target value Y to be predicted will be the following:

Y = {0, 0, 0, 1, 1}

Logistic Regression Example from DS Central

Our primary objective in machine learning is to learn using this data to find the function (and parameters) that will allow us to predict values Y from just using X as input.

It’s worth noting that the predicted values are defined as ŷ , which are specifically the values computed with our inference function, distinct to the “true” or “actual” values of Y that we defined above.

The functions that we will be using for logistic regression will be the following:

Let’s break down this function:

z — is our linear mapping function
ŷ —is the resulting predicted outputs
Xᵀ —The transpose of the matrix containing our inputs xi and xj
σ — The sigmoid function which is covered in more detail below

And the parameters that we’ll be looking to learn with our machine learning algorithm are:

W— The weights that will be applied to the inputs
b — The bias that will be added

This is now the inference function that will allow us to process predictions from new data points. If we say for example that we have a new unseen set of inputs X = { (0, 1) }, and we assume that the learned parameters were W = (1, 1), b = 0after running our machine learning algorithm through our training data (which we’ll do later on), then we’ll be able to run this through our prediction function by substituting the values as follows:

Gradient descent visualized from ML Academy

For this we will need to use the derivatives of each of the formulas. The first one, which is the derivative of our linear mapping function z is:

∂z = z(x) - y

Where the variables are defined as follows:

∂z — The derivative of the linear mapping function z(x)
z(x) — the result of the linear mapping function applied to input x
y — the actual value label expected for that input x

Similarly the derivatives for w and b respectively are the following:

∂w = (x - ∂z)/m
∂b = ∂z/m

In this case m is the total number of input elements.

We will now be able to re-adjust the parameters using the above as follows:

w = w - θ · ∂w
b = b - θ · ∂b

In this case θ is the learning rate, which as the name suggests controls the ratio by which the parameters will be modified on each iteration. Intuitively, the smaller, the more iterations it will be required for the algorithm to converge, however if the learning is too big, it will overshoot, leading to never being able to converge (from the image above you can imagine it will keep bouncing from side to side never reaching the bottom).

In order for us to calculate loss, we will be using the log loss function, known also as cross-entropy loss function. This function is defined as follows:

Log loss (cross entropy loss) function

Intuitive diagram to visualize cost function from ML Mastery

The loss function will provide us an idea of the improvement of our algorithm across iterations.

Note: In this post we won’t delve into much detail, nor best practices on machine learning, however at the end of the article we will be listing a broad range of sources for people interested to take their machine learning (or GPU compute) knowledge to the next level.

Machine Learning GPU Code Implementation

Now that we have covered some of the core concepts, we will be able to learn about the implementation of the shader, which is the code that will be executed in the GPU.

First we need to define all the input and output buffers as follows:

layout(set = 0, binding = 0) buffer bxi { float xi[]; };
layout(set = 0, binding = 1) buffer bxj { float xj[]; };
layout(set = 0, binding = 2) buffer by { float y[]; };
layout(set = 0, binding = 3) buffer bwin { float win[]; };
layout(set = 0, binding = 4) buffer bwouti { float wouti[]; };
layout(set = 0, binding = 5) buffer bwoutj { float woutj[]; };
layout(set = 0, binding = 6) buffer bbin { float bin[]; };
layout(set = 0, binding = 7) buffer bbout { float bout[]; };
layout(set = 0, binding = 8) buffer blout { float lout[]; };

In more detail:

The input Xas arrays xi and xj will hold the micro-batch of inputs
The array ywill hold all the expected labels for micro-batch inputs
The two input weight parameters wini and woutj will be used for calculating predictions
The input parameter b which will be used for calculating the predictions
The output weights wouti and woutjcontains weights and will store the derivative of W for all micro-batches that should be subtracted
Similarly the output bias array contains the derivatives of bfor all micro-batches that should be subtracted in batch
Finally lout contains the output array where losses will be returned

layout (constant_id = 0) const uint M = 0;

float m = float(M);

Now that we have all the input and output parameters defined, we can start the main function, which will contain the implementation of our machine learning training algorithm.

We will first start by keeping track of the current index of the global invocation. Since the GPU executes in parallel, each of these runs will be running directly in parallel, so this allows the current execution to consistently keep track of what iteration index is currently being executed.

// ...code from previous blocks

int main() {

  uint idx = gl_GlobalInvocationID.x;

  // ...code from latter blocks
}

We now can start preparing all the variables that we’ll be using throughout the algorithms. All our inputs are buffer arrays, so we’ll want to store them in vec2 and float variables.

// ...code from previous blocks

int main() {

  // ...code from previous blocks

    vec2 wCurr = vec2(win[0], win[1]);
    float bCurr = bin[0];

    vec2 xCurr = vec2(xi[idx], xj[idx]);
    float yCurr = y[idx];

  // ...code from latter blocks
}

Now we get into the more fun part —implementing the inference function. Below we will implement the inference function to calculate ŷ, which involves both the linear mapping function, as well as the sigmoid function.

// ...code from previous blocks

float sigmoid(float z) {
    return 1.0 / (1.0 + exp(-z));
}

float inference(vec2 x, vec2 w, float b) {
    // Compute the linear mapping function
    float z = dot(w, x) + b;
    // Calculate the y-hat with sigmoid
    float yHat = sigmoid(z);
    return yHat;
}


int main() {

  // ...code from previous blocks

    float yHat = inference(xCurr, wCurr, bCurr);

  // ...code from latter blocks
}

Now that we have yHat, we can now use it to calculate the derivatives (∂z, ∂w and ∂b), which in this case are the derivative of the currently-executed index input element.

// ...code from previous blocks

int main() {

  // ...code from previous blocks

    float dZ = yHat - yCurr;
    vec2 dW = (1. / m) * xCurr * dZ;
    float dB = (1. / m) * dZ;

  // ...code from latter blocks
}

We can now pass the derivatives as outputs, so the parameters can be re-adjusted for the next iteration.

// ...code from previous blocks


int main() {

  // ...code from previous blocks

    wouti[idx] = dW.x;
    woutj[idx] = dW.y;
    bout[idx] = dB;

  // ...code from latter blocks
}

Finally we’re able to calculate the loss and add it to the output lout array.

// ...code from previous blocks

float calculateLoss(float yHat, float y) {
    return -(y * log(yHat)  +  (1.0 - y) * log(1.0 - yHat));
}

int main() {

  // ...code from previous blocks

    lout[idx] = calculateLoss(yHat, yCurr);
}

That’s it, we’ve now finished the shader that will enable us to train a Logistic Regression algorithm in the GPU — you can find the full code for the shader in the GPU logistic regression example repository.

Now we’ll cover the Kompute code required to run this code against a dataset to train our first model and find the parameters.

Machine Learning Orchestration from Kompute

In order to run the shader we created above in the GPU using Kompute, we will follow the following steps:

Import Kompute and create our main function
Create all the Kompute Tensors required
Create the Kompute Manager and initialize a Kompute Sequence
Execute the Kompute Tensor GPU initialization via Kompute Sequence
Record batch algorithm execution in Kompute Sequence
Iterate 100 times: Run micro-batch execution & update weights
Print resulting parameters to use for further inference

1. Import Kompute and create our main function

We will be importing the single header of Kompute — it’s also possible to use the more granular class-based headers if required. We will also create some of the base configuration variables; namely ITERATIONS and learningRate which will be used in latter code blocks.

#include "kompute/Kompute.cpp"

int main() {
  uint32_t ITERATIONS = 100;
  float learningRate = 0.1;

    // code from latter blocks
}

2. Create all the Kompute Tensors required

// ...code from previous blocks

int main() {
    // ...code from previous blocks

    std::shared_ptr<kp::Tensor> xI{ new kp::Tensor({ 0, 1, 1, 1, 1 })};
    std::shared_ptr<kp::Tensor> xJ{ new kp::Tensor({ 0, 0, 0, 1, 1 })};

    std::shared_ptr<kp::Tensor> y{ new kp::Tensor({ 0, 0, 0, 1, 1 })};

    std::shared_ptr<kp::Tensor> wIn{ new kp::Tensor({ 0.001, 0.001 })};
    std::shared_ptr<kp::Tensor> wOutI{ new kp::Tensor({ 0, 0, 0, 0, 0 })};
    std::shared_ptr<kp::Tensor> wOutJ{ new kp::Tensor({ 0, 0, 0, 0, 0 })};

    std::shared_ptr<kp::Tensor> bIn{ new kp::Tensor({ 0 })};
    std::shared_ptr<kp::Tensor> bOut{ new kp::Tensor({ 0, 0, 0, 0, 0 })};

    std::shared_ptr<kp::Tensor> lOut{ new kp::Tensor({ 0, 0, 0, 0, 0 })};

    // ...code from latter blocks
}

We also store them in a parameter vector for easier access:

// ...code from previous blocks

int main() {
    // ...code from previous blocks

    std::vector<std::shared_ptr<kp::Tensor>> params =
        {xI, xJ, y, wIn, wOutI, wOutJ, bIn, bOut, lOut};

    // ...code from latter blocks
}

3. Create the Kompute Manager and initialize a Kompute Sequence

If you remember from the previous example, we were able to execute commands directly using the Kompute Manager. However we are able to use the Kompute Sequence resource if we want further granularity to record command batches that can be submitted and loaded into the GPU before processing. For this, we will create a Kompute Manager, then create a Kompute Sequence through it.

// ...code from previous blocks

int main() {
    // ...code from previous blocks

    // We first create a manager selecting device 0
    kp::Manager mgr;

    // We then create a sequence and request the pointer (as it's provided as a weak_ptr)
    std::shared_ptr<kp::Sequence> sq = mgr.getOrCreateManagedSequence("createTensors").lock()

    // ...code from latter blocks
}

4. Execute the Kompute Tensor GPU initialization via Kompute Sequence

We can now start by running instructions on GPU resources — namely we will start by initialising and map the Tensors with their respective GPU memory. Here you will see how Kompute Sequences provide you with further granularity on the command execution, but it won’t be until the ML inference section that you will see the flexibility of Kompute Sequence.

Let’s get started by recording commands, namely the OpTensorCreate command, and then evaluating the operation across all the tensors above. This operation will create the respective Vulkan memory/buffer resources.

// ...code from previous blocks

int main() {
    // ...code from previous blocks

    // Being listening for recording sequence
    sq->begin();

    // Register a recording command operation
    sq->record<kp::OpTensorCreate>(params);

    // Stop listening for recording sequence
    sq->end();

    // Evaluate the currently recorded sequence
    sq->eval();

    // ...code from latter blocks
}

5. Record batch algorithm execution in Kompute Sequence

You will also notice that we will be recording three types of Kompute Operations, namely:

// ...code from previous blocks

int main() {
    // ...code from previous blocks

    // Clear previous commands and begin listening for recording sequence
    sq->begin();

    // Ensure the input vectors are sync in GPU device memory
    sq->record<kp::OpTensorSyncDevice>({wIn, bIn});

    // Record the execution of the logistic regression shader
    sq->record<kp::OpAlgoBase<>>(
            params,
            "test/shaders/glsl/test_logistic_regression.comp");

    // Sync the output vectors to device so they can be used for processing
    sq->record<kp::OpTensorSyncLocal>({wOutI, wOutJ, bOut, lOut});

    // Stop listening for recording sequence
    sq->end();

    // ...code from latter blocks
}

kp::OpTensorSyncDevice — This operation ensures that the Tensors are synchronized with their GPU memory by mapping their local data into the GPU data. In this case, these Tensors use Device-only memory for processing efficiency, so the mapping is performed with a staging Tensor inside the operation (which is re-used throughout the operations for efficiency). Here we’re only wanting to sync the input weights, as these will be updated locally with the respective derivatives.
kp::OpAlgoBase — This is the Kompute Operation that binds the shader that we wrote above with all the local CPU/host resources. This includes making available the Tensors. It’s worth mentioning that the index of the tensors provided as parameters is the order in which they are mapped in the shaders via their respective bindings (as you can see in the shaders each vector has the format layout(binding = NUMBER).
kp::OpTensorSyncLocal— This Kompute Operation performs a similar set of instructions as the sync operation above, but instead of copying the data to the GPU memory, it does the converse. This Kompute Operation maps the data in the GPU memory into the local Tensor vector so it’s accessible from the GPU/host. As you can see we’re only running this operation in the output tensors.

6. Iterate 100 times: Run micro-batch execution & update weights

// ...code from previous blocks

int main() {
    // ...code from previous blocks

    // Iterate across all expected iterations
    for (size_t i = 0; i < ITERATIONS; i++) {

        // Run evaluation of recorded commands
        sq->eval();

        // Update all model parameters
        for(size_t j = 0; j < bOut->size(); j++) {
            wIn->data()[0] -= learningRate * wOutI->data()[j];
            wIn->data()[1] -= learningRate * wOutJ->data()[j];
            bIn->data()[0] -= learningRate * bOut->data()[j];
        }
    }
    // ...code from latter blocks
}

7. Print resulting parameters to use for further inference

// ...code from previous blocks

int main() {
    // ...code from previous blocks

    std::cout << "RESULTS" << std::endl;
    std::cout << "w1: " << wIn->data()[0] << std::endl;
    std::cout << "w2: " << wIn->data()[1] << std::endl;
    std::cout << "b: " << bIn->data()[0] << std::endl;
}

And we’re done!

You are able to find this entire example in the example repository, which you’ll be able to run and extend. You will find all the complete files in the GPU Logistic Regression example repo, including the Kompute C++ Code, and the Shader File.

What next?

Kompute Documentation for more details and further examples
The Machine Learning Engineer Newsletter if you want to keep updated on articles around Machine Learning
Awesome Production Machine Learning list for open source tools to deploy, monitor, version and scale your machine learning
Introduction to ML for Coders course by FastAI to learn further machine learning concepts
Vulkan SDK Tutorial for a deep dive into the underlying Vulkan components

Supercharging Game Development with GPU Accelerated Machine Learning

Sat, 12 Sep 2020 00:00:00 GMT

Image by Author

Recently, the world has seen various defining milestones in both, the gaming industry and the AI sector. In only a couple of weeks, we have seen major financial announcements in the gaming industry including Unity’s $1.3 Billion IPO and Epic Games’ $1.78B Investment. The AI sector has also been catching up with its hype, reaching a $60+ Billion Market in 2020, and bringing mind blowing applications in the intersection of AI and Gaming, including AlphaGo’s famous victory over Champion Lee Sedol, as well as deep learning powered games such as AI Dungeon (and many more applications).

This article provides a technical deep dive into this intersection between the two fields, applied artificial intelligence and game development. We delve specifically into how you can leverage the power of the cross-vendor / mobile GPU frameworks for accelerated processing of machine learning and advanced data processing use-cases.

Simple game interface for this tutorial (Image by Author)

In this tutorial you’ll learn how to use the Kompute framework to build GPU optimized code inside the popular open source Godot Game Engine.

You will understand how machine learning and advanced GPU compute can be leveraged in game development through the Godot Game engine.

No background knowledge beyond programming experience is required, but if you are curious about the underlying AI / GPU compute concepts referenced, we suggest checking out our previous article, “Machine Learning in Mobile & Cross-Vendor GPUs Made Simple With Kompute & Vulkan”.

You can find the full code in the example folder in the repository, together with the GDNative Library implementation and the Godot Custom Module implementation.

Image from Godot Open Source GitHub Repository

With over 30k github stars and more than 1k contributors Godot is the most popular OSS game engine. Godot caters for 2D and 3D development, and has been used for a broad range of mobile, desktop, console and web compatible games / applications. Godot is built in C++ making it fast and light — it’s only a 40MB download.

Godot is very intuitive for newcomers through robust design principles, and support for high-level languages, including its domain-specific-language GdScript which has a Python-like syntax, making it very easy to adopt. It is also possible to develop using C#, Python, Visual Scripting, C++, etc.

In this tutorial we will be building a Game using the editor, using GdScript to trigger the ML training / inference, and C++ to develop the core processing components under-the-hood.

Enter Kompute & the Vulkan SDK

Playing “where’s waldo” with Khronos Membership (Image by Vincent Hindriksen via StreamHPC)

Large number of high profile (and new) frameworks have been adopting Vulkan as their core GPU processing SDK. The Godot Engine itself is working on a major 4.0 update that will bring Vulkan as its core rendering engine.

Kompute Documentation (Image by Author)

Kompute was not built to hide any of the core Vulkan concepts — the core Vulkan API is very well designed. Instead it augments Vulkan’s computing capabilities with a BYOV (bring your own Vulkan) design, enabling developers by reducing boilerplate code required and automating some of the more common workflows involved in writing Vulkan applications.

Artificial Intelligence in Game Development

To start with, we will need an interface that allows us to expose our Machine Learning logic, which will require primarily two functions:

train(…)— function which will allow the machine learning model to learn to predict outputs from the inputs provided
predict(...)— function that will predict the output of an unknown instance. This can be visualised in the two workflows outlined in the image below.

Data Science Process (Image by Author)

Particularly in game development, this would also be a common pattern for machine learning workflows, for both predictive and explanatory modelling use cases. This often consists of leveraging data generated by your users as they interact directly (or indirectly) with the game itself. This data can then serve as training features for machine learning models. Training of new models can be performed through manual “offline” workflows that data scientists would carry out, or alternatively through automated triggers retraining models.

Machine Learning Project in Godot

We will start by providing a high level overview of how our Kompute bindings will be used in our game in Godot. We will create a simple project and train the ML model we created, which will run in the GPU. In this section we will already have the custom KomputeModelML Godot class accessible — details on how to build it and import it into the godot project are covered in a latter section.

New Project Screen (Image by Author)

For this we are using Godot version 3.2.3 stable. Create a new project and a new 2D scene. You should see a blank project with a top level 2D node, as per the image on the left.

Now we’ll be able to start adding resources to our game. We will start by creating a simple UI interface / menu consisting of inputs for the machine learning model that reflects the workflows covered in the previous architecture diagram.

We will have two input LineEdit text boxes for the data X_i and X_j, and one input LineEdit text box for the Y expected predictions. The image below shows the node structure used to build the UI. You can also access the full godot project in the repo by importing the project.godot file. The nodes will be referenced to read the input data and display the output predictions (and learned parameters).

We will now be able to add the GdScript code in Godot which will allow us to read the inputs, train the model and perform predictions — this script is created under the Parentnode.

Below is the full script that we are using to perform the processing, which is using the KomputeModelMLcustom Godot class we build in the next section. We’ll break down each of the different areas of the code below.

extends Node2D

onready var xi_node = $UI/UIVBoxContainer/XIHBoxContainer/LineEdit
onready var xj_node = $UI/UIVBoxContainer/XJHBoxContainer/LineEdit
onready var y_node = $UI/UIVBoxContainer/XJHBoxContainer/LineEdit
onready var preds_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer2/PredictionsLabel
onready var w1_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer/Weight1Label
onready var w2_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer/Weight2Label
onready var bias_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer/BiasLabel

func compute_ml():

    var xi = str2var(xi_node.text)
    var xj = str2var(xj_node.text)
    var y = str2var(y_node.text)

    var s = KomputeModelML.new()

    s.train(y, xi, xj)

    var preds = s.predict(xi, xj)

    preds_node.text = str(preds)

    var params = s.get_params()

    w1_node.set_text(str(params[0]))
    w2_node.set_text(str(params[1]))
    bias_node.set_text(str(params[2]))

First we define variables to make the referencing of Godot Editor nodes simpler. This can be done with the dollar sign syntax $NODE/PATH as implemented in the snipped below.

extends Node2D

onready var xi_node = $UI/UIVBoxContainer/XIHBoxContainer/LineEdit
onready var xj_node = $UI/UIVBoxContainer/XJHBoxContainer/LineEdit
onready var y_node = $UI/UIVBoxContainer/XJHBoxContainer/LineEdit
onready var preds_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer2/PredictionsLabel
onready var w1_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer/Weight1Label
onready var w2_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer/Weight2Label
onready var bias_node = $UI/UIVBoxContainer/Panel/VBoxContainer/PredHBoxContainer/BiasLabel

The compute_ml() function below contains the logic for the machine learning training and predictions. We start by reading the inputs from the text boxes using the Godot editor node references.

# ... code from previous blocks

func compute_ml():

    var xi = str2var(xi_node.text)
    var xj = str2var(xj_node.text)
    var y = str2var(y_node.text)

    # ... code from latter blocks

We now can create an instance from the C++ class bindings that we will be constructing in the next section. This is the class that exposes the training and prediction functions that we will be using for machine learning inference.


func compute_ml():

    # ... code from previous code blocks

    # We create an instance
    var kml = KomputeModelML.new()

    # ... code from latter code blocks

We now can train our model by passing the inputs and the expected predictions. The ML model underneath is a logistic regression model, which will adjust its internal parameters to fit the inputs and outputs best, resulting in a model that is then able to predict unseen datapoints.


func compute_ml():

    # ... code from previous code blocks

    # We create an instance
    kml.train(y, xi, xj)

    # ... code from latter code blocks

Now that we have trained our model, we can perform predictions on unseen data-points. For simplicity, we will pass in the same inputs we used for testing, however you could pass completely new arrays and see what predictions it produces. We can then display the results in the preds_node reference node variable that we defined, which will show in the display.


# ... code from previous code blocks

func compute_ml():

    # ... code from previous code blocks

    var preds = kml.predict(xi, xj)

    preds_node.text = str(preds)

    # ... code from latter blocks

Finally we want to also display the learned parameters, which in this case it includes w1 , w2 and the bias . We are able to display the weights and bias in the respective labels.

# ... code from previous blocks

func compute_ml():

    # ... code from previous blocks

    var params = s.get_params()

    w1_node.set_text(str(params[0]))
    w2_node.set_text(str(params[1]))
    bias_node.set_text(str(params[2]))

    # ... code from latter blocks

The final thing to set up is to connect the “Kompute Train & Predict” button into the compute_ml function, which can be done by setting up a signal that points to the function itself via the editor.

Setting up the Signal to the compute_ml method in our script (Image by Author)

Once this is all set up, we can run the Game and trigger the training and prediction using the provided inputs. We can then see the learned parameters as well as the prediction outputs. It is also possible to see how the learned parameters change as we modify the inputs y, xi and xj.

Kompute ML Godot interface with resulting parameters (Image by Author)

Kompute ML Implementation

Now that there is an intuition on what is happening in the game level, we are able to code our underlying C++ class to create the Godot bindings using the Kompute framework.

Kompute Architecture Design (Image by Author)

We will be following the design principles of Kompute which are outlined in this accompanying diagram showing the different components. We will be following this workflow to load the data in the GPU and perform the training.

The header file for the Godot class binding implementation is outlined below, which we will break down in detail. As you can see creating a C++ class with bindings is quite intuitive, and you can see the same functions that we used to call from our Godot GdScript above.

KomputeLogisticRegression.hpp Implementation

#pragma once

#include <Godot.hpp>
#include <Node2D.hpp>
#include <Array.hpp>

#include "kompute/Kompute.hpp"

namespace godot {
class KomputeLogisticRegression : public Node2D {
private:
    GODOT_CLASS(KomputeLogisticRegression, Node2D);

    // For custom module use gdclass
    // GDCLASS(KomputeSummatorNode, Node);


public:
    KomputeLogisticRegression();

    void train(Array y, Array xI, Array xJ);

    Array predict(Array xI, Array xJ);

    Array get_params();


    static void _register_methods();

    // For custom module use bind methods instead
    //static void _bind_methods();


private:
    kp::Tensor mWeights;
    kp::Tensor mBias;
};

static std::string LR_SHADER = R"(
//... rest of header file

The initial section of the class header file includes:


#include "kompute/Kompute.hpp"

#include <Godot.hpp>
#include <Node2D.hpp>
#include <Array.hpp>

namespace godot {
class KomputeLogisticRegression : public Node2D {
private:
    GODOT_CLASS(KomputeLogisticRegression, Node2D);

Import of the Kompute.hppheader containing all the Kompute dependencies that we’ll use in this project
The top-level Godot.hpp import is required to ensure all Godot components are available.
Node2D is the resource that we will be inheriting from, but you can inherit from other classes in the inheritance tree depending on how you plan to use your custom Godot class
For data passing across the application, we will be using a Godot Array which will handle transfer between GdScript and the Naive C++, as well as a Kompute Tensor which will handle the GPU data management.
The GODOT_CLASS macro definition extends the class by adding extra Godot related functionality. As we will see below, you will need to use GDCLASS instead when building as a Godot custom module.

Following the base functionality, we have to define the core logic:


// ... previous code blocks

public:
    KomputeLogisticRegression();

    void train(Array y, Array xI, Array xJ);

    Array predict(Array xI, Array xJ);

    Array get_params();


// ... latter code blocks

void train(Array y, Array xI, Array xJ) —Trains the machine learning model using the GPU native code for the logistic regression model. It takes the input array(s) X, and the array ycontaining the expected outputs.
Array predict(Array xI, Array xJ) —Perform the inference request. In this implementation it is not using GPU code as generally there tends to be less performance gains through parallelization on the inference side. However there are still expected performance gains if multiple inputs are processed in parallel (which this function allows for).
Array get_params() —Returns an array containing the learned parameters in the format of [ <weight_1>, <weight_2>, <bias> ].

We then are able to declare the methods that will be bound between the C++ and the high level GdScript Godot Engine, and accessible to the editor and broader Game in general. We’ll look briefly below at the code required to register a function.


// ... previous code blocks

    static void _register_methods();

// ... latter code blocks

For data management we will be using Kompute tensors as well as Arrays — in this case we will only need to “learn” and “persist” the weights and biases of our logistic regression model.

// ... code from previous blocks

private:
    kp::Tensor mWeights;
    kp::Tensor mBias;
};

// ... code from latter blocks

Finally, we also define the shader code, which is basically the code that will be executed as machine code inside of the GPU. Kompute allows us to pass a string containing the code, however for production deployments it is possible to convert the shaders to binary, and also use the utilities available to convert into header files.

// ... code from previous blocks


static std::string LR_SHADER = R"(
    // ... shader code excluded for simplicity
)";


// ... code from latter blocks

If you are interested in the full implementation you can find all the files in the gdnative implementation and custom module implementation folders. Furthermore if you are interested in the theoretical and underlying foundational concepts of these techniques, this is covered fully in our previous post.

Compiling and Integrating into Godot

Now that we have the base code for our GPU optimized machine learning model, we are now able to proceed to running it in our Godot game engine. There are two main ways in which Godot allows us to add C++ code into our projects:

GdNative Library Build Instructions— In Godot you can add your own “GdNative scripts” which are basically C++ classes with bindings to GdScript, which means these can be used and referenced dynamically in the project. This approach works with the standard Godot installation and doesn’t require re-compiling the full editor like the 2nd option.
Custom ModuleBuild Instructions— In Godot only the underlying core classes are part of the “core” components; everything else — the UI, the editor, the GdScript language, the networking capabilities — are custom models. Writing a custom module is easy, and this is how we are able to expose some of the core Kompute functionality. This option requires the full Godot C++ project to be recompiled with the custom module.

Once you set one of these methods through the instructions, you will be able to access the custom object from within Godot. There are some nuanced implementation differences between each approach.

We won’t be covering specific build details in the blog post, but you will be able to find the exact code and build instructions in the GdNative Library / Custom Module links above.

What’s next?

Congratulations, you’ve made it all the way to the end! Although there was a broad range of topics covered in this post, there is a massive amount of concepts that were skimmed through. These include the underlying Vulkan concepts, GPU computing fundamentals, machine learning best practices, and more advanced Kompute concepts. Luckily, there are resources online to expand your knowledge on each of these. Here are some links I recommend for further reading:

“Machine Learning in Mobile & Cross-Vendor GPUs Made Simple With Kompute & Vulkan” article with a deeper dive in theory and concepts
Kompute Documentation for more details and further examples
The Machine Learning Engineer Newsletter if you want to keep updated on articles around Machine Learning
Awesome Production Machine Learning list for open source tools to deploy, monitor, version and scale your machine learning
Introduction to ML for Coders course by FastAI to learn further machine learning concepts
Vulkan SDK Tutorial for a deep dive into the underlying Vulkan components