【KernelGen】Add index_copy_ operator by Schopenhauer-loves-Hegel · Pull Request #1743 · flagos-ai/FlagGems

Schopenhauer-loves-Hegel · 2026-02-27T06:15:39Z

PR Category

Operator

Type of Change

New Feature

Description

Add index_copy_ operator implementation with Triton kernel.

Implementation mode: manual_kernel
Accuracy test: 24/24 passed

Issue

N/A

Progress

Change is properly reviewed (1 reviewer required, 2 recommended).
Change is responded to an issue.
Change is fully covered by a UT.

Performance

torch.float16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0150	0.5200	0.029
[256, 256]	0.0160	0.5170	0.031
[1024, 1024]	0.0230	0.5180	0.044
[4096, 4096]	0.2280	0.5690	0.401
[1024, 65536]	1.0030	1.3320	0.753
[10000, 256]	0.0280	0.5120	0.055
[10000, 65536]	9.6590	11.7830	0.820

torch.float32

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0130	0.5120	0.025
[256, 256]	0.0140	0.5110	0.027
[1024, 1024]	0.0280	0.5650	0.049
[4096, 4096]	0.3140	0.5730	0.548
[1024, 65536]	1.3980	1.3760	1.016
[10000, 256]	0.0370	0.5160	0.072
[10000, 65536]	13.5900	12.2410	1.110

Overall: median speedup = 0.064x, mean speedup = 0.356x (14 data points)

Generated by auto_gen tool with Claude Code

Schopenhauer-loves-Hegel · 2026-02-27T06:34:32Z

Closing in favor of grouped PRs to reduce review overhead.

CLAassistant · 2026-04-15T20:29:53Z

All committers have signed the CLA.

factnn · 2026-04-17T09:05:58Z

index_copy_ Benchmark Speedup Results

float16

Size	Torch (ms)	Gems (ms)	Speedup
[64, 64]	0.016	0.717	0.022
[256, 256]	0.016	0.702	0.023
[1024, 1024]	0.027	0.731	0.036
[4096, 4096]	0.228	0.839	0.272
[1024, 65536]	1.005	1.383	0.726
[10000, 65536]	9.645	11.881	0.812

float32

Size	Torch (ms)	Gems (ms)	Speedup
[64, 64]	0.013	0.722	0.018
[256, 256]	0.014	0.756	0.018
[1024, 1024]	0.028	0.737	0.038
[4096, 4096]	0.315	0.828	0.381
[1024, 65536]	1.393	1.460	0.954
[10000, 65536]	13.617	12.335	1.104

Note: Significant performance regression on small tensors due to kernel launch overhead. Large tensors (float32) show comparable or better performance.

tengqm

overall looks good

tengqm · 2026-04-26T15:03:12Z

+
+
+@pytest.mark.index_copy
+def test_index_copy_perf():


Suggested change

def test_index_copy_perf():

def test_index_copy():

tengqm · 2026-04-26T15:03:36Z

+        op_name="index_copy",
+        torch_op=torch.index_copy,
+        input_fn=index_copy_input_fn,
+        dtypes=[torch.float16, torch.float32],


why skip bf16?

tengqm · 2026-04-26T15:03:57Z

+        op_name="index_copy_",
+        torch_op=torch.Tensor.index_copy_,
+        input_fn=index_copy_input_fn,
+        dtypes=[torch.float16, torch.float32],


same question here

shall we specify inplace=True?

tengqm · 2026-04-26T15:04:24Z

+class TensorSelectBenchmark(GenericBenchmark2DOnly):
+    pass


This additional layer makes no sense.

tengqm · 2026-04-26T15:07:44Z

@@ -0,0 +1,45 @@
+# Feb 27 KernelGen PR Status


Remove this file.

tengqm · 2026-05-10T00:51:21Z

+    def __call__(self, *args, **kwargs):
+        key = f"{self.arg_key(*args)}"
+        if key in self.overloads:
+            overload = self.overloads[key]


If you return here with return overload(*args, **kwargs) here, you don't need the else on line 176.

No. You didn't return, still.

tengqm · 2026-05-11T10:20:25Z

+
+
+def _inplace_input_fn(shape, dtype, device):
+    yield from _tensor_input_fn(shape, dtype, device)


Why this wrapper?

tengqm · 2026-05-11T10:22:43Z

+    def __call__(self, *args, **kwargs):
+        key = f"{self.arg_key(*args)}"
+        if key in self.overloads:
+            overload = self.overloads[key]


No. You didn't return, still.

…d wrapper, rename benchmark Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…HAPES Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Schopenhauer-loves-Hegel requested review from 0x45f, huangyiqun, kiddyjinjin and zhangpeiyang1 as code owners February 27, 2026 06:15

