[Bug]: FP8 block-scale Linear hits a device-side assert on SM107: DeepGEMM JIT compiles for sm_107f and its CUTLASS disables the TMEM paths #20508

System Info

• TensorRT-LLM main @ e88d19c, PyTorch backend, CUDA 13.2 toolchain, torch 2.14
• DeepGEMM at the pinned b64107f2b9 (3rdparty/fetch_content.json)
• GPU with compute capability 10.7 (SM107)

Who can help?

Upstream root cause: deepseek-ai/DeepGEMM#461 (DeepGEMM's vendored CUTLASS 4.2.1, which is also what TensorRT-LLM's pinned b64107f2b9 ships, has no SM107 enablement; that issue asks for a CUTLASS 4.8.0 bump, and reports the same assert from vLLM). This issue is about the TensorRT-LLM side: the default FP8 block-scale Linear path on SM107 hits it.

Related: #18758 maps SM107 to the sm_100f NVRTC target for the trtllm-gen FMHA JIT for the same kind of reason. #19067 covers the build side of SM107 (107-real folds to 100f).

Reproduction

An FP8 block-scale Linear with default settings (DeepGEMM enabled) on SM107:

import torch
from tensorrt_llm._torch.modules.linear import Linear
from tensorrt_llm.models.modeling_utils import QuantConfig
from tensorrt_llm.quantization import QuantAlgo

m, k, n = 12, 512, 1024
lin = Linear(k, n, bias=False, dtype=torch.bfloat16,
             quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES)).cuda()
with torch.no_grad():
    lin.weight.copy_(torch.randint(-32, 33, lin.weight.shape, device="cuda").to(lin.weight.dtype))
    lin.weight_scale.fill_(1e-3)
lin.post_load_weights()
out = lin(torch.randn(m, k, device="cuda", dtype=torch.bfloat16))
torch.cuda.synchronize()

Linear.forward → FP8BlockScalesLinearMethod.apply → torch.ops.trtllm.fp8_swap_ab_gemm → deep_gemm.fp8_gemm_nt. The JIT-compiled sm100_fp8_fp4_gemm_1d1d_impl kernel hits a device-side assert:

.../deep_gemm/include/cute/arch/tmem_allocator_sm100.hpp:143: void cute::TMEM::Allocator2Sm::allocate(int, unsigned int *): block: [13,0,0], thread: [80,0,0] Assertion `0 && "Attempting to use TMEM allocation PTX without CUTE_ARCH_TCGEN05_TMEM_ENABLED"` failed.
...
torch.AcceleratorError: CUDA error: device-side assert triggered

The same happens for every shape I tried: (m, k, n) = (12, 512, 1024), (1, 7168, 2048) and (128, 7168, 2048). In the unit tests it shows up as test_kimi_kda_fused_verify_parity.py::test_fused_vs_sequential_two_rounds[full-rank-fp8]. That test aborts the process, so the rest of the file then fails with it.

Expected behavior

The DeepGEMM FP8 GEMM runs on SM107, as it does on SM100/SM103. Alternatively, Linear does not route SM107 to DeepGEMM.

Actual behavior

With DG_JIT_DEBUG=1, DeepGEMM compiles the kernel with --gpu-architecture=sm_107f. By default, deep_jit::cuda::Device::get_arch() returns <major><minor>f for every arch with major ≥ 10.

The CUTLASS bundled with DeepGEMM enables the tcgen05/TMEM CuTe paths only for the 100f, 101f and 103f families (CUDA_ARCH_FAMILY(1000/1010/1030) in cutlass/arch/config.h, and then CUTE_ARCH_TCGEN05_TMEM_ENABLED in cute/arch/config.hpp). Under sm_107f, __CUDA_ARCH_FAMILY_SPECIFIC__ is 1070, so CUTE_ARCH_TCGEN05_TMEM_ENABLED is not defined. The TMEM allocator then compiles to its assert(0) fallback. The CUTLASS used by TensorRT-LLM itself has no 1070 family either.

What I checked:

• Rewriting --gpu-architecture=sm_107f to sm_100f in the JIT compile, with an nvcc wrapper under a custom CUDA_HOME, makes all three shapes pass on SM107. The max relative error against an fp32 reference matches the non-DeepGEMM path for each shape (for example 0.0413 vs 0.0413 at 12×512×1024). TensorRT-LLM's own .so files also run sm_100f code on SM107 (see [Bug]: -a 107-real is silently folded to 100, no SM107 code is generated #19067).
• Linear(..., disable_deep_gemm=True) avoids the problem on SM107 for all three shapes.

Additional notes

Suggested fixes, in order of preference:

1. Have the DeepGEMM JIT target sm_100f on SM107 until the pinned DeepGEMM carries a CUTLASS with SM107 support (Update CUTLASS to 4.8.0 for SM107 (Rubin) support deepseek-ai/DeepGEMM#461). For example, extend the existing 3rdparty/patches/deepgemm_*.patch mechanism so that get_arch() maps 107 → 100f. This matches what [None][fix] Map SM107 to the sm_100f NVRTC target in the trtllm-gen FMHA kernel loader #18758 does for the trtllm-gen FMHA JIT, and it is the change verified above.
2. Until then, route SM107 away from DeepGEMM for FP8 block-scale Linear, as disable_deep_gemm does today. In FP8BlockScalesLinearMethod.apply / transform_weights, is_sm_100f() alone currently sends SM107 to fp8_swap_ab_gemm.

I only checked fp8_gemm_nt. The other SM100 DeepGEMM kernels compiled through the same JIT (BF16 GEMM, MQA logits, and so on) use the same CuTe TMEM allocator and are probably affected the same way, but I have not run them.