[MoE Refactor] Create MK for TRTLLM Kernels (#32564)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Robert Shaw <rshaw@neuralmagic.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <rshaw@neuralmagic.com>
This commit is contained in:
@@ -127,6 +127,14 @@ def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch):
|
||||
)
|
||||
|
||||
|
||||
def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"deepseek-ai/DeepSeek-V3.1",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=triton"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"Known issue: lack of kernel support. "
|
||||
@@ -149,6 +157,14 @@ def test_deepseek_fp8_block_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatc
|
||||
)
|
||||
|
||||
|
||||
def test_deepseek_nvfp4_moe_flashinfer_vllm(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"nvidia/DeepSeek-R1-0528-FP4-v2",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=cutlass"],
|
||||
)
|
||||
|
||||
|
||||
def test_deepseek_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"nvidia/DeepSeek-R1-0528-FP4-v2",
|
||||
@@ -200,3 +216,67 @@ def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
|
||||
## NemoTron ##
|
||||
|
||||
|
||||
def test_nemotron_fp8_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_cutlass"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"FP8 MoE backend FLASHINFER_TRTLLM does not support the "
|
||||
"deployment configuration since kernel does not support "
|
||||
"no act_and_mul MLP layer."
|
||||
)
|
||||
)
|
||||
def test_nemotron_fp8_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"FP8 MoE backend TRITON does not support the "
|
||||
"deployment configuration since kernel does not support "
|
||||
"no act_and_mul MLP layer."
|
||||
)
|
||||
)
|
||||
def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=triton"],
|
||||
)
|
||||
|
||||
|
||||
def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_cutlass"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"FP4 MoE backend FLASHINFER_TRTLLM does not support the "
|
||||
"deployment configuration since kernel does not support "
|
||||
"hidden_dim % 512 != 0."
|
||||
)
|
||||
)
|
||||
def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch):
|
||||
can_initialize(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
|
||||
hf_overrides=HF_OVERRIDE_TEXT,
|
||||
extra_args=["--moe-backend=flashinfer_trtllm"],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user