Skip to content

cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 - #798

Open
jschmied wants to merge 6 commits into
ashhart:mainfrom
jschmied:cuda-moe-dense-window
Open

jschmied wants to merge 6 commits into
ashhart:mainfrom
jschmied:cuda-moe-dense-window

Conversation

@jschmied

Copy link
Copy Markdown
Contributor

What this changes

Stacked on #746 (the window test checks against its core/attention_ref.zig). Operator-only, no model calls them yet:

  • window_attention.cu/.zig: sliding-window attention, decode rows over a ring of keys (split by key tile, merged
    in tile order) and prompt rows, grouped KV heads, window 0 = full.
  • moe_glue.cu/.zig: top-k by logit + bias (ties: lower id), sigmoid weights, optionally normalised and scaled,
    a shared-expert slot; a warp form up to 512 experts, a block form beyond; the weighted combine of the slots.
  • dense_rows.cu/.zig: bf16 rows times bf16 weight rows into fp32 (routers, output heads), a row's bits at any row
    count. core/moe_ref.zig, core/dense_ref.zig: the float64 references.

Users: families with sliding-window layers or sigmoid-routed MoE on CUDA (Kolibri 1's shapes in the tests; GLM and
DeepSeek-style routing are the same top-k form).

Receipt

  • Environment: TensorFold c36b536d (1.0.9) plus cuda: tree attention and DFlash block attention for every family, checked against float64 #746 and this commit; one DGX Spark (GB10, sm_121, CUDA 13). No
    checkpoint.
  • tf-cuda-test window: 48/4x128 window 513 over rings of 640 and 600 (a ring end inside a tile), 8/2x64 window 100
    with 3 streams, 16/1x128 no window; prompt rows from position 0 and across the ring's wrap: worst 0.465, bits at
    any row count and chunk.
  • tf-cuda-test moe-glue: 384 experts top 6 + shared, 256 top 8 normalised x2.5, 512 top 32 (warp form's widest),
    513 top 6 (block form's narrowest), 4096 top 64; picks exact, weights and sums worst 0.996.
  • tf-cuda-test dense-rows: [384, 2560], [4103, 2560], [257, 7168]: worst 0.007, bits at every row count.
  • Mutations: window bound one key wider (74x), sigmoid exponent x1.0625 (5,527x), a dense row's first product
    dropped (17,529x).
  • sm_120: BobClawblaw ran window, dense-rows, moe-glue on an RTX 5090, all passing with mutations (Qwen3.8-27B on CUDA: the qwen27 family's engine for NVIDIA GPUs, every stream drafted #608).
  • Speed in this PR: none claimed, nothing calls it.
  • Tests run: zig build test and zig build test-cuda-host (27/27 steps, 350/350 tests). zig fmt --check clean.
    lean_check.py: 538 findings, as on main.
  • Not run, and why: Metal (unchanged); no model on main uses it.

Until #746 lands, this branch carries its commit; the change itself is the last commit (2710ae2d).

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (No
    serving path changes.)
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • Prompt processing is no slower than the last release.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

…e_attention.cu, block_attention.cu), checked against core/attention_ref.zig
… merge combines them in order; a slice takes its own one-warp block where four-warp blocks would leave SMs idle (TreeAttention.sms), the same bits
@jschmied
jschmied force-pushed the cuda-moe-dense-window branch from 2710ae2 to bd3e55c Compare October 13, 2026 17:57
…lves again, only one-warp blocks leave them to the merge (a small fixed slice area), so the scratch stays at one partial a group
@jschmied
jschmied force-pushed the cuda-moe-dense-window branch from bd3e55c to 60887de Compare October 13, 2026 18:56
…async stages), the same arithmetic in the same order
@jschmied
jschmied force-pushed the cuda-moe-dense-window branch from 60887de to e0a37c5 Compare October 13, 2026 19:16
…w), partials sized to the query heads, one-warp blocks with 3 tile stages where they fit one wave of the SMs, else 2
…, and dense rows for every family, checked against float64
@jschmied
jschmied force-pushed the cuda-moe-dense-window branch from e0a37c5 to 4ffe90d Compare October 13, 2026 20:40
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.

1 participant