Where the word came from

Monads came out of algebraic topology, not programming. In the late 1950s, Roger Godement described a construction he called the standard construction[^1]: a functor $T$ together with two natural transformations that let you wrap things up and flatten double wrappings. Topologists used it to build resolutions for computing cohomology. It was a tool, not yet a subject.

Through the 1960s the structure got several names. Peter Huber and others called it a triple, because it consists of three things: a functor and two transformations. In 1965 two papers appeared almost back to back, one by Heinrich Kleisli[^2] and one by Samuel Eilenberg and John Moore[^3], and they answered the same question in two different ways: does every monad come from a pair of adjoint functors? (Yes, and we’ll see both answers below.)

The word monad itself is usually credited to Jean Bénabou, and the story that gets told is that it came up around the category theory meetings of the mid-60s, where people were tired of “triple” and “standard construction” and wanted something better. It’s a blend of monoid and triad, and there’s a nod to Leibniz’s Monadology in it too. Saunders Mac Lane picked the name for Categories for the Working Mathematician[^4] in 1971, and that pretty much settled it, though you’ll still find “triple” in older papers.

:::margin Leibniz’s monads were indivisible units that make up the world. A category theorist’s monad is about as far from that as it gets, but the name stuck. :::

Then nothing much happened in programming for about twenty years. In 1989 Eugenio Moggi[^5] pointed out that many “impure” features of programs (exceptions, state, nondeterminism, input/output) could be described uniformly as monads, and that this gave a clean semantics for them. Philip Wadler[^6] took Moggi’s idea and showed functional programmers how to write code with it. By 1996, Haskell 1.3 used monads for I/O, and the idea escaped into the wider world.

What is a monad, informally

A lot of the time, functions don’t just return a value. They return a value plus something: a value that might be missing, a value with a log attached, a value that depends on some state, a list of possible values. Call that “something” a context, and write $T A$ for “a value of type $A$ in context $T$”.

A monad is what you need to chain functions like this together without the context getting in your way. That needs exactly two things:

1. A way to put a plain value into the context, without doing anything interesting. Call it $\eta$ (in code: return or pure).
2. A way to take a context-wrapped value and feed it into a function that produces another context-wrapped value. That’s bind, and we’ll spend a whole section on it.

Before going further, let’s look at two concrete examples, the abstract version will make much more sense once you’ve seen it in practice.

The Maybe monad

The simplest context is “this might not be there”. Mathematically, given a type $A$, we add one extra point to it, call it nothing:

$$ \text{Maybe},A = A + 1 $$

Here $+$ is disjoint union and $1$ is a one-element set. A value of $\text{Maybe},A$ is either $\text{just}(a)$ for some $a \in A$, or $\text{nothing}$.

Putting a value in the context is obvious: $\eta(a) = \text{just}(a)$.

Chaining is where it earns its keep. Suppose I have

$$ f : A \to \text{Maybe},B, \qquad g : B \to \text{Maybe},C $$

We can’t just compose them as $g \circ f$, because $f$ hands back a $\text{Maybe},B$ and $g$ wants a plain $B$. What we want is: if $f$ gave us $\text{just}(b)$, run $g$ on $b$; if $f$ gave us $\text{nothing}$, stop and return $\text{nothing}$. That rule is the entire Maybe monad. It’s the “early return on null” pattern you’ve written a thousand times, written once.

The State monad, and a detour through currying

The second example needs a small detour, because it relies on currying, and currying turns out to matter a lot later.

Currying is the observation that a function of two arguments is the same thing as a function that takes the first argument and returns a function of the second:

$$ \text{Hom}(A \times S,\ C) ;\cong; \text{Hom}(A,\ C^S) $$

Here $C^S$ means “functions from $S$ to $C$”. Going from left to right is called currying, going back is uncurrying. In category theory language, this says the functor $- \times S$ is left adjoint to the functor $(-)^S$. Keep that word, adjoint, in the back of your mind; it’s coming back.

Now, a stateful computation is one that reads some state $s \in S$, produces a result, and produces a new state. As a pure function:

$$ \text{State}_S,A = S \to (A \times S) = (A \times S)^S $$

Look at the shape of that: it’s $(- \times S)$ followed by $(-)^S$. It is literally the two halves of the currying adjunction glued together. That is not a coincidence, and it’s my favourite fact in this whole topic: every adjunction gives you a monad, and the State monad is the one that currying gives you.

For State, $\eta$ returns a value without touching the state:

$$ \eta(a) = \lambda s.\ (a, s) $$

and chaining runs the first computation, takes the result and the new state, and hands both to the next one. You get the feeling of mutable state while everything stays a pure function.

The mathematics

Definition

Let $\mathcal{C}$ be a category; for programmers, think of it as “types and functions”. A monad on $\mathcal{C}$ is a triple $(T, \eta, \mu)$ where

• $T : \mathcal{C} \to \mathcal{C}$ is a functor (it sends a type $A$ to $T A$, and a function $f : A \to B$ to $Tf : TA \to TB$; this is map),
• $\eta : \text{Id} \Rightarrow T$ is a natural transformation called the unit, with components $\eta_A : A \to TA$,
• $\mu : T^2 \Rightarrow T$ is a natural transformation called the multiplication, with components $\mu_A : TTA \to TA$ (this is flatten or join).

They have to satisfy three laws. Associativity:

$$ \mu_A \circ T\mu_A ;=; \mu_A \circ \mu_{TA} \quad : TTTA \to TA $$

and the two unit laws:

$$ \mu_A \circ T\eta_A ;=; \text{id}{TA} ;=; \mu_A \circ \eta{TA} $$

Read them in words. Associativity says that if you have three layers of context, it doesn’t matter whether you flatten the inner two first or the outer two first. The unit laws say that wrapping something in a trivial layer and then flattening gives you back what you started with, whichever side you add the layer on.

If those look like the laws of a monoid ($(xy)z = x(yz)$ and $1x = x = x1$), that’s because they are. A monad is a monoid where the elements are functors, “multiplication” is composition of functors, and the unit is the identity functor. Hence the infamous line: a monad is just a monoid in the category of endofunctors. It sounds like a joke, but it’s actually a correct and fairly compact definition.

:::margin This line is from Mac Lane, more or less. It became a meme thanks to James Iry’s A Brief, Incomplete, and Mostly Wrong History of Programming Languages. :::

Checking against Maybe: $\mu$ takes $\text{Maybe}(\text{Maybe},A) = (A + 1) + 1$ and collapses it to $A + 1$ by sending $\text{just}(\text{just}(a))$ to $\text{just}(a)$ and both kinds of $\text{nothing}$ to $\text{nothing}$. You can check the laws by hand in a minute.

For lists: $T A$ is finite lists of $A$, $\eta(a) = [a]$, and $\mu$ concatenates a list of lists. Associativity of $\mu$ is just associativity of concatenation.

The Kleisli category

Here’s the first big idea. In programming, we rarely work with $\mu$ directly. We work with functions of the form

$$ f : A \to TB $$

“functions that return a value in context”. Kleisli’s insight was that these functions form a category of their own, where the context is handled automatically.

The Kleisli category $\mathcal{C}_T$ has:

• the same objects as $\mathcal{C}$,
• a morphism $A \to B$ in $\mathcal{C}_T$ is a morphism $A \to TB$ in $\mathcal{C}$,
• the identity on $A$ is $\eta_A : A \to TA$,
• composition of $f : A \to TB$ and $g : B \to TC$ is

$$ g \circ_T f ;=; \mu_C \circ Tg \circ f \quad : A \to TC $$

Read the composition right to left: run $f$ to get a $TB$; map $g$ over it to get a $TTC$; flatten with $\mu$ to get a $TC$. That’s exactly the “chain these context-returning functions” operation from the Maybe example, now stated once for every monad. People write it as $g \mathrel{<!=!<} f$, or $f \mathrel{>!=!>} g$ in the other order, and call it the fish operator.

The monad laws are exactly what make this a category. Associativity of $\mu$ gives associativity of $\circ_T$, and the unit laws make $\eta$ an identity:

$$ \eta_B \circ_T f = f, \qquad f \circ_T \eta_A = f, \qquad h \circ_T (g \circ_T f) = (h \circ_T g) \circ_T f $$

This is, I think, the cleanest way to say what monads give programmers: a monad turns effectful functions into something you can compose like ordinary functions, with the same laws. That’s all. The rest is notation.

Kleisli also showed how this answers the 1965 question. There’s a functor $F_T : \mathcal{C} \to \mathcal{C}_T$ that sends a plain function $f$ to $\eta \circ f$, and a functor $U_T : \mathcal{C}_T \to \mathcal{C}$ that sends $A$ to $TA$. These are adjoint, $F_T \dashv U_T$, and the monad you get back from them is the original $T$.

Monad algebras

The second construction asks a different question. Instead of “what are the functions in context?”, it asks: what does it mean to be able to evaluate a context?

An algebra for the monad $T$ (often called a $T$-algebra) is an object $A$ together with a morphism

$$ a : TA \to A $$

that collapses the context, in a way that respects the monad:

$$ a \circ \eta_A = \text{id}_A, \qquad a \circ Ta = a \circ \mu_A $$

The first law says: if I wrap a value trivially and then evaluate, I get the value back. The second says: if I have two layers, I can evaluate the inner layer then the outer, or flatten first and evaluate once, and I get the same answer.

This is easiest to feel with the list monad. A list-algebra is a set $A$ with a map $a : \text{List},A \to A$ that turns any list of elements into a single element, consistently. Take $A = \mathbb{N}$ and $a = \text{sum}$:

$$ \text{sum}([n]) = n, \qquad \text{sum}([\text{sum}(xs_1), \dots, \text{sum}(xs_k)]) = \text{sum}(xs_1 \mathbin{+!!+} \cdots \mathbin{+!!+} xs_k) $$

Both laws hold. In fact, list-algebras are exactly monoids: give me a monoid and folding with its operation is an algebra; give me an algebra and $x \cdot y = a([x, y])$ with unit $a([,])$ is a monoid. Similarly, Maybe-algebras are pointed sets: the algebra has to say where $\text{nothing}$ goes, and that’s just picking a default element. If you’ve ever written “value or default”, you’ve written a Maybe-algebra.

A morphism of algebras $h : (A, a) \to (B, b)$ is a function $h : A \to B$ that commutes with evaluation:

$$ h \circ a = b \circ Th $$

For lists, that’s a monoid homomorphism.

The Eilenberg–Moore category

The algebras and their morphisms form a category, written $\mathcal{C}^T$ and called the Eilenberg–Moore category. Like Kleisli’s, it comes with an adjunction:

$$ F^T : \mathcal{C} \to \mathcal{C}^T, \quad A \mapsto (TA, \mu_A) $$

$$ U^T : \mathcal{C}^T \to \mathcal{C}, \quad (A, a) \mapsto A $$

$U^T$ forgets the algebra structure. $F^T$ sends $A$ to the free algebra on $A$, which is $TA$ evaluated by flattening. For lists, the free algebra on $A$ is the free monoid on $A$, which is just lists of $A$ under concatenation. Again $F^T \dashv U^T$, and again the monad you get back is $T$.

So we have two adjunctions that both produce $T$, and here is the beautiful part. Among all adjunctions that give rise to $T$, Kleisli’s is the smallest (it’s the initial one) and Eilenberg–Moore’s is the largest (it’s the terminal one). Every other adjunction producing $T$ sits between them. And the Kleisli category sits inside the Eilenberg–Moore one: it’s equivalent to the full subcategory of free algebras.

:::margin An adjunction whose comparison functor to the Eilenberg–Moore category is an equivalence is called monadic. Groups over sets are monadic; this is why “a group is an algebra for a monad” is a meaningful statement. :::

The way I think about it:

• the Kleisli category is about computations: arrows $A \to TB$, things that produce effects,
• the Eilenberg–Moore category is about handlers: maps $TA \to A$, things that consume effects and give you a plain answer.

That split is exactly how modern effect systems and algebraic effects are designed: a program is a Kleisli arrow, a handler is an algebra. Fifty-year-old mathematics, sitting quietly underneath.

Bind

Now back to code. Programmers rarely write $\mu$ or Kleisli composition directly. They write bind, usually spelled $\gg!=$:

$$ (\gg!=) : TA \to (A \to TB) \to TB $$

$$ m \gg!= f ;=; \mu_B\big((Tf)(m)\big) $$

Notice the currying again: bind takes a wrapped value and gives back a function that wants a continuation $f$. It says: here’s a value in context, and here’s what to do with the plain value once you have it; you figure out the context.

Bind, $\mu$ and Kleisli composition are all interdefinable, so a monad can be given by any of them. From bind you recover the others:

$$ \mu(mm) = mm \gg!= \text{id}, \qquad Tf(m) = m \gg!= (\eta \circ f), \qquad (f \mathrel{>!=!>} g)(a) = f(a) \gg!= g $$

And the monad laws, rewritten in terms of bind, become the three laws you’ll see in every Haskell tutorial:

$$ \eta(a) \gg!= f ;=; f(a) \qquad \text{(left identity)} $$

$$ m \gg!= \eta ;=; m \qquad \text{(right identity)} $$

$$ (m \gg!= f) \gg!= g ;=; m \gg!= (\lambda x.\ f(x) \gg!= g) \qquad \text{(associativity)} $$

They’re the same laws as before, just seen from the point of view of someone writing a pipeline. Left and right identity say $\eta$ doesn’t do anything. Associativity says that how you group the steps of a pipeline doesn’t matter, which is what lets you refactor a long chain into helper functions without changing its meaning.

Bind in TypeScript

TypeScript doesn’t have higher-kinded types, so we can’t write one generic Monad interface the way Haskell does. That’s fine. Writing a couple of concrete monads by hand is more illuminating anyway.

Here’s Maybe:

type Maybe<A> = { kind: 'just'; value: A } | { kind: 'nothing' }

const just = <A>(value: A): Maybe<A> => ({ kind: 'just', value })
const nothing: Maybe<never> = { kind: 'nothing' }

// η
const pure = just

// >>=
const bind = <A, B>(m: Maybe<A>, f: (a: A) => Maybe<B>): Maybe<B> =>
  m.kind === 'nothing' ? m : f(m.value)

That’s the entire monad. Now some functions that can fail:

const parseNumber = (s: string): Maybe<number> => {
  const n = Number(s)
  return Number.isNaN(n) ? nothing : just(n)
}

const reciprocal = (n: number): Maybe<number> =>
  n === 0 ? nothing : just(1 / n)

const safeSqrt = (n: number): Maybe<number> =>
  n < 0 ? nothing : just(Math.sqrt(n))

And chaining them:

const compute = (input: string): Maybe<number> =>
  bind(bind(parseNumber(input), reciprocal), safeSqrt)

compute('4')    // just(0.5)
compute('0')    // nothing   (division by zero)
compute('-4')   // nothing   (negative square root)
compute('hey')  // nothing   (not a number)

No null checks in the pipeline. Each function only deals with its own failure case, and bind handles the plumbing. If you’ve used optional chaining, or the then method on promises, you’ve been using something monad-shaped already. (Promises aren’t a lawful monad, because they flatten nested promises automatically, but they’re close enough that the intuition carries over.)

Here’s State, which is where the currying story pays off:

type State<S, A> = (s: S) => [A, S]

const pureS = <S, A>(a: A): State<S, A> => (s) => [a, s]

const bindS =
  <S, A, B>(m: State<S, A>, f: (a: A) => State<S, B>): State<S, B> =>
  (s) => {
    const [a, s1] = m(s)
    return f(a)(s1)
  }

const get = <S>(): State<S, S> => (s) => [s, s]
const put = <S>(s: S): State<S, void> => () => [undefined, s]

A tiny example: generating fresh variable names, the way a compiler would.

const fresh: State<number, string> = bindS(get<number>(), (n) =>
  bindS(put(n + 1), () => pureS(`x${n}`)),
)

const threeNames = bindS(fresh, (a) =>
  bindS(fresh, (b) =>
    bindS(fresh, (c) => pureS<number, string[]>([a, b, c])),
  ),
)

threeNames(0) // [['x0', 'x1', 'x2'], 3]

There’s no mutable counter anywhere. Every function is pure, and the state is threaded through by bind. The nesting is a bit ugly, though, and that ugliness is exactly what Haskell’s syntax was designed to get rid of.

Monads in Haskell

In Haskell, monad is a type class, and $\eta$ and bind are its methods:

class Applicative m => Monad m where
  (>>=)  :: m a -> (a -> m b) -> m b
  return :: a -> m a
  return = pure

Compare the type of bind to the math: $TA \to (A \to TB) \to TB$. It’s the same thing, character for character. The Applicative superclass is there because every monad is also a functor (that’s $T$ acting on functions) and an applicative (which lets you combine independent computations), and Haskell asks you to provide those too.

