[Issue]: torch.topk is slower and more unstable in TheRock

[rasmith]

Problem Description

I've been encountering that torch.topk is generally slower and more unstable in TheRock 7.14. Here are some benchmarks for some relevant vLLM shapes and dtypes (note it does seem to win consistently on the SparseMLA types). By unstable, I mean it returns different, but valid, results for the top-k values:

Top-k Performance Comparison: ROCm 7.2.3 vs TheRock 7.14

Times in microseconds (us). Speedup = ROCm7.2.3 / TheRock7.14 (>1.0x means TheRock is faster).

Category	Shape	Dtype	k	MI250 ROCm7.2.3	MI250 TheRock7.14	MI250 Speedup	MI300 ROCm7.2.3	MI300 TheRock7.14	MI300 Speedup	MI355 ROCm7.2	MI355 TheRock7.14	MI355 Speedup
MoE routing (DeepSeek-V3)	(1, 256)	fp32	8	31.6	41.7	0.76x	23.9	20.7	1.15x	18.4	19.9	0.92x
MoE routing (DeepSeek-V3)	(32, 256)	fp32	8	32.0	43.7	0.73x	23.1	23.9	0.97x	21.3	22.1	0.96x
MoE routing (DeepSeek-V3)	(256, 256)	fp32	8	39.3	49.4	0.80x	24.1	26.7	0.90x	24.4	25.7	0.95x
MoE routing (DeepSeek-V3)	(2048, 256)	fp32	8	123.3	136.6	0.90x	44.5	47.8	0.93x	45.4	47.7	0.95x
MoE routing (DeepSeek-V2)	(1, 160)	fp32	6	31.0	39.8	0.78x	21.7	18.6	1.17x	16.9	18.5	0.91x
MoE routing (DeepSeek-V2)	(2048, 160)	fp32	6	100.0	113.0	0.88x	40.8	40.8	1.00x	36.4	38.2	0.95x
MoE routing (cohere2)	(1, 8)	fp32	1	18.3	20.5	0.89x	16.3	11.1	1.47x	8.5	8.0	1.06x
MoE routing (cohere2)	(2048, 8)	fp32	1	34.7	44.3	0.78x	18.8	19.9	0.94x	15.7	16.4	0.96x
MoE routing (moondream)	(1, 16)	bf16	2	28.5	38.0	0.75x	20.6	16.6	1.24x	15.5	16.8	0.92x
MoE routing (moondream)	(2048, 16)	bf16	2	43.6	60.1	0.73x	26.9	28.8	0.93x	21.6	23.3	0.93x
MoE grouped (group scores)	(1, 8)	fp32	4	26.0	36.7	0.71x	21.1	16.8	1.26x	15.5	16.8	0.92x
MoE grouped (group scores)	(2048, 8)	fp32	4	46.3	60.1	0.77x	26.4	27.9	0.95x	22.1	22.9	0.97x
Gemma4 fallback	(1, 128)	bf16	8	28.3	37.5	0.75x	22.0	18.6	1.18x	17.1	18.6	0.92x
Gemma4 fallback	(16, 128)	bf16	8	27.9	36.3	0.77x	22.5	20.8	1.08x	18.9	20.2	0.94x
Gemma4 fallback	(2048, 128)	bf16	8	68.3	79.4	0.86x	27.8	30.0	0.93x	27.2	29.1	0.93x
Gemma4 fallback	(2048, 128)	fp16	8	62.6	73.2	0.86x	26.6	28.6	0.93x	26.0	27.4	0.95x
Gemma4 fallback	(250000, 128)	bf16	8	5275.5	5765.0	0.92x	1371.0	1492.3	0.92x	1433.1	1537.8	0.93x
Top-k sampling	(1, 32000)	fp32	50	73.2	67.0	1.09x	56.8	56.2	1.01x	49.9	50.7	0.98x
Top-k sampling	(1, 128256)	fp32	50	104.0	102.3	1.02x	77.5	75.5	1.03x	71.9	69.1	1.04x
Top-k sampling	(32, 128256)	fp32	50	221.2	253.6	0.87x	151.9	150.9	1.01x	134.9	134.2	1.01x
Top-k sampling	(256, 128256)	fp32	50	919.9	945.9	0.97x	299.7	298.1	1.01x	323.0	322.8	1.00x
Logprobs	(1, 128256)	fp32	5	92.0	102.9	0.89x	77.2	74.7	1.03x	71.3	72.1	0.99x
Logprobs	(1, 128256)	fp32	20	90.3	100.8	0.90x	69.7	76.2	0.91x	72.2	71.9	1.00x
Logprobs	(32, 128256)	fp32	20	221.0	251.8	0.88x	143.9	145.3	0.99x	129.5	129.7	1.00x
Sparse MLA attn	(1, 4096)	fp32	64	73.2	74.7	0.98x	46.6	42.5	1.10x	44.3	40.3	1.10x
Sparse MLA attn	(1, 16384)	fp32	64	65.3	50.1	1.30x	43.3	40.2	1.08x	40.4	39.8	1.02x
Sparse MLA attn	(32, 4096)	fp32	64	80.7	74.9	1.08x	53.3	47.6	1.12x	50.5	44.4	1.14x
Sparse MLA attn	(32, 16384)	fp32	64	121.0	134.8	0.90x	87.1	85.5	1.02x	84.2	83.2	1.01x

Operating System

Ubuntu 24.04

CPU

Various

GPU

MI250/MI300/MI355

ROCm Version

TheRock7.14

ROCm Component

No response

Steps to Reproduce

I have a small script to do the benchmarks.

test_topk_perf_vllm.py

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

rocminfo --support output

Paste output here

Additional Information

No response