July 20, 2026
Here are two ways to sum the squares of the elements a mask selects. They compute the same number:
import torch
torch.manual_seed(0)
data = torch.randn(1000)
mask = torch.rand(1000) < 0.7
by_multiply = (data * data * mask.float()).sum() # compute everything, zero out the rest
by_indexing = (data[mask] ** 2).sum() # pull out the selected elements first
print(f"multiply by mask : {by_multiply.item():.4f}")
print(f"boolean indexing : {by_indexing.item():.4f}")multiply by mask : 707.5559
boolean indexing : 707.5559Now the question that matters: on a GPU, which is faster?
The intuitive answer is the second. It’s obviously doing less work — if 30% of the elements are masked out, it squares 700 numbers instead of 1000, and it never wastes a multiply on a value it’s going to discard. Counting arithmetic, it wins comfortably.
On many eager CUDA workloads at moderate mask densities, the second one is slower despite doing less arithmetic — because producing a variable-length result costs a compaction and, in eager CUDA, a host-visible size dependency that the arithmetic saving doesn’t come close to paying for. Let me be careful about how strongly I state that: I wrote this post without a GPU to hand, so treat it as a hypothesis with a mechanism attached, which the experiment at the end is designed to test. At extreme sparsity, or when compaction saves substantial downstream work, the ordering can flip.
Understanding why is the subject of this post, and it rests on one fact about how the machine actually works:
A GPU is an asynchronous device sitting at the end of a queue. Your Python code doesn’t execute operations; it enqueues them and races ahead. The GPU works through that queue behind you. As long as the queue stays full, the machine stays busy — and anything that forces Python to consume a device result creates a host dependency on the preceding device work. While the submitting thread waits, it can’t enqueue anything that follows.
data[mask] produces a tensor whose size depends on the data. In ordinary eager CUDA execution, returning that compact tensor means establishing an exact, host-visible output size first — torch.nonzero, which exposes the same data-dependent count, documents this synchronization explicitly. (Compiled execution can instead carry supported data-dependent sizes symbolically, so profile the path you actually deploy.) The host then can’t proceed past that dependency until the device work completes, nor submit dependent work from that thread while it waits. You saved 300 multiplies, each a fraction of a nanosecond, and paid with a stall costing microseconds or worse. (Treat the exact kernel and synchronization sequence as version- and backend-dependent, and confirm it with a profiler rather than taking my word for the mechanism.)
That possible inversion — less arithmetic, more cost — is the thing to internalize, because your intuition prices operations by how much they compute, and the machine prices them by how much they make it wait.
TL;DR — A GPU executes asynchronously behind a queue, and performance can be limited by host synchronization, launch overhead, memory traffic, compute throughput, allocation, or shape dynamism. Which of those dominates is a property of the workload, not a fixed ranking — but FLOP counting is rarely where to start looking. Data-dependent output shapes are often harder to compile and replay efficiently, though modern
torch.compilecan represent many of them. This post builds that diagnostic model, shows thecumsumtrick that avoids a host round trip entirely, and — because none of it is worth believing unmeasured — builds the timing harness that keeps you honest, then uses it.
Several things compete to be your bottleneck, and which one wins is a property of the workload rather than a fixed law. What follows is a diagnostic order — what to suspect and measure first — not a ranking that holds everywhere.
1. Host dependencies. Anything that forces Python to consume a device result makes the submitting thread wait for the preceding device work to finish. (Already-queued independent work on other streams can still proceed — this stalls the thread, not necessarily the whole device.) The usual offenders:
.item(), .tolist(), float(t), int(t), print(t).cpu(), .numpy()x[mask] and torch.nonzero — in ordinary eager CUDA execution the exact output size has to be known before the compact tensor can be returned, and torch.nonzero documents that synchronization explicitlyif t > 0: — any Python control flow that branches on a tensor’s value2. Kernel-launch overhead. In eager CUDA code most computational operations launch one or more kernels, each carrying fixed overhead commonly measured on the microsecond scale, depending on hardware and launch path. (Not everything launches: metadata-only operations like transpose launch nothing, one library call may launch several kernels, and a compiled graph may fuse many expressions into one.) A Python loop issuing 128 tiny operations pays that overhead 128 times and never gives the machine enough work to fill it. This is the real reason the brute-force loop from the indexing post is unusable — not that its arithmetic is slow, but that it’s 128 Python-dispatched indexing operations where gather expresses the same work as one vectorized operator. (The precise kernel count on either side is implementation-dependent; the dispatch count is not.)
3. Memory bandwidth. Many kernels — elementwise, indexing, normalization, decode-step attention, anything with low data reuse — wait on DRAM rather than on arithmetic. The size of that gap depends on which units the kernel can actually use, and it is easy to quote misleadingly. A T4 offers roughly 65 fp16 Tensor Core TFLOPS and roughly 8.1 fp32 CUDA core TFLOPS against roughly 320 GB/s of bandwidth. That’s about 200 FLOPs per byte if you can reach the Tensor Core peak, and about 25 for ordinary fp32 elementwise work on CUDA cores — which is the relevant figure here, since an elementwise multiply cannot use Tensor Cores at all. Either way an elementwise op performs about one FLOP per byte, so it is memory-limited by a wide margin. This is also why fusion buys so much: not fewer FLOPs, fewer round trips to memory.
4. Shape dynamism. Static shapes are generally easier to specialize, fuse, allocate for, and replay. CUDA Graph replay in particular has strict requirements around stable execution structure and memory addresses. Modern torch.compile, by contrast, can represent many dynamic and even data-dependent output sizes using symbolic shapes — so this is not “dynamic shapes break the compiler.” Dynamic outputs can still reduce optimization opportunities, trigger recompilation, or fail when an unbacked value reaches unsupported control flow, and compiler behaviour is version-, configuration-, and backend-dependent. It’s part of why production inference often pads to fixed lengths and masks rather than slicing ragged: the “wasted” compute on padding buys predictability.
5. Allocation, which isn’t free either — every eager intermediate is a trip through the caching allocator.
6. Raw compute, which does sometimes win: a large GEMM or convolution can genuinely be arithmetic-bound. It’s simply rarer than beginners assume for the elementwise, indexing, normalization and decode-shaped work that fills most models.
The useful form of this isn’t a ranking to memorize but a set of questions to ask a slow program, roughly in this order: Is the host waiting on a device result? Is it drowning in tiny launches? Is it moving far more data than it reuses? Is it actually compute-bound? Is allocation or shape dynamism in the way? Is the compiler fusing what I think it is? Which one dominates is workload-dependent — the point of the model is to tell you what to measure first, which is why this post ends up building a measurement harness rather than a table of adjectives.
With the model in hand, the comparison from the opening becomes readable — and a third option appears that keeps a static shape while being numerically safer than multiplying by the mask:
by_where = torch.where(mask, data * data, torch.zeros_like(data)).sum()
print(f"multiply by mask : {by_multiply.item():.4f}")
print(f"boolean indexing : {by_indexing.item():.4f}")
print(f"torch.where : {by_where.item():.4f}")
torch.testing.assert_close(by_multiply, by_where)
print("multiply and where agree within floating-point tolerance ✓")multiply by mask : 707.5559
boolean indexing : 707.5559
torch.where : 707.5559
multiply and where agree within floating-point tolerance ✓(The boolean-index result differs in the final digit — not an error, just a different summation order over a different number of terms. Floating-point addition isn’t associative, which the numerical stability post has more to say about.)
(d*d*mask).sum() |
(d[mask]**2).sum() |
where(mask, d*d, 0).sum() |
|
|---|---|---|---|
| Output shape | static | data-dependent | static |
| Device→host dependency | none | commonly required in eager CUDA | none |
| Allocation | fixed-size intermediates (eager) | variable-size buffer plus intermediates | fixed-size intermediates (eager) |
| Wasted arithmetic | yes (×0 on masked) | none | yes (×0 on masked) |
| Compiler behaviour | static shape; generally fusible | data-dependent output shape — see below | static shape; generally fusible |
Safe against 0 × inf? |
no | yes | forward only — see below |
At moderate densities with a cheap downstream reduction, the middle column is the one I’d expect to lose despite doing the least arithmetic, because it avoids nanoseconds of multiplication and pays for a variable-length compaction. At extreme sparsity, or where compaction removes substantial downstream work, selection may win — establishing the crossover is exactly what the benchmark sweep at the end of this chapter is for.
Two honest notes on that table. “Fixed-size intermediates” is not “no allocation” — in eager mode data * data and torch.zeros_like(data) each materialize a full-size tensor; the difference from the middle column is that their sizes are known in advance, which is exactly what lets a compiler fuse them away. And the compiler row deserves care: I originally wrote graph break there, and that was wrong. Current torch.compile can capture supported data-dependent-output operations using unbacked symbolic shapes, and fullgraph=True enables that capture by default — I confirmed it compiles boolean-mask indexing on this build. That’s not a promise it always works: unsupported operations, control flow on unbacked values, or backend limits still fail, and fullgraph=True turns those into compilation errors rather than proving everything is supported. Under some default configurations you get graph breaks or reduced optimization instead. The honest statement is that a data-dependent output shape makes capture and fusion configuration- and version-dependent, where a static shape makes them routine. It’s worth separating four things that get blurred together: CUDA-graph replay constraints, torch.compile’s dynamic-shape support, data-dependent output shapes, and graph breaks. They interact; they are not the same thing.
The third column is the one I’d reach for, though for a reason worth separating from performance. Multiplying by the mask has a flaw the table hints at: if a masked-out slot holds inf or nan — a padding sentinel, an upstream overflow — then inf × 0 = nan and the entire reduction is poisoned. torch.where selects, so the garbage never enters the sum. Except where is eager: both branches are always evaluated, and if gradients flow, the chain rule reconstructs the same 0 × inf one level down, giving you a finite forward pass and a nan gradient. The form that survives both directions sanitizes the input before the dangerous operation:
sanitize-then-square gives the same value, static shape, no compaction ✓It keeps the static output shape, avoids the data-dependent compaction, and fixes the gradient. The three-rung derivation, with the verified nan gradient, is in the numerical stability post.
So for the regime motivating this example, the implementation I’d benchmark first is sanitize-then-operate on a static shape. Its numerical advantage over multiply-by-mask is established right here and holds regardless of hardware; its performance advantage over compaction is the hypothesis the sweep at the end is meant to test. Real code needs both kinds of reasoning, and it’s worth keeping straight which one you’ve actually demonstrated.
Once you can see stalls, you start finding them in code that looks innocent. Here’s one:
It writes 0, 1, 2, … into the masked-in positions in order, and it’s doing more work than it looks. The right-hand data[mask] builds a variable-length selection purely so Python can ask it for a length — and a host-visible length guarantees at least one synchronization. The left-hand side is then a separate masked write. (I originally described this as “two nonzero calls, two syncs”; that was more confident than the semantics support — the assignment target is a __setitem__, not a second read, and how many kernels and sync points that costs internally is backend-dependent.) Separately, the arange is built on the CPU, since no device= was passed, so it also needs copying across.
The obvious rewrite removes the redundant masked read and the CPU-to-device helper copy, but still leaves one explicit host dependency:
You can get to zero, and the trick generalizes far beyond this example. The observation is that a cumulative sum over a boolean mask is a running count of how many Trues you’ve seen — so subtracting one gives each True its 0-based rank, which is exactly the arange you wanted:
mask [True, False, True, True]
mask.cumsum(0) [1, 1, 2, 3]
minus 1 [0, 0, 1, 2]At the k-th True, cumsum - 1 equals k. The values sitting under the False positions are meaningless, but harmless — where doesn’t select them. (It does still evaluate the rank expression at every position, since where is eager; that’s fine here because the unselected values are merely irrelevant rather than dangerous, which is exactly the distinction the masking section turned on.)
torch.manual_seed(0)
for _ in range(500):
n = torch.randint(4, 64, (1,)).item()
d = torch.randn(n)
mk = torch.rand(n) < 0.6
if mk.sum() == 0:
continue
reference = d.clone() # reference: variable-length indexing
reference[mk] = torch.arange(reference[mk].size(0)).float()
fast = torch.where(mk, (mk.cumsum(0) - 1).float(), d) # no host-visible size
assert torch.equal(reference, fast)
print("500 random trials: cumsum + where reproduces the indexed assignment exactly ✓")500 random trials: cumsum + where reproduces the indexed assignment exactly ✓One thing to be clear about: that’s a CPU correctness test, and its scaffolding — the .item(), the if mk.sum() == 0 branch, the reference[mk].size(0) — would all create host dependencies on CUDA. Only the candidate expression, torch.where(mk, (mk.cumsum(0) - 1).float(), d), is the proposed hot-path code. It has no nonzero, no data-dependent shape, and no host round trip. It’s a scan followed by a pointwise selection rather than a single fused kernel — cumsum is a scan, and I shouldn’t imply a compiler will fold it into the where — but it stays entirely on-device and keeps a static output shape, which is the property that matters.
The objection people raise is that cumsum is a sequential scan and therefore slow. GPU scan implementations are highly parallel, and keeping a static dense output can be favourable when the alternative is a host-visible variable-length result — but whether it wins for a given shape and consumer is an empirical question I haven’t measured, so take it as a reason to try both rather than a verdict. More importantly the pattern generalizes: turn a masked position’s rank into a value with cumsum, then select with where. That’s how you number valid tokens or build segment IDs. Note what it does not do: the result stays dense, carrying rank-derived values in the masked-in slots. If what you genuinely need is a compact, variable-length tensor, this doesn’t replace nonzero — it replaces the cases where you only thought you needed one.
Everything above is unfalsifiable until you can time it, and the obvious way to time GPU code is wrong in a way that will actively deceive you:
The CPU enqueues and returns immediately. Your timer captures the enqueue. You will observe an impossible speedup and believe it — this is the single most common way people convince themselves an optimization worked.
The fix has three mandatory parts:
import torch, time, statistics
def benchmark_cuda(fn, warmup=10, iters=100):
for _ in range(warmup): # 1. warm up: CUDA context, autotuning, allocator growth
fn()
torch.cuda.synchronize() # 2. drain the queue BEFORE starting the clock
t0 = time.perf_counter()
for _ in range(iters):
fn()
torch.cuda.synchronize() # 3. drain the queue BEFORE stopping it
return (time.perf_counter() - t0) / itersFor numbers you intend to publish rather than just to build intuition, graduate to torch.cuda.Event(enable_timing=True) pairs, which timestamp on the device itself:
Note the end.synchronize() — recording an event doesn’t make its timestamp readable; you have to wait for it, just more narrowly than a full-device torch.cuda.synchronize(). Better still, use torch.utils.benchmark.Timer, which handles warmup, repetition, synchronization, thread control and statistics for you — it exists precisely because everyone’s hand-rolled version has a subtle bug in it. And record the environment beside the result: device, torch version, thread count, dtype, shapes. A timing without its context is an anecdote.
The GPU-specific effects need a GPU, but two costs this series has been hand-waving about are measurable anywhere — and they’re a good excuse to practise the discipline above rather than merely describe it. First, the environment, because a timing without it is an anecdote:
import platform
from torch.utils.benchmark import Timer
print(f"python {platform.python_version()} | torch {torch.__version__} | {platform.machine()} | "
f"{'cuda' if torch.cuda.is_available() else 'cpu'}")
print(f"PyTorch default intra-op threads: {torch.get_num_threads()}")
print(f"Timer benchmark threads : 1 (its default — see the note below)")python 3.9.6 | torch 2.8.0 | arm64 | cpu
PyTorch default intra-op threads: 4
Timer benchmark threads : 1 (its default — see the note below)I want to be concrete about why this matters, because an earlier draft of this post timed each operation once and published the ratio. Timing x.t().contiguous() a single time gave me speedup factors of 1130×, then 5×, then 531× on three consecutive runs of the same code — and I had confidently written one of those numbers into the prose. A post that tells you to warm up and repeat should not publish single-shot measurements. So rather than hand-roll it a second time, here is the library the post just recommended, doing the job:
x = torch.randn(2048, 2048)
xt = x.t() # zero-copy view, non-contiguous
# NOTE: no num_threads= argument. Timer defaults to 1 thread deliberately, so a
# comparison isn't dominated by threadpool oversubscription — passing
# get_num_threads() on a many-core box can inflate a small copy by orders of magnitude.
noop = Timer("x.contiguous()", globals={"x": x}).blocked_autorange(min_run_time=0.5)
copy = Timer("xt.contiguous()", globals={"xt": xt}).blocked_autorange(min_run_time=0.5)
print(f"contiguous() on contiguous data : {noop.median*1e3:8.4f} ms")
print(f"contiguous() after transpose : {copy.median*1e3:8.4f} ms")
print(f"tensor payload materialized : {x.numel()*4/2**20:.0f} MiB")contiguous() on contiguous data : 0.0000 ms
contiguous() after transpose : 1.5721 ms
tensor payload materialized : 16 MiBThe no-op is essentially free; the real rearrangement costs milliseconds. Those numbers are specific to the machine printed above — the point is the shape of the result, not the multiplier. When a reshape after a permute isn’t stride-compatible, this is the bill it quietly pays — the hidden copy the strides post warned about, now with a measurement attached. (Not every such combination copies: some permuted layouts remain view-compatible for the shape you’re asking for, which is exactly why that post insists you verify aliasing — comparing a.untyped_storage().data_ptr(), which identifies the allocation, rather than a.data_ptr(), which identifies the first element — instead of inferring it from syntax.)
The same for expand versus repeat:
a = torch.randn(1024, 1)
exp = Timer("a.expand(1024, 1024)", globals={"a": a}).blocked_autorange(min_run_time=0.5)
rep = Timer("a.repeat(1, 1024)", globals={"a": a}).blocked_autorange(min_run_time=0.5)
print(f"expand : {exp.median*1e6:9.3f} us (stride {a.expand(1024,1024).stride()}, "
f"allocates no repeated element storage)")
print(f"repeat : {rep.median*1e6:9.3f} us (allocates {1024*1024*4/2**20:.0f} MiB)")
assert torch.equal(a.expand(1024, 1024), a.repeat(1, 1024))
print("identical values, very different construction cost ✓")expand : 0.374 us (stride (1, 0), allocates no repeated element storage)
repeat : 59.736 us (allocates 4 MiB)
identical values, very different construction cost ✓Identical contents, substantially different construction cost and allocation behaviour — the exact ratio varies a lot by machine, so read the shape of the result rather than the multiplier. For a read-only consumer that handles zero-stride inputs efficiently — the index tensor you expand before a gather, say — expand avoids a copy you never needed. That’s a claim about construction, though: a downstream kernel may reject zero strides, materialize internally anyway, or run faster on contiguous data, and a tensor reused many times can amortize the copy. If it matters, benchmark the consumer end-to-end rather than the construction.
(NumPy users: the sync issue doesn’t exist for you, since NumPy is synchronous — arr[mask] is just a copy. But launch overhead has an analogue in Python-level loop overhead, and the bandwidth and contiguity arguments transfer directly; np.ascontiguousarray is contiguous(), and np.broadcast_to is expand.)
I don’t have a GPU in the environment this post was written in, so rather than invent numbers, here’s the experiment with a falsifiable prediction attached. Run it on a free Colab T4:
N = 10_000_000
data = torch.randn(N, device='cuda')
mask = torch.rand(N, device='cuda') < 0.7
t_mul = benchmark_cuda(lambda: (data * data * mask.float()).sum())
t_index = benchmark_cuda(lambda: (data[mask] ** 2).sum())
t_where = benchmark_cuda(lambda: torch.where(mask, data*data, torch.zeros_like(data)).sum())One axis isn’t enough to learn anything, though. A real sweep varies at least: N across a few orders of magnitude; mask density from 0.1% to 90%, since that’s the variable most likely to flip the result; eager versus compiled; a cheap downstream reduction versus an expensive one, since compaction only pays when it saves later work; and dtype.
My prediction, written down before running it so it can be wrong: where ≈ mul ≪ index at moderate densities, with the gap widening as N grows — and index closing or winning at extreme sparsity, or where compaction avoids substantial downstream computation. If it comes out otherwise, the cost model in this post needs revising, not the measurement. For a post with this title, one real GPU sweep is the missing piece, and it’s the first thing I’d add.
The reason to carry this cost model around is that it catches real bugs in code that has nothing to do with masking. In the KV cache post, the natural way to extend a cache is:
torch.cat allocates a new tensor and copies the entire existing cache into it — every layer, every decode step. The step-t copy moves t × bytes_per_token, so summed across a generation from n₀ to n₁ tokens the total traffic is roughly
\[\text{bytes copied} \;\approx\; \text{bytes/token} \times \frac{n_1^2 - n_0^2}{2}\]
36 KiB/token -> 4.5 GiB copied during one generationFour and a half gigabytes of data the GPU already had, moved for no reason, inside a loop that is already memory-bandwidth-bound. Nothing about that is visible if you’re counting FLOPs; it’s obvious the moment you count bytes and allocations. Avoiding that repeated full-cache copy is one of the reasons preallocated and static caches matter — though paged cache systems are solving a broader set of problems than this one, including incremental allocation, fragmentation, and sharing cache blocks across common prefixes.
The GPU is an asynchronous device at the end of a queue, and much of its utilization depends on keeping enough independent work available while avoiding unnecessary host dependencies and memory traffic.
.item(), nonzero, boolean-mask indexing, branching on a tensor value — makes the host wait, and an operation doing less arithmetic can easily cost far more.nonzero, ask whether cumsum plus where gets you the same answer without the stall.If one sentence survives: your intuition prices operations by the arithmetic they perform; the machine prices them by the stalls and the bytes they cause, and the two orderings are frequently opposites.
That’s the tensor-fundamentals arc: layout, axis algebra, indexing, numerics, and now execution. The last post in the group is a set of problems where shape alone can’t establish correctness — some fail semantically, others numerically, and others through device, layout, or execution contracts.
torch.utils.benchmark — the timing harness you should use instead of a hand-rolled one.torch.compile tutorial and the graph-break FAQ — dynamic outputs, unbacked symbolic shapes, graph breaks, and where the compiler’s current limits actually are.