vLLM e PagedAttention: serving high-throughput

# Pseudocódigo de PagedAttention CUDA kernel (simplificado)
# Cada thread block processa attention para 1 sequência, varrendo seus blocos físicos
def paged_attention_kernel(
    Q,                # [num_seqs, num_heads, head_dim]
    K_cache,          # [num_blocks, num_heads, head_dim, block_size]
    V_cache,          # idem
    block_tables,     # [num_seqs, max_num_blocks_per_seq] → physical block id
    seq_lens,         # [num_seqs]
):
    seq_idx = block_idx_x
    head_idx = block_idx_y

    q = Q[seq_idx, head_idx]                         # carrega Q em registradores
    block_table = block_tables[seq_idx]
    seq_len = seq_lens[seq_idx]

    # online softmax (FlashAttention-style)
    m = -inf; l = 0.0; o = 0.0

    for logical_block_idx in range(ceil(seq_len / block_size)):
        physical_block_id = block_table[logical_block_idx]
        K_block = K_cache[physical_block_id, head_idx]  # [head_dim, block_size]
        V_block = V_cache[physical_block_id, head_idx]

        scores = q @ K_block / sqrt(head_dim)         # [block_size]
        m_new = max(m, scores.max())
        l_new = exp(m - m_new) * l + sum(exp(scores - m_new))
        o = exp(m - m_new) * o + V_block @ exp(scores - m_new)
        m, l = m_new, l_new

    output[seq_idx, head_idx] = o / l

# Ativar prefix caching
vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --enable-prefix-caching \
    --block-size 16 \
    --max-model-len 32768 \
    --gpu-memory-utilization 0.92

# Llama 3.1 70B em 4× A100 80GB com TP=4
vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --quantization awq_marlin \           # INT4 AWQ + Marlin kernel
    --kv-cache-dtype fp8_e5m2 \           # KV cache em FP8 (Hopper+)
    --max-model-len 32768 \
    --max-num-seqs 128 \
    --enable-prefix-caching \
    --enable-chunked-prefill \
    --gpu-memory-utilization 0.93 \
    --port 8000

# Throughput esperado (chat tipico 256 in / 256 out):
# - sem otimizações: ~50 req/min
# - com PagedAttention + cont batching: ~600 req/min
# - + AWQ INT4 + FP8 KV + prefix cache: ~1200-1800 req/min

# Cliente OpenAI apontando para vLLM
from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="not-needed",   # vLLM ignora; configure auth no proxy
)

resp = client.chat.completions.create(
    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
    messages=[
        {"role": "system", "content": "Você é tutor da FFV Academy."},
        {"role": "user", "content": "Explique PagedAttention em 3 bullets."},
    ],
    temperature=0.5,
    stream=True,
    extra_body={
        "min_p": 0.05,
        "repetition_penalty": 1.0,
        "use_beam_search": False,
    },
)
for chunk in resp:
    print(chunk.choices[0].delta.content or "", end="", flush=True)