Skip to content

[None][feat] VisualGen: Add CuTe DSL JIT kernels for FMHA#15831

Merged
chang-l merged 5 commits into
NVIDIA:mainfrom
xrq-phys:ruqingx/cutedsl_fmha_jit
Jul 20, 2026
Merged

[None][feat] VisualGen: Add CuTe DSL JIT kernels for FMHA#15831
chang-l merged 5 commits into
NVIDIA:mainfrom
xrq-phys:ruqingx/cutedsl_fmha_jit

Conversation

@xrq-phys

@xrq-phys xrq-phys commented Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Summary by CodeRabbit

  • New Features

    • Added support for runtime-compiled attention kernels, with caching for repeated runs.
    • Expanded attention support to include block-scaled Q/K paths and new FP4-based quantization options.
    • Improved handling of causal, windowed, and optional LSE outputs in attention workflows.
  • Bug Fixes

    • Tightened validation for attention configurations to better reject unsupported combinations.
    • Removed outdated file-tracking and packaging rules for certain kernel binaries.

Description

Replace the packaged Blackwell CuTe DSL FMHA binaries with runtime JIT compilation and port the forward kernels and scheduling helpers into VisualGen using public CuTe DSL interfaces.

Add dense QK16PV8 and block-scaled MXFP8/NVFP4 attention paths, including quantized Q/K scale handling and QK smoothing. Extend the VisualGen attention configuration with validated backend-specific recipes while keeping internal kernel controls private.

Update kernel integration, packaging metadata, and unit and performance coverage for the JIT and quantized execution paths.

Performance (B300, TFLOPS, higher is better)

Sequence cubin (QK16PV8) JIT (QK16PV8) Difference (QK16PV8) JIT(MXFP8) JIT(NVFP4)
4,096 1,869.2 1,814.3 -2.9% 1,976.3 2,056.2
32,768 1,924.1 1,838.0 -4.5% 2,206.6 2,483.6
44,160 1,908.9 1,822.1 -4.5% 2,194.1 2,458.5
65,536 1,890.2 1,855.9 -1.8% 2,188.8 2,482.4
98,304 1,850.5 1,840.0 -0.6% 2,184.7 2,490.3
131,072 1,843.5 1,839.1 -0.2% 2,201.0 2,495.5

Above results use CuTe DSL 4.5.0, as pinned by requirements.txt.
Measured from CUPTI device-kernel timestamps collected by PyTorch Profiler/Kineto. Compilation, quantization, allocation, and model execution are excluded.

if CuTe DSL version was upgraded from 4.5.0 to 4.6.0

Sequence cubin (QK16PV8) JIT (QK16PV8) JIT(MXFP8) JIT(NVFP4)
44,160 1,909.3 1,818.0 2,243.9 2,540.1
131,072 1,842.6 1,839.5 2,265.4 2,513.3

Test Coverage

  • tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py
    • Verifies runtime JIT compilation and compile-cache population on Blackwell GPUs.
    • Compares dense BF16 and QK16PV8 FMHA outputs against PyTorch SDPA across causal/non-causal attention, MHA/GQA, multiple batch sizes, and sequence lengths.
    • Compares block-scaled MXFP8 and NVFP4 outputs against a dequantized BF16 SDPA reference across causal/non-causal and MHA/GQA configurations.
  • tests/unittest/_torch/visual_gen/test_visual_gen_args.py
    • Validates supported CUTEDSL recipes for QK16PV8, MXFP8, and NVFP4.
    • Verifies qk_sf_vec handling and rejection of unsupported backend/quantization combinations.
  • tests/unittest/_torch/visual_gen/test_attention_integration.py
    • Exercises CUTEDSL through the integrated VisualGen attention module.
    • Checks self-attention, cross-attention, and Wan-shaped attention against reference implementations.
    • Extends CUTEDSL coverage beyond head dimension 128 where supported.
  • tests/unittest/_torch/visual_gen/test_attention_perf.py
    • Exercises the JIT CUTEDSL backend in the existing performance matrix, including newly supported head dimensions.

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

@xrq-phys
xrq-phys requested review from a team as code owners July 1, 2026 10:57
@xrq-phys

xrq-phys commented Jul 1, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@coderabbitai

