Heuristic:FMInference FlexLLMGen Sequence Length Alignment
| Knowledge Sources | |
|---|---|
| Domains | Optimization, LLM_Inference |
| Last Updated | 2026-02-09 12:00 GMT |
Overview
Pad tokenized input sequences to 256-token boundaries for more efficient GPU memory allocation and batch processing in FlexLLMGen.
Description
When batching prompts for inference, the HELM integration automatically adjusts the padding length to align to 256-token multiples. After tokenizing all prompts and finding the actual maximum sequence length, the system rounds up to the nearest multiple of 256 (capped by the original pad_to_seq_len). This alignment improves GPU memory allocator efficiency because CUDA memory allocators handle power-of-two or aligned sizes more efficiently. It also avoids excessive padding when the specified pad_to_seq_len is much larger than the actual longest sequence.
Usage
Use this heuristic when batching variable-length prompts for FlexLLMGen inference, especially in HELM benchmark evaluation or any scenario with heterogeneous prompt lengths. The technique reduces wasted memory from over-padding while keeping tensor sizes aligned for GPU efficiency.
The Insight (Rule of Thumb)
- Action: After tokenizing all prompts, compute the actual maximum sequence length. Round up to the nearest multiple of 256:
pad_len = ceil(max_len / 256) * 256. - Value: 256-token alignment boundary.
- Trade-off: Wastes at most 255 tokens of padding per sequence (minimal overhead) but gains consistent memory allocation patterns.
- Cap: Never exceed the originally specified
pad_to_seq_lento prevent exceeding the model's maximum context window.
Reasoning
GPU memory allocators (CUDA) use a binning strategy where allocations are rounded to discrete sizes. Tensors aligned to large power-of-two boundaries (like 256) fit neatly into bins, reducing fragmentation. When padding is too tight (e.g., padding to exactly 137 tokens), each batch may have a slightly different tensor size, causing repeated allocation/deallocation of non-standard sizes. The 256-token alignment ensures that batches with similar-length prompts get identical tensor shapes, enabling memory reuse.
Additionally, this auto-adjustment catches cases where the user specifies a very large pad_to_seq_len (e.g., 2048) but the actual data only needs 300 tokens. Without adjustment, 85% of each sequence would be wasted padding.
Code Evidence
Sequence length alignment from flexllmgen/apps/helm_run.py:131-148:
def get_batches(scenario_state, tokenizer, batch_size, pad_to_seq_len):
prompts = []
for r in scenario_state.request_states:
prompts.append(r.request.prompt)
# Tokenize
input_ids = tokenizer(prompts, padding="max_length",
return_tensors="np",
max_length=pad_to_seq_len).input_ids
assert len(input_ids.shape) == 2, f"Please use a longer pad_to_seq_len. current = {pad_to_seq_len}"
max_seq_len = max(np.sum(input_ids != tokenizer.pad_token_id, axis=1))
if max_seq_len != pad_to_seq_len:
pad_to_seq_len = min(int(math.ceil(max_seq_len / 256) * 256), pad_to_seq_len)
input_ids = tokenizer(prompts, padding="max_length",
return_tensors="np",
max_length=pad_to_seq_len).input_ids