feat(kernel): Add gluon kernel for DSA by Max191 · Pull Request #627 · lightseekorg/tokenspeed

Max191 · 2026-07-09T23:37:11Z

Summary

Adds a gluon kernel implementation for DSA, including the scoring, topk, and sparse attention kernels.

Benchmark results

Benchmarks are compared against the triton implementation using GLM shapes at different sequence lengths:

Decode

Seq len	Candidate rows	Triton scores ms	Triton TopK ms	Triton sparse attn ms	Triton total ms	Gluon scores ms	Gluon TopK ms	Gluon sparse attn ms	Gluon total ms	Speedup
512	512	0.019	0.044	0.305	0.368	0.020	0.018	0.185	0.223	1.65x
1k	1,024	0.022	0.045	0.335	0.402	0.022	0.019	0.183	0.225	1.79x
2k	2,048	0.022	0.045	0.389	0.456	0.022	0.019	0.180	0.221	2.06x
8k	8,192	0.022	0.100	0.389	0.511	0.022	0.064	0.181	0.267	1.91x
16k	16,384	0.022	0.172	0.391	0.585	0.022	0.101	0.180	0.303	1.93x
32k	32,768	0.021	0.318	0.398	0.737	0.020	0.200	0.181	0.401	1.84x
64k	65,536	0.022	0.262	0.392	0.675	0.021	0.198	0.181	0.399	1.69x
90k	90,000	0.022	0.283	0.405	0.711	0.022	0.203	0.181	0.407	1.75x
128k	131,072	0.024	0.318	0.411	0.753	0.024	0.206	0.181	0.411	1.83x
256k	262,144	0.033	0.440	0.404	0.877	0.033	0.212	0.187	0.432	2.03x
512k	524,288	0.045	0.660	0.400	1.105	0.045	0.212	0.185	0.442	2.50x
1m	1,048,576	0.069	1.139	0.442	1.650	0.069	0.216	0.188	0.473	3.49x

Prefill

Seq len	Candidate rows	Triton scores ms	Triton TopK ms	Triton sparse attn ms	Triton total ms	Gluon scores ms	Gluon TopK ms	Gluon sparse attn ms	Gluon total ms	Speedup
512	30,752	0.021	0.070	1.376	1.467	0.021	0.019	0.188	0.228	6.44x
1k	63,520	0.021	0.070	1.407	1.499	0.021	0.019	0.185	0.225	6.67x
2k	129,056	0.025	0.071	1.456	1.552	0.024	0.019	0.183	0.225	6.90x
8k	522,272	0.042	0.127	1.832	2.001	0.040	0.062	0.188	0.291	6.89x
16k	1,046,560	0.064	0.200	1.846	2.109	0.064	0.095	0.188	0.347	6.08x
32k	2,095,136	0.114	0.344	1.853	2.311	0.108	0.225	0.188	0.520	4.44x
64k	4,192,288	0.218	0.291	1.873	2.383	0.195	0.231	0.190	0.615	3.87x
90k	5,757,984	0.289	0.314	1.873	2.476	0.260	0.249	0.191	0.700	3.54x
128k	8,386,592	0.406	0.370	1.888	2.664	0.372	0.282	0.190	0.844	3.15x
256k	16,775,200	0.839	0.628	1.962	3.430	0.779	0.507	0.196	1.482	2.31x
512k	33,552,416	1.592	1.137	1.941	4.670	1.490	0.967	0.197	2.654	1.76x
1m	67,106,848	3.138	2.045	1.956	7.139	2.988	1.864	0.196	5.048	1.41x

Test Plan

Adds new unit tests. All benchmarks above also match numerics between gluon and triton to within 10e-8.

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: eb17ddb706

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-07-10T15:53:53Z

+    softmax_scale: float,
+    q_len_per_req: int = 1,
+    index_k_cache: torch.Tensor | None = None,
+    plan: object | None = None,


Accept forwarded seq_lens_2d in Gluon decode top-k

When this kernel is selected through tokenspeed_kernel.ops.attention.dsa_decode_topk, the public wrapper always forwards seq_lens_2d=... (even when it is None). This implementation does not accept that keyword, so the new GFX950 default selection raises TypeError: unexpected keyword argument 'seq_lens_2d' before launching for every decode top-k call; mirror the Triton/DeepGEMM wrapper signature or absorb the unused kwarg.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-07-10T15:53:53Z

