Asynchronous Data Copies in CuTe DSL

I’ve always been interested in writing performant code, and recently I’ve been learning how to do that for NVIDIA GPUs. Specifically, I’ve been looking into CuTe DSL, which is a domain-specific language based on Python by NVIDIA. It offers a less verbose and slightly less complicated way to program GPUs compared to writing CUDA C++, but also is more expressive and powerful than Triton.

My ultimate goal is to learn how to write something like FlashAttention 4 in CuTe DSL. However, even a “simple” general matrix multiply (GEMM) kernel for an A100 is about 800 lines of code with many abstractions mixed in, and is difficult for me to grasp, despite the many examples provided on GitHub. The docs are also rather terse. So I thought I’d incrementally work up to complicated kernels, starting from a simple copy kernel to understand how to do asynchronous data copies (which are needed for the fancier kernels).

Prerequisites: I assume a basic familiarity with core GPU programming concepts (see, e.g., sections 1 and 2 of ThunderKittens), as well as some knowledge of CuTe DSL, at the level of the elementwise-addition kernel from the docs. Simon Veitner’s An applied introduction to CuTeDSL and Thread-value layouts in CuTe are great supplements to that example.

Basic info

Our goal in this post is to write a kernel that copies data from a source matrix $A$ to a destination matrix $B$, both of shape $M \times N$. The general idea is to split the source and destination matrices into tiles of a fixed size, and have one thread block (a group of threads that execute on the same streaming multiprocessor) perform the copy for one tile. We will copy source tiles from global memory (GMEM) to shared memory (SMEM) asynchronously, and then copy these to destination tensor blocks on GMEM synchronously. The asynchronous copy is not strictly necessary in this example, but it will be for more complicated kernels where we need to overlap computation with memory operations.

Setup

import torch

import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack


if __name__ == "__main__":
    # Create source and destination matrices
    M, N = 8192, 8192
    src = torch.randn(M, N, device="cuda", dtype=torch.float16)
    dst = torch.empty_like(src)

    src_ = from_dlpack(src, assumed_align=16)
    dst_ = from_dlpack(dst, assumed_align=16)

    # Compile and run the kernel
    op = TensorCopyAsync()
    compiled = cute.compile(op, src_, dst_)
    compiled(src_, dst_)

    torch.testing.assert_close(dst, src)

This is straightforward. We create source and destination matrices, both of type float16. The from_dlpack function converts a PyTorch tensor (or a tensor from another DLPack-compatible framework like JAX) into a CuTe tensor with static shape and stride. The TensorCopyAsync class holds the kernel, which we will write next, and after we compile and run it, we verify that dst matches src.

Host code

We initialize the class by providing the number of rows we want in a tile (tile_m) and the number of threads in a thread block (num_threads). We also define a synchronization barrier for the thread block (cta_sync_barrier), where cta stands for “Cooperative Thread Array”, which is CUDA nomenclature for a thread block.

import cutlass
import cutlass.pipeline as pipeline


class TensorCopyAsync:
    def __init__(self, tile_m: int = 32, num_threads: int = 512):
        if num_threads % tile_m != 0:
            raise ValueError("num_threads must be divisible by tile_m")

        self._tile_m = tile_m
        self._num_threads = num_threads
        self.cta_sync_barrier = pipeline.NamedBarrier(
            barrier_id=1, num_threads=num_threads
        )

For the __call__ method, we first calculate the size of the tile based on how many elements we want a single thread to move. Since a single copy instruction can move 32, 64, or 128 bits, here we choose 128, which corresponds to 128 // 16 = 8 elements per thread.

@cute.jit
def __call__(self, mSrc: cute.Tensor, mDst: cute.Tensor):
    copy_bits = 128
    vector_elems = copy_bits // mSrc.element_type.width
    threads_per_row = self._num_threads // self._tile_m
    tile_n = threads_per_row * vector_elems

    if cutlass.const_expr(
        mSrc.shape[0] % self._tile_m != 0 or mSrc.shape[1] % tile_n != 0
    ):
        raise ValueError(
            f"mSrc/mDst shape must be divisible by ({self._tile_m}, {tile_n})"
        )

Here, each thread will copy vector_elems (which is 8) elements which are laid out sequentially in a single row.

