Making WAF ML models go brrr: saving decades of processing time

Shaving 82% off WAF ML inference time

Cloudflare’s WAF Attack Score is a machine learning layer on top of the Web Application Firewall, designed to catch attack bypasses and unknown threats that rule-based detections miss. It has been used to spot zero-day vulnerabilities before public disclosure, and it now runs real-time inference on tens of millions of requests per second across millions of Internet properties. That scale creates constant pressure to make the inference path faster.

Through a series of targeted optimizations, we reduced average execution time from 1519 microseconds to 275 microseconds — a 5.5x speedup and roughly 82% less processing time. The changes touched both feature extraction and model inference, and they are now live in production.

To understand where the time went, it helps to look at the pipeline. The WAF Attack Score system processes an HTTP request in five stages:

1. HTTP Request Content — raw request input is the starting point.
2. Normalization & Transformation — the data is standardized and cleaned, with content substitutions and de-duplication applied.
3. Feature Extraction — the transformed content is tokenized to produce both statistical and structural signals.
4. Machine Learning Model Inference — pre-trained models map those features to classifications such as XSS, SQLi, or RCE, or to a score.
5. Classification Output in WAF — a score from 1 (likely malicious) to 99 (likely clean) is assigned to the request, informing the WAF’s action.

Feature extraction and inference were the two biggest contributors to latency, so those are the areas we focused on.

Feature extraction: cutting redundant work

Profiling showed that much of the execution time was spent in tokenization and feature generation. The original implementation performed a series of transformations and extractions sequentially, recomputing values that could be cached or derived more cheaply. By restructuring the order of operations and eliminating duplicated passes over the same content, we removed a substantial portion of the overhead without changing the features themselves.

One key win came from reducing the number of times the content was copied and re-parsed. The previous path created multiple intermediate representations of the normalized request body. The optimized path reuses buffers and avoids unnecessary allocation, which particularly helps when an individual request contains long payloads or many parameters.

Model inference: faster scoring

Once features are extracted, the model inference step maps them to an attack classification. A large share of that processing time comes from scoring every extracted feature set against the full model ensemble. The optimization here focused on streamlining the scoring loop: the model now consumes features in a tightly packed format that matches the layout the scoring engine expects, removing the need for on-the-fly conversion and indexing lookups during inference.

We also evaluated the balance between the number of features passed to the model and the marginal gain in classification accuracy. By trimming the feature set to those that carry the most signal, we reduced the mathematical work per request while keeping attack scores within the same tolerance as before. Model accuracy was validated against a held-out set of malicious and benign traffic to confirm no measurable regression.

The latency measurements below show the impact of these combined changes across the full WAF Attack Score inference path. The figures were captured on production traffic samples before and after the rollout.

The result of 275 microseconds per request opens headroom for customers who run the WAF attack score on high-traffic endpoints, and it makes it practical to apply ML scoring to more of the request stream without proportionally increasing CPU cost. The same optimizations apply regardless of whether the final output is a categorical attack label or a numeric score, so the gains are uniform across all WAF ML use cases.

Next steps

The implementation details of each optimization are documented in the sections that follow. The work discussed covers the feature extraction pipeline and the inference path, which together delivered the bulk of the latency reduction.

Feature extraction: from hash maps to branchless lookups

For the WAF Attack Score ML model, feature extraction turns an input string into a float tensor of shape 1 x m by tokenizing the input—specifically, by sliding a 3-byte window over the bytes and mapping each ngram to its tensor index via Rust's std::collections::HashMap. Benchmarking four input sizes (44 to 9,482 bytes, covering typical request bodies, user agents, and URIs) with Criterion.rs on a 13th Gen Intel Core i7-13800H laptop showed throughput of 28–36 MiB/s, with processing time scaling roughly linearly with input length. Flamegraphs from a 100-second pprof session confirmed that the overwhelming majority of time went into two operations: HashMap lookups inside the tensor_populate_ngrams function and string replacement via Regex::replace_all.

The initial implementation performed one HashMap lookup per 3-byte window. For the avg-body-1000 benchmark, that meant 1,469 lookups per input. HashMap lookups are fast in isolation, but at this volume they become cache-unfriendly.

Optimization attempt #1: Aho-Corasick

