perf(qwen36): qualify M4 head-256 attention - #640
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dc4c3391b1
ℹ️ 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".
| (try? sysctlString("hw.model")) == qualifiedMachineModel, | ||
| (try? sysctlString("machdep.cpu.brand_string")) == "Apple M4 Max" | ||
| else { return false } | ||
| let bytesPerGB = UInt64(1024 * 1024 * 1024) | ||
| return ProcessInfo.processInfo.physicalMemory / bytesPerGB >= qualifiedMemoryGB |
There was a problem hiding this comment.
Enforce the GPU-core qualification in production
On a 128 GB Mac16,5 M4 Max with 32 GPU cores—the explicit negative case in QwenHead256AttentionTuningTests.unqualifiedHardware—this production predicate still returns true because it checks only the model identifier, brand, and memory. Such a provider therefore receives the unqualified qL=512 policy even though the test and release report require qL=128 outside the exact 40-core configuration; use the same full HardwareInfo qualification on the production path or otherwise verify the GPU-core count.
AGENTS.md reference: docs/AGENTS.md:L19-L22
Useful? React with 👍 / 👎.
| String(describing: type(of: model))) | ||
| } | ||
|
|
||
| let resolvedAttentionExecutionPolicy = attentionExecutionPolicy | ||
| ?? QwenHead256AttentionTuning.currentMachinePolicy( | ||
| isQwen35: isQwen35, |
There was a problem hiding this comment.
Gate the wide block on the D256 shape
On the qualified machine, every Qwen35MoEModel receives qL=512 regardless of its attention head dimension, although this change qualifies only the D256 Qwen3.6 artifact. The accepted qwen3_5_moe family is not shape-invariant—the repository's production-factory fixture uses the same model class with head_dim: 64—so loading another supported Qwen3.5/3.6 target silently applies the unmeasured wider block instead of retaining qL=128. Derive eligibility from the model's layer kinds/head dimension rather than the family type alone.
Useful? React with 👍 / 👎.
| - qualified M4 Max providers use composed attention with qL=512 for speed; | ||
| - `force_fused` explicitly requests the D256 Steel kernel when bounded | ||
| transient memory is more important than throughput; | ||
| - automatic fused selection remains unqualified and fails closed; | ||
| - other hardware retains the historical qL=128 posture pending measurement. |
There was a problem hiding this comment.
Add canonical code citations to the behavior claims
These public behavioral claims about automatic qL selection and fused execution provide no canonical source paths or line references. Add citations to the production tuning and factory wiring so the report remains traceable to the implementation, as required for documentation that describes behavior.
AGENTS.md reference: docs/AGENTS.md:L19-L22
Useful? React with 👍 / 👎.
| public let activeMemoryBeforeBytes: Int | ||
| public let peakMemoryBytes: Int | ||
| public let transientPeakBytes: Int |
There was a problem hiding this comment.
Record the attention arm in benchmark evidence
When this benchmark is used for the qL=128 versus qL=512 or composed versus fused comparisons introduced here, the JSON records memory and the resolved KV backend but not the requested or resolved attention control/query-block size. Because those arms are selected through process environment, a missing, malformed, or unforwarded variable can make both runs execute the same arm while the strict Python validator still accepts them as a valid comparison. Serialize the effective attention policy and block size in each sample or report and require the expected values during validation.
Useful? React with 👍 / 👎.
Summary
Mac16,5M4 Max (40 GPU cores, 128 GB)DARKBLOOM_CBV2_ATTN_EXECUTION=fusedas an explicit bounded-memory D256 arm; automatic fused selection stays closed because it is slowerDependency PRs:
Before
After
flowchart LR A[Qwen3.6 prefill] --> B{exact Mac16,5 40-core 128GB?} B -- yes --> C[composed qL512] B -- no --> D[historical qL128] E{execution=fused?} -- yes --> F[validated D256 Steel bounded-memory arm] E -- no --> C C --> G[faster qualified TTFT] F --> H[lower transient, slower, opt-in] I[release metallib] --> J[FP16 + BF16 D256 symbol gates]Code Flow
flowchart LR subgraph Before A1[EngineV2Factory] --> B1[CBv2LayerCache default policy] B1 --> C1[AttentionV1 static qL128] D1[fetch-metallib] --> E1[existing symbol contract] end subgraph After A2[QwenHead256AttentionTuning] --> B2[EngineV2Factory injected policy] B2 --> C2[composed qL512 or explicit fused] C2 --> D2[MLXFast force-aware API] E2[SchedulerPrefillBenchmark v3] --> F2[TTFT + transient memory evidence] G2[fetch-metallib + release smoke] --> H2[D256 FP16/BF16 symbols] endM4 Max Results
Production CBv2, release build, contiguous KV, exact Qwen artifact, same final dependency chain:
At 32k, qL512 adds 40.6 MiB transient (+2.2%). Explicit forced-fused uses 1.668 GB transient versus 1.974 GB (-15.5%) but is slower, so it is not selected automatically.
Four alternate fused speed kernels were implemented and measured (Q8 head-sharded, Q32/D128-sharded, two-head GQA, eight-head GQA). All passed numerical/causal gates but lost to composed attention and were removed; no losing experimental route ships.
Full evidence:
docs/reports/2026-08-17-qwen36-head256-attention.md.Verification
go test ./...bbbdecbdacb2406f0fde6454493d4a9de82685047c69531977d90e3829c88ec7Follow-up
M3 Max qualification will be added after hardware access is provided. M5/NAX remains a separate follow-up and is not enabled by this PR.
Need help on this PR? Tag
@codesmith-botwith what you need. Autofix is disabled.