The next order of business is to define the copy operations we want. We first want to copy data from src to a staging tile in SMEM, then from that tile back to dst on GMEM. In CuTe, we specify the copy instruction by constructing a CopyAtom, which is a Python class that holds information related to the instruction, such as which copy operation to use and the data type of the elements to be copied. Here we define the two copy atoms: atom_async_copy is for the GMEM to SMEM copy, whereas atom_store is for the reverse.

atom_async_copy = cute.make_copy_atom(
    cute.nvgpu.cpasync.CopyG2SOp(),  # Asynchronous GMEM -> SMEM copy operation
    mSrc.element_type,
    num_bits_per_copy=copy_bits,
)
atom_store = cute.make_copy_atom(
    cute.nvgpu.CopyUniversalOp(),  # General-purpose copy operation
    mDst.element_type,
    num_bits_per_copy=copy_bits,
)

The rest of the code defines a thread-value layout for the tiles, and divides the source and destination tensors according to the layout with cute.zipped_divide. Note the use of cute.make_tiled_copy_tv to create tiled versions of the copy instructions we defined earlier.

thr_layout = cute.make_layout(
    (self._tile_m, threads_per_row), stride=(threads_per_row, 1)
)
val_layout = cute.make_layout((1, vector_elems))
tiled_copy_load = cute.make_tiled_copy_tv(
    atom_async_copy, thr_layout, val_layout
)
tiled_copy_store = cute.make_tiled_copy_tv(atom_store, thr_layout, val_layout)

sSrc_layout = cute.make_layout((self._tile_m, tile_n), stride=(tile_n, 1))
gSrc = cute.zipped_divide(mSrc, (self._tile_m, tile_n))
gDst = cute.zipped_divide(mDst, (self._tile_m, tile_n))

We finally launch the kernel with one tile assigned to one thread block. We use cute.ceil_div to calculate how many tiles we need to cover the entire matrix.

tiles_mn = cute.ceil_div(mSrc.shape, (self._tile_m, tile_n))
self.kernel(
    gSrc,
    gDst,
    sSrc_layout,
    tiled_copy_load,
    tiled_copy_store,
).launch(
    grid=[tiles_mn[1], tiles_mn[0], 1],
    block=[self._num_threads, 1, 1],
)

Device code

For the kernel itself, we start by accessing the source and destination tiles assigned to this thread block.

@cute.kernel
def kernel(
    self,
    gSrc: cute.Tensor,
    gDst: cute.Tensor,
    sSrc_layout: cute.Layout,
    tiled_copy_load: cute.TiledCopy,
    tiled_copy_store: cute.TiledCopy,
):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()

    blkSrc = gSrc[((None, None), (bidy, bidx))]
    blkDst = gDst[((None, None), (bidy, bidx))]

This is followed by allocating memory in SMEM, which we will use for staging.

smem = cutlass.utils.SmemAllocator()
sSrc = smem.allocate_tensor(gSrc.element_type, sSrc_layout, 16)

We now need to create copy instructions for the elements of the tile that only the current thread is responsible for. We do this in CuTe using the following snippet (the function names are self-explanatory):

thr_copy_load = tiled_copy_load.get_slice(tidx)
tSgSrc = thr_copy_load.partition_S(blkSrc)
tSsSrc = thr_copy_load.partition_D(sSrc)
cute.copy(tiled_copy_load, tSgSrc, tSsSrc)

When using cute.copy with an asynchronous copy instruction, we also have to use the following code to instruct the GPU to execute it.

cute.arch.cp_async_commit_group()  # Submit this thread's queued operations
cute.arch.cp_async_wait_group(0)  # Wait until this thread's operations are done
self.cta_sync_barrier.arrive_and_wait()  # Sync between all threads in the thread block

We finally copy data from the SMEM tensor to dst.

thr_copy_store = tiled_copy_store.get_slice(tidx)
tSsStore = thr_copy_store.partition_S(sSrc)
tSgDst = thr_copy_store.partition_D(blkDst)
cute.copy(tiled_copy_store, tSsStore, tSgDst)

That’s it! On profiling this kernel with Nsight Compute on an A100, it achieves 87% of the bandwidth of the GPU, which is pretty good.

Parting thoughts

