Skip to content

[Engine] Packed prefill resolves automatically, for Huginn too - #149

Open
0z5a wants to merge 1 commit into
InferMatrixLab:mainfrom
0z5a:engine/auto-packed-prefill
Open

0z5a wants to merge 1 commit into
InferMatrixLab:mainfrom
0z5a:engine/auto-packed-prefill

Conversation

@0z5a

@0z5a 0z5a commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Summary

ExecutionConfig.prefill_uva defaults to None and resolves at engine construction: packed (variable-length) prefill when the attention backend packs prefill (Triton opts in, FlashAttention-4 natively), the KV layout is LAST_EXITED and prefill is not staged (wavefront chunking and mixed mode keep the per-row path their schedulers were written against). An explicit True/False always wins; --no-prefill-uva keeps per-row prefill. A CPU engine resolves to per-row prefill and records the decision.

Huginn (boundary KV) packs prefill too: prelude and coda already build packed metadata at depth 0. VLLM_RLT_HUGINN_PACKED_PREFILL=0 opts out. The GSM8K harness accepts Huginn checkpoints (mean_recurrence depth) so the gate below could run.

Packed bf16 prefill uses different arithmetic from per-row prefill, so a few greedy sequences change; this change was gated on accuracy rather than identity.

Measurements

GH200, 2048-token prompts (Ouro async + graphs, Huginn sync + graphs), one run per level:

C=8 C=32 identical
Ouro-1.4B tok/s, packed vs per-row 2.42× 2.00× (TTFT −61%) 37/40
Huginn-0125 tok/s, packed vs per-row 2.07× 2.08× (TTFT p50 41.9 → 18.1 s) —
  • Huginn, 512-token prompts: 1.20× / 1.33× at 8 / 32; 36/41 identical.
  • Ouro, 512-token prompts: within noise of per-row prefill (0.98× / 1.19× at 8 / 32, two kept runs; a per-row rerun gave 0.98× / 1.10×). Without UVA staging packed prefill fell to 0.88× / 0.93×.
  • Ouro, 64 agent sessions: 66.0 vs 54.7 tok/s with per-row prefill on the same tree (1.21×). RTX 5090, 32 sessions on main: 2.17× (per-row prefill managed about 3K prompt tokens/s).
  • GSM8K gate for Huginn (greedy; the harness timed out both arms, so the 70 questions both completed): per-row 20, packed 21 correct, 47/70 sequences identical.

Platforms. GH200: torch 2.14.0+cu130, Triton 3.8.0, NIXL 1.4.1. RTX 5090: torch 2.13.0+cu130, Triton 3.7.1. Models: Ouro-1.4B (24 layers, 4 loops, fixed depth) and Huginn-0125 (2 prelude + 4 recurrent + 2 coda layers, 32 loops). Burst workloads issue C requests at once with random-token prompts; session workloads replay multi-turn agent traces with the prefix cache on.

Tests

  • CPU suite: 830 passed; ruff passes. New test covers the resolution rules (CUDA/CPU, layouts, boundary KV and its opt-out, staged prefill, explicit settings, CPU engine recording).
  • GPU tests on this branch (GH200, head 7cc7dd0, flash-attn-3 3.0.1 and flash-attn-4 4.0.0b30 installed, pytest --run-gpu tests/test_engine.py tests/test_split_prefill.py tests/test_input_preparation.py tests/test_huginn.py): 180 passed, 1 failed: test_cuda_split_prefill_survives_decode_graph_capture_and_replay (num_preluded_tokens 4 instead of 2). It fails the same way on main (66920da) on this host, so this branch adds no failure; the FlashAttention-4 split-prefill case, skipped without the package, ran and passed. Without FlashAttention installed, the FlashAttention tests failed at import on this branch and on main (66920da) alike.

Notes

The GPU numbers come from a development branch that contains all changes of this series; most cells are single runs, and on the shared GH200 host two runs of main at 32 requests differed by up to 16%, so smaller differences are ties. This branch alone was checked with the CPU suite, ruff and the GPU tests above.

Series. Draft PRs from one study of serving looped models. They are independent except #152, which is stacked on #147; several touch the same files, so later ones may need a rebase after earlier ones merge.

ExecutionConfig.prefill_uva now defaults to None and resolves at engine
construction: packed (variable-length) prefill when the attention backend
packs prefill (Triton opts in, FlashAttention-4 natively), the KV layout is
LAST_EXITED and prefill is not staged (wavefront chunking and mixed mode keep
the per-row path their schedulers were written against). An explicit
True/False always wins; --no-prefill-uva keeps per-row prefill. A CPU engine
resolves to per-row prefill and records the decision.

Huginn (boundary KV) now packs prefill too: prelude and coda already build
packed metadata at depth 0. VLLM_RLT_HUGINN_PACKED_PREFILL=0 opts out. The
GSM8K harness accepts Huginn checkpoints (mean_recurrence depth) for the gate.

Packed bf16 prefill uses different arithmetic from per-row prefill, so a few
greedy sequences change; it was gated on accuracy instead of identity.

GH200, 2048-token prompts, async + graphs (Ouro) / sync + graphs (Huginn):
  Ouro-1.4B: 2.42x / 2.00x throughput at 8 / 32 requests, TTFT -61% at 32;
    37/40 requests identical.
  Huginn-0125: 2.08x at 32 requests, median TTFT 41.9 -> 18.1 s; 512-token
    prompts 1.20x / 1.33x at 8 / 32, 36/41 identical.
  Ouro, 512-token prompts: within run-to-run noise of per-row prefill
    (0.98x / 1.19x at 8 / 32, two kept runs; a per-row rerun gave 0.98x / 1.10x).
GSM8K gate (Huginn, greedy, the 70 questions both arms completed): per-row 20,
  packed 21 correct, 47/70 sequences identical.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant