Skip to content

Commit feceaa4

Browse files
titaiwangmsCopilot
andcommitted
Harden Qwen packed hybrid runtime contract
Verify native packed-state dataflow, operand layouts, and Phase 1 attributes before emitting runtime metadata. Route paged context explicitly, expand malformed-graph coverage, and fix the CI dependency and stale VibeVoice streaming denylist regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
1 parent 8093505 commit feceaa4

10 files changed

Lines changed: 843 additions & 59 deletions

File tree

‎src/mobius/components/_attention.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -593,9 +593,18 @@ def forward(
593593
position_embeddings: tuple,
594594
past_key_value: tuple | None = None,
595595
static_cache: StaticCacheState | None = None,
596+
paged_state: PagedAttentionState | None = None,
596597
):
598+
if paged_state is not None:
599+
if past_key_value is not None or static_cache is not None:
600+
raise ValueError(
601+
"Paged attention state cannot be combined with dense cache state"
602+
)
603+
return self.forward_paged(op, hidden_states, position_embeddings, paged_state)
597604
if isinstance(past_key_value, PagedAttentionState):
598-
return self.forward_paged(op, hidden_states, position_embeddings, past_key_value)
605+
raise TypeError(
606+
"PagedAttentionState must be passed through paged_state, not past_key_value"
607+
)
599608
# Q projection (doubled) → split into Q and gate per head
600609
q_gate = self.q_proj(op, hidden_states)
601610
# Reshape to per-head view so split separates Q/gate within each head

‎src/mobius/components/_attention_test.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
make_config,
1616
)
1717
from mobius.components._attention import Attention, FusedQKVAttention, Qwen35Attention
18+
from mobius.components._paged_attention import PagedAttentionState
1819

1920

2021
class TestAttention:
@@ -215,6 +216,22 @@ def test_forward_builds_graph(self):
215216
assert count_op_type(graph, "Attention") >= 1
216217
assert count_op_type(graph, "Sigmoid") >= 1
217218

219+
def test_legacy_paged_state_fails_fast(self):
220+
config = make_config(partial_rotary_factor=0.5)
221+
attn = Qwen35Attention(config)
222+
builder, op, _ = create_test_builder()
223+
hidden = create_test_input(builder, "hidden", [1, 8, 64])
224+
paged_state = PagedAttentionState(*([None] * 6))
225+
226+
with pytest.raises(TypeError, match="paged_state"):
227+
attn(
228+
op,
229+
hidden,
230+
attention_bias=None,
231+
position_embeddings=(),
232+
past_key_value=paged_state,
233+
)
234+
218235

219236
class TestGQAContextDispatch:
220237
"""Tests for the GQAContext direct GroupQueryAttention emission path."""

‎src/mobius/components/_gated_deltanet.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,11 @@ def forward_paged(
348348
z = self.in_proj_z(op, hidden_states) # (N, Hv*Dv)
349349
raw_b = op.Cast(self.in_proj_b(op, hidden_states), to=ir.DataType.FLOAT)
350350
raw_a = op.Cast(self.in_proj_a(op, hidden_states), to=ir.DataType.FLOAT)
351+
# The generic builder does not infer Cast output types until the
352+
# production optimization pipeline runs. Keep the native ABI explicit
353+
# even for directly-built graphs.
354+
raw_b.type = ir.TensorType(ir.DataType.FLOAT)
355+
raw_a.type = ir.TensorType(ir.DataType.FLOAT)
351356

352357
# Invoke the child module so nn realizes its qualified conv1d.weight.
353358
conv_out, new_conv_state = self.conv1d(

0 commit comments

Comments
 (0)