The full code is available here. I tried experimenting with different tile sizes and number of threads, but 87% is the most I could achieve. If you have ideas on how to improve the throughput even more, please send me an email or open a GitHub issue!

Next up is a GEMM kernel for the Ampere generation (following the example here).

Steering Language Models with Sequential Monte Carlo

A good chunk of my work as part of my PhD involves using sequential Monte Carlo (SMC) methods to solve decision-making problems. SMC algorithms are used to efficiently sample from sequences of distributions, and while they are mostly used in physics, signal processing and Bayesian statistics, recently they have also found uses in inference-time alignment of generative models. In this post, I’ll show how I used SMC to “steer” a tiny pre-trained language model into writing sad stories.

Background: How LLMs Generate Text

A language model works over a vocabulary $\mathsf{S}$ of tokens. At each step $t \in \mathbb{N}$, it predicts the probability distribution of the next token given the history so far:

$$s_{1:t-1} \mapsto p(\cdot \mid s_{1:t-1}) \in \mathcal{P}(\mathsf{S}),$$

where $\mathcal{P}(\mathsf{S})$ is the set of distributions over the vocabulary $\textsf{S}$ and $s_{1:t-1} \coloneqq \{ s_{1}, s_{2}, \dots, s_{t-1} \}$. Generating text with the language model is the task of sampling a sequence $s_{1:T}$ from the joint distribution

$$\mathbb{P}_{T}(s_{1:T}) \coloneqq \prod_{t=1}^{T} p(s_{t} \mid s_{1:t-1}), \quad \text{where } s_{1:0} \coloneqq \emptyset.$$

Inference-Time Alignment of LLMs

Inference-time alignment refers to modifying a pre-trained model’s sampling behavior without changing its parameters. Instead of fine-tuning the model weights, we intervene at sampling time and adjust how likely the model is to pick certain continuations. Sequential Monte Carlo provides a natural framework for this: by iteratively reweighting and resampling partial generations based on a reward signal, we can nudge the model toward desired behaviors while retaining some of its inherent randomness and diversity.

Formally, we specify our preferences through a sequence of reward functions

$$r_{t}: \mathsf{S}^{t} \to \mathbb{R}, \quad t \geq 1,$$

which score partial sequences. This setup is quite general: $r_t$ could encode stylistic preferences, safety constraints, or factuality. In my case, $r_{t}$ is just a neural network that has been trained to output a “sadness score” in $[0, 1]$ for a given phrase. We then define a new distribution $\mathbb{Q}_{T}$ over sequences, which is tilted towards the cumulative reward:

$$\mathbb{Q}_{T}(s_{1:T}) \propto \mathbb{P}_{T}(s_{1:T}) \cdot \exp\left\{\sum_{t=1}^{T} \eta \cdot r_{t}(s_{1:t})\right\}, \quad \eta > 0.$$

Intuitively, $\mathbb{Q}_T$ reweights the model’s likelihoods so that sequences with higher cumulative reward become exponentially more probable.

Sampling from $\mathbb{Q}_{T}$ is straightforward with SMC. The recipe is as follows, with steps 1 and 2 repeated for all $n \in \{1, \dots, N\}, N \in \mathbb{N}$:

1. Propose: Sample a token $s_{t}^{n} \sim p(\cdot \mid s_{1:t-1}^{n})$ and append to sequence, $s_{1:t}^{n} = (s_{1:t-1}^{n}, s_{t}^{n})$.
2. Weight: Compute unnormalized weight $w_{t}^{n} = \exp \{ \eta \cdot r_{t}(s_{1:t}^{n}) \}$, then normalize: $W_{t}^{n} = w_{t}^{n} / \sum_{m=1}^{N} w_{t}^{m}$.
3. Resample: Draw $N$ new particles from $\{s_{1:t}^n\}_{n=1}^N$ with replacement, proportionally to $W_{t}^{n}$. (This duplicates the ‘sad’ phrases and prunes away overly cheerful ones.)
4. Repeat: Until $t = T$.

Over time, the population of particles gradually concentrates on high-reward trajectories, in this case the sadder continuations.

Sob Story Time

To illustrate the method, I’m using TinyStories-33M, a language model trained on short children’s stories (Eldan and Li 2023). Importantly, this model is not fine-tuned for sadness (or anything else), and is just a pure text predictor. The code accompanying this post is available on GitHub.

I gave TinyStories the prompt:

“When the prince came home, he saw”

Here’s what the base model produced (one sample from $\mathbb{P}_{T}$):

“When the prince came home, he saw the heavy bag of jewelry. He wanted to buy it and wear it. He asked the king to sell it to him. The kind king said ‘Yes!’, and …”

Hmm, way too cheerful for our tastes. Now here’s what happens after steering with SMC (one sample from $\mathbb{Q}_{T}$):

“When the prince came home, he saw the sad family sitting by the stove. He felt very sad too. He had lost his rare treasure box and now it was gone forever.”

That’s more like it!

Bonus: Steering as Optimization

A simple but neat result is that the steered distribution $\mathbb{Q}_{T}$ is the minimizer of

$$\mathcal{L}(Q) = \mathbb{E}_{Q} \left[ -\sum_{t=1}^{T} r_{t}(s_{1:t}) \right] + \frac{1}{\eta} \, \mathcal{D}_{\mathrm{KL}} \left[ Q \Vert \mathbb{P}_{T} \right], \quad Q \in \mathcal{P}(\mathsf{S}),$$

where $\mathcal{D}_{\mathrm{KL}}$ is the Kullback-Leibler (KL) divergence (see, e.g., Bissiri et al. 2016). The first term on the RHS is responsible for maximizing the cumulative reward (sadness), while the second term is a regularizer forcing $Q$ to stay close to the base model $\mathbb{P}_{T}$. This regularization term prevents collapse into a small number of “super sad” trajectories, thus preserving diversity of model outputs.

This optimization perspective also makes clear the connection to reinforcement learning from human feedback (RLHF, Ziegler et al. 2020), where the same objective $\mathcal{L}(Q)$ is minimized by fine-tuning the model weights. Here, we’re skipping the optimization and directly sampling from the minimizer with SMC.

Parting Notes

The algorithm presented here is the simplest version of SMC (known as the “bootstrap” particle filter), and in practice more sophisticated techniques are required to actually deliver on the promise of preserving output diversity. These techniques include adaptive resampling schedules and twisting, see, e.g., Naesseth et al. (2019). For a state-of-the-art application of SMC to language models, I recommend Zhao et al. (2024).

References

Using LaTeX Snippets in Markdown Files in Neovim

I frequently write notes in Markdown using Neovim, and many of those notes contain math. Unfortunately, all of my handy LaTeX snippets created with LuaSnip don’t work in Markdown, which makes it a pain to write anything beyond the simplest equations. In this post, I document how I got my snippets working in Markdown.

Basic Setup

The first step is to tell LuaSnip to load LaTeX snippets in Markdown files. This saves us from having to duplicate all our LaTeX snippets. Since I use lazy.nvim as my package manager, I just had to tweak the config function like this:

config = function()
  local ls = require("luasnip")

  ls.config.set_config({
    ...
  })

  -- Make TeX snippets available in markdown
  ls.filetype_extend("markdown", { "tex" })

  ...
end,

It turns out that this is all we need to get basic snippet functionality working. However, many of my snippets use conditional expansion, with snippets only expanding if the cursor is inside a math environment. For example, typing ff automatically expands to \frac{}{} when I’m inside dollar signs or an align environment, but it won’t expand when I type ‘affirm’ or ‘puff’ in normal text. Since a lot of my go-to snippets rely on this, that’s the next thing I set up.

Conditional Expansion in Math Environments

When I write LaTeX files, I rely on a package called VimTex for syntax highlighting and a bunch of other useful features. One of its perks is that it can detect math environments, which we can use as a condition for snippet expansion with the following Lua function (see the article by ejmastnak for more details):

tex_utils.in_math = function()
  -- This function requires the VimTeX plugin.
  return vim.fn["vimtex#syntax#in_mathzone"]() == 1
end

But VimTeX is a filetype-specific plugin, so it only loads for .tex files. To get the same functionality in Markdown, we need to create a nvim/after/syntax/markdown.vim file with the following contents:

" From https://github.com/lervag/vimtex/issues/2395
if exists('b:current_syntax')
  unlet b:current_syntax
