Repository navigation
Conversation
This was referenced Nov 1, 2026
FLUX.2-klein, Z-Image-Turbo, LTX-2.5 and Wan2.2 ship with attention_backend: auto, the engine's ragged attention resource with the platform's default kernel (FlashInfer FA3 on an H100, AITER's asm flash attention on MI300X/MI325X); sdpa stays for parity runs and `flashinfer` means what it did. Wan2.2 declares a self- and a cross-attention resource and its step (one segment per guidance branch), and its DiT takes the engine's callables instead of calling SDPA itself. The AITER ragged manager no longer stages every layer's output through one shared buffer and clones it out: the kernel allocates per call, which under capture is a graph-pool block per layer, so the per-layer copy (32 x 37 MB a step on klein-9B) is gone. A ROCm test scores the path at the served DiT shapes against fp64 SDPA.
…b, off on AMD GPUs)
A bundle is a directory per (GPU product, torch, HIP or CUDA, Triton) key holding one uncompressed tar per cache (Inductor, Triton, MIOpen) and a manifest. Packing appends only new or changed files under a lock; unpacking extracts the bundle for this stack into the cache dirs a server here resolves and leaves a marker that each worker reads at startup to log the key and whether the bundle matches. mstar serve --compile-bundle unpacks before starting.
The runner hints dynamo that each static input's bucket-varying dim is dynamic (maybe_mark_dynamic, so a forward that needs a fixed size still compiles per bucket) and compiles under the platform's overrides for any count. Size 1 keeps a graph of its own. Each bucket is still captured separately. Meant for encoders over fixed-size windows, whose per-bucket Inductor compiles are most of a cold start.
MIOpen keeps each database in one file (find and perf DBs as key=value text, kernels in SQLite), so packing caches warmed separately replaced the bundle's records with the last host's. A changed database is now merged with the bundle's last copy, this host's record winning on the same key, and appended only when the merge adds something. MIOpen files are compared by content: two warm-ups can leave databases of the same size and mtime.
amd-flashinfer 0.6.18+amd.1's generated AITER prefill config names __hip_bfloat16 before anything declares it, because the JIT's -DUSE_ROCM=1 turns off the hip_bf16.h include in torch's BFloat16.h. Every bf16 model on FlashInfer's auto backend dies at setup with a ninja log. Give the JIT the header through FLASHINFER_EXTRA_CUDAFLAGS on ROCm; 0.6.18+amd.2 includes it itself.
… pool's workers The pool is built under a config patch at the top of warmup, since the setting is part of a graph's cache key and an env default would recompile everything a golden cache holds.
A submodule can compile at load, and klein does, so the pool built at the top of warmup came too late: the first compile had already built it under inductor's defaults, quiesce included. The call moves to the top of EngineManager.build, so no model has to ask for it.
…r copied per layer
… request through the fused prologue
…draft passes at one or two rows
…attention keeps replaced row starts
…d head, refuse a context past the rope table
…p an unused availability check
merceod
force-pushed
the
users/atindra/pr-36-t6-fused-draft
branch
from
November 1, 2026 15:06
38a1d22 to
3707ec6
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #1090 (pr-24-t3-specdec). Picks the commits of #1091, #1096, #1099, #1101 first; they drop out of this diff as those merge.
Why
#1090's draft passes ran over the verify's (rows, k+1) tokens and threw k of them away after attention, and the step's bookkeeping was 88 small torch kernels plus 148 copies and fills. A 9B k=2 step at c=1 spent 886 us per drafted token on drafting (one GPU trace, shared node). After that, the 248k x 4096 vocabulary head is most of a draft pass.
What changed
Stacked on #1090, with #1091 (bringing #1085), #1096, #1099 (bringing #1092) and #1101 (bringing #1100) picked first: #1099 for 7915843, the fp8 op the draft head uses, #1096 for the compile_dynamic capture configs, #1101 because this PR's decode configs give a seeded start 22 compile-pool jobs instead of 1, which exposes the hang #1101 fixes.
Qwen3_5MTPHead.draft,SpecDecodeResource.attend_draft,attend_rows): pass 0 projects all k+1 tokens and attends one query a row, later passes run one token a row.mstar::qkv_norm_rope_cache: q/k norm, rope and the MTP layer's KV write in one kernel.mstar::spec_verify_inputsandmstar::spec_acceptreplace the bookkeeping with two Triton kernels.record_draftrunssample_unfilteredat temperature 0 in the graph, into the next step's token register. GDN verify rows run on FlyDSL's chain kernel.decode_rows=Truecompiles a verify bucket as decode by its rows, and the MI325X decode GEMV configs gain rows 3, 6, 12 and the fc.draft_head: fp8: an fp8 copy of the head with one scale for all of it, AITER wvSplitKQ at 1-2 rows,_scaled_mmfrom 3. 415 to 213 us at one row, +1 GB of VRAM on 9B.draft_head_variants.pymeasures acceptance per head variant.decode_rowsdefaults to False and only the ROCminductor_overridesread it, the rest sits behindnum_draft_tokens, which no CUDA deployment declares. rust/ unchanged. Docs updated.Verification
set_sync_debug_mode("error"), and the fp8 head at 1, 2 and 4 rows. Three defects found reading the diff (argmax buffers sized bymax_requests, a crash on an ungated head, a context past the rope table) are fixed with tests in d1e608d.Not in this PR
draft_head: bf16turns it off.event_syncbursts show up on plain arms too. FlyDSL's standalone GDR test script segfaults in a bare process after an AiterVerify case, never in pytest or a server.