Add Phi-4 Flash SambaY support - #729
Open
justinchuby wants to merge 1 commit into
Open
Conversation
Implement the pinned Phi-4 mini flash reasoning architecture with source-faithful Mamba, differential attention, YOCO state topology, dynamic cache ABI, exact scoped CUDA parity coverage, and explicit runtime/cache limitations. Co-authored-by: Copilot App <[email protected]> Copilot-Session: 489bb9f5-59db-436a-a63b-eb8164e879e3 Signed-off-by: Justin Chu <[email protected]>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| @@ -0,0 +1,344 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,124 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,477 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,477 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| transient layer-16 memory and layer-17 KV within the same graph invocation. | ||
| """ | ||
|
|
||
| from __future__ import annotations |
| @@ -0,0 +1,204 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,250 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,250 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| def _new_remote_cache(reference_model: Any, config: Any, batch: int, max_length: int) -> Any: | ||
| """Instantiate the source's fixed-capacity cache so decode can grow past prefill.""" | ||
| module = importlib.import_module(type(reference_model.model).__module__) | ||
| cache_type = getattr(module, "SambaYCache") |
| "configuration_phi4flash", "modeling_phi4flash" | ||
| ) | ||
| ) | ||
| reference_class = getattr(reference_module, "Phi4FlashForCausalLM") |
Performance Comparison
|
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The pinned CUDA workflow can inadvertently reinstall a non-CUDA torch build due to torch==2.6.0 being pinned in requirements/ci/phi4flash.txt after installing a CUDA wheel from the PyTorch CUDA index.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds first-class support for the phi4flash (Phi-4 Flash Reasoning) architecture to mobius, including its heterogeneous SambaY cache ABI and a pinned-source CUDA parity harness that validates the exported ONNX graph against the upstream trust_remote_code implementation.
Changes:
- Introduces
Phi4FlashConfig,Phi4FlashCausalLMModel, andPhi4FlashCausalLMTask, registers them in the model/task registries, and documents model catalog support. - Adds new reusable building blocks required by Phi-4 Flash (portable differential GQA attention and a state-threaded Mamba-1 block), plus targeted graph/unit coverage.
- Adds a scoped self-hosted CUDA workflow + pinned dependency stack and a CUDA integration test for L3 numerical parity against the pinned remote implementation.
File summaries
| File | Description |
|---|---|
| tests/synthetic_parity_test.py | Skips phi4flash in synthetic parity due to pinned CUDA reference dependency needs. |
| tests/phi4flash_integration_test.py | New pinned-source CUDA L3 parity test validating logits, cache slots, and post-layer residual streams. |
| tests/model_coverage_test.py | Adds coverage notes for phi4flash and its Phi4FlashForCausalLM alias. |
| tests/build_graph/recurrent_test.py | Excludes phi4flash from generic SSM tests and adds dedicated Phi-4 Flash graph/cache-contract assertions. |
| tests/build_graph/_support.py | Includes phi4flash and Phi4FlashForCausalLM in the hybrid/dedicated test matrix. |
| tests/_test_configs.py | Adds a tiny Phi4FlashConfig entry to the SSM config matrix for L1 graph build coverage. |
| testdata/cases/causal-lm/phi4-mini-flash-reasoning.yaml | New (skipped) L4/L5 case entry documenting limitations and why golden runs aren’t committed yet. |
| src/mobius/tasks/_phi4flash.py | New task wiring for Phi-4 Flash’s dynamic heterogeneous SambaY cache ABI + metadata emission. |
| src/mobius/tasks/init.py | Exposes Phi4FlashCausalLMTask and registers the phi4flash-text-generation task key. |
| src/mobius/models/phi4flash.py | New Phi-4 Flash SambaY model implementation (Mamba + differential attention + YOCO shared memory/KV topology). |
| src/mobius/models/phi4flash_test.py | Unit coverage for config extraction, weight routing, cache ABI, attention behavior, and ORT GenAI rejection. |
| src/mobius/models/init.py | Exports Phi4FlashCausalLMModel. |
| src/mobius/integrations/ort_genai/auto_export.py | Explicitly rejects ORT GenAI config generation for phi4flash due to unsupported cache ABI. |
| src/mobius/components/_mamba_block.py | Adds StatefulMambaBlock plus helpers and related activation/memory-mixer utilities used by Phi-4 Flash. |
| src/mobius/components/_mamba_block_test.py | Adds graph-construction tests for StatefulMambaBlock and FloatSwiGLU. |
| src/mobius/components/_differential_attention.py | New portable differential grouped-query attention component with padded/unpadded branching. |
| src/mobius/components/_differential_attention_test.py | New tests covering graph structure and basic CPU execution for differential attention branches. |
| src/mobius/components/init.py | Re-exports DifferentialGQAAttention, StatefulMambaBlock, FloatSwiGLU, and GatedMemoryMixer. |
| src/mobius/_registry.py | Registers phi4flash and the Phi4FlashForCausalLM alias with test model metadata. |
| src/mobius/_configs/_base.py | Adds Phi4FlashConfig with schedule derivation, validation, and HF config extraction. |
| src/mobius/_configs/init.py | Exports Phi4FlashConfig. |
| src/mobius/main.py | Adds CLI/runtime guardrails rejecting static/paged cache modes and ORT GenAI runtime for phi4flash. |
| requirements/ci/phi4flash.txt | New pinned CUDA reference dependency list for the Phi-4 Flash parity workflow. |
| docs/model-catalog.md | Documents phi4flash in the hybrid SSM+attention model catalog. |
| .github/workflows/phi4flash_parity.yml | New scoped self-hosted CUDA workflow to run the pinned-source L3 parity test. |
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # Exact CUDA reference stack for microsoft/Phi-4-mini-flash-reasoning | ||
| # at 1dff8163d28ec880ca2411c474ddc0a927792810. This is deliberately scoped | ||
| # to the Phi-4 Flash parity workflow: it must not constrain public extras. | ||
| torch==2.6.0 |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
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
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
microsoft/Phi-4-mini-flash-reasoning@1dff8163d28ec880ca2411c474ddc0a927792810with source-faithful SambaY Mamba, differential-attention, YOCO shared-memory/KV topology, and dynamic heterogeneous cache ABIValidation
9966 passed, 68 skipped28 passed, 4 skippedLimitations