Heuristic:Mit han lab Llm awq Kernel Selection Thresholds
| Knowledge Sources | |
|---|---|
| Domains | Optimization, Inference |
| Last Updated | 2026-02-15 01:00 GMT |
Overview
Runtime kernel selection heuristics for TinyChat inference: sequence length threshold (8192) for fused vs Flash Attention, batch size threshold (8) for GEMV vs GEMM, and split_k auto-tuning over powers of 2.
Description
TinyChat uses three adaptive kernel selection strategies during inference. First, the attention kernel switches between a fused FasterTransformer kernel (for sequences <= 8192) and Flash Attention (for longer sequences). Second, the quantized linear kernel dispatches to either a low-latency GEMV kernel (for small batches < 8 tokens) or a standard GEMM kernel (for larger batches). Third, the split_k parameter for quantized GEMM operations is auto-tuned at model load time by benchmarking all powers of 2 from 1 to 32 and selecting the fastest.
Usage
These heuristics are applied automatically during TinyChat model initialization and inference. Understanding them helps when debugging performance issues or when deploying on non-standard hardware. The split_k auto-tuning adds a one-time latency cost at model load but improves all subsequent inference throughput. The 8192 sequence threshold determines whether Flash Attention needs to be installed.
The Insight (Rule of Thumb)
- Attention Kernel (8192 threshold):
- Sequences <= 8192: Use fused FasterTransformer kernel (custom CUDA, lower latency for decode)
- Sequences > 8192: Use Flash Attention (handles arbitrary lengths, avoids OOM)
- Trade-off: Fused kernel is faster for short sequences but pre-allocates KV cache for full max_seq_len
- Linear Kernel (batch size < 8):
- Batch tokens < 8: Use specialized GEMV kernel (optimized for single/few-token decode)
- Batch tokens >= 8: Use GEMM kernel (better for batched/prefill workloads)
- Trade-off: GEMV has lower latency per token; GEMM has better throughput for batches
- Split_k Auto-tuning:
- Grid: `[1, 2, 4, 8, 16, 32]` (powers of 2)
- Default: `split_k_iters = 8` (before tuning)
- Measurement: 1000 iterations with 1000 warmup iterations
- Optimization: Tuned per unique `(in_features, out_features)` pair, then applied to all matching layers
- Trade-off: Higher split_k increases parallelism but adds synchronization overhead
- Device Warmup:
- 8192x8192 random matrix multiplication for 100 iterations before benchmarking
- Purpose: Ensures CUDA/cuDNN kernels are compiled and caches are warm
Reasoning
The 8192 sequence length threshold reflects the memory trade-off of the FasterTransformer fused kernel, which pre-allocates KV cache tensors of shape `(batch, heads, max_seq_len, head_dim)`. At 8192 tokens this is manageable (~256MB for a 7B model), but longer sequences would waste excessive memory. Flash Attention computes attention without materializing the full attention matrix, making it essential for long contexts. The batch size threshold of 8 separates the decode phase (typically 1 token at a time) from prefill (many tokens). GEMV kernels are memory-bandwidth-bound and optimal for small batches, while GEMM kernels are compute-bound and optimal for large batches. The split_k parameter controls how the K-dimension of matrix multiplication is partitioned across thread blocks; the optimal value depends on the specific matrix dimensions and GPU architecture, hence auto-tuning.
# From tinychat/modules/fused_attn.py:358-390
# Attention kernel selection based on sequence length
if self.kv_max_seq_len <= 8192:
self.forward = self.short_forward # FasterTransformer fused kernel
else:
# Flash Attention for long sequences to avoid OOM
...
# From awq/quantize/qmodule.py:206-212
# Batch size threshold for kernel dispatch
if inputs.numel() / inputs.shape[-1] < 8:
# small batch size, use inference kernel (GEMV)
...
else:
# large batch size, use normal inference (GEMM)
...
# From tinychat/utils/tune.py:58-64
# Split_k auto-tuning over powers of 2
for split_k_iters in [1, 2, 4, 8, 16, 32]:
module.split_k_iters = split_k_iters
cur_latency = _time_module(module, inputs, measure_iters)
if best_split_k_iter is None or best_latency >= cur_latency:
best_split_k_iter = split_k_iters
best_latency = cur_latency
# From tinychat/utils/tune.py:10-13
# Device warmup before benchmarking
def device_warmup(device: str):
warm_up = torch.randn((8192, 8192)).to(device)
for i in range(100):
torch.mm(warm_up, warm_up)