From 924707a673a7295e0372eaa906f41cb26a6d95e4 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 21:11:04 +0000 Subject: [PATCH] fix: add FFNType/RouterMode to LayerSpec in e2e test --- tests/e2e/test_one_layer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/e2e/test_one_layer.py b/tests/e2e/test_one_layer.py index 3517b4f8..71f0cecb 100644 --- a/tests/e2e/test_one_layer.py +++ b/tests/e2e/test_one_layer.py @@ -25,11 +25,11 @@ def test_swa_cache_gather(): from dsv4.cache.state_cache import StateCachePool from dsv4.cache.handle import LayerCacheHandle from dsv4.model.config import DSV4Config - from dsv4.model.layer_schedule import LayerSpec, AttentionType + from dsv4.model.layer_schedule import LayerSpec, AttentionType, FFNType, RouterMode config = DSV4Config.flash() # Layer 0 of Flash is SWA - spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA) + spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA, ffn=FFNType.MOE, router_mode=RouterMode.HASH) schema = build_schema(config, spec) state = StateCachePool(schema, max_requests=2, device='cuda') @@ -67,14 +67,14 @@ def test_swa_cache_gather(): def test_swa_attention_forward(): """Test SWA attention through the full AttentionSubBlock forward.""" from dsv4.model.config import DSV4Config - from dsv4.model.layer_schedule import LayerSpec, AttentionType + from dsv4.model.layer_schedule import LayerSpec, AttentionType, FFNType, RouterMode from dsv4.layers.attention import AttentionSubBlock from dsv4.cache.schema import build_schema from dsv4.cache.state_cache import StateCachePool from dsv4.cache.handle import LayerCacheHandle config = DSV4Config.flash() - spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA) + spec = LayerSpec(layer_idx=0, attn=AttentionType.SWA, ffn=FFNType.MOE, router_mode=RouterMode.HASH) attn = AttentionSubBlock(config, spec) schema = build_schema(config, spec)