feat(step3): Step-3.5/3.7-Flash family scaffold + config decoders#16
Merged
Conversation
TheTom
changed the title
TheTom
changed the title
4 tasks
Contributor
Author
|
Accompanying metaltile kernels PR: 0xClandestine/metaltile#242 — ships the three net-new primitives this scaffold's forward path will call into ( |
ekryski
marked this pull request as draft
June 3, 2026 23:29
TheTom
changed the title
TheTom
force-pushed
the
tom/wip/step3
branch
2 times, most recently
from
June 10, 2026 13:53
7a57f8e to
d10a964
Compare
Commit message hygiene checkAll commit messages and PR text are clean. ✅ |
…ard yet)
Claim the Step-3 family namespace and wire the loader so a Step-3.5-Flash
or Step-3.7-Flash checkpoint is recognized end-to-end. Concrete forward
+ per-layer code lands in follow-ups; today every load path throws
`Step3Error.notYetImplemented` after the architecture is validated.
### Architecture surface decoded by this scaffold
- **45-layer hybrid text decoder** with a `[full, swa, swa, swa]`
layer-type pattern repeating (12 full-attn + 33 sliding-window-512).
GQA per-layer-type — full-attn carries 64 q-heads, SWA carries 96
q-heads (decoded into `perLayerHeads` overrides), shared 8 KV groups
at `head_dim=128`.
- **Partial RoPE on full-attn only** (`partial_rotary_factors=0.5`).
Per-layer-type RoPE base (5M full / 10K SWA) handled via
`perLayerRopeTheta`. **YARN scaling scoped to selected layer-types**
via `yarn_only_types` — read into `yarnLayerTypes`.
- **MoE on layers 3-44** (288 experts, top-k=8, intermediate=1280, one
always-on shared expert at dim=1280). Dense MLP on layers 0-2 with
intermediate=11264. **Sigmoid + router-bias gate** (DeepSeek-V3
style) routed via `routerBias`/`routerNormTopK`/`routerScalingFactor`.
- **Clamped SwiGLU** (gpt-oss-style activation clip,
`clip(silu(g), max=L) * clip(x, -L, L)`) per-layer via
`perLayerSwigluLimit` + `perLayerSharedSwigluLimit`. Sparse map —
only the layers with a non-zero limit get entries.
- **Head-wise attention output gate** (`use_head_wise_attn_gate`) —
flagged in config; new op needed at forward time.
- **Perception-Encoder vision tower** (47 layers, hidden=1536, 16
heads, head_dim=96, patch=14, image=728). 2D RoPE + additive abs
pos-embed, quick-GELU, LayerScale (`ls_init_value=0.1`), pre-norm.
Two stride-2 Conv2d downsamplers shrink 52×52 → 13×13 = 169 patch
tokens; final projector maps `hidden*4 → text.hidden`.
### Files
- `Sources/FFAI/Models/Step3.swift` — family entry (`enum Step3`,
`Step3Variant` protocol, `Step3Error`). Lists the
text-only/conditional architecture strings + the VL subset.
- `Sources/FFAI/Models/Text/Step3Text.swift` — `Step3TextConfig`
decoder (verbose per-layer overrides), `Step3Hybrid` variant that
decodes the config then throws `notYetImplemented`, and the
placeholder `Step3Model` type.
- `Sources/FFAI/Models/Vision/Step3Vision.swift` — `Step3VLVisionConfig`
decoder + `enum Step3VL` orchestrator. Same pattern: decode +
variant lookup + throw.
- `Sources/FFAI/Loader/Model.swift` — new dispatch arm that hands
Step-3 VL architectures to `Step3VL.load`, matching the
Gemma4VL/Qwen3VL/Qwen3VLMoe shape directly above it.
### Compile
The three new files compile clean against the rest of FFAI. The
`Sources/FFAI/Ops/Ops.swift` build error on this branch is a
pre-existing main-branch regen drift (`ffai_sdpa_decode_d{64,256,512}_*`
grew `has_sink: UInt32` + `sink_logit: Float` constexpr params in the
generated `MetalTileSwift/Generated/MetalTileKernels.swift`, and the
sdpaDecode hand-written dispatch sites weren't updated). Unblocking
that is its own PR.
### Next
1. Forward path on Step3Hybrid — wire the layer stack, hybrid
KV-cache, sigmoid+bias MoE router, clamped-SwiGLU per-layer.
2. Head-wise attention gate Op (small new metaltile kernel or fuse
into o_proj input).
3. Perception-Encoder vision tower forward + projector.
4. Text-only Step-3 dispatch arm in the loader (today only the VL
arm is wired since `Step3Model` doesn't yet conform to
`LanguageModel`).
…on smoke test - Reword the WIP status/doc comments to factual phrasing (the scaffold state — forward stubbed, throws notYetImplemented — is unchanged). - Drop the model-specific default maxTokens on Step3 + Step3Hybrid (falls to the GenerationParameters default), mirroring the DSv4 review fix. - Add Step3RecognitionTests (CI-runnable, no checkpoint): family recognition, hybrid full/sliding config-decode, and default-params convention.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Claim the Step-3 family namespace in FFAI and wire the loader so a Step-3.5-Flash / Step-3.7-Flash checkpoint is recognized end-to-end. The matching GPU kernels are merged in metaltile (#242). Concrete forward + per-layer code lands in follow-ups — every load path validates the config, then throws
Step3Error.notYetImplemented. The scaffold gives a stable type surface to land the forward paths against without churning dispatch.What's in the box
Sources/FFAI/Models/Step3.swiftenum Step3(modelTypes / architectures / VL architectures / variant dispatch),Step3Variantprotocol,Step3ErrorSources/FFAI/Models/Text/Step3Text.swiftStep3TextConfigdecoder (per-layer overrides for the hybrid stack),Step3Hybridvariant, placeholderStep3ModelSources/FFAI/Models/Vision/Step3Vision.swiftStep3VLVisionConfigdecoder (Perception-Encoder ViT + downsamplers + projector),enum Step3VLorchestratorSources/FFAI/Loader/Model.swiftStep3VL.load, matching the Gemma4VL / Qwen3VL patternFile layout follows the per-family convention (
Models/<Family>.swift+Models/Text/<Family>Text.swift+Models/Vision/<Family>Vision.swift), as in Gemma 4 / Qwen 3.Architecture surface decoded by this scaffold
Text backbone — 45-layer hybrid:
[full, swa, swa, swa]pattern (12 full-attention + 33 sliding-window-512); asymmetric GQA per layer-type (full = 64 q-heads / SWA = 96, both 8 KV, head_dim 128); partial RoPE on full-attn only (factor 0.5) with per-layer-type RoPE base + YARN scoped to the full-attn type; MoE on layers 3-44 (288 experts, top-8, +1 shared) with a sigmoid + router-bias gate (DeepSeek-V3 style); per-layer clamped SwiGLU; head-wise attention output gate.Vision tower — Perception-Encoder: 47 layers, hidden 1536, head_dim 96, patch 14, 2D RoPE + abs pos-embed, quick-GELU, LayerScale; post-encoder 2× stride-2 conv shrink → 169 patch tokens;
vit_large_projectormapshidden*4 → text.hidden.Tests
Step3RecognitionTests(CI-runnable, no checkpoint): family recognition by model_type / architecture, the hybrid full/sliding config-decode, and the default-generation-parameters convention.Next steps (separate PRs)
Step3Hybridforward — layer stack, hybrid full/SWA KV cache, sigmoid+bias MoE router, clamped-SwiGLU (kernels merged in metaltile #242).ffai_attn_head_gate, metaltile #242).Step3Modelconforms toLanguageModel(only the VL arm is wired today).Merge sequencing
Targets
dev(rebased 2026-06-10).dev's tip is currently red from the metaltile binding drift (ffai_sdpa_decode_d{64,256}_*grewhas_sink/sink_logit); this branch carries the interim 6-siteOps.swiftfix so its CI builds green against current bindings. #17 carries the same fix — whichever merges first makes the other's copy a no-op on rebase.