Maybe is in the standard library, but here’s what its instance looks like:

instance Monad Maybe where
  Nothing >>= _ = Nothing
  Just a  >>= f = f a

The real treat is do-notation, which is syntactic sugar for chains of bind:

compute :: String -> Maybe Double
compute input = do
  n <- readMaybe input
  r <- reciprocal n
  safeSqrt r

The compiler turns that into

readMaybe input >>= \n -> reciprocal n >>= \r -> safeSqrt r

Each arrow in do-notation is a bind, and the rest of the block is the continuation. It reads like imperative code, but it’s a pure expression, and the monad decides what “next line” means. For Maybe, “next line” means “if we didn’t fail”. For lists, it means “for each possibility”. For State, it means “with the updated state”.

Here’s the fresh-name example again, with a State monad written from scratch:

newtype State s a = State { runState :: s -> (a, s) }

instance Functor (State s) where
  fmap f (State g) = State $ \s -> let (a, s') = g s in (f a, s')

instance Applicative (State s) where
  pure a = State $ \s -> (a, s)
  State f <*> State g = State $ \s ->
    let (h, s1) = f s
        (a, s2) = g s1
    in (h a, s2)

instance Monad (State s) where
  State g >>= f = State $ \s ->
    let (a, s') = g s
    in runState (f a) s'

get :: State s s
get = State $ \s -> (s, s)

put :: s -> State s ()
put s = State $ \_ -> ((), s)

fresh :: State Int String
fresh = do
  n <- get
  put (n + 1)
  return ("x" ++ show n)

threeNames :: State Int [String]
threeNames = do
  a <- fresh
  b <- fresh
  c <- fresh
  return [a, b, c]

-- runState threeNames 0 == (["x0","x1","x2"], 3)

Compare this to the TypeScript version. Same idea, same bind, but the nesting is gone.

And finally, the one that made Haskell take monads seriously: IO. A value of type IO a is a description of a program that, when run, interacts with the world and produces an $a$. Bind lets you sequence these descriptions without ever leaving pure code:

main :: IO ()
main = do
  putStrLn "what's your name?"
  name <- getLine
  putStrLn ("hello, " ++ name)

Conceptually, IO is the State monad where the state is the entire outside world. You can’t copy the world or go back in time with it, and the monad structure makes sure you never try.

Further Reading

1. Eric Lippert’s blog series on monad
2. Monads by Mark Seeman (this is part of his series of blogs on design patterns and category theory)
3. Jeremy Bowers’ blog on Monads

[^1]: Godement, R. (1958). Topologie algébrique et théorie des faisceaux. Hermann.

[^2]: Kleisli, H. (1965). Every standard construction is induced by a pair of adjoint functors. Proceedings of the AMS, 16(3), 544–546.

[^3]: Eilenberg, S., & Moore, J. C. (1965). Adjoint functors and triples. Illinois Journal of Mathematics, 9(3), 381–398.

[^4]: Mac Lane, S. (1971). Categories for the Working Mathematician. Springer.

[^5]: Moggi, E. (1991). Notions of computation and monads. Information and Computation, 93(1), 55–92.

[^6]: Wadler, P. (1992). The essence of functional programming. POPL '92.

Introduction

The instruction name can broken into fairly understandable bits

GF2-P8-AFFINE-Q-B, where

1. GF2 stands for $\mathbb{GF}(2)$, the Galois field with two elements
2. P8 is for power of 8 ($2^8$). It specifies the field sizea as it works with 8-bit bytes, so operation operating within the finite field $\mathbb{GF}(2^8)$
3. Affine indicates affine transformations, which means it multiplies and 8-bit element by an 8 $\times$ 8 bit matrix and then adds (XORs) and 8-bit vector constant.
4. Qword variant, the transformation matrix is supplied as a qword 64-bit; each 64-bit chunk contains the 8 $\times$ 8 matrix representations
5. B, the elements being transforrmed are bytes.

A high level interpretation is: perform an affine transformations over $\mathbb{GF}(2)$ on packed 8-bit elements.

Here each byte is treated as an 8-bit vector over $\mathbb{GF}(2)$. Mathematically operation can be represented as: $$ y = Ax \oplus b $$ ,the matmul is operated over Galois field, whichm means instruction is essentially performing 8 XOR equations in parallel for every byte.

A byte as a vector over 8-bit Galois field

Consider an 8-bit byte

$$ x = x_7x_6x_5x_4x_3x_2x_1x_0 $$

where each bit satisfies

$$ x_i \in \mathbb{GF}(2)={0,1} $$

We can represent the byte as an 8-dimensional vector over $\mathbb{GF}(2)$:

$$ \mathbf{x}

\begin{bmatrix} x_7\ x_6\ x_5\ x_4\ x_3\ x_2\ x_1\ x_0 \end{bmatrix} \in \mathbb{GF}(2)^8 $$

An $8\times8$ matrix whose entries are elements of $GF(2)$ can then operate on this vector:

$$ A \in \mathbb{GF}(2)^{8\times8} $$

The resulting vector is

$$ \mathbf{y}=A\mathbf{x} $$

Because the arithmetic is performed over $\mathbb{GF}(2)$, addition is XOR and multiplication is AND.

For example, if

$$ A= \begin{bmatrix} 1&0&1&0&0&0&0&0\ 0&1&0&1&0&0&0&0\ 0&0&1&0&1&0&0&0\ 0&0&0&1&0&1&0&0\ 0&0&0&0&1&0&1&0\ 0&0&0&0&0&1&0&1\ 1&0&0&0&0&0&1&0\ 0&1&0&0&0&0&0&1 \end{bmatrix} $$

then

$$ \mathbf{y}

\begin{bmatrix} x_7\oplus x_5\ x_6\oplus x_4\ x_5\oplus x_3\ x_4\oplus x_2\ x_3\oplus x_1\ x_2\oplus x_0\ x_7\oplus x_1\ x_6\oplus x_0 \end{bmatrix} $$

Thus, an $8\times8$ binary matrix can be viewed as eight XOR equations, one for each output bit.

For an affine transformation, an additional constant vector $\mathbf{b}\in \mathbb{GF}(2)^8$ is added:

$$ \mathbf{y}=A\mathbf{x}\oplus\mathbf{b} $$

just a draft, will continue later

This post explores how a simple reparameterization enables gradients to flow through categorical sampling, unlocking end-to-end training with discrete latent variables.

In models like DALL-E[^1], the goal is to maximize the joint likelihood $\ln p_{\theta,\psi}(x,y)$ of images and text. Using the standard ELBO derivation, we get a familiar-looking lower bound:

$$ ln p_{\theta,\psi}(x,y) \geq \mathbb{E}{z \sim q{\phi}(z|x)}[\ln p_{\theta}(x|y,z)] - \beta D_{KL}(q_{\phi}(y,z|x) \parallel p_{\psi}(y,z)) $$

Here’s the twist: unlike a VAE’s continuous latent space, (z) here is sampled from a categorical distribution $q_{\phi}(z|x)$. It’s discrete. And that breaks everything. The sampling operation creates a hard stop in the computational graph—no gradients can flow back through it. So while we could compute the ELBO, we couldn’t optimize it end-to-end.

The solution? A clever reparameterization trick, but for categoricals instead of Gaussians. That’s where the Gumbel-Softmax comes in.

The Core Idea: Continuous Relaxation

The key insight comes from Maddison et al.'s Concrete distribution[^2]. Instead of sampling a one-hot vector directly, we sample from a continuous distribution over the simplex $\Delta^{k-1}$ (the set of $k$-dimensional probability vectors). Think of it as sampling near the vertices of a simplex, where each vertex represents a pure category.

The process is surprisingly elegant. Given unnormalized logits $\pi_1, \dots, \pi_k$:

1. Add Gumbel noise: Sample $g_i \sim \text{Gumbel}(0,1)$ and add it to each $\log \pi_i$. This creates perturbed logits that are stochastic but still differentiable.
2. Softmax with temperature: Apply a tempered softmax to get a “soft” one-hot vector:

$$ X_i = \frac{\exp\left(\log \pi_i + g_i\right)/\tau}{\sum_{j=1}^k \exp\left(\log \pi_j + g_j\right)/\tau} $$

The temperature $\tau > 0$ controls how sharply the distribution peaks. As $\tau \to 0$, the samples approach true one-hot vectors; as $\tau$ grows, they become more uniform. During training, you typically anneal $\tau$ from high to low, trading off gradient variance for fidelity.

The resulting Concrete distribution over the simplex has this joint density:

$$ p_{\pi,\tau}(X) = (k-1)! , \tau^{k-1} \prod_{i=1}^k \left(\frac{\pi_i X_i^{-\tau-1}}{\sum_{j=1}^k \pi_j X_j^{-\tau}}\right) $$

The formula is less important than what it does: it gives us a way to sample from a categorical distribution while keeping the gradients intact.

In practice we often use the Gumbel–Softmax (Concrete) relaxation during training to obtain differentiable, soft one-hot samples (with a temperature schedule), then discretize (argmax) at inference; alternatively, people use a straight-through estimator (hard forward, soft backward) during training. This small hack enables many models with discrete choices VAEs with discrete latents, hard-attention models, etc. though large text-to-image systems sometimes rely on discrete codebooks or other techniques instead of Gumbel relaxations.

[^1]: Ramesh, A., et al. (2021). Zero-Shot Text-to-Image Generation*. OpenAI.

[^2]: Maddison, C. J., et al. (2017). The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables.

What is PTX and Why Does It Matter?

PTX, or Parallel Thread Execution, is an intermediate assembly language used by NVIDIA GPUs to execute programs written in CUDA. By directly interacting with PTX, developers gain fine-grained control over GPU operations, enabling them to bypass some of the inefficiencies introduced by higher-level abstractions like standard CUDA kernels While working at this level requires expertise, it opens up opportunities for groundbreaking optimizations—something DeepSeek has capitalized on extensively.

The Role of utils.cuh in DeepEP

This file contains critical implementations that enable the integration of custom PTX instructions into CUDA workflows. For instance:

• Custom PTX instructions : Within utils.cuh, you'll find code snippets where DeepSeek overrides default CUDA compiler decisions with hand-tuned PTX logic. These modifications target specific hardware features, ensuring optimal utilization of resources.

• Memory Transfer Optimization : One of the standout achievements in DeepEP is the reduction of unnecessary memory transfers between host and device. By embedding PTX instructions directly, DeepSeek minimizes latency and boosts overall efficiency

Enhanced Memory Access :

1. Relaxed Memory access: Instructions like ld.relaxed and st.relaxed minimize synchronization overhead by avoiding unnecessary memory barriers.
2. Non-Allocating Loads: The ld.global.nc.L1::no_allocate instruction bypasses L1 cache allocation, reducing cache pollution for certain workloads.
3. Volatile Loads: Functions like ld.volatile.global ensure data consistency by bypassing compiler optimizations.

Atomic Operations Atomic instructions are crucial for ensuring data integrity in parallel environments. DeepEP includes custom atomic functions like atom.add.release.sys.global, which combine atomicity with relaxed or release semantics for fine-grained control over memory ordering.

How Does It Work?

The customization kernel relies heavily on inline assembly within CUDA, allowing developers to inject raw PTX instructions directly into their code. For example:

__device__ __forceinline__ void st_na_relaxed(const int *ptr, int val) {
    asm volatile("st.relaxed.gpu.global.L1::no_allocate.b32 [%0], %1;" : : "l"(ptr), "r"(val));
}

This function performs a relaxed store operation, bypassing the L1 cache for improved performance in specific scenarios.

Additionally, macros like UNROLLED_WARP_COPY enable optimized data movement within warps by unrolling loops and minimizing instruction latency:

#define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC, ST_FUNC) \
{ \
    constexpr int kLoopStride = 32 * (UNROLL_FACTOR); \
    typename std::remove_reference<decltype(LD_FUNC((SRC) + 0))>::type unrolled_values[(UNROLL_FACTOR)]; \
    auto __src = (SRC); \
    auto __dst = (DST); \
    for (int __i = (LANE_ID); __i < ((N) / kLoopStride) * kLoopStride; __i += kLoopStride) { \
        _Pragma("unroll") \
        for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) \
            unrolled_values[__j] = LD_FUNC(__src + __i + __j * 32); \
        _Pragma("unroll") \
        for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) \
            ST_FUNC(__dst + __i + __j * 32, unrolled_values[__j]); \
    } \
}

This macro is particularly useful for optimizing memory-bound operations in deep learning models.