From aad1edc3be7762c9abea078c97c8e735f00dfaa0 Mon Sep 17 00:00:00 2001 From: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> Date: Mon, 25 May 2026 12:09:41 +0900 Subject: [PATCH 1/3] visual_gen: add CuTe DSL attention via exported binaries Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> --- .gitattributes | 1 + .gitignore | 1 + docs/source/models/visual-generation.md | 72 ++- examples/visual_gen/README.md | 21 +- examples/visual_gen/visual_gen_flux.py | 22 +- examples/visual_gen/visual_gen_wan_i2v.py | 31 +- examples/visual_gen/visual_gen_wan_t2v.py | 31 +- setup.py | 3 +- .../visual_gen/attention_backend/__init__.py | 2 + .../visual_gen/attention_backend/cute_dsl.py | 229 ++++++++ .../visual_gen/attention_backend/parallel.py | 6 +- .../visual_gen/attention_backend/utils.py | 17 +- .../visual_gen/cute_dsl_kernels/__init__.py | 14 + .../cute_dsl_kernels/blackwell/__init__.py | 14 + .../blackwell/attention/__init__.py | 18 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ..._causal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...h128_causal_nonpersistent_varlen_tvmffi.so | 3 + ..._nonpersistent_varlen_lse_skipsm_tvmffi.so | 3 + ...ocausal_nonpersistent_varlen_lse_tvmffi.so | 3 + ...usal_nonpersistent_varlen_skipsm_tvmffi.so | 3 + ...28_nocausal_nonpersistent_varlen_tvmffi.so | 3 + .../blackwell/attention/fmha.py | 494 ++++++++++++++++++ tensorrt_llm/visual_gen/args.py | 97 ++-- .../test_lists/test-db/l0_b200.yml | 3 +- .../visual_gen/test_attention_cute_dsl.py | 248 +++++++++ .../visual_gen/test_attention_integration.py | 151 ++++-- .../_torch/visual_gen/test_attention_perf.py | 109 +++- .../_torch/visual_gen/test_visual_gen_args.py | 22 +- 86 files changed, 1642 insertions(+), 156 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/__init__.py create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py create mode 100644 tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py diff --git a/.gitattributes b/.gitattributes index cab2fa37863e..177818296355 100644 --- a/.gitattributes +++ b/.gitattributes @@ -16,3 +16,4 @@ docs/source/blogs/media/tech_blog10_full_strategy_performance.png filter=lfs dif docs/source/blogs/media/tech_blog10_context_wait_performance.png filter=lfs diff=lfs merge=lfs -text cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/cubin/kernelMetaInfo_cubin.cpp filter=lfs diff=lfs merge=lfs -text cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cubin/xqa_kernel_cubin.cpp filter=lfs diff=lfs merge=lfs -text +tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/*/*/*.so filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index a09c60676271..962f47d13419 100644 --- a/.gitignore +++ b/.gitignore @@ -57,6 +57,7 @@ tensorrt_llm/flash_mla/ tensorrt_llm/flash_mla_cpp_tllm.*.so tensorrt_llm/flash_mla_cpp_tllm.pyi tensorrt_llm/runtime/kv_cache_manager_v2/**/*.so +!tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/**/*.so **/*__mypyc*.so tensorrt_llm/scripts *docs/cpp_docs* diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 0c40ffe811a2..bcdc3552ddc9 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -12,8 +12,9 @@ Visual generation models based on diffusion transformers (DiT) have become the s TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion models, with a pipeline architecture separate from the LLM inference path. Key capabilities include: - A shared pipeline abstraction covering the denoising loop, guidance strategies, and component loading. -- Pluggable attention backends (PyTorch SDPA and TRT-LLM optimized kernels). +- Pluggable attention backends: PyTorch SDPA (`VANILLA`), TRT-LLM kernels (`TRTLLM`), TRT-LLM CuTe DSL kernels (`CUTEDSL`, Blackwell-class GPUs), and Flash Attention 4 (`FA4`). - Quantization support (dynamic and static) using the [ModelOpt](https://github.com/NVIDIA/TensorRT-Model-Optimizer) configuration format. +- Quantized attention support: `QK16PV8` to quantize Bmm2 on `CUTEDSL`, `SAGE` to run SageAttention on `TRTLLM` (requires Blackwell SM100). - Multi-GPU parallelism (CFG parallel, Ulysses sequence parallel). - **TeaCache** — a runtime caching optimization that skips transformer steps when timestep embeddings change slowly. - `trtllm-serve` integration with OpenAI-compatible API endpoints for image and video generation. @@ -107,6 +108,73 @@ args = VisualGenArgs( ) ``` +### Quantized Attention + +In addition to linear-layer quantization, VisualGen exposes two **attention-level** quantization presets that operate inside the attention kernel. They are configured through `AttentionConfig.quant_attention_config` (or the `--quant_attention_mode` flag in the example scripts) and are mutually exclusive with each other. + +- **QK16PV8** (`CUTEDSL` backend): Keeps Q & K in BF16 and quantizes only V to FP8 (E4M3, per-tensor), thus Bmm1 will be carried out in BF16 with Bmm2 in FP8. Targets Blackwell-class GPUs (`sm_100a` / `sm_103a`) with `head_dim = 128`. +- **SAGE** (`TRTLLM` backend): Quantizes Q, K, and V with per-block scaling factors. Q/K are stored as INT8 or FP8 (e4m3) and V as FP8 (e4m3); block sizes are tunable per axis (typically `(q, k, v) = (1, 4, 1)` for Wan-1.3B and `(1, 16, 1)` for larger Wan / FLUX checkpoints). Supported recipes are validated at runtime. + +CLI usage (Wan T2V, SageAttention): + +```bash +python visual_gen_wan_t2v.py \ + --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --prompt "A cute cat playing piano" \ + --attention_backend TRTLLM \ + --quant_attention_mode SAGE \ + --output_path output_sage.mp4 +``` + +CLI usage (Wan T2V, QK16PV8): + +```bash +python visual_gen_wan_t2v.py \ + --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --prompt "A cute cat playing piano" \ + --attention_backend CUTEDSL \ + --quant_attention_mode QK16PV8 \ + --output_path output_qk16pv8.mp4 +``` + +Programmatic equivalent (SageAttention): + +```python +from tensorrt_llm import VisualGenArgs + +args = VisualGenArgs( + model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + attention_config={ + "backend": "TRTLLM", + "quant_attention_config": { + "qk_dtype": "int8", + "q_block_size": 1, + "k_block_size": 16, + "v_block_size": 1, + }, + }, +) +``` + +Programmatic equivalent (QK16PV8): + +```python +from tensorrt_llm import VisualGenArgs + +args = VisualGenArgs( + model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + attention_config={ + "backend": "CUTEDSL", + "quant_attention_config": { + "qk_dtype": "bf16", + "q_block_size": 0, + "k_block_size": 0, + "v_block_size": 0, + }, + }, +) +``` + ### TeaCache TeaCache caches transformer outputs when timestep embeddings change slowly between denoising steps, skipping redundant computation. Enable via `VisualGenArgs.cache_config` (YAML or programmatic): @@ -126,7 +194,7 @@ The `teacache_thresh` parameter controls the similarity threshold. Cache-DiT is - **CFG Parallelism** (`--cfg_size 2`): Splits positive/negative guidance prompts across GPUs. - **Ulysses Parallelism** (`--ulysses_size N`): Splits the sequence dimension across GPUs for longer sequences. - **Parallel VAE** (`--parallel_vae_size N`): Shards the final VAE decode along a spatial axis across GPUs, useful to reduce VAE latency and improve GPU utilization (Constraint: `parallel_vae_size ≤ world_size`). Currently only supported for WAN models. -- **Attention Parallel**: There are 2 methods supported to run attention parallel. Both of these methods require the attention backend to support LSE (only FA4 currently) - +- **Attention Parallel**: There are 2 methods supported to run attention parallel. Both of these methods require the attention backend to support LSE (`FA4` and `CUTEDSL`) - - **Attention2D Parallelism** (`--attn2d_row_size N`, `--attn2d_col_size M`): Shards the sequence axis across a 2D `N x M` device mesh, all-gathering Q along rows and K/V along columns so each rank computes a sub-block of the attention matrix (total CP degree = `N * M`; not currently combinable with Ulysses). - **Ring Attention Parallelism** (`--ring_size N`): Shards the sequence axis across a 1D ring of `N` ranks and streams K/V blocks around the ring so each rank computes its attention output without materializing the full K/V (mutually exclusive with Attention2D). ## Developer Guide diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index b9a27d3e5fbf..8a784e46c89c 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -94,7 +94,18 @@ python visual_gen_wan_t2v.py \ --prompt "A cute cat playing piano" \ --height 480 --width 832 --num_frames 33 \ --attention_backend TRTLLM \ - --enable_sage_attention \ + --quant_attention_mode SAGE \ + --output_path output.mp4 +``` + +**With CuTe DSL Qk16Pv8:** +```bash +python visual_gen_wan_t2v.py \ + --model_path ${MODEL_ROOT}/Wan2.1-T2V-1.3B-Diffusers \ + --prompt "A cute cat playing piano" \ + --height 480 --width 832 --num_frames 33 \ + --attention_backend CUTEDSL \ + --quant_attention_mode QK16PV8 \ --output_path output.mp4 ``` @@ -114,7 +125,7 @@ WAN supports two parallelism modes that can be combined: - **CFG Parallelism**: Split positive/negative prompts across GPUs - **Sequence Parallelism**: - *Ulysses*: Split sequence along head dimension across GPUs; requires `ulysses_size` to divide the model's head count - - *Attention2D*: 2D mesh sequence parallelism; no head-count constraint; requires `--attention_backend FA4` + - *Attention2D*: 2D mesh sequence parallelism; no head-count constraint; requires `--attention_backend FA4` or `--attention_backend CUTEDSL` - Combining Ulysses and Attention2D is not yet supported @@ -284,8 +295,8 @@ python visual_gen_ltx2.py \ | `--image_cond_strength` | — | ✓ | 1.0 | Image conditioning strength | | `--enable_teacache` | ✓ | ✓ | — | False | Cache optimization | | `--teacache_thresh` | ✓ | ✓ | — | 0.2 | TeaCache similarity threshold | -| `--attention_backend` | ✓ | ✓ | — | VANILLA | `VANILLA`, `TRTLLM`, or `FA4` | -| `--enable_sage_attention` | ✓ | ✓ | — | False | SageAttention (requires `TRTLLM` attention backend) | +| `--attention_backend` | ✓ | ✓ | — | VANILLA | `VANILLA`, `TRTLLM`, `FA4`, or `CUTEDSL` | +| `--quant_attention_mode` | ✓ | ✓ | — | NO_QUANT | `NO_QUANT`, `QK16PV8` (requires `CUTEDSL` backend), or `SAGE` (requires `TRTLLM` backend) | | `--cfg_size` | — | ✓ | — | 1 | CFG parallelism | | `--ulysses_size` | ✓ | ✓ | — | 1 | Ulysses parallelism | | `--parallel_vae_size` | - | ✓ | — | 1 | Parallelism used for VAE | @@ -322,7 +333,7 @@ python visual_gen_ltx2.py \ - Sequence length must be divisible by `ulysses_size` **Attention2D Errors:** -- Requires `--attention_backend FA4` +- Requires `--attention_backend FA4` or `--attention_backend CUTEDSL` - Combining with `--ulysses_size` is not yet supported - Total GPUs = `cfg_size × attn2d_row_size × attn2d_col_size` - Sequence length must be divisible by `attn2d_row_size × attn2d_col_size` diff --git a/examples/visual_gen/visual_gen_flux.py b/examples/visual_gen/visual_gen_flux.py index 00ccbb0dd8bb..e9f860151669 100755 --- a/examples/visual_gen/visual_gen_flux.py +++ b/examples/visual_gen/visual_gen_flux.py @@ -209,15 +209,18 @@ def parse_args(): "--attention_backend", type=str, default="VANILLA", - choices=["VANILLA", "TRTLLM", "FA4"], + choices=["VANILLA", "TRTLLM", "FA4", "CUTEDSL"], help="Attention backend (VANILLA: PyTorch SDPA, TRTLLM: optimized kernels, " - "FA4: Flash Attention 4). " + "FA4: Flash Attention 4, CUTEDSL: CuTe DSL kernels). " "Note: TRTLLM falls back to VANILLA for cross-attention.", ) parser.add_argument( - "--enable_sage_attention", - action="store_true", - help="Enable SageAttention (per-block quantized Q/K/V). Requires TRTLLM backend.", + "--quant_attention_mode", + default="NO_QUANT", + choices=["NO_QUANT", "QK16PV8", "SAGE"], + help="Quantized attention presets: NO_QUANT: no quantization; " + "QK16PV8: quantize P@V(Bmm2) only (requires backend=CUTEDSL); " + "SAGE: Sage Attention algorithm (requires backend=TRTLLM).", ) # Parallelism @@ -342,7 +345,7 @@ def build_visual_gen_args(args) -> VisualGenArgs: cache_kwargs = {} attention_cfg: dict = {"backend": args.attention_backend} - if args.enable_sage_attention: + if args.quant_attention_mode == "SAGE": attention_cfg["quant_attention_config"] = { "qk_dtype": "int8", "q_block_size": 1, @@ -350,6 +353,12 @@ def build_visual_gen_args(args) -> VisualGenArgs: "v_block_size": 1, } logger.info("SageAttention: INT8 Q/K, blocks (1, 16, 1)") + elif args.quant_attention_mode == "QK16PV8": + attention_cfg["quant_attention_config"] = { + "qk_dtype": "bf16", + "v_dtype": "fp8", + } + logger.info("QK16PV8: BF16 Q/K, FP8 V (per-tensor)") kwargs = dict( revision=args.revision, @@ -450,6 +459,7 @@ def main(): "model_path": args.model_path, "linear_type": args.linear_type, "attention_backend": args.attention_backend, + "quant_attention_mode": args.quant_attention_mode, "height": args.height, "width": args.width, "steps": args.steps, diff --git a/examples/visual_gen/visual_gen_wan_i2v.py b/examples/visual_gen/visual_gen_wan_i2v.py index 671b02250779..b9a02b95b8d5 100644 --- a/examples/visual_gen/visual_gen_wan_i2v.py +++ b/examples/visual_gen/visual_gen_wan_i2v.py @@ -191,21 +191,18 @@ def parse_args(): "--attention_backend", type=str, default="VANILLA", - choices=["VANILLA", "TRTLLM", "FA4"], + choices=["VANILLA", "TRTLLM", "FA4", "CUTEDSL"], help="Attention backend (VANILLA: PyTorch SDPA, TRTLLM: optimized kernels, " - "FA4: Flash Attention 4). " + "FA4: Flash Attention 4, CUTEDSL: CuTe DSL kernels). " "Note: TRTLLM falls back to VANILLA for cross-attention.", ) - - # SageAttention (requires --attention_backend TRTLLM) parser.add_argument( - "--enable_sage_attention", - action="store_true", - help=( - "Enable SageAttention (per-block quantized Q/K/V). Requires TRTLLM backend. " - "Block layout is chosen from --model_path: (1, 4, 1) for Wan2.1, " - "(1, 16, 1) otherwise." - ), + "--quant_attention_mode", + default="NO_QUANT", + choices=["NO_QUANT", "QK16PV8", "SAGE"], + help="Quantized attention presets: NO_QUANT: no quantization; " + "QK16PV8: quantize P@V(Bmm2) only (requires backend=CUTEDSL); " + "SAGE: Sage Attention algorithm (requires backend=TRTLLM). ", ) # Parallelism @@ -331,10 +328,8 @@ def _wan_needs_fine_grained_sage(model_path: str) -> bool: def main(): args = parse_args() - attention_cfg = { - "backend": args.attention_backend, - } - if args.enable_sage_attention: + attention_cfg = {"backend": args.attention_backend} + if args.quant_attention_mode == "SAGE": k_block_size = 4 if _wan_needs_fine_grained_sage(args.model_path) else 16 attention_cfg["quant_attention_config"] = { "qk_dtype": "int8", @@ -343,6 +338,12 @@ def main(): "v_block_size": 1, } logger.info(f"SageAttention: INT8 Q/K, blocks (1, {k_block_size}, 1)") + elif args.quant_attention_mode == "QK16PV8": + attention_cfg["quant_attention_config"] = { + "qk_dtype": "bf16", + "v_dtype": "fp8", + } + logger.info("QK16PV8: BF16 Q/K, FP8 V (per-tensor)") if args.enable_cache_dit: cache_kwargs = {"cache_config": _cache_dit_config_from_args(args)} diff --git a/examples/visual_gen/visual_gen_wan_t2v.py b/examples/visual_gen/visual_gen_wan_t2v.py index ce4a32e766d6..f60226ecb906 100755 --- a/examples/visual_gen/visual_gen_wan_t2v.py +++ b/examples/visual_gen/visual_gen_wan_t2v.py @@ -185,21 +185,18 @@ def parse_args(): "--attention_backend", type=str, default="VANILLA", - choices=["VANILLA", "TRTLLM", "FA4"], + choices=["VANILLA", "TRTLLM", "FA4", "CUTEDSL"], help="Attention backend (VANILLA: PyTorch SDPA, TRTLLM: optimized kernels, " - "FA4: Flash Attention 4). " + "FA4: Flash Attention 4, CUTEDSL: CuTe DSL kernels). " "Note: TRTLLM falls back to VANILLA for cross-attention.", ) - - # SageAttention (requires --attention_backend TRTLLM) parser.add_argument( - "--enable_sage_attention", - action="store_true", - help=( - "Enable SageAttention (per-block quantized Q/K/V). Requires TRTLLM backend. " - "Block layout is chosen from --model_path: (1, 4, 1) for Wan2.x 1.3B, " - "(1, 16, 1) otherwise." - ), + "--quant_attention_mode", + default="NO_QUANT", + choices=["NO_QUANT", "QK16PV8", "SAGE"], + help="Quantized attention presets: NO_QUANT: no quantization; " + "QK16PV8: quantize P@V(Bmm2) only (requires backend=CUTEDSL); " + "SAGE: Sage Attention algorithm (requires backend=TRTLLM). ", ) # Parallelism @@ -357,10 +354,8 @@ def main(): else: parallel_str = "None" - attention_cfg = { - "backend": args.attention_backend, - } - if args.enable_sage_attention: + attention_cfg = {"backend": args.attention_backend} + if args.quant_attention_mode == "SAGE": k_block_size = 4 if _wan_needs_fine_grained_sage(args.model_path) else 16 attention_cfg["quant_attention_config"] = { "qk_dtype": "int8", @@ -369,6 +364,12 @@ def main(): "v_block_size": 1, } logger.info(f"SageAttention: INT8 Q/K, blocks (1, {k_block_size}, 1)") + elif args.quant_attention_mode == "QK16PV8": + attention_cfg["quant_attention_config"] = { + "qk_dtype": "bf16", + "v_dtype": "fp8", + } + logger.info("QK16PV8: BF16 Q/K, FP8 V (per-tensor)") if args.enable_cache_dit: cache_kwargs = {"cache_config": _cache_dit_config_from_args(args)} diff --git a/setup.py b/setup.py index c6f27bc0a229..5656531f5237 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -171,6 +171,7 @@ def has_ext_modules(self): # Include CUDA source for fused MoE align extension so runtime JIT can find it in wheels '_torch/auto_deploy/custom_ops/fused_moe/moe_align_kernel.cu', '_torch/auto_deploy/custom_ops/fused_moe/triton_fused_moe_configs/*', + '_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/**/*.so', 'usage/schemas/*.json', ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index 1391a65a4259..eed0c24d8745 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -20,6 +20,7 @@ simplified metadata that doesn't require KV caching. """ +from .cute_dsl import CuTeDSLAttention from .flash_attn4 import FlashAttn4Attention from .interface import AttentionBackend, AttentionTensorLayout from .parallel import Attention2DAttention, RingAttention, UlyssesAttention @@ -33,6 +34,7 @@ "AttentionTensorLayout", "get_visual_gen_attention_backend", "create_attention", + "CuTeDSLAttention", "FlashAttn4Attention", "TrtllmAttention", "TrtllmAttentionMetadata", diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py new file mode 100644 index 000000000000..013065600207 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py @@ -0,0 +1,229 @@ +# 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. +""" +CuTe DSL (NVIDIA kernels) Backend for Visual Generation Models + +Uses pre-compiled cubins derived from CUTLASS CuTe DSL FMHA. +Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16. +""" + +import math +from typing import Optional, Tuple + +import torch + +from tensorrt_llm.visual_gen.args import QuantAttentionConfig + +from ...attention_backend.interface import PredefinedAttentionMask +from .interface import AttentionBackend, AttentionTensorLayout + +_cute_dsl_import_error = None +try: + import tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.attention as cute_dsl + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.attention.fmha import ( + _cute_runtime_import_error, + ) + + if _cute_runtime_import_error is not None: + raise ImportError(_cute_runtime_import_error) +except (ImportError, OSError) as e: + cute_dsl = None + _cute_dsl_import_error = e + + +class CuTeDSLAttention(AttentionBackend): + """ + CuTe DSL (NVIDIA kernels) backend for diffusion models. + + Uses pre-compiled cubin kernels (head_dim=128 only). + """ + + def __init__( + self, + layer_idx: int = 0, + num_heads: int = 8, + head_dim: int = 64, + num_kv_heads: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + quant_attention_config: Optional[QuantAttentionConfig] = None, + skip_softmax_threshold_scale: Optional[float] = None, + **kwargs, + ): + # Only head_dim=128 cubins are packaged. + if head_dim != 128: + raise ValueError(f"CUTEDSL cubins require head_dim=128, got head_dim={head_dim}.") + self.layer_idx = layer_idx + self.num_heads = num_heads + self.head_dim = head_dim + self.num_kv_heads = num_kv_heads or num_heads + self.dtype = dtype + self.quant_attention_config = quant_attention_config + self.skip_softmax_threshold_scale = skip_softmax_threshold_scale + self.scale = 1.0 / math.sqrt(head_dim) + + # CuTe DSL expects [B, S, H, D] format + self._preferred_layout = AttentionTensorLayout.NHD + + def _prepare_inputs( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_mask: PredefinedAttentionMask, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool, torch.dtype]: + """Cast inputs to CuTeDSL-compatible dtype and resolve causal flag.""" + if _cute_dsl_import_error is not None: + raise ImportError( + f"CuTe DSL kernels are not available. Import error: {_cute_dsl_import_error}" + ) from _cute_dsl_import_error + + is_causal = attention_mask == PredefinedAttentionMask.CAUSAL + + # Packaged cubins support float16 and bfloat16 only. + origin_dtype = q.dtype + if q.dtype not in (torch.float16, torch.bfloat16): + q = q.to(torch.bfloat16) + k = k.to(torch.bfloat16) + v = v.to(torch.bfloat16) + return q, k, v, is_causal, origin_dtype + + # cute_dsl.cute_dsl_fmha_fwd is already decorated with @torch.compiler.disable + # Allow torch.compile to fuse preceding linear/norm with quantization of V / seq-preprocess + def _fwd( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + is_causal: bool, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size, seq_len_q, num_heads, _ = q.shape + _, seq_len_kv, _, value_head_dim = v.shape + out = torch.empty( + batch_size, + seq_len_q, + num_heads, + value_head_dim, + dtype=q.dtype, + device=q.device, + ) + lse = torch.empty( + batch_size, + seq_len_q, + num_heads, + dtype=torch.float32, + device=q.device, + ) + + # Options that instructs quantization of V + scale_v = kwargs.get("scale_v", 1.0) + if self.quant_attention_config is not None: + v_qscale = 448.0 / v.abs().amax().clamp(min=1e-3) + v = (v * v_qscale).to(torch.float8_e4m3fn) + scale_v = scale_v / v_qscale + + # Sequence preproc. + qo_indptr_host = [i * seq_len_q for i in range(batch_size + 1)] + qo_indptr = torch.tensor(qo_indptr_host).to(device=q.device, dtype=torch.int32) + kv_indptr_host = [i * seq_len_kv for i in range(batch_size + 1)] + kv_indptr = torch.tensor(kv_indptr_host).to(device=q.device, dtype=torch.int32) + + # Skip softmax. + skip_softmax_threshold_scale = self.skip_softmax_threshold_scale + if skip_softmax_threshold_scale is not None and skip_softmax_threshold_scale <= 0.0: + skip_softmax_threshold_scale = None + + cute_dsl.cute_dsl_fmha_fwd( + q.flatten(0, 1).contiguous(), + k.flatten(0, 1).contiguous(), + v.flatten(0, 1).contiguous(), + out.flatten(0, 1), + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + is_causal=is_causal, + sm_scale=self.scale, + lse=lse.flatten(0, 1).contiguous(), + scale_q=kwargs.get("scale_q", 1.0), + scale_k=kwargs.get("scale_k", 1.0), + scale_v=scale_v, + scale_o=kwargs.get("scale_o", 1.0), + max_qo_len=seq_len_q, + max_kv_len=seq_len_kv, + is_persistent=False, + skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale, + ) + return out, lse + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + **kwargs, + ) -> torch.Tensor: + """ + Forward pass using CuTe DSL (NVIDIA kernels). + + Dimensions are derived from tensor shapes (NHD layout: ``[B, S, H, D]``). + + Args: + q: Query tensor [batch_size, seq_len, num_heads, head_dim] + k: Key tensor [batch_size, seq_len_kv, num_kv_heads, head_dim] + v: Value tensor [batch_size, seq_len_kv, num_kv_heads, head_dim] + attention_mask: Attention mask type (CAUSAL or FULL) + + Returns: + Output tensor [batch_size, seq_len, num_heads, head_dim] + """ + output, _ = self.forward_with_lse(q, k, v, attention_mask=attention_mask, **kwargs) + return output + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass returning both output and log-sum-exp (LSE). + + Returns: + output: [batch_size, seq_len, num_heads, head_dim] + lse: [batch_size, num_heads, seq_len] - log-sum-exp per query position, + always in float32. Used for numerically stable combination of + partial attention results in Attention2D parallelism. + """ + q, k, v, is_causal, origin_dtype = self._prepare_inputs(q, k, v, attention_mask) + output, lse = self._fwd(q, k, v, is_causal, **kwargs) + if output.dtype != origin_dtype: + output = output.to(origin_dtype) + return output, lse.transpose(1, 2) + + @classmethod + def support_lse(cls) -> bool: + return True + + @property + def preferred_layout(self) -> AttentionTensorLayout: + """Return the preferred tensor layout for this backend.""" + return self._preferred_layout + + @classmethod + def support_fused_qkv(cls) -> bool: + return False diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index 0cd02787cd8b..f55564a2acd6 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -16,7 +16,7 @@ Parallelism Wrappers Wraps any attention backend with a parallelism strategy. Not a standalone -backend — compose around a real backend (VANILLA/TRTLLM/FA4). +backend — compose around a real backend (VANILLA/TRTLLM/FA4/CUTEDSL). """ @@ -256,8 +256,8 @@ class Attention2DAttention(AttentionBackend): Supported inner backends ------------------------ The inner backend must support LSE output (``support_lse() -> True``) — required - for the reduce-scatter combine step. Currently only the FA4 backend - (``FlashAttn4Attention``) meets this requirement. + for the reduce-scatter combine step. Currently the FA4 and CUTEDSL + backends meet this requirement. Note: ``AttentionTensorLayout.NHD`` and ``AttentionTensorLayout.HND`` are both handled transparently; transposition is applied before the inner forward and diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index 197374084c49..a3ce3df89cd9 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -37,7 +37,7 @@ def get_visual_gen_attention_backend( Get diffusion attention backend class by name. Args: - backend_name: Backend identifier ("VANILLA", "TRTLLM", "FA4") + backend_name: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL") Returns: Diffusion attention backend class @@ -48,9 +48,12 @@ def get_visual_gen_attention_backend( - "TRTLLM": Optimized for self-attention (requires same Q/KV seq lengths) Better performance but requires fused QKV - "FA4": Flash Attention 4; provides higher speedup on Blackwell GPUs (sm100) - Requires flash-attn package with cute interface + Requires flash-attn package with cute interface + - "CUTEDSL": CuTe DSL FMHA kernels; uses packaged cubins when present, + otherwise compiles from CuTe DSL source """ # Lazy imports to avoid circular dependency + from .cute_dsl import CuTeDSLAttention from .flash_attn4 import FlashAttn4Attention from .trtllm import TrtllmAttention from .vanilla import VanillaAttention @@ -63,6 +66,8 @@ def get_visual_gen_attention_backend( return TrtllmAttention elif backend_name == "FA4": return FlashAttn4Attention + elif backend_name == "CUTEDSL": + return CuTeDSLAttention else: # Default to VANILLA for maximum compatibility return VanillaAttention @@ -89,7 +94,7 @@ def create_attention( internally, simplifying the forward() call. Args: - backend: Backend identifier ("VANILLA", "TRTLLM", "FA4") + backend: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL") layer_idx: Layer index in the model num_heads: Number of attention heads head_dim: Dimension per head @@ -111,9 +116,9 @@ def create_attention( """ attn_cls = get_visual_gen_attention_backend(backend) - # Extract quant_attention_config from AttentionConfig and pass to TRTLLM backend. - # AttentionConfig validation disables unsupported recipes by normalizing - # quant_attention_config to None. + # Extract quant_attention_config from AttentionConfig and pass to backends + # that support it (TRTLLM SAGE recipes, CUTEDSL QK16PV8). AttentionConfig + # validation rejects unsupported (backend, recipe) combinations upstream. if attention_config is not None and attention_config.quant_attention_config is not None: kwargs["quant_attention_config"] = attention_config.quant_attention_config if backend.upper() == "TRTLLM": diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/__init__.py @@ -0,0 +1,14 @@ +# 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. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/__init__.py @@ -0,0 +1,14 @@ +# 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. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/__init__.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/__init__.py new file mode 100644 index 000000000000..fdf84901e4c5 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/__init__.py @@ -0,0 +1,18 @@ +# 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. + +from .fmha import cute_dsl_fmha_fwd + +__all__ = ["cute_dsl_fmha_fwd"] diff --git a/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 b/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 new file mode 100644 index 000000000000..93e56313b689 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bde8e845611c0ff97738d3822da3ab406169596065bc703a8c8f1af58aeb4fbd +size 711184 diff --git a/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 b/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 new file mode 100644 index 000000000000..8f7b825447a4 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d962a68d6d0c89ca5e71d0514b2ea5f1b1da3099c7c89ea4c05cf7c5468f6b0c +size 686552 diff --git a/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 b/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 new file mode 100644 index 000000000000..6f2a475d0922 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4bc6c9ff687d1a49bdff39dffcf9e08f74c0114dc63921b458695d1b1add12df +size 702960 diff --git a/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 b/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 new file mode 100644 index 000000000000..8c94270cef08 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d71eff9707d7e4464e6b2847ece251c352c783b60056086841cfc1ad5767489 +size 682416 diff --git a/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 b/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 new file mode 100644 index 000000000000..b12c8c09af05 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:117b99d4733a4cf500515a22f5fb2c87b5c89ffdf7736bc302fff48b27b68289 +size 645672 diff --git a/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 b/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 new file mode 100644 index 000000000000..be106aa226d8 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aed8d36370827808578ee1b2a7927a8bb964772ba1446a84d7eb66ae11857c28 +size 625128 diff --git a/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 b/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 new file mode 100644 index 000000000000..df218788b1ad --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:909f85670f5cfeae51fb7e7ca2ec6ee648f272c72c0e65545c2f94a084bbd8d7 +size 637440 diff --git a/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 b/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 new file mode 100644 index 000000000000..67cd56345e01 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cb1bd11b767de8d4b039a4456f799a0461cc0cfa423e4ce97d6bddf1f289a81d +size 616896 diff --git a/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 b/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 new file mode 100644 index 000000000000..05f14144f81e --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4451124313741a07ee5c616c5fbaf3db110a1ab531bfc5955eeb2a126eb33d3e +size 711096 diff --git a/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 b/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 new file mode 100644 index 000000000000..86aef559d4ab --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0bad02325167d25248a381c9e8cfa4479d93e4f7f2d2bbfff0816db8b03e52ea +size 690552 diff --git a/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 b/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 new file mode 100644 index 000000000000..2bd77cbe433b --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:38e71c573fad3a61051b4037cd38ca42daaa2cc6f2c1a007af74c7d5e0107b71 +size 702872 diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so new file mode 100644 index 000000000000..a15c97e7c4ff --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40cb6caa255f53e90c8c84fb5b5cf13b4776863d76f1d0ead0553fb3b574b169 +size 686424 diff --git a/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 b/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 new file mode 100644 index 000000000000..9b939a3c3917 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6d48995a1b50324fb7f103a543bd8aa2d48ab8617a972e7bc594a834f78e7888 +size 645576 diff --git a/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 b/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 new file mode 100644 index 000000000000..6273d4c74ad7 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2fff35882b5e0bfbebf2b313a67c7a80398a68df411a1fb30542f11a2e9212f1 +size 629136 diff --git a/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 b/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 new file mode 100644 index 000000000000..ba3f0e067efe --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1436b06493487e07c260e5ce283dcd85bfd9551db272d96d2e395788ed34ba5b +size 637352 diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so new file mode 100644 index 000000000000..daf2c2df15b8 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_100a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:460134c26759d391ac0d4df94627e63794898425f887a210871259fc030d2157 +size 620904 diff --git a/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 b/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 new file mode 100644 index 000000000000..10e981459b3c --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a6ea5c21b35d5562eec533b601348f527da2ed159c8d6f2a413f02d8b3dfce4 +size 657936 diff --git a/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 b/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 new file mode 100644 index 000000000000..e153954ec5ee --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ed589e4c6f369afb05903be7bee828ab5f8293736cca734390d6eab9ac609f56 +size 637400 diff --git a/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 b/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 new file mode 100644 index 000000000000..a8b2717c21dc --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:99801a8fb3c547c9f5778764b14049229e07a7395a2f26e098112ca3968884eb +size 653808 diff --git a/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 b/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 new file mode 100644 index 000000000000..0691e9326f5a --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:834109925ad81a8b50d44a436c34b9b9ff2f37bf87dc1c3420da10082f951f71 +size 629168 diff --git a/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 b/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 new file mode 100644 index 000000000000..a33382575553 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:962849852e70fd20d8f086b9257173ddda4f6ca654ed4fa165768a5398e50e75 +size 596520 diff --git a/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 b/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 new file mode 100644 index 000000000000..6a76fed84298 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8a040d0f9333c99a78cdf2baa42f2dc71cf2d8252ac3d1b7005f3afd7b4b536 +size 575976 diff --git a/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 b/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 new file mode 100644 index 000000000000..46d5cb25f686 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3c241a35bb7ab27b9b4fe26c03f1a8ac4ba1eeb28cfc68fcc0f0d5133a74f393 +size 588288 diff --git a/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 b/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 new file mode 100644 index 000000000000..073aa5c4860d --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:285541f2d3d081f8726d79c7e6f484af55a06e52e7580fb58a360f9aad5e846b +size 567744 diff --git a/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 b/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 new file mode 100644 index 000000000000..02af20762774 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b7284ff86f4d04fedf14e87947411e1765b55851ee9d9f12450f0314644abd8e +size 657848 diff --git a/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 b/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 new file mode 100644 index 000000000000..4eb755d3c573 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7f26d9fb521fea7260587558e4a5ec016c2c21b53f0fca580e65c46f7599c898 +size 641400 diff --git a/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 b/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 new file mode 100644 index 000000000000..a82de70cf9be --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d951a82391c4df1933bf1ff13bf7c455d90c95b19ebf3360c34a3151763ad9cc +size 653720 diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so new file mode 100644 index 000000000000..07ab962faa79 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ce18821322bdb617a5c26e52465ff1c02d6bd03f6f24ce7604a538a2862e107a +size 633176 diff --git a/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 b/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 new file mode 100644 index 000000000000..5396cdeac4a3 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5af2f95b9ff7328f1dce2863345ddbf72a39ec4cc0c3ce7601ad904d1c279bb3 +size 592328 diff --git a/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 b/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 new file mode 100644 index 000000000000..ec033517f17f --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:afe75060cd7f3dd1a37704e2a4fe70aec59bcf60e83f6cc77c510fae8adb1209 +size 575888 diff --git a/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 b/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 new file mode 100644 index 000000000000..490ac3c2eec1 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:807fc151395aabe1f4fb2bd0451d6d4f8def14e591cff121dc5426c2ee24b841 +size 588200 diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so new file mode 100644 index 000000000000..d708b622be2a --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/aarch64/sm_103a/cute_dsl_fmha_bf16_h128_nocausal_nonpersistent_varlen_tvmffi.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e0faf35561762e27fbf861fe3385298f08da3aa9cc9564b9d3cd99c7fb2f74e5 +size 567656 diff --git a/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 b/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 new file mode 100644 index 000000000000..334d4de2da9a --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:20f59a88225f55116a6269dde1eb677bd527cb91eb1200058765c8c64356b9fb +size 716248 diff --git a/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 b/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 new file mode 100644 index 000000000000..ffb16280b537 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f7666214fc5357f11f4db5cdae11c11f2c3a42679b30c1534535e12916aed578 +size 694232 diff --git a/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 b/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 new file mode 100644 index 000000000000..7a3a7cf23675 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cd18be0827da75961d58cb09d0c63fe0826ad9d1a83e16a44b74db43fb7b183b +size 705936 diff --git a/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 b/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 new file mode 100644 index 000000000000..86656cfff50e --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d3a3da59ea39eb5d41e345a04ba0ab573fb6ef7d950091d4757ac27af2c736e5 +size 683984 diff --git a/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 b/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 new file mode 100644 index 000000000000..4ce0c3cc1752 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ecd8a2048373f40ddf71a9706f7a16aef9fb0e5b6487937be695f6e7d0744161 +size 651128 diff --git a/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 b/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 new file mode 100644 index 000000000000..4b6997fefec4 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3407c3c3f177c8ab0b02c69c7cae9968710041d06acef9a10fc5ffa534ec274f +size 630200 diff --git a/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 b/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 new file mode 100644 index 000000000000..2c01e5f00009 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f64ccf41ae8680369363660c1c350386a8b9784141853832fb49d2767e283a9e +size 639728 diff --git a/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 b/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 new file mode 100644 index 000000000000..986c6a5db0d5 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4eca236d41155a0e2c5b0f3289c5399b60589ffb5b679abff456202566e2bdf0 +size 618192 diff --git a/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 b/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 new file mode 100644 index 000000000000..14609fab4bbc --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2d5ad69ef4a76110e48c4f1b6642e3e8b321dc1d794e83957b6228a8c68edfd4 +size 714872 diff --git a/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 b/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 new file mode 100644 index 000000000000..6c4ef0b8a4c7 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1896d588ebbbb6ffde8331432dafccbf996507504b6b86a0ba17bf539f2e8f9 +size 697240 diff --git a/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 b/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 new file mode 100644 index 000000000000..40433af18781 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7ce41d196fa54fb89e5848891e1e402d6511d0e5635098e98f00b0312ce3ad77 +size 704296 diff --git a/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 b/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 new file mode 100644 index 000000000000..14d97d8463c3 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2f45b423c9980f32c116718dca45f1b05e65ad9abd1e0a70e37a07c9106c6827 +size 686888 diff --git a/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 b/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 new file mode 100644 index 000000000000..faabf383e21f --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:150cd0b8fa1834d00aa2de8ead15916bc08c4e4f5ba002683ceccca5dc15a793 +size 650496 diff --git a/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 b/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 new file mode 100644 index 000000000000..2500b613f8e4 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c0eab2d74850c29abd02aa927cc53da3ac730ebed1f8d62e5836c68c3c86a42 +size 635008 diff --git a/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 b/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 new file mode 100644 index 000000000000..8c867ef9b986 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d044e658d6d3fda52a038a0d0d8b9cce502f7e383bb96deec344d10e217a6e98 +size 640056 diff --git a/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 b/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 new file mode 100644 index 000000000000..5449e7917a22 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:20bc7f069ea456ff8274d71dccaceeb979ba877e19a29d956753ebfa8e6000ad +size 624504 diff --git a/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 b/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 new file mode 100644 index 000000000000..e22d76d563cc --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:77fee6189e75f9b0233f7ccce9eec7d2f70b277f2f62d599c01f057e1c179235 +size 665400 diff --git a/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 b/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 new file mode 100644 index 000000000000..cc6eab46060c --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:182a606d2ca126fc83c2f914a6e5547d84a9aa668e3526710ca3e7396aac2f8e +size 642776 diff --git a/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 b/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 new file mode 100644 index 000000000000..9a06ad3879df --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:96052f92a4a4e9b8edc152a2e4e60d0b3c9ffc3e8cad0f72716f4be1072cea5e +size 656112 diff --git a/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 b/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 new file mode 100644 index 000000000000..04c07f32cdd9 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2c4ca62124eb16d29903dde0f9cf9a46b8a72d3480fee3447736b6c4f2d71895 +size 632560 diff --git a/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 b/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 new file mode 100644 index 000000000000..c27e7a6a19dc --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c9b8919f036acd199b4aed790563b63de6d1cfa3ef37bfeec00751a25ee5c29d +size 600408 diff --git a/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 b/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 new file mode 100644 index 000000000000..5f9f6cd54c9d --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eff71a1a8838b98ec684ada9e693db4f39a69527d8f072bb52ad3b136e2ad966 +size 580200 diff --git a/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 b/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 new file mode 100644 index 000000000000..62349117d032 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2e7d9c9e4ca723169a115029633811b97d6a57880274f668678218466bef9ee0 +size 591488 diff --git a/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 b/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 new file mode 100644 index 000000000000..4b09aab76f59 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:30f0681109e2520f4135d1bb4910d31c11b6d0aae64e09efc6705ba461e2eede +size 568848 diff --git a/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 b/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 new file mode 100644 index 000000000000..17b50b8f2a9b --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6a761c98700c565c98a0e7c7a67a569cf165cdf1fe1262c541ec066f60ffa7c1 +size 663704 diff --git a/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 b/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 new file mode 100644 index 000000000000..a4d1b3ecae30 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f21a9a2fc02f9ada3d126ee7c054f130fad0b16b62daadd63c881dc63991cf6b +size 644472 diff --git a/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 b/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 new file mode 100644 index 000000000000..883e738091b4 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b9edebdc6f4c23e407f057fba17144122f625db87a29bdd61d14951b8d05aadf +size 655176 diff --git a/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 b/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 new file mode 100644 index 000000000000..a22d88458722 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9becebe7efb3c0342b41668b6d82888e5a100bccdf906a25c10402ee76b56115 +size 635160 diff --git a/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 b/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 new file mode 100644 index 000000000000..836913616f49 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8239bc1f03bef805f46a356a9dcffa449dfa1f112ce8d39297f7cc272da60e48 +size 599536 diff --git a/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 b/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 new file mode 100644 index 000000000000..226972ea4776 --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:deb5f65d94e50b7466c027004c9fbfe0c7a5b9b67eec6c139496b70a7a0a85eb +size 581984 diff --git a/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 b/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 new file mode 100644 index 000000000000..bca39e1a5e2e --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dca9eeb79fe6a289db4cee5e9c2e8700dd9e138c149da1ec1299ea471efc3c8e +size 589096 diff --git a/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 b/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 new file mode 100644 index 000000000000..da1001774b1e --- /dev/null +++ b/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 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c24bd23172aa6d42fc3df1f7221e9ca3ce9c2b4e608386879f1734d58956fa45 +size 571544 diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py new file mode 100644 index 000000000000..3de28390c569 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/fmha.py @@ -0,0 +1,494 @@ +# 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. + +from __future__ import annotations + +import functools +import math +import platform +from pathlib import Path +from typing import Optional, Union + +import torch + +_cute_runtime_import_error = None +try: + import cutlass + from cuda.bindings import driver as cuda_driver + from cutlass import cute + from cutlass.cute import typing as cute_typing + from cutlass.cute.runtime import from_dlpack +except (ImportError, OSError) as e: + cutlass = None + cute = None + from_dlpack = None + cuda_driver = None + _cute_runtime_import_error = e + + +CUBINS_ROOT = Path(__file__).resolve().parent / "cubins" +SUPPORTED_GPU_ARCHS = ("sm_100a", "sm_103a") + + +def _dtype_to_str(dtype: torch.dtype) -> str: + float8_e4m3fn = getattr(torch, "float8_e4m3fn", None) + dtype_map = { + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.float32: "fp32", + } + if float8_e4m3fn is not None: + dtype_map[float8_e4m3fn] = "e4m3" + try: + return dtype_map[dtype] + except KeyError as exc: + raise ValueError(f"Unsupported CuTe DSL FMHA dtype: {dtype}") from exc + + +def _gpu_cpu_arch() -> str: + arch = platform.machine().lower() + if arch in ("amd64", "x64"): + return "x86_64" + if arch in ("arm64",): + return "aarch64" + return arch + + +def _get_gpu_arch(device: torch.device | str | None = None) -> str: + capability = torch.cuda.get_device_capability(device) + gpu_arch = f"sm_{capability[0]}{capability[1]}a" + if gpu_arch not in SUPPORTED_GPU_ARCHS: + supported = ", ".join(SUPPORTED_GPU_ARCHS) + raise ValueError( + f"Unsupported GPU architecture {gpu_arch}. Supported architectures: {supported}." + ) + return gpu_arch + + +def _get_cubins_dir(gpu_arch: str) -> Path: + return CUBINS_ROOT / _gpu_cpu_arch() / gpu_arch + + +def _get_variant_name( + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, + out_dtype: torch.dtype, + head_dim: int, + is_causal: bool, + is_persistent: bool = True, + varlen: bool = False, + with_lse: bool = False, + enable_skip_softmax: bool = False, + enable_tvm_ffi: bool = False, +) -> str: + qk_str = _dtype_to_str(qk_dtype) + pv_str = _dtype_to_str(pv_dtype) + out_str = _dtype_to_str(out_dtype) + if qk_dtype != pv_dtype: + dtype_str = f"{qk_str}_{pv_str}_{out_str}" + elif qk_dtype != out_dtype: + dtype_str = f"{qk_str}_{out_str}" + else: + dtype_str = qk_str + + causal_str = "causal" if is_causal else "nocausal" + persist_str = "persistent" if is_persistent else "nonpersistent" + varlen_str = "_varlen" if varlen else "" + lse_str = "_lse" if with_lse else "" + skip_str = "_skipsm" if enable_skip_softmax else "" + ffi_str = "_tvmffi" if enable_tvm_ffi else "" + return ( + f"cute_dsl_fmha_{dtype_str}_h{head_dim}_{causal_str}_{persist_str}" + f"{varlen_str}{lse_str}{skip_str}{ffi_str}" + ) + + +def _get_candidate_paths(variant_name: str, gpu_arch: str | None) -> list[Path]: + names = [f"{variant_name}.so", f"{variant_name}.o"] + if gpu_arch is not None: + return [_get_cubins_dir(gpu_arch) / name for name in names] + + host_dir = CUBINS_ROOT / _gpu_cpu_arch() + return [ + host_dir / supported_gpu_arch / name + for supported_gpu_arch in SUPPORTED_GPU_ARCHS + for name in names + ] + + +def _resolve_cubin_path( + variant_name: str, + gpu_arch: str | None = None, +) -> Path: + tried = [] + for candidate in _get_candidate_paths(variant_name, gpu_arch): + tried.append(candidate) + if candidate.exists(): + return candidate.resolve() + + searched = "\n".join(f" - {path}" for path in tried) + default_dir = ( + _get_cubins_dir(gpu_arch) + if gpu_arch is not None + else CUBINS_ROOT / _gpu_cpu_arch() / "" + ) + raise FileNotFoundError( + f"Could not find packaged CuTe DSL FMHA cubins for '{variant_name}'.\n" + f"Expected a .so or .o under {default_dir}.\nSearched:\n{searched}" + ) + + +def _check_cute_runtime_available() -> None: + if cute is not None: + return + raise ImportError( + f"CuTe DSL runtime is not available. Import error: {_cute_runtime_import_error}" + ) from _cute_runtime_import_error + + +def _load_cubin_from_path( + path: str | Path, + variant_name: str | None = None, + enable_tvm_ffi: bool = True, +): + _check_cute_runtime_available() + + cubin_path = Path(path).expanduser().resolve() + if variant_name is None: + variant_name = cubin_path.stem + + module = cute.runtime.load_module(str(cubin_path), enable_tvm_ffi=enable_tvm_ffi) + try: + return getattr(module, variant_name) + except AttributeError as exc: + raise AttributeError( + f"Loaded {cubin_path}, but symbol '{variant_name}' was not found. " + "The symbol name must match the .so filename stem / function prefix." + ) from exc + + +@functools.lru_cache(maxsize=None) +def _load_cute_dsl_fmha_cubin_cached( + variant_name: str, + gpu_arch: str | None, + enable_tvm_ffi: bool, +): + path = _resolve_cubin_path(variant_name, gpu_arch) + return _load_cubin_from_path(path, variant_name, enable_tvm_ffi) + + +def get_cute_dsl_fmha_cubin( + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, + out_dtype: torch.dtype, + head_dim: int, + is_causal: bool, + is_persistent: bool = True, + enable_tvm_ffi: bool = True, + varlen: bool = False, + with_lse: bool = False, + enable_skip_softmax: bool = False, + gpu_arch: str | None = None, +): + variant_name = _get_variant_name( + qk_dtype, + pv_dtype, + out_dtype, + head_dim, + is_causal, + is_persistent, + varlen, + with_lse, + enable_skip_softmax, + enable_tvm_ffi, + ) + return _load_cute_dsl_fmha_cubin_cached( + variant_name, + gpu_arch, + enable_tvm_ffi, + ) + + +def _check_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, +) -> bool: + if q.dim() not in (3, 4) or k.dim() != q.dim() or v.dim() != q.dim(): + raise ValueError("Expected 3D or 4D Q/K/V tensors with matching ranks.") + if q.dtype != k.dtype: + raise ValueError(f"Q/K dtype mismatch: {q.dtype} vs {k.dtype}") + if q.shape[-1] != k.shape[-1]: + raise ValueError(f"Q/K head dim mismatch: {q.shape[-1]} vs {k.shape[-1]}") + if k.shape[:-1] != v.shape[:-1]: + raise ValueError(f"K/V shape mismatch: {k.shape[:-1]} vs {v.shape[:-1]}") + expected_o_shape = (*q.shape[:-1], v.shape[-1]) + if tuple(o.shape) != expected_o_shape: + raise ValueError(f"Output shape mismatch: {tuple(o.shape)} vs {expected_o_shape}") + return q.dim() == 3 + + +def _to_cint_contiguous(tensor: torch.Tensor) -> torch.Tensor: + return tensor.to(torch.int32).contiguous() + + +def _to_cute_tensor(tensor: torch.Tensor, leading_dim: int): + float8_e4m3fn = getattr(torch, "float8_e4m3fn", None) + if tensor.dtype == float8_e4m3fn: + cute_tensor = from_dlpack(tensor.view(torch.int8), assumed_align=16) + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + cute_tensor.element_type = cutlass.Float8E4M3FN + return cute_tensor + return from_dlpack(tensor, assumed_align=16).mark_layout_dynamic(leading_dim=leading_dim) + + +def _get_runtime_problem( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: Optional[torch.Tensor], + qo_indptr: Optional[torch.Tensor], + kv_indptr: Optional[torch.Tensor], + max_qo_len: Optional[int], + max_kv_len: Optional[int], + varlen: bool, +) -> tuple[ + int, + int, + int, + int, + int, + int, + int, + int, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + if varlen: + if qo_indptr is None or kv_indptr is None or max_qo_len is None or max_kv_len is None: + raise ValueError("Varlen FMHA requires indptr tensors and max sequence lengths.") + if qo_indptr.dim() != 1 or kv_indptr.dim() != 1: + raise ValueError("Varlen FMHA indptr tensors must be 1D.") + if qo_indptr.numel() != kv_indptr.numel() or qo_indptr.numel() < 2: + raise ValueError("Varlen FMHA indptr tensors must have matching non-empty sizes.") + total_q, num_heads_q, head_dim = q.shape + _, num_heads_kv, _ = k.shape + batch_size = qo_indptr.numel() - 1 + max_s_q = max_qo_len + max_s_k = max_kv_len + q_4d = q.unsqueeze(0) + k_4d = k.unsqueeze(0) + v_4d = v.unsqueeze(0) + o_4d = o.unsqueeze(0) + lse_3d = lse.unsqueeze(0) if lse is not None else None + else: + batch_size, max_s_q, num_heads_q, head_dim = q.shape + _, max_s_k, num_heads_kv, _ = k.shape + total_q = batch_size * max_s_q + q_4d = q + k_4d = k + v_4d = v + o_4d = o + lse_3d = lse + value_head_dim = v.shape[-1] + return ( + batch_size, + max_s_q, + total_q, + max_s_k, + num_heads_q, + num_heads_kv, + head_dim, + value_head_dim, + q_4d, + k_4d, + v_4d, + o_4d, + lse_3d, + ) + + +@torch.compiler.disable +def cute_dsl_fmha_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + qo_indptr: Optional[torch.Tensor] = None, + kv_indptr: Optional[torch.Tensor] = None, + is_causal: bool = False, + sm_scale: Optional[float] = None, + window_left: int = -1, + window_right: int = -1, + lse: Optional[torch.Tensor] = None, + scale_q: Union[float, torch.Tensor] = 1.0, + scale_k: Union[float, torch.Tensor] = 1.0, + scale_v: Union[float, torch.Tensor] = 1.0, + scale_o: Union[float, torch.Tensor] = 1.0, + enable_tvm_ffi: bool = True, + is_persistent: bool = False, + max_qo_len: Optional[int] = None, + max_kv_len: Optional[int] = None, + kernel_fn=None, + skip_softmax_threshold_scale_factor: Optional[float] = None, +) -> None: + varlen = _check_inputs(q, k, v, o) + + scale_q = float(scale_q.item()) if isinstance(scale_q, torch.Tensor) else scale_q + scale_k = float(scale_k.item()) if isinstance(scale_k, torch.Tensor) else scale_k + scale_v = float(scale_v.item()) if isinstance(scale_v, torch.Tensor) else scale_v + scale_o = float(scale_o.item()) if isinstance(scale_o, torch.Tensor) else scale_o + + ( + batch_size, + max_s_q, + total_q, + max_s_k, + num_heads_q, + num_heads_kv, + head_dim, + value_head_dim, + q_4d, + k_4d, + v_4d, + o_4d, + lse_3d, + ) = _get_runtime_problem(q, k, v, o, lse, qo_indptr, kv_indptr, max_qo_len, max_kv_len, varlen) + use_skip_softmax = ( + skip_softmax_threshold_scale_factor is not None and skip_softmax_threshold_scale_factor > 0 + ) + problem_size = ( + batch_size, + max_s_q, + total_q, + max_s_k, + num_heads_q, + num_heads_kv, + head_dim, + value_head_dim, + ) + + if kernel_fn is None: + kernel_fn = get_cute_dsl_fmha_cubin( + q.dtype, + v.dtype, + o.dtype, + head_dim, + is_causal, + is_persistent=is_persistent, + varlen=varlen, + enable_tvm_ffi=enable_tvm_ffi, + with_lse=lse is not None, + enable_skip_softmax=use_skip_softmax, + gpu_arch=_get_gpu_arch(q.device), + ) + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + scale_softmax = scale_q * scale_k * sm_scale + scale_softmax_log2 = scale_softmax * math.log2(math.exp(1.0)) + scale_output = scale_v / scale_o + + skip_threshold_log2 = None + if use_skip_softmax: + threshold = skip_softmax_threshold_scale_factor / max_s_k + skip_threshold_log2 = cute_typing.Float32(math.log2(threshold)) + + ws_left = None if window_left == -1 else cute_typing.Int32(window_left) + ws_right = None if window_right == -1 else cute_typing.Int32(window_right) + if is_causal and ws_right is None: + ws_right = cute_typing.Int32(0) + + # CUBIN path + num_head_groups = num_heads_q // num_heads_kv + q_5d = q_4d.unflatten(2, (num_heads_kv, num_head_groups)) + k_5d = k_4d.unsqueeze(3) + v_5d = v_4d.unsqueeze(3) + o_5d = o_4d.unflatten(2, (num_heads_kv, num_head_groups)) + lse_4d = lse_3d.unflatten(2, (num_heads_kv, num_head_groups)) if lse is not None else None + + if enable_tvm_ffi: + qo_indptr_i32 = _to_cint_contiguous(qo_indptr) if varlen else None + kv_indptr_i32 = _to_cint_contiguous(kv_indptr) if varlen else None + kernel_fn( + q_5d, + k_5d, + v_5d, + o_5d, + problem_size, + qo_indptr_i32, + kv_indptr_i32, + lse_4d, + None, # sink + cute_typing.Float32(scale_softmax_log2), + cute_typing.Float32(scale_softmax), + cute_typing.Float32(scale_output), + skip_threshold_log2, + ws_left, + ws_right, + None, + None, + False, # reserved + ) + return + + q_cute = _to_cute_tensor(q_5d, leading_dim=4) + k_cute = _to_cute_tensor(k_5d, leading_dim=4) + v_cute = _to_cute_tensor(v_5d, leading_dim=4) + o_cute = _to_cute_tensor(o_5d, leading_dim=4) + + cum_seqlen_q_cute = None + cum_seqlen_k_cute = None + if varlen: + qo_indptr_i32 = _to_cint_contiguous(qo_indptr) + kv_indptr_i32 = _to_cint_contiguous(kv_indptr) + cum_seqlen_q_cute = from_dlpack(qo_indptr_i32, assumed_align=16).mark_layout_dynamic( + leading_dim=0 + ) + cum_seqlen_k_cute = from_dlpack(kv_indptr_i32, assumed_align=16).mark_layout_dynamic( + leading_dim=0 + ) + + lse_iter = None + if lse is not None: + lse_cute = from_dlpack(lse_4d, assumed_align=16).mark_layout_dynamic(leading_dim=2) + lse_iter = lse_cute.iterator + + stream = cuda_driver.CUstream(torch.cuda.current_stream(q).cuda_stream) + kernel_fn( + q_cute.iterator, + k_cute.iterator, + v_cute.iterator, + o_cute.iterator, + problem_size, + cum_seqlen_q_cute, + cum_seqlen_k_cute, + lse_iter, + None, # sink_iter + cute_typing.Float32(scale_softmax_log2), + cute_typing.Float32(scale_softmax), + cute_typing.Float32(scale_output), + skip_threshold_log2, + ws_left, + ws_right, + None, + None, + False, # reserved + stream, + ) diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 8f66a5ad31ab..c4a324779f0c 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -44,51 +44,43 @@ class QuantAttentionConfig(StrictBaseModel): - """Attention quantization recipe (TRTLLM backend only). + """Attention quantization recipe (TRTLLM / CUTEDSL backends). - Describes user intent for quantized attention: per-axis dtype and - per-block layout for Q, K, V. Providing this config to - ``AttentionConfig`` enables quantized attention; setting - ``AttentionConfig.quant_attention_config = None`` disables it. + Describes user intent for quantized attention: per-bmm dtype and per-block layout for Q, K, V. + Providing this config to AttentionConfig enables quantized attention; setting + AttentionConfig.quant_attention_config = None disables it. - Bare ``QuantAttentionConfig()`` is a valid finest-granularity recipe - (``qk_dtype="int8"``, ``v_dtype="fp8"``, ``q=k=v=1``); tune - ``k_block_size`` up (4 or 16) for speed. - - The ``"fp8"`` dtype maps to ``e4m3`` (``torch.float8_e4m3fn``) in - the current kernel; ``e5m2`` is not supported on this path. - - Unsupported recipes are rejected by ``AttentionConfig``'s validator - with a ``ValueError``. + Bare QuantAttentionConfig() is a valid Qk16Pv8 recipe. + Unsupported recipes are rejected by AttentionConfig's validator with a ValueError. """ - qk_dtype: Literal["int8", "fp8"] = Field( - "int8", + qk_dtype: Literal["bf16", "int8", "fp8"] = Field( + "bf16", status="prototype", - description="Q/K quantization dtype: 'int8' or 'fp8' (e4m3 in practice).", + description="Q/K quantization dtype; bf16 leaves Q/K unquantized.", ) v_dtype: Literal["fp8"] = Field( "fp8", status="prototype", - description="V quantization dtype. The current kernel always stores V in FP8 (e4m3).", + description="V quantization dtype. The current kernels always load V in FP8 (e4m3).", ) q_block_size: int = Field( - 1, - ge=1, + 0, + ge=0, status="prototype", - description="Elements per quantization block for Q.", + description="Elements per quantization block for Q; 0 for per-tensor quantization.", ) k_block_size: int = Field( - 1, - ge=1, + 0, + ge=0, status="prototype", - description="Elements per quantization block for K.", + description="Elements per quantization block for K; 0 for per-tensor quantization.", ) v_block_size: int = Field( - 1, - ge=1, + 0, + ge=0, status="prototype", - description="Elements per quantization block for V.", + description="Elements per quantization block for V; 0 for per-tensor quantization.", ) @@ -102,16 +94,16 @@ class QuantAttentionConfig(StrictBaseModel): class AttentionConfig(StrictBaseModel): """Configuration for Attention layers.""" - backend: Literal["VANILLA", "TRTLLM", "FA4"] = Field( + backend: Literal["VANILLA", "TRTLLM", "FA4", "CUTEDSL"] = Field( "VANILLA", status="prototype", - description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, FA4", + description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, FA4, CUTEDSL", ) quant_attention_config: Optional[QuantAttentionConfig] = Field( None, status="prototype", description=( - "Quantized-attention recipe (TRTLLM backend only). " + "Quantized-attention recipe (TRTLLM / CUTEDSL backends). " "Set to a QuantAttentionConfig instance to enable quantized " "attention; leave as None to disable." ), @@ -132,35 +124,48 @@ class AttentionConfig(StrictBaseModel): @model_validator(mode="after") def _validate_quant_attention_config(self) -> "AttentionConfig": - SUPPORTED_QUANT_RECIPES = { + # SAGE recipes target the TRTLLM backend (per-block Q/K/V scales). + SAGE_RECIPES = { ("int8", "fp8", (1, 1, 1)), ("int8", "fp8", (1, 4, 1)), ("int8", "fp8", (1, 16, 1)), ("fp8", "fp8", (1, 1, 1)), ("fp8", "fp8", (1, 4, 1)), } + # QK16PV8 (CUTEDSL backend): Q/K kept in bf16, V quantized to FP8. + QK16PV8_DTYPES = { + ("bf16", "fp8", (0, 0, 0)), + } if self.quant_attention_config is None: return self - if self.backend != "TRTLLM": - raise ValueError( - f"quant_attention_config requires backend='TRTLLM', " - f"got backend='{self.backend}'. Either set backend='TRTLLM' " - f"or remove quant_attention_config." - ) - - q = self.quant_attention_config + q_config = self.quant_attention_config recipe = ( - q.qk_dtype, - q.v_dtype, - (q.q_block_size, q.k_block_size, q.v_block_size), + q_config.qk_dtype, + q_config.v_dtype, + (q_config.q_block_size, q_config.k_block_size, q_config.v_block_size), ) - if recipe not in SUPPORTED_QUANT_RECIPES: + if self.backend == "TRTLLM": + if recipe not in SAGE_RECIPES: + raise ValueError( + f"Unsupported quant_attention_config={self.quant_attention_config!r} " + f"for backend='TRTLLM'. Supported SAGE recipes " + f"(qk_dtype, v_dtype, (q_block, k_block, v_block)): " + f"{sorted(SAGE_RECIPES)}." + ) + elif self.backend == "CUTEDSL": + if recipe not in QK16PV8_DTYPES: + raise ValueError( + f"Unsupported quant_attention_config={self.quant_attention_config!r} " + f"for backend='CUTEDSL'. Supported (qk_dtype, v_dtype): " + f"{sorted(QK16PV8_DTYPES)}." + ) + else: raise ValueError( - f"Unsupported quant_attention_config={self.quant_attention_config!r}. " - f"Supported recipes (qk_dtype, v_dtype, (q_block, k_block, v_block)): " - f"{sorted(SUPPORTED_QUANT_RECIPES)}." + f"quant_attention_config requires backend in ('TRTLLM', 'CUTEDSL'), " + f"got backend='{self.backend}'. Either change backend or " + f"remove quant_attention_config." ) return self diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 3b04ba55c708..1d1ebe9e2ae2 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -187,9 +187,10 @@ l0_b200: - unittest/_torch/visual_gen/test_teacache.py - unittest/_torch/visual_gen/test_cache_dit.py - unittest/_torch/visual_gen/test_quant_ops.py + - unittest/_torch/visual_gen/test_attention_cute_dsl.py + - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_attention_perf.py - - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py - unittest/_torch/visual_gen/test_trtllm_serve_e2e.py - unittest/_torch/visual_gen/test_model_loader.py diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py new file mode 100644 index 000000000000..eb6a6ec1093a --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py @@ -0,0 +1,248 @@ +# 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. + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.attention import cute_dsl_fmha_fwd +from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.attention.fmha import ( + get_cute_dsl_fmha_cubin, +) + + +def test_cute_dsl_cubin_kernel_can_import_and_load() -> None: + gpu_arch = _require_supported_gpu_arch() + try: + kernel = get_cute_dsl_fmha_cubin( + torch.bfloat16, + torch.bfloat16, + torch.bfloat16, + 128, + is_causal=False, + is_persistent=False, + varlen=True, + enable_tvm_ffi=True, + gpu_arch=gpu_arch, + ) + except ImportError as exc: + pytest.skip(str(exc)) + + assert kernel is not None + + +def _require_supported_gpu_arch() -> str: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for CuTe DSL FMHA kernels.") + + compute_capability = torch.cuda.get_device_capability() + gpu_arch = f"sm_{compute_capability[0]}{compute_capability[1]}a" + if gpu_arch not in ("sm_100a", "sm_103a"): + pytest.skip("CuTe DSL FMHA smoke tests require a supported Blackwell-class GPU.") + + return gpu_arch + + +def _make_indptr(lens: list[int], device: torch.device) -> torch.Tensor: + lens_tensor = torch.tensor(lens, dtype=torch.int32, device=device) + return torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + lens_tensor.cumsum(0).int(), + ] + ) + + +def _make_tensor( + max_len: int, + total_len: int, + num_heads: int, + head_dim: int, + dtype: torch.dtype, + ref_dtype: torch.dtype, + scale: float, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + tensor_f32 = torch.randn( + max_len + total_len, + num_heads, + head_dim, + dtype=torch.float32, + device=device, + ) + tensor_f32 *= 0.1 + if dtype == torch.float8_e4m3fn: + tensor = (tensor_f32 / scale).to(dtype) + ref = (tensor.float() * scale).to(ref_dtype) + else: + tensor = tensor_f32.to(dtype) + ref = tensor.to(ref_dtype) + return tensor[max_len:], ref[max_len:] + + +def _sdpa_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + is_causal: bool, + sm_scale: float, +) -> torch.Tensor: + out = [] + for batch_idx in range(qo_indptr.numel() - 1): + q_start = int(qo_indptr[batch_idx].item()) + q_end = int(qo_indptr[batch_idx + 1].item()) + kv_start = int(kv_indptr[batch_idx].item()) + kv_end = int(kv_indptr[batch_idx + 1].item()) + + q_i = q[q_start:q_end].transpose(0, 1).unsqueeze(0) + k_i = k[kv_start:kv_end].transpose(0, 1).unsqueeze(0) + v_i = v[kv_start:kv_end].transpose(0, 1).unsqueeze(0) + if q_i.shape[1] != k_i.shape[1]: + repeat_factor = q_i.shape[1] // k_i.shape[1] + k_i = k_i.repeat_interleave(repeat_factor, dim=1) + v_i = v_i.repeat_interleave(repeat_factor, dim=1) + out_i = F.scaled_dot_product_attention( + q_i, + k_i, + v_i, + is_causal=is_causal, + scale=sm_scale, + ) + out.append(out_i.squeeze(0).transpose(0, 1)) + return torch.cat(out, dim=0) + + +@pytest.mark.parametrize( + ("q_lens", "kv_lens", "is_causal"), + [ + pytest.param([256], [256], False, id="single_nocausal"), + pytest.param([512], [512], True, id="single_causal"), + pytest.param([64, 128], [128, 512], False, id="varlen_nocausal"), + ], +) +@pytest.mark.parametrize( + ("qk_dtype", "pv_dtype", "out_dtype", "atol", "rtol"), + [ + pytest.param( + torch.bfloat16, + torch.float8_e4m3fn, + torch.bfloat16, + 5e-2, + 5e-2, + id="bf16_fp8_bf16", + ), + pytest.param( + torch.bfloat16, + torch.bfloat16, + torch.bfloat16, + 2e-2, + 2e-2, + id="bf16", + ), + ], +) +@pytest.mark.parametrize( + ("num_heads", "num_heads_kv"), + [ + pytest.param(1, 1, id="mha_1h"), + pytest.param(4, 4, id="mha_4h"), + pytest.param(4, 2, id="gqa_4h_2kv"), + ], +) +@pytest.mark.parametrize("head_dim", [128]) +def test_cute_dsl_fmha_context_forward_cubin_smoke( + q_lens: list[int], + kv_lens: list[int], + is_causal: bool, + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, + out_dtype: torch.dtype, + num_heads: int, + num_heads_kv: int, + head_dim: int, + atol: float, + rtol: float, +) -> None: + gpu_arch = _require_supported_gpu_arch() + device = torch.device("cuda:0") + sm_scale = head_dim**-0.5 + scale_v = 0.06 if pv_dtype == torch.float8_e4m3fn else 1.0 + + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + qo_indptr = _make_indptr(q_lens, device) + kv_indptr = _make_indptr(kv_lens, device) + total_q = int(qo_indptr[-1].item()) + total_kv = int(kv_indptr[-1].item()) + max_qo_len = max(q_lens) + max_kv_len = max(kv_lens) + + q, q_ref = _make_tensor( + max_qo_len, total_q, num_heads, head_dim, qk_dtype, out_dtype, 1.0, device + ) + k, k_ref = _make_tensor( + max_kv_len, total_kv, num_heads_kv, head_dim, qk_dtype, out_dtype, 1.0, device + ) + v, v_ref = _make_tensor( + max_kv_len, total_kv, num_heads_kv, head_dim, pv_dtype, out_dtype, scale_v, device + ) + out_storage = torch.empty( + max_qo_len + total_q, + num_heads, + head_dim, + dtype=out_dtype, + device=device, + ) + out = out_storage[max_qo_len:] + kernel_fn = get_cute_dsl_fmha_cubin( + qk_dtype, + pv_dtype, + out_dtype, + head_dim, + is_causal, + is_persistent=False, + varlen=True, + enable_tvm_ffi=True, + gpu_arch=gpu_arch, + ) + + cute_dsl_fmha_fwd( + q, + k, + v, + out, + qo_indptr, + kv_indptr, + is_causal=is_causal, + sm_scale=sm_scale, + scale_v=scale_v, + max_qo_len=max_qo_len, + max_kv_len=max_kv_len, + kernel_fn=kernel_fn, + ) + torch.cuda.synchronize() + + out_ref = _sdpa_ref( + q_ref, + k_ref, + v_ref, + qo_indptr, + kv_indptr, + is_causal, + sm_scale, + ) + torch.testing.assert_close(out, out_ref, atol=atol, rtol=rtol) diff --git a/tests/unittest/_torch/visual_gen/test_attention_integration.py b/tests/unittest/_torch/visual_gen/test_attention_integration.py index 893304b4088f..52ebb45062c2 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -7,6 +7,7 @@ """ from types import SimpleNamespace +from typing import Optional import pytest import torch @@ -18,6 +19,7 @@ # ============================================================================ # Flash Attention 4 availability # ============================================================================ +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, @@ -29,6 +31,7 @@ from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig _flash_attn4_available = _fa4_fwd is not None +_cute_dsl_available = _cute_dsl_import_error is None # ============================================================================ # Original naive implementations for comparison @@ -142,6 +145,20 @@ def create_model_config( return config +def _require_attention_backend(attn_backend: str, head_dim: Optional[int] = None) -> None: + if attn_backend == "FA4" and not _flash_attn4_available: + pytest.fail("FlashAttention 4 backend is required for FA4 attention test") + if attn_backend == "CUTEDSL" and not _cute_dsl_available: + pytest.fail("CuTe DSL backend is required for CUTEDSL attention test") + if attn_backend == "CUTEDSL": + compute_capability = torch.cuda.get_device_capability() + gpu_arch = f"sm_{compute_capability[0]}{compute_capability[1]}a" + if gpu_arch not in ("sm_100a", "sm_103a"): + pytest.skip("CUTEDSL attention test requires a supported Blackwell-class GPU") + if head_dim is not None and head_dim != 128: + pytest.skip("CUTEDSL attention test requires head_dim=128") + + def copy_weights_self_attention(naive: NaiveWanSelfAttention, integrated: Attention): """Copy weights from naive to integrated self-attention.""" # QKV projection: naive has to_qkv, integrated has qkv_proj @@ -213,11 +230,22 @@ def generate_rope_embeddings( # ============================================================================ # Test functions # ============================================================================ -@pytest.mark.parametrize("attn_backend", ["VANILLA", "TRTLLM", "FA4"]) -def test_self_attention_equivalence(attn_backend: str): +@pytest.mark.parametrize("head_dim", [32, 128]) +@pytest.mark.parametrize( + ("attn_backend", "quant_attention_config"), + [ + ("VANILLA", None), + ("TRTLLM", None), + ("FA4", None), + ("CUTEDSL", None), + ("CUTEDSL", QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8")), + ], +) +def test_self_attention_equivalence( + head_dim: int, attn_backend: str, quant_attention_config: "QuantAttentionConfig | None" +): """Test that integrated self-attention produces same output as naive.""" - if attn_backend == "FA4" and not _flash_attn4_available: - pytest.fail("FlashAttention 4 backend is required for FA4 self-attention test") + _require_attention_backend(attn_backend, head_dim) print("\n" + "=" * 60) print("Testing Self-Attention Equivalence") @@ -226,9 +254,8 @@ def test_self_attention_equivalence(attn_backend: str): # Config batch_size = 2 seq_len = 16 - hidden_size = 128 num_heads = 4 - head_dim = hidden_size // num_heads + hidden_size = head_dim * num_heads device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.bfloat16 # Use bf16 since flashinfer doesn't support fp32 @@ -238,7 +265,13 @@ def test_self_attention_equivalence(attn_backend: str): # Create models naive = NaiveWanSelfAttention(hidden_size, num_heads, head_dim, dtype=dtype).to(device) - model_config = create_model_config(hidden_size, num_heads, head_dim, attn_backend=attn_backend) + model_config = create_model_config( + hidden_size, + num_heads, + head_dim, + attn_backend=attn_backend, + quant_attention_config=quant_attention_config, + ) integrated = Attention( hidden_size, num_heads, qkv_mode=QKVMode.FUSE_QKV, config=model_config ).to(device) # self attention @@ -266,13 +299,14 @@ def test_self_attention_equivalence(attn_backend: str): # Compare (using looser tolerance for bf16) max_diff = (out_naive - out_integrated).abs().max().item() mean_diff = (out_naive - out_integrated).abs().mean().item() - is_close = torch.allclose(out_naive, out_integrated, rtol=1e-2, atol=1e-2) + tol = 1e-2 if quant_attention_config is None else 2e-2 + is_close = torch.allclose(out_naive, out_integrated, rtol=tol, atol=tol) print("\nResults:") print(f" Output shape: naive={out_naive.shape}, integrated={out_integrated.shape}") print(f" Max absolute difference: {max_diff:.2e}") print(f" Mean absolute difference: {mean_diff:.2e}") - print(f" Outputs match (rtol=1e-2, atol=1e-2): {is_close}") + print(f" Outputs match (rtol={tol}, atol={tol}): {is_close}") if is_close: print(" ✅ PASS: Self-attention outputs match!") @@ -300,6 +334,10 @@ def test_sage_attention_self_attention(qk_dtype: str, batch_size: int, seq_len: 3. Outputs are finite (no NaN/Inf) 4. Approximate agreement with naive (cosine similarity > 0.99) """ + compute_capability = torch.cuda.get_device_capability() + gpu_arch = f"sm_{compute_capability[0]}{compute_capability[1]}a" + if qk_dtype == "int8" and gpu_arch not in ["sm_100a"]: + pytest.skip("Int8 kernels are only available for SM100 devices.") print("\n" + "=" * 60) print(f"Testing SageAttention (qk_dtype={qk_dtype}, B={batch_size}, S={seq_len})") print("=" * 60) @@ -386,11 +424,21 @@ def test_sage_attention_self_attention(qk_dtype: str, batch_size: int, seq_len: return cos_sim > 0.99 -@pytest.mark.parametrize("attn_backend", ["VANILLA", "FA4"]) -def test_cross_attention_equivalence(attn_backend: str): +@pytest.mark.parametrize("head_dim", [32, 128]) +@pytest.mark.parametrize( + ("attn_backend", "quant_attention_config"), + [ + ("VANILLA", None), + ("FA4", None), + ("CUTEDSL", None), + ("CUTEDSL", QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8")), + ], +) +def test_cross_attention_equivalence( + head_dim: int, attn_backend: str, quant_attention_config: "QuantAttentionConfig | None" +): """Test that integrated cross-attention produces same output as naive.""" - if attn_backend == "FA4" and not _flash_attn4_available: - pytest.fail("FlashAttention 4 backend is required for FA4 cross-attention test") + _require_attention_backend(attn_backend, head_dim) print("\n" + "=" * 60) print("Testing Cross-Attention Equivalence") @@ -400,9 +448,8 @@ def test_cross_attention_equivalence(attn_backend: str): batch_size = 2 seq_len = 16 encoder_seq_len = 24 # Different from query seq_len - hidden_size = 128 num_heads = 4 - head_dim = hidden_size // num_heads + hidden_size = num_heads * head_dim device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.bfloat16 # Use bf16 since flashinfer doesn't support fp32 @@ -414,7 +461,13 @@ def test_cross_attention_equivalence(attn_backend: str): # Create models naive = NaiveWanCrossAttention(hidden_size, num_heads, head_dim, dtype=dtype).to(device) - model_config = create_model_config(hidden_size, num_heads, head_dim, attn_backend=attn_backend) + model_config = create_model_config( + hidden_size, + num_heads, + head_dim, + attn_backend=attn_backend, + quant_attention_config=quant_attention_config, + ) integrated = Attention( hidden_size, num_heads, qkv_mode=QKVMode.SEPARATE_QKV, config=model_config ).to(device) # cross attention @@ -441,13 +494,14 @@ def test_cross_attention_equivalence(attn_backend: str): # Compare (using looser tolerance for bf16) max_diff = (out_naive - out_integrated).abs().max().item() mean_diff = (out_naive - out_integrated).abs().mean().item() - is_close = torch.allclose(out_naive, out_integrated, rtol=1e-2, atol=1e-2) + tol = 1e-2 if quant_attention_config is None else 2e-2 + is_close = torch.allclose(out_naive, out_integrated, rtol=tol, atol=tol) print("\nResults:") print(f" Output shape: naive={out_naive.shape}, integrated={out_integrated.shape}") print(f" Max absolute difference: {max_diff:.2e}") print(f" Mean absolute difference: {mean_diff:.2e}") - print(f" Outputs match (rtol=1e-2, atol=1e-2): {is_close}") + print(f" Outputs match (rtol={tol}, atol={tol}): {is_close}") if is_close: print(" ✅ PASS: Cross-attention outputs match!") @@ -467,12 +521,25 @@ def test_cross_attention_equivalence(attn_backend: str): (1, 2048, 512, 12, 128), ], ) -def test_fa4_cross_attention_wan_shapes( - batch: int, seq_len_q: int, seq_len_kv: int, num_heads: int, head_dim: int +@pytest.mark.parametrize( + ("attn_backend", "quant_attention_config"), + [ + ("FA4", None), + ("CUTEDSL", None), + ("CUTEDSL", QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8")), + ], +) +def test_fast_cross_attention_wan_shapes( + batch: int, + seq_len_q: int, + seq_len_kv: int, + num_heads: int, + head_dim: int, + attn_backend: str, + quant_attention_config: "QuantAttentionConfig | None", ): - """Test FA4 cross-attention correctness at Wan-realistic shapes.""" - if not _flash_attn4_available: - pytest.fail("FlashAttention 4 backend is required for FA4 Wan-shape tests") + """Test fast cross-attention correctness at Wan-realistic shapes.""" + _require_attention_backend(attn_backend, head_dim) hidden_size = num_heads * head_dim device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -489,15 +556,21 @@ def test_fa4_cross_attention_wan_shapes( device ) - cfg_fa4 = create_model_config(hidden_size, num_heads, head_dim, attn_backend="FA4") - fa4_model = Attention(hidden_size, num_heads, qkv_mode=QKVMode.SEPARATE_QKV, config=cfg_fa4).to( - device + cfg_fast = create_model_config( + hidden_size, + num_heads, + head_dim, + attn_backend=attn_backend, + quant_attention_config=quant_attention_config, ) + fast_model = Attention( + hidden_size, num_heads, qkv_mode=QKVMode.SEPARATE_QKV, config=cfg_fast + ).to(device) copy_weights_cross_attention(naive, ref) - copy_weights_cross_attention(naive, fa4_model) + copy_weights_cross_attention(naive, fast_model) ref.eval() - fa4_model.eval() + fast_model.eval() torch.manual_seed(42) hidden_states = torch.randn(batch, seq_len_q, hidden_size, device=device, dtype=dtype) @@ -505,12 +578,13 @@ def test_fa4_cross_attention_wan_shapes( with torch.no_grad(): out_ref = ref(hidden_states, encoder_hidden_states) - out_fa4 = fa4_model(hidden_states, encoder_hidden_states) + out_fast = fast_model(hidden_states, encoder_hidden_states) - max_diff = (out_ref - out_fa4).abs().max().item() - is_close = torch.allclose(out_ref, out_fa4, rtol=1e-2, atol=1e-2) + max_diff = (out_ref - out_fast).abs().max().item() + tol = 1e-2 if quant_attention_config is None else 2e-2 + is_close = torch.allclose(out_ref, out_fast, rtol=tol, atol=tol) print(f" Max diff: {max_diff:.2e}, match: {is_close}") - assert is_close, f"FA4 cross-attn mismatch at Wan shapes: max_diff={max_diff:.2e}" + assert is_close, f"{attn_backend} cross-attn mismatch at Wan shapes: max_diff={max_diff:.2e}" def test_trtllm_cached_prepare(): @@ -686,8 +760,11 @@ def run_all_tests(): results = {} # Run self-attention tests with different backends - for backend in ["VANILLA", "TRTLLM"] + (["FA4"] if _flash_attn4_available else []): - results[f"self_attention_{backend}"] = test_self_attention_equivalence(backend) + fast_backends = ["FA4"] if _flash_attn4_available else [] + if _cute_dsl_available: + fast_backends.append("CUTEDSL") + for backend in ["VANILLA", "TRTLLM"] + fast_backends: + results[f"self_attention_{backend}"] = test_self_attention_equivalence(backend, "NO_QUANT") # Run SageAttention self-attention tests (subset for manual runner) for batch_size in [1, 2]: @@ -699,9 +776,11 @@ def run_all_tests(): ) # Run cross-attention tests - results["cross_attention_VANILLA"] = test_cross_attention_equivalence("VANILLA") + results["cross_attention_VANILLA"] = test_cross_attention_equivalence("VANILLA", "NO_QUANT") if _flash_attn4_available: - results["cross_attention_FA4"] = test_cross_attention_equivalence("FA4") + results["cross_attention_FA4"] = test_cross_attention_equivalence("FA4", "NO_QUANT") + if _cute_dsl_available: + results["cross_attention_CUTEDSL"] = test_cross_attention_equivalence("CUTEDSL", "NO_QUANT") # Run TRTLLM-specific caching tests results["trtllm_cached_prepare"] = test_trtllm_cached_prepare() diff --git a/tests/unittest/_torch/visual_gen/test_attention_perf.py b/tests/unittest/_torch/visual_gen/test_attention_perf.py index 27cffc6213d8..90a8ac3608c0 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_perf.py +++ b/tests/unittest/_torch/visual_gen/test_attention_perf.py @@ -39,6 +39,7 @@ # ============================================================================ # Flash Attention 4 availability # ============================================================================ +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error 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.flash_attn4 import ( _flash_attn_fwd_import_error as _fa4_import_error, @@ -51,6 +52,30 @@ from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig _flash_attn4_available = _fa4_fwd is not None +_cute_dsl_available = _cute_dsl_import_error is None + + +def _require_attention_backend(backend: str, head_dim: Optional[int] = None) -> None: + if backend == "FA4" and not _flash_attn4_available: + pytest.fail( + "FlashAttention 4 backend is required for FA4 attention perf test" + + (f": {_fa4_import_error}" if _fa4_import_error else "") + ) + if backend == "CUTEDSL" and not _cute_dsl_available: + pytest.fail( + "CuTe DSL backend is required for CUTEDSL attention perf test" + + (f": {_cute_dsl_import_error}" if _cute_dsl_import_error else "") + ) + if backend == "CUTEDSL": + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for CUTEDSL attention perf test") + compute_capability = torch.cuda.get_device_capability() + gpu_arch = f"sm_{compute_capability[0]}{compute_capability[1]}a" + if gpu_arch not in ("sm_100a", "sm_103a"): + pytest.skip("CUTEDSL attention perf test requires a supported Blackwell-class GPU") + if head_dim is not None and head_dim != 128: + pytest.skip("CUTEDSL attention perf test requires head_dim=128") + # NVTX support for profiling try: @@ -277,10 +302,21 @@ def create_attention_model( return model def create_cross_attention_model( - self, hidden_size: int, num_heads: int, head_dim: int, backend: str + self, + hidden_size: int, + num_heads: int, + head_dim: int, + backend: str, + quant_attention_config: "QuantAttentionConfig | None" = None, ) -> Attention: """Create a WAN cross-attention model with specified backend.""" - config = create_model_config(hidden_size, num_heads, head_dim, attn_backend=backend) + config = create_model_config( + hidden_size, + num_heads, + head_dim, + attn_backend=backend, + quant_attention_config=quant_attention_config, + ) model = Attention(hidden_size, num_heads, qkv_mode=QKVMode.SEPARATE_QKV, config=config).to( self.device ) @@ -318,6 +354,7 @@ def benchmark_cross_attn_single( seq_len_kv: int, head_dim: int, backend: str, + quant_attention_config: "QuantAttentionConfig | None" = None, verbose: bool = True, ) -> Optional[Dict]: """Benchmark a single cross-attention configuration. @@ -328,7 +365,13 @@ def benchmark_cross_attn_single( hidden_size = num_heads * head_dim try: - model = self.create_cross_attention_model(hidden_size, num_heads, head_dim, backend) + model = self.create_cross_attention_model( + hidden_size, + num_heads, + head_dim, + backend, + quant_attention_config=quant_attention_config, + ) hidden_states = torch.randn( batch_size, seq_len_q, hidden_size, device=self.device, dtype=self.dtype @@ -473,6 +516,7 @@ def benchmark_comparison( seq_len: int, head_dim: int, description: str = "", + quant_attention_config: "QuantAttentionConfig | None" = None, verbose: bool = True, ) -> Dict[str, Optional[Dict]]: """Benchmark and compare all backends for a given configuration.""" @@ -486,7 +530,13 @@ def benchmark_comparison( results = {} for backend in self.backends: results[backend] = self.benchmark_single( - batch_size, num_heads, seq_len, head_dim, backend, verbose + batch_size, + num_heads, + seq_len, + head_dim, + backend, + quant_attention_config=quant_attention_config, + verbose=verbose, ) # Print comparison @@ -643,19 +693,36 @@ def setup(self): benchmark_iterations=20, ) - @pytest.mark.parametrize("backend", ["VANILLA", "TRTLLM", "FA4"]) - def test_self_attention_perf(self, backend: str): + @pytest.mark.parametrize("head_dim", [64, 128]) + @pytest.mark.parametrize( + ("backend", "quant_attention_config"), + [ + ("VANILLA", None), + ("TRTLLM", None), + ("FA4", None), + ("CUTEDSL", None), + ("CUTEDSL", QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8")), + ], + ) + def test_self_attention_perf( + self, + head_dim: int, + backend: str, + quant_attention_config: "QuantAttentionConfig | None", + ): """Test that attention backend runs without errors.""" - if backend == "FA4" and not _flash_attn4_available: - pytest.fail( - "FlashAttention 4 backend is required for FA4 self-attention perf test" - + (f": {_fa4_import_error}" if _fa4_import_error else "") - ) + _require_attention_backend(backend, head_dim) - batch_size, num_heads, seq_len, head_dim = 1, 24, 1024, 64 + batch_size, num_heads, seq_len = 1, 24, 1024 result = self.benchmark.benchmark_single( - batch_size, num_heads, seq_len, head_dim, backend, verbose=True + batch_size, + num_heads, + seq_len, + head_dim, + backend, + quant_attention_config=quant_attention_config, + verbose=True, ) assert result is not None, f"{backend} benchmark failed to produce results" @@ -814,7 +881,13 @@ def test_fa4_vs_vanilla_cross_attn_wan_shapes( results = {} for backend in ["VANILLA", "FA4"]: results[backend] = self.benchmark.benchmark_cross_attn_single( - batch, num_heads, seq_len_q, seq_len_kv, head_dim, backend, verbose=True + batch, + num_heads, + seq_len_q, + seq_len_kv, + head_dim, + backend, + verbose=True, ) vanilla = results.get("VANILLA") @@ -850,7 +923,13 @@ def test_fa4_cross_attn_quick( """Quick FA4 cross-attention correctness and timing check.""" for backend in ["VANILLA", "FA4"]: result = self.benchmark.benchmark_cross_attn_single( - batch, num_heads, seq_len_q, seq_len_kv, head_dim, backend, verbose=True + batch, + num_heads, + seq_len_q, + seq_len_kv, + head_dim, + backend, + verbose=True, ) assert result is not None, f"{backend} cross-attn failed" assert result["avg_ms"] > 0 diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 1b64933a4f00..cc3328740223 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -73,8 +73,8 @@ def test_legacy_linear_field_rejected(self): class TestAttentionConfigQuantValidation: """Unsupported quantized-attention recipes are rejected with ValueError.""" - def test_quant_config_rejected_on_non_trtllm_backend(self): - with pytest.raises(ValidationError, match="requires backend='TRTLLM'"): + def test_quant_config_rejected_on_unsupported_backend(self): + with pytest.raises(ValidationError, match="requires backend in"): AttentionConfig( backend="VANILLA", quant_attention_config=QuantAttentionConfig(), @@ -84,13 +84,25 @@ def test_quant_config_rejected_when_unsupported(self): with pytest.raises(ValidationError, match="Unsupported quant_attention_config"): AttentionConfig( backend="TRTLLM", - quant_attention_config=QuantAttentionConfig(k_block_size=127), + quant_attention_config=QuantAttentionConfig( + qk_dtype="int8", q_block_size=1, k_block_size=127, v_block_size=1 + ), ) - def test_supported_quant_configs(self): + def test_supported_quant_config_sage(self): attention = AttentionConfig( backend="TRTLLM", - quant_attention_config=QuantAttentionConfig(), + quant_attention_config=QuantAttentionConfig( + qk_dtype="int8", q_block_size=1, k_block_size=16, v_block_size=1 + ), + ) + + assert attention.quant_attention_config is not None + + def test_supported_quant_config_cute(self): + attention = AttentionConfig( + backend="CUTEDSL", + quant_attention_config=QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8"), ) assert attention.quant_attention_config is not None From 6703dc1becf5471de7fefe67e3f9bfe916289bd9 Mon Sep 17 00:00:00 2001 From: RuQing Xu <7891482+xrq-phys@users.noreply.github.com> Date: Fri, 29 May 2026 11:21:24 +0900 Subject: [PATCH 2/3] Drop example code Co-authored-by: Zhenhua Wang <4936589+zhenhuaw-me@users.noreply.github.com> Signed-off-by: RuQing Xu <7891482+xrq-phys@users.noreply.github.com> --- docs/source/models/visual-generation.md | 25 ++----------------------- 1 file changed, 2 insertions(+), 23 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index bcdc3552ddc9..c1260f91de22 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -115,29 +115,8 @@ In addition to linear-layer quantization, VisualGen exposes two **attention-leve - **QK16PV8** (`CUTEDSL` backend): Keeps Q & K in BF16 and quantizes only V to FP8 (E4M3, per-tensor), thus Bmm1 will be carried out in BF16 with Bmm2 in FP8. Targets Blackwell-class GPUs (`sm_100a` / `sm_103a`) with `head_dim = 128`. - **SAGE** (`TRTLLM` backend): Quantizes Q, K, and V with per-block scaling factors. Q/K are stored as INT8 or FP8 (e4m3) and V as FP8 (e4m3); block sizes are tunable per axis (typically `(q, k, v) = (1, 4, 1)` for Wan-1.3B and `(1, 16, 1)` for larger Wan / FLUX checkpoints). Supported recipes are validated at runtime. -CLI usage (Wan T2V, SageAttention): -```bash -python visual_gen_wan_t2v.py \ - --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ - --prompt "A cute cat playing piano" \ - --attention_backend TRTLLM \ - --quant_attention_mode SAGE \ - --output_path output_sage.mp4 -``` - -CLI usage (Wan T2V, QK16PV8): - -```bash -python visual_gen_wan_t2v.py \ - --model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ - --prompt "A cute cat playing piano" \ - --attention_backend CUTEDSL \ - --quant_attention_mode QK16PV8 \ - --output_path output_qk16pv8.mp4 -``` - -Programmatic equivalent (SageAttention): +Python API for SageAttention: ```python from tensorrt_llm import VisualGenArgs @@ -156,7 +135,7 @@ args = VisualGenArgs( ) ``` -Programmatic equivalent (QK16PV8): +Python API for QK16PV8: ```python from tensorrt_llm import VisualGenArgs From 501f8a97427362ede382fd89ba620250f29fb71d Mon Sep 17 00:00:00 2001 From: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> Date: Fri, 29 May 2026 11:27:47 +0900 Subject: [PATCH 3/3] Drop example-side changes. Will prioritize API side usage. Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> --- examples/visual_gen/README.md | 21 ++++----------- examples/visual_gen/visual_gen_flux.py | 22 +++++----------- examples/visual_gen/visual_gen_wan_i2v.py | 31 +++++++++++------------ examples/visual_gen/visual_gen_wan_t2v.py | 31 +++++++++++------------ 4 files changed, 41 insertions(+), 64 deletions(-) diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index 8a784e46c89c..b9a27d3e5fbf 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -94,18 +94,7 @@ python visual_gen_wan_t2v.py \ --prompt "A cute cat playing piano" \ --height 480 --width 832 --num_frames 33 \ --attention_backend TRTLLM \ - --quant_attention_mode SAGE \ - --output_path output.mp4 -``` - -**With CuTe DSL Qk16Pv8:** -```bash -python visual_gen_wan_t2v.py \ - --model_path ${MODEL_ROOT}/Wan2.1-T2V-1.3B-Diffusers \ - --prompt "A cute cat playing piano" \ - --height 480 --width 832 --num_frames 33 \ - --attention_backend CUTEDSL \ - --quant_attention_mode QK16PV8 \ + --enable_sage_attention \ --output_path output.mp4 ``` @@ -125,7 +114,7 @@ WAN supports two parallelism modes that can be combined: - **CFG Parallelism**: Split positive/negative prompts across GPUs - **Sequence Parallelism**: - *Ulysses*: Split sequence along head dimension across GPUs; requires `ulysses_size` to divide the model's head count - - *Attention2D*: 2D mesh sequence parallelism; no head-count constraint; requires `--attention_backend FA4` or `--attention_backend CUTEDSL` + - *Attention2D*: 2D mesh sequence parallelism; no head-count constraint; requires `--attention_backend FA4` - Combining Ulysses and Attention2D is not yet supported @@ -295,8 +284,8 @@ python visual_gen_ltx2.py \ | `--image_cond_strength` | — | ✓ | 1.0 | Image conditioning strength | | `--enable_teacache` | ✓ | ✓ | — | False | Cache optimization | | `--teacache_thresh` | ✓ | ✓ | — | 0.2 | TeaCache similarity threshold | -| `--attention_backend` | ✓ | ✓ | — | VANILLA | `VANILLA`, `TRTLLM`, `FA4`, or `CUTEDSL` | -| `--quant_attention_mode` | ✓ | ✓ | — | NO_QUANT | `NO_QUANT`, `QK16PV8` (requires `CUTEDSL` backend), or `SAGE` (requires `TRTLLM` backend) | +| `--attention_backend` | ✓ | ✓ | — | VANILLA | `VANILLA`, `TRTLLM`, or `FA4` | +| `--enable_sage_attention` | ✓ | ✓ | — | False | SageAttention (requires `TRTLLM` attention backend) | | `--cfg_size` | — | ✓ | — | 1 | CFG parallelism | | `--ulysses_size` | ✓ | ✓ | — | 1 | Ulysses parallelism | | `--parallel_vae_size` | - | ✓ | — | 1 | Parallelism used for VAE | @@ -333,7 +322,7 @@ python visual_gen_ltx2.py \ - Sequence length must be divisible by `ulysses_size` **Attention2D Errors:** -- Requires `--attention_backend FA4` or `--attention_backend CUTEDSL` +- Requires `--attention_backend FA4` - Combining with `--ulysses_size` is not yet supported - Total GPUs = `cfg_size × attn2d_row_size × attn2d_col_size` - Sequence length must be divisible by `attn2d_row_size × attn2d_col_size` diff --git a/examples/visual_gen/visual_gen_flux.py b/examples/visual_gen/visual_gen_flux.py index e9f860151669..00ccbb0dd8bb 100755 --- a/examples/visual_gen/visual_gen_flux.py +++ b/examples/visual_gen/visual_gen_flux.py @@ -209,18 +209,15 @@ def parse_args(): "--attention_backend", type=str, default="VANILLA", - choices=["VANILLA", "TRTLLM", "FA4", "CUTEDSL"], + choices=["VANILLA", "TRTLLM", "FA4"], help="Attention backend (VANILLA: PyTorch SDPA, TRTLLM: optimized kernels, " - "FA4: Flash Attention 4, CUTEDSL: CuTe DSL kernels). " + "FA4: Flash Attention 4). " "Note: TRTLLM falls back to VANILLA for cross-attention.", ) parser.add_argument( - "--quant_attention_mode", - default="NO_QUANT", - choices=["NO_QUANT", "QK16PV8", "SAGE"], - help="Quantized attention presets: NO_QUANT: no quantization; " - "QK16PV8: quantize P@V(Bmm2) only (requires backend=CUTEDSL); " - "SAGE: Sage Attention algorithm (requires backend=TRTLLM).", + "--enable_sage_attention", + action="store_true", + help="Enable SageAttention (per-block quantized Q/K/V). Requires TRTLLM backend.", ) # Parallelism @@ -345,7 +342,7 @@ def build_visual_gen_args(args) -> VisualGenArgs: cache_kwargs = {} attention_cfg: dict = {"backend": args.attention_backend} - if args.quant_attention_mode == "SAGE": + if args.enable_sage_attention: attention_cfg["quant_attention_config"] = { "qk_dtype": "int8", "q_block_size": 1, @@ -353,12 +350,6 @@ def build_visual_gen_args(args) -> VisualGenArgs: "v_block_size": 1, } logger.info("SageAttention: INT8 Q/K, blocks (1, 16, 1)") - elif args.quant_attention_mode == "QK16PV8": - attention_cfg["quant_attention_config"] = { - "qk_dtype": "bf16", - "v_dtype": "fp8", - } - logger.info("QK16PV8: BF16 Q/K, FP8 V (per-tensor)") kwargs = dict( revision=args.revision, @@ -459,7 +450,6 @@ def main(): "model_path": args.model_path, "linear_type": args.linear_type, "attention_backend": args.attention_backend, - "quant_attention_mode": args.quant_attention_mode, "height": args.height, "width": args.width, "steps": args.steps, diff --git a/examples/visual_gen/visual_gen_wan_i2v.py b/examples/visual_gen/visual_gen_wan_i2v.py index b9a02b95b8d5..671b02250779 100644 --- a/examples/visual_gen/visual_gen_wan_i2v.py +++ b/examples/visual_gen/visual_gen_wan_i2v.py @@ -191,18 +191,21 @@ def parse_args(): "--attention_backend", type=str, default="VANILLA", - choices=["VANILLA", "TRTLLM", "FA4", "CUTEDSL"], + choices=["VANILLA", "TRTLLM", "FA4"], help="Attention backend (VANILLA: PyTorch SDPA, TRTLLM: optimized kernels, " - "FA4: Flash Attention 4, CUTEDSL: CuTe DSL kernels). " + "FA4: Flash Attention 4). " "Note: TRTLLM falls back to VANILLA for cross-attention.", ) + + # SageAttention (requires --attention_backend TRTLLM) parser.add_argument( - "--quant_attention_mode", - default="NO_QUANT", - choices=["NO_QUANT", "QK16PV8", "SAGE"], - help="Quantized attention presets: NO_QUANT: no quantization; " - "QK16PV8: quantize P@V(Bmm2) only (requires backend=CUTEDSL); " - "SAGE: Sage Attention algorithm (requires backend=TRTLLM). ", + "--enable_sage_attention", + action="store_true", + help=( + "Enable SageAttention (per-block quantized Q/K/V). Requires TRTLLM backend. " + "Block layout is chosen from --model_path: (1, 4, 1) for Wan2.1, " + "(1, 16, 1) otherwise." + ), ) # Parallelism @@ -328,8 +331,10 @@ def _wan_needs_fine_grained_sage(model_path: str) -> bool: def main(): args = parse_args() - attention_cfg = {"backend": args.attention_backend} - if args.quant_attention_mode == "SAGE": + attention_cfg = { + "backend": args.attention_backend, + } + if args.enable_sage_attention: k_block_size = 4 if _wan_needs_fine_grained_sage(args.model_path) else 16 attention_cfg["quant_attention_config"] = { "qk_dtype": "int8", @@ -338,12 +343,6 @@ def main(): "v_block_size": 1, } logger.info(f"SageAttention: INT8 Q/K, blocks (1, {k_block_size}, 1)") - elif args.quant_attention_mode == "QK16PV8": - attention_cfg["quant_attention_config"] = { - "qk_dtype": "bf16", - "v_dtype": "fp8", - } - logger.info("QK16PV8: BF16 Q/K, FP8 V (per-tensor)") if args.enable_cache_dit: cache_kwargs = {"cache_config": _cache_dit_config_from_args(args)} diff --git a/examples/visual_gen/visual_gen_wan_t2v.py b/examples/visual_gen/visual_gen_wan_t2v.py index f60226ecb906..ce4a32e766d6 100755 --- a/examples/visual_gen/visual_gen_wan_t2v.py +++ b/examples/visual_gen/visual_gen_wan_t2v.py @@ -185,18 +185,21 @@ def parse_args(): "--attention_backend", type=str, default="VANILLA", - choices=["VANILLA", "TRTLLM", "FA4", "CUTEDSL"], + choices=["VANILLA", "TRTLLM", "FA4"], help="Attention backend (VANILLA: PyTorch SDPA, TRTLLM: optimized kernels, " - "FA4: Flash Attention 4, CUTEDSL: CuTe DSL kernels). " + "FA4: Flash Attention 4). " "Note: TRTLLM falls back to VANILLA for cross-attention.", ) + + # SageAttention (requires --attention_backend TRTLLM) parser.add_argument( - "--quant_attention_mode", - default="NO_QUANT", - choices=["NO_QUANT", "QK16PV8", "SAGE"], - help="Quantized attention presets: NO_QUANT: no quantization; " - "QK16PV8: quantize P@V(Bmm2) only (requires backend=CUTEDSL); " - "SAGE: Sage Attention algorithm (requires backend=TRTLLM). ", + "--enable_sage_attention", + action="store_true", + help=( + "Enable SageAttention (per-block quantized Q/K/V). Requires TRTLLM backend. " + "Block layout is chosen from --model_path: (1, 4, 1) for Wan2.x 1.3B, " + "(1, 16, 1) otherwise." + ), ) # Parallelism @@ -354,8 +357,10 @@ def main(): else: parallel_str = "None" - attention_cfg = {"backend": args.attention_backend} - if args.quant_attention_mode == "SAGE": + attention_cfg = { + "backend": args.attention_backend, + } + if args.enable_sage_attention: k_block_size = 4 if _wan_needs_fine_grained_sage(args.model_path) else 16 attention_cfg["quant_attention_config"] = { "qk_dtype": "int8", @@ -364,12 +369,6 @@ def main(): "v_block_size": 1, } logger.info(f"SageAttention: INT8 Q/K, blocks (1, {k_block_size}, 1)") - elif args.quant_attention_mode == "QK16PV8": - attention_cfg["quant_attention_config"] = { - "qk_dtype": "bf16", - "v_dtype": "fp8", - } - logger.info("QK16PV8: BF16 Q/K, FP8 V (per-tensor)") if args.enable_cache_dit: cache_kwargs = {"cache_config": _cache_dit_config_from_args(args)}