[AMD] Register 7 JIT kernel unit tests for AMD nightly CI (#28967)
This commit is contained in:
@@ -10,10 +10,11 @@ import pytest
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk
|
from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
|
|
||||||
register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-large")
|
register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-b200")
|
register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-b200")
|
||||||
|
register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)
|
||||||
|
|
||||||
|
|
||||||
def _ref(score, seq_lens, block_size, topk):
|
def _ref(score, seq_lens, block_size, topk):
|
||||||
|
|||||||
@@ -9,10 +9,11 @@ import pytest
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
|
from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
|
|
||||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
|
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
|
||||||
|
register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)
|
||||||
|
|
||||||
dev = "cuda"
|
dev = "cuda"
|
||||||
HEAD_DIM = 128
|
HEAD_DIM = 128
|
||||||
|
|||||||
@@ -98,10 +98,11 @@ from sglang.jit_kernel.tests.deepseek_v4.common import (
|
|||||||
make_legacy_context,
|
make_legacy_context,
|
||||||
to_seq_extend,
|
to_seq_extend,
|
||||||
)
|
)
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
from sglang.utils import is_in_ci
|
from sglang.utils import is_in_ci
|
||||||
|
|
||||||
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
|
register_amd_ci(est_time=25, suite="nightly-amd-kernel-1-gpu", nightly=True)
|
||||||
|
|
||||||
Mode = Literal["decode", "prefill"]
|
Mode = Literal["decode", "prefill"]
|
||||||
ShapeTier = Literal["ci", "smoke", "full"]
|
ShapeTier = Literal["ci", "smoke", "full"]
|
||||||
|
|||||||
@@ -22,7 +22,7 @@ from typing import Optional, Tuple
|
|||||||
import pytest
|
import pytest
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from sglang.jit_kernel.fused_store_index_cache import (
|
from sglang.jit_kernel.fused_store_index_cache import (
|
||||||
@@ -50,6 +50,7 @@ except ImportError:
|
|||||||
|
|
||||||
register_cuda_ci(est_time=24, suite="base-b-kernel-unit-1-gpu-large")
|
register_cuda_ci(est_time=24, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||||
|
register_amd_ci(est_time=24, suite="nightly-amd-kernel-1-gpu", nightly=True)
|
||||||
|
|
||||||
PAGE_SIZE = 64
|
PAGE_SIZE = 64
|
||||||
HEAD_DIM = 128
|
HEAD_DIM = 128
|
||||||
|
|||||||
@@ -11,7 +11,7 @@ import sys
|
|||||||
import pytest
|
import pytest
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
|
from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
|
||||||
@@ -28,6 +28,7 @@ except ImportError:
|
|||||||
|
|
||||||
register_cuda_ci(est_time=6, suite="base-b-kernel-unit-1-gpu-large")
|
register_cuda_ci(est_time=6, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||||
|
register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)
|
||||||
|
|
||||||
|
|
||||||
def _make_tensors(N, T, H, HV, K, V, device="cuda", seed=2025):
|
def _make_tensors(N, T, H, HV, K, V, device="cuda", seed=2025):
|
||||||
|
|||||||
@@ -5,9 +5,10 @@ import pytest
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.mem_cache.memory_pool_host import AsymmetricMHATokenToKVPoolHost
|
from sglang.srt.mem_cache.memory_pool_host import AsymmetricMHATokenToKVPoolHost
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
|
|
||||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
|
register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)
|
||||||
|
|
||||||
# These tests use AsymmetricMHATokenToKVPoolHost methods and let that class call
|
# These tests use AsymmetricMHATokenToKVPoolHost methods and let that class call
|
||||||
# the real sgl-kernel transfer ops.
|
# the real sgl-kernel transfer ops.
|
||||||
|
|||||||
@@ -5,9 +5,10 @@ import torch
|
|||||||
|
|
||||||
from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import mla_kv_pack_quantize_fp8
|
from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import mla_kv_pack_quantize_fp8
|
||||||
from sglang.jit_kernel.utils import get_ci_test_range
|
from sglang.jit_kernel.utils import get_ci_test_range
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||||
|
|
||||||
register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
|
register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
|
||||||
|
register_amd_ci(est_time=60, suite="nightly-amd-kernel-1-gpu", nightly=True)
|
||||||
|
|
||||||
DEVICE = "cuda"
|
DEVICE = "cuda"
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user