SigReg Kernel

Today I will try to build a Triton kernel for the SigReg regularization. I have a small knowledge in Triton but it's a good way to start learning (thx to Codex).

Baseline

Ok first, let's start by building the roofline model and see what we have to deal with. First we have 2 things to do:

• Evaluate flops per element
• Evaluate bytes per element

Here is a table for symbol references:

Symbol	Description
B	batch size
S	number of slices
D	feature dimension
P	number of integration points
q	number of bytes per scalar

Many of the operations are approximations. Like in matmul there is actually $N-1$ operations but it can be approximated as $N$.

This global estimation is very naive since PyTorch surely does some optimization under the hood, especially for reads and writes.

1. Projection Matrix Generation

First is the projection matrix generation, the one used to project each element into $N$ slices.

projection_matrix = self._get_synced_projection_matrix(feature_dim, device)

There is mainly 2 operations contributing:

projection_matrix = torch.randn(
    (feature_dim, self.num_slices),
    device=device,
    generator=self._generator,
    requires_grad=False,
)
projection_matrix = F.normalize(projection_matrix, p=2, dim=0)

Flops

• randn is pure rng so it is counted as 0 FLOPs.
• Normalization does square + sum over all elements $\rightarrow 2 \times D \times S$, then we have the square root over all reduced elements (one by column) $\rightarrow S$ operations and we have a division for each element $\rightarrow D \times S$.

So we have $FLOPs=S \times (1 + 3 \times D)$

Division is actually slower than add and multiply but in roofline model we count the number of arithmetic operations so a division is 1 FLOP.

Bytes

• randn does 1 write
• normalize does 1 read and 1 write

So we have $Bytes = 3 \times D \times S \times q$

2. Projection Matrix Multiplication

The next step is to multiply our input (usually an embedding) with the projection matrix.

# Project data: [B, D] @ [D, S] -> [B, S]
x_projected = x @ projection_matrix

Flops

• For one output element we have $D$ add and multiply. So each element computation produce $ 2 \times D$ FLOPs.
• The output size will be $B \times S$

So we have $FLOPs=2 \times D \times B \times S$

Bytes

• 2 read, $[B, D]$ and $[D, S]$
• 1 write $[B, S]$

So we have $Bytes = (B\times D + D \times S + B \times S) \times q$

3. Frequency Projection

For the Epps-Pulley test, we need to evaluate against multiple integration points / frequencies for evaluating the characteristic function.

args = x_projected.unsqueeze(-1) * self.integration_points

Flops

• unsqueeze is 0 FLOPs, basically just a view/metadata operation
• We then have an element wise multiply with broadcasting (broadcasting is 0 FLOPs). Each slice is multiplied by an integration point and that for each element in the batch $\rightarrow$ $B \times S \times P$

So we have $FLOPs=B \times S \times P$

Bytes

• 1 read for x_projected $[B, S]$ and 1 read for integration_points $[P]$
• 1 write for args $[B, S, P]$

So we have $Bytes = (B \times S \times P + B \times S + P) \times q$

4. Cosine and Sine

To compute the characteristic function we need cosine and sine values of our integration points.

cos_vals = torch.cos(args)
sin_vals = torch.sin(args)

