[KMCompiler]Optimize index_copy and index_copy_ generic kernels

[Onisen7]

PR Category

Operator

Type of Change

Performance Optimization

Description

This PR optimizes index_copy and index_copy_ across the generic path and
the Ascend backend.

For the generic implementation in src/flag_gems/ops/index_copy_.py, this PR:

• selects the Triton block size according to the workload size;
• uses a lightweight Triton copy kernel for contiguous out-of-place inputs;
• keeps a redispatch fallback for non-contiguous inputs;
• masks invalid stores and provides debug-time device assertions;
• normalizes negative dimensions before launch;
• skips the index-copy kernel launch for empty source tensors; and
• reuses generated kernels without specializing the Python wrapper cache for
  every complete input shape.

For the Ascend backend, this PR adds an Ascend-specific implementation in:

• src/flag_gems/runtime/backend/_ascend/ops/index_copy_.py
• src/flag_gems/runtime/backend/_ascend/ops/__init__.py

The Ascend implementation:

• uses a flat Triton kernel for workloads with inner_size <= 4;
• uses a row-oriented 2D Triton kernel for wider suffixes;
• uses 64-bit offset arithmetic for large tensors;
• keeps runtime dimensions out of Triton specialization;
• selects the flat-kernel block size according to the workload size;
• skips kernel launches for empty index or inner dimensions;
• launches kernels under the input device context;
• masks invalid indices to prevent out-of-bounds stores; and
• supports both out-of-place index_copy and in-place index_copy_.

Issue

N/A

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue. (N/A: no linked issue)
• Change is fully covered by a UT.

Testing

Functional test command:

python -m pytest -s tests/test_index_copy_.py

Result: 24 passed on NVIDIA H20, Hygon, PPU, MetaX MC550, Iluvatar, and
Ascend 910B.

Ascend test environment:

• Device: Ascend 910B
• CANN: 8.5.0
• Python: 3.11.13
• flagtree: 0.6.0+ascend3.2

Performance

Generic benchmark command:

unset TRITON_DEBUG
python -m pytest -s benchmark/test_index_copy.py

Speedup = PyTorch latency / FlagGems latency.

Platform	index_copy before	index_copy after	index_copy_ before	index_copy_ after	Combined before	Combined after
NVIDIA H20	0.243x	0.803x	0.227x	1.576x	0.235x	1.190x
Hygon	0.257x	0.581x	0.174x	1.136x	0.215x	0.859x
PPU	0.284x	1.331x	0.242x	1.771x	0.263x	1.551x
MetaX MC550	0.444x	1.122x	0.432x	1.578x	0.438x	1.350x
Iluvatar	0.209x	1.053x	0.195x	1.560x	0.202x	1.307x

Ascend benchmark command:

python -m pytest -s benchmark/test_index_copy.py

The default benchmark ran in kernel mode at the comprehensive level.

The generic implementation completed the functional test with 24 passed,
but did not complete the default Ascend benchmark. The new Ascend-specific
implementation completed all 48 benchmark rows: 2 APIs, 3 dtypes, and 8 input
configurations.

Operator	float16 mean	float32 mean	bfloat16 mean	Overall mean
index_copy	1.045x	0.872x	1.051x	0.989x
index_copy_	1.430x	1.345x	1.428x	1.401x
Combined	-	-	-	1.195x

The Ascend values are arithmetic means of the configured benchmark rows.
Performance varies by input shape, so the reported means do not imply that
every individual configuration is faster than the native implementation. Full
per-shape benchmark logs are available if additional review is needed.

[CLAassistant]

All committers have signed the CLA.

[Caeruleann]

/test | index_copy:H100

[github-actions]

❌ On-demand test failed

Operator: | index_copy
Runner: H100
Backend: H100

The test failed to complete. Check the workflow run for details.

[Caeruleann]

/test index_copy:h20

[github-actions]

Test Environment

Env	Setting	Env	Setting
Time	2026-08-15 15:41:38 -> 2026-08-15 15:40:28	FlagGems	5.3.4.post1.dev82+gb3d986b0c+gb3d986b0c -> 5.3.4.post1.dev82+gb3d986b0c+gb3d986b0c
Arch	x86_64	OS	ubuntu 24.04
Python	3.12.13	Torch	2.11.0+cu130
FlagTree	0.6.1	Triton	3.6.0
Vendor	nvidia	Device	cuda

Test Result

Accuracy Result

Summary

Case	Status	Duration	Total	Passed	Skipped	Failed
Before	Passed	20.45 sec	12	12	0	0
After	Passed	21.70 sec	12	12	0	0

Performance Result

Summary

Case	Status	Duration	Reason
Before	Passed	41.53 seconds	All tests completed.
After	Passed	38.38 seconds	All tests completed.

Benchmark data

Dtype	Speedup	Base	Gems	Shape
fp16	0.381 -> 1.107	-	-	-
	0.884 -> 0.966	35.287 -> 35.289	39.927 -> 36.513	[[1073741824],0,[536870912],[536870912]]
	0.022 -> 1.259	0.009 -> 0.009	0.395 -> 0.007	[[64,64],1,[32],[64,32]]
	0.305 -> 0.856	0.134 -> 0.136	0.441 -> 0.159	[[4096,4096],1,[2048],[4096,2048]]
	0.166 -> 1.227	0.074 -> 0.074	0.446 -> 0.060	[[64,512,512],1,[256],[64,256,512]]
	0.528 -> 1.228	4.101 -> 4.097	7.759 -> 3.338	[[1024,1024,1024],1,[512],[1024,512,1024]]
fp32	0.426 -> 1.167	-	-	-
	0.893 -> 0.989	37.327 -> 37.332	41.823 -> 37.750	[[1073741824],0,[536870912],[536870912]]
	0.023 -> 1.301	0.009 -> 0.009	0.382 -> 0.007	[[64,64],1,[32],[64,32]]
	0.363 -> 0.989	0.158 -> 0.158	0.435 -> 0.160	[[4096,4096],1,[2048],[4096,2048]]
	0.216 -> 1.284	0.096 -> 0.096	0.442 -> 0.075	[[64,512,512],1,[256],[64,256,512]]
	0.637 -> 1.271	5.489 -> 5.481	8.621 -> 4.311	[[1024,1024,1024],1,[512],[1024,512,1024]]
bf16	0.383 -> 1.107	-	-	-
	0.884 -> 0.967	35.278 -> 35.305	39.915 -> 36.521	[[1073741824],0,[536870912],[536870912]]
	0.023 -> 1.259	0.009 -> 0.009	0.390 -> 0.007	[[64,64],1,[32],[64,32]]
	0.314 -> 0.855	0.135 -> 0.133	0.429 -> 0.156	[[4096,4096],1,[2048],[4096,2048]]
	0.168 -> 1.228	0.074 -> 0.074	0.441 -> 0.060	[[64,512,512],1,[256],[64,256,512]]
	0.527 -> 1.228	4.097 -> 4.099	7.768 -> 3.338	[[1024,1024,1024],1,[512],[1024,512,1024]]

---

📎 Download the full log here