Remove deep layers during context encoding (prefill) while keeping them for token generation (decode). Identifies layer importance asymmetry via virtual gates; achieves 1.37x prefill speedup without retraining on any pre-trained model.
Remove deep layers during context encoding (prefill) while keeping them for token generation (decode). Identifies layer importance asymmetry via virtual gates; achieves 1.37x prefill speedup without retraining on any pre-trained model.
POP: Stage-Aware Layer Pruning for Inference
Standard inference removes entire layers, but different layers matter differently across inference stages. During prefill (context encoding), deep layers are largely redundant; during decode (token generation), they're critical. POP exploits this asymmetry by removing layers only during prefill, using virtual gates to identify which layers to prune without requiring model retraining.
The key insight is that layer importance changes dramatically between prefill and decode. This suggests a pragmatic pruning strategy that's stage-aware rather than uniform.
Core Concept
POP operates on the observation that:
Prefill stage: Context encoding where tokens attend over full sequence; deep layers contribute minimally
Decode stage: Single-token generation where all layers are needed for quality
By identifying prunable layers via a virtual gate mechanism, POP removes them during prefill only, maintaining full model capacity for decode.
Architecture Overview
Virtual Gate Estimator: Lightweight mechanism to estimate layer importance without retraining
Prunable Layer Identification: Determines which deep layers can be removed during prefill
Stage-Aware Execution: Disable prunable layers only in prefill; enable in decode
"""
Estimate importance of layer by approximating loss change if removed.
Uses: ΔL ≈ (1/2) g^T H^{-1} g where H is Hessian, g is gradient
Args:
hidden_states: Input to layer
layer_idx: Which layer to evaluate
loss_fn: Loss function for measurement
Returns:
Importance score [0, 1]
"""
self
# Compute gradient of loss w.r.t. layer output
True
# Fisher approximation: diag(H) ≈ E[g^2]
2
0
1
# Approximate importance: sum of |gradient| weighted by inverse Fisher
sum
abs
0
1
self
return
min
1.0
def
rank_layers_by_importance
self,
val_dataset: List[torch.Tensor],
loss_fn
List
Tuple
int
float
"""
Rank all layers by importance.
Returns:
List of (layer_idx, importance_score) tuples, sorted by importance
"""
for
in
range
len
self
for
in
self
sum
len
# Sort by importance (ascending)
sorted
lambda
1
return
Step 2: Identify Prunable Layers
Determine which layers can be safely removed during prefill.
# Prunable layer identificationclassPrunableLayerAnalyzer:
def__init__(self, model: nn.Module, pruning_ratio: float = 0.3):
"""
Identify which layers are prunable during prefill.
Args:
model: Language model
pruning_ratio: Target fraction of layers to prune
"""self.model = model
self.pruning_ratio = pruning_ratio
self.num_layers = len(model.layers)
self.target_prune_count = int(self.num_layers * pruning_ratio)
defidentify_prunable_layers(
self,
val_dataset: List[dict]
) -> List[int]:
"""
Determine which deep layers are least important.
Args:
val_dataset: Validation data for importance scoring
Returns:
List of layer indices to prune during prefill
"""
gate = VirtualGate(self.model)
defloss_fn(output):
# Simple loss: just magnitude as proxy for importancereturn torch.mean(output ** 2)
# Get hidden states from validation set
hidden_states_list = []
for sample in val_dataset:
with torch.no_grad():
hidden = self.model.embed_tokens(sample["input_ids"])
hidden_states_list.append(hidden)
# Rank layers
ranked_layers = gate.rank_layers_by_importance(
hidden_states_list,
loss_fn
)
# Select bottom N layers (least important)
prunable = [layer_idx for layer_idx, _ in ranked_layers[:self.target_prune_count]]
# Constraint: only consider deep layers (last 50%)
deep_layer_threshold = self.num_layers // 2
prunable = [l for l in prunable if l >= deep_layer_threshold]
returnsorted(prunable)
Step 3: Implement Stage-Aware Execution
Create inference wrapper that applies pruning only during prefill.
# Stage-aware model wrapperclassPrefillOnlyPrunedModel(nn.Module):
def__init__(self, model: nn.Module, prunable_layers: List[int]):
"""
Wrap model to prune layers only during prefill.
Args:
model: Base language model
prunable_layers: Layers to skip during prefill
"""super().__init__()
self.model = model
self.prunable_layers = set(prunable_layers)
self.is_prefill = Truedefforward(
self,
input_ids: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
past_key_values: Optional[List] = None,
use_cache: bool = False):
"""
Forward pass with stage-aware pruning.
Args:
input_ids: [batch, seq_len]
attention_mask: Optional mask
past_key_values: KV cache for decode
use_cache: Whether to return cache
Returns:
logits and cache if use_cache=True
"""# Determine stage# Prefill: seq_len > 1 or no past_key_values# Decode: seq_len == 1 and past_key_values present
is_prefill = (input_ids.shape[1] > 1) or (past_key_values isNone)
self.is_prefill = is_prefill
# Embedding
hidden_states = self.model.embed_tokens(input_ids)
new_cache = [] if use_cache elseNone# Process through layersfor layer_idx, layer inenumerate(self.model.layers):
# Skip prunable layers during prefill onlyif is_prefill and layer_idx inself.prunable_layers:
# Skip layer but maintain hidden state (residual path)if new_cache isnotNone:
new_cache.append(None)
continue# Apply layer normallyif past_key_values isnotNone:
past = past_key_values[layer_idx]
else:
past = None
hidden_states, cache_out = layer(
hidden_states,
attention_mask=attention_mask,
past_key_values=past,
use_cache=use_cache
)
if use_cache:
new_cache.append(cache_out)
# Output projection
logits = self.model.lm_head(hidden_states)
if use_cache:
return logits, new_cache
else:
return logits
defgenerate(self, input_ids: torch.Tensor,
max_new_tokens: int = 128) -> torch.Tensor:
"""Generate tokens using stage-aware pruning."""
generated = input_ids.clone()
cache = Nonefor _ inrange(max_new_tokens):
# Get logits for next tokenif cache isNone:
# Prefill: process all context
logits, cache = self.forward(
generated,
use_cache=True
)
next_logits = logits[:, -1, :]
else:
# Decode: process only last token
logits, cache = self.forward(
generated[:, -1:],
past_key_values=cache,
use_cache=True
)
next_logits = logits[:, 0, :]
# Sample next token
next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
generated = torch.cat([generated, next_token], dim=1)
return generated
Key results: 1.37× prefill speedup on Llama-3.1, Qwen3-VL, Gemma-3 without retraining. Works on any model. Particularly effective for multimodal inference.
