fix(deepseek-v4): per-layer CSA/HCA compressed-record symbolic dim - #622
Draft
justinchuby wants to merge 1 commit into
Draft
fix(deepseek-v4): per-layer CSA/HCA compressed-record symbolic dim#622justinchuby wants to merge 1 commit into
justinchuby wants to merge 1 commit into
Conversation
The native CompressedSparseAttention exporter stamped a single shared symbolic
dim (`past_compressed_records` / `present_compressed_records`) on every CSA
layer's compressed-record axis, on the assumption that "every CSA layer
advances its cache together". That is false for the official interleaved
schedule: a ratio-4 CSA layer pools one record per ~4 tokens while a ratio-128
HCA layer pools one per ~128, so their record counts diverge as the sequence
grows. A shared symbol then forces ORT to bind the same dim to two different
sizes (e.g. 2 and 0 at prefill), which it correctly rejects:
symbol component.model.past_compressed_records bound to conflicting sizes
2 and 0 across bound inputs
Give each layer its own record-axis symbolic dim
(`past_compressed_records.{layer_id}` / `present_compressed_records.{layer_id}`
/ `selected_records.{layer_id}`) via new `CsaLayerPlan` properties. Within a
layer the attention cache and the learned-index cache still share the one
per-layer symbol because they advance in lockstep at the same ratio -- a real
constraint worth expressing -- but layers of different ratios no longer alias.
Extends `test_native_csa_emits_both_ratios_for_interleaved_schedule` to assert
the ratio-4 and ratio-128 layers carry distinct record axes and that the
ratio-4 attention/index caches share their layer's axis; updates the ratio-128
IO test to the per-layer symbol name. No BC shim (dev-time API change).
Surfaced while bringing up the onnx-genai native-decode CSA/HCA E2E proof
against a tiny alternating ratio-4/ratio-128 fixture built by this exporter.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Member
Author
Review verdict — APPROVE (numerics/correctness) · Chew (Code Reviewer, Numerics)Reviewed HEAD What it fixes#593 stamped ONE shared symbolic dim on the record axis of every CSA/HCA layer. On the interleaved ratio-4 + ratio-128 schedule the record counts diverge, so ORT rejects binding one symbol to two sizes ("bound to conflicting sizes 2 and 0"). The PR moves symbol creation from outside → inside the per-layer loop and derives per-layer names Gate results
Bonus: the old code also shared Non-blocking conditions (not defects)
ScopeUnblocks correct export of the mixed ratio-4/ratio-128 DeepSeek-V4-Flash CSA/HCA compressed-attention state (required by onnx-genai #2063 E2E). It does not touch or enable routed top-k MoE — block-FP8 / planar-FP4 expert weights still typed-reject at the runtime-capability gate ( — Chew · numerics/precision review · reproduced on host (no GPU needed) |
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.
Problem
The native
pkg.nxrt::CompressedSparseAttentionexporter (#593) stamped asingle shared symbolic dim (
past_compressed_records/present_compressed_records) on every CSA layer's compressed-record axis,on the assumption that "every CSA layer advances its cache together."
That assumption is false for the official DeepSeek-V4-Flash interleaved
schedule (21 ratio-4 CSA + 20 ratio-128 HCA). A ratio-4 layer pools one
compressed record per ~4 tokens; a ratio-128 layer pools one per ~128. Their
record counts diverge as the sequence grows, so a shared symbol forces ORT to
bind the same dim to two different sizes and it (correctly) rejects the graph:
This makes any exported graph with a mixed ratio schedule unloadable — the
exact schedule the real checkpoint uses.
Fix
Give each layer its own record-axis symbolic dim via new
CsaLayerPlanproperties:
past_compressed_records.{layer_id}present_compressed_records.{layer_id}selected_records.{layer_id}Within a layer, the attention cache and the learned-index cache still share
that one per-layer symbol — they advance in lockstep at the same ratio, which is
a real invariant worth expressing. Across layers of different ratios the
axes no longer alias.
Tests
test_native_csa_emits_both_ratios_for_interleaved_schedule— extended toassert the ratio-4 and ratio-128 layers carry distinct record axes, and
that within the ratio-4 layer the attention/index caches share their layer's
axis.
deepseek_v4_flash_test.py(45 tests) passes;ruff format --checkandruff checkclean on the touched files.No BC shim (dev-time API change, per the no-backward-compat directive).
Downstream
This is the prerequisite for loading a mixed-ratio CSA/HCA graph in onnx-genai's
native decode runtime; the companion onnx-genai E2E proof (stacked on #2063)
regenerates its tiny fixture from this fix.