biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: add FFNType/RouterMode to LayerSpec in e2e test

Browse Source
This commit is contained in:
biondizzle
2026-05-30 21:11:04 +00:00
parent e2e21c6350
commit 924707a673
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/e2e/test_one_layer.py
View File

@@ -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)