Skip to content

perf(qwen36): qualify M4 head-256 attention - #640

Open
Gajesh2007 wants to merge 1 commit into
masterfrom
perf/qwen36-head256-prefill
Open

perf(qwen36): qualify M4 head-256 attention#640
Gajesh2007 wants to merge 1 commit into
masterfrom
perf/qwen36-head256-prefill

Conversation

@Gajesh2007

@Gajesh2007 Gajesh2007 commented Aug 17, 2026

Copy link
Copy Markdown
Member

Summary

  • pin the complete force-aware MLX -> mlx-c -> mlx-swift -> mlx-swift-lm chain
  • qualify composed qL=512 prefill only for the exact tested Mac16,5 M4 Max (40 GPU cores, 128 GB)
  • retain qL=128 everywhere else and preserve the operator override
  • expose DARKBLOOM_CBV2_ATTN_EXECUTION=fused as an explicit bounded-memory D256 arm; automatic fused selection stays closed because it is slower
  • add active/peak/transient memory evidence to scheduler-prefill schema 3 and update its strict consumer fixtures
  • require both FP16 and BF16 D256 symbols in source-matched and final signed metallibs
  • extend the real Qwen artifact canary through default, forced-fused, VLM, tools, MTP, and cancellation paths
  • prepare provider v0.8.6 and keep the coordinator fallback version synchronized

Dependency PRs:

Before

flowchart LR
  A[Qwen3.6 prefill, 512-token scheduler chunk] --> B[split into 4 x qL128]
  B --> C[composed D256 attention]
  C --> D[baseline TTFT]
  E[long-context memory pressure] --> C
  F[release metallib] --> G[no D256 symbol contract]
Loading

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]
Loading

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]
  end
Loading

M4 Max Results

Production CBv2, release build, contiguous KV, exact Qwen artifact, same final dependency chain:

Prompt qL128 control qualified qL512 Improvement
8k 6,206.8 ms 6,056.6 ms 2.4%
32k 32,879.5 ms 31,660.9 ms 3.7% / 1.22 s

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

  • coordinator: go test ./...
  • provider: 2,094 tests in 214 suites
  • mlx-swift: 541 tests, 6 skipped
  • mlx-swift-lm: 528 tests, 1 skipped
  • benchmark contract: 62 Python tests
  • exact BF16 Hq16/Hkv2 qL 4/5/8/9 route tests
  • default and explicit forced-fused real-artifact production canary
  • release build with source-matched no-JIT metallib
  • metallib SHA-256 bbbdecbdacb2406f0fde6454493d4a9de82685047c69531977d90e3829c88ec7
  • release workflow YAML and metallib helper syntax checks

Follow-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.


View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need. Autofix is disabled.

@vercel

vercel Bot commented Aug 17, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
d-inference Ready Ready Preview Aug 17, 2026 10:19am
d-inference-console-ui-dev Ready Ready Preview Aug 17, 2026 10:19am
d-inference-landing Ready Ready Preview Aug 17, 2026 10:19am

Request Review

@chatgpt-codex-connector chatgpt-codex-connector 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.

💡 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".

Comment on lines +12 to +16
(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

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge 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 👍 / 👎.

Comment on lines 582 to +587
String(describing: type(of: model)))
}

let resolvedAttentionExecutionPolicy = attentionExecutionPolicy
?? QwenHead256AttentionTuning.currentMachinePolicy(
isQwen35: isQwen35,

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge 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 👍 / 👎.

Comment on lines +26 to +30
- 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.

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge 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 👍 / 👎.

Comment on lines +23 to +25
public let activeMemoryBeforeBytes: Int
public let peakMemoryBytes: Int
public let transientPeakBytes: Int

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge 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 👍 / 👎.

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