Saturation Bitmasks for Sparse Inference

Overview

profile-saturation-bitmasks is the final profiler in the saturation measurement pipeline. Unlike the CSV report from profile-phi3-saturation, this tool generates:

1. saturation-maps.rs — Ready-to-embed Rust code for .rknot encoder
2. saturation-manifest.json — Go/no-go decision metadata

These outputs enable zero-cost runtime filtering in sparse FFN kernels.

Why Bitmasks?

Previously, saturation detection happened at inference time. Now we:

1. Profile once during model preparation
2. Generate bitmasks (1 bit per frozen neuron)
3. Embed in .rknot metadata
4. Lookup at runtime: if saturation_mask[neuron_id] { skip; }

Benefit: O(1) lookup without recomputation, ~1KB per layer overhead.

Usage

cargo run --release --bin profile-saturation-bitmasks -- \
  --knot /path/to/phi3-mini.knot \
  --num-prompts 100 \
  --output-dir ./saturation-output

Arguments

Outputs

1. saturation-maps.rs

Rust code ready for inclusion in the .rknot encoder pipeline.

Example snippet:

use distributed_inference::rknot::saturation_metadata::{
    SaturationMaps, LayerFrozenBitmask, LayerSaturationStat
};
use std::collections::HashMap;

pub fn phi3_saturation_maps() -> SaturationMaps {
    let mut maps = SaturationMaps::new(32, "variance_gradient_tracking".to_string())
        .with_model_name("phi3-mini".to_string());

    // Layer 0: 8.2% frozen (671/8192)
    {
        let mut mask = LayerFrozenBitmask::new(0, 8192);
        mask.frozen_mask[0] = 0x42; // neurons 0-7
        mask.frozen_mask[15] = 0xff; // neurons 120-127
        // ... more bytes ...
        maps = maps.insert_layer_mask(mask);
    }

    // Layer 1: 12.1% frozen (991/8192)
    {
        let mut mask = LayerFrozenBitmask::new(1, 8192);
        // ... initialization ...
        maps = maps.insert_layer_mask(mask);
    }

    // ... all 32 layers ...

    maps
}

Integration into encode-rknot:

// In encode-rknot pipeline:
let sat_maps = phi3_saturation_maps();
rknot_header.saturation_maps = Some(sat_maps);

2. saturation-manifest.json

Decision metadata for go/no-go approval.

Example structure:

{
  "model": "phi3-mini",
  "saturation_method": "variance + gradient tracking",
  "profiling_date": "2026-05-17T14:32:45.123456Z",
  "sampling_size": 100,
  "intermediate_size": 8192,
  "num_layers": 32,
  "var_threshold": 0.01,
  "grad_threshold": 0.001,
  "activation_threshold": 0.1,
  "layers": [
    {
      "layer_id": 0,
      "pct_frozen_ffn": "8.2",
      "estimated_speedup": "1.09x",
      "avg_variance_ffn": "0.123456",
      "saturated_pct": "7.5"
    },
    {
      "layer_id": 1,
      "pct_frozen_ffn": "12.1",
      "estimated_speedup": "1.14x",
      "avg_variance_ffn": "0.087654",
      "saturated_pct": "11.2"
    },
    // ... layers 2-31 ...
  ],
  "saturation_trend": {
    "early_layers_0_7_pct": "8.7",
    "mid_layers_8_15_pct": "14.2",
    "late_layers_16_31_pct": "22.3",
    "trend": "increases_with_depth"
  },
  "summary": {
    "total_estimated_speedup": "1.42x",
    "total_bitmask_bytes": 32896,
    "bytes_per_layer": "1028",
    "recommendation": "PROCEED: Apply selective pruning; estimated improvement warranted",
    "next_step": "Encode saturation_maps into .rknot header via encode-rknot pipeline"
  }
}

Profiling Methodology

For each prompt in the test set:

1. Forward pass through all FFN layers
2. Activate recording on gate/up projections
3. Compute SwiGLU output: silu(gate) * up
4. Track statistics per neuron:
   • Mean activation
   • Variance (activation patterns)
   • Min/max bounds
   • Fraction below threshold

After profiling (100+ diverse prompts):

5. Classify saturation per neuron:

   • Low variance (< 0.01): frozen activation pattern
   • Zero gradient (|mean| < 0.001): no effective signal
   • Low activity (< 0.1 mag > 80% of time): redundant

6. Generate bitmask: 1 bit = 1 frozen neuron

Go/No-Go Decision

PROCEED (avg_speedup > 1.05)

Recommendation: Apply selective pruning; estimated improvement warranted

• Implement sparse FFN kernel using bitmasks
• Encode saturation_maps into .rknot metadata
• Runtime filters skip frozen neurons

CONDITIONAL (1.02 < avg_speedup ≤ 1.05)

Recommendation: Marginal improvement; consider if inference is latency-critical

• Profile more diverse prompts for confidence
• Consider other optimizations (quantization, attention sparsity)
• May proceed if latency is top priority

SKIP (avg_speedup ≤ 1.02)

Recommendation: Insufficient saturation detected; skip sparse inference for this model

• Use dense model weights
• Apply other optimization techniques (KV-cache, token pruning, etc.)

Architecture