+    out: torch.Tensor | None = None,
+    lens_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del plan, q_len_per_req


Honor q_len_per_req for request-scoped decode inputs

For spec-verify decode with q_len_per_req > 1, the public contract passes seq_lens and block_table per request, and each token's visible length is derived from seq_lens[req] - (q_len_per_req - 1) + j. Dropping q_len_per_req makes the wrapper require one seq_lens/block_table row per token and the kernel indexes them by token, so valid multi-token requests that this kernel advertises in its traits are either rejected or read the wrong request row.

Useful? React with 👍 / 👎.

borontion · 2026-07-10T18:36:28Z

+        [BLOCK_H, D_V],
+        [1, 0],
+    )
+    sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(


why are we using swizzled shared not padded shared here?

I didn't optimize too much yet. I'll try padded shared when I move on to optimizing the kernel further, but I just wanted to land a basic implementation first.

yeah sounds good. we can double check layout part later

borontion · 2026-07-10T18:53:30Z

+        device=q.device, dtype=torch.int64
+    ).contiguous()
+    row_starts = row_starts.to(device=q.device, dtype=torch.int32).contiguous()
+    row_ends = row_ends.to(device=q.device, dtype=torch.int32).contiguous()


is that true we need to make all of the contiguous? this may have overhead. we can not support stride tensor?

Right now the kernel assumes flat layout, so for correctness, we should technically be making them contiguous here, but I'll add this to my list to look into for optimization work.

borontion · 2026-07-10T19:08:27Z

+    else:
+        max_query_rows = max(1, int(max_logits_bytes) // (max(seq_len_sum, 1) * 4))
+    block_n = 32
+    select_warps = 8


any reason we use 8 warps instead of 4 here?

Codex claims this is a tuning decision, although I don't really believe Codex did proper benchmarking/tuning to determine this :P

I'll look into this as well when I optimize.

borontion · 2026-07-10T19:15:23Z

+    equal_i32 = equal.to(gl.int32)
+    tile_greater = gl.sum(greater_i32, axis=0).to(gl.int32)
+    tile_equal = gl.sum(equal_i32, axis=0).to(gl.int32)
+    greater_start = gl.atomic_add(


we have buffer_atomic_add will that be better?

I'll experiment with this as well during the optimization phase. Apparently, buffer_atomic_add does not support signed integers (I'd have to use unsigned). I think unsigned should work here too, but I'll wait until I start collecting profiles before I make the change.

borontion

thanks Max. we can continue optimizing it

Max191 added 12 commits July 10, 2026 15:44

Add GLM-style DSA correctness tests

5b80ee4

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Add Gluon DSA top-k path

f774151

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Add Gluon DSA selected attention path

f38070c

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Port Gluon DSA dense path to tiled MFMA

4ada190

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Implement Gluon DSA top-k selector

8248e61

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Avoid Gluon DSA TopK long-row compile cliff

9a280c2

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Route long-row Gluon DSA TopK through Triton

ff356fe

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Keep long-row DSA TopK in Gluon

f82285a

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Generalize Gluon DSA attention support

277e511

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Split Gluon DSA TopK scoring kernels

3c0d825

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Move Gluon DSA tests to AMD package

1b61464

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Select Gluon DSA by default in golden test

eb17ddb

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

Max191 force-pushed the Max191/gluon-dsa-kernel branch from abecc91 to eb17ddb Compare July 10, 2026 15:48

Max191 marked this pull request as ready for review July 10, 2026 15:50

Max191 requested review from a team as code owners July 10, 2026 15:50

chatgpt-codex-connector Bot reviewed Jul 10, 2026

View reviewed changes

borontion reviewed Jul 10, 2026

View reviewed changes

borontion approved these changes Jul 14, 2026

View reviewed changes

Uh oh!

Conversation

Max191 commented Jul 9, 2026

Summary

Benchmark results

Decode

Prefill

Test Plan

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 10, 2026

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!

borontion Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

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!

borontion left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

borontion Jul 14, 2026 •

edited

Loading

borontion left a comment •

edited

Loading