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

Schopenhauer-loves-Hegel · 2026-02-27T06:14:09Z

PR Category

Operator

Type of Change

New Feature

Description

Add feature_dropout operator implementation with Triton kernel.

Implementation mode: manual_kernel
Accuracy test: 96/96 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.bfloat16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[8, 64]	0.0128	0.0233	0.549
[16, 128]	0.0127	0.0223	0.571
[32, 256]	0.0128	0.0196	0.650
[4, 64, 32, 32]	0.0162	0.0214	0.757
[8, 128, 16, 16]	0.0162	0.0229	0.709
[16, 256, 8, 8]	0.0162	0.0216	0.750
[2, 512, 14, 14]	0.0161	0.0216	0.747
[4, 1024, 7, 7]	0.0190	0.0263	0.722

torch.float16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[8, 64]	0.0156	0.0389	0.401
[16, 128]	0.0156	0.0177	0.882
[32, 256]	0.0128	0.0231	0.553
[4, 64, 32, 32]	0.0180	0.0216	0.833
[8, 128, 16, 16]	0.0162	0.0220	0.739
[16, 256, 8, 8]	0.0163	0.0204	0.795
[2, 512, 14, 14]	0.0174	0.0219	0.798
[4, 1024, 7, 7]	0.0174	0.0229	0.761

torch.float32

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[8, 64]	0.0140	0.0201	0.696
[16, 128]	0.0125	0.0264	0.474
[32, 256]	0.0140	0.0195	0.720
[4, 64, 32, 32]	0.0184	0.0272	0.676
[8, 128, 16, 16]	0.0155	0.0220	0.703
[16, 256, 8, 8]	0.0155	0.0203	0.762
[2, 512, 14, 14]	0.0164	0.0220	0.749
[4, 1024, 7, 7]	0.0151	0.0228	0.663

Overall: median speedup = 0.721x, mean speedup = 0.694x (24 data points)

Generated by auto_gen tool with Claude Code

Schopenhauer-loves-Hegel · 2026-02-27T06:35:01Z

Closing in favor of grouped PRs to reduce review overhead.

CLAassistant · 2026-04-15T20:27:04Z

All committers have signed the CLA.

factnn · 2026-04-17T08:58:32Z

feature_dropout Benchmark Speedup Results

float16

Size	Torch (ms)	Gems (ms)	Speedup
[8, 64]	0.016	0.020	0.786
[16, 128]	0.016	0.018	0.898
[4, 64, 32, 32]	0.022	0.020	1.085
[8, 128, 16, 16]	0.017	0.020	0.835
[2, 512, 14, 14]	0.016	0.041	0.395
[4, 1024, 7, 7]	0.016	0.037	0.432

float32

Size	Torch (ms)	Gems (ms)	Speedup
[8, 64]	0.014	0.037	0.384
[4, 64, 32, 32]	0.015	0.051	0.301
[8, 128, 16, 16]	0.016	0.044	0.375
[4, 1024, 7, 7]	0.017	0.048	0.344

bfloat16

Size	Torch (ms)	Gems (ms)	Speedup
[8, 64]	0.014	0.045	0.304
[4, 64, 32, 32]	0.018	0.045	0.392
[16, 256, 8, 8]	0.016	0.034	0.482
[4, 1024, 7, 7]	0.016	0.042	0.382

Note: feature_dropout shows performance regression across most sizes, likely due to kernel launch overhead on small tensors.

tengqm

Overall lgtm.

tengqm

lgtm

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

factnn · 2026-05-08T08:40:46Z

UT Results

Branch rebased to upstream/master and tests passed:

python -m pytest tests/ -m feature_dropout --ref cpu -x -q
72 passed

…ebase

…i#1732) Co-authored-by: factnn <1050552884@qq.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: bin913 <842884726@qq.com>

Co-authored-by: factnn <1050552884@qq.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: bin913 <842884726@qq.com>

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

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 feature_dropout operator~~ 【KernelGen】Add feature_dropout operator Mar 23, 2026

factnn force-pushed the auto-gen/feature_dropout branch from ff45ed6 to a5ccfe4 Compare April 15, 2026 20:26

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

factnn force-pushed the auto-gen/feature_dropout branch from a5ccfe4 to fb25338 Compare April 17, 2026 23:55

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

factnn force-pushed the auto-gen/feature_dropout branch 2 times, most recently from 954672d to 91fede8 Compare April 23, 2026 03:59

tengqm reviewed Apr 23, 2026

View reviewed changes

Comment thread benchmark/test_feature_dropout_perf.py Outdated

tengqm reviewed Apr 23, 2026

View reviewed changes

Comment thread src/flag_gems/__init__.py

Comment thread tests/test_feature_dropout.py Outdated

Comment thread tests/test_feature_dropout.py Outdated

Comment thread tests/test_feature_dropout.py Outdated

factnn requested review from bin913, douxetpur and w1120029931-bit as code owners April 24, 2026 07:41

factnn force-pushed the auto-gen/feature_dropout branch from 1ca7b61 to 373d609 Compare April 24, 2026 07:48

factnn force-pushed the auto-gen/feature_dropout branch from 373d609 to 3bbee51 Compare April 26, 2026 06:31

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

tengqm reviewed Apr 26, 2026

View reviewed changes

Comment thread benchmark/test_feature_dropout_perf.py Outdated

tengqm reviewed Apr 26, 2026

View reviewed changes

Comment thread benchmark/test_feature_dropout.py

tengqm reviewed Apr 26, 2026

View reviewed changes

Comment thread src/flag_gems/ops/feature_dropout.py

tengqm reviewed Apr 26, 2026

View reviewed changes

Comment thread PR_status.md Outdated

github-actions Bot removed the documentation Improvements or additions to documentation label Apr 27, 2026

tengqm previously approved these changes Apr 27, 2026

View reviewed changes

tengqm reviewed Apr 27, 2026

View reviewed changes

factnn dismissed tengqm’s stale review via 07c7c3d April 28, 2026 10:28

factnn force-pushed the auto-gen/feature_dropout branch from ec84d8d to 07c7c3d Compare April 28, 2026 10:29

tengqm reviewed Apr 28, 2026

View reviewed changes

Comment thread benchmark/test_feature_dropout.py Outdated

tengqm previously approved these changes Apr 28, 2026

View reviewed changes

factnn dismissed tengqm’s stale review via aedc370 April 28, 2026 10:43

factnn force-pushed the auto-gen/feature_dropout branch from 0b82cb6 to aedc370 Compare April 28, 2026 10:43

feat: add feature_dropout operator with tests and benchmark

1f90724

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

factnn force-pushed the auto-gen/feature_dropout branch from e282855 to 1f90724 Compare May 8, 2026 07:16

factnn and others added 2 commits May 8, 2026 22:29

Merge remote-tracking branch 'upstream/master' into feature_dropout_r…

e88e6fc

…ebase

Merge branch 'master' into auto-gen/feature_dropout

7d65cdf

bin913 approved these changes May 9, 2026

View reviewed changes

bin913 merged commit f92157a into flagos-ai:master May 9, 2026
18 checks passed

tengqm pushed a commit that referenced this pull request May 17, 2026

feat: add feature_dropout operator with tests and benchmark (#1732)

8342309

Co-authored-by: factnn <1050552884@qq.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: bin913 <842884726@qq.com>

Navigation Menu

Uh oh!

【KernelGen】Add feature_dropout operator - #1732

【KernelGen】Add feature_dropout operator#1732
bin913 merged 3 commits into
flagos-ai:masterfrom
Schopenhauer-loves-Hegel:auto-gen/feature_dropout

Schopenhauer-loves-Hegel commented Feb 27, 2026

Uh oh!

Schopenhauer-loves-Hegel commented Feb 27, 2026

Uh oh!

CLAassistant commented Apr 15, 2026 •

edited

Loading

Uh oh!

factnn commented Apr 17, 2026

Uh oh!

Uh oh!

tengqm left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tengqm left a comment

Uh oh!

Uh oh!

factnn commented May 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

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

feature_dropout Benchmark Speedup Results

float16

float32

bfloat16

Uh oh!

Uh oh!

tengqm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tengqm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

factnn commented May 8, 2026

UT Results

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

CLAassistant commented Apr 15, 2026 •

edited

Loading