The first attempt replaced HashMap lookups with the Aho-Corasick library, which builds a finite state machine for linear-time multi-pattern search. Given the recommendation to prefer the DFA variant when search speed is critical, the searcher was configured with AhoCorasickBuilder::kind(AhoCorasickKind::DFA), then queried via find_overlapping_iter:

static ref NORM_VOCAB_AC: AhoCorasick = AhoCorasick::builder().kind(Some(AhoCorasickKind::DFA)).build(&[    
    "abc",
    "def",
    "wuq",
    "ijf",
    "iru",
    "piw",
    "mjw",
    "isn",
    "od ",
    "pro",
    ...
]).unwrap();

That change produced a substantial latency drop across all benchmark cases compared to the HashMap baseline. The DFA's predictable memory access pattern proved far more cache-friendly than hash probing.

Optimization attempt #2: match statements

The next step was to discard the Aho-Corasick machinery entirely and let the compiler generate lookup code from a Rust match statement on the ngram bytes. As the Godbolt compiler explorer output shows, the compiler turns this into a jump table with byte-wise comparisons—efficient, with minimal branching. This approach shaved another 7–18% off latency, depending on the input case.

Optimization attempt #3: replacing regex with a single-pass iterator

The remaining bottleneck was Regex::replace_all, used to collapse every sequence of lowercase letters in the input into a single "#" byte before ngram extraction. That regex was doing far more work than the task required:

1. Replace each lowercase run with "#".
2. Slide a 3-byte window over the replaced bytes to form ngrams.
3. Look up the ngram index and increment the corresponding tensor slot.

The operation is simple enough to implement with a single pass and no intermediate allocations. A custom WindowedReplacer iterator was introduced, leveraging Rust's iterator optimization and avoiding the allocation of a replaced string buffer. Benchmark results showed this doubled pre-processing speed relative to the match-based version and made the whole pipeline four to five times faster than the original HashMap + regex implementation.

Optimization attempt #4: branchless ngram lookups

The match-based lookup still contained many cmp instructions, creating branch points that can trigger CPU branch mispredictions. With thousands of unique ngrams, eliminating these branches is non-trivial. The solution exploits the fact that each ngram is exactly 3 bytes.

A direct lookup table of size 256 × 256 × 256 would require 64 MB (storing a u16 tensor index), which is impractical. Restricting to ASCII bytes (0..127) would cut that to 8 MB but reintroduces a branch for high-bit bytes. The workaround: since only a small subset of byte values actually appears in ngrams, precompute offset tables for those unique byte values rather than storing the full Cartesian product. With N unique ngram bytes, each byte's position in a 3-byte ngram can be encoded as an offset into an N × N × N space:

const NGRAM_OFFSETS: [[u32; 256]; 3] = [
    [
        // offsets of first byte in ngram
    ],
    [
        // offsets of second byte in ngram
    ],
    [
        // offsets of third byte in ngram
    ],
];

The ngram index is then computed with a simple const function, and the tensor index is looked up from a compile-time-generated const array. A final get_unchecked_mut avoids runtime boundary checks. This version is entirely branchless, and the lookup arrays total roughly 500 KiB, comfortably fitting in modern CPU L2/L3 caches.

#[inline]
const fn ngram_index(ngram: [u8; 3]) -> usize {
    (NGRAM_OFFSETS[0][ngram[0] as usize]
        + NGRAM_OFFSETS[1][ngram[1] as usize]
        + NGRAM_OFFSETS[2][ngram[2] as usize]) as usize
}

One last tweak: processing ngrams six at a time (8 input bytes per iteration) lets the compiler unroll the loop and auto-vectorize the second loop, taking advantage of parallel execution units.

const CHUNK_SIZE: usize = 6;

let chunks_max_offset =
    ((input.len().saturating_sub(2)) / CHUNK_SIZE) * CHUNK_SIZE;
for i in (0..chunks_max_offset).step_by(CHUNK_SIZE) {
    for ngram in input[i..i + CHUNK_SIZE + 2].windows(3) {
        update_tensor_with_ngram(tensor, ngram.try_into().unwrap());
    }
}

The combined result is a six- to twelve-fold latency improvement over the baseline, with larger gains on longer inputs where branch mispredictions and cache misses would otherwise dominate. The final implementation retains the single-pass non-allocating replacement iterator and the branchless ngram lookup with offset tables.

