[Bug]: Fused DSpark RMSNorm/RoPE is never used on SM107, although its error messages list SM107 as supported #20365

System Info

• GPU: SM107, 8 GPUs over 2 nodes
• TensorRT-LLM main @ 15f060c, PyTorch backend
• Model: Kimi K3 with the MLA-backboned standalone DSpark drafter ([TRTLLM-16564][feat] MLA-backboned standalone DSpark drafter (Inferact/Kimi-K3-DSpark) #19040)

Who can help?

@dc3671 (author of #19040)

Information / Reproduction

is_fused_dspark_rmsnorm_rope_supported() in tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py gates on _get_dspark_arch_str(), whose table only has SM100 and SM103:

_DSV4_DSPARK_ARCH_BY_SM = {
    100: "sm_100",
    103: "sm_103",
}

so on SM107 it always returns False, and _rmsnorm_rope_batched() / _mla_rope_into() in modeling_dspark.py silently fall back to the eager fp32 RMSNorm/RoPE chain (three times per drafter layer and step). At the same time cute_dsl_dspark_rmsnorm_rope and cute_dsl_dspark_rope_into reject with "requires ... an SM100, SM103, or SM107 GPU", and cute_dsl_mla_decode_unavailability_reason() accepts SM107 for the same drafter, so SM107 looks intended.

The table is shared with the DSV4 DSpark attention kernel (dspark_attention_custom_op.py), which is a tcgen05 kernel and is reported as "SM100 or SM103" there. The RMSNorm/RoPE kernels (cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py) are plain SIMT CuTe kernels (warp reductions, no TMA/UMMA), so they do not need that restriction.

Expected behavior

The fused RMSNorm/RoPE path is used on SM107.

Actual behavior

The eager fallback is used on SM107.

Additional notes

With SM107 admitted (separately from the attention gate) on our SM107 system:

• outputs are bit-identical to the eager chain (max |diff| = 0) for the three drafter call shapes (context KV [1, N, 576] with norm_dim 512, query RoPE into the absorbed query [1, 8, H, 192] -> [..., 576], block latent [1, 8, 576]), H = 8 / 16, N = 1 / 3 / 8;
• one decode step's worth of these calls in a CUDA graph (5 layers x 3 calls) takes ~4.4x less time;
• end-to-end decode step of Kimi K3 + DSpark: -0.3 ms (~2%) at 2K / 32K / 128K context.

I will open a PR that gives the RMSNorm/RoPE ops their own architecture set (100, 103, 107) and leaves the attention gate unchanged.