[None][feat] Add GLM-5.3-Flash (glm5_next) support

cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.h

@@ -31,7 +31,8 @@ constexpr int kCompactHeadsWorkThreshold = 144;
 constexpr bool isSupportedHeadCount(int numHeads)
 {
     return numHeads == 1 || numHeads == 2 || numHeads == 3 || numHeads == 4 || numHeads == 6 || numHeads == 8
-        || numHeads == 12 || numHeads == 16 || numHeads == 24 || numHeads == 32 || numHeads == 48 || numHeads == 96;
+        || numHeads == 12 || numHeads == 16 || numHeads == 24 || numHeads == 32 || numHeads == 48 || numHeads == 64
+        || numHeads == 96;
 }
 
 //! Select the compact-head kernel within the measured KDA decode work threshold.

cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu

@@ -1560,6 +1560,7 @@ void dispatch_kda_decode_heads(KdaDecodeLaunchParams const& p)
     case 24: dispatch_kda_decode_layout<kCompact, 24>(p); break;
     case 32: dispatch_kda_decode_layout<kCompact, 32>(p); break;
     case 48: dispatch_kda_decode_layout<kCompact, 48>(p); break;
+    case 64: dispatch_kda_decode_layout<kCompact, 64>(p); break;
     case 96: dispatch_kda_decode_layout<kCompact, 96>(p); break;
     default:
         if constexpr (kCompact)

cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeOptimized.cu

@@ -740,6 +740,7 @@ void dispatchKdaDecodeOptimizedHeads(KdaDecodeParams const& params, cudaStream_t
     case 24: TLLM_CUDA_CHECK((launchKernelSchedule<kSchedule, 24, kUpdateConvCache>(params, stream))); break;
     case 32: TLLM_CUDA_CHECK((launchKernelSchedule<kSchedule, 32, kUpdateConvCache>(params, stream))); break;
     case 48: TLLM_CUDA_CHECK((launchKernelSchedule<kSchedule, 48, kUpdateConvCache>(params, stream))); break;
+    case 64: TLLM_CUDA_CHECK((launchKernelSchedule<kSchedule, 64, kUpdateConvCache>(params, stream))); break;
     case 96: TLLM_CUDA_CHECK((launchKernelSchedule<kSchedule, 96, kUpdateConvCache>(params, stream))); break;
     default: TLLM_CHECK_WITH_INFO(false, "Optimized KDA decode does not support numHeads=%d", params.numHeads);
     }

cpp/tensorrt_llm/thop/kdaDecodeOp.cpp

@@ -97,9 +97,9 @@ void validate_kda_decode_fusion_inputs(at::Tensor x_q, at::Tensor x_k, at::Tenso
     int const HV = static_cast<int>(x_v.size(2));
     TORCH_CHECK(B > 0, "KDA decode requires a non-empty batch");
     bool const supportedHeads = H == 1 || H == 2 || H == 3 || H == 4 || H == 6 || H == 8 || H == 12 || H == 16
-        || H == 24 || H == 32 || H == 48 || H == 96;
+        || H == 24 || H == 32 || H == 48 || H == 64 || H == 96;
     TORCH_CHECK(
-        H == HV && supportedHeads, "KDA decode fusion CUDA supports H == HV in {1,2,3,4,6,8,12,16,24,32,48,96}");
+        H == HV && supportedHeads, "KDA decode fusion CUDA supports H == HV in {1,2,3,4,6,8,12,16,24,32,48,64,96}");
     TORCH_CHECK(x_k.size(1) == B && x_k.size(2) == H && x_v.size(1) == B,
         "x_q, x_k, and x_v batch/head dimensions are inconsistent");
     TORCH_CHECK(HV % H == 0, "HV must be divisible by H");