diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index 11d249c09245..cd1094eeb003 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -2468,11 +2468,6 @@ class LlmRequest : public GenericLlmRequest LlmRequest(LlmRequest&& request) = default; LlmRequest(LlmRequest const& request) = default; - /// @brief Create a Response from the current state of the request - /// @details Note that there is some dependency on the order of operations in this method. Modify with care! - /// @return An optional Response - std::optional createResponse(bool useFastLogits = false, int32_t mpiWorldRank = 0); - std::optional createResult(bool useFastLogits = false, int32_t mpiWorldRank = 0); void createSerializedResult( @@ -2488,10 +2483,6 @@ class LlmRequest : public GenericLlmRequest std::shared_ptr createChildRequest(RequestIdType requestId); - void movePromptEmbeddingTableToGpu(runtime::BufferManager const& manager); - - void moveLoraWeightsToGpu(runtime::BufferManager const& manager); - // Remove LoRA weights and LoRA config tensors void removeLoraTensors(); }; diff --git a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp index 958d1092b754..ecb3dd0182f8 100644 --- a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp +++ b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp @@ -52,17 +52,6 @@ runtime::SizeType32 GenericLlmRequest::getBeamWidthByIter(bool template class GenericLlmRequest; -std::optional LlmRequest::createResponse(bool useFastLogits, int32_t mpiWorldRank) -{ - auto requestId = isChild() ? mParentRequestId : mRequestId; - auto result = createResult(useFastLogits, mpiWorldRank); - if (result.has_value()) - { - return executor::Response(requestId, result.value(), mClientId); - } - return std::nullopt; -} - void LlmRequest::createSerializedResult( std::vector& serializedResult, bool& isFinal, bool useFastLogits, int32_t mpiWorldRank) { @@ -355,29 +344,6 @@ std::shared_ptr LlmRequest::createChildRequest(RequestIdType request return childReq; } -void LlmRequest::movePromptEmbeddingTableToGpu(runtime::BufferManager const& manager) -{ - if (!mPromptEmbeddingTable.has_value() - || mPromptEmbeddingTable.value()->getMemoryType() == runtime::MemoryType::kGPU) - { - return; - } - - TensorPtr gpuPromptEmbeddingTable = manager.copyFrom(*mPromptEmbeddingTable.value(), runtime::MemoryType::kGPU); - mPromptEmbeddingTable = gpuPromptEmbeddingTable; -} - -void LlmRequest::moveLoraWeightsToGpu(runtime::BufferManager const& manager) -{ - if (!mLoraWeights.has_value() || mLoraWeights.value()->getMemoryType() == runtime::MemoryType::kGPU) - { - return; - } - // TODO for tp / pp models we only need to move the bit that belong on the local device - TensorPtr gpuLoraWeights = manager.copyFrom(*mLoraWeights.value(), runtime::MemoryType::kGPU); - mLoraWeights = gpuLoraWeights; -} - void LlmRequest::removeLoraTensors() { mLoraWeights.reset(); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index c75b8893d445..85cc16fc9352 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -484,8 +484,6 @@ void initBindings(nb::module_& m) .def("validate", &tb::LlmRequest::validate, nb::arg("max_input_len"), nb::arg("max_seq_len"), nb::arg("max_draft_len"), nb::arg("vocab_size_padded"), nb::arg("max_endocer_input_len") = std::nullopt, nb::arg("enable_kv_cache_reuse") = false) - .def("create_response", &tb::LlmRequest::createResponse, nb::arg("use_fast_logits") = false, - nb::arg("mpi_world_rank") = 0) .def("create_child_request", &tb::LlmRequest::createChildRequest, nb::arg("child_id")) .def("create_result", &tb::LlmRequest::createResult, nb::arg("use_fast_logits") = false, nb::arg("mpi_world_rank") = 0) @@ -497,8 +495,6 @@ void initBindings(nb::module_& m) self.createSerializedResult(serialized_result, is_final, use_fast_logits, mpi_world_rank); return std::make_tuple(nb::bytes(serialized_result.data(), serialized_result.size()), is_final); }) - .def("move_prompt_embedding_table_to_gpu", &tb::LlmRequest::movePromptEmbeddingTableToGpu, nb::arg("manager")) - .def("move_lora_weights_to_gpu", &tb::LlmRequest::moveLoraWeightsToGpu, nb::arg("manager")) .def("finish_by_reason", &tb::LlmRequest::finishByReason, nb::arg("finish_reason")) .def("set_first_scheduled_time", &tb::LlmRequest::setFirstScheduledTime) .def("update_perf_metrics", &tb::LlmRequest::updatePerfMetrics, nb::arg("iter_counter")) diff --git a/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp b/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp index 600202c8b026..a0110b0ad93d 100644 --- a/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp @@ -463,7 +463,7 @@ class ParamTest : public LlmRequestTest, public ::testing::WithParamInterface(GetParam())}; bool const excludeInputFromOutput{std::get<1>(GetParam())}; @@ -517,8 +517,8 @@ TEST_P(ParamTest, createResponse) for (auto& llmReq : llmRequests) { - auto response = llmReq->createResponse(); - EXPECT_FALSE(response); + auto resultOpt = llmReq->createResult(); + EXPECT_FALSE(resultOpt); } SizeType32 constexpr numIterations{5}; @@ -536,15 +536,14 @@ TEST_P(ParamTest, createResponse) } llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - auto response = llmReq->createResponse(); - EXPECT_TRUE(streaming == response.has_value()); + auto resultOpt = llmReq->createResult(); + EXPECT_TRUE(streaming == resultOpt.has_value()); for (int beamIdx = 0; beamIdx < numReturnBeams; ++beamIdx) { if (streaming) { - EXPECT_EQ(response.value().getRequestId(), requestId); - auto result = response.value().getResult(); + auto const& result = resultOpt.value(); EXPECT_EQ(result.outputTokenIds.size(), numReturnBeams); auto const& beamTokens = result.outputTokenIds.at(beamIdx); if (returnAllGeneratedTokens) @@ -564,8 +563,8 @@ TEST_P(ParamTest, createResponse) } } - response = llmReq->createResponse(); - EXPECT_FALSE(response); + resultOpt = llmReq->createResult(); + EXPECT_FALSE(resultOpt); } } @@ -584,21 +583,17 @@ TEST_P(ParamTest, createResponse) for (auto seqIdx = 0; seqIdx < numSequences; seqIdx++) { auto llmReq = llmRequests.at(seqIdx); - auto response = llmReq->createResponse(); + auto resultOpt = llmReq->createResult(); if (!streaming && llmRequests.at(seqIdx)->getState() != tb::LlmRequestState::kGENERATION_COMPLETE) { - EXPECT_FALSE(response); + EXPECT_FALSE(resultOpt); continue; } - EXPECT_TRUE(response) << "seqIdx " << seqIdx; - EXPECT_FALSE(response.value().hasError()) << "seqIdx " << seqIdx; + EXPECT_TRUE(resultOpt) << "seqIdx " << seqIdx; - // All response should have the same request id of the original request. - EXPECT_EQ(response.value().getRequestId(), requestId); - - auto result = response.value().getResult(); + auto const& result = resultOpt.value(); EXPECT_EQ(result.outputTokenIds.size(), numReturnBeams); // Only the first sequence has finished. @@ -661,11 +656,10 @@ TEST_P(ParamTest, createResponse) for (auto seqIdx = 1; seqIdx < numSequences; seqIdx++) { - auto response = llmRequests.at(seqIdx)->createResponse(); - EXPECT_TRUE(response) << "seqIdx " << seqIdx; - EXPECT_FALSE(response.value().hasError()) << "seqIdx " << seqIdx; + auto resultOpt = llmRequests.at(seqIdx)->createResult(); + EXPECT_TRUE(resultOpt) << "seqIdx " << seqIdx; - auto result = response.value().getResult(); + auto const& result = resultOpt.value(); // All sequences have finished. EXPECT_TRUE(result.isSequenceFinal) << "seqIdx " << seqIdx; EXPECT_TRUE(result.isFinal) << "seqIdx " << seqIdx; @@ -703,12 +697,12 @@ TEST_F(LlmRequestTest, createResultDisaggContextComplete) // This is the regression case — without the fix, createResult() returns nullopt // because DISAGG_CONTEXT_COMPLETE was not handled by createResult's early guard // or its context-phase branch. - auto response = llmReq.createResult(/*useFastLogits=*/false, /*mpiWorldRank=*/0); - ASSERT_TRUE(response.has_value()) << "createResult() must not return nullopt for DISAGG_CONTEXT_COMPLETE"; - EXPECT_TRUE(response->contextPhaseParams.has_value()) + auto resultOpt = llmReq.createResult(/*useFastLogits=*/false, /*mpiWorldRank=*/0); + ASSERT_TRUE(resultOpt.has_value()) << "createResult() must not return nullopt for DISAGG_CONTEXT_COMPLETE"; + EXPECT_TRUE(resultOpt->contextPhaseParams.has_value()) << "contextPhaseParams must be populated for context-only DISAGG_CONTEXT_COMPLETE requests"; - EXPECT_EQ(response->contextPhaseParams->getReqId(), requestId); - EXPECT_TRUE(response->isSequenceFinal); + EXPECT_EQ(resultOpt->contextPhaseParams->getReqId(), requestId); + EXPECT_TRUE(resultOpt->isSequenceFinal); } TEST_F(LlmRequestTest, generationOnlyRequestAdoptsContextPhaseDraftTokens) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index d77b9fdebaac..66e7b30f141d 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -89,6 +89,7 @@ Models that select the V2 manager by default: | DeepSeek-V4 | Sparse attention attaches auxiliary per-layer buffers | | GPT-OSS | Sliding window on every other layer (VSWA), so the sliding-window and full-attention pools are sized independently | | Gemma3 / Gemma4 (text and multimodal) | Alternating sliding-window and full-attention layers (VSWA); same independent pool sizing | +| Llama / Llama4 | Uniform KV pool layout; chunked attention does not partition the pools | Separately, Gemma4 hybrid attention and sparse-attention models are routed to V2 unconditionally: their per-layer buffer layouts cannot be represented by V1's diff --git a/jenkins/BoltProfileGen.groovy b/jenkins/BoltProfileGen.groovy index 22be6bf98a6f..343276b83cc2 100644 --- a/jenkins/BoltProfileGen.groovy +++ b/jenkins/BoltProfileGen.groovy @@ -84,6 +84,15 @@ PROMOTE = (params.promote ?: "false").toString() // bundle -> bolted tarball on the cluster ("consume immediately after generating"). // The merge job (slurm_merge.sh) runs apply_bolt.py when BOLT_APPLY=1. APPLY_PROFILES = (params.applyProfiles ?: "true").toString() +// publishBoltedCanonical: after merge (and, if applicable, promote), push the +// just-BOLTed tarball back to the input artifactPath under the CANONICAL name so +// downstream consumers of that build get BOLT transparently, preserving the +// original as unbolted-. Requires APPLY_PROFILES=true (the bolted tarball +// is produced by the merge job's BOLT_APPLY=1 step) and PROMOTE=true (only the +// postmerge producer repushes a canonical). Default OFF -- the rollout is flipped +// on in a follow-up change -- so this is inert until then. Resolution mirrors the +// other toggles: param, then env. +PUBLISH_BOLTED_CANONICAL = (params.boltPublishCanonical ?: env.boltPublishCanonical ?: "false").toString() == "true" // Multiply each workload's client `iterations` (num_requests = concurrency * // iterations) to lengthen the measured serving window without editing the shared // perf-sanity configs. Default 64: the point where the aggregated workloads' @@ -372,8 +381,18 @@ def submitProfileGen(pipeline) return 0 } '''.stripIndent() + // If a prior run against THIS artifactPath already published the BOLTed build + // as canonical, publishBoltedCanonical preserved the original as + // unbolted-. Prefer it so we profile/BOLT the un-BOLTed input rather + // than double-BOLTing. Inert for a fresh artifactPath (no unbolted- object), + // where canonical IS the un-BOLTed build. + def unboltedTarUrl = "${URM_ARTIFACTORY_BASE}/${ARTIFACT_PATH}/unbolted-${BOLT_TARNAME}" def tarStage = """ URL='${llmTarfile}' + if curl -fsI '${unboltedTarUrl}' >/dev/null 2>&1; then + echo '[INFO] found unbolted-${BOLT_TARNAME}; using it as the un-BOLTed input (avoids double-BOLT on reuse)' + URL='${unboltedTarUrl}' + fi DEST='${ws}/builds/${BOLT_TARNAME}' PARTS=16 """.stripIndent() + boltFetchLib + ''' @@ -509,6 +528,26 @@ def submitProfileGen(pipeline) pipeline.echo("PROMOTE=false: skipping Artifactory promote of ${bundle}") } + // Publish the freshly BOLTed build back to the input artifactPath under the + // CANONICAL name (unbolted- preserved), so downstream consumers of + // that build get BOLT transparently. The merge job (APPLY_PROFILES=true) + // already produced ${outDir}/bolt-${BOLT_TARNAME} natively on the aarch64 + // cluster node; without this it is reclaimed by the retention sweep below. + // Runs BEFORE that sweep. Coupled to PROMOTE so only the postmerge producer + // (which promotes the bundle) repushes a canonical -- a premerge + // generate-and-consume run (promote=false) never does, even with the toggle + // on. Gated off by default (PUBLISH_BOLTED_CANONICAL). + if (PUBLISH_BOLTED_CANONICAL && APPLY_PROFILES == "true" && PROMOTE == "true") { + stage("Publish BOLTed build as canonical") { + publishBoltedCanonical(pipeline, remote, + "${outDir}/bolt-${BOLT_TARNAME}", // bolted build from the merge job + "${ws}/builds/${BOLT_TARNAME}", // un-BOLTed input on the cluster + "${URM_ARTIFACTORY_BASE}/${ARTIFACT_PATH}", BOLT_TARNAME) + } + } else if (!PUBLISH_BOLTED_CANONICAL) { + pipeline.echo("PUBLISH_BOLTED_CANONICAL=false: not repushing a BOLTed canonical tarball.") + } + // Retention: best-effort purge of workspaces older than 7 days so scratch // doesn't grow unbounded across runs. Depth 4 under the bolt-ci root maps // to ///, i.e. one per-run workspace. @@ -732,6 +771,67 @@ def promoteBundle(pipeline, remote, String bundle) pipeline.echo("Promoted. latest = ${base}/latest.tar.gz") } +// --------------------------------------------------------------------------- +// Publish the just-BOLTed build under the CANONICAL name at the input artifactPath, +// mirroring Build.groovy's premerge apply/consume convention: BOLT is the DEFAULT, +// so canonical becomes the BOLTed build and the original is preserved as +// unbolted-. The merge job (BOLT_APPLY=1) produced the bolted tarball +// natively on the aarch64 cluster node into ${outDir}/bolt-. +// +// Runs CLUSTER-SIDE (the frontend reaches Artifactory and the tarballs live on +// scratch). Fails loudly if the bolted tarball is missing, and is idempotent on a +// reused artifactPath: the unbolted- snapshot is skipped if it already exists, so a +// re-run never overwrites the real unbolted with an already-BOLTed canonical. +// +// Credentials use the same stdin-fed netrc pattern as promoteBundle(): the secret +// travels over ssh stdin (read by `cat > netrc`), never embedded in the base64'd +// script, so it can't leak via the decoded remote command or `ps`. +// --------------------------------------------------------------------------- +def publishBoltedCanonical(pipeline, remote, String boltedLocal, String unboltedLocal, String artifactBase, String tarName) +{ + def host = URM_ARTIFACTORY_BASE.replaceFirst(/^https?:\/\//, "").tokenize('/').first() + def netrc = "${boltedLocal}.netrc" + def canonicalUrl = "${artifactBase}/${tarName}" + def unboltedUrl = "${artifactBase}/unbolted-${tarName}" + pipeline.echo("Publishing BOLTed build as canonical ${canonicalUrl} (un-BOLTed preserved at ${unboltedUrl})") + // Remote side carries no secret: it reads the netrc from stdin. Clean up via + // trap, not a trailing rm, so a failed curl under `set -e` still removes the + // plaintext netrc; umask covers the window before chmod. + def publish = """ + set -e + umask 077 + trap 'rm -f "${netrc}"' EXIT + test -f "${boltedLocal}" || { echo "[ERROR] bolted tarball missing: ${boltedLocal} (did the merge run with BOLT_APPLY=1?)" >&2; exit 1; } + test -f "${unboltedLocal}" || { echo "[ERROR] un-BOLTed tarball missing: ${unboltedLocal}" >&2; exit 1; } + cat > "${netrc}" + chmod 600 "${netrc}" + # 1) Preserve the un-BOLTed build as unbolted-${tarName}, unless already there. + # A PUT-upload, NOT /api/copy: the target repo is virtual, so a server-side + # copy returns 409. Idempotent via the existence check: on a re-run the + # canonical is ALREADY the BOLTed build, and re-publishing would otherwise + # overwrite the real unbolted with a bolted copy. + if curl -fsI --netrc-file "${netrc}" "${unboltedUrl}" >/dev/null 2>&1; then + echo "[INFO] unbolted-${tarName} already exists; skipping upload" + else + echo "[INFO] uploading un-BOLTed tarball -> unbolted-${tarName}" + curl -fS --netrc-file "${netrc}" --retry 5 --retry-all-errors \\ + --connect-timeout 30 --speed-time 300 --speed-limit 1024 \\ + -T "${unboltedLocal}" "${unboltedUrl}" + fi + # 2) Overwrite canonical ${tarName} with the BOLTed build (the release name). + curl -fS --netrc-file "${netrc}" --retry 5 --retry-all-errors \\ + --connect-timeout 30 --speed-time 300 --speed-limit 1024 \\ + -T "${boltedLocal}" "${canonicalUrl}" + """.stripIndent() + pipeline.withCredentials([pipeline.usernamePassword(credentialsId: 'urm-artifactory-creds', + usernameVariable: 'ART_USER', passwordVariable: 'ART_PASS')]) { + def feed = "printf 'machine ${host} login %s password %s\\n' \"\$ART_USER\" \"\$ART_PASS\" | " + Utils.exec(pipeline, timeout: false, numRetries: 2, noNVDFEvent: true, + script: feed + Utils.sshUserCmd(remote, b64BashRemoteCmdStdin(publish, "${boltedLocal}.publish.sh"))) + } + pipeline.echo("Published. canonical = ${canonicalUrl}") +} + pipeline { agent { @@ -798,6 +898,11 @@ pipeline { choices: ["true", "false"], description: "Re-BOLT the input tarball with the just-generated bundle in the merge job, as a same-commit check that the profiles apply. No extra GPU allocation." ) + choice( + name: "boltPublishCanonical", + choices: ["false", "true"], + description: "After merge (requires applyProfiles=true), push the BOLTed build back to the input artifactPath under the canonical name (original preserved as unbolted-). Default false; the rollout is turned on in a follow-up change." + ) string( name: "slurmPlatform", defaultValue: "", diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index 70a824cd2187..7c653c6e2f5d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -602,11 +602,15 @@ def _is_supported_with_reason( and 0 < meta.num_contexts <= 4 and attn.head_dim != 512 and not has_q_only + and meta.num_generations == 0 ): # NVBug 6579626: the per-layer host overhead of the FlashInfer # TRTLLM-Gen context path regresses TTFT for small BF16 batches. # Let the FMHA selector choose the fallback implementation. H512 # and Q-only cached-KV requests cannot use that fallback. + # Keep mixed batches on FlashInfer: the monolithic fallback also + # switches decoding to C++ kernels that generation-only warmup + # does not cover, causing first-use JIT stalls (NVBug 6716104). return False, ( "small-batch BF16 context attention uses the fallback FMHA for " "performance because the FlashInfer TRTLLM-Gen context path " diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py new file mode 100644 index 000000000000..7636f5660b17 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused paged-cache scatter for the MiniMax-M3 MSA backend. + +One Triton launch writes a layer's new-token main K, main V, and (sparse +layers) index-K into their paged HND caches at the step's write slots. +The legacy path costs three aten advanced-indexing writes per layer plus +their index preprocessing; at 60 layers per forward step, all captured +into decode CUDA graphs, the launch count dominates the cost. The kernel +derives each token's (page, within-page) split from ``out_cache_loc`` +in-register, so it needs no precomputed index tensors at all. + +Sources may be strided row views (slices of the fused QKV projection); +only the innermost [num_heads * head_dim] extent must be contiguous. +Stores cast to the cache dtype, which folds the FP8 KV-cache cast in. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fused_paged_scatter_kernel( + k_src, + v_src, + idx_src, + k_cache, + v_cache, + idx_cache, + out_cache_loc, + k_src_row_stride, + v_src_row_stride, + idx_src_row_stride, + kc_stride_page, + kc_stride_head, + kc_stride_tok, + vc_stride_page, + vc_stride_head, + vc_stride_tok, + ic_stride_page, + ic_stride_tok, + tokens_per_block, + H: tl.constexpr, + D: tl.constexpr, + HAS_IDX: tl.constexpr, +): + # int64 throughout: t * row_stride can exceed 2^31 elements on large + # eager prefill steps (num_tokens up to max_num_tokens times the fused + # QKV row stride), and the slot * page-stride products likewise. + t = tl.program_id(0).to(tl.int64) + slot = tl.load(out_cache_loc + t).to(tl.int64) + page = slot // tokens_per_block + within = slot % tokens_per_block + d = tl.arange(0, D) + for h in tl.static_range(H): + k_vals = tl.load(k_src + t * k_src_row_stride + h * D + d) + v_vals = tl.load(v_src + t * v_src_row_stride + h * D + d) + k_dst = k_cache + page * kc_stride_page + h * kc_stride_head + within * kc_stride_tok + d + v_dst = v_cache + page * vc_stride_page + h * vc_stride_head + within * vc_stride_tok + d + tl.store(k_dst, k_vals.to(k_cache.dtype.element_ty)) + tl.store(v_dst, v_vals.to(v_cache.dtype.element_ty)) + if HAS_IDX: + i_vals = tl.load(idx_src + t * idx_src_row_stride + d) + i_dst = idx_cache + page * ic_stride_page + within * ic_stride_tok + d + tl.store(i_dst, i_vals.to(idx_cache.dtype.element_ty)) + + +def _row_stride_if_fusable(src: torch.Tensor, inner: int) -> Optional[int]: + """Row stride (elements) if `src` is a [T, inner] row view with contiguous + rows (e.g. a column slice of the fused QKV projection); None otherwise.""" + if src.dim() != 2 or src.shape[1] != inner or src.stride(1) != 1: + return None + return src.stride(0) + + +def fused_write_layer_caches( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + idx_cache: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + idx_k: Optional[torch.Tensor], +) -> bool: + """Fused single-launch write of new-token K/V (+index-K) into paged HND + caches. Returns False when a layout or device precondition fails, so the caller can + keep the legacy per-cache writes. + + `k_cache`/`v_cache` are [num_pages, num_kv_heads, tokens_per_block, + head_dim] HND views; `idx_cache` is the MQA index-K view with one head. + `k`/`v` are the layer's new-token values as [T, H*D] row views; + `idx_k` is [T, D]. Their inner dimension must be contiguous. + """ + if not k_cache.is_cuda or any( + tensor.device != k_cache.device for tensor in (k, v, v_cache, out_cache_loc) + ): + return False + if k_cache.dim() != 4 or v_cache.dim() != 4: + return False + if v_cache.shape != k_cache.shape: + return False + if k_cache.stride(-1) != 1 or v_cache.stride(-1) != 1: + return False + num_pages, num_heads, tokens_per_block, head_dim = k_cache.shape + if (head_dim & (head_dim - 1)) != 0: + return False + inner = num_heads * head_dim + k_stride = _row_stride_if_fusable(k, inner) + v_stride = _row_stride_if_fusable(v, inner) + if k_stride is None or v_stride is None: + return False + + has_idx = idx_k is not None + idx_stride = 0 + ic_stride_page = 0 + ic_stride_tok = 0 + if has_idx: + if idx_cache is None or idx_cache.dim() != 4 or idx_cache.stride(-1) != 1: + return False + if idx_k.device != k_cache.device or idx_cache.device != k_cache.device: + return False + if int(idx_cache.shape[1]) != 1 or int(idx_cache.shape[3]) != head_dim: + return False + if int(idx_cache.shape[2]) != tokens_per_block: + return False + idx_stride = _row_stride_if_fusable(idx_k, head_dim) + if idx_stride is None: + return False + ic_stride_page = idx_cache.stride(0) + ic_stride_tok = idx_cache.stride(2) + + num_tokens = int(out_cache_loc.shape[0]) + if num_tokens == 0: + return True + if k.shape[0] < num_tokens or v.shape[0] < num_tokens: + return False + if has_idx and idx_k.shape[0] < num_tokens: + return False + + _fused_paged_scatter_kernel[(num_tokens,)]( + k, + v, + idx_k if has_idx else k, # unused when HAS_IDX=False + k_cache, + v_cache, + idx_cache if has_idx else k_cache, # unused when HAS_IDX=False + out_cache_loc, + k_stride, + v_stride, + idx_stride, + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + ic_stride_page, + ic_stride_tok, + tokens_per_block, + H=num_heads, + D=head_dim, + HAS_IDX=has_idx, + num_warps=2, + ) + return True + + +__all__ = ["fused_write_layer_caches"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py index 28716a6dece5..b3aa264d893e 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py @@ -140,7 +140,11 @@ def write_msa_phase_kv( phases cover the step between them and neither repeats the other's write. k and v are the phase's token slice, and token_offset its first token on - the step's token axis, which is what msa_out_cache_loc is indexed by. + the step's token axis, which is what msa_out_cache_loc is indexed by. A + phase handed no K/V (k and v None) has nothing to write: that is how the + MiniMax-M3 model layer, which stores the whole step's K/V itself through + MiniMaxM3MsaSparseAttention.write_layer_caches ahead of its indexer, + tells both libraries the cache is already resident. """ if attention_input_type != AttentionInputType.mixed: raise NotImplementedError( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 5a3f652fbe8c..78944e3c0dab 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -1208,6 +1208,57 @@ def support_fused_rope(cls) -> bool: # index branches explicitly. return False + def write_layer_caches( + self, + k: torch.Tensor, + v: torch.Tensor, + idx_k: Optional[torch.Tensor], + metadata, + ) -> None: + """Write this layer's new-token K, V and (bf16 indexer) index-K. + + One fused kernel launch when the source/cache layouts allow it, else + the legacy per-cache writes. The model layer calls this first, so the + index-K cache is populated before run_indexer's proxy pass reads it, + and then hands forward() k=v=None: write_msa_phase_kv writes nothing + for a phase without live K/V, so neither FMHA library repeats the + write. `idx_k` is None on the FP8 indexer path, where the fused + producer has already inserted E4M3 index-K into the side cache. + `metadata` only supplies the step's write slots (msa_out_cache_loc, + filled by prepare()) and the cache manager. + """ + from .kernels.msa_scatter import fused_write_layer_caches + + layer_idx = self.layer_idx + buffers = metadata.kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") + k_view, v_view = buffers[:, 0], buffers[:, 1] + idx_cache = metadata.msa_idx_k_cache(layer_idx) if idx_k is not None else None + num_tokens = int(k.shape[0]) + out_cache_loc = metadata.msa_out_cache_loc[:num_tokens] + if fused_write_layer_caches(k_view, v_view, idx_cache, out_cache_loc, k, v, idx_k): + return + num_kv_heads = int(k_view.shape[1]) + head_dim = int(k_view.shape[3]) + write_kv_slots( + k_view, + out_cache_loc, + k.reshape(num_tokens, num_kv_heads, head_dim), + layout="HND", + ) + write_kv_slots( + v_view, + out_cache_loc, + v.reshape(num_tokens, num_kv_heads, head_dim), + layout="HND", + ) + if idx_k is not None: + write_kv_slots( + idx_cache, + out_cache_loc, + idx_k.reshape(num_tokens, 1, int(idx_cache.shape[-1])), + layout="HND", + ) + def run_indexer( self, idx_q: torch.Tensor, @@ -1215,6 +1266,7 @@ def run_indexer( metadata, *, idx_sm_scale: Optional[float] = None, + idx_k_prewritten: bool = False, ) -> torch.Tensor: """Write the index-K cache and return the selected block indices. @@ -1222,6 +1274,8 @@ def run_indexer( forward_args.sparse_backend_args. Returns [total_q, num_kv_heads, topk]. The generation rows are scored by the CuTe DSL kernel and any context rows by the fmha_sm100 proxy pass, over the plan prepare() built. + `idx_k_prewritten` marks that the fused per-layer cache write + (write_layer_caches) already stored this layer's index-K. """ config = self.m3_config idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 @@ -1254,13 +1308,18 @@ def run_indexer( "The MiniMax-M3 BF16 indexer requires BF16 index-Q and a live " f"BF16 index-K tensor; got Q={idx_q_view.dtype}, K={live_k_dtype}." ) - idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim) - metadata.msa_write_idx_k(self.layer_idx, idx_k_view) + # The fused per-layer write (write_layer_caches, signalled by + # idx_k_prewritten) may already have stored this live bf16 index-K + # ahead of the proxy pass; write it here only when it did not. + if not idx_k_prewritten: + idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim) + metadata.msa_write_idx_k(self.layer_idx, idx_k_view) # The FP8 indexer mirrors vLLM's unscaled E4M3 contract: normalized # index Q/K are cast directly and the proxy accumulates their QK scores # in FP32. Block ordering is invariant to the omitted positive scale. # The fused production path arrives here with E4M3 Q and an already - # populated cache; the BF16 path writes its live K above. + # populated cache; the BF16 path writes its live K above unless the + # fused per-layer write already did. # Inputs for the CuTe DSL scorer, which takes this step's generation # span. Left None on a pure-prefill step, which has no span, so the diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index febfe31293e5..52d66018ebb6 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -16,6 +16,7 @@ import enum import os import threading +from contextlib import nullcontext from dataclasses import replace from functools import lru_cache from typing import ClassVar, List, Mapping, Optional, Tuple, Union @@ -34,13 +35,16 @@ from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, DynamicTensorSpec, OptimizationProfile, TunableRunner, - TuningConfig) + TuningConfig, autotune) from ..cublaslt_utils import IS_CUBLASLT_AVAILABLE from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl from .fast_custom_op import fast_custom_op if IS_FLASHINFER_AVAILABLE: + from flashinfer import autotune as _flashinfer_autotune + from flashinfer import mm_mxfp8 as _flashinfer_mm_mxfp8 + from flashinfer import mxfp8_quantize as _flashinfer_mxfp8_quantize from flashinfer.fp4_quantization import nvfp4_quantize as _flashinfer_nvfp4_quantize from ..modules.multi_stream_utils import do_multi_stream @@ -48,7 +52,9 @@ from ..utils import (ActivationType, deep_gemm_gen_tuning_buckets, fp4_scale_infer_shape, get_last_power_of_2_num_tokens_buckets, - is_nvfp4_marlin_supported_sm, last_positive_power_of_2) + get_power_of_2_num_tokens_buckets, + is_nvfp4_marlin_supported_sm, last_positive_power_of_2, + next_positive_power_of_2) if IS_CUTLASS_DSL_AVAILABLE: from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ @@ -57,6 +63,9 @@ # BufferKind is bound from C++; see cpp/tensorrt_llm/thop/outputTensor.h (torch_ext::BufferKind). from tensorrt_llm.bindings.internal.thop import BufferKind +IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE = (IS_FLASHINFER_AVAILABLE + and IS_CUTLASS_DSL_AVAILABLE) + # Used to WAR an issue in torch.bmm that it would break the graph when the out is not contiguous. @torch.library.custom_op("trtllm::bmm_out", mutates_args=("out", )) @@ -567,6 +576,9 @@ def _( _MXFP8_AUTOTUNED_OP = "trtllm::mxfp8_mxfp8_gemm_autotuned::gemm" +_MXFP8_QUANTIZE_AUTOTUNED_OP = "trtllm::mxfp8_quantize_autotuned::quantize" +_FLASHINFER_MXFP8_GEMM_AUTOTUNED_OP = ( + "trtllm::flashinfer_mxfp8_gemm_autotuned::gemm") def _map_to_mxfp8_large_m_bucket(num_tokens: int) -> int: @@ -718,6 +730,122 @@ def _( return act.new_empty((act.size(0), weight.size(0)), dtype=output_dtype) +class MXFP8QuantizeRunner(TunableRunner): + """Profile the native and FlashInfer CuTeDSL MXFP8 activation quantizers.""" + + TRTLLM = -1 # -1 is the AutoTuner fallback tactic. + CUTE_DSL = 0 + + tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_power_of_2_num_tokens_buckets, next_positive_power_of_2), )) + + def __init__(self, input_dtype: torch.dtype) -> None: + self.input_dtype = input_dtype + + def unique_id(self) -> Tuple[torch.dtype]: + return (self.input_dtype, ) + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + return [self.TRTLLM, self.CUTE_DSL] + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = TRTLLM, + ) -> Tuple[torch.Tensor, torch.Tensor]: + activation = inputs[0] + if tactic == self.CUTE_DSL: + return _flashinfer_mxfp8_quantize( + activation, + is_sf_swizzled_layout=True, + alignment=32, + enable_pdl=None, + backend="cute-dsl", + ) + return torch.ops.trtllm.mxfp8_quantize(activation, True) + + +class FlashInferMXFP8GemmRunner(TunableRunner): + """Profile the FlashInfer CUTLASS and CuTeDSL MXFP8 GEMMs.""" + + CUTLASS = -1 # -1 is the AutoTuner fallback tactic. + CUTE_DSL = 0 + + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_power_of_2_num_tokens_buckets, + next_positive_power_of_2), ), + constraint_specs=(ConstraintSpec(1, 0, _mxfp8_scale_infer_shape), ), + ) + + def __init__(self, output_dtype: torch.dtype) -> None: + self.output_dtype = output_dtype + + def unique_id(self) -> Tuple[torch.dtype]: + return (self.output_dtype, ) + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + return [self.CUTLASS, self.CUTE_DSL] + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = CUTLASS, + ) -> torch.Tensor: + act, act_scale, weight, weight_scale = inputs + if tactic == self.CUTLASS: + backend = "cutlass" + context = nullcontext() + else: + # Skip FlashInfer's CuTeDSL config sweep; use its heuristic config. + backend = "cute-dsl" + context = _flashinfer_autotune(tune_mode=False, + skip_ops="mxfp8_gemm") + with context: + return _flashinfer_mm_mxfp8( + act, + weight.t(), + act_scale, + weight_scale, + out_dtype=self.output_dtype, + use_8x4_sf_layout=False, + backend=backend, + ) + + +def _choose_mxfp8_tactic(custom_op: str, runner: TunableRunner, + inputs: List[torch.Tensor], tune: bool) -> int: + """Pick this token bucket's tactic; ``tune`` profiles on a cache miss.""" + with autotune(tune_mode=tune, skip_dynamic_tuning_buckets=True): + _, tactic = AutoTuner.get().choose_one(custom_op, [runner], + runner.tuning_config, inputs) + return tactic + + +def mxfp8_quantize_gemm_autotuned( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool = False, +) -> torch.Tensor: + """MXFP8 quantize + GEMM with each stage's tuned backend for this bucket.""" + quantize_runner = MXFP8QuantizeRunner(input.dtype) + quantize_inputs = [input] + act, act_scale = quantize_runner( + quantize_inputs, + tactic=_choose_mxfp8_tactic(_MXFP8_QUANTIZE_AUTOTUNED_OP, + quantize_runner, quantize_inputs, tune)) + gemm_runner = FlashInferMXFP8GemmRunner(output_dtype) + gemm_inputs = [act, act_scale, weight, weight_scale] + return gemm_runner(gemm_inputs, + tactic=_choose_mxfp8_tactic( + _FLASHINFER_MXFP8_GEMM_AUTOTUNED_OP, gemm_runner, + gemm_inputs, tune)) + + class FP4GemmRunner(TunableRunner): runner_dict = dict() tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index f23a7e1b3c41..ba1a38fa6956 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -982,6 +982,16 @@ def __init__(self, apply_routing=False, moe_backend=model_config.moe_backend, use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm) + # For MIXED_PRECISION, resolve the per-expert quant config (e.g. W4A8_AWQ) + # instead of using the ambiguous global MIXED_PRECISION config. + # For other cases (e.g. nvfp4, unquantized MTP layers), use + # override_quant_config as-is — it already encodes exclusions like MTP. + if (override_quant_config is not None and + override_quant_config.quant_algo == QuantAlgo.MIXED_PRECISION): + expert_quant_config = self._get_experts_quant_config( + model_config, layer_idx) + else: + expert_quant_config = override_quant_config self.experts = create_moe( num_experts=num_experts, routing_method=self.gate.routing_method, @@ -991,17 +1001,13 @@ def __init__(self, reduce_results= False, # In both low‑latency and attention‑DP modes, FusedMoE skips the in‑op all‑reduce. model_config=model_config, - override_quant_config=override_quant_config, + override_quant_config=expert_quant_config, aux_stream_dict=aux_stream_dict, layer_idx=layer_idx, # DS-R1 W4A8 is only supported through custom quantization script from # examples/quantization/quantize_mixed_precision_moe.py - weight_loading_mode=( - MoEWeightLoadingMode.W4A8_CUSTOM - if self._get_experts_quant_config( - model_config, - layer_idx).layer_quant_mode.is_int4_weight_only_per_group() - else MoEWeightLoadingMode.VANILLA), + weight_loading_mode=self._expert_weight_loading_mode( + expert_quant_config), ) self.mapping = model_config.mapping @@ -1096,6 +1102,19 @@ def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: return model_config.quant_config_dict.get( f"model.layers.{layer_idx}.mlp.experts", model_config.quant_config) + @staticmethod + def _expert_weight_loading_mode( + expert_quant_config: Optional[QuantConfig]) -> MoEWeightLoadingMode: + # W4A8_CUSTOM only for the int4-weight-per-group experts produced by + # examples/quantization/quantize_mixed_precision_moe.py; everything else + # (including the unquantized case, where there is no expert quant config) + # uses VANILLA. Guard None: override_quant_config is Optional, so the + # resolved expert config can be None on an unquantized layer. + if (expert_quant_config is not None and expert_quant_config. + layer_quant_mode.is_int4_weight_only_per_group()): + return MoEWeightLoadingMode.W4A8_CUSTOM + return MoEWeightLoadingMode.VANILLA + @staticmethod def _get_shared_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: @@ -1302,13 +1321,14 @@ def __init__(self, "TRTLLM_DEEPSEEK_EAGER_FUSION_DISABLED", "0") == "0" self.enable_fusion &= not self.enable_attention_dp - # FIXME: incompatible with mixed quantization mode quant_config = self._get_decoder_layer_quant_config( model_config, layer_idx) - self.is_nvfp4 = quant_config.layer_quant_mode.has_nvfp4() - assert ( - quant_config.quant_algo - is not QuantAlgo.MIXED_PRECISION), "MIXED_PRECISION is ambiguous" + # For MIXED_PRECISION, the global quant_algo doesn't map to a single + # QuantMode. Per-module configs (e.g. expert W4A8_AWQ vs attention + # FP8_BLOCK_SCALES) are resolved individually where needed, so we + # conservatively set layer-level flags here. + self.is_nvfp4 = (quant_config.quant_algo != QuantAlgo.MIXED_PRECISION + and quant_config.layer_quant_mode.has_nvfp4()) self.allreduce = None self.moe_allreduce = None diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 18e1951e70f2..46df832ac489 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import os from typing import Any, Dict, List, Literal, Optional, Tuple, Union @@ -1134,6 +1148,14 @@ def forward( @register_auto_model("LlamaForCausalLM") class LlamaForCausalLM(SpecDecOneEngineForCausalLM[LlamaModel, LlamaConfig]): + @classmethod + def get_preferred_kv_cache_manager_version( + cls, + pretrained_config: Any = None, + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for Llama.""" + return "V2" + @classmethod def get_preferred_transceiver_runtime( cls, @@ -1505,6 +1527,22 @@ def call_with_text_prompt( class Llama4ForConditionalGeneration(SpecDecOneEngineForCausalLM[Llama4Model, Llama4Config]): + @classmethod + def get_preferred_kv_cache_manager_version( + cls, + pretrained_config: Any = None, + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for Llama4.""" + return "V2" + + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config: Any = None, + ) -> Optional[Literal["CPP", "PYTHON"]]: + """Prefer the Python transceiver for Llama4 NIXL disaggregated serving.""" + return "PYTHON" + def __init__( self, model_config: ModelConfig[Llama4Config], diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 2a79aa5868f5..878e209c95ad 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1403,20 +1403,34 @@ def _msa_attention_core( The backend runs the sparse GQA or dense paged GQA through its inherited FMHA forward; this layer selects the top-k blocks (sparse only) and builds the forward_args the FMHA reads. + + This layer owns the cache write: write_layer_caches stores the + new-token K/V (and, on the bf16 indexer path, index-K) in one launch + before the indexer's proxy pass reads the index-K cache. forward() + then receives k=v=None, which is the backend's contract for "K/V are + already resident", so neither FMHA phase writes them again. """ if self.is_sparse_attention_layer: assert idx_q is not None + # On the FP8 indexer path idx_k is None: the fused producer already + # inserted E4M3 index-K into the side cache, so only K/V are written. + self.attn.write_layer_caches(k, v, idx_k, attn_metadata) # Publish the selected blocks so the FMHA runs the sparse path. - kv_block_indexes = self.attn.run_indexer(idx_q, idx_k, attn_metadata) + # idx_k_prewritten: index-K is already in the cache (written above + # on bf16, or by the FP8 producer), so run_indexer must not write it. + kv_block_indexes = self.attn.run_indexer( + idx_q, idx_k, attn_metadata, idx_k_prewritten=True + ) forward_args = AttentionForwardArgs( output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=kv_block_indexes), ) else: assert idx_q is None and idx_k is None + self.attn.write_layer_caches(k, v, None, attn_metadata) # No top-k selection means the FMHA attends the full page table. forward_args = AttentionForwardArgs(output=output) - self.attn.forward(q, k, v, attn_metadata, forward_args=forward_args) + self.attn.forward(q, None, None, attn_metadata, forward_args=forward_args) return output def _sparse_forward( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 37cd852fe58a..b3e29a427482 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -16,7 +16,8 @@ from torch.nn.parameter import Parameter import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils -from tensorrt_llm._torch.custom_ops.torch_custom_ops import BufferKind +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + BufferKind, mxfp8_quantize_gemm_autotuned) from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm._utils import is_device_integrated, mpi_disabled from tensorrt_llm.bindings import ipc_nvls_supported @@ -3280,7 +3281,8 @@ class MXFP8LinearMethod(LinearMethodBase): - FlashInfer: reuse the CUTLASS-layout activations, weights, and scales with ``mm_mxfp8``. MiniMax-M3 enables this path automatically only while tuning or capturing decode CUDA graphs; eager execution remains - on the native TensorRT-LLM op. + on the native TensorRT-LLM op. With ``tune_decode_graph_backends``, + decode graphs use per-bucket autotuned quantize and GEMM backends. ``TRTLLM_MXFP8_GEMM_BACKEND`` can explicitly select ``trtllm``, ``flashinfer``, or ``auto``. The reference layout is 2D [O,K/32]; both @@ -3310,6 +3312,7 @@ def __init__(self) -> None: f"'flashinfer', or 'auto', got {self.backend!r}") self._flashinfer_mxfp8 = None self._flashinfer_autotuned = False + self.tune_decode_graph_backends = False if self.backend == "flashinfer": self._load_flashinfer(required=True) elif self.backend == "auto" and not self._load_flashinfer( @@ -3379,6 +3382,7 @@ def disable_flashinfer_auto(self) -> None: if self.backend == "auto": self.backend = "trtllm" self._flashinfer_autotuned = False + self.tune_decode_graph_backends = False @classmethod def _swizzled_scale_size(cls, out_features: int, in_features: int) -> int: @@ -3420,43 +3424,55 @@ def apply(self, module: Linear, input: torch.Tensor, input = input.reshape(-1, input.shape[-1]) if self.use_cutlass: - # Dynamic MXFP8 activation quantization (swizzled SF layout), then - # the CUTLASS block-scaled e4m3xe4m3 GEMM. - act_e4m3, act_sf = torch.ops.trtllm.mxfp8_quantize( - input.contiguous(), True) - use_flashinfer = self.backend == "flashinfer" or ( - self.backend == "auto" and - (_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() or - (self._flashinfer_autotuned - and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get()))) - if use_flashinfer: - flashinfer_mxfp8 = self._flashinfer_mxfp8 - assert flashinfer_mxfp8 is not None - output = flashinfer_mxfp8( - act_e4m3, - module.weight.t(), - act_sf, - module.weight_scale, - out_dtype=module.dtype, - use_8x4_sf_layout=False, - backend="cutlass", - ) - else: - # globalScale is the alpha multiplier; pure MXFP8xMXFP8 uses 1.0. - global_scale = torch.ones([1], - dtype=torch.float32, - device=input.device) - gemm = (torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned - if self.needs_native_autotune else - torch.ops.trtllm.mxfp8_mxfp8_gemm) - output = gemm( - act_e4m3, - act_sf, + input = input.contiguous() + if (self.tune_decode_graph_backends + and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get()): + # Tune only in the warmup-only pass (flashinfer_mxfp8_autotune). + output = mxfp8_quantize_gemm_autotuned( + input, module.weight, module.weight_scale, - global_scale, module.dtype, + tune=_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get(), ) + else: + # Dynamic MXFP8 activation quantization (swizzled SF layout), + # then the CUTLASS block-scaled e4m3xe4m3 GEMM. + act_e4m3, act_sf = torch.ops.trtllm.mxfp8_quantize(input, True) + use_flashinfer = self.backend == "flashinfer" or ( + self.backend == "auto" and + (_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() or + (self._flashinfer_autotuned + and _FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVE.get()))) + if use_flashinfer: + flashinfer_mxfp8 = self._flashinfer_mxfp8 + assert flashinfer_mxfp8 is not None + output = flashinfer_mxfp8( + act_e4m3, + module.weight.t(), + act_sf, + module.weight_scale, + out_dtype=module.dtype, + use_8x4_sf_layout=False, + backend="cutlass", + ) + else: + # globalScale is the alpha multiplier; pure MXFP8xMXFP8 + # uses 1.0. + global_scale = torch.ones([1], + dtype=torch.float32, + device=input.device) + gemm = (torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if self.needs_native_autotune else + torch.ops.trtllm.mxfp8_mxfp8_gemm) + output = gemm( + act_e4m3, + act_sf, + module.weight, + module.weight_scale, + global_scale, + module.dtype, + ) if bias is not None: output = output + bias else: diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index 3bdefda1184c..7f4efc789b07 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -469,10 +469,13 @@ def _get_quant_config_dict(self, model_config: ModelConfig) -> Optional[Dict]: Extract quantization configuration from model_config """ - if model_config.quant_config is None: + # Prefer the resolved per-module override (e.g. W4A8_AWQ for experts) + # over the global config which may be MIXED_PRECISION. + quant_config = getattr(self, "_override_quant_config", None) or model_config.quant_config + if quant_config is None: return None - quant_mode = model_config.quant_config.layer_quant_mode + quant_mode = quant_config.layer_quant_mode return { "has_fp8_qdq": quant_mode.has_fp8_qdq() if hasattr(quant_mode, "has_fp8_qdq") diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 0fae89ad1b9b..9fb1b4947240 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2147,7 +2147,8 @@ def _release_megamoe_profiling_scratch(): def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: """Runs forward passes to populate the autotuner cache.""" - from ..custom_ops.torch_custom_ops import MXFP8GemmRunner + from ..custom_ops.torch_custom_ops import ( + IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE, MXFP8GemmRunner) from ..modules.linear import (MXFP8LinearMethod, flashinfer_mxfp8_autotune) @@ -2178,8 +2179,13 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager) -> None: getattr(module, "_use_flashinfer_mxfp8_decode_graph_default", False) for module in self.model.modules())) if use_mxfp8_flashinfer_graph_default: + # CuTeDSL is the alternative backend; PP has no graph-pass handoff. + tune_with_cute_dsl = (IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE + and not self.mapping.has_pp()) for quant_method in mxfp8_methods: quant_method.enable_flashinfer_auto() + quant_method.tune_decode_graph_backends = ( + tune_with_cute_dsl and quant_method.uses_flashinfer) flashinfer_mxfp8_methods = [ method for method in mxfp8_methods if method.needs_flashinfer_autotune diff --git a/tests/AGENTS.md b/tests/AGENTS.md new file mode 100644 index 000000000000..d1b3cfebd44f --- /dev/null +++ b/tests/AGENTS.md @@ -0,0 +1,28 @@ +# AGENTS.md — tests + +Scope: this directory and everything under it. Supplements the repo-root `AGENTS.md`. + +## Whether a new test runs in CI + +This is a reminder for whoever adds a test, not a hard gate. Not every test runs +in CI — a test is a CI test, a QA test, or a local developer test: + +- If a newly added test needs to be protected by CI, it must be included in at + least one CI test list under `tests/integration/test_lists/test-db/` (e.g. + `l0_cpu.yml`, `l0_b200.yml`, `l0_dgx_b200.yml`). A CI test that no list + references — directly or through a listed parent directory — is never + collected: it reports nothing, so a regression in it merges unnoticed, a + silent false green. +- A test that is not included in any CI test list is treated as intended for + local developer testing or QA rather than CI. + +So when you add a test, confirm with the author/reviewer whether it should run +in CI and, if so, which platform's test list(s) should carry it. Leaving a new +test out of every CI list is fine for QA-only or local-only tests — but make it +a deliberate, confirmed decision rather than a forgotten step. + +Choose the list by where a CI test runs: CPU-only tests go in `l0_cpu.yml`; +single-GPU tests in the per-GPU list for their target (e.g. `l0_b200.yml`); +multi-GPU tests in a multi-GPU list (e.g. `l0_dgx_b200.yml`, +`l0_gb300_multi_gpus.yml`). If you are unsure which list(s) a test belongs in, +ask the author/reviewer rather than leaving it unlisted. diff --git a/tests/CLAUDE.md b/tests/CLAUDE.md new file mode 100644 index 000000000000..40c6cf85d526 --- /dev/null +++ b/tests/CLAUDE.md @@ -0,0 +1,2 @@ +# In ./CLAUDE.md +@AGENTS.md diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py index 1de19eddc13a..146e547ad4fa 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py @@ -39,11 +39,11 @@ _lpips_model_path, _lpips_pinned_fp32_matmul_precision, _preserve_lpips_candidate_on_failure, + _require_exists, _run_lpips_eval, _run_reusable_image_lpips_eval, _run_single_device_feature_generator, _save_lpips_video_mp4, - _skip_if_missing, _validate_single_feature_config, ) @@ -175,7 +175,7 @@ def _run_cosmos3_lpips_pipeline(num_frames, video=None): ) model_path = _lpips_model_path(COSMOS3_NANO_MODEL_SUBPATH) - _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Nano checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() args = VisualGenArgs( model=model_path, @@ -234,7 +234,7 @@ def _generate_cosmos3_lpips_video(output_path): def _cosmos3_v2v_lpips_reference_bytes(): - _skip_if_missing(COSMOS3_LPIPS_V2V_REFERENCE_MP4, "Cosmos3 V2V LPIPS reference fixture") + _require_exists(COSMOS3_LPIPS_V2V_REFERENCE_MP4, "Cosmos3 V2V LPIPS reference fixture") with open(COSMOS3_LPIPS_V2V_REFERENCE_MP4, "rb") as f: return f.read() @@ -271,7 +271,7 @@ def _generate_cosmos3_feature_image(case, output_path): pipeline = None try: model_path = _lpips_model_path(COSMOS3_NANO_MODEL_SUBPATH) - _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Nano checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() # Pin fp32-matmul arithmetic only for the profiles whose goldens are # re-baselined under it. NVFP4's golden is waived (nvbugs/6572800) and @@ -435,7 +435,7 @@ def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): the shared FP8 dynamic-quant config. """ model_path = _lpips_model_path("Cosmos3-Nano") - _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Nano checkpoint", is_dir=True) out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_example") os.makedirs(out_dir, exist_ok=True) @@ -477,7 +477,7 @@ def test_cosmos3_t2i_4step_example(_visual_gen_deps, llm_root, llm_venv): checkpoint's fixed distilled schedule; the run must produce an image. """ model_path = _lpips_model_path("Cosmos3-Super-Text2Image-4Step") - _skip_if_missing(model_path, "Cosmos3-Super-Text2Image-4Step checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Super-Text2Image-4Step checkpoint", is_dir=True) out_dir = os.path.join( llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_t2i_4step_example" @@ -545,7 +545,7 @@ def test_cosmos3_i2v_4step_example(_visual_gen_deps, llm_root, llm_venv): must produce a video. """ model_path = _lpips_model_path("Cosmos3-Super-Image2Video-4Step") - _skip_if_missing(model_path, "Cosmos3-Super-Image2Video-4Step checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Super-Image2Video-4Step checkpoint", is_dir=True) out_dir = os.path.join( llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_i2v_4step_example" @@ -601,7 +601,7 @@ def _run_cosmos3_i2v_4step_lpips_pipeline(image_path): ) model_path = _lpips_model_path(COSMOS3_I2V_4STEP_MODEL_SUBPATH) - _skip_if_missing(model_path, "Cosmos3-Super-Image2Video-4Step checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Super-Image2Video-4Step checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() args = VisualGenArgs( model=model_path, @@ -721,7 +721,7 @@ def test_cosmos3_edge_i2v_example(_visual_gen_deps, llm_root, llm_venv): video. """ model_path = _lpips_model_path("Cosmos3-Edge") - _skip_if_missing(model_path, "Cosmos3-Edge checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Edge checkpoint", is_dir=True) out_dir = os.path.join( llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_edge_i2v_example" @@ -775,7 +775,7 @@ def test_cosmos3_edge_policy_droid_example(_visual_gen_deps, llm_root, llm_venv) from safetensors.torch import load_file model_path = _lpips_model_path("Cosmos3-Edge-Policy-DROID") - _skip_if_missing(model_path, "Cosmos3-Edge-Policy-DROID checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Edge-Policy-DROID checkpoint", is_dir=True) out_dir = os.path.join( llm_venv.get_working_directory(), @@ -878,7 +878,7 @@ def _run_cosmos3_edge_lpips_pipeline(**forward_kwargs): ) model_path = _lpips_model_path("Cosmos3-Edge") - _skip_if_missing(model_path, "Cosmos3-Edge checkpoint", is_dir=True) + _require_exists(model_path, "Cosmos3-Edge checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() args = VisualGenArgs( model=model_path, diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_flux.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_flux.py index 5eb7cc48a46e..e7f643ce23c0 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_flux.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_flux.py @@ -33,14 +33,13 @@ _fixed_nvfp4_quantization_backend, _golden_media_path, _lpips_deterministic_algorithms, - _lpips_model_path, _preserve_lpips_candidate_on_failure, _run_lpips_eval, _run_reusable_image_lpips_eval, _run_single_device_feature_generator, - _skip_if_missing, _validate_single_feature_config, ) +from test_common.llm_data import get_checkpoint FLUX_LPIPS_PROMPT = "a tiny astronaut hatching from an egg on the moon" FLUX_LPIPS_HEIGHT = 256 @@ -134,7 +133,6 @@ def _generate_flux_lpips_image(model_path, output_path): from tensorrt_llm.media.encoding import save_image from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs - _skip_if_missing(model_path, "FLUX checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() with _lpips_deterministic_algorithms(): args = VisualGenArgs( @@ -163,8 +161,7 @@ def _generate_flux_image(case, output_path): from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.media.encoding import save_image - model_path = _lpips_model_path(case.checkpoint_subdir) - _skip_if_missing(model_path, f"{case.checkpoint_subdir} checkpoint", is_dir=True) + model_path = get_checkpoint(case.checkpoint_subdir) _disable_inductor_compile_worker_quiesce() pipeline = None with _lpips_deterministic_algorithms(), _fixed_nvfp4_quantization_backend(case.features): @@ -210,7 +207,7 @@ def test_flux1_lpips_against_golden(tmp_path): golden_path = _golden_media_path( tmp_path, "flux1_lpips_golden.png", "FLUX.1 LPIPS golden image" ) - _generate_flux_lpips_image(_lpips_model_path("FLUX.1-dev"), generated_path) + _generate_flux_lpips_image(get_checkpoint("FLUX.1-dev"), generated_path) score = _run_lpips_eval( tmp_path, "flux1", @@ -228,7 +225,7 @@ def test_flux2_lpips_against_golden(tmp_path): golden_path = _golden_media_path( tmp_path, "flux2_lpips_golden.png", "FLUX.2 LPIPS golden image" ) - _generate_flux_lpips_image(_lpips_model_path("FLUX.2-dev"), generated_path) + _generate_flux_lpips_image(get_checkpoint("FLUX.2-dev"), generated_path) score = _run_lpips_eval( tmp_path, "flux2", @@ -269,8 +266,7 @@ def test_flux_accuracy_against_golden(request, tmp_path, case, _visual_gen_lpips def test_flux1_example(_visual_gen_deps, llm_root, llm_venv): """Run the FLUX.1 example with the supported single-GPU NVFP4 config.""" - model_path = _lpips_model_path("FLUX.1-dev") - _skip_if_missing(model_path, "FLUX.1-dev checkpoint", is_dir=True) + model_path = get_checkpoint("FLUX.1-dev") out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "flux1_example") os.makedirs(out_dir, exist_ok=True) @@ -299,8 +295,7 @@ def test_flux1_example(_visual_gen_deps, llm_root, llm_venv): def test_flux2_example(_visual_gen_deps, llm_root, llm_venv): """Run the FLUX.2 example with the supported single-GPU NVFP4 config.""" - model_path = _lpips_model_path("FLUX.2-dev") - _skip_if_missing(model_path, "FLUX.2-dev checkpoint", is_dir=True) + model_path = get_checkpoint("FLUX.2-dev") out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "flux2_example") os.makedirs(out_dir, exist_ok=True) @@ -329,8 +324,7 @@ def test_flux2_example(_visual_gen_deps, llm_root, llm_venv): def test_flux2_reference_image_example(_visual_gen_deps, llm_root, llm_venv, tmp_path): """Run the FLUX.2 example with the existing reference-image request argument.""" - model_path = _lpips_model_path("FLUX.2-dev") - _skip_if_missing(model_path, "FLUX.2-dev checkpoint", is_dir=True) + model_path = get_checkpoint("FLUX.2-dev") reference_path = _golden_media_path( tmp_path, "flux2_lpips_golden.png", "FLUX.2 reference image" ) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py index 87d758dff00e..263fb89e4022 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py @@ -33,10 +33,10 @@ _lpips_deterministic_algorithms, _lpips_model_path, _preserve_lpips_candidate_on_failure, + _require_exists, _run_lpips_eval, _run_reusable_image_lpips_eval, _run_single_device_feature_generator, - _skip_if_missing, _validate_single_feature_config, ) @@ -117,7 +117,7 @@ def _generate_glm_image_lpips_image(model_path, output_path): from tensorrt_llm.media.encoding import save_image from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs - _skip_if_missing(model_path, "GLM-Image checkpoint", is_dir=True) + _require_exists(model_path, "GLM-Image checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() with _lpips_deterministic_algorithms(): args = VisualGenArgs( @@ -147,7 +147,7 @@ def _generate_glm_image_feature_image(case, output_path): from tensorrt_llm.media.encoding import save_image model_path = _lpips_model_path(case.checkpoint_subdir) - _skip_if_missing(model_path, f"{case.checkpoint_subdir} checkpoint", is_dir=True) + _require_exists(model_path, f"{case.checkpoint_subdir} checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() pipeline = None with _lpips_deterministic_algorithms(), _fixed_nvfp4_quantization_backend(case.features): diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_hunyuan.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_hunyuan.py index eb49f2919a19..fb6d62e320d3 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_hunyuan.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_hunyuan.py @@ -26,12 +26,11 @@ _disable_inductor_compile_worker_quiesce, _golden_media_path, _lpips_deterministic_algorithms, - _lpips_model_path, _preserve_lpips_candidate_on_failure, _run_lpips_eval, _save_lpips_video_mp4, - _skip_if_missing, ) +from test_common.llm_data import get_checkpoint HUNYUAN_T2V_MODEL_SUBPATH = "HunyuanVideo-1.5-Diffusers-480p_t2v" @@ -65,8 +64,7 @@ def _run_hunyuan_lpips_pipeline(): VisualGenArgs, ) - model_path = _lpips_model_path(HUNYUAN_T2V_MODEL_SUBPATH) - _skip_if_missing(model_path, "HunyuanVideo 1.5 480p T2V checkpoint", is_dir=True) + model_path = get_checkpoint(HUNYUAN_T2V_MODEL_SUBPATH) _disable_inductor_compile_worker_quiesce() with _lpips_deterministic_algorithms(): args = VisualGenArgs( @@ -136,8 +134,7 @@ def test_hunyuan_t2v_example(_visual_gen_deps, llm_root, llm_venv): ``configs/hunyuan-t2v-fp8-1gpu.yaml`` work together as documented in the README, at the example's own 480p defaults. """ - model_path = _lpips_model_path(HUNYUAN_T2V_MODEL_SUBPATH) - _skip_if_missing(model_path, "HunyuanVideo 1.5 480p T2V checkpoint", is_dir=True) + model_path = get_checkpoint(HUNYUAN_T2V_MODEL_SUBPATH) out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "hunyuan_example") os.makedirs(out_dir, exist_ok=True) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py index d3c24615dbc9..24a616fb9747 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_ltx2.py @@ -35,16 +35,15 @@ _fixed_nvfp4_quantization_backend, _golden_media_path, _lpips_deterministic_algorithms, - _lpips_model_path, _preserve_lpips_candidate_on_failure, _run_lpips_eval, _run_reusable_video_lpips_eval, _run_single_device_feature_generator, _save_lpips_video_mp4, - _skip_if_missing, _validate_single_feature_config, _visual_gen_output_path, ) +from test_common.llm_data import get_checkpoint LTX2_LPIPS_NUM_FRAMES = 49 LTX2_LPIPS_NUM_INFERENCE_STEPS = 8 @@ -128,20 +127,19 @@ def _ltx2_lpips_text_encoder_path(): for candidate in candidates: if os.path.isdir(candidate): return candidate - return candidates[0] + raise FileNotFoundError( + f"LTX-2 text encoder not found under any of: {candidates}. " + f"Stage '{LTX2_TEXT_ENCODER_SUBPATH}' under LLM_MODELS_ROOT to run this test." + ) def _generate_ltx2_feature_video(case, output_path): from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - checkpoint_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") + checkpoint_path = get_checkpoint("LTX-2/ltx-2-19b-dev.safetensors") text_encoder_path = _ltx2_lpips_text_encoder_path() - spatial_upsampler_path = _lpips_model_path("LTX-2", "ltx-2-spatial-upscaler-x2-1.0.safetensors") - distilled_lora_path = _lpips_model_path("LTX-2", "ltx-2-19b-distilled-lora-384.safetensors") - _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") - _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) - _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") - _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") + spatial_upsampler_path = get_checkpoint("LTX-2/ltx-2-spatial-upscaler-x2-1.0.safetensors") + distilled_lora_path = get_checkpoint("LTX-2/ltx-2-19b-distilled-lora-384.safetensors") _disable_inductor_compile_worker_quiesce() pipeline = None with ( @@ -231,14 +229,10 @@ def _generate_ltx2_lpips_video(output_path, *, enable_cuda_graph=False): from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import CudaGraphConfig, TorchCompileConfig, VisualGenArgs - checkpoint_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") + checkpoint_path = get_checkpoint("LTX-2/ltx-2-19b-dev.safetensors") text_encoder_path = _ltx2_lpips_text_encoder_path() - spatial_upsampler_path = _lpips_model_path("LTX-2", "ltx-2-spatial-upscaler-x2-1.0.safetensors") - distilled_lora_path = _lpips_model_path("LTX-2", "ltx-2-19b-distilled-lora-384.safetensors") - _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") - _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) - _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") - _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") + spatial_upsampler_path = get_checkpoint("LTX-2/ltx-2-spatial-upscaler-x2-1.0.safetensors") + distilled_lora_path = get_checkpoint("LTX-2/ltx-2-19b-distilled-lora-384.safetensors") _disable_inductor_compile_worker_quiesce() # TorchCompileConfig(enable=False) does not suppress nested @torch.compile decorators. @@ -286,15 +280,10 @@ def _generate_ltx2_cuda_graph_trtllm_backend_video(output_path): TorchCompileConfig, ) - scratch_space = conftest.llm_models_root() - checkpoint_path = os.path.join(scratch_space, LTX2_MODEL_CHECKPOINT_PATH) + checkpoint_path = get_checkpoint(LTX2_MODEL_CHECKPOINT_PATH) text_encoder_path = _ltx2_lpips_text_encoder_path() - spatial_upsampler_path = os.path.join(scratch_space, LTX2_UPSAMPLER_SUBPATH) - distilled_lora_path = os.path.join(scratch_space, LTX2_DISTILLED_LORA_SUBPATH) - _skip_if_missing(checkpoint_path, "LTX-2 checkpoint") - _skip_if_missing(text_encoder_path, "LTX-2 text encoder", is_dir=True) - _skip_if_missing(spatial_upsampler_path, "LTX-2 spatial upsampler") - _skip_if_missing(distilled_lora_path, "LTX-2 distilled LoRA") + spatial_upsampler_path = get_checkpoint(LTX2_UPSAMPLER_SUBPATH) + distilled_lora_path = get_checkpoint(LTX2_DISTILLED_LORA_SUBPATH) _disable_inductor_compile_worker_quiesce() visual_gen_args = VisualGenArgs( @@ -435,10 +424,8 @@ def test_ltx2_example(_visual_gen_deps, llm_root, llm_venv): ``--text_encoder_path`` because the shared YAML intentionally omits it to keep the config model-path-agnostic. """ - model_path = _lpips_model_path("LTX-2", "ltx-2-19b-dev.safetensors") - _skip_if_missing(model_path, "LTX-2 checkpoint") + model_path = get_checkpoint("LTX-2/ltx-2-19b-dev.safetensors") text_encoder_path = _ltx2_lpips_text_encoder_path() - _skip_if_missing(text_encoder_path, "LTX-2 text encoder (gemma-3-12b-it)", is_dir=True) out_dir = os.path.join(llm_venv.get_working_directory(), "visual_gen_output", "ltx2_example") os.makedirs(out_dir, exist_ok=True) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 0fbed4eb379a..b5a1c211665f 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -39,8 +39,8 @@ _run_lpips_eval, _run_wan_lpips_pipeline, _save_lpips_video_mp4, - _skip_if_missing, ) +from test_common.llm_data import get_checkpoint def _parallel_config(**kwargs): @@ -279,8 +279,7 @@ def wan22_within_build_reference(tmp_path_factory): pytest.skip("Required modules not available") if torch.cuda.device_count() < 1: pytest.skip("Within-build reference generation requires a GPU") - model_path = _lpips_model_path("Wan2.2-T2V-A14B-Diffusers") - _skip_if_missing(model_path, "Wan checkpoint", is_dir=True) + model_path = get_checkpoint("Wan2.2-T2V-A14B-Diffusers") tllm_site = _validated_tllm_site( os.path.dirname(os.path.dirname(os.path.abspath(tllm_bindings.__file__))) ) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_qwen_image.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_qwen_image.py index 056ae743d218..59a14e675787 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_qwen_image.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_qwen_image.py @@ -23,7 +23,6 @@ import pytest import torch -from defs import conftest from defs.common import venv_check_call from defs.examples.visual_gen.visual_gen_test_utils import ( FeatureConfigState, @@ -38,14 +37,13 @@ _fixed_nvfp4_quantization_backend, _golden_media_path, _lpips_deterministic_algorithms, - _lpips_model_path, _preserve_lpips_candidate_on_failure, _run_lpips_eval, _run_reusable_image_lpips_eval, _run_single_device_feature_generator, - _skip_if_missing, _validate_single_feature_config, ) +from test_common.llm_data import get_checkpoint # QwenImage (text-to-image) — default-setting LPIPS golden. # Params mirror the QwenImage 20B reference defaults (pipeline_qwen_image.py). @@ -121,7 +119,6 @@ def _generate_qwenimage_lpips_image(model_path, output_path, *, enable_cuda_grap from tensorrt_llm.media.encoding import save_image from tensorrt_llm.visual_gen.args import CudaGraphConfig, TorchCompileConfig, VisualGenArgs - _skip_if_missing(model_path, "QwenImage checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() args = VisualGenArgs( model=model_path, @@ -210,7 +207,6 @@ def _generate_qwen_image_layered_lpips_image(model_path, input_path, output_path from tensorrt_llm.media.encoding import save_image from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs - _skip_if_missing(model_path, "Qwen-Image-Layered checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() args = VisualGenArgs( model=model_path, @@ -246,8 +242,7 @@ def _generate_qwenimage_feature_image(case, output_path): from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.media.encoding import save_image - model_path = _lpips_model_path(QWEN_IMAGE_MODEL_SUBPATH) - _skip_if_missing(model_path, "QwenImage checkpoint", is_dir=True) + model_path = get_checkpoint(QWEN_IMAGE_MODEL_SUBPATH) _disable_inductor_compile_worker_quiesce() pipeline = None with _lpips_deterministic_algorithms(), _fixed_nvfp4_quantization_backend(case.features): @@ -327,7 +322,7 @@ def test_qwenimage_lpips_against_golden(_visual_gen_deps, tmp_path): golden_path = _golden_media_path( tmp_path, "qwenimage_lpips_golden.png", "QwenImage LPIPS golden image" ) - _generate_qwenimage_lpips_image(_lpips_model_path(QWEN_IMAGE_MODEL_SUBPATH), generated_path) + _generate_qwenimage_lpips_image(get_checkpoint(QWEN_IMAGE_MODEL_SUBPATH), generated_path) score = _run_lpips_eval( tmp_path, "qwenimage", @@ -349,7 +344,7 @@ def test_qwen_image_layered_lpips_against_golden(tmp_path): _copy_qwen_image_layered_lpips_input(tmp_path, input_path) _write_qwen_image_layered_lpips_golden_grid(tmp_path, golden_path) _generate_qwen_image_layered_lpips_image( - _lpips_model_path(QWEN_IMAGE_LAYERED_MODEL_SUBPATH), + get_checkpoint(QWEN_IMAGE_LAYERED_MODEL_SUBPATH), input_path, generated_path, ) @@ -375,14 +370,11 @@ def test_qwen_image_example(_visual_gen_deps, llm_root, llm_venv): ``configs/qwen-image-fp8-1gpu.yaml`` work together as documented. Uses the local Qwen-Image checkpoint and the shared FP8 blockwise dynamic-quant config. """ - scratch_space = conftest.llm_models_root() - model_path = os.path.join(scratch_space, QWEN_IMAGE_MODEL_SUBPATH) - _skip_if_missing(model_path, "Qwen-Image checkpoint", is_dir=True) + model_path = get_checkpoint(QWEN_IMAGE_MODEL_SUBPATH) model_index_path = os.path.join(model_path, "model_index.json") - if not os.path.isfile(model_index_path): - pytest.skip( - f"Qwen-Image checkpoint is incomplete: {model_path} (missing {model_index_path})" - ) + assert os.path.isfile(model_index_path), ( + f"Qwen-Image checkpoint is incomplete: {model_path} (missing {model_index_path})" + ) out_dir = os.path.join( llm_venv.get_working_directory(), "visual_gen_output", "qwen_image_example" @@ -414,15 +406,11 @@ def test_qwen_image_example(_visual_gen_deps, llm_root, llm_venv): def test_qwen_image_layered_example(_visual_gen_deps, tmp_path, llm_root, llm_venv): """Run examples/visual_gen/models/qwen_image_layered.py end-to-end.""" - scratch_space = conftest.llm_models_root() - model_path = os.path.join(scratch_space, QWEN_IMAGE_LAYERED_MODEL_SUBPATH) - _skip_if_missing(model_path, "Qwen-Image-Layered checkpoint", is_dir=True) + model_path = get_checkpoint(QWEN_IMAGE_LAYERED_MODEL_SUBPATH) model_index_path = os.path.join(model_path, "model_index.json") - if not os.path.isfile(model_index_path): - pytest.skip( - f"Qwen-Image-Layered checkpoint is incomplete: {model_path} " - f"(missing {model_index_path})" - ) + assert os.path.isfile(model_index_path), ( + f"Qwen-Image-Layered checkpoint is incomplete: {model_path} (missing {model_index_path})" + ) input_path = tmp_path / "qwen_image_layered_input.png" _copy_qwen_image_layered_lpips_input(tmp_path, input_path) @@ -467,16 +455,11 @@ def test_qwen_image_edit_example(_visual_gen_deps: Any, llm_root: str, llm_venv: Validates that the Qwen-Image-Edit example script and ``configs/qwen-image-edit-2511-fp8-1gpu.yaml`` work together as documented. """ - model_path = os.environ.get("QWEN_IMAGE_EDIT_MODEL_PATH") or os.path.join( - conftest.llm_models_root(), QWEN_IMAGE_EDIT_MODEL_SUBPATH - ) - _skip_if_missing(model_path, "Qwen-Image-Edit-2511 checkpoint", is_dir=True) + model_path = get_checkpoint(QWEN_IMAGE_EDIT_MODEL_SUBPATH) model_index_path = os.path.join(model_path, "model_index.json") - if not os.path.isfile(model_index_path): - pytest.skip( - f"Qwen-Image-Edit-2511 checkpoint is incomplete: {model_path} " - f"(missing {model_index_path})" - ) + assert os.path.isfile(model_index_path), ( + f"Qwen-Image-Edit-2511 checkpoint is incomplete: {model_path} (missing {model_index_path})" + ) out_dir = os.path.join( llm_venv.get_working_directory(), "visual_gen_output", "qwen_image_edit_example" diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py index 394312fefd79..0b5271c66a63 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py @@ -70,10 +70,10 @@ _run_reusable_video_lpips_eval, _run_single_device_feature_generator, _save_lpips_video_mp4, - _skip_if_missing, _validate_single_feature_config, _visual_gen_output_path, ) +from test_common.llm_data import get_checkpoint WAN_T2V_MODEL_SUBPATH = "Wan2.1-T2V-1.3B-Diffusers" WAN22_T2V_MODEL_SUBPATH = "Wan2.2-T2V-A14B-Diffusers" @@ -331,8 +331,7 @@ def test_fastwan_lpips_against_golden(request, tmp_path, fastwan_video_path): def _generate_wan_feature_video(case, output_path): from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - model_path = _lpips_model_path(case.checkpoint_subdir) - _skip_if_missing(model_path, f"{case.checkpoint_subdir} checkpoint", is_dir=True) + model_path = get_checkpoint(case.checkpoint_subdir) _disable_inductor_compile_worker_quiesce() pipeline = None with ( diff --git a/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py index 3f415a0dcb78..43681b608d08 100644 --- a/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py +++ b/tests/integration/defs/examples/visual_gen/visual_gen_test_utils.py @@ -355,14 +355,14 @@ def _lpips_model_path(*parts): return os.path.join(_llm_models_root(), *parts) -def _skip_if_missing(path, label, is_dir=False): +def _require_exists(path, label, is_dir=False): exists = os.path.isdir(path) if is_dir else os.path.exists(path) if not exists: - pytest.skip(f"{label} not found: {path}") + raise FileNotFoundError(f"{label} not found: {path}") def _extract_visual_gen_lpips_golden_media(tmp_path): - _skip_if_missing(VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP, "VisualGen LPIPS golden media zip") + _require_exists(VISUAL_GEN_LPIPS_GOLDEN_MEDIA_ZIP, "VisualGen LPIPS golden media zip") extract_dir = tmp_path / "visual_gen_lpips_golden_media" if extract_dir.exists(): return extract_dir @@ -377,7 +377,7 @@ def _extract_visual_gen_lpips_golden_media(tmp_path): def _golden_media_path(tmp_path, media_name, label): path = _extract_visual_gen_lpips_golden_media(tmp_path) / media_name - _skip_if_missing(path, label) + _require_exists(path, label) return path @@ -745,7 +745,7 @@ def _run_wan_lpips_pipeline( from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs - _skip_if_missing(model_path, "Wan checkpoint", is_dir=True) + _require_exists(model_path, "Wan checkpoint", is_dir=True) _disable_inductor_compile_worker_quiesce() args_kwargs = dict( model=model_path, diff --git a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py index 6319aa85976e..cb3826957661 100644 --- a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py +++ b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py @@ -73,6 +73,13 @@ "iterIntraDeviceCopyBytes", ] +SECONDARY_FIELDS = { + "secondaryMaxNumBlocks", + "secondaryFreeNumBlocks", + "secondaryUsedNumBlocks", +} +NON_SECONDARY_FIELDS = set(ALL_FIELDS) - SECONDARY_FIELDS + TEST_NAMES = { 1: "Cold start", 2: "Partial block reuse", @@ -383,26 +390,34 @@ def test_rapid_fire(self, llm_instance, all_collected, request): assert total_alloc > 0, "iterAllocTotalBlocks = 0 across all entries" def test_field_completeness(self, llm_instance, all_collected, request): - """Field completeness — verify all 18 fields present across all collected stats.""" + """Field completeness — verify fields in their V2 window and cold-pool views.""" # If running standalone (no prior tests), generate some traffic if not all_collected: llm_instance.generate(["Hello world"], SamplingParams(max_tokens=16)) collect_stats(llm_instance, all_collected) entries_with_kv = 0 - missing_fields = set() for s in all_collected: ki = s.get("kvCacheIterationStats") if ki: entries_with_kv += 1 + # V2 reports secondary gauges by cold pool group, not by window. for ws, v in ki.items(): - for field in ALL_FIELDS: - if field not in v: - missing_fields.add(field) + missing_fields = NON_SECONDARY_FIELDS - v.keys() + assert not missing_fields, ( + f"Missing kvCacheIterationStats fields for window {ws}: " + f"{sorted(missing_fields)}" + ) + + for group, v in s.get("kvCacheIterationStatsByColdPoolGroup", {}).items(): + missing_fields = SECONDARY_FIELDS - v.keys() + assert not missing_fields, ( + f"Missing kvCacheIterationStatsByColdPoolGroup fields for group {group}: " + f"{sorted(missing_fields)}" + ) print(f" Entries with kvCacheIterationStats: {entries_with_kv}/{len(all_collected)}") assert entries_with_kv > 0, "no entries contain kvCacheIterationStats" - assert len(missing_fields) == 0, f"Missing fields: {sorted(missing_fields)}" # --------------------------------------------------------------------------- diff --git a/tests/integration/defs/perf/README_release_test.md b/tests/integration/defs/perf/README_release_test.md index 13aa38c5a6fc..ca094112159a 100644 --- a/tests/integration/defs/perf/README_release_test.md +++ b/tests/integration/defs/perf/README_release_test.md @@ -66,12 +66,7 @@ benchmark_cmd = [ #### Backend Selection ```python -if self._config.backend != "pytorch": - benchmark_cmd += [ - f"--backend=tensorrt", f"--engine_dir={engine_dir}" - ] -else: - benchmark_cmd += ["--backend=pytorch"] +benchmark_cmd += [f"--backend={self._config.backend}"] ``` #### Optional Parameter Configuration diff --git a/tests/integration/defs/perf/pytorch_model_config.py b/tests/integration/defs/perf/pytorch_model_config.py index 87c25d1b2cff..5e05e983517a 100644 --- a/tests/integration/defs/perf/pytorch_model_config.py +++ b/tests/integration/defs/perf/pytorch_model_config.py @@ -14,7 +14,7 @@ # limitations under the License. # -*- coding: utf-8 -*- """ -Model pytorch/TRT yaml config for trtllm-bench perf tests +Model PyTorch YAML config for trtllm-bench performance tests. """ from ..conftest import llm_models_root @@ -61,17 +61,12 @@ def get_model_yaml_config(model_label: str, Returns: dict: yaml config """ - if 'pytorch' in model_label: - # Pytorch backend config - base_config = { - 'print_iter_log': True, - 'cuda_graph_config': { - 'enable_padding': True, - }, - } - else: - # TRT backend config - base_config = {} + base_config = { + 'print_iter_log': True, + 'cuda_graph_config': { + 'enable_padding': True, + }, + } if 'kv_cache_dtype' in model_label: base_config.update({ diff --git a/tests/integration/defs/perf/test_perf.py b/tests/integration/defs/perf/test_perf.py index e2eeebd6211f..9fd4708e3928 100644 --- a/tests/integration/defs/perf/test_perf.py +++ b/tests/integration/defs/perf/test_perf.py @@ -431,8 +431,7 @@ def __init__( runtime: str = "python", api: str = "", streaming: str = "", - backend: str = "", - mode: str = "plugin", + backend: str = "pytorch", data_type: str = "float16", max_batch_size: int = 512, max_num_tokens: int = 2048, @@ -451,8 +450,6 @@ def __init__( tp_size: int = 1, pp_size: int = 1, num_gpus: int = 1, - # only for torch-backend currently - extra: bool = False, moe_backend: str | None = None, # _autodeploy backend specific parameters ad_compile_backend: str = "torch-cudagraph", @@ -461,16 +458,14 @@ def __init__( ): # The model name. self.model_name = model_name - # Python, cpp, bench, or serve runtime. + # Python, bench, or serve runtime. self.runtime = runtime # API Type: only executor is allowed self.api = api - # Backend Type: pytorch or cpp + # Backend type: pytorch or _autodeploy self.backend = backend # Streaming responses self.streaming = streaming - # Plugin or OOTB mode. - self.mode = mode # Activation dtype. self.data_type = data_type # Percentage of weights that resides on GPU. @@ -509,15 +504,11 @@ def __init__( self.pp_size = pp_size # Number of GPUs. self.num_gpus = num_gpus - # Extra flag to enable pytorch_model_config reading for TRT backend - self.extra = extra self.moe_backend = moe_backend # _autodeploy backend specific parameters self.ad_compile_backend = ad_compile_backend self.extra_runtime = extra_runtime self.skip_loading_weights = skip_loading_weights - # Just build engines - self.build_only = False def to_string(self, custom_server_name: str = None, @@ -534,26 +525,22 @@ def to_string(self, if device_subtype: entries.append(f"subtype:{device_subtype}") - if self.runtime == "cpp": # bertBenchmark runtime - entries.append(f"cpp") - elif self.runtime == "serve": - entries.append(f"serve") + if self.runtime == "serve": + entries.append("serve") if self.backend == 'pytorch': - entries.append(f"pytorch") + entries.append("pytorch") if self.streaming == "streaming": - entries.append(f"streaming") - elif self.runtime == "bench": # trtllm-bench - entries.append(f"bench") + entries.append("streaming") + elif self.runtime == "bench": + entries.append("bench") if self.backend == 'pytorch': - entries.append(f"pytorch") + entries.append("pytorch") elif self.backend == '_autodeploy': - entries.append(f"_autodeploy") + entries.append("_autodeploy") if self.streaming == "streaming": - entries.append(f"streaming") + entries.append("streaming") - # Add mode and dtype. - if self.runtime not in ("bench", "serve"): - entries.append(self.mode) + # Add dtype. entries.append(self.data_type) if self.gpu_weights_percent != -1: @@ -572,9 +559,6 @@ def to_string(self, if self.kv_cache_free_gpu_mem_fraction != 0.9: entries.append(f"kv_frac:{self.kv_cache_free_gpu_mem_fraction}") - if self.build_only: - entries.append(f"build_only") - if self.batch_sizes[0] > 0: # Add batch size(s). if custom_bs is None: @@ -642,10 +626,6 @@ def to_string(self, if self.num_gpus > 1: entries.append(f"gpus:{self.num_gpus}") - # Add extra flag for llm-api-config.yml. - if self.extra: - entries.append("extra") - if self.moe_backend is not None: entries.append(f"moe:{self.moe_backend}") @@ -676,10 +656,8 @@ def load_from_str(self, test_param_labels) -> None: self.api = labels.pop(0) if labels[0] == "exe" else "" self.backend = labels.pop(0) if labels[0] in ["pytorch", "_autodeploy" - ] else "" + ] else "pytorch" self.streaming = labels.pop(0) if labels[0] == "streaming" else "" - if self.runtime not in ("bench", "serve"): - self.mode = labels.pop(0) self.data_type = labels.pop(0) if labels[0].startswith("gwp"): self.gpu_weights_percent = float(labels.pop(0).replace("gwp:", "")) @@ -698,32 +676,26 @@ def load_from_str(self, test_param_labels) -> None: self.kv_cache_free_gpu_mem_fraction = float( labels.pop(0).replace("kv_frac:", "")) - if labels[0] == "build_only": - self.build_only = True - labels.pop(0) - - if not self.build_only: - if labels[0].startswith("bs:"): - self.batch_sizes = [ - int(x) for x in labels.pop(0).replace("bs:", "").split("+") - ] - else: - self.batch_sizes = [0] - - if labels[0].startswith("input_output_len"): - io_lens = labels.pop(0).replace("input_output_len:", - "").split("+") - self.input_lens = [int(x.split(",")[0]) for x in io_lens] - self.output_lens = [int(x.split(",")[1]) for x in io_lens] - elif labels[0].startswith("input_len"): - self.input_lens = [ - int(x) - for x in labels.pop(0).replace("input_len:", "").split("+") - ] - self.output_lens = [] - else: - raise RuntimeError( - f"Unexpected test name label for seq lens: {labels[0]}!") + if labels[0].startswith("bs:"): + self.batch_sizes = [ + int(x) for x in labels.pop(0).replace("bs:", "").split("+") + ] + else: + self.batch_sizes = [0] + + if labels[0].startswith("input_output_len"): + io_lens = labels.pop(0).replace("input_output_len:", "").split("+") + self.input_lens = [int(x.split(",")[0]) for x in io_lens] + self.output_lens = [int(x.split(",")[1]) for x in io_lens] + elif labels[0].startswith("input_len"): + self.input_lens = [ + int(x) + for x in labels.pop(0).replace("input_len:", "").split("+") + ] + self.output_lens = [] + else: + raise RuntimeError( + f"Unexpected test name label for seq lens: {labels[0]}!") if len(labels) > 0: self.num_beams = 1 if not labels[0].startswith("beams:") else int( @@ -766,11 +738,6 @@ def load_from_str(self, test_param_labels) -> None: self.num_gpus = 1 if not labels[0].startswith("gpus:") else int( labels.pop(0).replace("gpus:", "")) - if len(labels) > 0: - self.extra = True if labels[0] == "extra" else False - if self.extra: - labels.pop(0) - self.moe_backend = None if labels and labels[0].startswith("moe:"): self.moe_backend = labels.pop(0).removeprefix("moe:") @@ -798,16 +765,17 @@ def validate(self): VALID_RUNTIMES = ["serve", "bench"] assert self.runtime in VALID_RUNTIMES, \ f"Unsupported runtime '{self.runtime}'; only 'serve' and 'bench' are supported." + assert self.backend in ["pytorch", "_autodeploy"], \ + f"Unsupported backend '{self.backend}'." + if self.runtime == "serve": + assert self.backend == "pytorch", \ + "The serve runtime supports only the pytorch backend." if self.moe_backend is not None: assert self.moe_backend, "moe backend must not be empty!" assert self.backend == "pytorch", \ "moe backend overrides require the pytorch backend!" - # Validate plugin mode. - VALID_MODES = ["plugin", "ootb", "ootb_except_mha"] - assert self.mode in VALID_MODES, f"Invalid mode {self.mode}!" - # Validate dtype. VALID_DTYPES = ["float32", "float16", "bfloat16", "float8", "float4"] assert self.data_type in VALID_DTYPES, f"Invalid data_type {self.data_type}!" @@ -851,26 +819,24 @@ def validate(self): assert self.num_gpus == self.tp_size * self.pp_size, f"Num of GPU shall be equal to TP*PP: {self.num_gpus}, {self.tp_size}, {self.pp_size}" if self.gpu_weights_percent != -1: assert 0 <= self.gpu_weights_percent <= 1, f"Invalid gpu_weights_percent: {self.gpu_weights_percent}!" - if not self.build_only: - assert len(self.input_lens) > 0, f"Empty input_lens!" - if self.is_bert_like(): - assert len( - self.output_lens - ) == 0, f"BERT-like models must not have output_lens!" - else: - assert len( - self.output_lens - ) > 0, f"GPT-like models and enc-dec models must have output_lens!" - - # BERT with small BS is very unstable. Try to avoid it. - if self.is_bert_like(): - if self.runtime == "trtllm-bench": - self.batch_sizes[ - 0] = self.max_batch_size if self.max_batch_size > 0 else 1 - print(f"batch_sizes: {self.batch_sizes}") - assert all( - [b >= 32 for b in self.batch_sizes] - ), f"BERT with small BS is very unstable! Please increase to at least 32." + assert len(self.input_lens) > 0, f"Empty input_lens!" + if self.is_bert_like(): + assert len(self.output_lens + ) == 0, f"BERT-like models must not have output_lens!" + else: + assert len( + self.output_lens + ) > 0, f"GPT-like models and enc-dec models must have output_lens!" + + # BERT with small BS is very unstable. Try to avoid it. + if self.is_bert_like(): + if self.runtime == "bench": + self.batch_sizes[ + 0] = self.max_batch_size if self.max_batch_size > 0 else 1 + print(f"batch_sizes: {self.batch_sizes}") + assert all( + [b >= 32 for b in self.batch_sizes] + ), f"BERT with small BS is very unstable! Please increase to at least 32." # Skip if not enough GPUs. TRTLLM_TOTAL_GPU_COUNT overrides # auto-detection for multi-node setups. @@ -948,7 +914,7 @@ def is_enc_dec(self) -> bool: def get_fixed_dataset_sequence_length(self) -> int | None: """Return the common total length when every dataset shape is fixed.""" - if self.build_only or not self.output_lens: + if not self.output_lens: return None if len(self.input_lens) != len(self.output_lens): @@ -1018,18 +984,7 @@ def set_runtime_configs(self, else: raise RuntimeError(f"Invalid runtime {self._config.runtime}.") - if self._config.runtime == "bench": - build_script = "trtllm-bench" - elif self._config.runtime == "serve": - build_script = None - elif self._config.runtime == "aggr_server": - build_script = None - elif self._config.runtime == "multi_node_disagg_server": - build_script = None - else: - raise RuntimeError( - f"Invalid runtime {self._config.runtime}: engine-build flows " - "were removed with the legacy TensorRT backend.") + build_script = "trtllm-bench" if self._config.runtime == "bench" else None self._build_script = build_script self._benchmark_script = benchmark_script @@ -1082,36 +1037,6 @@ def _get_model_yaml_config(self) -> dict: kv_cache_config.setdefault('avg_seq_len', fixed_sequence_length) return config - def get_trtllm_bench_build_command(self, engine_dir) -> list: - model_dir = self.get_trtllm_bench_model() - if model_dir == "": - pytest.skip("Model Name is not supported by trtllm-bench") - # Legacy "_hf" label; weights load from --model_path. - model_name = self._config.model_name - if not model_name.endswith("_hf"): - model_name = model_name + "_hf" - build_cmd = [ - self._build_script, "--log_level=info", f"--workspace={engine_dir}", - f"--model={model_name}", f"--model_path={model_dir}", "build", - f"--tp_size={self._config.tp_size}", - f"--pp_size={self._config.pp_size}" - ] - max_seq_len = max(self._config.input_lens) + max( - self._config.output_lens) - build_cmd.append(f"--max_seq_len={max_seq_len}") - # Add max_batch_size and max_num_tokens to ensure build matches runtime configuration - # Note: trtllm-bench requires both to be specified together (option group constraint) - assert self._config.max_batch_size > 0, f"max_batch_size must be > 0, got {self._config.max_batch_size}" - assert self._config.max_num_tokens > 0, f"max_num_tokens must be > 0, got {self._config.max_num_tokens}" - build_cmd.append(f"--max_batch_size={self._config.max_batch_size}") - build_cmd.append(f"--max_num_tokens={self._config.max_num_tokens}") - if self._config.quantization: - build_cmd.append( - f"--quantization={self._config.quantization.upper()}") - if self._config.model_name in TRUST_REMOTE_CODE_MODELS: - build_cmd.append(f"--trust_remote_code=True") - return build_cmd - def get_prepare_data_command(self, engine_dir, input_len, output_len) -> list: data_cmd = [] @@ -1227,14 +1152,7 @@ def get_trtllm_bench_command(self, engine_dir): f"--report_json={report_path}", f"--kv_cache_free_gpu_mem_fraction={self._config.kv_cache_free_gpu_mem_fraction}", ] - if self._config.backend == "pytorch": - benchmark_cmd += ["--backend=pytorch"] - elif self._config.backend == "_autodeploy": - benchmark_cmd += ["--backend=_autodeploy"] - else: - benchmark_cmd += [ - f"--backend=tensorrt", f"--engine_dir={engine_dir}" - ] + benchmark_cmd += [f"--backend={self._config.backend}"] if self._config.num_reqs > 0: benchmark_cmd += [f"--num_requests={self._config.num_reqs}"] if self._config.concurrency != -1: @@ -1250,16 +1168,15 @@ def get_trtllm_bench_command(self, engine_dir): if self._config.num_gpus > 1: benchmark_cmd += [f"--warmup={2 * self._config.num_gpus}"] - #Add extra-llm-api-config.yml for pytorch backend and tensorrt backend with extra flag - if self._config.backend == "pytorch" or (self._config.backend == "" - and self._config.extra): + # Add extra-llm-api-config.yml for the PyTorch backend. + if self._config.backend == "pytorch": pytorch_config_path = os.path.join(engine_dir, "extra-llm-api-config.yml") if not os.path.exists(pytorch_config_path): os.makedirs(os.path.dirname(pytorch_config_path), exist_ok=True) config = self._get_model_yaml_config() if config: - print_info(f"pytorch/TRT model config: {config}") + print_info(f"PyTorch model config: {config}") with open(pytorch_config_path, 'w') as f: yaml.dump(config, f, default_flow_style=False) benchmark_cmd += [f"--config={pytorch_config_path}"] @@ -1567,14 +1484,9 @@ def get_commands(self): server_env=server_env, server_timeout=server_timeout) - # Construct engine build command. + # PyTorch and AutoDeploy load model checkpoints directly. build_cmd = [] - if self._config.runtime == "bench": - if self._config.backend in ["pytorch", "_autodeploy"]: - pass - else: - build_cmd = self.get_trtllm_bench_build_command(engine_dir) - else: + if self._config.runtime != "bench": pytest.skip("only support trtllm-bench and serve runtime") # Construct prepare synthetic data command @@ -1861,8 +1773,6 @@ def _get_metrics(self) -> List[PerfTestMetric]: metric_type), cmd_idx=cmd_idx, )) - if self._config.build_only: - return metrics # Then, construct inference latency and gpu mem usage metrics, for each # bs and each seq len. diff --git a/tests/integration/defs/perf/utils.py b/tests/integration/defs/perf/utils.py index 2e03a7d63680..003dac0a0d47 100644 --- a/tests/integration/defs/perf/utils.py +++ b/tests/integration/defs/perf/utils.py @@ -322,17 +322,6 @@ def run_cmd(self, cmd_idx: int, venv) -> str: print(f'Augmented LD_LIBRARY_PATH={envs["LD_LIBRARY_PATH"]}') benchmark_cmd = mpi_cmd + command output += _run_command_with_captured_output(benchmark_cmd, env=envs) - match = re.search(r'--engine_dir=([^\s]+)', current_cmd_str) - if match: - engine_dir = match.group(1) - print_info(f'writing config.json in {engine_dir} to output log') - with open(os.path.join(engine_dir, "config.json"), "r") as f: - config_content = f.read() - output += "\n" + "=" * 50 + "\n" - output += "ENGINE CONFIG:\n" - output += "=" * 50 + "\n" - output += config_content - output += "\n" + "=" * 50 + "\n" return output def get_cmd_str(self, cmd_idx) -> List[str]: diff --git a/tests/integration/defs/stress_test/stress_test.py b/tests/integration/defs/stress_test/stress_test.py index 46ff03b6a976..7f38babfe929 100644 --- a/tests/integration/defs/stress_test/stress_test.py +++ b/tests/integration/defs/stress_test/stress_test.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Stress test script for inference of model using TensorRT LLM with PyTorch/TRT backend. +Stress test script for model inference using the TensorRT LLM PyTorch backend. This script is used for stress testing inference performance using trtllm-serve and aiperf. The script supports three test modes: @@ -115,12 +115,10 @@ class ModelConfig: model_dir: str tp_size: int memory_requirement: int - backend: Optional[str] = None def __str__(self) -> str: model_name = os.path.basename(self.model_dir) - backend_str = f"_{self.backend}" if self.backend else "" - return f"{model_name}_tp{self.tp_size}{backend_str}" + return f"{model_name}_tp{self.tp_size}" @property def model_name(self) -> str: @@ -419,7 +417,6 @@ def is_port_available(port: int, "test_mode", ["stress-test", "stress-stage-alone", "stress-test-with-accuracy"], ids=lambda x: x) -@pytest.mark.parametrize("backend", ["trt", "pytorch"], ids=lambda x: x) @pytest.mark.parametrize("capacity_scheduler_policy", ["GUARANTEED_NO_EVICT", "MAX_UTILIZATION"], ids=lambda x: x) @@ -464,9 +461,9 @@ def is_port_available(port: int, memory_requirement=172032), ], ids=lambda x: f"{os.path.basename(x.model_dir)}_tp{x.tp_size}") -def test_run_stress_test(config, stress_time_timeout, backend, - capacity_scheduler_policy, test_mode): - """Run the stress test with the provided configuration, backend, and test mode. +def test_run_stress_test(config, stress_time_timeout, capacity_scheduler_policy, + test_mode): + """Run the stress test with the provided configuration and test mode. This test function calls the stress_test function with the given parameters. The function should start with test_ prefix to be recognized as a test function by pytest. @@ -474,18 +471,9 @@ def test_run_stress_test(config, stress_time_timeout, backend, Args: config: Model configuration for the test (injected by pytest.mark.parametrize) stress_time_timeout: Tuple of (stress_time, stress_timeout) in seconds - backend: Backend to use ("trt" or "pytorch") capacity_scheduler_policy: Scheduler policy ("GUARANTEED_NO_EVICT", "MAX_UTILIZATION") test_mode: Test mode ("stress-test" or "stress-stage-alone") """ - # Create a new ModelConfig with the backend parameter - # Convert 'trt' to None as expected by the ModelConfig - - new_config = ModelConfig(model_dir=config.model_dir, - tp_size=config.tp_size, - memory_requirement=config.memory_requirement, - backend=backend) - # Extract stress_time and stress_timeout from the tuple stress_time, stress_timeout = stress_time_timeout @@ -493,9 +481,8 @@ def test_run_stress_test(config, stress_time_timeout, backend, server_config = ServerConfig( capacity_scheduler_policy=capacity_scheduler_policy) - # Call the existing stress_test function with the new config and test mode - stress_test(new_config, test_mode, server_config, stress_time, - stress_timeout) + # Call the existing stress_test function with the config and test mode + stress_test(config, test_mode, server_config, stress_time, stress_timeout) def stress_test(config, @@ -684,14 +671,13 @@ def stress_test(config, print_warning(f"Failed to detect GPU architecture: {e}. " "Using default MOE backend (CUTLASS).") - if config.backend == "pytorch": - extra_llm_options.update({ - "cuda_graph_config": { - "enable_padding": True, - "batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128], - }, - "print_iter_log": True, - }) + extra_llm_options.update({ + "cuda_graph_config": { + "enable_padding": True, + "batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128], + }, + "print_iter_log": True, + }) with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as temp_file: @@ -711,7 +697,7 @@ def stress_test(config, "--pp_size", str(test_server_config.pp_size), "--backend", - config.backend, + "pytorch", ] # Only add ep_size parameter if it's not None diff --git a/tests/integration/defs/visual_gen/AGENTS.md b/tests/integration/defs/visual_gen/AGENTS.md new file mode 100644 index 000000000000..2d0eb482d0b3 --- /dev/null +++ b/tests/integration/defs/visual_gen/AGENTS.md @@ -0,0 +1,20 @@ +# AGENTS.md — VisualGen integration tests + +Scope: this directory (`tests/integration/defs/visual_gen/`). Supplements the +repo-root `AGENTS.md`. + +## `pytest.skip` policy + +A VisualGen test that cannot run because a resource CI is expected to provide is +missing must **fail loudly**, not `pytest.skip` — a silent skip reports green and +hides a non-working test from CI. This covers missing model checkpoints, unbuilt or +unimportable first-party modules, and missing compiled TRT-LLM ops. For example, +resolve checkpoints with `from test_common.llm_data import get_checkpoint` — the same +helper the unit suite uses, which raises `FileNotFoundError` when the model is absent — +and do not fall back to `@cached_in_llm_models_root(..., fail_if_path_is_invalid=False)`, +which green-skips a missing checkpoint for the broader LLM suite. + +Skips are only for genuine environment gating a run legitimately can't satisfy: +platform, hardware capability (CUDA / SM / arch), GPU count, and per-test +config/workload preconditions. When a condition couples such a gate with a disallowed +one, keep the gate as a skip and fail only on the resource half. diff --git a/tests/integration/defs/visual_gen/CLAUDE.md b/tests/integration/defs/visual_gen/CLAUDE.md new file mode 100644 index 000000000000..40c6cf85d526 --- /dev/null +++ b/tests/integration/defs/visual_gen/CLAUDE.md @@ -0,0 +1,2 @@ +# In ./CLAUDE.md +@AGENTS.md diff --git a/tests/integration/defs/visual_gen/test_visual_gen_benchmark.py b/tests/integration/defs/visual_gen/test_visual_gen_benchmark.py index d81a103ef635..7e0e53e613bd 100644 --- a/tests/integration/defs/visual_gen/test_visual_gen_benchmark.py +++ b/tests/integration/defs/visual_gen/test_visual_gen_benchmark.py @@ -31,8 +31,8 @@ import pytest import requests import yaml +from test_common.llm_data import get_checkpoint -from defs import conftest from tensorrt_llm._utils import get_free_port # --------------------------------------------------------------------------- @@ -42,16 +42,9 @@ _WAN_T2V_MODEL = "Wan2.1-T2V-1.3B-Diffusers" -def _wan_t2v_path() -> Path: - """Resolve the Wan T2V model path, or call pytest.skip if unavailable.""" - root = Path(conftest.llm_models_root()) - model_path = root / _WAN_T2V_MODEL - if not model_path.is_dir(): - pytest.skip( - f"Wan T2V model not found: {model_path} " - f"(set LLM_MODELS_ROOT or place {_WAN_T2V_MODEL} under scratch)" - ) - return model_path +def _wan_t2v_path() -> str: + """Resolve the Wan T2V model path, failing loudly if it is not staged.""" + return get_checkpoint(_WAN_T2V_MODEL) # Common small-scale generation params for fast CI diff --git a/tests/integration/test_lists/qa/llm_perf_core.yml b/tests/integration/test_lists/qa/llm_perf_core.yml index bfac70bd8c90..1bf324ff4b93 100644 --- a/tests/integration/test_lists/qa/llm_perf_core.yml +++ b/tests/integration/test_lists/qa/llm_perf_core.yml @@ -8,7 +8,7 @@ llm_perf_core: # 3: H100, test cases # 4: H100, GB200, B200, B300, GB300, RTX6000-Server test cases # 5: GB200, B200, B300, GB300, RTX6000-Server test cases -# 6: B200, GB200, B300, GB300 test cases +# 6: B200, GB200, B300, GB300, VR200 test cases # 7: B200, B300 test cases # 8: B300 test cases # 9: H100, B200, B300, RTX6000-Server test cases @@ -80,13 +80,19 @@ llm_perf_core: # 4: H100, GB200, B200, B300, GB300, RTX6000-Server test cases - condition: + terms: + compute_capability: + - 9.0 + - 10.0 + - 10.3 + - 12.0 ranges: system_gpu_count: gte: 2 - compute_capability: - gte: 9.0 - lte: 12.0 tests: + # compute_capability is a list, not a range: a range from Hopper to + # RTX-6000-Server (12.0) also spans Rubin (10.7), and these cases have not been + # run there. A new architecture opts in by being added to the list. - perf/test_perf.py::test_perf[gpt_oss_20b_fp4-bench-pytorch-float4-input_output_len:512,512] - perf/test_perf.py::test_perf[qwen3_235b_a22b_fp8-bench-pytorch-float8-input_output_len:1000,2000-con:256-ep:8-gpus:8] #nemotron_nano_12b_v2 @@ -117,25 +123,32 @@ llm_perf_core: # 5: GB200, B200, B300, GB300, RTX6000-Server test cases - condition: + terms: + compute_capability: + - 10.0 + - 10.3 + - 12.0 ranges: system_gpu_count: gte: 4 - compute_capability: - gte: 10.0 - lte: 12.0 tests: + # compute_capability is a list, not a range: a range from Hopper to + # RTX-6000-Server (12.0) also spans Rubin (10.7), and these cases have not been + # run there. A new architecture opts in by being added to the list. - perf/test_perf.py::test_perf[qwen3_235b_a22b_fp4-bench-pytorch-float4-input_output_len:1000,2000-con:512-ep:4-gpus:4] -# 6: B200, GB200, B300, GB300 test cases +# 6: B200, GB200, B300, GB300, VR200 test cases - condition: ranges: system_gpu_count: gte: 4 compute_capability: gte: 10.0 - lte: 10.3 + lte: 10.7 tests: + # Upper bound is 10.7 so Rubin shares this set: a VR200 NVL72 node exposes 4 GPUs, + # the same shape these cases are sized for, so it runs them rather than a copy. #qwen3.5_27b (dense BF16 52G, 1-GPU): maxbs 512 mamba states need >95G - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:128,128] - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:500,2000] diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 0d3daf885a0c..85c16ae8d2bd 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -73,6 +73,7 @@ l0_cpu: - unittest/_torch/visual_gen/test_teacache.py - unittest/_torch/visual_gen/test_tensor_payload.py - unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py + - unittest/_torch/visual_gen/test_utils.py - unittest/_torch/visual_gen/test_visual_gen_args.py - unittest/_torch/visual_gen/test_visual_gen_params.py - unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -155,4 +156,5 @@ l0_cpu: - unittest/usage - unittest/utils/test_logger.py - unittest/utils/test_confidential_compute.py + - unittest/utils/test_llm_data.py - unittest/visual_gen/test_output.py diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 5e28af2b8dbf..4f947b5ed9ae 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -41,6 +41,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[fp8-latency] SKI accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload SKIP (https://nvbugs/6767938) accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6665906) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] SKIP (https://nvbugs/5838199) +cpp/test_multi_gpu.py::test_cache_transceiver[8proc-nixl_kvcache-90] SKIP (https://nvbugs/5838199) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90] SKIP (https://nvbugs/5838199) disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6105768) disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_cache_aware_balance[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) @@ -88,6 +89,8 @@ examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_ examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[nvfp4] SKIP (https://nvbugs/6572800) examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_lpips_against_golden SKIP (https://nvbugs/6759067) examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_edit_example SKIP (https://nvbugs/6726626) +examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_layered_example SKIP (https://nvbugs/6771010) +examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_layered_lpips_against_golden SKIP (https://nvbugs/6771010) examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[cuda-graph] SKIP (https://nvbugs/6572800) examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[nvfp4] SKIP (https://nvbugs/6572800) examples/visual_gen/test_visual_gen_wan.py::test_fastwan_lpips_against_golden SKIP (https://nvbugs/6572800) diff --git a/tests/test_common/llm_data.py b/tests/test_common/llm_data.py index 51f4555cb225..5c09c27ee70c 100644 --- a/tests/test_common/llm_data.py +++ b/tests/test_common/llm_data.py @@ -86,6 +86,17 @@ def llm_models_root(check: bool = False) -> Optional[Path]: return root if root.exists() else None +def get_checkpoint(model_subdir: str) -> str: + """Resolve a checkpoint under LLM_MODELS_ROOT, or fail loudly if missing.""" + root = llm_models_root(check=True) + path = root / model_subdir + if not path.exists(): + raise FileNotFoundError( + f"Checkpoint not found: {path}. Stage '{model_subdir}' under LLM_MODELS_ROOT to run this test." + ) + return str(path) + + def llm_datasets_root() -> str: return os.path.join(llm_models_root(check=True), "datasets") diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 9af80597eed9..a604a28cfe1c 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -4,8 +4,10 @@ These validate MiniMax-M3 backend selection, indexer/cache integration, decode scratch-buffer sizing, and the paged HND contract passed to the packaged MSA -kernel. Generic block-sparse MQA/GQA numerical coverage lives in the parent -``test_sparse_mqa_gqa.py`` module. +kernel; the CUDA-gated fused cache-scatter test checks the single-launch +K/V/index-K write against the legacy per-cache path. Generic block-sparse +MQA/GQA numerical coverage lives in the parent ``test_sparse_mqa_gqa.py`` +module. """ import sys @@ -20,10 +22,16 @@ MiniMaxM3KVCacheManagerV2, MiniMaxM3MsaSparseAttention, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, msa_paged_kv, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import MsaDecodeSpan from tensorrt_llm._torch.attention.backends.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -1580,3 +1588,136 @@ def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch): assert metadata.msa_seq_lens_cuda.tolist() == [4, 6] assert metadata.msa_out_cache_loc.tolist() == [-1] * 5 assert metadata.msa_n_valid_blocks.tolist() == [0] * 5 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + ("src_dtype", "cache_dtype"), + [ + # bf16 K/V into a bf16 cache: the plain path. + (torch.bfloat16, torch.bfloat16), + # bf16 K/V into an fp8 cache: the kernel folds in the E4M3 cast. + (torch.bfloat16, torch.float8_e4m3fn), + # fp8 K/V into an fp8 cache: production with an FP8 KV cache, where + # the fused QK-norm+RoPE kernel already emits E4M3 k/v + # (MiniMaxM3Attention._emit_fp8_main_qkv), so the kernel stores + # without a cast. + (torch.float8_e4m3fn, torch.float8_e4m3fn), + ], +) +@pytest.mark.parametrize("with_idx", [True, False]) +def test_fused_scatter_matches_reference(src_dtype, cache_dtype, with_idx): + """The fused per-layer cache scatter must match the legacy write_kv_slots + path exactly on production-shaped inputs: non-contiguous HND cache views + carved from a pooled allocation and strided source rows sliced from a fused + projection, for every source/cache dtype pairing the model produces. + Asserting on the whole pool also catches stray writes outside the targeted + slots.""" + torch.manual_seed(0) + device = "cuda" + num_pages, num_kv_heads, tokens_per_block, head_dim = 6, 4, 32, 128 + num_tokens = 17 + inner = num_kv_heads * head_dim + + # Paged HND caches carved from a pool with a coalescing axis, so the + # views are non-contiguous like production get_buffers(...) output. + pool = torch.zeros( + num_pages, 2, num_kv_heads, tokens_per_block, head_dim, dtype=cache_dtype, device=device + ) + k_cache, v_cache = pool[:, 0], pool[:, 1] + idx_pool = torch.zeros( + num_pages, 2, 1, tokens_per_block, head_dim, dtype=torch.bfloat16, device=device + ) + idx_cache = idx_pool[:, 0] + + # Strided sources: rows sliced out of a wider fused-projection tensor. + # randn has no fp8 variant, so generate bf16 and cast the whole buffer, as + # the fused producer would, before slicing the column views. + qkv = torch.randn(num_tokens, 3 * inner + 64, dtype=torch.bfloat16, device=device) + # The index branch stays bf16 on the bf16 indexer path even when the main + # K/V are fp8, so carve index-K out before casting. + idx_k = qkv[:, 2 * inner : 2 * inner + head_dim] if with_idx else None + qkv = qkv.to(src_dtype) + k = qkv[:, :inner] + v = qkv[:, inner : 2 * inner] + + slots = torch.randperm(num_pages * tokens_per_block, device=device)[:num_tokens].to(torch.int32) + + ref_pool = pool.clone() + ref_idx_pool = idx_pool.clone() + write_kv_slots( + ref_pool[:, 0], slots, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + write_kv_slots( + ref_pool[:, 1], slots, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ) + if with_idx: + write_kv_slots( + ref_idx_pool[:, 0], slots, idx_k.reshape(num_tokens, 1, head_dim), layout="HND" + ) + + assert fused_write_layer_caches( + k_cache, v_cache, idx_cache if with_idx else None, slots, k, v, idx_k + ) + + torch.testing.assert_close(pool.to(torch.float32), ref_pool.to(torch.float32)) + torch.testing.assert_close(idx_pool, ref_idx_pool) + + +@pytest.mark.parametrize("sparse", [True, False]) +def test_msa_attention_core_owns_the_cache_write(sparse): + """The model layer's MSA core must write the caches exactly once and in + the right place: write_layer_caches runs before run_indexer (whose proxy + pass reads the index-K cache), run_indexer is told index-K is already + resident, and forward() receives k=v=None so no FMHA phase writes K/V + again.""" + from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3Attention + + num_tokens, width = 3, 128 + topk_indices = torch.zeros(num_tokens, 1, 16, dtype=torch.int32) + events = [] + + class FakeBackend: + layer_idx = 7 + + def write_layer_caches(self, k, v, idx_k, metadata): + events.append(("write", k, v, idx_k, metadata)) + + def run_indexer(self, idx_q, idx_k, metadata, *, idx_k_prewritten=False): + events.append(("indexer", idx_q, idx_k, metadata, idx_k_prewritten)) + return topk_indices + + def forward(self, q, k, v, metadata, forward_args=None): + events.append(("forward", q, k, v, metadata, forward_args)) + + layer = SimpleNamespace(is_sparse_attention_layer=sparse, attn=FakeBackend()) + q, k, v = (torch.zeros(num_tokens, width) for _ in range(3)) + idx_q = torch.zeros(num_tokens, width) if sparse else None + idx_k = torch.zeros(num_tokens, width) if sparse else None + metadata = object() + output = torch.empty(num_tokens, width) + + result = MiniMaxM3Attention._msa_attention_core(layer, q, k, v, idx_q, idx_k, metadata, output) + + assert result is output + names = [event[0] for event in events] + if sparse: + assert names == ["write", "indexer", "forward"] + _, indexer_q, indexer_k, indexer_metadata, prewritten = events[1] + assert indexer_q is idx_q and indexer_k is idx_k and indexer_metadata is metadata + assert prewritten is True + else: + assert names == ["write", "forward"] + + _, written_k, written_v, written_idx_k, write_metadata = events[0] + assert written_k is k and written_v is v and write_metadata is metadata + assert written_idx_k is idx_k + + _, forward_q, forward_k, forward_v, forward_metadata, forward_args = events[-1] + assert forward_q is q and forward_metadata is metadata + assert forward_k is None and forward_v is None + assert forward_args.output is output + if sparse: + assert forward_args.sparse_backend_args.topk_indices is topk_indices + else: + assert forward_args.sparse_backend_args is None diff --git a/tests/unittest/_torch/attention/test_combined_fmha.py b/tests/unittest/_torch/attention/test_combined_fmha.py index 9f76ea59da41..9f60870a8524 100644 --- a/tests/unittest/_torch/attention/test_combined_fmha.py +++ b/tests/unittest/_torch/attention/test_combined_fmha.py @@ -117,6 +117,70 @@ def test_flashinfer_fp8_mode_remains_implementation_local() -> None: assert not fmha._use_fp8_context_fmha(output, AttentionInputType.generation_only) +@pytest.mark.parametrize( + "num_contexts,num_generations,dtype,expect_fallback", + [ + (1, 0, torch.bfloat16, True), + (4, 0, torch.bfloat16, True), + (5, 0, torch.bfloat16, False), + (1, 1, torch.bfloat16, False), + (2, 9, torch.bfloat16, False), + (2, 19, torch.bfloat16, False), + (2, 38, torch.bfloat16, False), + (4, 1, torch.bfloat16, False), + (0, 1, torch.bfloat16, False), + (1, 0, torch.float16, False), + ], +) +def test_small_context_fallback_preserves_mixed_batch_generation( + num_contexts: int, + num_generations: int, + dtype: torch.dtype, + expect_fallback: bool, +) -> None: + attn = FakeAttention() + attn.head_dim = 256 + attn.sparse_params = None + attn.position_embedding_type = 0 + fmha = FlashInferTrtllmGenFmha(attn) + metadata = SimpleNamespace( + num_contexts=num_contexts, + num_generations=num_generations, + helix_position_offsets=None, + num_sparse_topk=0, + use_spec_decoding=False, + kv_cache_block_offsets=object(), + kv_cache_manager=None, + tokens_per_block=32, + is_cross=False, + beam_width=1, + ) + if num_contexts == 0: + input_type = AttentionInputType.generation_only + phases = (None, FmhaPhase.GENERATION) + elif num_generations == 0: + input_type = AttentionInputType.context_only + phases = (None, FmhaPhase.CONTEXT) + else: + input_type = AttentionInputType.mixed + phases = (None, FmhaPhase.CONTEXT, FmhaPhase.GENERATION) + num_tokens = num_contexts + num_generations + q = torch.empty((num_tokens, 3 * attn.head_dim), dtype=dtype) + forward_args = AttentionForwardArgs( + output=torch.empty((num_tokens, attn.head_dim), dtype=dtype), + attention_input_type=input_type, + is_fused_qkv=True, + ) + + for phase in phases: + supported, reason = fmha._is_supported_with_reason( + q, None, None, attn, metadata, forward_args, phase=phase + ) + assert supported is not expect_fallback, reason + if expect_fallback: + assert "small-batch BF16 context attention" in reason + + def test_triton_custom_mask_rejects_whole_request_probe() -> None: fmha = object.__new__(TritonCustomMaskFmha) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv3.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv3.py new file mode 100644 index 000000000000..201d22142c6a --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv3.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Per-module quant config resolution for DeepSeek MIXED_PRECISION checkpoints. + +DeepSeek-R1-W4AFP8 ships an hf_quant_config with the global +quant_algo=MIXED_PRECISION, which does not map to a single QuantMode. The MoE +experts must resolve their own per-module config (W4A8_AWQ) instead of using +the global one. These tests exercise that resolution on CPU without weights. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.models.modeling_deepseekv3 import Deepseekv3MoE +from tensorrt_llm._torch.moe.fused_moe import MoEWeightLoadingMode +from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + +EXPERTS_KEY = "model.layers.{}.mlp.experts" + + +@pytest.fixture +def mixed_precision_config(): + return QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION) + + +@pytest.fixture +def w4a8_awq_config(): + return QuantConfig(quant_algo=QuantAlgo.W4A8_AWQ, group_size=128) + + +def test_w4a8_awq_config_is_int4_weight_only_per_group(w4a8_awq_config): + # This predicate is what selects MoEWeightLoadingMode.W4A8_CUSTOM for the + # experts in Deepseekv3MoE, so pin it down explicitly. + assert w4a8_awq_config.layer_quant_mode.is_int4_weight_only_per_group() + + +def test_experts_quant_config_resolved_from_per_module_dict( + mixed_precision_config, w4a8_awq_config +): + model_config = SimpleNamespace( + quant_config=mixed_precision_config, + quant_config_dict={EXPERTS_KEY.format(0): w4a8_awq_config}, + ) + + resolved = Deepseekv3MoE._get_experts_quant_config(model_config, 0) + + assert resolved is w4a8_awq_config + assert resolved.quant_algo == QuantAlgo.W4A8_AWQ + assert resolved.layer_quant_mode.is_int4_weight_only_per_group() + assert not mixed_precision_config.layer_quant_mode.is_int4_weight_only_per_group() + + +def test_experts_quant_config_falls_back_to_global_for_unlisted_layer( + mixed_precision_config, w4a8_awq_config +): + model_config = SimpleNamespace( + quant_config=mixed_precision_config, + quant_config_dict={EXPERTS_KEY.format(0): w4a8_awq_config}, + ) + + assert Deepseekv3MoE._get_experts_quant_config(model_config, 1) is mixed_precision_config + + +def test_experts_quant_config_falls_back_to_global_without_dict(mixed_precision_config): + model_config = SimpleNamespace(quant_config=mixed_precision_config, quant_config_dict=None) + + assert Deepseekv3MoE._get_experts_quant_config(model_config, 0) is mixed_precision_config + + +def test_expert_weight_loading_mode_w4a8_custom_for_w4a8_awq(w4a8_awq_config): + # The resolved W4A8_AWQ expert config selects the custom loading mode; this + # is the assignment the MoE construction makes from the resolved config. + assert ( + Deepseekv3MoE._expert_weight_loading_mode(w4a8_awq_config) + is MoEWeightLoadingMode.W4A8_CUSTOM + ) + + +def test_expert_weight_loading_mode_vanilla_for_non_int4(mixed_precision_config): + # Neither the ambiguous MIXED_PRECISION global nor a plain FP8 config is + # int4-weight-per-group, so both fall to VANILLA. + assert ( + Deepseekv3MoE._expert_weight_loading_mode(mixed_precision_config) + is MoEWeightLoadingMode.VANILLA + ) + assert ( + Deepseekv3MoE._expert_weight_loading_mode(QuantConfig(quant_algo=QuantAlgo.FP8)) + is MoEWeightLoadingMode.VANILLA + ) + + +def test_expert_weight_loading_mode_none_is_vanilla(): + # override_quant_config is Optional, so the resolved expert config is None on + # an unquantized layer; the mode selection must not dereference it. + assert Deepseekv3MoE._expert_weight_loading_mode(None) is MoEWeightLoadingMode.VANILLA + + +def _bare_configurable_moe(override_quant_config): + moe = object.__new__(ConfigurableMoE) + moe._override_quant_config = override_quant_config + return moe + + +def test_quant_config_dict_prefers_override_over_mixed_precision_global( + mixed_precision_config, w4a8_awq_config +): + moe = _bare_configurable_moe(w4a8_awq_config) + model_config = SimpleNamespace(quant_config=mixed_precision_config) + + result = ConfigurableMoE._get_quant_config_dict(moe, model_config) + + assert result == { + "has_fp8_qdq": False, + "has_nvfp4": False, + "has_w4afp8": True, + "has_fp8_block_scales": False, + } + # The global MIXED_PRECISION mode would not have flagged w4afp8. + assert not mixed_precision_config.layer_quant_mode.is_int4_weight_only_per_group() + + +def test_quant_config_dict_falls_back_to_global_without_override(): + fp8_config = QuantConfig(quant_algo=QuantAlgo.FP8) + moe = _bare_configurable_moe(None) + model_config = SimpleNamespace(quant_config=fp8_config) + + result = ConfigurableMoE._get_quant_config_dict(moe, model_config) + + assert result == { + "has_fp8_qdq": True, + "has_nvfp4": False, + "has_w4afp8": False, + "has_fp8_block_scales": False, + } + + +def test_quant_config_dict_is_none_when_unquantized(): + moe = _bare_configurable_moe(None) + model_config = SimpleNamespace(quant_config=None) + + assert ConfigurableMoE._get_quant_config_dict(moe, model_config) is None diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index 4d18f25151e0..6ebb346d8610 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -20,10 +20,14 @@ import pytest import torch +import tensorrt_llm._torch.custom_ops.torch_custom_ops as custom_ops_module import tensorrt_llm._torch.modules.linear as linear_module from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE, + FlashInferMXFP8GemmRunner, MXFP8GemmRunner, + MXFP8QuantizeRunner, _get_mxfp8_large_m_tuning_buckets, _map_to_mxfp8_large_m_bucket, ) @@ -529,3 +533,94 @@ def test_mxfp8_flashinfer_decode_graph_matches_native(monkeypatch, batch_size): assert flashinfer_gemm.call_count == 1 graph.replay() torch.testing.assert_close(graph_output, native_output, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif( + not _mxfp8_cutlass_op_available(), + reason="MXFP8xMXFP8 GEMM op not compiled or sm < 100", +) +@pytest.mark.parametrize("batch_size", (1, 8, 16, 32)) +def test_mxfp8_decode_graph_backend_tuning_matches_native(monkeypatch, batch_size): + """Per-bucket tuned decode graphs must match the native op. + + Profile the quantizer and GEMM backends for a decode bucket during the + warmup-only pass, then capture the same shape and replay it. This covers + the CuTeDSL scale layouts and the in-process winner cache used by capture. + """ + if not IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE: + pytest.skip("FlashInfer CuTeDSL MXFP8 kernels are not available") + + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) + torch.manual_seed(0) + out_f, in_f = 256, 512 + weight = torch.randn(out_f, in_f, dtype=torch.bfloat16) + weight_e4m3, weight_scale = quant_bf16_to_mxfp8(weight, 32) + x = torch.randn(batch_size, in_f, dtype=torch.bfloat16, device="cuda") + quant_config = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) + + native = Linear( + in_features=in_f, + out_features=out_f, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + ).cuda() + tuned = Linear( + in_features=in_f, + out_features=out_f, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + ).cuda() + weights = [{"weight": weight_e4m3, "weight_scale_inv": weight_scale}] + native.load_weights(weights) + tuned.load_weights(weights) + native_output = native(x) + + method = tuned.quant_method + assert isinstance(method, MXFP8LinearMethod) + assert method.enable_flashinfer_auto() + method.tune_decode_graph_backends = True + + # Record the (quantize, GEMM) tactics each pass selects. + chosen_tactics = [] + choose_tactic = custom_ops_module._choose_mxfp8_tactic + + def record_tactic(*args, **kwargs): + tactic = choose_tactic(*args, **kwargs) + chosen_tactics.append(tactic) + return tactic + + monkeypatch.setattr(custom_ops_module, "_choose_mxfp8_tactic", record_tactic) + + # Both CuTeDSL candidates must run and match, independent of which + # backend wins the profiling below. + act, act_scale = MXFP8QuantizeRunner(x.dtype)([x], tactic=MXFP8QuantizeRunner.CUTE_DSL) + cute_dsl_output = FlashInferMXFP8GemmRunner(tuned.dtype)( + [act, act_scale, tuned.weight, tuned.weight_scale], + tactic=FlashInferMXFP8GemmRunner.CUTE_DSL, + ) + torch.testing.assert_close(cute_dsl_output, native_output, rtol=2e-2, atol=2e-2) + + # Warmup-only pass: profile both backends of each stage for this bucket. + with flashinfer_mxfp8_autotune(), flashinfer_mxfp8_decode_graph_capture(): + warmup_output = tuned(x) + torch.testing.assert_close(warmup_output, native_output, rtol=2e-2, atol=2e-2) + assert len(chosen_tactics) == 2 + warmup_tactics = tuple(chosen_tactics) + chosen_tactics.clear() + + # Capture pass: the in-process winners are reused without profiling. + static_x = x.clone() + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + with flashinfer_mxfp8_decode_graph_capture(): + graph_output = tuned(static_x) + assert tuple(chosen_tactics) == warmup_tactics + + # Replay on a fresh input so the check cannot pass on the captured result. + replay_x = torch.randn_like(x) + static_x.copy_(replay_x) + graph.replay() + torch.testing.assert_close(graph_output, native(replay_x), rtol=2e-2, atol=2e-2) diff --git a/tests/unittest/_torch/visual_gen/AGENTS.md b/tests/unittest/_torch/visual_gen/AGENTS.md new file mode 100644 index 000000000000..ddfc9d13201f --- /dev/null +++ b/tests/unittest/_torch/visual_gen/AGENTS.md @@ -0,0 +1,20 @@ +# AGENTS.md — VisualGen unit tests + +Scope: this directory and everything under it (`tests/unittest/_torch/visual_gen/`, +including `multi_gpu/`, `kernels/`, `sparse_attention/`). Supplements the repo-root +`AGENTS.md`. + +## `pytest.skip` policy + +A VisualGen test that cannot run because a resource CI is expected to provide is +missing must **fail loudly**, not `pytest.skip` — a silent skip reports green and +hides a non-working test from CI. This covers missing model checkpoints, unbuilt or +unimportable first-party modules, and missing compiled TRT-LLM ops. For example, +resolve checkpoints from `LLM_MODELS_ROOT` via +`from utils.llm_data import get_checkpoint`, which raises `FileNotFoundError` when the +model is absent — don't guard the test with `if not os.path.exists(ckpt): pytest.skip(...)`. + +Skips are only for genuine environment gating a run legitimately can't satisfy: +platform (`sys.platform`), hardware capability (CUDA / SM / arch), GPU count, and +per-test config/workload preconditions. When a condition couples such a gate with a +disallowed one, keep the gate as a skip and fail only on the resource half. diff --git a/tests/unittest/_torch/visual_gen/CLAUDE.md b/tests/unittest/_torch/visual_gen/CLAUDE.md new file mode 100644 index 000000000000..40c6cf85d526 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/CLAUDE.md @@ -0,0 +1,2 @@ +# In ./CLAUDE.md +@AGENTS.md diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py index e9397c775f3f..f71f61ecd714 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py @@ -33,19 +33,11 @@ import torch.nn as nn import torch.nn.functional as F -try: - from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask - from tensorrt_llm._torch.visual_gen.attention_backend import Attention2DAttention - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( - _flash_attn_fwd as _fa4_fwd, - ) - from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False - _fa4_fwd = None +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.visual_gen.attention_backend import Attention2DAttention +from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention +from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd +from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout # FA4 kernel availability (separate from the combine kernel used by Attention2DAttention) _flash_attn4_available = _fa4_fwd is not None @@ -154,8 +146,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py index eb984ecf982b..22801ba49903 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py @@ -23,30 +23,23 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - AttentionConfig, - DiffusionModelConfig, - TorchCompileConfig, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( - COSMOS3_EDGE_BACKBONE_TYPE, - Cosmos3VFMTransformer, - ) - from tensorrt_llm.models.modeling_utils import QuantConfig - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False - # Module-level configs below reference this; every test skips in this - # branch, but the definitions still have to import cleanly. - COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" +from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + TorchCompileConfig, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + COSMOS3_EDGE_BACKBONE_TYPE, + Cosmos3VFMTransformer, +) +from tensorrt_llm.models.modeling_utils import QuantConfig # Attention2D (attn2d) wraps the compute backend in Attention2DAttention, which # requires (a) an LSE-capable inner backend — only FA4, VANILLA does not support # LSE — and (b) the ``flash_attn_combine`` JIT kernel. Detect both up front so the -# attn2d tests skip cleanly when the kernels are not built (e.g. non-Blackwell CI). +# attn2d tests fail loudly when the kernels are missing on their Blackwell runner — +# a build/dependency problem, not a reason to skip. try: from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( _flash_attn_fwd as _fa4_fwd, @@ -55,7 +48,7 @@ _flash_attn_combine as _fa_combine, ) - _ATTN2D_AVAILABLE = MODULES_AVAILABLE and _fa4_fwd is not None and _fa_combine is not None + _ATTN2D_AVAILABLE = _fa4_fwd is not None and _fa_combine is not None except ImportError: _ATTN2D_AVAILABLE = False @@ -212,8 +205,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") backend = "nccl" if use_cuda else "gloo" @@ -822,62 +813,50 @@ def _logic_cosmos3_attn2d_ulysses_vs_single_gpu(rank, world_size): class TestCosmos3TransformerParallel: """Cosmos3 TP / Ulysses / CFG parity vs single-GPU (synthetic weights, no checkpoint).""" - def _skip_if_unavailable(self): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - def test_tp2_vs_single_gpu(self): - self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_tp_vs_single_gpu) def test_ulysses2_vs_single_gpu(self): - self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_vs_single_gpu) def test_edge_tp2_vs_single_gpu(self): """Edge's non-gated relu² MLP shards without the gate_up fusion the Qwen recipe uses, so column/row splitting takes a different path.""" - self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_edge_tp_vs_single_gpu) def test_edge_ulysses2_vs_single_gpu(self): """Edge under sequence sharding: no und Q/K norm, and the reasoner's keys are normed only where the generator consumes them.""" - self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_edge_ulysses_vs_single_gpu) def test_ulysses2_audio_vs_single_gpu(self): """Ulysses parity with the audio modality on: video + audio tokens are sharded together across the sequence dimension.""" - self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_audio_vs_single_gpu) def test_ulysses2_action_vs_single_gpu(self): """Ulysses parity with action tokens appended to the GEN sequence.""" - self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_action_vs_single_gpu) @pytest.mark.gpu4 def test_tp2_ulysses2_vs_single_gpu(self): - self._skip_if_unavailable() run_test_in_distributed(world_size=4, test_fn=_logic_cosmos3_tp_ulysses_vs_single_gpu) @pytest.mark.gpu4 def test_cfg2_ulysses2_vs_single_gpu(self): - self._skip_if_unavailable() run_test_in_distributed(world_size=4, test_fn=_logic_cosmos3_cfg_ulysses_vs_single_gpu) def test_attn2d_2x1_vs_single_gpu(self): - self._skip_if_unavailable() - if not _ATTN2D_AVAILABLE: - pytest.skip("FA4 / flash_attn_combine JIT kernels not available") + assert _ATTN2D_AVAILABLE, ( + "FA4 / flash_attn_combine JIT kernels not available; expected on the Blackwell CI runner" + ) run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_attn2d_vs_single_gpu) @pytest.mark.gpu4 def test_attn2d_2x1_ulysses2_vs_single_gpu(self): - self._skip_if_unavailable() - if not _ATTN2D_AVAILABLE: - pytest.skip("FA4 / flash_attn_combine JIT kernels not available") + assert _ATTN2D_AVAILABLE, ( + "FA4 / flash_attn_combine JIT kernels not available; expected on the Blackwell CI runner" + ) run_test_in_distributed(world_size=4, test_fn=_logic_cosmos3_attn2d_ulysses_vs_single_gpu) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py index 2e1918235321..dd33d28f2375 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux2_transformer_parallel.py @@ -34,17 +34,12 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - AttentionConfig, - DiffusionModelConfig, - TorchCompileConfig, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + TorchCompileConfig, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping try: from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( @@ -93,8 +88,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True, **kwargs): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") backend = "nccl" if use_cuda else "gloo" @@ -272,10 +265,7 @@ def _logic_flux2_transformer_parallel_vs_single_gpu( torch.manual_seed(SEED_WEIGHTS) dist_config = _make_model_config(pretrained_cfg, backend="FA4", **parallel_cfg_kwargs) - try: - dist_model = Flux2Transformer2DModel(dist_config).to(device).to(dtype) - except (ImportError, ValueError, NotImplementedError) as e: - pytest.skip(f"[{label}] Parallel backend unavailable: {e}") + dist_model = Flux2Transformer2DModel(dist_config).to(device).to(dtype) dist_model.load_state_dict(ref_state) torch.manual_seed(SEED_INPUT) @@ -329,11 +319,10 @@ def _logic_flux2_transformer_parallel_vs_single_gpu( @pytest.mark.integration @pytest.mark.flux2 class TestFlux2TransformerParallel: - def _skip_if_unavailable(self): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - if not _flash_attn4_available: - pytest.skip("FlashAttn4 JIT kernels not available") + def _require_flash_attn4(self) -> None: + assert _flash_attn4_available, ( + "FlashAttn4 JIT kernels not available; expected on the Blackwell CI runner" + ) @pytest.mark.parametrize( "label,parallel_cfg_kwargs", @@ -341,7 +330,7 @@ def _skip_if_unavailable(self): ids=[name for name, _ in _FLUX2_8GPU_PARALLEL_COMBINATIONS], ) def test_parallel_all_combinations_vs_single_gpu_8gpu(self, label, parallel_cfg_kwargs): - self._skip_if_unavailable() + self._require_flash_attn4() run_test_in_distributed( world_size=8, test_fn=_logic_flux2_transformer_parallel_vs_single_gpu, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py index ee832051bb74..ffab5efaf4b3 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_tp.py @@ -34,30 +34,25 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - AttentionConfig, - DiffusionModelConfig, - TorchCompileConfig, - create_attention_metadata_state, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm._torch.visual_gen.models.flux.joint_proj import ( - FluxJointAttnMLPProj, - FluxJointQKVMLPProj, - ) - from tensorrt_llm.models.modeling_utils import QuantConfig - - from .tp_shard_utils import ( - copy_tp_parameter, - shard_dim1, - shard_fused_gate_up, - shard_fused_qkv_by_heads, - ) +from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + TorchCompileConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm._torch.visual_gen.models.flux.joint_proj import ( + FluxJointAttnMLPProj, + FluxJointQKVMLPProj, +) +from tensorrt_llm.models.modeling_utils import QuantConfig - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from .tp_shard_utils import ( + copy_tp_parameter, + shard_dim1, + shard_fused_gate_up, + shard_fused_qkv_by_heads, +) @pytest.fixture(autouse=True, scope="module") @@ -103,9 +98,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): """Run a test function in a distributed environment.""" - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py index d85820285819..34afe03f0d5b 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py @@ -22,18 +22,13 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm.models.modeling_utils import QuantConfig - from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig @pytest.fixture(autouse=True, scope="module") @@ -79,9 +74,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): """Run a test function in a distributed environment.""" - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py index 935f2153deba..92215e3183fb 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_async_ulysses.py @@ -35,18 +35,13 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm.models.modeling_utils import QuantConfig - from tensorrt_llm.visual_gen.args import AttentionConfig, ParallelConfig, TorchCompileConfig - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.visual_gen.args import AttentionConfig, ParallelConfig, TorchCompileConfig @pytest.fixture(autouse=True, scope="module") @@ -86,8 +81,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, fn_args): def run_test_in_distributed(world_size: int, test_fn: Callable, *fn_args): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py index 0c21c7e5c093..211d19bbae71 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_ulysses.py @@ -23,19 +23,14 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.rope import LTXRopeType - from tensorrt_llm.models.modeling_utils import QuantConfig - from tensorrt_llm.visual_gen.args import AttentionConfig, ParallelConfig, TorchCompileConfig - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.rope import LTXRopeType +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.visual_gen.args import AttentionConfig, ParallelConfig, TorchCompileConfig @pytest.fixture(autouse=True, scope="module") @@ -75,8 +70,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, fn_args): def run_test_in_distributed(world_size: int, test_fn: Callable, *fn_args): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_attention.py index 727827d11815..7747f0af6c7f 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_attention.py @@ -17,15 +17,9 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp +from diffusers.models.autoencoders.autoencoder_kl_wan import WanAttentionBlock -try: - from diffusers.models.autoencoders.autoencoder_kl_wan import WanAttentionBlock - - from tensorrt_llm._torch.visual_gen.modules.vae import ParallelVaeAttentionBlock - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.modules.vae import ParallelVaeAttentionBlock @pytest.fixture(autouse=True, scope="module") @@ -65,8 +59,6 @@ def _distributed_worker(rank, world_size, test_fn, port): def _run(world_size: int, test_fn: Callable): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Need {world_size} GPUs, have {torch.cuda.device_count()}") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py index 1ef923ab7019..4a8eeb7e61bf 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py @@ -21,25 +21,16 @@ import torch.distributed as dist import torch.multiprocessing as mp import torch.nn as nn +from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d -try: - from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d - - from tensorrt_llm._torch.visual_gen.models.wan.parallel_vae import WanCausalConvHalo - from tensorrt_llm._torch.visual_gen.modules.vae import ( - HaloExchangeConv, - HaloExchangeConv2dStride2, - ) - from tensorrt_llm._torch.visual_gen.modules.vae.conv import ( - _cat_spatial_halos, - _halo_exchange_buffer, - _physical_to_logical_channels_last, - _spatial_channels_last_format, - ) - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.models.wan.parallel_vae import WanCausalConvHalo +from tensorrt_llm._torch.visual_gen.modules.vae import HaloExchangeConv, HaloExchangeConv2dStride2 +from tensorrt_llm._torch.visual_gen.modules.vae.conv import ( + _cat_spatial_halos, + _halo_exchange_buffer, + _physical_to_logical_channels_last, + _spatial_channels_last_format, +) @pytest.fixture(autouse=True, scope="module") @@ -79,8 +70,6 @@ def _distributed_worker(rank, world_size, test_fn, port): def _run(world_size: int, test_fn: Callable): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Need {world_size} GPUs, have {torch.cuda.device_count()}") # Spawn distributed workers via a helper that retries with a fresh master @@ -345,7 +334,6 @@ def test_conv2d_stride2_offset_group_4gpu(self) -> None: _run(4, _logic_halo_conv2d_stride2_offset_group) -@pytest.mark.skipif(not MODULES_AVAILABLE, reason="Required modules not available") class TestHaloExchangeValidation: def test_missing_required_adjacent_group_fails_at_construction(self) -> None: conv = nn.Conv2d(4, 4, kernel_size=3, padding=1) @@ -359,7 +347,6 @@ def test_kernel_one_does_not_require_adjacent_group(self) -> None: HaloExchangeConv(conv, chunk_dim=3, adj_groups=[None], rank=0, world_size=2) -@pytest.mark.skipif(not MODULES_AVAILABLE, reason="Required modules not available") class TestSpatialChannelsLastFormat: """CPU unit tests for the halo channels-last layout helper. diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_group_norm.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_group_norm.py index 3dde825d65d5..e3eec279711c 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_group_norm.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_group_norm.py @@ -19,12 +19,7 @@ import torch.multiprocessing as mp import torch.nn as nn -try: - from tensorrt_llm._torch.visual_gen.modules.vae import GroupNormParallel - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.modules.vae import GroupNormParallel @pytest.fixture(autouse=True, scope="module") @@ -59,8 +54,6 @@ def _distributed_worker(rank, world_size, test_fn, port): def _run(world_size: int, test_fn: Callable): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Need {world_size} GPUs, have {torch.cuda.device_count()}") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py index befca713cb34..f7ef7160d253 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py @@ -20,20 +20,14 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp +from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan -try: - from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan - - from tensorrt_llm._torch.visual_gen.models.wan.parallel_vae import ( - ParallelVAE_TrtllmWan, - ParallelVAE_Wan, - ) - from tensorrt_llm._torch.visual_gen.models.wan.wan_vae import WanVAE, WanVAEConfig - from tensorrt_llm._torch.visual_gen.modules.vae.parallel_vae_interface import ParallelVAEFactory - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.models.wan.parallel_vae import ( + ParallelVAE_TrtllmWan, + ParallelVAE_Wan, +) +from tensorrt_llm._torch.visual_gen.models.wan.wan_vae import WanVAE, WanVAEConfig +from tensorrt_llm._torch.visual_gen.modules.vae.parallel_vae_interface import ParallelVAEFactory @pytest.fixture(autouse=True, scope="module") @@ -73,8 +67,6 @@ def _distributed_worker(rank, world_size, test_fn, port): def _run(world_size: int, test_fn: Callable): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Need {world_size} GPUs, have {torch.cuda.device_count()}") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py index 850f5b6089b8..74412b67946c 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_qwen_image_attention_parallel.py @@ -24,20 +24,15 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( - Attention2DAttention, - RingAttention, - UlyssesAttention, - ) - from tensorrt_llm._torch.visual_gen.config import AttentionConfig, DiffusionModelConfig - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenJointAttention - from tensorrt_llm.visual_gen.args import ParallelConfig - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( + Attention2DAttention, + RingAttention, + UlyssesAttention, +) +from tensorrt_llm._torch.visual_gen.config import AttentionConfig, DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenJointAttention +from tensorrt_llm.visual_gen.args import ParallelConfig try: from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( @@ -85,8 +80,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs): def _run_test_in_distributed(world_size: int, test_fn: Callable, **kwargs) -> None: - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") from ._visual_gen_dist_utils import spawn_with_retry diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py index 3f49e635c6ac..dc512935fa26 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py @@ -41,19 +41,11 @@ import torch.nn.functional as F from torch.distributed.device_mesh import init_device_mesh -try: - from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask - from tensorrt_llm._torch.visual_gen.attention_backend import RingAttention, UlyssesAttention - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention - from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( - _flash_attn_fwd as _fa4_fwd, - ) - from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False - _fa4_fwd = None +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.visual_gen.attention_backend import RingAttention, UlyssesAttention +from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention +from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd +from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout _flash_attn4_available = _fa4_fwd is not None @@ -141,8 +133,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py index 1be7502b652b..f5523fe02973 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_tp_attention.py @@ -38,17 +38,12 @@ import torch.multiprocessing as mp import torch.nn.functional as F -try: - from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl - from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode - from tensorrt_llm.mapping import Mapping - from tensorrt_llm.visual_gen.args import AttentionConfig - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.visual_gen.args import AttentionConfig @pytest.fixture(autouse=True, scope="module") @@ -88,8 +83,6 @@ def _distributed_worker(rank, world_size, test_fn, port, *args): def _run(world_size: int, test_fn: Callable, *args): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Need {world_size} GPUs, have {torch.cuda.device_count()}") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py index e5d4f2660321..0f22d38a843f 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py @@ -29,13 +29,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.distributed import all_to_all_4d - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False - +from tensorrt_llm._torch.distributed import all_to_all_4d # Loop count must comfortably exceed kNumSlots so the ring wraps at least # twice. kNumSlots is 3 today; 8 iterations = ~2.67 full rotations. @@ -257,8 +251,6 @@ def _worker_multi_pg(rank, world_size, port): def _run(world_size: int, test_fn: Callable): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py index 2b900c603ce6..8415bd6bc52a 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py @@ -18,19 +18,13 @@ import torch.multiprocessing as mp import torch.nn.functional as F -# Try to import the modules - skip tests if not available -try: - from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask - from tensorrt_llm._torch.distributed import all_to_all_4d, all_to_all_5d - from tensorrt_llm._torch.visual_gen.attention_backend import UlyssesAttention, VanillaAttention - from tensorrt_llm._torch.visual_gen.attention_backend.interface import ( - AttentionBackend, - AttentionTensorLayout, - ) - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.distributed import all_to_all_4d, all_to_all_5d +from tensorrt_llm._torch.visual_gen.attention_backend import UlyssesAttention, VanillaAttention +from tensorrt_llm._torch.visual_gen.attention_backend.interface import ( + AttentionBackend, + AttentionTensorLayout, +) @pytest.fixture(autouse=True, scope="module") @@ -82,9 +76,6 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = Should accept (rank, world_size) as arguments. use_cuda: Whether to use CUDA (requires sufficient GPUs) """ - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py index 3f63447dab5f..e2669bc337c5 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.py @@ -36,17 +36,13 @@ import torch.multiprocessing as mp import torch.nn.functional as F -try: - from tensorrt_llm._torch.visual_gen.attention_backend import UlyssesAttention - from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention - from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state - from tensorrt_llm.visual_gen.args import QuantAttentionConfig +from tensorrt_llm._torch.visual_gen.attention_backend import UlyssesAttention +from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state +from tensorrt_llm.visual_gen.args import QuantAttentionConfig - MODULES_AVAILABLE = True - ATTENTION_META_DICT = threading.local() - ATTENTION_META_DICT.metadata = create_attention_metadata_state() -except ImportError: - MODULES_AVAILABLE = False +ATTENTION_META_DICT = threading.local() +ATTENTION_META_DICT.metadata = create_attention_metadata_state() def _cuda_cc(): @@ -94,8 +90,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port): def run_test_in_distributed(world_size: int, test_fn: Callable): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if not torch.cuda.is_available(): pytest.skip("CUDA required for SageAttention") if _cuda_cc()[0] != 10: diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py index 44ac859a70b5..75b4d6a58ff0 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py @@ -12,12 +12,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping @pytest.fixture(autouse=True, scope="module") @@ -50,8 +45,6 @@ def _worker(rank, world_size, test_fn, port): def _run_multi_gpu(world_size, test_fn): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if not torch.cuda.is_available() or torch.cuda.device_count() < world_size: pytest.skip(f"Requires {world_size} GPUs, have {torch.cuda.device_count()}") # Spawn distributed workers via a helper that retries with a fresh master @@ -596,7 +589,6 @@ def _logic_tp2_attn2d_2x1_groups(rank, world_size): assert x.item() == 2.0, f"Rank {rank}: attn2d_col all_reduce expected 2, got {x.item()}" -@pytest.mark.skipif(not MODULES_AVAILABLE, reason="Modules not available") class TestMultiGPU: def test_default_order_cfg2_ulysses2(self): _run_multi_gpu(4, _logic_default_order_cfg2_ulysses2) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py index 042525f82e30..96fba95677f5 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py @@ -31,20 +31,15 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm.models.modeling_utils import QuantConfig - from tensorrt_llm.visual_gen.args import ( - AttentionConfig, - ParallelConfig, - TeaCacheConfig, - TorchCompileConfig, - ) - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + ParallelConfig, + TeaCacheConfig, + TorchCompileConfig, +) @pytest.fixture(autouse=True, scope="module") @@ -84,8 +79,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, fn_args): def run_test_in_distributed(world_size: int, test_fn: Callable, *fn_args): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") # Spawn distributed workers via a helper that retries with a fresh master diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py index e89aa09ddeb3..d0d26f646190 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py @@ -16,15 +16,10 @@ Run with: pytest tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN21_1_3B=/path/to/1.3b \\ - pytest tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py -v -s """ import gc import os -from pathlib import Path from typing import Callable os.environ["TLLM_DISABLE_MPI"] = "1" @@ -35,23 +30,16 @@ import torch.distributed as dist import torch.multiprocessing as mp import torch.nn.functional as F +from diffusers import DiffusionPipeline +from utils.llm_data import get_checkpoint -try: - from pathlib import Path - - from diffusers import DiffusionPipeline - - from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader - from tensorrt_llm.visual_gen.args import ( - AttentionConfig, - ParallelConfig, - TorchCompileConfig, - VisualGenArgs, - ) - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader +from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + ParallelConfig, + TorchCompileConfig, + VisualGenArgs, +) try: from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( @@ -61,7 +49,7 @@ _flash_attn_combine as _fa_combine, ) - _ATTN2D_AVAILABLE = MODULES_AVAILABLE and _fa4_fwd is not None and _fa_combine is not None + _ATTN2D_AVAILABLE = _fa4_fwd is not None and _fa_combine is not None except ImportError: _ATTN2D_AVAILABLE = False @@ -72,28 +60,7 @@ def _cleanup_mpi_env(): os.environ.pop("TLLM_DISABLE_MPI", None) -# ============================================================================= -# Path helpers (mirror tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py) -# ============================================================================= - - -def _llm_models_root() -> str: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN21_1_3B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN21_1_3B", "Wan2.1-T2V-1.3B-Diffusers") +WAN21_1_3B_SUBDIR = "Wan2.1-T2V-1.3B-Diffusers" # ============================================================================= @@ -146,8 +113,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs): def run_test_in_distributed(world_size: int, test_fn: Callable, **kwargs): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") # Spawn distributed workers via a helper that retries with a fresh master @@ -279,7 +244,7 @@ def _free(*objs) -> None: # ============================================================================= -def _logic_wan_cfg_ulysses_pvae(rank: int, world_size: int, *, checkpoint_path: str) -> None: +def _logic_wan_cfg_ulysses_pvae(rank: int, world_size: int, *, checkpoint_subdir: str) -> None: """End-to-end pipeline run with cfg=2, ulysses=2, parallel_vae=2. All ranks participate in the TRTLLM forward (CFG parallel, Ulysses, parallel @@ -288,6 +253,7 @@ def _logic_wan_cfg_ulysses_pvae(rank: int, world_size: int, *, checkpoint_path: """ assert world_size == 4, f"This test is hardcoded to world_size=4, got {world_size}" + checkpoint_path = get_checkpoint(checkpoint_subdir) trtllm_pipe = PipelineLoader(_build_parallel_args(checkpoint_path)).load(skip_warmup=True) trtllm_video = _capture_trtllm_video( trtllm_pipe, @@ -346,11 +312,12 @@ def _logic_wan_cfg_ulysses_pvae(rank: int, world_size: int, *, checkpoint_path: def _logic_wan_cfg2_attn2d2x1_ulysses2_pvae8( - rank: int, world_size: int, *, checkpoint_path: str + rank: int, world_size: int, *, checkpoint_subdir: str ) -> None: """End-to-end pipeline: cfg=2, ulysses=2, attn2d=2×1, async_ulysses, parallel_vae=8 (8 GPUs).""" assert world_size == 8, f"This test is hardcoded to world_size=8, got {world_size}" + checkpoint_path = get_checkpoint(checkpoint_subdir) trtllm_pipe = PipelineLoader(_build_cfg2_attn2d2x1_ulysses2_pvae8_args(checkpoint_path)).load( skip_warmup=True ) @@ -420,32 +387,20 @@ class TestWanPipelineParallel: def test_cfg2_ulysses2_pvae2(self): """world=4, cfg=2, ulysses=2, parallel_vae=2 vs HF reference.""" - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - if not os.path.exists(WAN21_1_3B_PATH): - pytest.skip( - f"Checkpoint not found: {WAN21_1_3B_PATH}. Set DIFFUSION_MODEL_PATH_WAN21_1_3B." - ) run_test_in_distributed( world_size=4, test_fn=_logic_wan_cfg_ulysses_pvae, - checkpoint_path=WAN21_1_3B_PATH, + checkpoint_subdir=WAN21_1_3B_SUBDIR, ) def test_cfg2_attn2d2x1_ulysses2_pvae8(self): """world=8, cfg=2, ulysses=2, attn2d=2×1, async_ulysses, parallel_vae=8 vs HF reference.""" - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if not _ATTN2D_AVAILABLE: pytest.skip("FA4 / flash_attn_combine JIT kernels not available") - if not os.path.exists(WAN21_1_3B_PATH): - pytest.skip( - f"Checkpoint not found: {WAN21_1_3B_PATH}. Set DIFFUSION_MODEL_PATH_WAN21_1_3B." - ) run_test_in_distributed( world_size=8, test_fn=_logic_wan_cfg2_attn2d2x1_ulysses2_pvae8, - checkpoint_path=WAN21_1_3B_PATH, + checkpoint_subdir=WAN21_1_3B_SUBDIR, ) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py index fbedc1aa1410..6f3aef886476 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py @@ -34,21 +34,16 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - AttentionConfig, - DiffusionModelConfig, - TorchCompileConfig, - create_attention_metadata_state, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm.models.modeling_utils import QuantConfig - - from .tp_shard_utils import copy_tp_parameter +from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + TorchCompileConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm.models.modeling_utils import QuantConfig - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from .tp_shard_utils import copy_tp_parameter @pytest.fixture(autouse=True, scope="module") @@ -104,9 +99,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): """Run a test function in a distributed environment.""" - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py index 814c88990f34..0bda70d2729d 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_transformer_parallel.py @@ -35,19 +35,14 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.config import ( - AttentionConfig, - DiffusionModelConfig, - TorchCompileConfig, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - - from .tp_shard_utils import copy_tp_parameter +from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + TorchCompileConfig, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from .tp_shard_utils import copy_tp_parameter try: from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( @@ -101,8 +96,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True, **kwargs): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") backend = "nccl" if use_cuda else "gloo" @@ -312,10 +305,7 @@ def _logic_wan_transformer_parallel_vs_single_gpu( torch.manual_seed(SEED_WEIGHTS) dist_config = _make_model_config(pretrained_cfg, backend="FA4", **parallel_cfg_kwargs) - try: - dist_model = WanTransformer3DModel(dist_config).to(device).to(dtype) - except (ImportError, ValueError, NotImplementedError) as e: - pytest.skip(f"[{label}] Parallel backend unavailable: {e}") + dist_model = WanTransformer3DModel(dist_config).to(device).to(dtype) if dist_config.visual_gen_mapping.tp_size > 1: _copy_ref_weights_to_tp(ref_model, dist_model, pretrained_cfg) @@ -379,10 +369,7 @@ def _logic_wan_transformer_parallel_forward_sanity( torch.manual_seed(SEED_WEIGHTS) config = _make_model_config(pretrained_cfg, backend="FA4", **parallel_cfg_kwargs) - try: - model = WanTransformer3DModel(config).to(device).to(dtype) - except (ImportError, ValueError, NotImplementedError) as e: - pytest.skip(f"[{label}] Parallel backend unavailable: {e}") + model = WanTransformer3DModel(config).to(device).to(dtype) _stabilize_model_weights(model) torch.manual_seed(SEED_INPUT) @@ -416,11 +403,10 @@ def _logic_wan_transformer_parallel_forward_sanity( class TestWanTransformerParallel: """Transformer-only WAN correctness across parallel topologies.""" - def _skip_if_unavailable(self): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") - if not _flash_attn4_available: - pytest.skip("FlashAttn4 JIT kernels not available") + def _require_flash_attn4(self) -> None: + assert _flash_attn4_available, ( + "FlashAttn4 JIT kernels not available; expected on the Blackwell CI runner" + ) @pytest.mark.parametrize( "label,parallel_cfg_kwargs", @@ -428,7 +414,7 @@ def _skip_if_unavailable(self): ids=[name for name, _ in _WAN_8GPU_PARALLEL_COMBINATIONS], ) def test_parallel_all_combinations_vs_single_gpu_8gpu(self, label, parallel_cfg_kwargs): - self._skip_if_unavailable() + self._require_flash_attn4() run_test_in_distributed( world_size=8, test_fn=_logic_wan_transformer_parallel_vs_single_gpu, @@ -437,7 +423,7 @@ def test_parallel_all_combinations_vs_single_gpu_8gpu(self, label, parallel_cfg_ ) def test_parallel_attn2d_2x2_forward_sanity_4gpu(self): - self._skip_if_unavailable() + self._require_flash_attn4() run_test_in_distributed( world_size=4, test_fn=_logic_wan_transformer_parallel_forward_sanity, @@ -446,7 +432,7 @@ def test_parallel_attn2d_2x2_forward_sanity_4gpu(self): ) def test_parallel_attn2d_2x2_vs_single_gpu_4gpu(self): - self._skip_if_unavailable() + self._require_flash_attn4() run_test_in_distributed( world_size=4, test_fn=_logic_wan_transformer_parallel_vs_single_gpu, @@ -456,9 +442,10 @@ def test_parallel_attn2d_2x2_vs_single_gpu_4gpu(self): def test_parallel_attn2d_2x2_ulysses2_vs_single_gpu_8gpu(self): """world=8, attn2d=2×2, ulysses=2 vs single-GPU FA4 reference.""" - self._skip_if_unavailable() - if not _attn2d_available: - pytest.skip("FA4 / flash_attn_combine JIT kernels not available") + self._require_flash_attn4() + assert _attn2d_available, ( + "FA4 / flash_attn_combine JIT kernels not available; expected on the Blackwell CI runner" + ) run_test_in_distributed( world_size=8, test_fn=_logic_wan_transformer_parallel_vs_single_gpu, @@ -471,7 +458,7 @@ def test_parallel_attn2d_2x2_ulysses2_vs_single_gpu_8gpu(self): ) def test_parallel_ring4_vs_single_gpu_4gpu(self): - self._skip_if_unavailable() + self._require_flash_attn4() run_test_in_distributed( world_size=4, test_fn=_logic_wan_transformer_parallel_vs_single_gpu, @@ -480,7 +467,7 @@ def test_parallel_ring4_vs_single_gpu_4gpu(self): ) def test_parallel_ring2_ul2_vs_single_gpu_4gpu(self): - self._skip_if_unavailable() + self._require_flash_attn4() run_test_in_distributed( world_size=4, test_fn=_logic_wan_transformer_parallel_vs_single_gpu, diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py index 00dd752d902e..a04eca96a0c1 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py @@ -24,7 +24,6 @@ import gc import json import os -from pathlib import Path from types import SimpleNamespace from typing import Callable @@ -34,28 +33,23 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp +from utils.llm_data import get_checkpoint -try: - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import ( - VSAMetadataBuilder, - _cute_dsl_import_error, - set_vsa_forward_context, - ) - from tensorrt_llm._torch.visual_gen.config import ( - AttentionConfig, - DiffusionModelConfig, - TorchCompileConfig, - ) - from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping - from tensorrt_llm._utils import get_free_port - from tensorrt_llm.visual_gen.sparse_attention import VideoSparseAttentionConfig +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import ( + VSAMetadataBuilder, + _cute_dsl_import_error, + set_vsa_forward_context, +) +from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + TorchCompileConfig, +) +from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping +from tensorrt_llm._utils import get_free_port +from tensorrt_llm.visual_gen.sparse_attention import VideoSparseAttentionConfig - MODULES_AVAILABLE = True - _cute_dsl_available = _cute_dsl_import_error is None -except ImportError: - MODULES_AVAILABLE = False - _cute_dsl_available = False - _cute_dsl_import_error = None +_cute_dsl_available = _cute_dsl_import_error is None @pytest.fixture(autouse=True, scope="module") @@ -95,8 +89,6 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs): def run_test_in_distributed(world_size: int, test_fn: Callable, **kwargs): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") port = get_free_port() @@ -205,26 +197,7 @@ def _log(rank: int, msg: str) -> None: # ============================================================================= -def _llm_models_root() -> str | None: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - return str(root) if root.exists() else None - - -def _checkpoint(env_var: str, default_name: str) -> str | None: - if env_var in os.environ: - return os.environ[env_var] - models_root = _llm_models_root() - return os.path.join(models_root, default_name) if models_root is not None else None - - -WAN21_VSA_14B_PATH = _checkpoint( - "DIFFUSION_MODEL_PATH_WAN21_VSA_T2V_14B_720P", - "Wan2.1-VSA-T2V-14B-720P-Diffusers", -) +WAN21_VSA_14B_SUBDIR = "Wan2.1-VSA-T2V-14B-720P-Diffusers" # ============================================================================= @@ -238,9 +211,11 @@ def _logic_vsa_ulysses_real_model( *, ulysses_size: int, cfg_size: int = 1, - checkpoint_path: str, + checkpoint_subdir: str, label: str, ) -> None: + checkpoint_path = get_checkpoint(checkpoint_subdir) + from tensorrt_llm._torch.visual_gen.checkpoints.weight_loader import WeightLoader from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel @@ -254,20 +229,14 @@ def _logic_vsa_ulysses_real_model( _log(rank, f"[{label}] building ref_model (world_size=1)") ref_config = _make_vsa_model_config(pretrained_dict) - try: - ref_model = WanTransformer3DModel(ref_config).to(device).to(dtype) - except (ImportError, ValueError, NotImplementedError) as e: - pytest.skip(f"[{label}] VSA ref model unavailable: {e}") + ref_model = WanTransformer3DModel(ref_config).to(device).to(dtype) _log(rank, f"[{label}] ref_model created") _log(rank, f"[{label}] building dist_model (ul={ulysses_size} cfg={cfg_size})") dist_config = _make_vsa_model_config( pretrained_dict, ulysses_size=ulysses_size, cfg_size=cfg_size ) - try: - dist_model = WanTransformer3DModel(dist_config).to(device).to(dtype) - except (ImportError, ValueError, NotImplementedError) as e: - pytest.skip(f"[{label}] VSA parallel model unavailable: {e}") + dist_model = WanTransformer3DModel(dist_config).to(device).to(dtype) _log(rank, f"[{label}] dist_model created") _log(rank, f"[{label}] loading checkpoint weights") @@ -338,15 +307,8 @@ class TestWanVsaUlyssesRealModel: """ def _skip_if_unavailable(self): - if not MODULES_AVAILABLE: - pytest.skip("Required modules not available") if not _cute_dsl_available: pytest.skip(f"CUTEDSL not available (requires Blackwell GPU): {_cute_dsl_import_error}") - if WAN21_VSA_14B_PATH is None or not os.path.isdir(WAN21_VSA_14B_PATH): - pytest.skip( - "Wan2.1-VSA-T2V-14B-720P checkpoint not found; " - "set DIFFUSION_MODEL_PATH_WAN21_VSA_T2V_14B_720P or LLM_MODELS_ROOT" - ) def test_real_model_vsa_cfg2_ulysses4_vs_single_gpu(self): """world=8, cfg=2, ulysses=4, real VSA 14B weights, VSA sparsity=0.9.""" @@ -356,7 +318,7 @@ def test_real_model_vsa_cfg2_ulysses4_vs_single_gpu(self): test_fn=_logic_vsa_ulysses_real_model, ulysses_size=4, cfg_size=2, - checkpoint_path=WAN21_VSA_14B_PATH, + checkpoint_subdir=WAN21_VSA_14B_SUBDIR, label="real,cfg=2,ul=4", ) diff --git a/tests/unittest/_torch/visual_gen/test_cache_dit.py b/tests/unittest/_torch/visual_gen/test_cache_dit.py index e836cad6f995..35f6c61dcf6b 100644 --- a/tests/unittest/_torch/visual_gen/test_cache_dit.py +++ b/tests/unittest/_torch/visual_gen/test_cache_dit.py @@ -4,21 +4,21 @@ """Tests for Cache-DiT in visual generation. Wan 2.2 step-split logic is covered with small CPU-side tests. Wan, FLUX, and LTX-2 -integration tests run on GPU only when cache_dit is installed, CUDA is available, and -checkpoints can be resolved (TRTLLM_CACHE_DIT_* env vars or the fallbacks inside each test). +integration tests run on GPU when CUDA is available; cache-dit is a required dependency +(requirements.txt), so a missing install fails loudly rather than skipping. The checkpoints +they need are resolved under LLM_MODELS_ROOT and fail loudly when unavailable. """ from __future__ import annotations import contextlib -import importlib.util import logging -import os from types import SimpleNamespace from unittest.mock import patch import pytest import torch +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.cache.cache_dit_enablers import split_wan22_inference_steps @@ -26,125 +26,16 @@ # Integration prerequisites # --------------------------------------------------------------------------- -_CACHE_DIT_SPEC = importlib.util.find_spec("cache_dit") - -requires_cache_dit = pytest.mark.skipif( - _CACHE_DIT_SPEC is None, - reason="optional dependency cache-dit not installed (pip install cache-dit)", -) - requires_cuda = pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA not available", ) -_WAN_SUBPATH = "Wan2.1-T2V-1.3B-Diffusers" -_FLUX_SUBPATH = "FLUX.1-dev" -_QWEN_IMAGE_SUBPATH = "qwen-image" -_LTX2_DIR = "LTX-2" -_LTX2_WEIGHTS_FILE = "ltx-2-19b-dev.safetensors" -_LTX2_TEXT_ENCODER_SUBPATH = "gemma-3-12b-it" -# Same share as other Gemma checkpoints (e.g. google/gemma-3-1b-it -> gemma/gemma-3-1b-it). -_LTX2_TEXT_ENCODER_RELATIVE_PATHS = ( - _LTX2_TEXT_ENCODER_SUBPATH, - os.path.join("gemma", _LTX2_TEXT_ENCODER_SUBPATH), -) -# Default NFS layout used on CI runners (override via TRTLLM_CACHE_DIT_*_CHECKPOINT). -_CI_DEFAULT_LLM_MODELS = "/home/scratch.trt_llm_data_ci/llm-models" -_DEFAULT_WAN_CHECKPOINT = os.path.join(_CI_DEFAULT_LLM_MODELS, _WAN_SUBPATH) -_DEFAULT_FLUX_CHECKPOINT = os.path.join(_CI_DEFAULT_LLM_MODELS, _FLUX_SUBPATH) -_DEFAULT_LTX2_CHECKPOINT = os.path.join(_CI_DEFAULT_LLM_MODELS, _LTX2_DIR, _LTX2_WEIGHTS_FILE) -_DEFAULT_LTX2_TEXT_ENCODER = os.path.join(_CI_DEFAULT_LLM_MODELS, _LTX2_TEXT_ENCODER_SUBPATH) -_DEFAULT_QWEN_IMAGE_CHECKPOINT = os.path.join(_CI_DEFAULT_LLM_MODELS, _QWEN_IMAGE_SUBPATH) - - -def _resolve_qwen_image_checkpoint() -> str | None: - """Qwen-Image: explicit env, CI default tree, then LLM_MODELS_ROOT.""" - explicit = os.environ.get("TRTLLM_CACHE_DIT_QWEN_IMAGE_CHECKPOINT", "").strip() - if explicit and os.path.isdir(explicit): - return os.path.abspath(explicit) - if os.path.isdir(_DEFAULT_QWEN_IMAGE_CHECKPOINT): - return os.path.abspath(_DEFAULT_QWEN_IMAGE_CHECKPOINT) - root = os.environ.get("LLM_MODELS_ROOT", "").strip() - if root: - cand = os.path.join(root, _QWEN_IMAGE_SUBPATH) - if os.path.isdir(cand): - return os.path.abspath(cand) - return None - - -def _resolve_wan_checkpoint() -> str | None: - """Wan 2.1 1.3B: explicit env, then CI default tree, then LLM_MODELS_ROOT.""" - explicit = os.environ.get("TRTLLM_CACHE_DIT_WAN_CHECKPOINT", "").strip() - if explicit and os.path.isdir(explicit): - return os.path.abspath(explicit) - if os.path.isdir(_DEFAULT_WAN_CHECKPOINT): - return os.path.abspath(_DEFAULT_WAN_CHECKPOINT) - root = os.environ.get("LLM_MODELS_ROOT", "").strip() - if root: - cand = os.path.join(root, _WAN_SUBPATH) - if os.path.isdir(cand): - return os.path.abspath(cand) - return None - - -def _resolve_flux_checkpoint() -> str | None: - """FLUX.1 dev tree: explicit env, FLUX1_MODEL_PATH, CI default, then LLM_MODELS_ROOT.""" - explicit = os.environ.get("TRTLLM_CACHE_DIT_FLUX_CHECKPOINT", "").strip() - if explicit and os.path.isdir(explicit): - return os.path.abspath(explicit) - flux1 = os.environ.get("FLUX1_MODEL_PATH", "").strip() - if flux1 and os.path.isdir(flux1): - return os.path.abspath(flux1) - if os.path.isdir(_DEFAULT_FLUX_CHECKPOINT): - return os.path.abspath(_DEFAULT_FLUX_CHECKPOINT) - root = os.environ.get("LLM_MODELS_ROOT", "").strip() - if root: - cand = os.path.join(root, _FLUX_SUBPATH) - if os.path.isdir(cand): - return os.path.abspath(cand) - return None - - -def _resolve_ltx2_checkpoint() -> str | None: - """LTX-2 weights file: explicit env, LTX2_MODEL_PATH, CI default, then LLM_MODELS_ROOT (same tiers as Wan/Flux).""" - explicit = os.environ.get("TRTLLM_CACHE_DIT_LTX2_CHECKPOINT", "").strip() - if explicit: - if os.path.isfile(explicit): - return os.path.abspath(explicit) - if os.path.isdir(explicit): - cand = os.path.join(explicit, _LTX2_WEIGHTS_FILE) - if os.path.isfile(cand): - return os.path.abspath(cand) - ltx2_model = os.environ.get("LTX2_MODEL_PATH", "").strip() - if ltx2_model and os.path.isfile(ltx2_model): - return os.path.abspath(ltx2_model) - if os.path.isfile(_DEFAULT_LTX2_CHECKPOINT): - return os.path.abspath(_DEFAULT_LTX2_CHECKPOINT) - root = os.environ.get("LLM_MODELS_ROOT", "").strip() - if root: - cand = os.path.join(root, _LTX2_DIR, _LTX2_WEIGHTS_FILE) - if os.path.isfile(cand): - return os.path.abspath(cand) - return None - - -def _resolve_ltx2_text_encoder() -> str | None: - """Gemma text encoder directory: explicit env, CI default, then LLM_MODELS_ROOT.""" - explicit = os.environ.get("TRTLLM_CACHE_DIT_LTX2_TEXT_ENCODER", "").strip() - if explicit and os.path.isdir(explicit): - return os.path.abspath(explicit) - for rel in _LTX2_TEXT_ENCODER_RELATIVE_PATHS: - cand = os.path.join(_CI_DEFAULT_LLM_MODELS, rel) - if os.path.isdir(cand): - return os.path.abspath(cand) - root = os.environ.get("LLM_MODELS_ROOT", "").strip() - if root: - for rel in _LTX2_TEXT_ENCODER_RELATIVE_PATHS: - cand = os.path.join(root, rel) - if os.path.isdir(cand): - return os.path.abspath(cand) - return None +_WAN_SUBDIR = "Wan2.1-T2V-1.3B-Diffusers" +_FLUX_SUBDIR = "FLUX.1-dev" +_QWEN_IMAGE_SUBDIR = "qwen-image" +_LTX2_CHECKPOINT_SUBDIR = "LTX-2/ltx-2-19b-dev.safetensors" +_LTX2_TEXT_ENCODER_SUBDIR = "gemma/gemma-3-12b-it" @contextlib.contextmanager @@ -276,7 +167,6 @@ def test_boundary_one_all_timesteps_low(self): # --------------------------------------------------------------------------- -@requires_cache_dit @requires_cuda class TestCacheDiTRealPipelineForward: """Wan, FLUX.1, and LTX-2 use the CI llm-models tree when checkpoints are present. @@ -341,13 +231,7 @@ def _load_visual_gen_pipeline( return loader.load(skip_warmup=True) def test_wan_cache_dit_skips_blocks_after_forward(self): - ckpt = _resolve_wan_checkpoint() - if ckpt is None: - pytest.skip( - "Wan 2.1 1.3B not found: set TRTLLM_CACHE_DIT_WAN_CHECKPOINT, " - f"install under {_DEFAULT_WAN_CHECKPOINT}, " - f"or under $LLM_MODELS_ROOT/{_WAN_SUBPATH}" - ) + ckpt = get_checkpoint(_WAN_SUBDIR) pipeline = None with _suppress_stdlib_logging_for_cache_dit(): @@ -388,13 +272,7 @@ def test_wan_cache_dit_enabled_after_torch_compile(self): torch.compile would contribute nothing — a perf-only regression that is invisible to the correctness assertions of the compile-off tests. """ - ckpt = _resolve_wan_checkpoint() - if ckpt is None: - pytest.skip( - "Wan 2.1 1.3B not found: set TRTLLM_CACHE_DIT_WAN_CHECKPOINT, " - f"install under {_DEFAULT_WAN_CHECKPOINT}, " - f"or under $LLM_MODELS_ROOT/{_WAN_SUBPATH}" - ) + ckpt = get_checkpoint(_WAN_SUBDIR) pipeline = None with _suppress_stdlib_logging_for_cache_dit(): @@ -440,13 +318,7 @@ def test_wan_cache_dit_enabled_after_torch_compile(self): torch._dynamo.reset() def test_flux_cache_dit_skips_blocks_after_forward(self): - ckpt = _resolve_flux_checkpoint() - if ckpt is None: - pytest.skip( - "FLUX.1-dev not found: set TRTLLM_CACHE_DIT_FLUX_CHECKPOINT or FLUX1_MODEL_PATH, " - f"install under {_DEFAULT_FLUX_CHECKPOINT}, " - f"or under $LLM_MODELS_ROOT/{_FLUX_SUBPATH}" - ) + ckpt = get_checkpoint(_FLUX_SUBDIR) pipeline = None with _suppress_stdlib_logging_for_cache_dit(): @@ -454,7 +326,7 @@ def test_flux_cache_dit_skips_blocks_after_forward(self): pipeline = self._load_visual_gen_pipeline(ckpt) name = pipeline.__class__.__name__ if name not in ("FluxPipeline", "Flux2Pipeline"): - pytest.skip(f"Checkpoint resolved to {name}, not a FLUX visual_gen pipeline") + pytest.fail(f"Checkpoint resolved to {name}, not a FLUX visual_gen pipeline") assert pipeline.cache_accelerator is not None assert pipeline.cache_accelerator.is_enabled() @@ -481,25 +353,8 @@ def test_flux_cache_dit_skips_blocks_after_forward(self): self._teardown_cache_dit(pipeline) def test_ltx2_cache_dit_skips_blocks_after_forward(self): - ckpt = _resolve_ltx2_checkpoint() - text_enc = _resolve_ltx2_text_encoder() - if ckpt is None or text_enc is None: - missing = [] - if ckpt is None: - missing.append("LTX-2 checkpoint") - if text_enc is None: - missing.append("Gemma text encoder") - pytest.skip( - f"Missing {' and '.join(missing)}: set " - "TRTLLM_CACHE_DIT_LTX2_CHECKPOINT (file or directory with " - f"{_LTX2_WEIGHTS_FILE}) and TRTLLM_CACHE_DIT_LTX2_TEXT_ENCODER, " - "or LTX2_MODEL_PATH, or stage under CI tree " - f"{_DEFAULT_LTX2_CHECKPOINT} and one of " - f"{', '.join(os.path.join(_CI_DEFAULT_LLM_MODELS, r) for r in _LTX2_TEXT_ENCODER_RELATIVE_PATHS)} " - f"(same as Wan/Flux under {_CI_DEFAULT_LLM_MODELS}), " - f"or $LLM_MODELS_ROOT/{_LTX2_DIR}/{_LTX2_WEIGHTS_FILE} and " - f"$LLM_MODELS_ROOT/<{' or '.join(_LTX2_TEXT_ENCODER_RELATIVE_PATHS)}>" - ) + ckpt = get_checkpoint(_LTX2_CHECKPOINT_SUBDIR) + text_enc = get_checkpoint(_LTX2_TEXT_ENCODER_SUBDIR) pipeline = None with _suppress_stdlib_logging_for_cache_dit(): @@ -507,7 +362,7 @@ def test_ltx2_cache_dit_skips_blocks_after_forward(self): pipeline = self._load_visual_gen_pipeline(ckpt, text_encoder_path=text_enc) name = pipeline.__class__.__name__ if name != "LTX2Pipeline": - pytest.skip(f"Checkpoint resolved to {name}, not LTX2Pipeline") + pytest.fail(f"Checkpoint resolved to {name}, not LTX2Pipeline") assert pipeline.cache_accelerator is not None assert pipeline.cache_accelerator.is_enabled() @@ -537,13 +392,7 @@ def test_ltx2_cache_dit_skips_blocks_after_forward(self): self._teardown_cache_dit(pipeline) def test_qwen_image_cache_dit_skips_blocks_after_forward(self): - ckpt = _resolve_qwen_image_checkpoint() - if ckpt is None: - pytest.skip( - "Qwen-Image not found: set TRTLLM_CACHE_DIT_QWEN_IMAGE_CHECKPOINT, " - f"install under {_DEFAULT_QWEN_IMAGE_CHECKPOINT}, " - f"or under $LLM_MODELS_ROOT/{_QWEN_IMAGE_SUBPATH}" - ) + ckpt = get_checkpoint(_QWEN_IMAGE_SUBDIR) pipeline = None with _suppress_stdlib_logging_for_cache_dit(): @@ -551,7 +400,7 @@ def test_qwen_image_cache_dit_skips_blocks_after_forward(self): pipeline = self._load_visual_gen_pipeline(ckpt) name = pipeline.__class__.__name__ if name != "QwenImagePipeline": - pytest.skip(f"Checkpoint resolved to {name}, not QwenImagePipeline") + pytest.fail(f"Checkpoint resolved to {name}, not QwenImagePipeline") assert pipeline.cache_accelerator is not None assert pipeline.cache_accelerator.is_enabled() @@ -589,7 +438,6 @@ def forward(self, hidden_states, encoder_hidden_states): return hidden_states, encoder_hidden_states -@requires_cache_dit class TestFluxEnablerCompiledBlockCheckFlags: """enable_cache_dit_for_flux must disable cache_dit's inspect-based forward-pattern checks exactly when blocks are torch.compile wrappers. @@ -648,7 +496,6 @@ def test_eager_blocks_keep_pattern_checks(self, is_flux2): # --------------------------------------------------------------------------- -@requires_cache_dit class TestCacheDiTEnablerRegistry: def test_enabler_keys_name_registered_pipeline_classes(self): """Every CUSTOM_CACHE_DIT_ENABLERS key must name a real registered diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py index 26f6b1184ab7..ae59e715af03 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -8,10 +8,6 @@ K-norm, native flow schedule parity against cosmos-framework, strict weight loading, and per-family defaults. Checkpoint-gated tests cover the real checkpoint (tokenizer, recipe/scheduler wiring, load + forward). - -Override checkpoint: - DIFFUSION_MODEL_PATH_COSMOS3_EDGE=/path/to/Cosmos3-Edge \\ - pytest tests/unittest/_torch/visual_gen/test_cosmos3_edge.py -v """ import gc @@ -25,6 +21,7 @@ import pytest import torch from diffusers import UniPCMultistepScheduler +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig @@ -78,16 +75,8 @@ def _cleanup_gpu(): def _require_edge_checkpoint() -> str: """Resolve the Edge checkpoint lazily so unit tests collect and run on - machines without model storage; only checkpoint-gated tests skip.""" - path = os.environ.get("DIFFUSION_MODEL_PATH_COSMOS3_EDGE") - if not path: - root = Path(os.environ.get("LLM_MODELS_ROOT", "/home/scratch.trt_llm_data_ci/llm-models/")) - if not root.exists(): - root = Path("/scratch/trt_llm_data/llm-models/") - path = str(root / "Cosmos3-Edge") - if not os.path.isdir(path): - pytest.skip(f"Checkpoint not found: {path}") - return path + machines without model storage; only checkpoint-gated tests fail.""" + return get_checkpoint("Cosmos3-Edge") def _reduced_edge_config() -> SimpleNamespace: diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index b9fcd3daa2e8..592661d9ec86 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -18,10 +18,6 @@ Run prompt metadata unit tests (no GPU): pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -k FormatPromptWithMetadata - -Override checkpoint: - DIFFUSION_MODEL_PATH_COSMOS3=/path/to/Cosmos3-Nano \\ - pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -s """ import gc @@ -37,6 +33,7 @@ import PIL.Image import pytest import torch +from utils.llm_data import get_checkpoint import tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 as pipe_mod from tensorrt_llm._torch.visual_gen.models.cosmos3.action import resolve_action_content_size @@ -87,24 +84,6 @@ def _cleanup_gpu(): torch.cuda.empty_cache() -def _llm_models_root() -> str: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -COSMOS3_NANO_PATH = _checkpoint("DIFFUSION_MODEL_PATH_COSMOS3", "Cosmos3-Nano") - PROMPT = "A serene mountain lake at sunrise with mist rising from the water." NUM_STEPS = 4 SEED = 42 @@ -128,11 +107,9 @@ def _checkpoint(env_var: str, default_name: str) -> str: def _require_checkpoint() -> str: - if not COSMOS3_NANO_PATH or not os.path.exists(COSMOS3_NANO_PATH): - pytest.skip(f"Checkpoint not found: {COSMOS3_NANO_PATH}") if not torch.cuda.is_available(): pytest.skip("CUDA not available") - return COSMOS3_NANO_PATH + return get_checkpoint("Cosmos3-Nano") def _load_pipeline(checkpoint_path: str, **visual_gen_kwargs): @@ -232,14 +209,14 @@ def _assert_valid_audio( def _require_audio_pipeline(pipeline) -> None: if not getattr(pipeline, "audio_gen", False): - pytest.skip("Checkpoint does not enable audio generation") + pytest.fail("Checkpoint does not enable audio generation") if not hasattr(pipeline, "audio_tokenizer"): - pytest.skip("Audio tokenizer was not loaded for this pipeline") + pytest.fail("Audio tokenizer was not loaded for this pipeline") def _require_action_pipeline(pipeline) -> None: if not getattr(pipeline, "action_gen", False): - pytest.skip("Checkpoint does not enable action generation") + pytest.fail("Checkpoint does not enable action generation") def _assert_valid_action(action: torch.Tensor, *, raw_action_dim: int, chunk_size: int): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_t2v_offload.py b/tests/unittest/_torch/visual_gen/test_cosmos3_t2v_offload.py index c43fd275d723..2601fb828e28 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_t2v_offload.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_t2v_offload.py @@ -5,17 +5,13 @@ Loads the full model and calls pipeline.forward() with cpu_offload_config.enable=True; offloaded output is checked against the non-offload TRT-LLM baseline. CUDA graphs -with offloading must raise. These skip unless the checkpoint is available. +with offloading must raise. Fast, GPU-free offload unit tests (including the Cosmos3-specific offload wiring) live in test_offloading.py. Run: pytest tests/unittest/_torch/visual_gen/test_cosmos3_t2v_offload.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_COSMOS3=/path/to/cosmos3 \\ - pytest tests/unittest/_torch/visual_gen/test_cosmos3_t2v_offload.py -v -s """ import os @@ -30,11 +26,11 @@ os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") import gc -from pathlib import Path import pytest import torch import torch.nn.functional as F +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( COSMOS3_GENERATOR_OFFLOAD_COMPONENT, @@ -60,24 +56,7 @@ def _cleanup_gpu(): torch.cuda.empty_cache() -def _llm_models_root() -> Path: - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - else: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return root - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or str(_llm_models_root() / default_name) - - -COSMOS3_PATH = _checkpoint("DIFFUSION_MODEL_PATH_COSMOS3", "Cosmos3-Nano") +COSMOS3_SUBDIR = "Cosmos3-Nano" # Mirrors a validated manual offload run; reduced steps keep the test fast. INFER_PROMPT = "A cute cat playing piano." @@ -106,8 +85,6 @@ def _make_pipeline(checkpoint_path: str, *, enable_offload: bool = False): from tensorrt_llm.visual_gen.args import CpuOffloadConfig as ArgsOffloadConfig from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), @@ -139,7 +116,7 @@ def _capture_video(pipe) -> torch.Tensor: @pytest.fixture def cosmos3_offload_pipeline(): - pipeline = _make_pipeline(COSMOS3_PATH, enable_offload=True) + pipeline = _make_pipeline(get_checkpoint(COSMOS3_SUBDIR), enable_offload=True) yield pipeline del pipeline gc.collect() @@ -163,13 +140,13 @@ def test_cosmos3_offload_forward(self, cosmos3_offload_pipeline): print(f" stages: {pipeline.offloader.stages()}") def test_cosmos3_offload_matches_baseline(self): - offload_pipe = _make_pipeline(COSMOS3_PATH, enable_offload=True) + offload_pipe = _make_pipeline(get_checkpoint(COSMOS3_SUBDIR), enable_offload=True) offload_video = _capture_video(offload_pipe) del offload_pipe gc.collect() torch.cuda.empty_cache() - baseline_pipe = _make_pipeline(COSMOS3_PATH, enable_offload=False) + baseline_pipe = _make_pipeline(get_checkpoint(COSMOS3_SUBDIR), enable_offload=False) baseline_video = _capture_video(baseline_pipe) del baseline_pipe gc.collect() @@ -201,10 +178,8 @@ def test_cosmos3_raises_if_cuda_graph_and_offload_enabled(self): from tensorrt_llm.visual_gen.args import CpuOffloadConfig as ArgsOffloadConfig from tensorrt_llm.visual_gen.args import CudaGraphConfig, TorchCompileConfig, VisualGenArgs - if not os.path.exists(COSMOS3_PATH): - pytest.skip(f"Checkpoint not found: {COSMOS3_PATH}") args = VisualGenArgs( - model=COSMOS3_PATH, + model=get_checkpoint(COSMOS3_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), cuda_graph_config=CudaGraphConfig(enable=True), cpu_offload_config=ArgsOffloadConfig(enable=True, stages=INFER_OFFLOAD_STAGES), diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index e145f760d2bc..9a8105b14d57 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -11,21 +11,17 @@ Run all: pytest tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py -v -s - -Override checkpoint: - DIFFUSION_MODEL_PATH_COSMOS3=/path/to/Cosmos3-Nano \\ - pytest tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py -v -s """ import gc import os -from pathlib import Path from types import SimpleNamespace os.environ["TLLM_DISABLE_MPI"] = "1" import pytest import torch +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig, DiffusionPipelineConfig @@ -57,24 +53,6 @@ def _cleanup_gpu(): torch.cuda.empty_cache() -def _llm_models_root() -> str: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch/trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -COSMOS3_NANO_PATH = _checkpoint("DIFFUSION_MODEL_PATH_COSMOS3", "Cosmos3-Nano") - DEVICE = "cuda" DTYPE = torch.bfloat16 _NUM_TRAIN_TIMESTEPS = 1000.0 @@ -97,12 +75,11 @@ def _transformer_config_path(checkpoint_dir: str) -> str: def _require_checkpoint() -> str: - if not COSMOS3_NANO_PATH or not os.path.isdir(COSMOS3_NANO_PATH): - pytest.skip(f"Checkpoint not found: {COSMOS3_NANO_PATH}") - config_path = _transformer_config_path(COSMOS3_NANO_PATH) + checkpoint_dir = get_checkpoint("Cosmos3-Nano") + config_path = _transformer_config_path(checkpoint_dir) if not os.path.isfile(config_path): - pytest.skip(f"Transformer config not found: {config_path}") - return COSMOS3_NANO_PATH + pytest.fail(f"Transformer config not found: {config_path}") + return checkpoint_dir def _load_model_config(checkpoint_dir: str) -> DiffusionModelConfig: diff --git a/tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py b/tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py index 846ab546733a..b9e997437df1 100644 --- a/tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py @@ -15,11 +15,7 @@ """Pipeline-level tests for WanDMDPipeline. Run: - DIFFUSION_MODEL_PATH_FASTWAN=/path/to/checkpoint \\ - pytest tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py -v -s - -Override checkpoint path via DIFFUSION_MODEL_PATH_FASTWAN env var or place -the checkpoint at $LLM_MODELS_ROOT/FastWan2.2-TI2V-5B-FullAttn-Diffusers + pytest tests/unittest/_torch/visual_gen/test_fastwan_pipeline.py -v -s """ import gc @@ -29,6 +25,7 @@ import pytest import torch +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs @@ -40,25 +37,13 @@ def _cleanup_mpi_env(): os.environ.pop("TLLM_DISABLE_MPI", None) -def _checkpoint_path() -> str: - if "DIFFUSION_MODEL_PATH_FASTWAN" in os.environ: - return os.environ["DIFFUSION_MODEL_PATH_FASTWAN"] - root = os.environ.get( - "LLM_MODELS_ROOT", - "/home/scratch.trt_llm_data_ci/llm-models", - ) - return os.path.join(root, "FastWan2.2-TI2V-5B-FullAttn-Diffusers") - - -FASTWAN_PATH = _checkpoint_path() +FASTWAN_SUBDIR = "FastWan2.2-TI2V-5B-FullAttn-Diffusers" @pytest.fixture(scope="module") def fastwan_pipeline(): - if not os.path.exists(FASTWAN_PATH): - pytest.skip(f"Checkpoint not found: {FASTWAN_PATH}") args = VisualGenArgs( - model=FASTWAN_PATH, + model=get_checkpoint(FASTWAN_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -109,10 +94,8 @@ class TestFastWanFP8: """ def test_fp8_trtllm(self): - if not os.path.exists(FASTWAN_PATH): - pytest.skip(f"Checkpoint not found: {FASTWAN_PATH}") args = VisualGenArgs( - model=FASTWAN_PATH, + model=get_checkpoint(FASTWAN_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), quant_config={"quant_algo": "FP8", "dynamic": True}, attention_config=AttentionConfig(backend="TRTLLM"), diff --git a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py index eb99617ebdaf..786bb2af4d94 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py @@ -13,9 +13,9 @@ - Multi-GPU parallelism (Ulysses sequence parallelism, 2+ GPUs) """ +import functools import gc import os -from pathlib import Path import numpy as np import PIL.Image @@ -24,6 +24,7 @@ import torch.distributed as dist import torch.multiprocessing as mp import torch.nn.functional as F +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader @@ -43,28 +44,16 @@ ] -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if it is set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "You shall set LLM_MODELS_ROOT env or be able to access scratch.trt_llm_data to run this test" - ) - return str(root) +# Checkpoint paths for integration tests, resolved lazily on first test access so +# a missing checkpoint fails the tests that need it rather than whole-file collection. +@functools.lru_cache +def _flux1_path() -> str: + return get_checkpoint("FLUX.1-dev") -# Checkpoint paths for integration tests -FLUX1_CHECKPOINT_PATH = os.environ.get( - "FLUX1_MODEL_PATH", - os.path.join(_llm_models_root(), "FLUX.1-dev"), -) -FLUX2_CHECKPOINT_PATH = os.environ.get( - "FLUX2_MODEL_PATH", - os.path.join(_llm_models_root(), "FLUX.2-dev"), -) +@functools.lru_cache +def _flux2_path() -> str: + return get_checkpoint("FLUX.2-dev") def _get_flux_transformer_inputs(transformer, device="cuda", dtype=torch.bfloat16): @@ -135,28 +124,6 @@ def _find_first_quantizable_linear(transformer): return None, None -@pytest.fixture -def flux1_checkpoint_exists(): - """Check if FLUX.1 checkpoint is available locally.""" - if not FLUX1_CHECKPOINT_PATH or not os.path.exists(FLUX1_CHECKPOINT_PATH): - pytest.skip( - f"FLUX.1 checkpoint not found at {FLUX1_CHECKPOINT_PATH}. " - "Set FLUX1_MODEL_PATH or stage checkpoint under LLM_MODELS_ROOT." - ) - return True - - -@pytest.fixture -def flux2_checkpoint_exists(): - """Check if FLUX.2 checkpoint is available locally.""" - if not FLUX2_CHECKPOINT_PATH or not os.path.exists(FLUX2_CHECKPOINT_PATH): - pytest.skip( - f"FLUX.2 checkpoint not found at {FLUX2_CHECKPOINT_PATH}. " - "Set FLUX2_MODEL_PATH or stage checkpoint under LLM_MODELS_ROOT." - ) - return True - - # ============================================================================= # Pipeline Loading Tests # ============================================================================= @@ -166,10 +133,10 @@ class TestFluxPipelineLoading: """Integration tests for FLUX pipeline loading.""" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_load_flux1_pipeline_basic(self, flux1_checkpoint_exists): + def test_load_flux1_pipeline_basic(self): """Test loading FLUX.1 pipeline.""" args = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), ) pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_COMPONENTS) @@ -184,10 +151,10 @@ def test_load_flux1_pipeline_basic(self, flux1_checkpoint_exists): torch.cuda.empty_cache() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_load_flux2_pipeline_basic(self, flux2_checkpoint_exists): + def test_load_flux2_pipeline_basic(self): """Test loading FLUX.2 pipeline.""" args = VisualGenArgs( - model=FLUX2_CHECKPOINT_PATH, + model=_flux2_path(), ) pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_COMPONENTS) @@ -202,10 +169,10 @@ def test_load_flux2_pipeline_basic(self, flux2_checkpoint_exists): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("backend", ["VANILLA", "TRTLLM"]) - def test_load_flux1_with_attention_backend(self, flux1_checkpoint_exists, backend: str): + def test_load_flux1_with_attention_backend(self, backend: str): """Test loading FLUX.1 with different attention backends.""" args = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), attention_config=AttentionConfig(backend=backend), ) @@ -228,10 +195,10 @@ class TestFluxQuantization: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"]) - def test_load_flux1_with_quantization(self, flux1_checkpoint_exists, quant_algo: str): + def test_load_flux1_with_quantization(self, quant_algo: str): """Test loading FLUX.1 with FP8 quantization and verify FP8 weights.""" args = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), quant_config={"quant_algo": quant_algo, "dynamic": True}, ) @@ -277,10 +244,10 @@ def test_load_flux1_with_quantization(self, flux1_checkpoint_exists, quant_algo: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"]) - def test_load_flux2_with_quantization(self, flux2_checkpoint_exists, quant_algo: str): + def test_load_flux2_with_quantization(self, quant_algo: str): """Test loading FLUX.2 with FP8 quantization and verify FP8 weights.""" args = VisualGenArgs( - model=FLUX2_CHECKPOINT_PATH, + model=_flux2_path(), quant_config={"quant_algo": quant_algo, "dynamic": True}, ) @@ -334,7 +301,7 @@ class TestFluxFP8NumericalCorrectness: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"]) - def test_fp8_vs_bf16_single_layer(self, flux1_checkpoint_exists, quant_algo: str): + def test_fp8_vs_bf16_single_layer(self, quant_algo: str): """Test FP8 vs BF16 numerical accuracy on a single Linear layer. Pattern (matching Wan test_fp8_vs_bf16_numerical_correctness): @@ -346,7 +313,7 @@ def test_fp8_vs_bf16_single_layer(self, flux1_checkpoint_exists, quant_algo: str # Load BF16 pipeline (reference) print(f"\n[Compare {quant_algo}] Loading BF16 pipeline...") args_bf16 = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), ) pipeline_bf16 = PipelineLoader(args_bf16).load( skip_warmup=True, skip_components=SKIP_COMPONENTS @@ -355,7 +322,7 @@ def test_fp8_vs_bf16_single_layer(self, flux1_checkpoint_exists, quant_algo: str # Load FP8 pipeline print(f"[Compare {quant_algo}] Loading {quant_algo} pipeline...") args_fp8 = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), quant_config={"quant_algo": quant_algo, "dynamic": True}, ) pipeline_fp8 = PipelineLoader(args_fp8).load( @@ -417,7 +384,7 @@ def test_fp8_vs_bf16_single_layer(self, flux1_checkpoint_exists, quant_algo: str @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES", "FP8_PER_CHANNEL_PER_TOKEN"]) - def test_fp8_vs_bf16_full_transformer_e2e(self, flux1_checkpoint_exists, quant_algo: str): + def test_fp8_vs_bf16_full_transformer_e2e(self, quant_algo: str): """End-to-end test: Compare full FLUX.1 transformer FP8 vs BF16 output. Runs the entire transformer (19 dual + 38 single blocks) and compares outputs. @@ -426,7 +393,7 @@ def test_fp8_vs_bf16_full_transformer_e2e(self, flux1_checkpoint_exists, quant_a # Load BF16 transformer (reference) print("\n[E2E] Loading BF16 transformer...") args_bf16 = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), ) pipeline_bf16 = PipelineLoader(args_bf16).load( skip_warmup=True, skip_components=SKIP_COMPONENTS @@ -436,7 +403,7 @@ def test_fp8_vs_bf16_full_transformer_e2e(self, flux1_checkpoint_exists, quant_a # Load FP8 transformer print(f"[E2E] Loading {quant_algo} transformer...") args_fp8 = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), quant_config={"quant_algo": quant_algo, "dynamic": True}, ) pipeline_fp8 = PipelineLoader(args_fp8).load( @@ -528,7 +495,7 @@ class TestFluxFP8Memory: """Test FP8 memory reduction for FLUX models.""" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_fp8_vs_bf16_memory_comparison(self, flux1_checkpoint_exists): + def test_fp8_vs_bf16_memory_comparison(self): """Test FP8 uses ~2x less memory than BF16 (matching Wan test).""" def get_module_memory_gb(module): @@ -539,7 +506,7 @@ def get_module_memory_gb(module): torch.cuda.reset_peak_memory_stats() args_bf16 = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), ) pipeline_bf16 = PipelineLoader(args_bf16).load( skip_warmup=True, skip_components=SKIP_COMPONENTS @@ -558,7 +525,7 @@ def get_module_memory_gb(module): torch.cuda.reset_peak_memory_stats() args_fp8 = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), quant_config={"quant_algo": "FP8", "dynamic": True}, ) pipeline_fp8 = PipelineLoader(args_fp8).load( @@ -594,7 +561,7 @@ class TestFluxAttentionBackend: """Test VANILLA vs TRTLLM attention backend numerical correctness.""" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_attention_backend_comparison(self, flux1_checkpoint_exists): + def test_attention_backend_comparison(self): """Test that VANILLA and TRTLLM backends produce similar outputs. FLUX uses joint self-attention (same seq_len for Q and KV), so both @@ -605,7 +572,7 @@ def test_attention_backend_comparison(self, flux1_checkpoint_exists): # (two full transformers don't fit in GPU memory simultaneously) print("\n[Attention Backend Test] Loading baseline transformer (VANILLA)...") args_baseline = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), attention_config=AttentionConfig(backend="VANILLA"), ) pipeline_baseline = PipelineLoader(args_baseline).load( @@ -627,7 +594,7 @@ def test_attention_backend_comparison(self, flux1_checkpoint_exists): # Load and run TRTLLM backend print("[Attention Backend Test] Loading TRTLLM transformer...") args_trtllm = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), attention_config=AttentionConfig(backend="TRTLLM"), ) pipeline_trtllm = PipelineLoader(args_trtllm).load( @@ -689,14 +656,14 @@ class TestFluxE2E: """End-to-end pipeline tests: full generation compared to HuggingFace reference.""" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_flux1_e2e_vs_hf(self, flux1_checkpoint_exists): + def test_flux1_e2e_vs_hf(self): """Full FLUX.1 pipeline (all components) generates image matching HF reference.""" from diffusers import FluxPipeline as HFFluxPipeline # 1. Generate HF reference image - hf_pipe = HFFluxPipeline.from_pretrained( - FLUX1_CHECKPOINT_PATH, torch_dtype=torch.bfloat16 - ).to("cuda") + hf_pipe = HFFluxPipeline.from_pretrained(_flux1_path(), torch_dtype=torch.bfloat16).to( + "cuda" + ) hf_result = hf_pipe( prompt="a tiny astronaut hatching from an egg on the moon", height=256, @@ -712,7 +679,7 @@ def test_flux1_e2e_vs_hf(self, flux1_checkpoint_exists): # 2. Load TRT-LLM pipeline args = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), ) pipeline = PipelineLoader(args).load() @@ -739,14 +706,14 @@ def test_flux1_e2e_vs_hf(self, flux1_checkpoint_exists): torch.cuda.empty_cache() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_flux2_e2e_vs_hf(self, flux2_checkpoint_exists): + def test_flux2_e2e_vs_hf(self): """Full FLUX.2 pipeline (all components) generates image matching HF reference.""" from diffusers import Flux2Pipeline as HFFlux2Pipeline # 1. Generate HF reference image - hf_pipe = HFFlux2Pipeline.from_pretrained( - FLUX2_CHECKPOINT_PATH, torch_dtype=torch.bfloat16 - ).to("cuda") + hf_pipe = HFFlux2Pipeline.from_pretrained(_flux2_path(), torch_dtype=torch.bfloat16).to( + "cuda" + ) hf_result = hf_pipe( prompt="a tiny astronaut hatching from an egg on the moon", height=256, @@ -762,7 +729,7 @@ def test_flux2_e2e_vs_hf(self, flux2_checkpoint_exists): # 2. Load TRT-LLM pipeline args = VisualGenArgs( - model=FLUX2_CHECKPOINT_PATH, + model=_flux2_path(), ) pipeline = PipelineLoader(args).load() @@ -790,7 +757,7 @@ def test_flux2_e2e_vs_hf(self, flux2_checkpoint_exists): torch.cuda.empty_cache() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_flux2_reference_image_e2e_vs_hf(self, flux2_checkpoint_exists): + def test_flux2_reference_image_e2e_vs_hf(self): """FLUX.2 reference-image generation matches the diffusers pipeline.""" from diffusers import Flux2Pipeline as HFFlux2Pipeline @@ -823,9 +790,9 @@ def test_flux2_reference_image_e2e_vs_hf(self, flux2_checkpoint_exists): }, ] - hf_pipe = HFFlux2Pipeline.from_pretrained( - FLUX2_CHECKPOINT_PATH, torch_dtype=torch.bfloat16 - ).to("cuda") + hf_pipe = HFFlux2Pipeline.from_pretrained(_flux2_path(), torch_dtype=torch.bfloat16).to( + "cuda" + ) hf_images = [] for case in cases: hf_result = hf_pipe( @@ -839,7 +806,7 @@ def test_flux2_reference_image_e2e_vs_hf(self, flux2_checkpoint_exists): gc.collect() torch.cuda.empty_cache() - pipeline = PipelineLoader(VisualGenArgs(model=FLUX2_CHECKPOINT_PATH)).load() + pipeline = PipelineLoader(VisualGenArgs(model=_flux2_path())).load() for case, hf_image_batch in zip(cases, hf_images): result = pipeline.forward( **case, @@ -871,13 +838,8 @@ class TestFluxBatchGeneration: @pytest.fixture(scope="class") def flux1_pipeline(self): """Load FLUX.1 TRT-LLM pipeline once for all FLUX.1 batch tests.""" - if not FLUX1_CHECKPOINT_PATH or not os.path.exists(FLUX1_CHECKPOINT_PATH): - pytest.skip( - f"FLUX.1 checkpoint not found at {FLUX1_CHECKPOINT_PATH}. " - "Set FLUX1_MODEL_PATH or LLM_MODELS_ROOT." - ) args = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -889,13 +851,8 @@ def flux1_pipeline(self): @pytest.fixture(scope="class") def flux2_pipeline(self): """Load FLUX.2 TRT-LLM pipeline once for all FLUX.2 batch tests.""" - if not FLUX2_CHECKPOINT_PATH or not os.path.exists(FLUX2_CHECKPOINT_PATH): - pytest.skip( - f"FLUX.2 checkpoint not found at {FLUX2_CHECKPOINT_PATH}. " - "Set FLUX2_MODEL_PATH or LLM_MODELS_ROOT." - ) args = VisualGenArgs( - model=FLUX2_CHECKPOINT_PATH, + model=_flux2_path(), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -1117,7 +1074,7 @@ class TestFluxParallelism: torch.cuda.is_available() and torch.cuda.device_count() < 2, reason="Ulysses parallel test requires at least 2 GPUs", ) - def test_ulysses_2gpu_correctness(self, flux1_checkpoint_exists): + def test_ulysses_2gpu_correctness(self): """Test Ulysses (ulysses_size=2) correctness against single-GPU baseline. Similar pattern to WAN's test_cfg_2gpu_correctness: @@ -1137,7 +1094,7 @@ def test_ulysses_2gpu_correctness(self, flux1_checkpoint_exists): # Load single-GPU reference print("\n[1/3] Loading single-GPU reference (ulysses_size=1) on GPU 0...") args_baseline = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), ) pipeline_baseline = PipelineLoader(args_baseline, device="cuda:0").load( skip_warmup=True, skip_components=SKIP_COMPONENTS @@ -1173,7 +1130,7 @@ def test_ulysses_2gpu_correctness(self, flux1_checkpoint_exists): master_port = get_free_port() mp.spawn( _run_ulysses_worker, - args=(2, master_port, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict), + args=(2, master_port, _flux1_path(), inputs_cpu, return_dict), nprocs=2, join=True, ) @@ -1298,7 +1255,7 @@ class TestFluxCombinedOptimizations: torch.cuda.is_available() and torch.cuda.device_count() < 2, reason="Combined optimization test requires at least 2 GPUs", ) - def test_all_optimizations_combined(self, flux1_checkpoint_exists): + def test_all_optimizations_combined(self): """Test FP8 + TeaCache + TRTLLM attention + Ulysses=2 combined correctness. Validates that all optimizations work together correctly. @@ -1317,7 +1274,7 @@ def test_all_optimizations_combined(self, flux1_checkpoint_exists): # Load baseline on GPU 0 (no optimizations) print("\n[1/3] Loading baseline on GPU 0 (BF16, no optimizations)...") args_baseline = VisualGenArgs( - model=FLUX1_CHECKPOINT_PATH, + model=_flux1_path(), ) pipeline_baseline = PipelineLoader(args_baseline, device="cuda:0").load( skip_warmup=True, skip_components=SKIP_COMPONENTS @@ -1355,7 +1312,7 @@ def test_all_optimizations_combined(self, flux1_checkpoint_exists): master_port = get_free_port() mp.spawn( _run_all_optimizations_worker, - args=(2, master_port, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict), + args=(2, master_port, _flux1_path(), inputs_cpu, return_dict), nprocs=2, join=True, ) diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py index 618b450d2699..a79241418bbe 100644 --- a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py @@ -496,11 +496,11 @@ def _assert_nvfp4_blocks_quantized(pipe) -> None: def _skip_if_no_fp8_ops() -> None: try: if not hasattr(torch.ops, "tensorrt_llm"): - pytest.skip("tensorrt_llm torch ops not available") + pytest.fail("tensorrt_llm torch ops not available") _ = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor _ = torch.ops.tensorrt_llm.quantize_e4m3_activation except (AttributeError, RuntimeError) as e: - pytest.skip(f"FP8 quantization ops not available: {e}") + pytest.fail(f"FP8 quantization ops not available: {e}") def _skip_if_no_nvfp4_ops() -> None: @@ -509,7 +509,7 @@ def _skip_if_no_nvfp4_ops() -> None: try: _ = torch.ops.trtllm.fp4_quantize except (AttributeError, RuntimeError) as e: - pytest.skip(f"fp4_quantize op not available: {e}") + pytest.fail(f"fp4_quantize op not available: {e}") def _transformer_inputs(transformer, device: str = "cuda", dtype=torch.bfloat16, seed: int = 42): diff --git a/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_pipeline.py b/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_pipeline.py index 00e78c399371..04ec5bcc95b8 100644 --- a/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_pipeline.py @@ -365,11 +365,11 @@ def _assert_nvfp4_blocks_quantized(pipe) -> None: def _skip_if_no_fp8_ops() -> None: try: if not hasattr(torch.ops, "tensorrt_llm"): - pytest.skip("tensorrt_llm torch ops not available") + pytest.fail("tensorrt_llm torch ops not available") _ = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor _ = torch.ops.tensorrt_llm.quantize_e4m3_activation except (AttributeError, RuntimeError) as e: - pytest.skip(f"FP8 quantization ops not available: {e}") + pytest.fail(f"FP8 quantization ops not available: {e}") def _skip_if_no_nvfp4_ops() -> None: @@ -378,7 +378,7 @@ def _skip_if_no_nvfp4_ops() -> None: try: _ = torch.ops.trtllm.fp4_quantize except (AttributeError, RuntimeError) as e: - pytest.skip(f"fp4_quantize op not available: {e}") + pytest.fail(f"fp4_quantize op not available: {e}") def _transformer_inputs(transformer, device: str = "cuda", dtype=torch.bfloat16, seed: int = 42): diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index f1064f780fa1..4a5c984f6e23 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -21,7 +21,7 @@ import pytest import torch import torch.nn.functional as F -from test_common.llm_data import llm_models_root +from test_common.llm_data import get_checkpoint from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig @@ -50,19 +50,11 @@ ] -_LTX2_BASE = os.path.join(str(llm_models_root(check=True)), "LTX-2") -_GEMMA3_DEFAULT = os.path.join(str(llm_models_root(check=True)), "gemma", "gemma-3-12b-it") - - -CHECKPOINT_PATH_BF16 = os.environ.get( - "LTX2_MODEL_PATH", - os.path.join(_LTX2_BASE, "ltx-2-19b-dev.safetensors"), -) -CHECKPOINT_PATH_FP8 = os.environ.get( - "LTX2_MODEL_PATH_FP8", - os.path.join(_LTX2_BASE, "ltx-2-19b-dev-fp8.safetensors"), -) -GEMMA3_PATH = os.environ.get("LTX2_TEXT_ENCODER_PATH", _GEMMA3_DEFAULT) +_LTX2_BF16_SUBDIR = "LTX-2/ltx-2-19b-dev.safetensors" +_LTX2_FP8_SUBDIR = "LTX-2/ltx-2-19b-dev-fp8.safetensors" +_LTX2_GEMMA3_SUBDIR = "gemma/gemma-3-12b-it" +_LTX2_UPSAMPLER_SUBDIR = "LTX-2/ltx-2-spatial-upscaler-x2-1.0.safetensors" +_LTX2_LORA_SUBDIR = "LTX-2/ltx-2-19b-distilled-lora-384.safetensors" def _ltx2_pipeline_config(**overrides): @@ -73,7 +65,7 @@ def _ltx2_pipeline_config(**overrides): load needs ``text_encoder_path`` set. Tests can pass extra keys via ``overrides`` (e.g. ``spatial_upsampler_path`` for two-stage). """ - cfg = {"text_encoder_path": GEMMA3_PATH} + cfg = {"text_encoder_path": get_checkpoint(_LTX2_GEMMA3_SUBDIR)} cfg.update(overrides) return cfg @@ -171,24 +163,14 @@ def _find_first_quantizable_linear(transformer): @pytest.fixture def ltx2_bf16_checkpoint_exists(): - """Check if LTX2 BF16 checkpoint is available locally.""" - if not CHECKPOINT_PATH_BF16 or not os.path.exists(CHECKPOINT_PATH_BF16): - pytest.skip( - f"LTX2 BF16 checkpoint not found at {CHECKPOINT_PATH_BF16}. " - "Set LTX2_MODEL_PATH or stage checkpoint under LLM_MODELS_ROOT/LTX-2/." - ) - return True + """Resolve the LTX2 BF16 checkpoint, failing loudly if it is not staged.""" + return get_checkpoint(_LTX2_BF16_SUBDIR) @pytest.fixture def ltx2_fp8_checkpoint_exists(): - """Check if LTX2 FP8 checkpoint is available locally.""" - if not CHECKPOINT_PATH_FP8 or not os.path.exists(CHECKPOINT_PATH_FP8): - pytest.skip( - f"LTX2 FP8 checkpoint not found at {CHECKPOINT_PATH_FP8}. " - "Set LTX2_MODEL_PATH_FP8 or stage checkpoint under LLM_MODELS_ROOT/LTX-2/." - ) - return True + """Resolve the LTX2 FP8 checkpoint, failing loudly if it is not staged.""" + return get_checkpoint(_LTX2_FP8_SUBDIR) # ============================================================================ @@ -204,7 +186,7 @@ class TestLTX2Quantization: def test_load_with_quantization(self, ltx2_bf16_checkpoint_exists, quant_algo: str): """Test loading LTX2 with FP8 quantization and verify FP8 weights.""" args = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_bf16_checkpoint_exists, quant_config={"quant_algo": quant_algo, "dynamic": True}, pipeline_config=_ltx2_pipeline_config(), ) @@ -260,7 +242,7 @@ def test_fp8_vs_bf16_single_layer(self, ltx2_bf16_checkpoint_exists, quant_algo: """ print(f"\n[Compare {quant_algo}] Loading BF16 pipeline...") args_bf16 = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_bf16_checkpoint_exists, pipeline_config=_ltx2_pipeline_config(), ) pipeline_bf16 = PipelineLoader(args_bf16).load( @@ -269,7 +251,7 @@ def test_fp8_vs_bf16_single_layer(self, ltx2_bf16_checkpoint_exists, quant_algo: print(f"[Compare {quant_algo}] Loading {quant_algo} pipeline...") args_fp8 = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_bf16_checkpoint_exists, quant_config={"quant_algo": quant_algo, "dynamic": True}, pipeline_config=_ltx2_pipeline_config(), ) @@ -337,7 +319,7 @@ def get_module_memory_gb(module): torch.cuda.reset_peak_memory_stats() args_bf16 = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_bf16_checkpoint_exists, pipeline_config=_ltx2_pipeline_config(), ) pipeline_bf16 = PipelineLoader(args_bf16).load( @@ -353,7 +335,7 @@ def get_module_memory_gb(module): torch.cuda.reset_peak_memory_stats() args_fp8 = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_bf16_checkpoint_exists, quant_config={"quant_algo": "FP8", "dynamic": True}, pipeline_config=_ltx2_pipeline_config(), ) @@ -390,7 +372,7 @@ def test_attention_backend_comparison(self, ltx2_bf16_checkpoint_exists): """ print("\n[Attention Backend Test] Loading baseline transformer (VANILLA)...") args_baseline = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_bf16_checkpoint_exists, attention_config=AttentionConfig(backend="VANILLA"), pipeline_config=_ltx2_pipeline_config(), ) @@ -417,7 +399,7 @@ def test_attention_backend_comparison(self, ltx2_bf16_checkpoint_exists): print("[Attention Backend Test] Loading TRTLLM transformer...") args_trtllm = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_bf16_checkpoint_exists, attention_config=AttentionConfig(backend="TRTLLM"), pipeline_config=_ltx2_pipeline_config(), ) @@ -1168,31 +1150,15 @@ def test_fp8_scale_packed_detection(self): # Two-Stage Pipeline Loading Tests (requires checkpoints) # ============================================================================ -UPSAMPLER_PATH = os.environ.get( - "LTX2_UPSAMPLER_CHECKPOINT", - os.path.join(_LTX2_BASE, "ltx-2-spatial-upscaler-x2-1.0.safetensors"), -) -LORA_PATH = os.environ.get( - "LTX2_LORA_CHECKPOINT", - os.path.join(_LTX2_BASE, "ltx-2-19b-distilled-lora-384.safetensors"), -) - -_HAS_TWO_STAGE_ASSETS = ( - os.path.exists(CHECKPOINT_PATH_BF16) - and os.path.exists(UPSAMPLER_PATH) - and os.path.exists(LORA_PATH) -) - @pytest.fixture def ltx2_two_stage_assets_exist(): - """Check if all two-stage assets (checkpoint + upsampler + LoRA) are available.""" - if not _HAS_TWO_STAGE_ASSETS: - pytest.skip( - f"Two-stage assets not found. Need checkpoint at {CHECKPOINT_PATH_BF16}, " - f"upsampler at {UPSAMPLER_PATH}, and LoRA at {LORA_PATH}." - ) - return True + """Resolve all two-stage assets (checkpoint + upsampler + LoRA), failing loudly if missing.""" + return SimpleNamespace( + checkpoint=get_checkpoint(_LTX2_BF16_SUBDIR), + upsampler=get_checkpoint(_LTX2_UPSAMPLER_SUBDIR), + lora=get_checkpoint(_LTX2_LORA_SUBDIR), + ) class TestLTX2TwoStageLoRAHelpers: @@ -1614,10 +1580,10 @@ def test_loads_as_two_stage_pipeline(self, ltx2_two_stage_assets_exist): ) args = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_two_stage_assets_exist.checkpoint, pipeline_config=_ltx2_pipeline_config( - spatial_upsampler_path=UPSAMPLER_PATH, - distilled_lora_path=LORA_PATH, + spatial_upsampler_path=ltx2_two_stage_assets_exist.upsampler, + distilled_lora_path=ltx2_two_stage_assets_exist.lora, ), ) @@ -1647,10 +1613,10 @@ def test_two_stage_lora_deltas_match_transformer(self, ltx2_two_stage_assets_exi ) args = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_two_stage_assets_exist.checkpoint, pipeline_config=_ltx2_pipeline_config( - spatial_upsampler_path=UPSAMPLER_PATH, - distilled_lora_path=LORA_PATH, + spatial_upsampler_path=ltx2_two_stage_assets_exist.upsampler, + distilled_lora_path=ltx2_two_stage_assets_exist.lora, ), ) @@ -1693,10 +1659,10 @@ def test_two_stage_with_quantization(self, ltx2_two_stage_assets_exist, quant_al ) args = VisualGenArgs( - model=CHECKPOINT_PATH_BF16, + model=ltx2_two_stage_assets_exist.checkpoint, pipeline_config=_ltx2_pipeline_config( - spatial_upsampler_path=UPSAMPLER_PATH, - distilled_lora_path=LORA_PATH, + spatial_upsampler_path=ltx2_two_stage_assets_exist.upsampler, + distilled_lora_path=ltx2_two_stage_assets_exist.lora, ), quant_config={"quant_algo": quant_algo, "dynamic": True}, ) diff --git a/tests/unittest/_torch/visual_gen/test_model_loader.py b/tests/unittest/_torch/visual_gen/test_model_loader.py index fb30573a38c5..0bc11b44ac14 100644 --- a/tests/unittest/_torch/visual_gen/test_model_loader.py +++ b/tests/unittest/_torch/visual_gen/test_model_loader.py @@ -5,36 +5,15 @@ import json import os -from pathlib import Path from types import SimpleNamespace import pytest import torch import torch.nn as nn +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if it is set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "You shall set LLM_MODELS_ROOT env or be able to access scratch.trt_llm_data to run this test" - ) - return str(root) - - -# Skip if checkpoint not available -# Set DIFFUSION_MODEL_PATH env var to run integration tests -CHECKPOINT_PATH = os.environ.get( - "DIFFUSION_MODEL_PATH", - os.path.join(_llm_models_root(), "Wan2.1-T2V-1.3B-Diffusers"), -) - # Skip heavy components (text_encoder ~44GB, vae ~300MB) to speed up tests # These components are loaded via diffusers and don't need quantization testing SKIP_HEAVY_COMPONENTS = [ @@ -46,30 +25,27 @@ def _llm_models_root() -> str: @pytest.fixture -def checkpoint_exists(): - return CHECKPOINT_PATH and os.path.exists(CHECKPOINT_PATH) +def checkpoint_path() -> str: + return get_checkpoint("Wan2.1-T2V-1.3B-Diffusers") -def test_meta_init_mode_creates_meta_tensors(checkpoint_exists): +def test_meta_init_mode_creates_meta_tensors(checkpoint_path) -> None: """Test that MetaInitMode creates tensors on meta device (no GPU memory).""" - if not checkpoint_exists: - pytest.skip("Checkpoint not available") - from tensorrt_llm._torch.models.modeling_utils import MetaInitMode from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig from tensorrt_llm._torch.visual_gen.models import AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs # Load config directly - args = VisualGenArgs(model=CHECKPOINT_PATH) + args = VisualGenArgs(model=checkpoint_path) config = DiffusionPipelineConfig.from_pretrained( - CHECKPOINT_PATH, + checkpoint_path, args=args, ) # Create pipeline WITH MetaInitMode with MetaInitMode(): - pipeline = AutoPipeline.from_config(config, CHECKPOINT_PATH) + pipeline = AutoPipeline.from_config(config, checkpoint_path) # Verify tensors are on meta device (no GPU memory allocated) param = next(pipeline.transformer.parameters()) @@ -485,18 +461,15 @@ def test_reject_invalid_vae_dynamic_quantization_modes(raw, message): DiffusionPipelineConfig.load_vae_conv_quant_config(raw) -def test_load_wan_pipeline_basic(checkpoint_exists): +def test_load_wan_pipeline_basic(checkpoint_path) -> None: """Test basic loading without quantization using VisualGenArgs.""" - if not checkpoint_exists: - pytest.skip("Checkpoint not available") - from tensorrt_llm._torch.visual_gen import PipelineLoader from tensorrt_llm.visual_gen.args import VisualGenArgs # Simple one-liner with VisualGenArgs # Skip text_encoder/vae to speed up test (focus on transformer) args = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, ) pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS) @@ -514,7 +487,7 @@ def test_load_wan_pipeline_basic(checkpoint_exists): assert param.dtype in [torch.float32, torch.bfloat16, torch.float16] -def test_load_wan_pipeline_with_fp8_dynamic_quant(checkpoint_exists): +def test_load_wan_pipeline_with_fp8_dynamic_quant(checkpoint_path) -> None: """Test loading with FP8 dynamic quantization using VisualGenArgs. Verifies the dynamic quantization flow: @@ -523,9 +496,6 @@ def test_load_wan_pipeline_with_fp8_dynamic_quant(checkpoint_exists): 3. BF16 checkpoint weights are quantized on-the-fly 4. Quantized weights are in FP8 format """ - if not checkpoint_exists: - pytest.skip("Checkpoint not available") - from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen import PipelineLoader from tensorrt_llm.visual_gen.args import VisualGenArgs @@ -533,7 +503,7 @@ def test_load_wan_pipeline_with_fp8_dynamic_quant(checkpoint_exists): # Use VisualGenArgs with FP8 quantization # Skip text_encoder/vae to speed up test (focus on transformer quantization) args = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, quant_config={"quant_algo": "FP8", "dynamic": True}, ) pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS) @@ -560,7 +530,7 @@ def test_load_wan_pipeline_with_fp8_dynamic_quant(checkpoint_exists): assert found_fp8_linear, "No FP8 Linear modules found in transformer" -def test_load_wan_pipeline_with_fp8_rowwise(checkpoint_exists): +def test_load_wan_pipeline_with_fp8_rowwise(checkpoint_path) -> None: """Test loading with FP8 row-wise (per-channel-per-token) dynamic quantization. Verifies: @@ -568,15 +538,12 @@ def test_load_wan_pipeline_with_fp8_rowwise(checkpoint_exists): 2. Linear weights are FP8 after loading 3. weight_scale is 1-D [out_features] — one scale per output row, not a scalar """ - if not checkpoint_exists: - pytest.skip("Checkpoint not available") - from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen import PipelineLoader from tensorrt_llm.visual_gen.args import VisualGenArgs args = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, quant_config={"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN", "dynamic": True}, ) pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS) @@ -608,18 +575,15 @@ def test_load_wan_pipeline_with_fp8_rowwise(checkpoint_exists): assert found_fp8_linear, "No FP8 Linear modules found in transformer" -def test_load_wan_pipeline_with_fp8_blockwise(checkpoint_exists): +def test_load_wan_pipeline_with_fp8_blockwise(checkpoint_path) -> None: """Test loading with FP8 blockwise quantization using VisualGenArgs.""" - if not checkpoint_exists: - pytest.skip("Checkpoint not available") - from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen import PipelineLoader from tensorrt_llm.visual_gen.args import VisualGenArgs # Skip text_encoder/vae to speed up test (focus on transformer quantization) args = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, quant_config={"quant_algo": "FP8_BLOCK_SCALES", "dynamic": True}, ) pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS) @@ -707,11 +671,8 @@ def test_visual_gen_args_to_quant_config(): assert dwq is True -def test_load_without_quant_config_no_fp8(checkpoint_exists): +def test_load_without_quant_config_no_fp8(checkpoint_path) -> None: """Test that loading without quant_config does NOT produce FP8 weights.""" - if not checkpoint_exists: - pytest.skip("Checkpoint not available") - from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen import PipelineLoader from tensorrt_llm.visual_gen.args import VisualGenArgs @@ -719,7 +680,7 @@ def test_load_without_quant_config_no_fp8(checkpoint_exists): # No quantization specified # Skip text_encoder/vae to speed up test (focus on transformer) args = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, ) pipeline = PipelineLoader(args).load(skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS) @@ -786,7 +747,7 @@ def _get_cuda_peak_memory_gb(): return torch.cuda.max_memory_allocated() / 1024**3 -def test_fp8_vs_bf16_memory_comparison(checkpoint_exists): +def test_fp8_vs_bf16_memory_comparison(checkpoint_path) -> None: """Test FP8 dynamic quant uses ~2x less memory than BF16, including peak memory. This test verifies that dynamic quantization doesn't create unnecessary @@ -796,9 +757,6 @@ def test_fp8_vs_bf16_memory_comparison(checkpoint_exists): - BF16: ~2.6 GB model memory, similar peak during loading - FP8: ~1.3 GB model memory, peak should be < 2x BF16 peak """ - if not checkpoint_exists: - pytest.skip("Checkpoint not available") - from tensorrt_llm._torch.visual_gen import PipelineLoader from tensorrt_llm.visual_gen.args import VisualGenArgs @@ -809,7 +767,7 @@ def test_fp8_vs_bf16_memory_comparison(checkpoint_exists): torch.cuda.reset_peak_memory_stats() args_bf16 = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, ) pipeline_bf16 = PipelineLoader(args_bf16).load( skip_warmup=True, skip_components=SKIP_HEAVY_COMPONENTS @@ -833,7 +791,7 @@ def test_fp8_vs_bf16_memory_comparison(checkpoint_exists): torch.cuda.reset_peak_memory_stats() args_fp8 = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, quant_config={"quant_algo": "FP8", "dynamic": True}, ) pipeline_fp8 = PipelineLoader(args_fp8).load( @@ -890,7 +848,7 @@ def test_fp8_vs_bf16_memory_comparison(checkpoint_exists): torch.cuda.reset_peak_memory_stats() args_fp8_block = VisualGenArgs( - model=CHECKPOINT_PATH, + model=checkpoint_path, quant_config={"quant_algo": "FP8_BLOCK_SCALES", "dynamic": True}, ) pipeline_fp8_block = PipelineLoader(args_fp8_block).load( diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_parity.py b/tests/unittest/_torch/visual_gen/test_qwen_image_parity.py index 9f2b5d8ea770..5d309aae9538 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_parity.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_parity.py @@ -9,28 +9,16 @@ ``transformer`` state_dict into both, runs them on identical inputs, and asserts cosine similarity above a per-module threshold. -Tests are skipped automatically unless ``QWEN_IMAGE_CKPT`` points to a -local Qwen-Image checkpoint. +The parity tests fail loudly if the Qwen-Image checkpoint is missing; +stage ``qwen-image`` under ``LLM_MODELS_ROOT`` to run them. """ import json -import os from pathlib import Path import pytest import torch - -_CKPT_ENV = "QWEN_IMAGE_CKPT" - - -def _ckpt_path() -> Path | None: - ckpt = os.environ.get(_CKPT_ENV) - if not ckpt: - return None - path = Path(ckpt) - if not (path / "transformer" / "config.json").is_file(): - return None - return path +from utils.llm_data import get_checkpoint def _load_transformer_state_dict(ckpt: Path) -> dict[str, torch.Tensor]: @@ -48,15 +36,6 @@ def _cosine(a: torch.Tensor, b: torch.Tensor) -> float: return torch.nn.functional.cosine_similarity(a, b, dim=0).item() -requires_ckpt = pytest.mark.skipif( - _ckpt_path() is None, - reason=( - f"Qwen-Image checkpoint not found at {_CKPT_ENV}. " - "Set QWEN_IMAGE_CKPT to a local Qwen/Qwen-Image checkpoint " - "to enable parity tests." - ), -) - requires_cuda = pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA is required for BF16 parity tests.", @@ -65,15 +44,13 @@ def _cosine(a: torch.Tensor, b: torch.Tensor) -> float: @pytest.fixture(scope="module") def transformer_state_dict() -> dict[str, torch.Tensor]: - ckpt = _ckpt_path() - assert ckpt is not None + ckpt = Path(get_checkpoint("qwen-image")) return _load_transformer_state_dict(ckpt) @pytest.fixture(scope="module") def transformer_config() -> dict: - ckpt = _ckpt_path() - assert ckpt is not None + ckpt = Path(get_checkpoint("qwen-image")) return json.loads((ckpt / "transformer" / "config.json").read_text()) @@ -82,7 +59,6 @@ def transformer_config() -> dict: # =========================================================================== -@requires_ckpt @requires_cuda @pytest.mark.parametrize("timestep_value", [0.001, 0.25, 0.5, 0.99]) def test_qwen_timestep_proj_embedding_parity(transformer_state_dict, timestep_value): @@ -196,7 +172,6 @@ def test_apply_rotary_emb_qwen_parity(): # =========================================================================== -@requires_ckpt @requires_cuda def test_pre_post_block_modules_parity(transformer_state_dict, transformer_config): """img_in, txt_in, txt_norm, norm_out, proj_out parity vs diffusers. @@ -302,7 +277,6 @@ def sd(prefix): # =========================================================================== -@requires_ckpt @requires_cuda def test_qwen_image_transformer_block_parity(transformer_state_dict, transformer_config): """One ``QwenImageTransformerBlock`` must match diffusers cos >= 0.999. @@ -417,7 +391,6 @@ def test_qwen_image_transformer_block_parity(transformer_state_dict, transformer # =========================================================================== -@requires_ckpt @requires_cuda @pytest.mark.slow def test_qwen_image_transformer_full_parity(transformer_state_dict, transformer_config): diff --git a/tests/unittest/_torch/visual_gen/test_utils.py b/tests/unittest/_torch/visual_gen/test_utils.py index 09705a5602c7..9570dbdce411 100644 --- a/tests/unittest/_torch/visual_gen/test_utils.py +++ b/tests/unittest/_torch/visual_gen/test_utils.py @@ -23,13 +23,8 @@ import torch.distributed as dist import torch.multiprocessing as mp -try: - from tensorrt_llm._torch.visual_gen.utils import SequenceSharder - from tensorrt_llm._utils import get_free_port - - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False +from tensorrt_llm._torch.visual_gen.utils import SequenceSharder +from tensorrt_llm._utils import get_free_port @pytest.fixture(autouse=True, scope="module") @@ -175,8 +170,6 @@ def _spawn_entry_combined(rank: int, world_size: int, port: int): def _run_dist(world_size: int, entry: Callable[[int, int, int], None]): - if not MODULES_AVAILABLE: - pytest.skip("SequenceSharder import failed") port = get_free_port() mp.spawn(entry, args=(world_size, port), nprocs=world_size, join=True) diff --git a/tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py index c18e9e8f67e9..23618a455c23 100644 --- a/tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py @@ -13,16 +13,10 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py -v -s -k 480p pytest tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py -v -s -k 720p - -Override checkpoint paths: - DIFFUSION_MODEL_PATH_WAN21_I2V_480P=/path/to/480p \\ - DIFFUSION_MODEL_PATH_WAN21_I2V_720P=/path/to/720p \\ - pytest tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py -v -s """ import gc import os -from pathlib import Path os.environ["TLLM_DISABLE_MPI"] = "1" @@ -32,6 +26,7 @@ import torch.nn.functional as F from diffusers import DiffusionPipeline from PIL import Image +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import ( @@ -49,33 +44,10 @@ def _cleanup_mpi_env(): # ============================================================================ -# Path helpers +# Checkpoints # ============================================================================ - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN21_I2V_480P_PATH = _checkpoint( - "DIFFUSION_MODEL_PATH_WAN21_I2V_480P", "Wan2.1-I2V-14B-480P-Diffusers" -) -WAN21_I2V_720P_PATH = _checkpoint( - "DIFFUSION_MODEL_PATH_WAN21_I2V_720P", "Wan2.1-I2V-14B-720P-Diffusers" -) +WAN21_I2V_480P_SUBDIR = "Wan2.1-I2V-14B-480P-Diffusers" # ============================================================================ # Test constants @@ -104,8 +76,6 @@ def _make_test_image(height: int, width: int) -> Image.Image: def _load_trtllm_pipeline(checkpoint_path: str): """Load TRTLLM WanImageToVideoPipeline without torch.compile or warmup.""" - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), @@ -267,7 +237,7 @@ class TestWan21_I2V_480P_PipelineCorrectness: def test_cosine_similarity(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN21_I2V_480P_PATH, + checkpoint_path=get_checkpoint(WAN21_I2V_480P_SUBDIR), height=480, width=832, num_frames=33, @@ -293,11 +263,8 @@ class TestWanI2VBatchGeneration: @pytest.fixture(scope="class") def i2v_full_pipeline(self): """Load full I2V pipeline (all components) for batch tests.""" - if not WAN21_I2V_480P_PATH or not os.path.exists(WAN21_I2V_480P_PATH): - pytest.skip("Checkpoint not available. Set DIFFUSION_MODEL_PATH_WAN21_I2V_480P.") - args = VisualGenArgs( - model=WAN21_I2V_480P_PATH, + model=get_checkpoint(WAN21_I2V_480P_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -358,10 +325,8 @@ class TestWan21I2VCombinedOptimizations: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_fp8_teacache_trtllm(self): - if not os.path.exists(WAN21_I2V_480P_PATH): - pytest.skip(f"Checkpoint not found: {WAN21_I2V_480P_PATH}") args = VisualGenArgs( - model=WAN21_I2V_480P_PATH, + model=get_checkpoint(WAN21_I2V_480P_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), quant_config={"quant_algo": "FP8", "dynamic": True}, attention_config=AttentionConfig(backend="TRTLLM"), diff --git a/tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py b/tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py index 798755df8c86..5377fa2ced82 100644 --- a/tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py @@ -19,12 +19,6 @@ pytest tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py -v -s -k wan21_i2v_480p pytest tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py -v -s -k wan21_i2v_720p pytest tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py -v -s -k wan22_raises - -Override checkpoint paths: - DIFFUSION_MODEL_PATH_WAN21_I2V_480P=/path/to/480p \\ - DIFFUSION_MODEL_PATH_WAN21_I2V_720P=/path/to/720p \\ - DIFFUSION_MODEL_PATH_WAN22_I2V=/path/to/wan22 \\ - pytest tests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.py -v -s """ import os @@ -32,12 +26,12 @@ os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path import numpy as np import pytest import torch from PIL import Image +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TeaCacheConfig, VisualGenArgs @@ -59,35 +53,12 @@ def _cleanup_gpu(): # ============================================================================ -# Path helpers +# Checkpoints # ============================================================================ - -def _llm_models_root() -> Path: - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - else: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return root - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or str(_llm_models_root() / default_name) - - -WAN21_I2V_480P_PATH = _checkpoint( - "DIFFUSION_MODEL_PATH_WAN21_I2V_480P", "Wan2.1-I2V-14B-480P-Diffusers" -) -WAN21_I2V_720P_PATH = _checkpoint( - "DIFFUSION_MODEL_PATH_WAN21_I2V_720P", "Wan2.1-I2V-14B-720P-Diffusers" -) - -WAN22_I2V_A14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN22_I2V", "Wan2.2-I2V-A14B-Diffusers") +WAN21_I2V_480P_SUBDIR = "Wan2.1-I2V-14B-480P-Diffusers" +WAN21_I2V_720P_SUBDIR = "Wan2.1-I2V-14B-720P-Diffusers" +WAN22_I2V_A14B_SUBDIR = "Wan2.2-I2V-A14B-Diffusers" INFER_NUM_FRAMES = 33 # (33-1)/4+1 = 9 latent frames; smallest realistic shape INFER_NUM_STEPS = 50 # Required for meaningful cache hits with calibrated coefficients @@ -100,8 +71,6 @@ def _checkpoint(env_var: str, default_name: str) -> str: def _make_pipeline(checkpoint_path: str, use_ret_steps: bool = False): - if not checkpoint_path or not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, cache_config=TeaCacheConfig( @@ -115,7 +84,7 @@ def _make_pipeline(checkpoint_path: str, use_ret_steps: bool = False): @pytest.fixture def wan21_i2v_480p_pipeline(): - pipeline = _make_pipeline(WAN21_I2V_480P_PATH) + pipeline = _make_pipeline(get_checkpoint(WAN21_I2V_480P_SUBDIR)) yield pipeline del pipeline torch.cuda.empty_cache() @@ -123,7 +92,7 @@ def wan21_i2v_480p_pipeline(): @pytest.fixture def wan21_i2v_480p_ret_steps_pipeline(): - pipeline = _make_pipeline(WAN21_I2V_480P_PATH, use_ret_steps=True) + pipeline = _make_pipeline(get_checkpoint(WAN21_I2V_480P_SUBDIR), use_ret_steps=True) yield pipeline del pipeline torch.cuda.empty_cache() @@ -131,7 +100,7 @@ def wan21_i2v_480p_ret_steps_pipeline(): @pytest.fixture def wan21_i2v_720p_pipeline(): - pipeline = _make_pipeline(WAN21_I2V_720P_PATH) + pipeline = _make_pipeline(get_checkpoint(WAN21_I2V_720P_SUBDIR)) yield pipeline del pipeline torch.cuda.empty_cache() @@ -244,12 +213,8 @@ class TestWan22_I2V_TeaCacheRaisesError: """Wan2.2-I2V-A14B must raise ValueError when TeaCache is enabled.""" def test_wan22_raises_if_teacache_enabled(self): - if not os.path.exists(WAN22_I2V_A14B_PATH): - pytest.skip( - f"Checkpoint not found: {WAN22_I2V_A14B_PATH} (set DIFFUSION_MODEL_PATH_WAN22_I2V)" - ) args = VisualGenArgs( - model=WAN22_I2V_A14B_PATH, + model=get_checkpoint(WAN22_I2V_A14B_SUBDIR), cache_config=TeaCacheConfig(), ) with pytest.raises(ValueError, match=r"Wan 2\.2 TeaCache requires explicit"): diff --git a/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py index 125b4710ddef..d3d59660cbd0 100644 --- a/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py @@ -13,16 +13,10 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py -v -s -k 1_3b pytest tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py -v -s -k 14b - -Override checkpoint paths: - DIFFUSION_MODEL_PATH_WAN21_1_3B=/path/to/1.3b \\ - DIFFUSION_MODEL_PATH_WAN21_14B=/path/to/14b \\ - pytest tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py -v -s """ import gc import os -from pathlib import Path os.environ["TLLM_DISABLE_MPI"] = "1" @@ -31,6 +25,7 @@ import torch import torch.nn.functional as F from diffusers import DiffusionPipeline +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader @@ -48,31 +43,6 @@ def _cleanup_mpi_env(): os.environ.pop("TLLM_DISABLE_MPI", None) -# ============================================================================ -# Path helpers -# ============================================================================ - - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN21_1_3B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN21_1_3B", "Wan2.1-T2V-1.3B-Diffusers") -WAN21_14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN21_14B", "Wan2.1-T2V-14B-Diffusers") - # ============================================================================ # Test constants # ============================================================================ @@ -91,8 +61,6 @@ def _checkpoint(env_var: str, default_name: str) -> str: def _load_trtllm_pipeline(checkpoint_path: str): """Load TRTLLM WanPipeline without torch.compile or warmup.""" - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), @@ -246,7 +214,7 @@ class TestWan21_1_3B_PipelineCorrectness: def test_cosine_similarity(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN21_1_3B_PATH, + checkpoint_path=get_checkpoint("Wan2.1-T2V-1.3B-Diffusers"), height=480, width=832, num_frames=9, @@ -262,7 +230,7 @@ class TestWan21_14B_PipelineCorrectness: def test_cosine_similarity(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN21_14B_PATH, + checkpoint_path=get_checkpoint("Wan2.1-T2V-14B-Diffusers"), height=720, width=1280, num_frames=9, @@ -286,11 +254,8 @@ class TestWanBatchGeneration: @pytest.fixture(scope="class") def wan21_full_pipeline(self): """Load full Wan 2.1 pipeline (all components) for batch tests.""" - if not WAN21_1_3B_PATH or not os.path.exists(WAN21_1_3B_PATH): - pytest.skip("Checkpoint not available. Set DIFFUSION_MODEL_PATH_WAN21_1_3B.") - args = VisualGenArgs( - model=WAN21_1_3B_PATH, + model=get_checkpoint("Wan2.1-T2V-1.3B-Diffusers"), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -347,10 +312,8 @@ class TestWan21T2VCombinedOptimizations: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_fp8_teacache_trtllm(self): - if not os.path.exists(WAN21_1_3B_PATH): - pytest.skip(f"Checkpoint not found: {WAN21_1_3B_PATH}") args = VisualGenArgs( - model=WAN21_1_3B_PATH, + model=get_checkpoint("Wan2.1-T2V-1.3B-Diffusers"), torch_compile_config=TorchCompileConfig(enable=False), quant_config={"quant_algo": "FP8", "dynamic": True}, attention_config=AttentionConfig(backend="TRTLLM"), @@ -397,10 +360,8 @@ def test_fp8_teacache_trtllm(self): def _make_wan21_t2v(quant_config=None): - if not os.path.exists(WAN21_1_3B_PATH): - pytest.skip(f"Checkpoint not found: {WAN21_1_3B_PATH}") kwargs = dict( - model=WAN21_1_3B_PATH, + model=get_checkpoint("Wan2.1-T2V-1.3B-Diffusers"), torch_compile_config=TorchCompileConfig(enable=False), ) if quant_config is not None: @@ -499,11 +460,11 @@ def test_fp8_weights_loaded(self, wan21_t2v_fp8): """FP8 transformer blocks have float8_e4m3fn weights and weight_scale.""" try: if not hasattr(torch.ops, "tensorrt_llm"): - pytest.skip("tensorrt_llm torch ops not available") + pytest.fail("tensorrt_llm torch ops not available") _ = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor _ = torch.ops.tensorrt_llm.quantize_e4m3_activation except (AttributeError, RuntimeError) as e: - pytest.skip(f"FP8 quantization ops not available: {e}") + pytest.fail(f"FP8 quantization ops not available: {e}") for name, module in wan21_t2v_fp8.transformer.named_modules(): if isinstance(module, Linear) and "blocks." in name: assert module.weight.dtype == torch.float8_e4m3fn, ( @@ -517,11 +478,11 @@ def test_fp8_block_scales_weights_loaded(self, wan21_t2v_fp8_block): """FP8_BLOCK_SCALES transformer blocks have float8_e4m3fn weights and weight_scale.""" try: if not hasattr(torch.ops, "tensorrt_llm"): - pytest.skip("tensorrt_llm torch ops not available") + pytest.fail("tensorrt_llm torch ops not available") _ = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor _ = torch.ops.tensorrt_llm.quantize_e4m3_activation except (AttributeError, RuntimeError) as e: - pytest.skip(f"FP8 quantization ops not available: {e}") + pytest.fail(f"FP8 quantization ops not available: {e}") for name, module in wan21_t2v_fp8_block.transformer.named_modules(): if isinstance(module, Linear) and "blocks." in name: assert module.weight.dtype == torch.float8_e4m3fn, ( @@ -538,7 +499,7 @@ def test_nvfp4_weights_loaded(self, wan21_t2v_nvfp4): try: _ = torch.ops.trtllm.fp4_quantize except (AttributeError, RuntimeError) as e: - pytest.skip(f"fp4_quantize op not available: {e}") + pytest.fail(f"fp4_quantize op not available: {e}") from tensorrt_llm.quantization.utils import fp4_utils for name, module in wan21_t2v_nvfp4.transformer.named_modules(): @@ -648,7 +609,7 @@ def test_nvfp4_e2e_accuracy(self, wan21_t2v_bf16, wan21_t2v_nvfp4): try: _ = torch.ops.trtllm.fp4_quantize except (AttributeError, RuntimeError) as e: - pytest.skip(f"fp4_quantize op not available: {e}") + pytest.fail(f"fp4_quantize op not available: {e}") hs, ts, enc = _transformer_inputs() diff --git a/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py b/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py index 76e19f8232e5..68fefbd2886e 100644 --- a/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py @@ -19,12 +19,6 @@ pytest tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py -v -s -k wan21_1_3b pytest tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py -v -s -k wan21_14b pytest tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py -v -s -k wan22_raises - -Override checkpoint paths: - DIFFUSION_MODEL_PATH_WAN21_1_3B=/path/to/1.3b \\ - DIFFUSION_MODEL_PATH_WAN21_14B=/path/to/14b \\ - DIFFUSION_MODEL_PATH_WAN22_T2V=/path/to/wan22 \\ - pytest tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.py -v -s """ import os @@ -32,10 +26,10 @@ os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path import pytest import torch +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TeaCacheConfig, VisualGenArgs @@ -57,30 +51,12 @@ def _cleanup_gpu(): # ============================================================================ -# Path helpers +# Checkpoints # ============================================================================ - -def _llm_models_root() -> Path: - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - else: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return root - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or str(_llm_models_root() / default_name) - - -WAN21_1_3B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN21_1_3B", "Wan2.1-T2V-1.3B-Diffusers") -WAN21_14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN21_14B", "Wan2.1-T2V-14B-Diffusers") -WAN22_A14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN22_T2V", "Wan2.2-T2V-A14B-Diffusers") +WAN21_1_3B_SUBDIR = "Wan2.1-T2V-1.3B-Diffusers" +WAN21_14B_SUBDIR = "Wan2.1-T2V-14B-Diffusers" +WAN22_A14B_SUBDIR = "Wan2.2-T2V-A14B-Diffusers" INFER_NUM_FRAMES = 33 # (33-1)/4+1 = 9 latent frames; smallest realistic shape INFER_NUM_STEPS = 50 # Required for meaningful cache hits with calibrated coefficients @@ -93,8 +69,6 @@ def _checkpoint(env_var: str, default_name: str) -> str: def _make_pipeline(checkpoint_path: str, use_ret_steps: bool = False): - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, cache_config=TeaCacheConfig( @@ -108,7 +82,7 @@ def _make_pipeline(checkpoint_path: str, use_ret_steps: bool = False): @pytest.fixture def wan21_1_3b_pipeline(): - pipeline = _make_pipeline(WAN21_1_3B_PATH) + pipeline = _make_pipeline(get_checkpoint(WAN21_1_3B_SUBDIR)) yield pipeline del pipeline torch.cuda.empty_cache() @@ -116,7 +90,7 @@ def wan21_1_3b_pipeline(): @pytest.fixture def wan21_1_3b_ret_steps_pipeline(): - pipeline = _make_pipeline(WAN21_1_3B_PATH, use_ret_steps=True) + pipeline = _make_pipeline(get_checkpoint(WAN21_1_3B_SUBDIR), use_ret_steps=True) yield pipeline del pipeline torch.cuda.empty_cache() @@ -124,7 +98,7 @@ def wan21_1_3b_ret_steps_pipeline(): @pytest.fixture def wan21_14b_pipeline(): - pipeline = _make_pipeline(WAN21_14B_PATH) + pipeline = _make_pipeline(get_checkpoint(WAN21_14B_SUBDIR)) yield pipeline del pipeline torch.cuda.empty_cache() @@ -228,10 +202,8 @@ class TestWan22_T2V_TeaCacheRaisesError: """Wan2.2-T2V-A14B must raise ValueError when TeaCache is enabled.""" def test_wan22_raises_if_teacache_enabled(self): - if not os.path.exists(WAN22_A14B_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_A14B_PATH}") args = VisualGenArgs( - model=WAN22_A14B_PATH, + model=get_checkpoint(WAN22_A14B_SUBDIR), cache_config=TeaCacheConfig(), ) with pytest.raises(ValueError, match=r"Wan 2\.2 TeaCache requires explicit"): diff --git a/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache_user_coefficients.py b/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache_user_coefficients.py index aa16c708b6a7..058961243ba6 100644 --- a/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache_user_coefficients.py +++ b/tests/unittest/_torch/visual_gen/test_wan21_t2v_teacache_user_coefficients.py @@ -20,10 +20,6 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan13b_teacache_coefficients.py -v -s - -Override checkpoint: - DIFFUSION_MODEL_PATH_WAN21_1_3B=/path/to/weights \\ - pytest tests/unittest/_torch/visual_gen/test_wan13b_teacache_coefficients.py -v -s """ import gc @@ -31,10 +27,9 @@ os.environ["TLLM_DISABLE_MPI"] = "1" -from pathlib import Path - import pytest import torch +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TeaCacheConfig, VisualGenArgs @@ -55,23 +50,7 @@ def _cleanup_gpu(): torch.cuda.empty_cache() -def _llm_models_root() -> Path: - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - else: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return root - - -WAN21_1_3B_PATH = os.environ.get( - "DIFFUSION_MODEL_PATH_WAN21_1_3B", - str(_llm_models_root() / "Wan2.1-T2V-1.3B-Diffusers"), -) +WAN21_1_3B_SUBDIR = "Wan2.1-T2V-1.3B-Diffusers" PROMPT = "a cat sitting on a windowsill" HEIGHT, WIDTH = 480, 832 @@ -98,11 +77,8 @@ def _llm_models_root() -> Path: def _run_forward(coefficients: list, thresh: float, label: str) -> dict: """Load the pipeline with the given user-supplied coefficients, run one forward pass.""" - if not os.path.exists(WAN21_1_3B_PATH): - pytest.skip(f"Checkpoint not found: {WAN21_1_3B_PATH}") - args = VisualGenArgs( - model=WAN21_1_3B_PATH, + model=get_checkpoint(WAN21_1_3B_SUBDIR), cache_config=TeaCacheConfig( teacache_thresh=thresh, coefficients=coefficients, diff --git a/tests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.py index e85fade5557a..374dc4e91986 100644 --- a/tests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.py @@ -19,19 +19,13 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN22_I2V=/path/to/wan22_i2v \\ - pytest tests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.py -v -s """ -import importlib import os os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path import numpy as np import pytest @@ -39,6 +33,7 @@ import torch.nn.functional as F from diffusers import DiffusionPipeline from PIL import Image +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import ( @@ -55,29 +50,7 @@ def _cleanup_mpi_env(): os.environ.pop("TLLM_DISABLE_MPI", None) -# ============================================================================ -# Path helpers -# ============================================================================ - - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN22_I2V_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN22_I2V", "Wan2.2-I2V-A14B-Diffusers") +WAN22_I2V_SUBDIR = "Wan2.2-I2V-A14B-Diffusers" # ============================================================================ # Test constants @@ -106,8 +79,6 @@ def _make_test_image(height: int, width: int) -> Image.Image: def _load_trtllm_pipeline(checkpoint_path: str): """Load TRTLLM WanImageToVideoPipeline (two-stage) without torch.compile or warmup.""" - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), @@ -286,7 +257,7 @@ class TestWan22_I2V_A14B_PipelineCorrectness: def test_cosine_similarity(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN22_I2V_PATH, + checkpoint_path=get_checkpoint(WAN22_I2V_SUBDIR), height=480, width=832, num_frames=9, @@ -311,10 +282,8 @@ def test_cosine_similarity(self): def _make_wan22_i2v(quant_config=None, attention_config=None): - if not os.path.exists(WAN22_I2V_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_I2V_PATH}") kwargs = dict( - model=WAN22_I2V_PATH, + model=get_checkpoint(WAN22_I2V_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), ) if quant_config is not None: @@ -355,7 +324,7 @@ class TestWan22TwoStageI2VFeatures: def test_fp8_on_both_stages(self, wan22_i2v_fp8): """FP8 quantization is applied to both transformer and transformer_2.""" if wan22_i2v_fp8.transformer_2 is None: - pytest.skip("Not a two-stage checkpoint") + pytest.fail("Not a two-stage checkpoint") def _has_fp8(module): return any( @@ -370,7 +339,7 @@ def _has_fp8(module): def test_trtllm_attention_both_stages(self, wan22_i2v_trtllm): """TRTLLM self-attention and VANILLA cross-attention on both stages.""" if wan22_i2v_trtllm.transformer_2 is None: - pytest.skip("Not a two-stage checkpoint") + pytest.fail("Not a two-stage checkpoint") for stage_name, transformer in [ ("transformer", wan22_i2v_trtllm.transformer), @@ -401,11 +370,8 @@ class TestWan22I2VBatchGeneration: @pytest.fixture(scope="class") def wan22_i2v_full_pipeline(self): """Load full Wan 2.2 I2V pipeline (all components) for batch tests.""" - if not WAN22_I2V_PATH or not os.path.exists(WAN22_I2V_PATH): - pytest.skip("Checkpoint not available. Set DIFFUSION_MODEL_PATH_WAN22_I2V.") - args = VisualGenArgs( - model=WAN22_I2V_PATH, + model=get_checkpoint(WAN22_I2V_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -459,16 +425,13 @@ def test_batch_prompt_shape(self, wan22_i2v_full_pipeline): @pytest.mark.integration @pytest.mark.wan_i2v -@pytest.mark.skipif(importlib.util.find_spec("cache_dit") is None, reason="cache_dit not installed") class TestWan22I2VCombinedOptimizations: """FP8 + CacheDiT + TRTLLM attention combined on Wan 2.2 I2V (480x832).""" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_fp8_cache_dit_trtllm(self): - if not os.path.exists(WAN22_I2V_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_I2V_PATH}") args = VisualGenArgs( - model=WAN22_I2V_PATH, + model=get_checkpoint(WAN22_I2V_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), quant_config={"quant_algo": "FP8", "dynamic": True}, attention_config=AttentionConfig(backend="TRTLLM"), diff --git a/tests/unittest/_torch/visual_gen/test_wan22_i2v_teacache.py b/tests/unittest/_torch/visual_gen/test_wan22_i2v_teacache.py index 7cc439162d4a..1f410a250cff 100644 --- a/tests/unittest/_torch/visual_gen/test_wan22_i2v_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_wan22_i2v_teacache.py @@ -25,10 +25,6 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan22_i2v_teacache.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN22_I2V=/path/to/wan22 \\ - pytest tests/unittest/_torch/visual_gen/test_wan22_i2v_teacache.py -v -s """ import os @@ -36,12 +32,12 @@ os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path import numpy as np import pytest import torch from PIL import Image +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TeaCacheConfig, VisualGenArgs @@ -63,28 +59,10 @@ def _cleanup_gpu(): # ============================================================================ -# Path helpers +# Checkpoints # ============================================================================ - -def _llm_models_root() -> Path: - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - else: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return root - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or str(_llm_models_root() / default_name) - - -WAN22_I2V_A14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN22_I2V", "Wan2.2-I2V-A14B-Diffusers") +WAN22_I2V_A14B_SUBDIR = "Wan2.2-I2V-A14B-Diffusers" INFER_NUM_FRAMES = 33 # (33-1)/4+1 = 9 latent frames; smallest realistic shape INFER_NUM_STEPS = 20 # Wan 2.2 has no reference hit rate; just enough to exercise both backends @@ -127,10 +105,8 @@ def _make_test_image(height: int, width: int) -> Image.Image: @pytest.fixture def wan22_i2v_pipeline(): - if not os.path.exists(WAN22_I2V_A14B_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_I2V_A14B_PATH}") args = VisualGenArgs( - model=WAN22_I2V_A14B_PATH, + model=get_checkpoint(WAN22_I2V_A14B_SUBDIR), cache_config=TeaCacheConfig( teacache_thresh=0.15, coefficients=WAN22_I2V_HIGH_NOISE_COEFFICIENTS, diff --git a/tests/unittest/_torch/visual_gen/test_wan22_t2v_offload.py b/tests/unittest/_torch/visual_gen/test_wan22_t2v_offload.py index f99119715d45..969c29307c06 100644 --- a/tests/unittest/_torch/visual_gen/test_wan22_t2v_offload.py +++ b/tests/unittest/_torch/visual_gen/test_wan22_t2v_offload.py @@ -16,10 +16,6 @@ Run all: pytest tests/unittest/_torch/visual_gen/test_wan22_t2v_offload.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN22_T2V=/path/to/wan22 \\ - pytest tests/unittest/_torch/visual_gen/test_wan22_t2v_offload.py -v -s """ import os @@ -27,12 +23,12 @@ os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path from typing import Optional import pytest import torch import torch.nn.functional as F +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import ( @@ -58,29 +54,7 @@ def _cleanup_gpu(): torch.cuda.empty_cache() -# ============================================================================ -# Path helpers -# ============================================================================ - - -def _llm_models_root() -> Path: - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - else: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return root - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or str(_llm_models_root() / default_name) - - -WAN22_A14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN22_T2V", "Wan2.2-T2V-A14B-Diffusers") +WAN22_A14B_SUBDIR = "Wan2.2-T2V-A14B-Diffusers" INFER_PROMPT = "A cat sitting on a sunny windowsill watching birds outside." INFER_NEGATIVE_PROMPT = "" @@ -102,8 +76,6 @@ def _make_pipeline( enable_offload: bool = False, quant_config: Optional[dict] = None, ): - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), @@ -122,12 +94,12 @@ def _require_fp8_quant_ops() -> None: _ = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor _ = torch.ops.tensorrt_llm.quantize_e4m3_activation except (AttributeError, RuntimeError) as e: - pytest.skip(f"FP8 quantization ops not available: {e}") + pytest.fail(f"FP8 quantization ops not available: {e}") @pytest.fixture def wan22_offload_pipeline(): - pipeline = _make_pipeline(WAN22_A14B_PATH, enable_offload=True) + pipeline = _make_pipeline(get_checkpoint(WAN22_A14B_SUBDIR), enable_offload=True) yield pipeline del pipeline torch.cuda.empty_cache() @@ -225,12 +197,12 @@ def test_wan22_offload(self, wan22_offload_pipeline): _assert_offload_forward(wan22_offload_pipeline, model="T2V-A14B") def test_wan22_offload_matches_baseline(self): - _assert_offload_matches_baseline(WAN22_A14B_PATH, model="Wan2.2-T2V-A14B") + _assert_offload_matches_baseline(get_checkpoint(WAN22_A14B_SUBDIR), model="Wan2.2-T2V-A14B") def test_wan22_fp8_offload_matches_baseline(self): _require_fp8_quant_ops() _assert_offload_matches_baseline( - WAN22_A14B_PATH, + get_checkpoint(WAN22_A14B_SUBDIR), model="Wan2.2-T2V-A14B FP8", quant_config={"quant_algo": "FP8", "dynamic": True}, ) @@ -242,10 +214,8 @@ class TestWanT2VOffloadCudaGraphRaisesError: """CUDA graphs plus offloading must raise NotImplementedError on pipeline load.""" def test_wan22_raises_if_cuda_graph_and_offload_enabled(self): - if not os.path.exists(WAN22_A14B_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_A14B_PATH}") args = VisualGenArgs( - model=WAN22_A14B_PATH, + model=get_checkpoint(WAN22_A14B_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), cuda_graph_config=CudaGraphConfig(enable=True), cpu_offload_config=CpuOffloadConfig(enable=True), diff --git a/tests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.py index 97abe4b574a6..f6a4930601a8 100644 --- a/tests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.py @@ -18,25 +18,20 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN22_T2V=/path/to/wan22 \\ - pytest tests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.py -v -s """ -import importlib import os os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path import numpy as np import pytest import torch import torch.nn.functional as F from diffusers import DiffusionPipeline +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import ( @@ -54,29 +49,7 @@ def _cleanup_mpi_env(): os.environ.pop("TLLM_DISABLE_MPI", None) -# ============================================================================ -# Path helpers -# ============================================================================ - - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN22_A14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN22_T2V", "Wan2.2-T2V-A14B-Diffusers") +WAN22_A14B_SUBDIR = "Wan2.2-T2V-A14B-Diffusers" # ============================================================================ # Test constants @@ -100,8 +73,6 @@ def _load_trtllm_pipeline( enable_offload: bool = False, ): """Load TRTLLM WanPipeline (two-stage) without torch.compile or warmup.""" - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") kwargs = dict( model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), @@ -280,7 +251,7 @@ class TestWan22_A14B_PipelineCorrectness: def test_cosine_similarity(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN22_A14B_PATH, + checkpoint_path=get_checkpoint(WAN22_A14B_SUBDIR), height=480, width=832, num_frames=9, @@ -290,7 +261,7 @@ def test_cosine_similarity(self): def test_cosine_similarity_with_offload(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN22_A14B_PATH, + checkpoint_path=get_checkpoint(WAN22_A14B_SUBDIR), height=480, width=832, num_frames=9, @@ -314,10 +285,8 @@ def test_cosine_similarity_with_offload(self): def _make_wan22_t2v(quant_config=None, attention_config=None): - if not os.path.exists(WAN22_A14B_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_A14B_PATH}") kwargs = dict( - model=WAN22_A14B_PATH, + model=get_checkpoint(WAN22_A14B_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), ) if quant_config is not None: @@ -358,7 +327,7 @@ class TestWan22TwoStageFeatures: def test_fp8_on_both_stages(self, wan22_t2v_fp8): """FP8 quantization is applied to both transformer and transformer_2.""" if wan22_t2v_fp8.transformer_2 is None: - pytest.skip("Not a two-stage checkpoint") + pytest.fail("Not a two-stage checkpoint") def _has_fp8(module): return any( @@ -373,7 +342,7 @@ def _has_fp8(module): def test_trtllm_attention_both_stages(self, wan22_t2v_trtllm): """TRTLLM self-attention and VANILLA cross-attention on both stages.""" if wan22_t2v_trtllm.transformer_2 is None: - pytest.skip("Not a two-stage checkpoint") + pytest.fail("Not a two-stage checkpoint") for stage_name, transformer in [ ("transformer", wan22_t2v_trtllm.transformer), @@ -403,11 +372,8 @@ class TestWan22T2VBatchGeneration: @pytest.fixture(scope="class") def wan22_t2v_full_pipeline(self): """Load full Wan 2.2 T2V pipeline (all components) for batch tests.""" - if not WAN22_A14B_PATH or not os.path.exists(WAN22_A14B_PATH): - pytest.skip("Checkpoint not available. Set DIFFUSION_MODEL_PATH_WAN22_T2V.") - args = VisualGenArgs( - model=WAN22_A14B_PATH, + model=get_checkpoint(WAN22_A14B_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -457,16 +423,13 @@ def test_batch_prompt_shape(self, wan22_t2v_full_pipeline): @pytest.mark.integration @pytest.mark.wan_t2v -@pytest.mark.skipif(importlib.util.find_spec("cache_dit") is None, reason="cache_dit not installed") class TestWan22T2VCombinedOptimizations: """FP8 + CacheDiT + TRTLLM attention combined on Wan 2.2 T2V (480x832).""" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_fp8_cache_dit_trtllm(self): - if not os.path.exists(WAN22_A14B_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_A14B_PATH}") args = VisualGenArgs( - model=WAN22_A14B_PATH, + model=get_checkpoint(WAN22_A14B_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), quant_config={"quant_algo": "FP8", "dynamic": True}, attention_config=AttentionConfig(backend="TRTLLM"), diff --git a/tests/unittest/_torch/visual_gen/test_wan22_t2v_teacache.py b/tests/unittest/_torch/visual_gen/test_wan22_t2v_teacache.py index f295fab62e88..0f41d4437a1a 100644 --- a/tests/unittest/_torch/visual_gen/test_wan22_t2v_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_wan22_t2v_teacache.py @@ -25,10 +25,6 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan22_t2v_teacache.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN22_T2V=/path/to/wan22 \\ - pytest tests/unittest/_torch/visual_gen/test_wan22_t2v_teacache.py -v -s """ import os @@ -36,10 +32,10 @@ os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path import pytest import torch +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TeaCacheConfig, VisualGenArgs @@ -61,28 +57,10 @@ def _cleanup_gpu(): # ============================================================================ -# Path helpers +# Checkpoints # ============================================================================ - -def _llm_models_root() -> Path: - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - else: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return root - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or str(_llm_models_root() / default_name) - - -WAN22_A14B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN22_T2V", "Wan2.2-T2V-A14B-Diffusers") +WAN22_A14B_SUBDIR = "Wan2.2-T2V-A14B-Diffusers" INFER_NUM_FRAMES = 33 # (33-1)/4+1 = 9 latent frames; smallest realistic shape INFER_NUM_STEPS = 20 # Wan 2.2 has no reference hit rate; just enough to exercise both backends @@ -114,10 +92,8 @@ def _checkpoint(env_var: str, default_name: str) -> str: @pytest.fixture def wan22_t2v_pipeline(): - if not os.path.exists(WAN22_A14B_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_A14B_PATH}") args = VisualGenArgs( - model=WAN22_A14B_PATH, + model=get_checkpoint(WAN22_A14B_SUBDIR), cache_config=TeaCacheConfig( teacache_thresh=0.15, coefficients=WAN22_T2V_HIGH_NOISE_COEFFICIENTS, diff --git a/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py index 087da72c71c3..388a70ffb2d7 100644 --- a/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py @@ -14,10 +14,6 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN22_TI2V_5B=/path/to/wan22_ti2v_5b \\ - pytest tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py -v -s """ import importlib @@ -27,7 +23,6 @@ import gc from contextlib import ExitStack -from pathlib import Path from unittest import mock import numpy as np @@ -37,6 +32,7 @@ from diffusers import WanImageToVideoPipeline as HFWanImageToVideoPipeline from diffusers import WanPipeline as HFWanPipeline from PIL import Image +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import ( @@ -53,32 +49,7 @@ def _cleanup_mpi_env(): os.environ.pop("TLLM_DISABLE_MPI", None) -# ============================================================================ -# Path helpers -# ============================================================================ - - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN22_TI2V_5B_PATH = _checkpoint( - "DIFFUSION_MODEL_PATH_WAN22_TI2V_5B", - "Wan2.2-TI2V-5B-Diffusers", -) +WAN22_TI2V_5B_SUBDIR = "Wan2.2-TI2V-5B-Diffusers" # ============================================================================ # Test constants @@ -109,8 +80,6 @@ def _make_test_image(height: int, width: int) -> Image.Image: def _load_trtllm_pipeline(checkpoint_path: str): """Load TRTLLM WanPipeline without torch.compile or warmup.""" - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") args = VisualGenArgs( model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), @@ -312,7 +281,7 @@ class TestWan22TI2V5B_T2V_PipelineCorrectness: def test_cosine_similarity(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN22_TI2V_5B_PATH, + checkpoint_path=get_checkpoint(WAN22_TI2V_5B_SUBDIR), mode="t2v", height=704, width=1280, @@ -330,7 +299,7 @@ class TestWan22TI2V5B_I2V_PipelineCorrectness: def test_cosine_similarity(self): _assert_pipeline_matches_hf( - checkpoint_path=WAN22_TI2V_5B_PATH, + checkpoint_path=get_checkpoint(WAN22_TI2V_5B_SUBDIR), mode="i2v", height=704, width=1280, @@ -352,11 +321,8 @@ class TestWan22TI2V5BBatchGeneration: @pytest.fixture(scope="class") def wan22_ti2v_5b_full_pipeline(self): """Load full Wan 2.2 TI2V-5B pipeline (all components) for batch tests.""" - if not WAN22_TI2V_5B_PATH or not os.path.exists(WAN22_TI2V_5B_PATH): - pytest.skip("Checkpoint not available. Set DIFFUSION_MODEL_PATH_WAN22_TI2V_5B.") - args = VisualGenArgs( - model=WAN22_TI2V_5B_PATH, + model=get_checkpoint(WAN22_TI2V_5B_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), ) pipeline = PipelineLoader(args).load(skip_warmup=True) @@ -444,16 +410,13 @@ def test_i2v_batch_prompt_shape(self, wan22_ti2v_5b_full_pipeline): @pytest.mark.integration @pytest.mark.wan_t2v @pytest.mark.wan_i2v -@pytest.mark.skipif(importlib.util.find_spec("cache_dit") is None, reason="cache_dit not installed") class TestWan22TI2V5BCombinedOptimizations: """FP8 + CacheDiT + TRTLLM attention combined on Wan 2.2 TI2V-5B (704x1280).""" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_fp8_cache_dit_trtllm(self): - if not os.path.exists(WAN22_TI2V_5B_PATH): - pytest.skip(f"Checkpoint not found: {WAN22_TI2V_5B_PATH}") args = VisualGenArgs( - model=WAN22_TI2V_5B_PATH, + model=get_checkpoint(WAN22_TI2V_5B_SUBDIR), torch_compile_config=TorchCompileConfig(enable=False), quant_config={"quant_algo": "FP8", "dynamic": True}, attention_config=AttentionConfig(backend="TRTLLM"), diff --git a/tests/unittest/_torch/visual_gen/test_wan_transformer.py b/tests/unittest/_torch/visual_gen/test_wan_transformer.py index b3132939f3d9..297f29b8e20d 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_wan_transformer.py @@ -16,11 +16,6 @@ Run one: pytest tests/unittest/_torch/visual_gen/test_wan_transformer.py -v -s -k t2v pytest tests/unittest/_torch/visual_gen/test_wan_transformer.py -v -s -k i2v - -Override checkpoint paths: - DIFFUSION_MODEL_PATH_WAN21_1_3B=/path/to/Wan2.1-T2V-1.3B-Diffusers \\ - DIFFUSION_MODEL_PATH_WAN21_I2V_480P=/path/to/Wan2.1-I2V-14B-480P-Diffusers \\ - pytest tests/unittest/_torch/visual_gen/test_wan_transformer.py -v -s """ import os @@ -28,13 +23,13 @@ os.environ["TLLM_DISABLE_MPI"] = "1" import gc -from pathlib import Path from types import SimpleNamespace import pytest import torch import torch.nn.functional as F from diffusers import WanTransformer3DModel as HFWanTransformer3DModel +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.config import ( @@ -61,34 +56,6 @@ def _cleanup_gpu(): torch.cuda.empty_cache() -# ============================================================================ -# Path helpers -# ============================================================================ - - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN21_1_3B_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN21_1_3B", "Wan2.1-T2V-1.3B-Diffusers") -WAN21_I2V_480P_PATH = _checkpoint( - "DIFFUSION_MODEL_PATH_WAN21_I2V_480P", "Wan2.1-I2V-14B-480P-Diffusers" -) - - COS_SIM_THRESHOLD = 0.99 DEVICE = "cuda" DTYPE = torch.bfloat16 @@ -365,9 +332,7 @@ class TestWanT2VTransformerCorrectness: @pytest.fixture(scope="class") def t2v_models(self): - if not os.path.exists(WAN21_1_3B_PATH): - pytest.skip(f"Checkpoint not found: {WAN21_1_3B_PATH}") - hf_model, our_model = _load_models(WAN21_1_3B_PATH) + hf_model, our_model = _load_models(get_checkpoint("Wan2.1-T2V-1.3B-Diffusers")) yield hf_model, our_model del hf_model, our_model torch.cuda.empty_cache() @@ -434,12 +399,7 @@ class TestWanI2VTransformerCorrectness: @pytest.fixture(scope="class") def i2v_models(self): - if not WAN21_I2V_480P_PATH or not os.path.exists(WAN21_I2V_480P_PATH): - pytest.skip( - "Checkpoint not found. " - "Set DIFFUSION_MODEL_PATH_WAN21_I2V_480P=/path/to/Wan2.1-I2V-14B-480P-Diffusers" - ) - hf_model, our_model = _load_models(WAN21_I2V_480P_PATH) + hf_model, our_model = _load_models(get_checkpoint("Wan2.1-I2V-14B-480P-Diffusers")) yield hf_model, our_model del hf_model, our_model torch.cuda.empty_cache() diff --git a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py index 979619e8e433..c1c5e08aa397 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py @@ -11,21 +11,17 @@ Run: pytest tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py -v -s - -Override checkpoint path: - DIFFUSION_MODEL_PATH_WAN21_VSA=/path/to/vsa \\ - pytest tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py -v -s """ import gc import os -from pathlib import Path os.environ["TLLM_DISABLE_MPI"] = "1" import pytest import torch import torch.nn.functional as F +from utils.llm_data import get_checkpoint from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader @@ -45,29 +41,7 @@ def _cleanup_mpi_env(): os.environ.pop("TLLM_DISABLE_MPI", None) -# ============================================================================ -# Path helpers -# ============================================================================ - - -def _llm_models_root() -> str: - """Return LLM_MODELS_ROOT path if set in env, assert when it's set but not a valid path.""" - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ["LLM_MODELS_ROOT"]) - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - assert root.exists(), ( - "Set LLM_MODELS_ROOT or ensure /home/scratch.trt_llm_data_ci/llm-models/ is accessible." - ) - return str(root) - - -def _checkpoint(env_var: str, default_name: str) -> str: - return os.environ.get(env_var) or os.path.join(_llm_models_root(), default_name) - - -WAN21_VSA_PATH = _checkpoint("DIFFUSION_MODEL_PATH_WAN21_VSA", "Wan2.1-VSA-T2V-14B-720P-Diffusers") +WAN21_VSA_SUBDIR = "Wan2.1-VSA-T2V-14B-720P-Diffusers" # ============================================================================ # Test constants @@ -85,14 +59,12 @@ def _checkpoint(env_var: str, default_name: str) -> str: # ============================================================================ -def _load_vsa_pipeline(checkpoint_path: str, vsa_sparsity: float = 0.0): +def _load_vsa_pipeline(checkpoint_subdir: str, vsa_sparsity: float = 0.0): """Load TRTLLM WanPipeline with CUTEDSL + VSA backend.""" - if not os.path.exists(checkpoint_path): - pytest.skip(f"Checkpoint not found: {checkpoint_path}") if not _cute_dsl_available: pytest.skip(f"CUTEDSL not available (requires Blackwell GPU): {_cute_dsl_import_error}") args = VisualGenArgs( - model=checkpoint_path, + model=get_checkpoint(checkpoint_subdir), attention_config=AttentionConfig( backend="CUTEDSL", sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=vsa_sparsity), @@ -137,7 +109,7 @@ def _cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> float: def _assert_vsa_matches_dense( - checkpoint_path: str, + checkpoint_subdir: str, height: int, width: int, num_frames: int, @@ -162,14 +134,14 @@ def _assert_vsa_matches_dense( ) # --- CuTe-DSL path --- - vsa_pipe = _load_vsa_pipeline(checkpoint_path, vsa_sparsity=vsa_sparsity) + vsa_pipe = _load_vsa_pipeline(checkpoint_subdir, vsa_sparsity=vsa_sparsity) vsa_video = _capture_trtllm_video(vsa_pipe, **common_kwargs) del vsa_pipe gc.collect() torch.cuda.empty_cache() # --- SDPA fallback reference (same VSA formulation, fine attn via SDPA) --- - sdpa_pipe = _load_vsa_pipeline(checkpoint_path, vsa_sparsity=vsa_sparsity) + sdpa_pipe = _load_vsa_pipeline(checkpoint_subdir, vsa_sparsity=vsa_sparsity) with patch.object(_vsa_module, "is_cute_supported", return_value=False): sdpa_video = _capture_trtllm_video(sdpa_pipe, **common_kwargs) del sdpa_pipe @@ -207,7 +179,7 @@ class TestWanVsa14B_PipelineCorrectness: def test_cosine_similarity(self): _assert_vsa_matches_dense( - checkpoint_path=WAN21_VSA_PATH, + checkpoint_subdir=WAN21_VSA_SUBDIR, height=720, width=1280, num_frames=9, @@ -223,7 +195,7 @@ class TestWanVsaSparse: """VSA at sparsity=0.9: config propagates, output is correctly shaped and finite.""" def test_sparse_vsa(self): - pipeline = _load_vsa_pipeline(WAN21_VSA_PATH, vsa_sparsity=0.9) + pipeline = _load_vsa_pipeline(WAN21_VSA_SUBDIR, vsa_sparsity=0.9) try: attn_cfg = pipeline.pipeline_config.primary_model_config.attention assert attn_cfg.backend == "CUTEDSL" diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index e9c9925dfb37..4a8788c50932 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -737,11 +737,13 @@ def test_explicit_value_overrides_model_preference(self, user_setting, assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is user_setting - def test_registered_models_prefer_v2(self): + def test_registered_models_prefer_v2(self) -> None: from tensorrt_llm._torch.models.modeling_utils import \ get_registered_model_class architectures = ( + "LlamaForCausalLM", + "Llama4ForConditionalGeneration", "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM", @@ -770,7 +772,7 @@ def test_registered_models_prefer_v2(self): assert model_cls is not None assert model_cls.get_preferred_kv_cache_manager_version() == "V2" - def test_registered_models_keep_v2_on_nixl(self): + def test_registered_models_keep_v2_on_nixl(self) -> None: """Models preferring V2 and the Python transceiver keep V2 on NIXL. Both sentinels start at 'auto'; production resolves the transceiver @@ -783,6 +785,7 @@ def test_registered_models_keep_v2_on_nixl(self): get_registered_model_class architectures = ( + "Llama4ForConditionalGeneration", "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM", diff --git a/tests/unittest/others/test_perf_fixed_sequence_length.py b/tests/unittest/others/test_perf_fixed_sequence_length.py index dac808d9b49a..706d67d5440e 100644 --- a/tests/unittest/others/test_perf_fixed_sequence_length.py +++ b/tests/unittest/others/test_perf_fixed_sequence_length.py @@ -64,20 +64,18 @@ def test_fixed_dataset_sequence_length( @pytest.mark.parametrize( - ("runtime", "model_name", "num_loras", "build_only"), + ("runtime", "model_name", "num_loras"), [ - ("bench", "", 0, True), - ("bench", "", 1, False), - ("serve", "qwen3_4b_eagle3", 0, False), - ("serve", "nemotron_3_nano_omni_nvfp4", 0, False), - ("serve", "nemotron_3_nano_omni_nvfp4_image", 0, False), + ("bench", "", 1), + ("serve", "qwen3_4b_eagle3", 0), + ("serve", "nemotron_3_nano_omni_nvfp4", 0), + ("serve", "nemotron_3_nano_omni_nvfp4_image", 0), ], ) def test_variable_dataset_does_not_infer_sequence_length( runtime: str, model_name: str, num_loras: int, - build_only: bool, ) -> None: config = perf_test.PerfTestConfig( model_name=model_name, @@ -86,7 +84,6 @@ def test_variable_dataset_does_not_infer_sequence_length( output_lens=[2000], num_loras=num_loras, ) - config.build_only = build_only assert config.get_fixed_dataset_sequence_length() is None diff --git a/tests/unittest/utils/llm_data.py b/tests/unittest/utils/llm_data.py index 118ba2fe05ef..a2dc64be8e44 100644 --- a/tests/unittest/utils/llm_data.py +++ b/tests/unittest/utils/llm_data.py @@ -1,6 +1,8 @@ -from test_common.llm_data import llm_datasets_root, llm_models_root +from test_common.llm_data import (get_checkpoint, llm_datasets_root, + llm_models_root) __all__ = [ + "get_checkpoint", "llm_datasets_root", "llm_models_root", ] diff --git a/tests/unittest/utils/test_llm_data.py b/tests/unittest/utils/test_llm_data.py new file mode 100644 index 000000000000..a6eb436ca572 --- /dev/null +++ b/tests/unittest/utils/test_llm_data.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the fail-loud checkpoint resolver in ``test_common.llm_data``.""" + +import pytest + +from utils.llm_data import get_checkpoint + + +@pytest.mark.cpu_only +def test_get_checkpoint_returns_staged_path(tmp_path, monkeypatch) -> None: + """A staged checkpoint under LLM_MODELS_ROOT resolves to its absolute path.""" + monkeypatch.setenv("LLM_MODELS_ROOT", str(tmp_path)) + (tmp_path / "my-model").mkdir() + + assert get_checkpoint("my-model") == str(tmp_path / "my-model") + + +@pytest.mark.cpu_only +def test_get_checkpoint_raises_on_missing(tmp_path, monkeypatch) -> None: + """A missing checkpoint fails loudly instead of silently skipping the test.""" + monkeypatch.setenv("LLM_MODELS_ROOT", str(tmp_path)) + + with pytest.raises(FileNotFoundError, match="absent-model"): + get_checkpoint("absent-model")