DLRM L2-Pinning

0. TL;DR — Why this paper is in our study

This page is part of the measurement study “Worth the Bytes?”, which argues that transient inter-kernel L2 reuse of activations in eager-mode LLM inference is worth less than 0.65% of step time. DLRM L2-pinning (L2P) is the most direct counterpart to that mechanism we could find in published systems work.

1. Overview — the EmbeddingBag bottleneck

Deep-learning recommendation models (DLRM) interleave dense MLPs with very large embedding tables. At inference the dominant cost is the EmbeddingBag gather-reduce, a stream of irregular HBM reads that is latency-bound on the memory system.

1.1 Why now / the bad scenario

Embedding tables keep growing under tight SLAs. A single EmbeddingBag with pooling factor 150 over twelve 500K×128 fp32 tables issues hundreds of thousands of scattered row reads per batch, most of which miss every cache and go to HBM.

1.2 Baseline pain & the strawman

The working set far exceeds the A100's 40 MB L2, and shared memory is tiny. The repo even ships shared-memory and L1-prefetch variants that the paper does not adopt — evidence the SRAM route was tried and found wanting.

1.3 The obstacle & positioning

Popularity is skewed, so a small hot slice dominates lookups, but the baseline kernel cannot keep it resident across launches. The paper's three software-only techniques raise occupancy, hide latency in-kernel, and keep the hot slice warm in L2 across kernels — the last being our mechanism.

1.4 Contributions

1. OptMT: occupancy/register-budget tuning via -maxrregcount 42.
2. RPF: in-kernel register prefetch into lv0/lv1 — intra-kernel latency hiding, no cross-kernel state.
3. L2P: a separate prefetch_pin_L2 kernel issuing prefetch.global.L2::evict_last on hot rows. This is the cross-kernel L2-residency mechanism our study measures.

2. The L2-Pinning mechanism (annotated source)

L2P uses two kernels separated by a stream sync, inside one ATen op. A producer pins hot rows in L2; a consumer EmbeddingBag kernel reads them. The producer computes nothing — it only issues prefetch hints that change L2 replacement state.

2.1 The pin kernel (the producer)

From opt_designs/l2_pinning/EmbeddingBag.cu, lines 131–148. The evict_last hint biases the replacement policy toward later eviction (advisory, not guaranteed):

template <typename scalar_t, typename index_t>
__global__ void prefetch_pin_L2(const scalar_t *weight, int64_t featureSize,
                                int64_t weight_stride0,
                                int64_t const* d_hot_idx_vec, size_t pinning_size) {
    int Xidx = blockIdx.x * blockDim.x + threadIdx.x;
    int Yidx = threadIdx.y;
    int stride = gridDim.x * blockDim.x;
    if (Yidx < featureSize) {
        for (int i = Xidx; i < pinning_size; i += stride) {
            int64_t const lv_hot_idx = d_hot_idx_vec[i];        // a hot row id
            const scalar_t *weightFeat = weight + Yidx;
            // insert the L2 prefetch inline PTX code
            asm volatile (
                "prefetch.global.L2::evict_last [%0];" ::"l"(&weightFeat[lv_hot_idx * weight_stride0])
            :);
        }
    }
}

The hot-index list is read at runtime from DATASET_FILE_PATH (lines 447–471); each reuse profile points at its own 60,000-line hot-index file.

2.2 The pin→sync→consume sequence (and the timer)

The launch sequence, lines 494–517. The timer position is the single most important measurement detail on this page:

      } else {
        // cudaEventRecord(start, stream);              <-- commented out
        // reduce_L2_size<<<grid, 256, 0, stream>>>(d_dummy_vec, size_dummy);
        dim3 block_2 = dim3(2, 128);
        prefetch_pin_L2<scalar_t, index_t><<<grid, block_2, 0, stream>>>(   // PRODUCER: pin hot rows
            weight.const_data_ptr<scalar_t>(), featureSize, weight.stride(0),
            d_hot_idx_vec, pinning_size);
        cudaStreamSynchronize(stream);                // wait for the pin to finish
        // cudaEventRecord(stop, stream);               <-- commented out
        // cudaEventSynchronize(stop);                  <-- commented out

        cudaEventRecord(start, stream);              // <== TIMER STARTS HERE, AFTER THE PIN

        EmbeddingBag_updateOutputKernel_sum_mean<scalar_t, index_t><<<grid, block, 0, stream>>>(
            /* CONSUMER: reads the now-warm rows from L2 */ ... );

        cudaStreamSynchronize(stream);
        cudaEventRecord(stop, stream);
        cudaEventSynchronize(stop);
        cudaEventElapsedTime(&time, start, stop);   // time = CONSUMER ONLY
      }

2.3 Why evict_last, and not a residency-control API

Ampere exposes two ways to bias L2 retention: the cudaStreamAttrValue access-policy-window (APW) API, or the per-instruction hint prefetch.global.L2::evict_last (warms a line and biases it toward later eviction). The core L2 trick is the per-instruction hint — finer-grained for a scattered hot-index list — but the measured wrapper also re-establishes an L2 access-policy window per batch (see §5) on top of re-reading the 60K hot-index file, so the artifact as run is PTX prefetch PLUS per-batch APW setup. The hint is advisory: demand traffic between pin and consume can evict hot rows before reuse, bounding the realized benefit.

2.4 Where the bytes live: the memory hierarchy

L2P targets the middle level — the device-wide L2 — so a later, independent kernel launch hits it:

3. Repo analysis (artifact)

Each technique is a hand-written replacement of aten/src/ATen/native/cuda/EmbeddingBag.cu, recompiled into a from-source PyTorch v2.1.0. The optimization lives entirely in the CUDA kernel and the build flags.

3.1 Key kernels & symbols

3.2 The four reuse profiles (confirmed counts)

The hotness sweep is what makes this artifact a stress test for our claim:

3.3 The measurement harness

The Python side is thin: save_emb.py regenerates test_500K_12.pt (~3 GB), load_emb.py builds the EmbeddingBag and replays the trace, and the swapped CUDA kernel prints its own cudaEvent latency. parser_Kernel_ET.py averages over 12×12 launches and drops the first (cold-miss) sample — so the reported number already assumes a warm-ish cache, itself favorable to the reuse story.

4. The ablation ladder — how big is the L2 slice?

The page's whole point is to isolate the L2-residency contribution. The ladder builds five variants across four profiles; differencing adjacent rungs attributes each technique's marginal effect to a hardware cause. Only the non-pinning rungs (1–3) give a clean wall-clock difference; the L2P rungs (4–5) are artifact-dominated (§5), so the L2P attribution below is read NOT from §4.1 latency but from the Nsight L2 comparison (§7) plus the paper's ablation:

4.1 Measured latency ladder (PACE A100, job 9942760)

The clean rungs (1–3) show the validating shape: a strong reuse gradient (baseline hot 3.616 ms vs random 5.875 ms, ~1.6×, consistent with cache/hotness effects — the L2-specific share is attributed by §7, not this wall-clock), OptMT alone essentially flat (3.616 → 3.598 ms hot), and RPF+OptMT helping the cache-UNfriendly profiles (low −11%, random −12%) by hiding the DRAM latency L2 cannot when there is no reuse.

5. Why no pin-charged wall-clock is reported

Recall the confirmed fact from §2.2: the timer excludes the pin kernel. The repo's reported L2P speedup is therefore

$$\text{Speedup}_{\text{repo}} = \frac{T_{\text{baseline}}}{T_{\text{consumer}}^{\text{L2P}}}$$

which charges nothing for producing the warm L2 state. The pin kernel is not free: it touches all 60,000 hot rows in HBM, incurring real traffic and launch+sync overhead. In load_emb.py's harness that setup cost is re-paid every batch and cannot be cleanly separated into a one-time pin time — so rather than report a per-batch setup artifact as a net pin-charged wall-clock, we read the L2P contribution from the Nsight L2-layer attribution (§7) plus the paper's own ~13.5% ablation figure.

6. Our reproduction on PACE (measured)

6.1 What we found (hypothesis & measured outcome)

We hypothesized the genuine cross-kernel L2-residency contribution is small and hotness-dependent; the measurements confirm both. (1) The L2 payoff is entirely hotness-dependent: ~92% hit / ~7.8 MB DRAM on reuse_high vs ~30% hit / ~145 MB DRAM (≈18× HBM) on random (FACT, §7). (2) L2P adds essentially nothing over RPF at the L2 layer in this no-polluter microbench (91.7% vs 91.8% hit) — consistent with the paper's ~13.5%-incremental figure (BOUND) and a parallel to our ‘dirty blob survives intact at zero pressure’ result.

6.2 Method — what we build and run

Single A100-80GB (sm_80), single-GPU, no TP. We keep the artifact's arch scope (TORCH_CUDA_ARCH_LIST="8.0", -maxrregcount 42); the repo has no H100 flags, so A100 is the only bit-faithful target. Each of the five variants is a separate from-source compile.

6.3 Confirmed static facts (verifiable now)

These are read directly from the artifact and need no GPU run:

• L2 mechanism = prefetch.global.L2::evict_last inline PTX on hot rows (l2_pinning/EmbeddingBag.cu:144).
• Timer measures the consumer ONLY; pin+sync precede cudaEventRecord(start) (lines 497–502).
• Pin-timing wiring is commented out (lines 494, 499–500) — they could have charged it and did not.
• RPF stages in registers (lv0/lv1) — the RPF rung is provably not our mechanism.
• Each intended profile has a 60,000-line hot-index file (random's is 60,000 random ids). Note: our reuse_low L2P run skipped pinning due to a leading-space path bug in l2_quick.sh (see §4.1 †).
• Build risk (confirmed): L2P/combined .cu use std::cout/setprecision/string/vector/getenv without explicit #include <iostream>/<iomanip>/<string>/<vector> — relying on PyTorch transitive includes.

6.4 How to reproduce (commands)

After patching the hardcoded paths and regenerating the model:

# 0. regenerate the gitignored model (12 x 500K x 128 fp32, ~3 GB)
cd scripts && python save_emb.py

# 1. baseline: build with NO maxrregcount, then sweep the 4 profiles
cp opt_designs/baseline/EmbeddingBag_timed.cu \
   pytorch/aten/src/ATen/native/cuda/EmbeddingBag.cu
TORCH_CUDA_ARCH_LIST="8.0" python setup.py install
bash quick.sh > dummy; grep 'Latency of Embedding Bag kernel' dummy \
  | cut -d' ' -f6 | cut -d'm' -f1 > log; python parser_Kernel_ET.py

# 2-3. OptMT / RPF: rebuild with CMAKE_CUDA_FLAGS="-maxrregcount 42", rerun quick.sh

# 4. L2P: rebuild with l2_pinning .cu, run l2_quick.sh (it sets DATASET_FILE_PATH)
cp opt_designs/l2_pinning/EmbeddingBag_timed.cu \
   pytorch/aten/src/ATen/native/cuda/EmbeddingBag.cu
CMAKE_CUDA_FLAGS="-maxrregcount 42" python setup.py install
bash l2_quick.sh > dummy; ...   # 60K hot rows pinned per profile

# 5. combined (RPF+L2P+OptMT): rebuild with combined .cu, run l2_quick.sh

7. Nsight attribution — isolating the L2 signal

Layer A (latency diffs) tells us how big each rung is; Layer B (Nsight Compute) tells us why. The validating test is to profile RPF+OptMT vs. L2P+OptMT on the hot and the random profile — comparing register-prefetch (RPF+OptMT) against explicit pinning (L2P+OptMT) as two alternative consumer-side strategies, i.e. whether explicit L2 pinning beats register-prefetch at the L2 layer (baseline is NOT the Nsight comparand) — and watch two metrics: the L2 sector hit rate on the consumer kernel, and the HBM bytes read. If explicit L2P delivers more cross-kernel L2 reuse than register-prefetch, the hit rate must be higher and HBM reads lower on hot; and the hotness sweep itself (hot vs random) shows whether the L2 payoff exists at all.

7.1 The metrics that separate L2 from everything else

7.2 The L2-attribution result (measured, job 9935235)

This contrast — ~92% steady-state L2 hit on reuse_high collapsing to ~30% with ≈18× HBM traffic on random, and L2P ≈ RPF at the L2 layer — validates the thesis: the inter-kernel L2 benefit is a property of the access pattern, not the kernel pair, and explicit pinning adds nothing over register-prefetch without cross-kernel pollution.

7b. Claims vs. implementation (honest reading)

Because the techniques live entirely in one swappable CUDA file, we can check each paper claim against the source. The status column reflects what the code does, cross-checked against our PACE A100 measurements where relevant.

7c. Positioning — where L2 pinning sits among reuse techniques

Cross-kernel reuse can be attempted at every level of the GPU memory hierarchy. Placing L2P among the alternatives clarifies why it is the right object of study for an inter-kernel L2 claim.

8. Tie-back to “Worth the Bytes?”

Our study expresses the L2-residency contribution as a share of step time. For DLRM L2P we read that contribution from the measured L2 layer rather than from a wall-clock pin time.

Because the L2P wall-clock is a pinning-setup artifact (§5), we do NOT compute a single pin-charged L2 share; we read the L2-layer contribution from the Nsight attribution (§7) and the paper's ~13.5% figure. The argument is corroborating, not a proof: L2P is among the most favorable cases imaginable, yet even here explicit pinning adds essentially nothing over RPF at the L2 layer (91.7% vs 91.8%) and the L2 payoff collapses with hotness (~92% → ~30% hit, ≈18× HBM on random). For naturally-transient eager-LLM activations this is consistent with a small, conditional contributor — supporting the plausibility of the <0.65% bound.

8.1 What this page proves now vs. later

8.2 Bottom line