[https://nvbugs/7029684][fix] Precompile MiniMax FMHA variants in wheels by ruodil · Pull Request #21443 · NVIDIA/TensorRT-LLM

ruodil · 2026-11-02T14:34:17Z

Description

A cold MiniMax M3 mixed context/decode step can block while a newly selected MSA FMHA variant compiles synchronously. In a four-GPU, 128-request 1024/128 reproduction at the same main commit, GB300 spent 101.746 seconds compiling variant 1_0_1_3_1_0_0 during a 102.014-second host step; B300 spent 139.614 seconds compiling the same variant during a 139.862-second host step. See NVBug 7029684.

Build the complete legal FMHA dispatch matrix (350 variants) and its plan, sparse top-k and reduction helpers into wheels. Compatible Linux wheels load packaged modules after source/header, TVM FFI, host architecture and GPU compatibility checks. Source installs and explicit diagnostic modes retain the existing JIT path. This changes kernel packaging and loading; it does not add warmup requests, prepopulate a runtime cache, or change host KV sizing.

The builder uses a private temporary directory, probes real nvcc support for CMake's configured targets, and verifies every expected export. Incremental wheel staging removes obsolete artifacts. Ordinary wheel staging rejects symlinked staging paths to avoid modifying a linked shared tree; editable mode retains setuptools' behavior.

Draft: controlled unprofiled A/B validation remains incomplete. A separate B300 launcher-only smoke now passes after isolating inherited Slurm identity; it does not measure MiniMax or FMHA performance. Previous full benchmark attempts stalled or were killed. The synchronous-compilation mechanism is established for the shortened main diagnostic. It does not yet establish the exact cause of the original release/1.3 478% comparison, or a controlled before/after improvement for that report. This PR targets main.

Validation

Current head: c215afcf31decdc23a1fd322b8044de73efa01ed. Production implementation remains unchanged since e276b6a6. The latest test-only fix enables the existing explicit MPI launcher in both the default cold-cache helper and both A/B arms. Eleven focused local checks passed, including the real command-generation/run wrapper and environment restoration; six command/environment cases failed before this fix. All configured pre-commit hooks passed on the changed file, including Ruff and formatting. Independent review found no actionable defects. These local checks do not establish GPU benchmark completion. Prior-head GitHub checks and the prior 101 CPU tests remain historical evidence; current-head GitHub checks and remote runtime validation must complete separately.
Full SM103a wheel builds succeeded on ARM64 and x86-64. All 350 variants and three helpers linked and passed export checks. FMHA build durations were 3,583.1 s on the GB300 build host and 2,110.9 s on the B300 build host; these are build costs, not inference latency.
Real GPU numerical and CUDA Graph tests: 10 passed, zero skipped, zero failures on each architecture. Actual JUnit was checked for B300 #4 and GB300 #4. These jobs use implementation commit e276b6a6; its production files are byte-identical to the current head, whose only additional file is the explicit performance regression test. Jenkins UNSTABLE reflects missing performance CSV ingestion for these numerical-only jobs; it is not a numerical-test failure. Numerical tests alone do not establish absence of JIT fallback.
Same-implementation four-GPU profiled runs completed 512 requests, 1024 input / 1024 output: B300 #3 reported 11,832.4 output tokens/s; GB300 #3 reported 12,515.6. Maximum recorded rank-0 host step was 1,270.8 ms / 991.7 ms; the 5-context / 3,968-context-token / 24-generation step was 217.0 ms / 255.8 ms. Recorded rank-0 KV offload/onboard bytes were zero. These runs include nsys and calibration, so their throughput is not a clean before/after comparison with the original report. B300's post-benchmark standalone preflight skipped the existing production MSA import initialization and failed; the pipeline retried with the same output/cache directory, so B300 Update TRT-LLM code #3 must not be treated as a cold-cache proof. GB300 logged UCX interface errors, a shutdown error and a pip click dependency conflict; their impact remains unquantified.
The default-harness clean test verifies compatible packaged modules, runs the production case with its default warmup, and fails if the benchmark creates a fresh MSA JIT cache. It is explicitly invoked QA coverage, not automatically collected CI coverage. Both update aarch64 libraries to main branch #6 jobs reached that check but failed before model loading because the job omitted LLM_MODELS_ROOT; no benchmark result was produced. The model-path-corrected tasks at df5c4c2 were: B300 #7 and GB300 #7. No product code or Jenkins configuration changed for this environment correction. Both update aarch64 libraries to release/0.5.0 branch #7 jobs subsequently ended UNSTABLE after three SIGKILL/exit-137 attempts, each about 31 minutes after pytest collection, and archived no complete JUnit/CSV. They provide no clean performance result.

The local checks cover CMake target propagation, actual setuptools wheel/editable staging, manifest and dispatch coverage, and real Ninja invocation with controlled compiler inputs. A separate real nvcc 13.4.59 / CUTLASS 4.4.2 single-variant build and TVM FFI load succeeded, but is not substituted for GPU evidence above.

Remaining validation includes the unprofiled same-head cold-cache benchmarks, controlled comparison without the packaged artifacts, package-size/load-time costs, and assessing the nonfatal attribution warning (ninja -t inputs returned no wheel-target inputs) seen in both builds. The current main perf harness already includes a traffic warmup wave; default-harness throughput must not be presented as evidence independent of that existing warmup. The proposed fix itself adds no warmup. A separate A/B test at current head disables that extra wave in both arms (retaining the original eight warmup requests), runs the installed package and a private source-identical copy with only _aot and bytecode omitted, uses distinct cold cache roots and output directories, and checks the expected absence/presence of compiled cache modules. The installed package is never modified. Same-head B300 #9 and GB300 #9 built both wheels successfully, but both jobs ended UNSTABLE after runtime attempts terminated with SIGKILL/exit 137 and PMIx errors; neither job archived completed JUnit or per-arm CSV results, so no A/B result is established. Each A/B pair executes within one GPU allocation; results are needed before claiming a measured improvement. The preceding #8 test version omitted creation of per-arm output directories; read-only review reproduced its CSV writer failure. The current test creates those directories and passed the improved smoke using the real CSV writer. Archive raw JUnit XML as well as both CSVs: CSV alone does not establish completed request count or cache assertions.

PR Checklist

Final CUDA wheel and GPU validation completed; review the remaining checklist items before marking ready.

Current validation blocker

The performance harness has an explicit 1,800-second no-output stall detector that kills its child process group (tests/integration/defs/perf/utils.py). The earlier benchmark exits alone did not establish OOM, a scheduler fault, or FMHA failure.

The completed B300 MPI diagnostic #10 used the already published matching cf96e9ec main wheel/source, without this PR or a model benchmark. test_mpi_comm_session_submit_runs_once_per_rank passed with explicit mpirun (17.25 seconds); test_mpi_session_basic failed after 300 seconds with MpiPoolSession(wait_shutdown=True): worker identity collection incomplete (0/4 valid identities). The faulthandler stack is in mpi_session.py::_collect_worker_identities, and raw JUnit XML was archived. This narrows the blocker to dynamic MPI worker startup on that allocation. It does not establish that all earlier A/B failures had the same cause or validate the FMHA fix.

Before another full A/B run, verify the explicit-launch benchmark path on the existing wheel and preserve incremental child logs. Unprofiled same-head cold-cache A/B evidence on B300/GB300, package-size/load-time assessment and attribution-warning assessment are still required. The PR remains Draft and the final validation checklist stays unchecked.

Launch validation update (2026-11-10)

The preceding B300 explicit-launch smoke #11 failed before model execution because in-container srun was missing and MPI had insufficient slots. Only the personal single-node smoke environment was corrected: PRTE_MCA_plm=ssh, PRTE_MCA_ras=^slurm, and PRTE_MCA_rmaps_default_mapping_policy=slot:oversubscribe.

The completed corrected B300 smoke #12 still failed both launcher tests (2 failed in 172.96 seconds), using the existing matching cf96e9ec main wheel/source without this PR, a wheel rebuild, or a model benchmark. The raw JUnit and console show child processes reporting mpi_rank: 0 / tllm_mpi_size: 1, followed by RuntimeError: only rank 0 can start multi-node session, got 1 and MPI abort. The job ended UNSTABLE; this is a failed launcher validation, not a successful FMHA benchmark.

The earlier missing-srun/slot errors have therefore been replaced by an MPI communicator/session initialization blocker on this allocation. The exact mismatch still needs investigation; the evidence does not establish an FMHA regression or a working full A/B launch path. No further full A/B run is submitted before the launcher smoke succeeds. Controlled B300/GB300 cold-cache A/B results, package-size/load-time costs, and attribution-warning assessment remain open. Draft and the unchecked final GPU-validation checklist are retained.

Launcher identity diagnostic (2026-11-12)

B300 launcher smoke #13 completed with 2 passed, 0 failed/errors/skipped in 43.998 seconds (raw JUnit). This is actual per-case runtime evidence from the existing matching cf96e9ec main wheel/source, without this PR, a wheel rebuild or a model benchmark. The installed launcher SHA-256 exactly matches that source. ModulePerfTest and artifact stages succeeded; overall Jenkins UNSTABLE is solely the missing performance CSV ingestion for this model-free smoke.

Per-rank input logs confirm the failure mechanism: inner MPI processes have OMPI_COMM_WORLD_RANK=0/1, OMPI_COMM_WORLD_SIZE=2, but both inherit outer SLURM_PROCID=0, SLURM_NTASKS=1. The installed shell launcher prefers the inherited Slurm identity, whereas Python sees the inner MPI communicator. A temporary launcher shim in the personal single-node smoke removes seven Slurm task variables immediately before the installed launcher. The outer pytest/srun environment and GPU allocation remain intact. Effective identities are now 0/2 and 1/2, and every server/task/worker exits zero. Incremental rank stdout/stderr, launcher path/hash, raw JUnit and exit status are archived.

The B300 launcher-only blocker is therefore resolved in the personal validation environment. Carry this proven scoped correction into controlled validation rather than repeating the uncorrected launch. Controlled same-head cold-cache B300/GB300 A/B, package-size/load-time costs and attribution-warning assessment are still required; none is established by this smoke. Draft and the final GPU/performance checklist remain unchanged.

AOT wheel reuse is still unverified: the previous PR-built wheels were not archived or published, and existing direct read access could not authenticate to their build storage. Git confirms that their production sources match the current PR, but file existence, hashes and packaged manifests remain unverified. A bounded metadata-only probe is prepared locally but has not been executed. Ordinary main wheels cannot substitute for the PR AOT wheels; no new controlled A/B job has been submitted. GB300 launcher confirmation also remains outstanding.