[zhonglinnian]

Summary

A TopK whose K is known at compile time no longer reads K back from the device. hip.top_k now lowers to a new wrap_top_k_static(..., int64_t k, ...) when its result extent along axis is static, which skips a device-to-host copy and a stream synchronization per call: a steady ~8 µs per TopK on gfx1151.

Related issue or design

None. The runtime entry point is extracted from hip-ir-runtime's hip-ep patch 0010-gguf-host-implementation-... by Yu Xing (@yuxing12), its Top-K hunks only.

Why

wrap_top_k takes K as a device tensor, so every call copies it to the host and synchronizes the stream before launching the kernel. ONNX TopK's results hold exactly K elements along axis; when that extent is static in the result type, K is already known and the round trip is pure overhead. ONNX models with a constant K, such as argmax-style sampling, hit this on every inference.

What

1. feat(top-k): add a static-K runtime entry point (Yu Xing, from patch 0010)
   • wrap_top_k_static takes K by value; both entry points share one launch helper. Mock stub and declaration added.
   • wrap_top_k now reads K with hipMemcpyAsync(hipMemcpyDefault) on the session stream before synchronizing it, instead of a blocking hipMemcpy on the null stream. The read is ordered after the kernel that produced K, and K may be in host or device memory.
2. perf(top-k): pass a statically known K by value
   • TopKLowering.cpp: when the results' extent along axis is static, call wrap_top_k_static with that extent as an i64; a dynamic extent keeps wrap_top_k. No op or attribute change.
   • An out-of-range axis is now rejected by the lowering instead of indexing past the shape.

Performance

Same branch, one ONNX TopK (K = 1, rows = 1) per graph; "static" has K as a constant initializer (wrap_top_k_static), "runtime" has K as a graph input with a symbolic result extent (wrap_top_k). Median session.run() time over 5 interleaved rounds of ~1 s each.

Row width	static (µs)	runtime (µs)	Saved
64	49.8	58.0	8.2 µs (14.1%)
256	79.9	88.4	8.5 µs (9.6%)
1024	191.4	199.9	8.5 µs (4.3%)
4096	583.1	591.2	8.1 µs (1.4%)

Hardware: gfx1151 (Strix Halo), Linux, TheRock 7.13, ORT 1.31.0, Release build. Round-to-round spread ≤ 1.3 µs. Both paths were confirmed with HIPDNN_EP_DEBUG=1 and timed with it off. The runtime model also sizes its outputs dynamically, so part of its extra time may come from that shape handling rather than the call alone.

Test plan

• check-hip-mlir-lit: 669 tests, 660 passed, 9 unsupported, 0 failed.
  • hip-to-llvm/test_topk.mlir: the static case expects wrap_top_k_static; a dynamic-batch case with a negative axis checks K = 5 is passed by value; a dynamic-extent case keeps wrap_top_k.
  • e2e/test_topk_model.mlir expects wrap_top_k_static.
• test/numeric/tests/test_topk.py on gfx1151: test_topk_axis1[largest=1], test_topk_full_k and the new test_topk_runtime_k (K as a graph input, so wrap_top_k) pass. test_topk_axis1[largest=0] fails, on both paths with identical wrong indices, so it is an existing hip_top_k issue with largest=0, not this change.
• pre-commit run --from-ref origin/main --to-ref HEAD: all hooks pass.

Notes for reviewers

• Commit 1 is Yu Xing's code as extracted from patch 0010; commit 2 is the hip-ep caller, which 0010 does not have (there only external frontends call wrap_top_k_static).
• The largest=0 failure above needs its own fix in the kernel.

Checklist

• The change is focused, or links a design/series explaining its scope.
• Relevant tests were added or updated and the results are documented.
• User-facing or design documentation was updated when needed.
• Substantial AI assistance is disclosed, and I reviewed and understand the result.

AI assistance: extracted, extended, tested and benchmarked with Claude Code; commits carry Co-Authored-By: Claude, and the follow-up Made-with: Claude Code.

🤖 Generated with Claude Code

[github-actions]

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

1. The change is focused. Substantial work links the relevant issue
   or design discussion.
2. The PR documents relevant test results and updates affected
   documentation.
3. If AI tools provided substantial assistance, the description
   explains what was assisted and how it was validated, and commit
   trailers identify the tool. The contributor has reviewed and
   understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

[github-actions]

L2 Accuracy Results (EP vs CPU)

Model	Combined L2	Total Elems	Skipped NaN/Inf
conv_test_hybrid	1.12256E-07	64	0
GroupQueryAttention_seq256	59.2561	2621440	0
MatMulNBits_o_seq128	259.905	368640	0
QMoE_seq128	24.2583	368640	0

Threshold: 0.01 | Run: 7910 - Commit: 9f4a85e

[github-actions]

Binary sizes

Baseline: run 40278658372 - b199536 (dll-sizes)

File	This (MB)	Main (MB)	Delta (MB)
hipgpu.dll	59.31	59.30	+0.01
custom_kernels_gfx1170.dll	5.55	5.54	0.00
custom_kernels_gfx11-generic.dll	5.79	5.78	+0.01

MorphiZen EP Performance Results

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.78	7.48	247	3	1158
GroupQueryAttention_seq128	11389.80	2.5514	1	5	229
matmul_down_seq128	507.68	3.26	119	3	275

EPContext Export Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.83	48.92	249	3	15504

EPContext Import Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.83	10.08	249	3	15639

OGA Benchmark Results

Model	Warmup	Reps	Prompt Len	Gen Tokens	TTFT (ms)	TPS	Peak Mem (GB)	GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32	1	5	128	128	91.6	74.8	1.25	12.24
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	1	5	128	128	116.9	47.6	1.14	5.23

OGA Wheel Smoke (Python benchmark_e2e.py)

Model	TTFT (ms)	TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	118	47.9

VLM Benchmark Results

Model	TTFT (ms)	TPS	Output
gemma-4-E4B-it-fp16-int4-kquant-prune-logits	3945.5	7.9	This is a high-quality, outdoor photograph featuring a beautiful, fluffy white dog, likely a West Highland White Terrier or a si
Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx	14486.3	14.0	The user wants a detailed description of the image. Let me carefully observe all elements before writing. Subject: A white do

Run: 7910 - Commit: 9f4a85e