fix a bug with headDim 256 nvfp4-kv kernels

Signed-off-by: Perkz Zheng <[email protected]>

Author: PerkzZheng

cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h

@@ -685,7 +685,8 @@ struct KernelParams
         // The number of elements in 128B for Q.
         int32_t numEltsIn128BKv = (128 * 8) / get_size_in_bits(kernelMeta.mDataTypeKv);
         // The number of head elts (per token) in each block of shared memory (see above explanation).
-        int32_t numEltsInClampedHeadDimKv = std::min(numEltsIn128BKv, maxHeadDimKv);
+        // HeadDim will be split into multiple headDimStages (128) if maxHeadDimKv > 128.
+        int32_t numEltsInClampedHeadDimKv = std::min({numEltsIn128BKv, maxHeadDimKv, 128});
 
         // Do we have to transform K/V before MMA?
         bool const transformsKv{kernelMeta.mDataTypeKv != kernelMeta.mDataTypeQ};