Request: Optimized attention backend (fwd/bwd) support for GPT-OSS with learnable sink

[CPFLAME]

Hi team,
I’d like to ask whether there is forward and backward support for GPT-OSS attention when using a learnable sink.

In GPT-OSS, the attention module includes a learnable sink parameter. As a result, we currently have to train in the eager attention mode. This makes memory usage scale quadratically with sequence length, which prevents training on long contexts. I found that some kernels provide forward support for a learnable sink, but I couldn’t find full backward support, so we still can’t use an optimized backend end-to-end for training.

Questions:

• Is there an existing or planned roadmap to support learnable-sink attention in both forward and backward for optimized backends on recent GPUs?
• If yes, could you share an ETA or pointers to the relevant branch/PR?
• A minimal example showing how to enable this with GPT-OSS (model/config flags, attention impl selection) would be very helpful.

Environment (for reference):

Model: GPT-OSS (20B,120B)
GPU: H20

Thanks a lot!

[chengtbf]

+1，we also need this feature

[Rachmanino]

Hi, thanks for your interests. I've implemented a simple example in #853 demonstrating how to support the forward process of attention sink and sliding window attention in TileLang. We'll also work on the backward process and other variants later. Any help or suggestions are welcomed!

[Rachmanino]

The backward kernel development is in #877

[Rachmanino]

Based on some mathmetical derivation, the forward and backward process can be easily adapted from the FA schedule. For best performance of training GQA models like gpt-oss on Hopper, attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py and attention_sink/example_gqa_bwd_bhsd.py is preferable. Also note that the backward performance can be further optimized with manual pipeline design, which will be considered later.

[LeiWang1999]

likely this issue has been addressed, and the performance also looks good.

Benchmark of forward process

Benchmark Environment

• Hardware: NVIDIA H800
• CUDA version: 12.9
• Triton Version: 3.4.0

Results

• dtype=float16
• batch_size=1, heads=64, kv_heads=8 (the setting of GPT-OSS-120B)
• Full attention is adopted.

SEQ_LEN	headdim	Triton TFLOPs	TileLang TFLOPs	Speedup
2048	64	231.55	277.07	1.20x
2048	128	313.55	393.98	1.26x
4096	64	272.17	337.30	1.24x
4096	128	356.35	461.54	1.30x
8192	64	289.93	353.81	1.22x
8192	128	392.18	482.50	1.23x
16384	64	299.52	377.44	1.26x
16384	128	404.64	519.02	1.28x

Closed.

[chengtbf]

Thank you very much. We will give a try to run this kernel.

[CPFLAME]

Dear authors,
First of all, thank you for releasing this excellent operator and the clear examples. Your work is highly valuable and much appreciated by practitioners—we were able to get it running quickly and see promising results.

While testing, I noticed a few opportunities that might further improve robustness and usability:

• Consider making the input v contiguous internally (e.g., calling v = v.contiguous() before use). In some models, invoking the operator directly during forward can otherwise trigger bugs due to non-contiguous tensors.
• It would be very helpful if the operator supported both fp16 and bf16 dtypes. Many training pipelines run with bf16, so dual support would broaden applicability.
• Adding a softmax_scale argument would help align the loss accuracy with reference implementations and improve numerical parity in our experiments.

Thank you again for the great work and for sharing it with the community.

[Rachmanino]

Dear authors, First of all, thank you for releasing this excellent operator and the clear examples. Your work is highly valuable and much appreciated by practitioners—we were able to get it running quickly and see promising results.

While testing, I noticed a few opportunities that might further improve robustness and usability:

• Consider making the input v contiguous internally (e.g., calling v = v.contiguous() before use). In some models, invoking the operator directly during forward can otherwise trigger bugs due to non-contiguous tensors.
• It would be very helpful if the operator supported both fp16 and bf16 dtypes. Many training pipelines run with bf16, so dual support would broaden applicability.
• Adding a softmax_scale argument would help align the loss accuracy with reference implementations and improve numerical parity in our experiments.

Thank you again for the great work and for sharing it with the community.

Thanks for your recognition, and glad to hear that our kernels help in your work. We'll improve them later according to your feedback.

[Rachmanino]

Hi @CPFLAME, we've supported bfloat16 and customized sm_scale for attention sink via pr #924. We also updated README to announce the performance of bf16 kernels.