chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

KDA: fuse qkv conv and support stride for fused_sigmoid_gating_delta_rule_update_kernel (#19506)

Browse Source
This commit is contained in:
strgrb
2026-03-04 22:45:53 +08:00
committed by GitHub
parent 6910c1b281
commit 738ebfd330
4 changed files with 63 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/configs/mamba_utils.py
View File

@@ -217,11 +217,13 @@ class KimiLinearStateShape:
conv_state_k_shape = (divide(proj_k_size, tp_world_size), conv_kernel_size - 1)
temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
conv_state_shape = conv_state_shape[1], conv_state_shape[0]
conv_state_k_shape = conv_state_k_shape[1], conv_state_k_shape[0]
conv_state_shape = (
conv_state_shape[1],
conv_state_shape[0] + conv_state_k_shape[0] * 2,
)
return KimiLinearStateShape(
conv=[conv_state_shape, conv_state_k_shape, conv_state_k_shape],
conv=[conv_state_shape],
temporal=temporal_state_shape,
num_heads=num_heads,
head_dim=head_dim,