github-actions Bot added benchmark ops/aten core tests size/Medium labels Feb 27, 2026

Schopenhauer-loves-Hegel closed this Feb 27, 2026

Schopenhauer-loves-Hegel reopened this Feb 27, 2026

Schopenhauer-loves-Hegel changed the title ~~Add index_copy_ operator~~ 【KernelGen】Add index_copy_ operator Mar 23, 2026

factnn force-pushed the auto-gen/index_copy_ branch from 445bd7c to 93c173f Compare April 15, 2026 20:29

github-actions Bot added size/Large KernelGen and removed size/Medium labels Apr 15, 2026

factnn force-pushed the auto-gen/index_copy_ branch from 93c173f to 513db40 Compare April 18, 2026 00:45

github-actions Bot added size/Medium and removed size/Large labels Apr 18, 2026

factnn force-pushed the auto-gen/index_copy_ branch 2 times, most recently from 7d42283 to 8b7e646 Compare April 23, 2026 04:40

tengqm reviewed Apr 23, 2026

View reviewed changes

Comment thread src/flag_gems/ops/index_copy_.py Outdated

Comment thread src/flag_gems/ops/index_copy_.py Outdated

factnn requested review from bin913, douxetpur and w1120029931-bit as code owners April 26, 2026 06:39

github-actions Bot added the documentation Improvements or additions to documentation label Apr 26, 2026

tengqm reviewed Apr 26, 2026

View reviewed changes

tengqm reviewed May 10, 2026

View reviewed changes

factnn force-pushed the auto-gen/index_copy_ branch from 43b3b14 to 0b5f06e Compare May 10, 2026 13:32

feat: add index_copy_ operator with tests and benchmark

af10c71

factnn force-pushed the auto-gen/index_copy_ branch from 0b5f06e to af10c71 Compare May 11, 2026 07:09

chore: add index_copy and index_copy_ to operators.yaml

604aaa2

tengqm reviewed May 11, 2026

View reviewed changes

factnn and others added 14 commits May 11, 2026 23:22

fix: address review comments - early return, try/except, remove unuse…

5f63ca2

…d wrapper, rename benchmark Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Merge branch 'master' into auto-gen/index_copy_

1036096

Merge branch 'master' into auto-gen/index_copy_

e43da15

Merge branch 'master' into auto-gen/index_copy_

cf0e82a

ci: trigger CI rerun

a2f9a0a

Merge branch 'master' into auto-gen/index_copy_

57ad597

Merge branch 'master' into auto-gen/index_copy_

a302ec8

Merge branch 'master' into auto-gen/index_copy_

6cdb75d

fix: use IndexCopyBenchmark class consistently in benchmark

cb449e7

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

fix: use explicit shapes in IndexCopyBenchmark instead of REDUCTION_S…

4949861

…HAPES Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

fix: use explicit INDEX_COPY_SHAPES instead of REDUCTION_SHAPES in tests

fb3397d

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Merge branch 'master' into auto-gen/index_copy_

55ff81a

Merge branch 'master' into auto-gen/index_copy_

91dc702

Merge branch 'master' into auto-gen/index_copy_

e363fbb

tengqm reviewed May 13, 2026

View reviewed changes

Comment thread src/flag_gems/ops/index_copy_.py

factnn added 2 commits May 14, 2026 10:05

Merge branch 'master' into auto-gen/index_copy_

4146e9b

Merge branch 'master' into auto-gen/index_copy_

dbe5926

factnn mentioned this pull request May 14, 2026

【KernelGen】Add floor operator #1736

Merged

3 tasks

factnn added 3 commits May 14, 2026 17:27

Merge branch 'master' into auto-gen/index_copy_

1d1b874

Merge branch 'master' into auto-gen/index_copy_

bfec166

Merge branch 'master' into auto-gen/index_copy_

b8af1c7

bin913 approved these changes May 15, 2026

View reviewed changes

Merge branch 'master' into auto-gen/index_copy_

c627a80

tengqm merged commit 8d23621 into flagos-ai:master May 19, 2026
18 checks passed



		def _inplace_input_fn(shape, dtype, device):
		yield from _tensor_input_fn(shape, dtype, device)

Uh oh!

Conversation

Schopenhauer-loves-Hegel commented Feb 27, 2026

PR Category

Type of Change

Description

Issue

Progress

Performance

Uh oh!

Schopenhauer-loves-Hegel commented Feb 27, 2026

Uh oh!

CLAassistant commented Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

factnn commented Apr 17, 2026

index_copy_ Benchmark Speedup Results

float16

float32

Uh oh!

tengqm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

CLAassistant commented Apr 15, 2026 •

edited

Loading