Conversation
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
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.
FA4's CTA count and exp2 emulation policy can depend on the GPU and attention shape. This opt-in draft tunes those choices through VisualGen's existing warmup
AutoTuner, caches the result for inference, and retains FA4's default heuristic for untuned inputs.The tuner covers dense noncausal FP16/BF16 MHA with head dimension 128 on SM100/SM103. It compares 1CTA and eligible 2CTA with a bounded set of exp2 policies, preserves output and float32 LSE, and includes automatic split-KV as a baseline candidate. Exact shapes, dtype, strides, device, dependency/build identity and dispatch policy separate cache entries. Enable with
TLLM_VISUAL_GEN_FA4_AUTOTUNE=1.FA_DISABLE_2CTA=1and FA4's CUDA 12 guard remain authoritative, complementing #19292.FA4 b19 lacks the required per-call controls.
3rdparty/patches/flash_attn_4_b19.patchadds CTA/exp2 arguments, validates eligible inputs and includes exp2 policy in the compilation key without changing shared tuning tables. Calls omitting the new arguments retain their original behavior.The build produces a separate patched FA4 dependency wheel:
940cd9680f3315f2f06b43ab5bea2c2cf2d96806and applies the patch through3rdparty/prepare_fa4.py. The revision corresponds to b19; all 50 upstream package Python files match the published wheel before patching.scripts/build_wheel.pyuses FA4's upstream setuptools backend to buildflash_attn_4-4.0.0b19+trtllm.1-py3-none-any.whl, installs it before generating TRT-LLM Python stubs, and emits it beside the TRT-LLM wheel. The package retainsflash_attn.cute, includes upstream license/author files and adds build provenance for autotuner cache separation. Kernels still JIT-compile on the target GPU.flash-attn-4==4.0.0b19+trtllm.1and contains no FA4 package files. Stock b19 cannot satisfy that dependency. Source bootstrap requirements retain upstream b19; reinstalling them accepts the installed patched local version.--find-links; Jenkins artifact archives and uploads carry both wheels. Test paths using--no-depsinstall both explicitly, and precompiled installation checks receive the companion artifact location. Python examples andtrtllm-serveuse the installed patched package without a runtime patch step.sana-sol-attn.patch, its vendor lock, or vendored Sol-Attn sources in the final PR diff.See build and upgrade notes and the standalone tuning demo. No upstream FA4 version upgrade is included. An upstream register-allocation fix alone does not replace this per-call API.
Validation:
trtllm-serveHTTP endpoint. Default and tuned outputs, and the outputs across these entry points, are bitwise identical (LPIPS 0). Actual pipeline warmup populates FA4 cache entries; steady inference does not profile new tactics.setup.pyand resolve the separate patched FA4 dependency with--find-links; stock b19 cannot satisfy the exact runtime pin. Source, patch and version mismatch rejection and incremental staging are covered. The actual CMake FetchContent/patch and upstream FA4 wheel build were also exercised in the preceding packaging update, with all package files verified through wheel installation.Keep this draft opt-in as a validation prototype. Before productizing runtime search, compare it with a simple validated static FA4 policy: the existing SM100 noncausal h128 2CTA table uses frequency 10, so frequency 16 versus that default and 1CTA is a focused next experiment. Retain runtime tuning only if workload-dependent winner changes deliver meaningful end-to-end gains beyond the simpler policy. A global FA4 table change also needs coverage beyond VisualGen self-attention. Before an index-based release, publish the companion FA4 wheel to the configured NVIDIA package index and verify resolution there. This PR builds and archives it; no package has been published. The release channel must support its downstream local version (
+trtllm.1).