Skip to content

[2/8] refactor(kernels): move attention kernels to miles/kernels/attention - #3606

Merged
Zhichenzzz merged 1 commit into
zhichen/kernels-reorg-quant-moe-activationfrom
zhichen/kernels-reorg-attention
Oct 8, 2026
Merged

Zhichenzzz merged 1 commit into
zhichen/kernels-reorg-quant-moe-activationfrom
zhichen/kernels-reorg-attention

Conversation

@Zhichenzzz

@Zhichenzzz Zhichenzzz commented Sep 22, 2026 •

Copy link
Copy Markdown
Collaborator

What

Files the attention kernels under miles/kernels/attention/. Moves and import updates only.

The GDN backend selector stays in miles_plugins/models: it only picks between fla and FlashQLA and owns no kernel. DeepSeek-V4.1's dsa_topk import follows the selector; left alone, it would break at import time without any textual conflict.

Tests

  • Pre-rebase run on H200: the GDN backend test, both manual DeepSeek-V4 TileLang suites (indexer and sparse MLA) and the GLM-5 / DSv4 model imports give results identical to main.
  • pre-commit clean.

This branch on H200 (45f9d98bc), against main

Test This branch main
py_compile + import of every changed module pass –
tests/fast/test_linear_attn_kernel_selection.py 3 passed 3 passed
DeepSeek-V4 TileLang indexer suite 45 passed 45 passed
DeepSeek-V4 TileLang sparse MLA suite 23 passed 23 passed

Stack

PR Change
1 #3605 non-attention kernels → miles/kernels
2 #3606 attention kernels → miles/kernels/attention
3 #3607 one DSA kernel pair (GLM-5, DeepSeek-V4, V4.1)
4 #3608 per-arch DSA sparse attention dispatch
5 #3609 shared head-sharded linear-attention layer (miles_plugins/models/linear_attn.py)
6 #3610 Qwen3.5 / Qwen3-Next GDN on it
7 #3632 Kimi-K3 KDA on it
8 #3634 short-conv and KDA backward backends

Rules for miles/kernels are in its README: only kernels we write, one directory per op, names say what is computed (never the backend), callers import the entry, a torch-reference test per op.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

Tip: disable this comment in your organization's Code Review settings.

@HJSang HJSang left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Just refactoring and no risk to merge

@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch from 3c45a02 to 817ed9e Compare September 30, 2026 21:57
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch from 817ed9e to 188e617 Compare October 2, 2026 07:51
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch 2 times, most recently from 4c14c11 to f82f79a Compare October 2, 2026 22:31
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch from f82f79a to 7170a9a Compare October 3, 2026 07:43
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch from 7170a9a to 69cd55c Compare October 7, 2026 22:54
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch from 69cd55c to 0bde626 Compare October 7, 2026 23:45
@Zhichenzzz Zhichenzzz changed the title refactor(kernels): file the attention kernels under miles/kernels/attention [2/8] refactor(kernels): file the attention kernels under miles/kernels/attention Oct 7, 2026
@Zhichenzzz Zhichenzzz changed the title [2/8] refactor(kernels): file the attention kernels under miles/kernels/attention [2/8] refactor(kernels): move attention kernels to miles/kernels/attention Oct 8, 2026
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch 2 times, most recently from 254e1b0 to d540607 Compare October 8, 2026 15:19
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch from d540607 to 6cdfb87 Compare October 8, 2026 16:47
…ention

Second half of filing every hand-written training kernel by op. Moves plus
import updates; kernel sources are unchanged.

- dsa/: the GLM-5 / DeepSeek-V3.2, DeepSeek-V4 and DeepSeek-V4.1 TileLang
  indexer and sparse-attention kernels, still one copy per model (the next
  PR merges them), and the top-k selector (from
  miles_plugins/models/dsa_topk.py). V4.1's indexer import of dsa_topk
  follows the selector; left alone it breaks at import time with no
  textual conflict.
- dsa/kpool.py: the pure-Triton half of GLM-5.3-Flash's key-pooled
  selection splits out of glm5_next/ops/kpool_indexer.py; the replay-aware
  kpool_select_topk stays in the plugin.
- qsa/: Qwen3.8-Next's Triton block-sparse attention, out of
  qwen3_8_next/ops/kernel/ (sparse_attention.py, block_sparse_attention.py),
  exported from the package.
- deepseek_v4/ops/kernel is gone; its precision-aligned ops move up to
  deepseek_v4/ops.
- README: the layout and the rules. The GDN backend selector stays in
  miles_plugins/models: it only picks between fla and FlashQLA and owns no
  kernel.
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/kernels-reorg-attention branch from 6cdfb87 to a789e40 Compare October 8, 2026 19:16
@Zhichenzzz
Zhichenzzz merged commit 574baf7 into main Oct 8, 2026
38 of 41 checks passed
@Zhichenzzz
Zhichenzzz deleted the zhichen/kernels-reorg-attention branch October 8, 2026 19:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants