anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Enhancement] Fix SP path for MoT fused kernels, address review feedback

Browse Source 
Rebase MoT fused kernels (PR #1328) onto current main and fix issues:

- Rewrite _forward_sp_gen() to use MoT unified API instead of deleted
  *_moe_gen layers (qkv_proj_moe_gen, o_proj_moe_gen, q_norm_moe_gen,
  k_norm_moe_gen), which caused AttributeError when SP was active
- Parameterize bias in MoT kernel tests (ZJY0516 review feedback)
- Add configs/ directory with README documenting the 3-tier config
  loading mechanism (ZJY0516 review feedback)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
This commit is contained in:
Zhengyuan SuandClaude Opus 4.6 authored and princepride committed 2026-05-29 02:34:52 +00:00
1 parent b52da5455f
commit e1834ab19b
3 files changed
+100 -57
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/kernels/mot/test_mot_linear.py
+12 -8
View File
@@ -295,12 +295,13 @@ def _run_timing(
# =========================================================================
@pytest.mark.parametrize("bias", [True, False])
@pytest.mark.parametrize(
"image_num, K, N, dtype",
[(num, 3584, 4608, "w16a16_bf16") for num in _IMAGE_NUM],
ids=lambda val: "",
)
def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str):
def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool):
dcfg = _parse_dtype(dtype)
torch.manual_seed(42)
@@ -310,7 +311,7 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str):
head_size=_BAGEL_HEAD_SIZE,
total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
bias=True,
bias=bias,
params_dtype=dcfg.torch_dtype,
disable_tp=True,
).cuda()
@@ -319,7 +320,7 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str):
head_size=_BAGEL_HEAD_SIZE,
total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
bias=True,
bias=bias,
params_dtype=dcfg.torch_dtype,
disable_tp=True,
).cuda()
@@ -328,8 +329,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str):
head_size=_BAGEL_HEAD_SIZE,
total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
bias=True,
vae_bias=True,
bias=bias,
vae_bias=bias,
params_dtype=dcfg.torch_dtype,
disable_tp=True,
).cuda()
@@ -379,6 +380,7 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str):
# =========================================================================
@pytest.mark.parametrize("bias", [True, False])
@pytest.mark.parametrize(
"image_num, K, N, dtype",
[(num, 3584, 3584, "w16a16_bf16") for num in _IMAGE_NUM],
@@ -389,6 +391,7 @@ def test_mot_o_proj(
K: int,
N: int,
dtype: str,
bias: bool,
):
dcfg = _parse_dtype(dtype)
torch.manual_seed(42)
@@ -397,7 +400,7 @@ def test_mot_o_proj(
text_linear = RowParallelLinear(
K,
N,
bias=False,
bias=bias,
input_is_parallel=True,
params_dtype=dcfg.torch_dtype,
disable_tp=True,
@@ -405,7 +408,7 @@ def test_mot_o_proj(
vae_linear = RowParallelLinear(
K,
N,
bias=False,
bias=bias,
input_is_parallel=True,
params_dtype=dcfg.torch_dtype,
disable_tp=True,
@@ -413,7 +416,8 @@ def test_mot_o_proj(
mot_linear = MoTRowParallelLinear(
K,
N,
bias=False,
bias=bias,
vae_bias=bias,
input_is_parallel=True,
params_dtype=dcfg.torch_dtype,
disable_tp=True,