[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:
Robert Shaw
2026-03-03 13:39:50 -05:00
committed by GitHub
parent 881a6b011b
commit 97995f6376
77 changed files with 2575 additions and 2087 deletions

View File

@@ -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"],
)