Skip to content

fix(engine): load prefill triton kernels at startup, not mid-request - #169

Open
jason-fxz wants to merge 2 commits into
mainfrom
fix/prefill-jit-warmup
Open

fix(engine): load prefill triton kernels at startup, not mid-request#169
jason-fxz wants to merge 2 commits into
mainfrom
fix/prefill-jit-warmup

Conversation

@jason-fxz

Copy link
Copy Markdown
Collaborator

Two changes, one goal: no triton compile / cuModuleLoadData on the request path.

  1. fla/l2norm: do_not_specialize=["T"], T was tl.constexpr: a new module per unseen prompt length, forever.
  2. _warmup_prefill runs on every attention backend (was triton-only, so fi/fa had no prefill warmup), over a config-derived ladder covering every in-repo size bucket instead of [80, 128]. New --skip-prefill-warmup; FREETOKEN_WARMUP_MAX_LEN caps the ladder.

Suspected trigger of the sm_89 silent hangs (#123, #72; #31 has a py-spy stack inside cuModuleLoadData): a first-launch module load mid-forward on the scheduler thread.

Tested on H100 (driver 580.95.05, torch 2.11.0+cu130), nvidia/Qwen3.6-35B-A3B-NVFP4, offload + fi: l2norm bitwise equal, same latency; a 11..2560-token sweep loads zero in-repo kernels at request time (main: 4-10 per new size); first-request TTFT on a cold cache 15.1 s -> 0.6 s; warmup adds 4.1 s to startup on a warm cache.

The hang does not reproduce on sm_90. This PR removes the trigger; confirming the fix needs
an sm_89 card.

Refs #123, #72.

@jason-fxz jason-fxz changed the title Fix/prefill jit warmup fix(engine): load prefill triton kernels at startup, not mid-request Aug 25, 2026
@gdevenyi

gdevenyi commented Sep 4, 2026

Copy link
Copy Markdown

Tested on sm_89 (2x RTX 6000 Ada), which is the card this PR most wants a data point from:

main 86214a9 + this PR
first request after boot, ~1k prompt, time to first token 0.50 s 0.49 s
single-stream / 8-conc decode 90.8 / 326.0 tok/s 90.1 / 326.9 tok/s
TTFT, ~1k prompt (warm) 0.81 s 0.81 s

Model RadixArk/Qwen3.8-Flash-Next-NVFP4 (qwen4_exp, qsa_sparse attention backend, TP=2, flags in the setup note below). The l2norm change is bitwise-neutral here as claimed, and the wider prefill warmup adds no visible boot time at this size. The cold first-request stall does not exist on this path to begin with: the QSA backend's kernels are already warmed at startup, so 0.5 s cold vs 0.8 s warm-with-load is just the empty machine. About the sm_89 hang you reference (#123, #72): this session booted the same model about 30 times on these Ada cards (TP=1 and TP=2, offload and hybrid) without a single hang, on main and with several PRs applied, so whatever triggers it needs more than sm_89 + first-launch JIT. Applies to main only with a 3-way merge (the branch is off an older base).

Setup: 2x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4, no NVLink), 2x Xeon Gold 6526Y, 503 GiB RAM, CUDA 13.3, torch 2.11+cu130, sgl_kernel 0.4.5; RadixArk/Qwen3.8-Flash-Next-NVFP4 served as one TP=2 instance (local qwen4_exp TP patch on main 86214a9, plus a local fix so --moe-cache-auto honours --num-tokens, see #383). Flags: --moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --moe-prefill-hit-d2d --max-running-requests 16 --cuda-graph-max-bs 16 --tp-size 2 --gpu 0,1. Single-stream = median of three 64-vs-256-token completion pairs; aggregate = 8 concurrent 256-token completions; TTFT on a ~1k-token prompt; one run per configuration, baseline spread about +-3%.


Update 2026-09-05, merged into the deploy branch. Tried on 2 x RTX 6000 Ada (sm_89) serving Qwen3.8-Flash-Next (RadixArk NVFP4) at TP=2, offload backend, fp8 KV pool of 8 x 262,144 tokens, merged onto my deploy branch (main af71ba4 + #385/#386/#389/#392/#354 and ten other open PRs), tests run on the box, then put in production.

On this model the QSA sparse backend got no prefill warmup at all before this PR (the call was gated on the triton backend), so the first request of each size bucket compiled on the request path after a Triton-cache-cold start. With this PR: Prefill warmup over 10 lengths: [9, 16, 17, 32, 33, 48, 65, 80, 103, 4096], 16.6 s at startup at TP=2 (offload, 8 running), then no compiles during the benchmark; serving numbers unchanged (99.7-102 tok/s single-stream, 344-355 at 8 concurrent). Merge note: it conflicts with #231 in engine.py only because both add a method at the same spot. The fla/l2norm despecialization is already on main, so only the warmup part applied here.

gdevenyi added a commit to gdevenyi/FreeToken that referenced this pull request Sep 5, 2026
…ernels at startup, not mid-request

Upstream FlashML-org#169 at a8326af, merged onto deploy/chatdnp for the PR sweep.
Conflict in engine/engine.py: both sides add a method at the same place (FlashML-org#231's
_emit_moe_stats, FlashML-org#169's _warmup_prefill_lens); kept both.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
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.

2 participants