This is the tricky part because cosine and sine are transcendental functions (they can't be written with basic operations add, multiply, substract and divide). In GPUs, these functions are handled with special-function units (SFU) and not classical fused multiply-add pipelines.
This mean that there is no "universal" flop counts for these functions. Their exact cost depends on the architecture and implementation, so I will not include them in the global FLOPs formula below.

Flops

• args is $[B, S, P]$ and we apply cos and sine on it.

I do not count these operations in the global FLOPs formula because they are transcendental operations and are better modeled separately as SFU work.

Bytes

• 1 read of args $[B, S, P]$
• 2 writes, one for cosine $[B, S, P]$ and one for sine $[B, S, P]$

So we have $Bytes = 3 \times B \times S \times P \times q$

For the diagnosis, it is better to say that this region is SFU bound instead of compute bound.

5. Mean Reduction

We average over the batch samples to estimate the empirical characteristic function of the distribution for each slice, which is then compared to the normal characteristic function in the Epps-Pulley test.

cos_mean = cos_vals.mean(dim=0)
sin_mean = sin_vals.mean(dim=0)

Flops

• mean is an addition for each batch elements and a division for added elements.

So we have $FLOPs=2\times B \times S \times P + 2\times S \times P$

Bytes

• 2 reads of $[B, S, P]$
• 2 writes of $[S, P]$

So we have $Bytes = (2\times B \times S \times P + 2 \times S \times P)\times q$

6. Squared Error

We need to know how far the sample distribution is from the normal distribution. Since there is no imaginary part in the normal distribution (symmetric around 0) the value is 0.

squared_error = (cos_mean - self.phi).square() + sin_mean.square()

Flops

• There is a substraction, 2 square and an addition so 4 FLOPs in total for each element.

So we have $FLOPs= 4\times S \times P$

Bytes

• 2 read of $[S, P]$ for cosine and sine
• 1 read of $P$ for $\Phi$
• 1 write of $[S, P]$

So we have $Bytes = (3 \times S \times P + P) \times q$

7. Weighted Integration

This step approximates the Epps-Pulley integral by taking a weighted sum over frequencies, producing one normality statistic per slice. The result is multiplied by the global batch size, matching the sample-size factor in the Epps-Pulley test.

stats = torch.matmul(squared_error, self.weights) * global_batch_size

Flops

• Matmul between $[S, P]$ and $P$, each element has $P$ add and multiply $\rightarrow$ $2 \times S \times P$ FLOPs
• The batch size scaling is applied to all elements $\rightarrow$ $S$ FLOPs

So we have $FLOPs= 2 \times S \times P + S$

Bytes

• 1 read of $[S, P]$ for squared errors
• 1 read of $P$ for weights
• 1 write of $S$

So we have $Bytes = (S \times P + P + S) \times q$

8. Final Mean

As a final step, we average the Epps-Pulley statistic across all random slices to produce a single scalar loss. Each slice gives one estimate of how far the projected distribution is from a standard normal, and the final mean aggregates those per-slice discrepancies into the training objective.

stats.mean()

Flops

• There is $S$ additions and 1 division

So we have $FLOPs = S + 1$

Bytes

• 1 read of $S$
• 1 write of 1 scalar

So we have $Bytes = (S + 1) \times q$

Bonus - DDP Synchronization

If we are using DDP, there is a synchronization step to average the empirical characteristic function estimate across all ranks.

cos_mean = self._all_reduce_mean(cos_mean)
sin_mean = self._all_reduce_mean(sin_mean)

with:

tensor = dist_fn.all_reduce(tensor, op=dist.ReduceOp.SUM)
return tensor / self.world_size

This part is not included in the global formula below because it is a communication cost and not only a compute cost. It is better to model it separately from the standard roofline.

Flops

• all_reduce is a communication primitive, so it is not counted as standard FLOPs.
• The division by world_size is applied to all elements for both cosine and sine tensors $\rightarrow 2 \times S \times P$

So we have $FLOPs=2 \times S \times P$

Bytes

• Each synchronized tensor is $[S, P]$, so the payload is $S \times P \times q$
• We synchronize both cosine and sine tensors $\rightarrow 2 \times S \times P \times q$

So we have $Bytes = 2 \times S \times P \times q$

If we want the actual communication traffic on the interconnect, it depends on the all-reduce algorithm. For a ring all-reduce with $W$ ranks, the traffic per rank is approximately $4 \times \frac{W-1}{W} \times S \times P \times q$.

9. Total

Let's now do a total sum of FLOPs and bytes moved:

$$FLOPs = S(D(2B + 3) + P(3B + 8) + 3) + 1$$ $$Bytes = (D(B + 4S) + (3P + 1)(2S(B + 1) + 1)) \times q$$

We can now compute the arithmetic intensity (AI) of our Sigreg function: $$AI = \frac{S(D(2B + 3) + P(3B + 8) + 3) + 1}{(D(B + 4S) + (3P + 1)(2S(B + 1) + 1)) \times q}$$

This AI formula excludes the cosine and sine operations from the FLOPs count. They still contribute to runtime, but since they are transcendental SFU operations their cost is better analyzed separately than modeled as a fixed FLOPs count.

This actually doesn't tell us much. Let's pick some standard values to simplify the formula. First we can select slices and integrations points from the paper:

• $S$ = 1024
• $P$ = 17

This gives (after a refactoring): $$AI = \frac{1024(D(2B + 3) + 51B + 139) + 1}{(D(B + 4096) + 52(2048B + 2049)) \times q}$$

Now let's do a sweep on some batch and dim values to see what the roofline model tell us. I will use B200 gpus with BF16 precision as a reference.

The result is very clear, we are memory bound (at least with the theoritical analysis). We can see that increasing batch size helps a lot and that the effect of the dimension is relative to the batch size (quite logic).

This mean that to actually improve our implementation we have to reduce the bytes moved. We can possibly try to fuse some operations and try to reduce writting intermediates results.

Benchmarking Regions

Now that we saw that our Sigreg implementation is theoritically memory-bound, let's actually benchmark it. We will split the sigreg function into multiple regions:

1. Projection Generation

This region isolates the cost of generating and normalizing the random projection matrix:

projection_matrix = torch.randn(...)
projection_matrix = F.normalize(projection_matrix, p=2, dim=0)

2. Projection GEMM

This region measures the matrix multiplication between the input embeddings and the projection matrix:

x_projected = x @ projection_matrix

3. Trig Reduction

This region groups together the frequency projection, cosine/sine evaluation and the batch reduction:

args = x_projected.unsqueeze(-1) * self.integration_points
cos_vals = torch.cos(args)
sin_vals = torch.sin(args)
cos_mean = cos_vals.mean(dim=0)
sin_mean = sin_vals.mean(dim=0)

4. Epilogue

This region isolates the final error computation and weighted integration:

squared_error = (cos_mean - self.phi).square() + sin_mean.square()
stats = torch.matmul(squared_error, self.weights) * global_batch_size
return stats.mean()

Benchmark Code

The full benchmark command is:

python3 -m benchmarks.benchmark_sigreg \
    --batch-size 32 \
    --feature-dim 2048 \
    --num-slices 1024 \
    --n-points 17 \
    --dtype bf16

Result

device=NVIDIA GeForce RTX 5070 Ti
shape=batch_size:32 feature_dim:2048 num_slices:1024 n_points:17 dtype:bf16 backend:torch

region                     time_ms  share_%
projection_generation       0.1382    43.28
projection_gemm             0.0407    12.73
trig_reduction              0.0476    14.92
epilogue                    0.0928    29.07

full_forward                0.3853

At the time of the benchmark I only have a 5070ti, I will do it in a B200 when I can.

Observation

The result is actually consistent with the roofline model.

The dominant regions are the projection generation and the epilogue:

• projection generation: 43.28%
• epilogue: 29.07%
• trig reduction: 14.92%
• projection GEMM: 12.73%

This makes sense because projection generation has low arithmetic intensity and requires writing then normalizing a full projection matrix. The epilogue is also a low-arithmetic-intensity region with little work per byte moved, so it naturally shows up as a major cost. The trig block is still relevant, but for this corrected batch-size configuration it is no longer one of the two dominant regions.

Roofline model per region

We should create a more fine grained roofline model with each region in it. These plots are not purely theoretical models anymore: they show the measured region points against the theoretical roofline. Let's use the parameters from our benchmark:

And if we use an approximation of 20 FLOPs per cos/sin we have:

For the measured torch baseline, projection generation and the epilogue look like the strongest runtime targets. This makes the epilogue a very reasonable first custom kernel target: it is both expensive enough to matter and isolated enough to implement and validate cleanly.

First Kernel

At first sight, a natural idea would be to fuse the trig reduction and the epilogue into one custom kernel (since GEMM is already well optimized and generation is mainly RNG bound). In practice, this is not possible in the current distributed formulation because there is a synchronization boundary between these two parts:

cos_mean = self._all_reduce_mean(cos_mean)
sin_mean = self._all_reduce_mean(sin_mean)

The epilogue depends on the global means across all GPUs, so we cannot fuse a local kernel through the all_reduce.

This means the realistic decomposition is:

1. a local kernel before communication
2. the DDP synchronization
3. a post-synchronization kernel for the epilogue

So the real candidates are:

• projection generation
• local trig reduction
• post-allreduce epilogue

At the moment, projection generation and the epilogue look like the strongest targets for the measured configuration. The epilogue remains a good first kernel because it is isolated, small, and easy to validate, while the trig reduction is still a valid later target but is more constrained by the synchronization boundary.

Kernel Interface

Ok, for the biggest first gains we will start by writing a kernel for the epilogue:

squared_error = (cos_mean - self.phi).square() + sin_mean.square()

# Weighted integration along the 'points' dimension
stats = torch.matmul(squared_error, self.weights) * global_batch_size

if self.clip_value is not None:
    stats = torch.clamp(stats, min=self.clip_value)

The global batch size is still part of the interface, but we pass it explicitly as a scalar input instead of recomputing it inside the kernel.

Before diving into the code, we should first start designing the interface.

For the first kernel, we will stop at stats and keep the final stats.mean() outside the kernel.

This is better because:

• each slice is independent at this stage
• we only reduce over the small P dimension inside the kernel
• we avoid adding a second reduction over S
• it is easier to validate against the PyTorch reference
• it still captures the most expensive part of the epilogue

Inputs

• cos_mean: $[S, P]$
• sin_mean: $[S, P]$
• phi: $[P]$
• weights: $[P]$
• global_batch_size: scalar
• clip_value: scalar
• do_clip: boolean

Output

• stats: $[S]$

Dtype

• inputs may be bf16 or fp32
• all arithmetic and reductions are performed in fp32
• output stats is fp32

Even if the inputs are stored in BF16, the epilogue should accumulate in FP32 because it performs subtractive error computation and multiple reductions before producing a scalar loss. Using FP32 reduces numerical noise in the regularizer and makes the first custom kernel much easier to validate against the PyTorch reference.

The equivalent python code would look like this:

def epilogue_stats_ref(
    cos_mean,
    sin_mean,
    phi,
    weights,
    global_batch_size,
    do_clip: bool,
    clip_value: float,
):
    cos_mean = cos_mean.float()
    sin_mean = sin_mean.float()
    phi = phi.float()
    weights = weights.float()

    squared_error = (cos_mean - phi).square() + sin_mean.square()
    stats = torch.matmul(squared_error, weights) * float(global_batch_size)

    if do_clip:
        stats = torch.clamp(stats, min=float(clip_value))

    return stats

Then the final loss is still computed outside the kernel:

loss = epilogue_stats_ref(...).mean()

Kernel Design

If we actually look at the code, we are mainly doing operations on slices (comparing it to the reference normal disitribution). So we can compute one stats[s] value for each slice s.

Since P = 17 is small, the natural strategy is to parallelize over slices and reduce over the P dimension inside the kernel.

Since each slice is independent, we can assign one execution group to one slice.

A simple first design is:

• one warp handles one slice
• each active lane handles one integration point p
• the warp reduces the partial results over P
• one lane writes the final stats[s]

This matches the computation well because P = 17, which fits comfortably inside one warp.

For one slice s, the kernel computes:

for p in range(P):
    delta = cos_mean[s, p] - phi[p]
    term = (delta * delta + sin_mean[s, p] * sin_mean[s, p]) * weights[p]

stats[s] = global_batch_size * sum_p(term)

if do_clip:
    stats[s] = max(stats[s], clip_value)

So the kernel only reduces over the P dimension and writes one scalar per slice.

I will not explicit the kernel code here, everything is in the kernel folder (you can simply ask Claude, ChatGPT or Gemini to explain it to you, the code is quite simple).

Benchmark the Kernel

Ok now let's benchmark the kernel compared to the pytorch implementation (done on 100k iterations):

device=NVIDIA GeForce RTX 5070 Ti
shape=num_slices:1024 n_points:17 global_batch_size:32 dtype:bf16 clip:True

correct=True max_abs_diff=0.00024414

impl            time_ms       GB/s    GFLOP/s    speedup
pytorch          0.1074      0.688      0.992      1.000
triton           0.0170      4.349      6.271      6.322

We can see that the Triton kernel gives a clear speedup on the isolated epilogue.

What about the sigreg benchmark ?

device=NVIDIA GeForce RTX 5070 Ti
shape=batch_size:32 feature_dim:2048 num_slices:1024 n_points:17 dtype:bf16 backend:triton

region                     time_ms  share_%
projection_generation       0.1615    56.91
projection_gemm             0.0388    13.68
trig_reduction              0.0598    21.09
epilogue                    0.0236     8.32

full_forward                0.3051

The full benchmark also improves in the expected direction:

• the epilogue region drops from 0.0928 ms to 0.0236 ms
• the epilogue share drops from 29.07% to 8.32%
• the full forward goes from 0.3853 ms to 0.3051 ms

But the end-to-end gain is smaller because the full SIGReg runtime is still dominated by projection generation, and the other regions are unchanged

Fuse trig reduction to attack the intermediate materialization

Now let's attack the trigonometry reduction kernel, this will help reduce intermediate materialization. Let's start by first reading the code:

# Epps-Pulley Statistical Testing
args = x_projected.unsqueeze(-1) * self.integration_points
cos_vals = torch.cos(args)
sin_vals = torch.sin(args)

# Average over batch dimension [batch, slices, points] -> [slices, points]
cos_mean = cos_vals.mean(dim=0)
sin_mean = sin_vals.mean(dim=0)

Interface

For the interface, the inputs are x_projected (result of the projection to slices) and integration_points (frequencies for the test). The outputs will be cos_mean and sin_mean.
For the dtype, x_projected may be stored in bf16 or fp32, but integration_points should stay in fp32. The trig argument computation, cosine/sine evaluation, and batch reduction are performed in fp32, and the outputs cos_mean and sin_mean are stored in fp32.

Here is the python equivalent:

def trig_sigreg_ref(
    x_projected,
    integration_points,
) -> tuple[torch.Tensor, torch.Tensor]:
    # Broadcast the integration points over the batch and slice dimensions.
    args = x_projected.float().unsqueeze(-1) * integration_points.float()

    # This is the exact local computation we want the future Triton kernel to
    # replace without materializing the full [B, S, P] intermediates.
    cos_mean = torch.cos(args).mean(dim=0)
    sin_mean = torch.sin(args).mean(dim=0)
    return cos_mean, sin_mean

Design

For the design, the output shape is $[S, P]$ so the natural first design is to have one warp per slice and one thread per integration point. Then the kernel will iterate over batch elements to compute the values:

for s in slices:
    for p in integration_points:
        acc_cos = 0
        acc_sin = 0
        for b in batch:
            arg = float(x_projected[b, s]) * float(t[p])
            acc_cos += cos(arg)
            acc_sin += sin(arg)
        cos_mean[s, p] = acc_cos / B
        sin_mean[s, p] = acc_sin / B

Benchmark

Ok now let's benchmark the kernel compared to the pytorch implementation (done on 100k iterations):

device=NVIDIA GeForce RTX 5070 Ti
shape=batch_size:32 num_slices:1024 n_points:17 dtype:bf16

correct_cos=True correct_sin=True max_abs_diff_cos=0.00000036 max_abs_diff_sin=0.00000010

impl            time_ms       GB/s    GFLOP/s    speedup
pytorch          0.0456      4.488     37.375      1.000
triton           0.0221      9.271     77.198      2.065

So the isolated trig kernel is faster, but the speedup is still smaller than for the epilogue kernel.

What about the full sigreg benchmark with both kernels?

device=NVIDIA GeForce RTX 5070 Ti
shape=batch_size:32 feature_dim:2048 num_slices:1024 n_points:17 dtype:bf16 backend:triton

region                     time_ms  share_%
projection_generation       0.1603    64.45
projection_gemm             0.0348    13.97
trig_reduction              0.0261    10.51
epilogue                    0.0275    11.07

full_forward                0.4290

Compared to the torch baseline:

• trig reduction drops from 0.0476 ms to 0.0261 ms
• epilogue drops from 0.0928 ms to 0.0275 ms
• but the full forward goes from 0.3853 ms to 0.4290 ms

So at this point the conclusion is more nuanced:

• both isolated kernels are correct
• both isolated kernels are faster than their PyTorch equivalents
• but replacing both regions is not yet enough to improve the full SIGReg forward in this benchmark

The main reason is that projection generation now dominates the runtime even more strongly, and the gain on the trig kernel is fairly modest compared to the epilogue.