perf(diffusion): fused qknorm rope Flux 2.0

[xuanzic]

Fused DiT QK-Norm + RoPE kernel for FLUX.2

Replaces FLUX.2's per-block chain of 4 per-head RMSNorms + 4 RoPEs (double-stream) and
2+2 (single-stream) with one fused TVM-FFI kernel call per tensor. Per-warp algorithm
adapted from TensorRT-LLM PR-13052
(fusedDiTQKNormRopeKernel.cu), simplified to a per-tensor launcher that
slots into our existing separate-Q/K/V graph without restructuring. Templated on
dtype — both fp16 (FLUX.2's default) and bf16 are exported.

Numbers — FLUX.2-dev 1024×1024, 28 steps, GB300

Build	per step	28 steps	engine plan
bf16 legacy	198.91 ms	5.57 s	75.85 GB
bf16 fused	182.64 ms (−8.2%)	5.11 s	60.04 GB (−21%)
fp16 legacy	193.73 ms	5.42 s	75.85 GB
fp16 fused	183.83 ms (−5.1%)	5.15 s	60.04 GB (−21%)
torch.compile bf16 (reference)	185.00 ms	5.18 s	—

Both fused paths beat torch.compile. Before this PR, TRT BF16 was 1.08× slower;
after, 1.01× faster. Engine plan shrinks 21% because the unfused chain materialized
intermediate weight tensors the fused path collapses.

Kernel correctness (vs PyTorch reference, realistic RoPE cos/sin):
bf16 max_abs 1.56e-2 (limit 3e-2) ✓ · fp16 max_abs 1.95e-3 (limit 5e-3) ✓

Kernel microbench (vs PyTorch eager bf16): 15.8× – 25.7× across FLUX.2 shapes.

What each commit does

• f076f00 build — pip-installs apache-tvm-ffi==0.1.12 in the Dockerfile; CMake
  uses importlib.util.find_spec() to discover tvm_ffi + tensorrt_libs paths in
  the venv. Previously these libs were silently undiscoverable in the stock dev image.
• 2049c14 plugin bf16 — adds dtype=2 → kBF16 in five sites of
  tvm_ffi_kernel_plugin.{cpp,h} (output spec, format combination, DLTensor fill).
• 1193351 Mistral bf16 (1/3) — work_np_dtype for bf16 now uses
  ml_dtypes.bfloat16 (was np.float16, causing strongly-typed TRT type mismatches).
• 71dc58f fused kernel — new dit_rms_norm_rope.cu (bf16-only at landing time)
  • flux2_dit_builder.py injection at 6 paired call sites behind
    TRTMC_FUSED_DIT_KERNELS=1 env flag. Default path byte-identical to before. Build
    gated by option(TRTMC_BUILD_DIFFUSION_KERNELS OFF). sm_90 + sm_100 + sm_103.
• fp16 + remaining bf16 bundle fixes — three independent fixes that
  together complete the picture; bundled because the bf16 fixes only become testable
  end-to-end once you can also build a fused engine in fp16 (FLUX.2's default):
  • Kernel templated on dtype via DTypeTraits<T> (specializations for
    __nv_bfloat16 and __half); two TVM-FFI exports
    trtmc.dit_rms_norm_rope_{bf16,fp16}; builder picks the suffix by _CAST_DTYPE.
  • graph_ops.add_constant: detect ml_dtypes.bfloat16 (a V16 numpy extension dtype
    that TRT bindings can't auto-convert) and route to the explicit
    trt.Weights(trt.bfloat16, ptr, count) constructor. That constructor does not
    copy, so we also keep the numpy buffer alive via a module-level
    _bf16_weight_keepalive dict keyed by id(network). Without the keepalive, GC
    of the local array caused exit-139 segfaults during TRT graph construction.
  • mistral_encoder_builder.py: np.finfo rejects ml_dtypes.bfloat16 (not inexact); switch to ml_dtypes.finfo for the attention-mask min value.

How to verify

./scripts/docker_build_gb300.sh   # rebuilds image with apache-tvm-ffi

docker run --rm --gpus all -v "$PWD":/workspace/tensorrt-model-connect \
  -w /workspace/tensorrt-model-connect trtmc-dev-gb300:latest \
  bash -c "cmake -S . -B build -DTRTMC_BUILD_DIFFUSION_KERNELS=ON \
                            -DTRTMC_ENABLE_TVM_FFI=ON -GNinja && \
           cmake --build build --target trtmc_dit_rms_norm_rope \
                                       trtmc_tvm_ffi_plugin -j"

# Build (set TRTMC_FUSED_DIT_KERNELS=1 to use fused; --precision fp16 or bf16)
# Bench: same docker but with LD_PRELOAD=…trtmc_tvm_ffi_plugin.so:…trtmc_dit_rms_norm_rope.so
# and `python3 scripts/bench_flux2_perf.py --bundle … --backends trt_denoiser --num-steps 28`

Expected fused step time: fp16 0.184 s/step, bf16 0.183 s/step on GB300.

Not in this PR — follow-ups

• Port fuse_scale_shift (AdaLN, 81 calls/step) — biggest remaining win (~25 ms/step).
• Port group_norm_silu (VAE) — runs once per gen, small.
• C++-runtime registration of the kernel — currently registers via Python
  ctypes.CDLL; the C++ trtmc binary needs a static-init wrapper to use the fused path.
• Permanent tests/test_dit_rms_norm_rope.py (mirror test_flashinfer_plugin_e2e.py).

Risk / backwards compat

• Default path byte-identical to pre-PR. TRTMC_BUILD_DIFFUSION_KERNELS defaults
  OFF; without TRTMC_FUSED_DIT_KERNELS=1, the legacy chain runs.
• Plugin bf16 mapping is additive (fp32/fp16 paths untouched).
• The _bf16_weight_keepalive dict lives for process lifetime, bounded by engine
  builds. Acceptable; long-lived processes that build many bf16 engines could
  explicitly clear entries on completion.

Architecture note

TvmFfiKernelPlugin looks up kernels by global name at engine deserialize time. Any
toolchain that emits a tvm_ffi.Function slots in identically — raw CUDA today,
CuTeDSL/Triton/FlashInfer JIT later via the same add_tvm_ffi_kernel(...) Python
helper. This PR is the first end-to-end shake-out of that plug for diffusion graphs.