Further gains might be available via manual SIMD intrinsics or cache pre-fetching, but the extraction pipeline is no longer the dominant cost. Attention now shifts to inference latency.

Speeding up model inference

The WAF Attack Score model runs on TensorFlow Lite 2.6.0. Because inputs are transformed into fixed-size tensors during pre-processing, inference time is independent of the original request length, which makes profiling straightforward.

Profiling showed that the bulk of inference time went into matrix multiplication, implemented as three nested loops in the reference TensorFlow Lite kernel. That implementation leaves a lot of performance on the table; standard optimizations—blocking for cache reuse, SIMD vectorization, and loop unrolling—can dramatically speed it up.

void PortableMatrixBatchVectorMultiplyAccumulate(const float* matrix,
                                                 int m_rows, int m_cols,
                                                 const float* vector,
                                                 int n_batch, float* result) {
  float* result_in_batch = result;
  for (int b = 0; b < n_batch; b++) {
    const float* matrix_ptr = matrix;
    for (int r = 0; r < m_rows; r++) {
      float dot_prod = 0.0f;
      const float* vector_in_batch = vector + b * m_cols;
      for (int c = 0; c < m_cols; c++) {
        dot_prod += *matrix_ptr++ * *vector_in_batch++;
      }
      *result_in_batch += dot_prod;
     ++result_in_batch;
    }
  }
}

Enabling SIMD

TensorFlow Lite can use SIMD instructions for matrix multiplication, but the feature has to be compiled in explicitly.

if [[ "$(uname -m)" == x86_64* ]]; then
    # On x86_64 target x86-64-v3 CPU to enable AVX2 and FMA.
    arguments+=("--copt=-march=x86-64-v3")
fi

Rebuilding the library with AVX2 support changed the matrix multiplication path to use 8x8 blocks with multiply-accumulate instructions. The profiler output shows the same operations dominating, but absolute times dropped sharply.

That is a substantial gain from a small build-config change. The next step was to look at the runtime itself.

Switching to XNNPACK

TensorFlow Lite ships a benchmark_model tool that includes a built-in profiler. Local builds can be produced with:

bazel build -j 4 --copt=-march=native -c opt tensorflow/lite/tools/benchmark:benchmark_model

Benchmarking different configurations showed that enabling XNNPACK delivered roughly a 50% latency reduction over the original TensorFlow Lite implementation. Upgrading from TensorFlow Lite 2.6.0 to 2.16.1 and enabling both SIMD and XNNPACK cut inference time by more than four-fold, a 77.17% reduction.

Caching inference results

Faster code is good, but code that never runs is better. Amdahl's Law caps the gains from optimizing any single stage, so avoiding redundant work matters. A naive key-value cache would exhaust server memory quickly given the cardinality of URLs, headers, and bodies. However, real-world request data follows a Zipfian distribution: a small set of inputs accounts for most traffic, while the long tail is sparse.

That makes a Least Recently Used (LRU) cache the right tool. Hot inputs stay resident, and cold ones get evicted. Cloudflare uses lua-resty-mlcache to share cached inference results across Nginx workers through a shared memory dictionary. This trades a modest amount of memory for significant CPU savings and achieves roughly a ~70% cache hit ratio.

End-to-end results

The optimizations were rolled out in stages to preserve correctness and stability. First came SIMD-enabled TensorFlow Lite, which cut average WAF ML execution time from 1519 to 884 μs, about 41.80%.

The next phase combined the TensorFlow Lite 2.16.1 upgrade, XNNPACK, and pre-processing optimizations. Average execution time dropped from 932 to 552 μs, roughly 40.77%. (The 932 μs baseline was slightly above the earlier 884 μs figure because more customers were using the feature by then.)

Finally, LRU caching shaved another ~50.18%, bringing average execution time from 552 to 275 μs.

Overall, WAF ML execution time fell by ~81.90%, from 1519 to 275 μs—5.5x faster. At Cloudflare's average of 9.5 million requests per second through WAF ML, saving 1244 microseconds per request adds up to roughly 32 years of processing time saved every day. That is on top of the 65 years per day already saved by the optimizations described in the earlier Bot Management post.