restore models/mamba3/ scaffold - #11
Conversation
Only mamba1 (models/mamba/) was meant to be dropped. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9caad91e19
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| e->setBuffer(cL, 0, 7); | ||
| e->setBuffer(cDQ, 0, 8); | ||
| e->setBuffer(cDV, 0, 9); | ||
| e->setBuffer(cCS, 0, 10); | ||
| e->dispatchThreadgroups( |
There was a problem hiding this comment.
Bind required SSM state buffers before dispatch
dispatch_block launches mamba3_ssm after binding only buffers 0–10, but the kernel signature also requires buffers 11–15 (h_state_in, h_state_out, a_cs_in, a_cs_out, state_flags in mamba3_ssm.metal). As written, this dispatch can fail Metal argument validation or run with undefined state pointers/flags, so the SSM stage is not safely executable.
Useful? React with 👍 / 👎.
| e->setBuffer(gate_buf, 0, 2); | ||
| e->setBuffer(cL, 0, 3); | ||
| e->setBuffer(cDV, 0, 4); | ||
| e->dispatchThreadgroups( |
There was a problem hiding this comment.
Fix post_ssm buffer index mapping to kernel ABI
The post_ssm dispatch does not match the kernel’s declared arguments: mamba3_post_ssm expects norm_w at buffer 2 and output at buffer 3, with constants at 4/5/6, but this call binds gate_buf at index 2 and shifts constants down. That makes the kernel read/write the wrong buffers and constants, yielding incorrect results or runtime faults.
Useful? React with 👍 / 👎.
| e->dispatchThreadgroups( | ||
| MTL::Size(BH, 1, 1), | ||
| MTL::Size(128, 1, 1)); |
There was a problem hiding this comment.
Dispatch SSM over all DV tiles
The restored mamba3_ssm kernel is DV-tiled (BV=32) and uses gid.y to select which output slice to compute, but dispatch_block always launches with y=1. For any DV > 32, only the first tile is produced and the remaining channels are left uncomputed, so outputs are truncated/garbage in normal wider configs.
Useful? React with 👍 / 👎.
PR #5 squash-merged the deletion of models/mamba3/ without including the follow-up restore commit (eb80c9e on dev-cleanup). This cherry-picks that restore commit onto a fresh branch off main.
10 files / ~1.4k LOC restored. Scaffold only; no working inference. Build verified clean (./build.sh produces both libsk.dylib + libsk.metallib).