endif

syn include @tex syntax/tex.vim
syn region markdownMath start="\\\@<!\$" end="\$" skip="\\\$" contains=@tex keepend
syn region markdownMath start="\\\@<!\$\$" end="\$\$" skip="\\\$" contains=@tex keepend

let b:current_syntax = 'markdown'

This snippet first clears any existing Markdown syntax settings, then tells Neovim to pull in LaTeX syntax rules from syntax/tex.vim (provided by VimTeX). It defines $...$ and $$...$$ as math environments, so anything inside them gets highlighted and treated like LaTeX math, even in a Markdown file.

Bonus for VimWiki Users

If you use VimWiki for note-taking and also rely on the <Tab> key for snippet expansion, you’ll need to free up <Tab> in insert mode, since VimWiki already uses it for shortcuts. I solved this by adding the following autocommand for VimWiki:

config = function()
  vim.api.nvim_create_autocmd("FileType", {
    pattern = "vimwiki",
    callback = function()
      -- Disable <Tab> in insert mode.
      vim.keymap.del("i", "<Tab>", { buffer = true })
    end,
  })
end,

Conclusion

That’s it, we’re all set! We’re now ready to give Giles Castel a run for his money in Markdown ;). That said, this setup is not perfect. The biggest limitation is that it only recognizes math environments defined with single or double dollar signs, so conditional snippets won’t work inside environments like align. This is not a huge deal for me though, since all serious math gets done in .tex files. If you’re curious, you can check out my full Neovim configuration here.

Expected Information as Expected Utility

Imagine a scientist planning a clinical trial to determine the optimal dosage of a new diabetes drug. Due to constraints on budget and patient safety, only a limited number of tests can be conducted, so the scientist must decide which dosage levels to test to most effectively learn the relationship between dose and patient response. To make these decisions systematically, the scientist can turn to Bayesian experimental design (BED), which is a principled framework for designing optimal experiments under uncertainty.

Suppose the scientist models the dose–response relationship using a parametric form governed by unknown parameters $\theta \in \Theta$—such as the maximum effect of the drug, the dose achieving half that effect, and the baseline response. To learn about these parameters, she must choose a design $\xi$, which could specify, for example, a set of dosage levels and how patients are allocated to them. A widely used objective in BED is to choose the design that maximizes the expected information gain (EIG) (Lindley 1956):

$$\mathrm{EIG}(\xi) \coloneqq \mathbb{E}_{p(y \mid \xi)} \Bigl[ \mathbb{H} \bigl[ p_{\theta}(\theta) \bigr] - \mathbb{H} \bigl[ p_{\theta \mid y, \xi}(\theta) \bigr] \Bigr].$$

It quantifies how much we expect to reduce our uncertainty about $\theta$, starting from some prior belief $p_{\theta}$, on applying the design $\xi$. Here, $\mathbb{H}$ denotes Shannon entropy (or differential entropy for continuous $\theta$), and the expectation is over possible outcomes $y$ with distribution $p(y \mid \xi)$.

The EIG is intuitive—it is the expected reduction in entropy from prior to posterior. But intuition aside, it raises a natural question. The reduction in Shannon entropy is only one of many ways to measure uncertainty reduction, or even more generally, the utility of an experiment $(\xi, y)$ (Huan et al. 2024). Is there, then, any fundamental reason to prefer the EIG over alternative utility measures?

It turns out the answer is yes. In Bernardo (1979), José M. Bernardo, a student of Dennis Lindley, gave a justification for the EIG grounded in decision theory. He showed that the EIG arises naturally from a decision problem under reasonable assumptions on the utility function. I find this result extremely remarkable, and this is an appreciation post about a half-century later :)

The Decision-Theoretic Setup

Consider a decision problem where a scientist has to report a distribution $q \in \mathcal{P}(\Theta)$ for $\theta$, where the decision space $\mathcal{P}(\Theta)$ is the set of probability measures on $\Theta$.¹ The utility obtained if she reports $q$ when the true parameter is $\theta$ is quantified by means of a utility function $u:\mathcal{P}(\Theta) \times \Theta \to \mathbb{R}$.² Suppose the scientist applies a design $\xi$ and observes an outcome $y$, after which she updates her belief to $p_{\theta \mid y, \xi}(\theta)$ using Bayes’ rule. Then, her expected utility when reporting $q$ is

$$\int u(q, \theta) \, p_{\theta \mid y, \xi}(d\theta).$$

Under the principle of maximum expected utility, the scientist should report the distribution $q^{*}$ that maximizes the above integral, which may not be her actual belief $p_{\theta \mid y, \xi}$. Hence, to discourage lying, we require that $u$ is a strictly proper utility function, meaning that

$$\sup_{q} \int u(q, \theta) \, p_{\theta \mid y, \xi}(d\theta) = \int u(p_{\theta \mid y, \xi}, \theta) \, p_{\theta \mid y, \xi}(d\theta),$$

with the supremum only attained at $q^{*} = p_{\theta \mid y, \xi}$. In other words, truth-telling should be the optimal strategy.

The second restriction we need on the utility function is that it is local, i.e., $u(q, \theta) = u\bigl(q(\theta), \theta\bigr)$ for all $\theta \in \Theta$. Locality means the utility function only depends on the density of the reported distribution at the true parameter $\theta$, which feels reasonable (though admittedly not as easy to motivate as properness). The striking result of Bernardo (1979) is that if the utility function $u$ is strictly proper, local, and sufficiently smooth (as a function of $q$), then it is of the form

$$u(q, \theta) = A \log q(\theta) + B(\theta),$$

where the constant $A$ and function $B$ can be arbitrary. Plugging this back into the requirement that $u$ is proper, we see that the maximum expected utility is

$$\begin{align*} \sup_{q} \int u(q, \theta) \, p_{\theta \mid y, \xi}(d\theta) & = \int \left( A \log p_{\theta \mid y, \xi}(\theta) + B(\theta) \right) p_{\theta \mid y, \xi}(d\theta) \\ & = -A \mathbb{H}\left[ p(\theta \mid x, \xi) \right] + \mathbb{E}_{\theta \mid y, \xi}[B(\theta)]. \end{align*}$$

Lo and behold, the Shannon entropy appears!

Back to Experimental Design

Now, the gain in expected utility for the scientist from the experiment $(\xi, y)$ is

$$\begin{align*} \mathrm{Gain} & = \textrm{Max utility after observing } y - \textrm{Max utility from the prior} \\ & = \sup_{q} \int u(q, \theta) \, p_{\theta \mid y, \xi}(d\theta) - \sup_{p} \int u(p, \theta) \, p_{\theta}(d\theta), \end{align*}$$

which simplifies to

$$A \Bigl[ \mathbb{H}\bigl[p_{\theta}(\theta)\bigr] - \mathbb{H}\bigl[ p_{\theta \mid y, \xi}(\theta) \bigr] \Bigr] + \mathbb{E}_{\theta \mid y, \xi}[B(\theta)] - \mathbb{E}_{\theta}[B(\theta)].$$

To assess the quality of a design $\xi$, we take the expectation over the observation distribution $p(y \mid \xi)$, leaving us with

$$\begin{align*} & A \cdot \mathbb{E}_{y \mid \xi}\Bigl[ \mathbb{H}\bigl[p_{\theta}(\theta)\bigr] - \mathbb{H}\bigl[ p_{\theta \mid y, \xi}(\theta) \bigr] \Bigr] + \underbrace{\mathbb{E}_{y \mid \xi} \Bigl[ \mathbb{E}_{\theta \mid y, \xi}[B(\theta)] - \mathbb{E}_{\theta}[B(\theta)] \Bigr]}_{=0} \\ & = A \cdot \mathrm{EIG}(\xi). \end{align*}$$

This shows that, for the purpose of choosing the best design, the choice of $A$ and $B$ is irrelevant, and the utility function can simply be the log density, $u(q, \theta) = \log q(\theta)$.

Concluding Remarks

To motivate his paper, Bernardo writes in the abstract:

… a scientist typically does not have, nor can be normally expected to have, a clear idea of the utility of his results.

By considering a decision problem of choosing the best distribution to report, and showing that the EIG arises naturally under reasonable assumptions on the utility function, Bernardo makes a compelling argument for using the EIG in precisely those settings where the scientist cannot quantify the utility of her results a priori.

References