[lkdvos]

⚠️ vibe-coded, not meant for merging, mostly to get an idea of how much this matters vs the current approach.

---

Replaces the per-block loop in mul! for CUDA tensors by a single grouped gemm call (cublasGemmGroupedBatchedEx, cuBLAS ≥ 12.5) for all small blocks, and individual gemm calls for the few large ones.

• Grouped gemm is real-only, so complex element types use the 1m embedding: α * op(A) is expanded into real 2m × 2k blocks [ar -ai; ai ar] by a small kernel, while B and C are simply reinterpreted as real. Complex β is applied upfront.
• Blocks with m * n * k above TensorKitCUDAExt.GROUPED_GEMM_THRESHOLD[] (1024^3 real, 256^3 complex) use individual cublasGemmEx calls, since tensor-core gemm beats the grouped (SIMT) kernel there. These calls upload α and β once, instead of once per block as in CUDA.jl's mul!.
• Complex products with an adjoint B, and cuBLAS < 12.5, use only the individual gemm calls.
• The block structure and all host-side arguments are cached per pair of spaces in a GroupedGemmPlan.
• The grouped call always uses plain 32F/64F compute types: TF32 crashes on H100 with a misaligned-address error for unaligned blocks.

Benchmarks

C = A * B, with the old mul! as baseline.

• RTX A6000: 24–190× faster for tensors with many small or medium blocks (U(1), SU(2), U(1)×U(1), Float32/ComplexF32).
• H100: 10–140× faster for many small blocks. With few, very large blocks the timings are within a few % of the old implementation (2–3% slower on the largest blocks, up to ~11% for ComplexF32 under TF32 math mode, where the direct cublasGemmEx path is slower than CUDA.jl's gemmEx! for reasons I haven't pinned down).
• Under TF32 math mode, a lower real threshold (~256³) would be better; the current default is tuned for the default math mode.

Tests

The existing "Multiplication and inverse: test via CPU" testset covers the grouped path (the test spaces have small blocks). A new testset covers the direct path and real/complex β for every combination of adjoints.

🤖 Generated with Claude Code

[kshyatt]

It can certainly be on a separate PR, but it looks like rocBLAS now has support for this too.

[kshyatt]

don't we already have this in the cutensormap.jl file?

[lkdvos]

For context, this isn't really something that should actually be merged, and I would much rather have this all go through MAK instead. @jack-dunham was playing around with our kernels compared to the cuTENSORBlockSparse stuff, and I just wanted to add something to see if this actually matters, so see this more as a prototype.

It's also very unreviewed 🙃

[kshyatt]

Could some of these move into TensorKit src so we can extend them for ROCm too?