WIP: Mamba-Attention Hybrid — 18L (15 SSM + 3 GQA), First SSM Entry#1382

Draft

johnlennyt5 wants to merge 32 commits intoopenai:mainfrom

johnlennyt5:arch1/mamba-hybrid

Commits on Apr 3, 2026

Add tech spec and masterplan for Mamba-Attention hybrid

johnlennyt5
and
claude
committed
[Epic 1 | Task 1.1.1] Copy SOTA baseline to root train_gpt.py
johnlennyt5
committed
[Epic 1 | Task 1.1.2] Add mamba-ssm, causal-conv1d, pytest to requirements
johnlennyt5
committed
[Epic 1 | Task 1.1.3] Add Mamba hyperparameters to Hyperparameters class
johnlennyt5
committed
[Epic 1 | Task 1.2.1-1.2.3] Implement MambaBlock class with __init__, forward, and _selective_scan

johnlennyt5
and
claude
committed
[Epic 1 | Task 1.2.4] Add MambaBlock unit tests (Story 2B.1) — 34 tests all passing

johnlennyt5
and
claude
committed
[Epic 1 | Task 1.3.1] Modify GPT.__init__ for hybrid Mamba/Attention dispatch

johnlennyt5
and
claude
committed
[Epic 1 | Task 1.3.2] Add hybrid forward dispatch and GPT integration tests

johnlennyt5
and
claude
committed
[Epic 2B | Task 2B.1-2B.2] Fix flash_attn mock to use real CPU fallback in all test files

johnlennyt5
and
claude
committed
[Epic 2 | Task 2.2.1-2.2.3] Configure optimizer and quantization pipeline for Mamba params

johnlennyt5
and
claude
committed
[Epic 2 | Task 2.3.1-2.3.2] Add c_proj small normal init for stable SSM output

johnlennyt5
and
claude
committed
[Epic 4 | Task 4.1.1] Extend _HessianGPT with Mamba hybrid support

johnlennyt5
and
claude
committed
[Epic 5 | Task 5.2.1] Use CastedLinear for MambaBlock in_proj/out_proj (QAT support)

johnlennyt5
and
claude
committed
[Epic 3 | Task 3.2.3] Use float32 accumulation in selective scan for bf16 stability

johnlennyt5
and
claude
committed
[Epic 3 | Task 3.1.1+3.1.3] Integrate mamba-ssm and causal-conv1d CUDA kernels

johnlennyt5
and
claude
committed
[Epic 3 | Task 3.2.2] Add gradient checkpointing for Mamba layers + training tests

johnlennyt5
and
claude
committed
[Epic 3 | Task 3.2.1] Add torch.compile compatibility for MambaBlock

johnlennyt5
and
claude
committed
[Epic 2B | Task 2.1.4+2.2.4] Add all-params gradient and optimizer step tests

johnlennyt5
and
claude
committed
[Epic 2B | Tasks 2B.3.3-2B.3.4, 2B.5.1-2B.5.3, 2B.7.1, 4.1.2] Complete remaining CPU tests

johnlennyt5
and
claude
committed

Commits on Apr 4, 2026

Commits on Apr 5, 2026

[Epic 2B | Task 2B.7.3] Add pre-launch local validation with real FineWeb data

johnlennyt5
and
claude
committed

Commits on Apr 6, 2026

Commits on Apr 7, 2026