coderabbitai Bot commented Jul 1, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Replaces precompiled cubin-based Blackwell FMHA kernels with runtime CuTe DSL JIT compilation and a compile cache, adds block-scaled (MXFP8/NVFP4) Q/K quantization support with new qk_sf_vec config, new tile-scheduler and mask helper modules, and updates tests and packaging accordingly.

Changes

CuTe DSL Runtime Compilation Migration

Layer / File(s) Summary
Remove precompiled cubin packaging
.gitattributes, .gitignore, setup.py
Removes LFS tracking, gitignore exception, and package_data entries for precompiled Blackwell attention .so cubins.
Runtime kernel import and compile cache
tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py, tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/__init__.py
Adds runtime import/JIT of Blackwell CuTe kernel classes, GPU arch validation, and a compile cache with clear_cute_dsl_fmha_cache(); updates package exports to expose kernel classes and block-scaled utilities.
Forward entry point and block-scaled quantization
tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py
Adds cute_dsl_fmha_fwd supporting dense and block-scaled Q/K paths and _quantize_blockscaled_one for MXFP8/NVFP4 quantization with scale-factor computation.
CuTeDSLAttention integration and preprocessing
tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py, tensorrt_llm/_torch/visual_gen/attention_backend/utils.py
Updates backend init/typing, adds _smooth_qk, updates _prepare_inputs, and reworks _fwd to integrate V FP8 quantization and the new forward call signature; updates related documentation/comments.
FMHA tile scheduler and mask/numeric helpers
tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/helpers/*
Adds static/persistent tile schedulers, fused attention mask logic, and FP32/FP8 packed conversion helpers used by the CuTe DSL kernels.
Quant attention config schema and validation
tensorrt_llm/visual_gen/args.py, tests/unittest/_torch/visual_gen/test_visual_gen_args.py
Adds fp4 dtype and qk_sf_vec field to QuantAttentionConfig, reworks recipe validation tables/messages, and adds corresponding tests.
CuTe DSL functional and smoke tests
tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py
Replaces cubin-based tests with compile-cache smoke test and rewritten context-forward/block-scaled-forward correctness tests against an SDPA reference.
Simplify test backend-requirement helpers
tests/unittest/_torch/visual_gen/test_attention_integration.py, tests/unittest/_torch/visual_gen/test_attention_perf.py
Removes the head_dim parameter and CUTEDSL-specific head_dim==128 skip condition from _require_attention_backend helpers and updates call sites.

Estimated code review effort: 4 (Complex) | ~75 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Attn as CuTeDSLAttention
  participant Fwd as cute_dsl_fmha_fwd
  participant Cache as _COMPILE_CACHE
  participant Kernel as BlackwellFusedMultiHeadAttentionForward

  Attn->>Fwd: call with q, k, v, qk_sf_vec, q_sf, k_sf
  Fwd->>Fwd: validate NHD inputs, map mask/window args
  Fwd->>Cache: lookup compiled kernel by cache key
  alt cache miss
    Fwd->>Kernel: JIT-compile kernel (dense or block-scaled)
    Kernel-->>Cache: store compiled kernel
  end
  Fwd->>Kernel: launch on current CUDA stream
  Kernel-->>Fwd: output tensor, optional lse
  Fwd-->>Attn: return output
Loading

Suggested reviewers: QiJune, PerkzZheng, arysef, nv-guomingz, 2ez4bz

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 48.84% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly matches the main change: adding CuTe DSL JIT FMHA kernels for VisualGen.
Description check ✅ Passed The description covers the change, test coverage, checklist, and relevant performance notes.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 10

🧹 Nitpick comments (2)
tests/unittest/_torch/visual_gen/test_visual_gen_args.py (1)

110-155: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add negative CUTEDSL mismatch coverage.

Coverage in tests/unittest/_torch/visual_gen/test_visual_gen_args.py is insufficient for the new validation matrix. Please add CUTEDSL rejection cases for mismatched qk_dtype/qk_sf_vec, e.g. FP4 with qk_sf_vec=0 or 32, FP8 with qk_sf_vec=16, and an unsupported vector size like 17.

As per path instructions, tests/** reviews should assess coverage and suggest concrete file names; coverage here needs this follow-up in the same test file.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/visual_gen/test_visual_gen_args.py` around lines 110 -
155, Add negative coverage in AttentionConfig tests for CUTEDSL quant validation
in the same test module, since the current tests only cover valid fp8/fp4
combinations. Extend the existing quant config cases around
test_supported_quant_config_cute_mxfp8, test_supported_quant_config_cute_nvfp4,
and test_fp4_qk_dtype_rejected_on_trtllm to assert CUTEDSL rejects mismatched
qk_dtype/qk_sf_vec inputs such as fp4 with qk_sf_vec=0 or 32, fp8 with
qk_sf_vec=16, and an unsupported qk_sf_vec like 17, using ValidationError with
the existing unsupported quant_attention_config-style validation.

Source: Path instructions

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py (1)

543-571: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add return annotations on the new methods.

__init__() and _smooth_qk() are newly added but still unannotated, which goes against the repo’s Python typing rule and makes this backend harder to type-check.

Suggested fix
     def __init__(
         self,
         layer_idx: int = 0,
         num_heads: int = 8,
@@
         quant_attention_config: QuantAttentionConfig | None = None,
         skip_softmax_threshold_scale: float | None = None,
         **kwargs,
-    ):
+    ) -> None:
@@
     def _smooth_qk(
         self,
         q: torch.Tensor,
         k: torch.Tensor,
         alpha: float = 0.7,
-    ):
+    ) -> tuple[torch.Tensor, torch.Tensor]:

As per coding guidelines, “Always annotate functions with return types (use None if no return).”

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py` around lines
543 - 571, The new methods in Cute DSL backend are missing return type
annotations, which violates the typing guideline. Add an explicit return
annotation of None to the __init__ method in cute_dsl.py, and add the
appropriate return annotation to _smooth_qk based on its actual return value so
the class remains type-checkable. Use the method names __init__ and _smooth_qk
to locate the changes.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py`:
- Around line 578-585: The smooth scale computation in the attention backend can
produce NaNs when a Q/K channel is entirely zero because `q_max.pow(alpha) /
k_max.pow(1 - alpha)` hits `0/0` and `clamp()` does not recover it. Update the
scale logic in `cute_dsl.py` to guard the `s` calculation in the same block that
uses `q_max` and `k_max`, so zero-valued channels produce a finite fallback
scale instead of NaN. Keep the existing `q` and `k` rescaling path intact, but
ensure `s` is sanitized before `q.reciprocal()` and `k * s` are applied.
- Around line 92-105: Add an explicit divisibility check in cute_dsl_fmha_fwd
before any head-group reshape logic, since the later num_heads_q // num_heads_kv
and q.view(..., num_heads_kv, num_head_groups, ...) assume exact divisibility.
If num_heads_q is not a multiple of num_heads_kv, raise a clear ValueError
alongside the existing tensor validation near the Q/K/V/O shape checks so the
failure is caught early and reported with a contract-specific message.

In `@tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/helpers/fmha_helpers.py`:
- Around line 1019-1028: The backward-inference mask offset in fmha_helpers is
using the wrong sign, so the WINDOW_MASK_BWD_INFERENCE path shifts the attention
window in the opposite direction. Update the offset logic in the helper that
reads `cute.arch.thread_idx()` and computes `offset` so that
`WINDOW_MASK_BWD_INFERENCE` follows the same direction as the other BWD helpers,
using the `seqlen_q` minus `seqlen_k` form while keeping the inference mask
branch behavior otherwise unchanged.
- Around line 1143-1148: In the FP8 conversion branch that sets cvt_instruction,
replace the assert False fallback with an explicit TypeError so unsupported
types cannot pass silently if Python assertions are disabled. Keep the existing
cutlass.const_expr(fp8_type == cutlass.Float8E4M3FN) check, and in the else path
raise a clear error from this helper indicating the FP8 element type is
unsupported.
- Around line 827-831: The trailing-tile check in get_trip_count for
RESIDUAL_MASK_BWD is using the K dimension, which can produce the wrong mask
count. Update the RESIDUAL_MASK_BWD branch in fmha_helpers.py to base the
trailing-tile decision on the Q sequence length and its corresponding tile size
in the get_trip_count logic, so residual backward masking matches Q tiles
instead of seqlen_k.

In
`@tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/helpers/static_persistent_tile_scheduler.py`:
- Around line 541-566: The work-item validity check in
StaticPersistentTileScheduler should also reject CTAs that land in the padded
tail of the final partial cluster, since current_work_linear_idx only validates
against problem_layout_ncluster_mnl. Update the logic in the work-tile
coordinate path that builds cur_cluster_coord and cur_tile_coord so it
additionally verifies the computed tile coordinate is within
problem_shape_ntile_mnl before returning WorkTileInfo. Use the existing
scheduler fields and helpers in this method, such as
problem_layout_ncluster_mnl, cluster_shape_mn, and cur_tile_coord, so tail CTAs
are marked invalid when the problem size is not divisible by the cluster shape.
- Around line 798-809: The runtime tile coordinate in
static_persistent_tile_scheduler.py is missing the computed cta_id_in_cluster,
causing CTAs within the same cluster to map to the same tile. Update the
cur_tile_coord construction in the scheduler logic to include cta_id_in_cluster
alongside cluster_tile_coord_mn, and verify the placement logic in the create()
and runtime coordinate path stays consistent for cluster_shape_mn values other
than (1, 1).

In `@tensorrt_llm/visual_gen/args.py`:
- Around line 61-62: Fix the field description string in args.py so the two
adjacent literals are separated by a space; update the description used for the
Q/K quantization dtype option in the relevant Args/pydantic field so generated
schema/docs render “unquantized. fp4” instead of “unquantized.fp4”.

In `@tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py`:
- Around line 210-215: Add a block-scaled attention test case with a sequence
length that is not a multiple of 128 so the _quantize_blockscaled_one()
pad/unpad path is exercised end to end. Update the parametrized cases in
test_attention_cute_dsl.py (the attention test using batch_size, seq_len_q,
seq_len_kv, is_causal) to include at least one non-128-multiple
seq_len_q/seq_len_kv value, and mirror the same coverage in the related block
around the other referenced lines so the quantized path is actually validated.
- Around line 41-60: The smoke test only verifies the first `cute_dsl_fmha_fwd`
launch populates `_COMPILE_CACHE` once, but it does not confirm identical calls
reuse the same cached kernel. Update `test_cute_dsl_jit_compile_smoke` to invoke
`cute_dsl_fmha_fwd(q, k, v, out, is_causal=False, sm_scale=sm_scale)` a second
time after synchronization and assert `_COMPILE_CACHE` remains at size 1, so the
test covers cache-hit behavior and catches accidental recompilation under a new
key.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py`:
- Around line 543-571: The new methods in Cute DSL backend are missing return
type annotations, which violates the typing guideline. Add an explicit return
annotation of None to the __init__ method in cute_dsl.py, and add the
appropriate return annotation to _smooth_qk based on its actual return value so
the class remains type-checkable. Use the method names __init__ and _smooth_qk
to locate the changes.

In `@tests/unittest/_torch/visual_gen/test_visual_gen_args.py`:
- Around line 110-155: Add negative coverage in AttentionConfig tests for
CUTEDSL quant validation in the same test module, since the current tests only
cover valid fp8/fp4 combinations. Extend the existing quant config cases around
test_supported_quant_config_cute_mxfp8, test_supported_quant_config_cute_nvfp4,
and test_fp4_qk_dtype_rejected_on_trtllm to assert CUTEDSL rejects mismatched
qk_dtype/qk_sf_vec inputs such as fp4 with qk_sf_vec=0 or 32, fp8 with
qk_sf_vec=16, and an unsupported qk_sf_vec like 17, using ValidationError with
the existing unsupported quant_attention_config-style validation.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c3ad8c1a-872a-45d0-8076-600b278c749d

📥 Commits

Reviewing files that changed from the base of the PR and between 98d9d18 and 9eb7623.

⛔ Files ignored due to path filters (64)
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_e4m3_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_lse_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_skipsm_tvmffi.so is excluded by !**/*.so
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/x86_64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so is excluded by !**/*.so
📒 Files selected for processing (16)
  • .gitattributes
  • .gitignore
  • setup.py
  • tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py
  • tensorrt_llm/_torch/visual_gen/attention_backend/utils.py
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/__init__.py
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha_blockscaled.py
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/helpers/__init__.py
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/helpers/fmha_helpers.py
  • tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/helpers/static_persistent_tile_scheduler.py
  • tensorrt_llm/visual_gen/args.py
  • tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py
  • tests/unittest/_torch/visual_gen/test_attention_integration.py
  • tests/unittest/_torch/visual_gen/test_attention_perf.py
  • tests/unittest/_torch/visual_gen/test_visual_gen_args.py
💤 Files with no reviewable changes (3)
  • .gitattributes
  • .gitignore
  • setup.py

Comment thread tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py
Comment thread tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py Outdated
Comment thread tensorrt_llm/visual_gen/args.py Outdated
Comment thread tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py
Comment thread tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #56918 [ run ] triggered by Bot. Commit: 9eb7623 Link to invocation

@xrq-phys
xrq-phys force-pushed the ruqingx/cutedsl_fmha_jit branch from 9eb7623 to a68cec4 Compare July 1, 2026 11:12
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #56918 [ run ] completed with state FAILURE. Commit: 9eb7623
/LLM/main/L0_MergeRequest_PR pipeline #45723 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Comment thread tensorrt_llm/_torch/visual_gen/attention_backend/utils.py Outdated
@zhenhuaw-me
zhenhuaw-me requested a review from o-stoner July 2, 2026 02:53
@zhenhuaw-me

Copy link
Copy Markdown
Member

@o-stoner could you take a look at this PR? I am thinking that if any cute DSL infra that we can converge for this PR and #14280

@xrq-phys
xrq-phys force-pushed the ruqingx/cutedsl_fmha_jit branch from 6e52ac7 to ad46d3e Compare July 2, 2026 15:01
@xrq-phys

xrq-phys commented Jul 2, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #57200 [ run ] triggered by Bot. Commit: ad46d3e Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #57200 [ run ] completed with state SUCCESS. Commit: ad46d3e
/LLM/main/L0_MergeRequest_PR pipeline #45973 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@xrq-phys

xrq-phys commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #57462 [ run ] triggered by Bot. Commit: ad46d3e Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #57462 [ run ] completed with state FAILURE. Commit: ad46d3e
/LLM/main/L0_MergeRequest_PR pipeline #46199 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@xrq-phys

xrq-phys commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #57479 [ run ] triggered by Bot. Commit: ad46d3e Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #57479 [ run ] completed with state SUCCESS. Commit: ad46d3e
/LLM/main/L0_MergeRequest_PR pipeline #46214 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@xrq-phys

xrq-phys commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #57509 [ run ] triggered by Bot. Commit: ad46d3e Link to invocation

@xrq-phys

xrq-phys commented Jul 4, 2026

Copy link
Copy Markdown
Collaborator Author

/bot help

@github-actions

github-actions Bot commented Jul 4, 2026

Copy link
Copy Markdown

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental) --high-priority]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Supports wildcard * for pattern matching (e.g., "*PerfSanity*" matches all stages containing PerfSanity). Examples: "A10-PyTorch-1, xxx", "PerfSanity". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Supports wildcard * for pattern matching. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx", --extra-stage "Post-Merge".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

@xrq-phys
xrq-phys force-pushed the ruqingx/cutedsl_fmha_jit branch from 5a11ff5 to 8fc990e Compare July 16, 2026 07:05
@xrq-phys
xrq-phys requested a review from a team as a code owner July 16, 2026 07:05
@xrq-phys

Copy link
Copy Markdown
Collaborator Author

Thanks so much, Bowen !

Right now I have no causal xattn use case on my immediate list, but will check alignment whenever a use case comes up. Some future possibilities (currently out of scope):

  • Fully fused MoT might want top-left alignment in the future
  • AR-DiT might need bottom-right in the future.

Numerical correctness of the ~7k-line CUTLASS-derived port rests entirely on the SDPA-reference unit tests

Thanks for pointing out. Currently VisualGen's attention backends are tested through SDPA-reference + single-block runs + offline video quality checks. IIRC there are some ongoing plans to add e2e video-gen to the tests but still iterating on the time/coverage trade-off.

@xrq-phys

Copy link
Copy Markdown
Collaborator Author

Likely cause: PR #15831 changes telemetry-visible VisualGen argument schema in tensorrt_llm/visual_gen/args.py, including QuantAttentionConfig.qk_dtype, field descriptions, and supported AttentionConfig quantization recipes. The committed golden manifest tensorrt_llm/usage/llm_args_golden_manifest.json was not updated, so the golden-file consistency test fails.

Could you take a look at that and make sure that it is not the change in the PR caused this? Thanks!

Thanks for the comment, @pengbowang-nv !

I tried to rerun python3 scripts/generate_llm_args_golden_manifest.py but the golden file didn't change. I didn't observe any VisualGen-quantization-related entries in tensorrt_llm/usage/llm_args_golden_manifest.json, either.

To make it 100% safe that the previous failure is only caused by absence of rebase, I'll re-trigger the test to see if [A10-PyTorch-2] passes.

@xrq-phys

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59656 [ run ] triggered by Bot. Commit: 8fc990e Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59656 [ run ] completed with state SUCCESS. Commit: 8fc990e
/LLM/main/L0_MergeRequest_PR pipeline #48091 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@xrq-phys

Copy link
Copy Markdown
Collaborator Author

@zhenhuaw-me

Hey. I've found the reason that caused skipSm's behavior mismatch against trtllmGen kernels.

We need to fix a default value related to a CuTeDSL-specific optimization that allows skipping row_max update when the new row_max is within 8 * old_row_max. This 8.0 rescale threshold is coupling with skipSm affecting its behavior. To resolve, we need to reset this threshold to 1.0.

This fix is not kernel-side but unfortunately the value is set internally by host side of the kernel module. I'll raise a MR to our internal variant as well, but for TRTLLM here, please allow for one more CI loop as I'm about to push new changes.

Thanks

@xrq-phys

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59719 [ run ] triggered by Bot. Commit: f6421b8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59719 [ run ] completed with state SUCCESS. Commit: f6421b8
/LLM/main/L0_MergeRequest_PR pipeline #48145 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

xrq-phys and others added 5 commits July 20, 2026 12:16
Replace the packaged Blackwell CuTe DSL FMHA binaries with runtime JIT compilation and port the forward kernels and scheduling helpers into VisualGen using public CuTe DSL interfaces.

Add dense QK16PV8 and block-scaled MXFP8/NVFP4 attention paths, including quantized Q/K scale handling and QK smoothing. Extend the VisualGen attention configuration with validated backend-specific recipes while keeping internal kernel controls private.

Update kernel integration, packaging metadata, and unit and performance coverage for the JIT and quantized execution paths.

Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com>
Co-authored-by: Zhenhua Wang <4936589+zhenhuaw-me@users.noreply.github.com>
Signed-off-by: RuQing Xu <7891482+xrq-phys@users.noreply.github.com>
Config: Use qk_dtype=mxfp8 instead of qk_dtype=fp8 + reusing sage block sizes for DSL MXFP8
Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com>
Signed-off-by: RuQing Xu <7891482+xrq-phys@users.noreply.github.com>
There is a default-to-on optimization specific to CuTeDSL: usually skip-correction only happens when new_row_max < old_row_max , but CuTeDSL fmha.py has an optimization to allow skip updating row_max whenever new_row_max < 8.0 * old_row_max.

Usually this improves the performance as it reduces sm-correction frequency, but the same threshold is interacting with skipSoftmax threshold in a way that it sticks row_max to a lower value, making local_max/row_max < skip_sm_threshold harder to meet.

This commit resets this 8.0 threshold to 1.0 whenever skip-softmax is enabled.

Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com>
@xrq-phys
xrq-phys force-pushed the ruqingx/cutedsl_fmha_jit branch from f6421b8 to 534cdfe Compare July 20, 2026 03:17
@xrq-phys

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60246 [ run ] triggered by Bot. Commit: 534cdfe Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60246 [ run ] completed with state FAILURE. Commit: 534cdfe
/LLM/main/L0_MergeRequest_PR pipeline #48607 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@xrq-phys

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60311 [ run ] triggered by Bot. Commit: 534cdfe Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60311 [ run ] completed with state SUCCESS. Commit: 534cdfe
/LLM/main/L0_MergeRequest_PR pipeline #48659 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@xrq-phys

Copy link
Copy Markdown
Collaborator Author

/bot skip --comment "Error not related to VisualGen changes. Multi-GPU tests have all previously succeeded."

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60358 [ skip ] triggered by Bot. Commit: 534cdfe Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60358 [ skip ] completed with state SUCCESS. Commit: 534cdfe
Skipping testing for commit 534cdfe

Link to invocation

@xrq-phys

Copy link
Copy Markdown
Collaborator Author

@zhenhuaw-me I'm feeling confident about merging this now.

@chang-l
chang-l merged commit c7871ca into NVIDIA:main Jul 20, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants