bugfix generic-k code in top-k with softmax

[t4c1]

Fixes a bug in generic top-k softmax EVT implementation that resulted in wrong results when k != 2 and k != 4.

This allows removal of the static assert requiring k to be either 2 or 4. Comment in example 61 is also fixed to reflect that any k value is now supported.

[alihassanijr]

@t4c1 Thank you for submitting this patch.

Just a note on the assertion, it's there more as a warning to users that the generic sort comes with serious performance implications. Due to the control flow associated with it, it introduces somewhat heavy branching and register spilling, resulting in a smaller improvement over the baseline.

I'll let @hwu36 chime in on whether we want to get rid of the assertion.

[hwu36]

@alihassanijr , maybe leave it there but make the text clearer about the consequences?

[alihassanijr]

@hwu36 leave the assert there with a better message, or remove and make the doc clearer about the consequences?

[hwu36]

Leave it there and make the message better so people willing to try can try it.

[alihassanijr]

Suggested change

	static_assert(is_same_v<ElementCompute, float>, "Fused Top-K + Softmax reduction requires FP32 accumulation.");
	static_assert(is_same_v<ElementCompute, float>, "Fused Top-K + Softmax reduction requires FP32 accumulation.");
	static_assert(TopK == 2 || TopK == 4,
	"Fused Top-K + Softmax reduction only allows K=2 and K=4, because those cases have been performance-optimized. Other values of K can be enabled by removing this assertion, but they may come with serious performance implications."
	);

[alihassanijr]

Suggested change

	2. Top-K value is static (meaning multiple kernels have to be compiled to support
	different values.)
	2. Top-K value is static (meaning multiple kernels have to be compiled to support
	different values.)
	* NOTE: Only K=2 and K=4 cases are performance-optimized and enabled by default.
	There is also a generic sort that supports all K values greater than 1, but it can lead to serious performance implications to the underlying kernel.
	If necessary, users can simply remove the K==2 || K ==4 assertion under cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp, and the generic sort will automatically be used for all other Ks.