SaturationProfiler

Core state machine (in src/saturation_profiler.rs):

pub struct SaturationProfiler {
    pub layers: Vec<LayerSaturationProfile>,
    pub activation_threshold: f32,
}

impl SaturationProfiler {
    pub fn generate_saturation_bitmasks(
        &self,
        var_threshold: f32,
        grad_threshold: f32,
    ) -> crate::rknot::saturation_metadata::SaturationMaps
}

SaturationMaps (Canonical)

Defined in src/rknot/saturation_metadata.rs:

pub struct SaturationMaps {
    pub ffn_frozen_per_layer: HashMap<u32, LayerFrozenBitmask>,
    pub stats: Vec<LayerSaturationStat>,
    pub profile_date: String,
    pub profile_method: String,
    pub model_name: Option<String>,
}

pub struct LayerFrozenBitmask {
    pub layer_id: u32,
    pub frozen_mask: Vec<u8>,  // ceil(num_neurons / 8)
    pub num_neurons: u32,
}

Serializable with serde_json for inclusion in .rknot header.

Integration Points

1. Encoding phase (encode-rknot):

   let sat_maps = include!("saturation-maps.rs")::phi3_saturation_maps();
   rknot.layers[i].manifest.saturation_maps = Some(sat_maps);

2. Loading phase (rknot loader):

   if let Some(sat_maps) = &rknot_header.saturation_maps {
       let is_frozen = sat_maps.is_ffn_neuron_frozen(layer_id, neuron_id);
   }

3. Inference phase (Phi3Pipeline sparse kernel):

   for neuron_id in 0..intermediate_size {
       if saturation_maps.is_ffn_neuron_frozen(layer_id, neuron_id) {
           continue; // skip computation
       }
       // normal FFN kernel
   }

Thresholds

Default saturation criteria (tunable via code):

Tunable by modifying the profiler invocation:

let var_threshold = 0.005;  // stricter
let grad_threshold = 0.0005;
profiler.generate_saturation_bitmasks(var_threshold, grad_threshold)

Performance Characteristics

Profiling Cost

• Time: ~2-5s per 100 prompts (Phi-3-mini)
• Memory: ~1GB for 32 layers × 8192 neurons

Runtime Cost (Inference)

• Storage: ~1KB per layer (~32KB total for 32 layers)
• Lookup: O(1) bitwise operation (~10ns per neuron)
• Speedup: 1.1x – 2.0x typical (varies by layer depth)

Sparse Kernel Assumptions

• Assumes contiguous FFN integer indexing
• Bit i in layer L = neuron i frozen status
• Little-endian bit order within bytes
• Compatible with SwiGLU gate + up projection

Examples

Example 1: Phi-3-mini with high saturation

Profile outputs saturation increasing with depth:

Layer 0:  8.2% frozen → 1.09x speedup
Layer 16: 18.3% frozen → 1.22x speedup
Layer 31: 28.5% frozen → 1.40x speedup

Average: 1.42x speedup
Recommendation: PROCEED

Action: Encode bitmasks into .rknot, implement sparse FFN kernel.

Example 2: Uniform saturation (low)

Profile outputs stable, low saturation:

All layers: 2-3% frozen → ~1.02x speedup

Average: 1.02x speedup
Recommendation: SKIP

Action: Use dense inference, apply other optimizations (KV-cache, attention sparsity).

Example 3: Bi-modal distribution (early vs late)

Profile outputs two clusters:

Early (0-7):   5% frozen → 1.05x speedup
Mid (8-15):   12% frozen → 1.14x speedup
Late (16-31):  25% frozen → 1.34x speedup

Average: 1.18x speedup
Recommendation: PROCEED (conditional on latency criticality)

Action: Consider selective sparse FFN for late layers only (lower overhead).

Troubleshooting

Bitmasks are all zeros

• Cause: Thresholds too strict
• Fix: Lower var_threshold or grad_threshold
• Check: Compare with CSV report from profile-phi3-saturation

Manifest JSON is malformed

• Cause: Serde serialization failure
• Fix: Ensure all Layer stats fields are properly computed
• Debug: Check stderr for serde_json errors

Generated Rust code doesn't compile

• Cause: Missing imports in saturation-maps.rs
• Fix: Verify use distributed_inference::rknot::saturation_metadata::*
• Check: Ensure BitVec initialization is syntactically correct

References

• RKNOT Format: src/rknot/format.rs (canonical Lean mirror)
• Saturation Metadata: src/rknot/saturation_metadata.rs
• SaturationProfiler: src/saturation_profiler.rs
• Sparse Inference Design: docs/SPARSE_FFN_KNOT.md (pending)

Next Steps

1. Run profiler on your model:

   cargo run --release --bin profile-saturation-bitmasks -- \
     --knot model.knot --num-prompts 100 --output-dir output/

2. Review saturation-manifest.json:

   • Check recommendation
   • Verify speedup vs latency tradeoff

3. If proceeding:

   • Include saturation-maps.rs in encode-rknot
   • Implement sparse FFN kernel using bitmasks
   • Benchmark real inference speedup

4. If skipping:

   • Document decision in model metadata
   • Apply alternative optimizations
   • Revisit with different prompt corpus if needed