minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

Merge branch 'main' into dsv41-pd

Browse Source
This commit is contained in:
minke.yu
2026-09-22 14:56:35 +08:00
parent c74a4037fb 8ac19cc19f
commit b081dd3d23
795 changed files with 55262 additions and 8438 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject.toml
+2 -2
View File
@@ -96,7 +96,7 @@ dependencies = [
"uvicorn",
"uvloop",
"watchfiles",
"xgrammar==0.2.1",
"xgrammar==0.2.7",
"xxhash",
"zstandard",
]
@@ -186,7 +186,7 @@ test = [
# Pinned: a bump moves scoring for every eval test at once, so re-baseline
# MODEL_SCORE_THRESHOLDS in test/registered/eval/test_text_models_gsm8k_eval.py
# and the mmlu thresholds of run_eval's other callers before changing it.
"sgl-eval==0.1.0",
"sgl-eval==0.1.2",
"sglang[fastokens]",
"tabulate",
]
python/pyproject_cpu.toml
+2 -2
View File
@@ -70,7 +70,7 @@ dependencies = [
"uvicorn",
"uvloop",
"xxhash",
"xgrammar==0.2.1",
"xgrammar==0.2.7",
"zstandard",
]
@@ -110,7 +110,7 @@ test = [
"pandas",
"peft>=0.18.0",
"sentence_transformers",
"sgl-eval==0.1.0",
"sgl-eval==0.1.2",
]
all = []
dev = ["sglang[test]"]
python/pyproject_npu.toml
+2 -2
View File
@@ -67,7 +67,7 @@ dependencies = [
"uvicorn",
"uvloop",
"xxhash",
"xgrammar==0.2.1",
"xgrammar==0.2.7",
]
[project.optional-dependencies]
@@ -106,7 +106,7 @@ test = [
"peft>=0.18.0",
"pytest",
"sentence_transformers",
"sgl-eval==0.1.0",
"sgl-eval==0.1.2",
"tabulate",
]
python/pyproject_other.toml
+2 -2
View File
@@ -73,7 +73,7 @@ runtime_common = [
"compressed-tensors",
"outlines==0.1.11",
"timm==1.0.16",
"xgrammar==0.2.1",
"xgrammar==0.2.7",
]
# srt_empty: device-agnostic install — pure Python packages only, no torch dependency chain.
@@ -208,7 +208,7 @@ test = [
"peft>=0.18.0,<0.19.0", # Pin to <0.19.0 due to torchao incompatibility
"pytest",
"sentence_transformers",
"sgl-eval==0.1.0",
"sgl-eval==0.1.2",
"tabulate",
]
python/pyproject_xpu.toml
+1 -1
View File
@@ -69,7 +69,7 @@ dependencies = [
"uvicorn",
"xxhash",
"uvloop",
# "xgrammar==0.2.1", xgrammar depends on CUDA PyTorch and Triton only
# "xgrammar==0.2.7", xgrammar depends on CUDA PyTorch and Triton only
]
[project.optional-dependencies]
python/sglang/benchmark/one_batch.py
+14 -39
View File
@@ -71,14 +71,13 @@ from sglang.srt.arg_groups.overrides import (
)
from sglang.srt.configs.hybrid_arch import mambaish_config
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import bootstrap
from sglang.srt.distributed.parallel_state import (
destroy_distributed_environment,
destroy_model_parallel,
)
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.entrypoints.engine import _set_envs_and_config
from sglang.srt.hardware_backend.mlx.runtime import use_mlx
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
from sglang.srt.layers.moe import initialize_moe_config
from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
@@ -94,6 +93,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import (
SpawnRanks,
get_device,
get_model,
get_parallel,
get_schedule,
@@ -317,48 +317,23 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
cfg = resolving_view(server_args)
suppress_other_loggers()
rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
moe_ep_rank = tp_rank // (cfg.tp_size // cfg.ep_size)
model_config = ModelConfig.from_server_args(server_args)
attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = (
compute_dp_attention_world_info(
cfg.enable_dp_attention,
tp_rank,
cfg.tp_size,
cfg.dp_size,
cfg.attn_cp_size,
)
)
ps = ParallelState(
tp_rank=tp_rank,
tp_size=cfg.tp_size,
pp_rank=0,
pp_size=1,
dp_rank=None,
dp_size=cfg.dp_size,
attn_tp_rank=attn_tp_rank,
attn_tp_size=attn_tp_size,
attn_cp_rank=0,
attn_cp_size=cfg.attn_cp_size,
attn_dcp_rank=tp_rank % cfg.dcp_size,
attn_dcp_size=cfg.dcp_size,
attn_dp_rank=attn_dp_rank,
attn_dp_size=attn_dp_size,
moe_ep_rank=moe_ep_rank,
moe_ep_size=cfg.ep_size,
moe_dp_rank=None,
moe_dp_size=cfg.moe_dp_size,
gpu_id=gpu_id,
)
runner_kwargs = dict(
model_config=model_config,
mem_fraction_static=cfg.mem_fraction_static,
gpu_id=gpu_id,
ps=ps,
nccl_port=port_args.nccl_port,
server_args=server_args,
)
bootstrap.init_parallel_runtime(
server_args=server_args,
model_config=model_config,
device=get_device().device,
dist_port=port_args.nccl_port,
)
_use_mlx = use_mlx()
if _use_mlx:
from sglang.srt.hardware_backend.mlx.model_runner_stub import (
@@ -571,7 +546,7 @@ def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner):
model_runner=model_runner,
dp_size=get_parallel().dp_size,
attn_tp_size=get_parallel().attn_tp_size,
attn_cp_size=model_runner.ps.attn_cp_size,
attn_cp_size=model_runner.attn_cp_size,
tp_group=model_runner.tp_group,
get_idle_batch=None,
disable_cuda_graph=cuda_graph_fully_disabled(),
@@ -709,13 +684,12 @@ def correctness_test(
gpu_id,
tp_rank,
):
# With the placement this process was spawned with, so a rank read here
# does not need a process group -- the same bundle the runner is handed.
publish(
server_args,
role="scheduler",
ranks=SpawnRanks(
world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=0)
world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=0),
gpu_id=gpu_id,
),
)
@@ -926,7 +900,8 @@ def latency_test(
server_args,
role="scheduler",
ranks=SpawnRanks(
world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=0)
world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=0),
gpu_id=gpu_id,
),
)
initialize_moe_config()
python/sglang/benchmark/one_batch_server.py
+16
View File
@@ -149,6 +149,7 @@ class BenchArgs:
cache_hit_rate: float = 0.0
backend: str = "sglang"
fake_prefill: bool = False
flush_hicache_storage: bool = False
server_args_for_metrics: Optional[List[str]] = None
lora_name: Optional[List[str]] = None
lora_request_distribution: str = "uniform"
@@ -348,6 +349,13 @@ class BenchArgs:
"Use with a decode server running --disaggregation-transfer-backend fake "
"to benchmark pure decode performance without a real prefill node.",
)
parser.add_argument(
"--flush-hicache-storage",
action="store_true",
default=BenchArgs.flush_hicache_storage,
help="Also clear the hierarchical cache's storage tier before each case; "
"/flush_cache resets the radix tree and the host tier only.",
)
parser.add_argument(
"--server-args-for-metrics",
type=str,
@@ -624,12 +632,17 @@ def run_one_case(
lora_zipf_alpha: float = BenchArgs.lora_zipf_alpha,
fixed_prompt_file: str = "",
apply_chat_template: bool = False,
flush_hicache_storage: bool = False,
):
if backend == "vllm":
# You need to have export VLLM_SERVER_DEV_MODE=1 in your environment to use this endpoint.
_flush_cache_with_retry(url, "/reset_prefix_cache")
else:
_flush_cache_with_retry(url, "/flush_cache")
# /flush_cache resets the radix tree and the host tier; a storage tier
# persists across it and would serve the same prompts on the next case.
if flush_hicache_storage:
_flush_cache_with_retry(url, "/hicache/storage-backend/clear")
if fixed_prompt_file:
tok_inner = getattr(tokenizer, "tokenizer", tokenizer)
@@ -1360,6 +1373,7 @@ def run_benchmark_internal(
backend=bench_args.backend,
model_name=model_name,
fake_prefill=bench_args.fake_prefill,
flush_hicache_storage=bench_args.flush_hicache_storage,
lora_name=bench_args.lora_name,
lora_request_distribution=bench_args.lora_request_distribution,
lora_zipf_alpha=bench_args.lora_zipf_alpha,
@@ -1406,6 +1420,7 @@ def run_benchmark_internal(
backend=bench_args.backend,
model_name=model_name,
fake_prefill=bench_args.fake_prefill,
flush_hicache_storage=bench_args.flush_hicache_storage,
lora_name=bench_args.lora_name,
lora_request_distribution=bench_args.lora_request_distribution,
lora_zipf_alpha=bench_args.lora_zipf_alpha,
@@ -1463,6 +1478,7 @@ def run_benchmark_internal(
backend=bench_args.backend,
model_name=model_name,
fake_prefill=bench_args.fake_prefill,
flush_hicache_storage=bench_args.flush_hicache_storage,
lora_name=bench_args.lora_name,
lora_request_distribution=bench_args.lora_request_distribution,
lora_zipf_alpha=bench_args.lora_zipf_alpha,
python/sglang/cli/utils.py
+12 -2
View File
@@ -4,8 +4,6 @@ import os
import subprocess
from functools import lru_cache
from huggingface_hub import HfApi
from sglang.srt.environ import envs
from sglang.utils import (
has_diffusion_overlay_registry_match,
@@ -24,7 +22,17 @@ def _is_overlay_diffusion_model(model_path: str) -> bool:
return has_diffusion_overlay_registry_match(model_path, _load_overlay_registry())
def _diffusion_deps_available() -> bool:
# Locating diffusers is cheap; importing the registry costs ~2 s and then
# fails anyway without it. A false positive is caught by the caller.
import importlib.util
return importlib.util.find_spec("diffusers") is not None
def _is_diffusion_model_from_registry(model_path: str) -> bool:
if not _diffusion_deps_available():
return False
try:
from sglang.multimodal_gen.registry import is_registered_diffusion_model_path
except ImportError:
@@ -49,6 +57,8 @@ def _is_diffusers_model_dir(model_dir: str) -> bool:
def _is_gated_diffusion_repo(repo_id: str) -> bool:
"""Query HF model card metadata to check if a gated repo is a diffusers model."""
try:
from huggingface_hub import HfApi # lazy: ~0.3 s at CLI entry otherwise
info = HfApi().model_info(repo_id)
return getattr(info, "library_name", None) == "diffusers"
except Exception:
python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh
+187 -1
View File
@@ -250,6 +250,103 @@ TOPK_KERNEL void topk_ragged_kernel(const __grid_constant__ TopKRaggedParams par
// PDL trigger secondary at the end the block typically has no use, so ignore it
}
#ifdef USE_ROCM
// Only the ROCm DSA prefill emits this layout today, so CUDA/XPU builds stay
// unchanged. Nothing below is AMD-specific; the guard can be dropped later.
/**
* \brief Parameters of the packed (DSA extend) layout.
*
* Same addressing as the ragged layout -- every row's window lives inside one
* batch-global score buffer starting at `row_starts[i]` -- but the selected
* columns are mapped through a page table before they are written out. Prefill
* expands one request into many query-token rows, so several score rows share a
* page-table row; `row_to_batch[i]` says which one.
*/
struct TopKPackedParams {
// NOTE: may write. The head of the window is masked in place, see the kernel.
float* __restrict__ scores;
const int32_t* __restrict__ seq_lens; // per-row window length
const int32_t* __restrict__ row_starts; // per-row score column offset
const int32_t* __restrict__ row_to_batch; // per-row page-table row; null => identity
const int32_t* __restrict__ page_table;
int32_t* __restrict__ page_indices;
int64_t score_stride;
int64_t page_table_stride;
uint32_t topk;
uint32_t page_bits;
SGL_DEVICE PageTransform get_transform(uint32_t bx) const {
const auto table_row = row_to_batch == nullptr ? bx : static_cast<uint32_t>(row_to_batch[bx]);
return {page_table + static_cast<int64_t>(table_row) * page_table_stride, page_bits, nullptr};
}
};
/**
* \brief Top-k over packed rows, emitting page-table indices.
* \tparam kPDL whether to use PDL to synchronize with the indexer kernel
*
* Dispatch mirrors `topk_ragged_kernel`: both are prefill kernels, so the level
* is picked per row at runtime and only the register and streaming
* implementations are instantiated (no plan, no cluster path).
*/
template <bool kPDL>
TOPK_KERNEL void topk_packed_kernel(const __grid_constant__ TopKPackedParams params) {
device::enable_smem_spilling();
constexpr uint32_t kVecSize = impl::TopKStreaming::kVecSize;
__shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
__shared__ int32_t s_topk_indices[kMaxTopK];
const auto bx = blockIdx.x;
const auto seq_len = static_cast<uint32_t>(params.seq_lens[bx]);
const auto row_start = static_cast<uint32_t>(params.row_starts[bx]);
const auto topk = params.topk;
const auto transform = params.get_transform(bx);
const auto out = params.page_indices + bx * static_cast<int64_t>(topk);
const auto score = params.scores + bx * params.score_stride;
auto problem = TopKProblem{
.in = score + row_start,
.out = out,
.topk = topk,
.seq_len = seq_len,
};
if (seq_len <= topk) {
return trivial_transform<kPDL, TopKMode::PAGE_TABLE>(problem, transform);
}
// Round the window down to a `kVecSize` boundary and mask the <= 3 columns
// that pulls in, with the same bias / input_start as `topk_ragged_kernel`.
const auto rem = row_start % kVecSize;
if (rem != 0) {
// The mask has to land after the indexer has retired
// Otherwise it may be accidentally overwritten by DG upstream
device::PDLWaitPrimary<kPDL>();
static_assert(kVecSize <= kBlockSize, "not enough threads ");
if (const auto tx = threadIdx.x; tx < rem) {
score[row_start - rem + tx] = impl::padding_value();
}
}
using device::topk::broadcast;
problem.in -= rem;
problem.out = s_topk_indices; // write into stage buffer in smem first
problem.seq_len = seq_len + rem;
problem.bias = broadcast(-static_cast<int32_t>(rem));
problem.input_start = broadcast(rem);
if (problem.seq_len <= Register2::kMaxSeqLen) {
Register2::forward<kPDL>(problem, &smem);
} else if (problem.seq_len <= Register4::kMaxSeqLen) {
Register4::forward<kPDL>(problem, &smem);
} else {
Streaming::forward<kPDL>(problem, &smem);
}
device::PDLTriggerSecondary<kPDL>();
__syncthreads();
paged_transform<TopKMode::PAGE_TABLE>(problem, out, transform);
}
#endif // USE_ROCM
/**
* \brief Main kernel for the short items and epilogue of long items.
* \tparam kPDL whether to use PDL to synchronize with the cluster kernel (if any)
@@ -328,7 +425,8 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKPagedParams params
#endif
#ifndef SGL_TOPK_V2_MAX_C16_OCC1
#define SGL_TOPK_V2_MAX_C16_OCC1 7
// Non-portable clusters require a positive device probe.
#define SGL_TOPK_V2_MAX_C16_OCC1 0
#endif
constexpr uint32_t kNumPersistentClusters = SGL_TOPK_V2_MAX_C8_OCC2;
@@ -785,6 +883,94 @@ struct TopKKernel {
.config({.use_pdl = kUsePDL})
.launch(topk_ragged_kernel<kUsePDL>, params);
}
#ifdef USE_ROCM // see the packed kernel above
/**
* \brief Packed (DSA extend prefill) variant of `transform_paged`: per-row
* window inside one batch-global score buffer, page-table output, no plan.
*
* `scores` is written in place exactly like `transform_ragged` does, so rows
* must not overlap and the buffer must have no consumer after this call.
*
* `row_to_batch` absent means the page table is indexed by score row; present,
* it maps each score row onto the table row of the request it belongs to,
* which is what prefill needs (one request expands into many query rows).
*/
static void transform_packed(
const tvm::ffi::TensorView scores,
const tvm::ffi::TensorView seq_lens,
const tvm::ffi::TensorView row_starts,
const tvm::ffi::TensorView page_table,
const tvm::ffi::TensorView page_indices,
const uint32_t page_size,
const tvm::ffi::Optional<tvm::ffi::TensorView> row_to_batch) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto L = SymbolicSize{"max_seq_len"};
auto S = SymbolicSize{"score_stride"};
auto R = SymbolicSize{"page_table_rows"};
auto K = SymbolicSize{"topk"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
TensorMatcher({B, L}) // score
.with_strides({S, 1})
.with_dtype<float>()
.with_device(device_)
.verify(scores);
TensorMatcher({B}) // seq_lens
.with_dtype<int32_t>()
.with_device(device_)
.verify(seq_lens);
TensorMatcher({B}) // row_starts
.with_dtype<int32_t>()
.with_device(device_)
.verify(row_starts);
TensorMatcher({R, -1}) // page_table
.with_strides({-1, 1})
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_table);
TensorMatcher({B, K}) // page_indices
.with_dtype<int32_t>()
.with_device(device_)
.verify(page_indices);
const int32_t* row_to_batch_ptr = nullptr;
if (row_to_batch.has_value()) {
TensorMatcher({B}) // row_to_batch
.with_dtype<int32_t>()
.with_device(device_)
.verify(row_to_batch.value());
row_to_batch_ptr = static_cast<const int32_t*>(row_to_batch.value().data_ptr());
} else {
RuntimeCheck(R.unwrap() == B.unwrap(), "page_table must have one row per score row unless row_to_batch is given");
}
RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (16-byte vectorized load)");
// The kernel masks the head of each window in place, so overlapping rows
// would clobber each other.
RuntimeCheck(S.unwrap() >= L.unwrap(), "scores rows must not overlap");
const auto topk = static_cast<uint32_t>(K.unwrap());
RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]");
const auto params = TopKPackedParams{
.scores = static_cast<float*>(scores.data_ptr()),
.seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
.row_starts = static_cast<const int32_t*>(row_starts.data_ptr()),
.row_to_batch = row_to_batch_ptr,
.page_table = static_cast<const int32_t*>(page_table.data_ptr()),
.page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
.score_stride = S.unwrap(),
.page_table_stride = page_table.stride(0),
.topk = topk,
.page_bits = static_cast<uint32_t>(std::countr_zero(page_size)),
};
LaunchKernel(static_cast<uint32_t>(B.unwrap()), kBlockSize, device_.unwrap())
.config({.use_pdl = kUsePDL})
.launch(topk_packed_kernel<kUsePDL>, params);
}
#endif // USE_ROCM
};
} // namespace sglang
python/sglang/kernels/jit/csrc/dsa/kpool_topk_transform.cuh
+4
View File
@@ -11,7 +11,11 @@
#include <bit>
#include <cstddef>
#include <cstdint>
#ifndef USE_ROCM
#include <cuda_fp16.h>
#else
#include <hip/hip_fp16.h>
#endif
namespace sglang {
namespace {
python/sglang/kernels/jit/csrc/kvcacheio/hicache_tma.cuh
+646
View File
@@ -0,0 +1,646 @@
// HiCache host<->device KV transfer staged through shared memory by the TMA
// bulk-copy engine (sm_90+).
//
// One CTA owns a ring of shared-memory stages. A single loader warp fills
// stages with `cp.async.bulk` (global -> shared, completion counted on an
// mbarrier), which keeps the whole ring in flight with no registers or issue
// slots; the register-staging kernel in hicache.cuh cannot hold enough host
// loads in flight per SM for that. Store warps drain filled stages. Row size is
// a runtime parameter (multiple of 16 B), so one compiled module serves every
// KV shape.
//
// Two hardware facts fix the shape of the kernel (numbers in the PR): every SM
// has a fixed write port to L2, so one CTA cannot exceed it and the block
// quota decides how much of the host link is used; and the TMA unit processes
// bulk ops at a fixed per-op rate, so a source run must move as one op
// (contiguous run -> one 1D bulk copy; strided page run -> one 2D tensor-map
// box), never one op per row. Revisit both if a future part widens the SM
// write port or the TMA op rate.
//
// Work unit ("chunk") = (buffer K|V, layer, run of consecutive positions of the
// index arrays). The loader prefetches the next chunk's indices before blocking
// on the ring so their latency overlaps the wait, and stashes the destination
// indices in smem for the store warps.
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/mbarrier.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/warp.cuh>
#include <dlpack/dlpack.h>
#include <algorithm>
#include <cstdint>
#include <cuda.h>
#include <cudaTypedefs.h>
namespace sglang {
namespace device::ptx {
// global -> shared::cta, completion counted on `bar` (arm with mbar_arrive_expect_tx).
SGL_DEVICE void bulk_g2s(void* dst_smem, const void* src_gmem, uint32_t bytes, uint64_t* bar) {
asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(
to_shared(dst_smem)),
"l"(src_gmem),
"r"(bytes),
"r"(to_shared(bar))
: "memory");
}
// 2D tiled tensor-map box at element coords (x, y) -> shared::cta.
SGL_DEVICE void bulk_tensor_2d_g2s(void* dst_smem, const CUtensorMap* map, int32_t x, int32_t y, uint64_t* bar) {
asm volatile(
"cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::
"r"(to_shared(dst_smem)),
"l"(map),
"r"(x),
"r"(y),
"r"(to_shared(bar))
: "memory");
}
// shared::cta -> global, tracked by the issuing thread's bulk group.
SGL_DEVICE void bulk_s2g(void* dst_gmem, const void* src_smem, uint32_t bytes) {
asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst_gmem),
"r"(to_shared(src_smem)),
"r"(bytes)
: "memory");
}
SGL_DEVICE void bulk_commit_group() {
asm volatile("cp.async.bulk.commit_group;" ::: "memory");
}
// Block until every committed bulk group has finished reading its smem source.
SGL_DEVICE void bulk_wait_group_read_all() {
asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory");
}
// Same, but the most recent group may still be reading.
SGL_DEVICE void bulk_wait_group_read_one() {
asm volatile("cp.async.bulk.wait_group.read 1;" ::: "memory");
}
// Block until every committed bulk group has fully landed in global memory.
SGL_DEVICE void bulk_wait_group_all() {
asm volatile("cp.async.bulk.wait_group 0;" ::: "memory");
}
SGL_DEVICE void fence_mbarrier_init() {
asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}
} // namespace device::ptx
struct HicacheTmaParams {
// Either a direct base pointer (`*_is_table == false`) or a device array of
// `num_layers` uint64 base pointers. `v_*` is unused when `has_v == false` (MLA).
const void* __restrict__ k_src;
const void* __restrict__ v_src;
void* __restrict__ k_dst;
void* __restrict__ v_dst;
const void* __restrict__ indices_src;
const void* __restrict__ indices_dst;
int64_t src_stride; // bytes between consecutive token rows
int64_t dst_stride;
uint32_t row_bytes; // bytes copied per token row, multiple of 16
uint32_t length; // number of token indices
uint32_t num_layers;
uint64_t units_per_row_magic; // ceil(2^32 / (row_bytes / 16)); see store loop
bool src_is_table;
bool dst_is_table;
bool has_v;
// Strided source rows ([K, V] views): one box per chunk instead of one op per row.
bool has_src_map;
CUtensorMap src_map[2];
};
// Rows of one chunk are spread over the loader lanes; each lane prefetches at
// most this many row indices, which bounds rows per chunk to 32x this.
inline constexpr uint32_t kHicacheTmaRowsPerLane = 4;
inline constexpr uint32_t kHicacheTmaMaxRows = kHicacheTmaRowsPerLane * device::kWarpThreads;
// Tensor-map boxes are limited to 256 elements per dimension; rows are mapped as
// 8-byte elements so this is the widest row a 2D box can cover.
inline constexpr uint32_t kHicacheTmaMaxMapRowBytes = 256 * 8;
// Rows per chunk: largest power of two that fits the stage, so chunks never
// straddle a (power-of-two) page and a page run stays one bulk op.
__host__ __device__ constexpr uint32_t hicache_tma_rows_per_chunk(uint32_t stage_bytes, uint32_t row_bytes) {
uint32_t rows = 1;
while (rows * 2 <= stage_bytes / row_bytes && rows * 2 <= kHicacheTmaMaxRows)
rows *= 2;
return rows;
}
template <uint32_t kStageBytes, uint32_t kNumStages>
struct HicacheTmaSmem {
alignas(128) uint8_t stages[kNumStages][kStageBytes];
int64_t dst_idx[kNumStages][kHicacheTmaMaxRows]; // destination row indices of the staged chunk
uint32_t dst_run[kNumStages]; // destination rows form one contiguous span
uint64_t full[kNumStages]; // loader -> storers: stage filled (count 1)
uint64_t empty[kNumStages]; // storers -> loader: stage drained (count kStoreWarps)
};
template <typename T, uint32_t kStageBytes, uint32_t kNumStages, uint32_t kStoreWarps>
__global__ void __launch_bounds__((1 + kStoreWarps) * device::kWarpThreads, 1)
hicache_tma_transfer_kernel(const __grid_constant__ HicacheTmaParams p) {
#if SGL_ARCH_HOPPER_OR_GREATER
using namespace device;
using Smem = HicacheTmaSmem<kStageBytes, kNumStages>;
extern __shared__ __align__(128) uint8_t smem_raw[];
auto& smem = *reinterpret_cast<Smem*>(smem_raw);
const uint32_t warp = threadIdx.x / kWarpThreads;
const uint32_t lane = threadIdx.x % kWarpThreads;
if (threadIdx.x == 0) {
for (uint32_t s = 0; s < kNumStages; ++s) {
ptx::mbar_init(&smem.full[s], 1);
ptx::mbar_init(&smem.empty[s], kStoreWarps);
}
ptx::fence_mbarrier_init();
}
__syncthreads();
const uint32_t rows_per_chunk = hicache_tma_rows_per_chunk(kStageBytes, p.row_bytes);
const uint32_t token_chunks = div_ceil(p.length, rows_per_chunk);
const uint32_t num_chunks = (p.has_v ? 2u : 1u) * p.num_layers * token_chunks;
const T* idx_src = static_cast<const T*>(p.indices_src);
const T* idx_dst = static_cast<const T*>(p.indices_dst);
struct ChunkInfo {
uint32_t t0; // first index position
uint32_t rows; // rows in this chunk
uint32_t layer;
bool is_v;
};
auto describe = [&](uint32_t chunk) {
const uint32_t tc = chunk % token_chunks;
const uint32_t rest = chunk / token_chunks;
const uint32_t t0 = tc * rows_per_chunk;
return ChunkInfo{t0, min(rows_per_chunk, p.length - t0), rest % p.num_layers, (rest / p.num_layers) != 0};
};
auto base_ptr = [&](const void* direct_or_table, bool is_table, uint32_t layer) -> const void* {
return is_table ? reinterpret_cast<const void*>(static_cast<const uint64_t*>(direct_or_table)[layer])
: direct_or_table;
};
if (warp == 0) {
// ---- loader: global -> smem ring via TMA bulk copies
struct Prefetch {
const void* src_base;
T src[kHicacheTmaRowsPerLane]; // rows lane, lane + 32, ...
T dst[kHicacheTmaRowsPerLane];
};
auto prefetch = [&](uint32_t chunk) {
const ChunkInfo c = describe(chunk);
Prefetch pf;
pf.src_base = base_ptr(c.is_v ? p.v_src : p.k_src, p.src_is_table, c.layer);
#pragma unroll
for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) {
const uint32_t r = lane + k * kWarpThreads;
pf.src[k] = r < c.rows ? idx_src[c.t0 + r] : T{0};
pf.dst[k] = r < c.rows ? idx_dst[c.t0 + r] : T{0};
}
return pf;
};
uint32_t chunk = blockIdx.x;
Prefetch next = chunk < num_chunks ? prefetch(chunk) : Prefetch{};
for (uint32_t it = 0; chunk < num_chunks; chunk += gridDim.x, ++it) {
const ChunkInfo c = describe(chunk);
const Prefetch cur = next;
if (chunk + gridDim.x < num_chunks) next = prefetch(chunk + gridDim.x);
// A run: every row sits at first + r (whole pages in order), on either side.
const T first = __shfl_sync(warp::kFullMask, cur.src[0], 0);
const T first_dst = __shfl_sync(warp::kFullMask, cur.dst[0], 0);
bool run = true, run_dst = true;
#pragma unroll
for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) {
const uint32_t r = lane + k * kWarpThreads;
run &= r >= c.rows || cur.src[k] == first + static_cast<T>(r);
run_dst &= r >= c.rows || cur.dst[k] == first_dst + static_cast<T>(r);
}
run = __all_sync(warp::kFullMask, run);
run_dst = __all_sync(warp::kFullMask, run_dst);
const bool contiguous = run && p.src_stride == p.row_bytes;
const bool boxed = run && !contiguous && p.has_src_map;
const uint32_t s = it % kNumStages;
ptx::mbar_wait_parity(&smem.empty[s], ((it / kNumStages) & 1) ^ 1); // fresh ring passes
#pragma unroll
for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) {
const uint32_t r = lane + k * kWarpThreads;
if (r < c.rows) smem.dst_idx[s][r] = static_cast<int64_t>(cur.dst[k]);
}
if (lane == 0) smem.dst_run[s] = run_dst && p.dst_stride == p.row_bytes;
// A box always lands rows_per_chunk rows (out-of-range rows are zero-filled).
const uint32_t tx_bytes = (boxed ? rows_per_chunk : c.rows) * p.row_bytes;
if (lane == 0) ptx::mbar_arrive_expect_tx(&smem.full[s], tx_bytes);
__syncwarp();
uint8_t* stage = smem.stages[s];
if (contiguous) {
if (lane == 0) {
ptx::bulk_g2s(
stage,
pointer::offset(cur.src_base, static_cast<int64_t>(first) * p.src_stride),
c.rows * p.row_bytes,
&smem.full[s]);
}
} else if (boxed) {
if (lane == 0) {
ptx::bulk_tensor_2d_g2s(stage, &p.src_map[c.is_v ? 1 : 0], 0, static_cast<int32_t>(first), &smem.full[s]);
}
} else {
#pragma unroll
for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) {
const uint32_t r = lane + k * kWarpThreads;
if (r < c.rows) {
ptx::bulk_g2s(
stage + r * p.row_bytes,
pointer::offset(cur.src_base, static_cast<int64_t>(cur.src[k]) * p.src_stride),
p.row_bytes,
&smem.full[s]);
}
}
}
}
} else {
// ---- storers: smem ring -> global, 16-byte units interleaved across all
// store threads. Row addressing comes from the loader-staged dst_idx.
constexpr uint32_t kUnroll = 4;
constexpr uint32_t kStoreThreads = kStoreWarps * kWarpThreads;
constexpr uint32_t kUnitsPerIter = kStoreThreads * kUnroll;
const uint32_t tid = threadIdx.x - kWarpThreads;
const uint32_t units_per_row = p.row_bytes / 16;
// unit -> (row, col) without a hardware divide: magic multiply plus a one-step fixup.
auto locate = [&](uint32_t u, uint32_t& row, uint32_t& col) {
row = static_cast<uint32_t>((static_cast<uint64_t>(u) * p.units_per_row_magic) >> 32);
int32_t rem = static_cast<int32_t>(u - row * units_per_row);
if (rem < 0) {
--row;
rem += units_per_row;
}
col = static_cast<uint32_t>(rem);
};
auto unit_dst = [&](void* dst_base, const int64_t* dst_idx, uint32_t u) -> uint4* {
uint32_t row, col;
locate(u, row, col);
return static_cast<uint4*>(pointer::offset(dst_base, dst_idx[row] * p.dst_stride, col * 16));
};
constexpr uint32_t kNoStage = ~0u;
uint32_t bulk_pending = kNoStage; // stage of the last bulk store not yet released (warp 0)
uint32_t it = 0;
for (uint32_t chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x, ++it) {
const ChunkInfo c = describe(chunk);
void* dst_base = const_cast<void*>(base_ptr(c.is_v ? p.v_dst : p.k_dst, p.dst_is_table, c.layer));
const uint32_t s = it % kNumStages;
const uint32_t n_units = c.rows * units_per_row;
const uint32_t n_full = n_units - n_units % kUnitsPerIter;
const uint4* stage = reinterpret_cast<const uint4*>(smem.stages[s]);
const int64_t* dst_idx = smem.dst_idx[s];
ptx::mbar_wait_parity(&smem.full[s], (it / kNumStages) & 1);
if (smem.dst_run[s]) {
// Contiguous span: one bulk store from store warp 0 runs at the SM's write
// port; it releases the previous bulk stage once that stage's smem read is
// done, keeping two stores in flight. The other store warps have nothing
// to read and release the stage right away.
if (tid < kWarpThreads) {
if (lane == 0) {
ptx::bulk_s2g(pointer::offset(dst_base, dst_idx[0] * p.dst_stride), stage, n_units * 16);
ptx::bulk_commit_group();
if (bulk_pending != kNoStage) {
ptx::bulk_wait_group_read_one();
ptx::mbar_arrive(&smem.empty[bulk_pending]);
}
}
bulk_pending = s;
} else if (lane == 0) {
ptx::mbar_arrive(&smem.empty[s]);
}
continue;
}
if (tid < kWarpThreads && bulk_pending != kNoStage) {
if (lane == 0) {
ptx::bulk_wait_group_read_all();
ptx::mbar_arrive(&smem.empty[bulk_pending]);
}
bulk_pending = kNoStage;
}
for (uint32_t u0 = tid; u0 < n_full; u0 += kUnitsPerIter) {
uint4 v[kUnroll];
uint4* dst[kUnroll];
#pragma unroll
for (uint32_t k = 0; k < kUnroll; ++k) {
v[k] = stage[u0 + k * kStoreThreads];
dst[k] = unit_dst(dst_base, dst_idx, u0 + k * kStoreThreads);
}
#pragma unroll
for (uint32_t k = 0; k < kUnroll; ++k)
__stcs(dst[k], v[k]);
}
for (uint32_t u = n_full + tid; u < n_units; u += kStoreThreads) {
__stcs(unit_dst(dst_base, dst_idx, u), stage[u]);
}
__syncwarp();
if (lane == 0) ptx::mbar_arrive(&smem.empty[s]);
}
if (tid == 0) ptx::bulk_wait_group_all(); // bulk stores must land before the grid completes
(void)bulk_pending; // the final stage is never reused
}
#endif
}
template <uint32_t kStageBytes, uint32_t kNumStages, uint32_t kStoreWarps, uint32_t kBlockQuota>
struct HiCacheTmaKernel {
using Smem = HicacheTmaSmem<kStageBytes, kNumStages>;
static_assert(kStageBytes % 128 == 0, "stage must stay 128-byte aligned for bulk copies");
static_assert(kNumStages >= 3 && kStoreWarps >= 1, "two bulk stores in flight plus one loading stage");
static constexpr uint32_t kThreads = (1 + kStoreWarps) * device::kWarpThreads;
template <typename T>
static constexpr auto kernel = hicache_tma_transfer_kernel<T, kStageBytes, kNumStages, kStoreWarps>;
static uint32_t rows_per_chunk(uint32_t row_bytes) {
return hicache_tma_rows_per_chunk(kStageBytes, row_bytes);
}
// Whether the device can hold the smem ring in one CTA; sm_90+ parts with
// small opt-in shared memory (consumer Blackwell) must keep the register kernel.
static bool fits_device(int64_t device_id) {
int max_smem = 0;
host::RuntimeDeviceCheck(
cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast<int>(device_id)));
return static_cast<std::size_t>(max_smem) >= sizeof(Smem);
}
// ceil(2^32 / units_per_row): (u * magic) >> 32 overestimates u / units_per_row
// by at most one for the unit counts a stage can hold; the kernel fixes that up.
static uint64_t units_per_row_magic(uint32_t row_bytes) {
const uint64_t upr = row_bytes / 16;
return ((uint64_t{1} << 32) + upr - 1) / upr;
}
static auto encode_tiled_fn() -> PFN_cuTensorMapEncodeTiled_v12000 {
static const auto fn = [] {
void* sym = nullptr;
cudaDriverEntryPointQueryResult status;
host::RuntimeDeviceCheck(
cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &sym, 12000, cudaEnableDefault, &status));
host::RuntimeCheck(status == cudaDriverEntryPointSuccess && sym != nullptr, "cuTensorMapEncodeTiled unavailable");
return reinterpret_cast<PFN_cuTensorMapEncodeTiled_v12000>(sym);
}();
return fn;
}
// 2D view of a strided row buffer: [rows][row_bytes / 8] uint64 elements with
// pitch `stride_bytes`; one box covers rows_per_chunk consecutive rows.
static void encode_src_map(CUtensorMap* map, const void* base, int64_t num_rows, uint32_t row_bytes, int64_t stride) {
const cuuint64_t gdim[2] = {row_bytes / 8, static_cast<cuuint64_t>(num_rows)};
const cuuint64_t gstride[1] = {static_cast<cuuint64_t>(stride)};
const cuuint32_t box[2] = {row_bytes / 8, rows_per_chunk(row_bytes)};
const cuuint32_t estride[2] = {1, 1};
const CUresult res = encode_tiled_fn()(
map,
CU_TENSOR_MAP_DATA_TYPE_UINT64,
2,
const_cast<void*>(base),
gdim,
gstride,
box,
estride,
CU_TENSOR_MAP_INTERLEAVE_NONE,
CU_TENSOR_MAP_SWIZZLE_NONE,
CU_TENSOR_MAP_L2_PROMOTION_NONE,
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
host::RuntimeCheck(res == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed: ", static_cast<int>(res));
}
static void launch(const HicacheTmaParams& params, bool use_int32, DLDevice device) {
using namespace host;
RuntimeCheck(params.row_bytes > 0 && params.row_bytes % 16 == 0, "HiCache TMA: row bytes must be a multiple of 16");
RuntimeCheck(params.row_bytes <= kStageBytes, "HiCache TMA: row bytes exceed the smem stage");
RuntimeCheck(
params.src_stride % 16 == 0 && params.dst_stride % 16 == 0, "HiCache TMA: strides must be multiples of 16");
if (params.length == 0 || params.num_layers == 0) return;
const uint32_t chunks =
(params.has_v ? 2u : 1u) * params.num_layers * div_ceil(params.length, rows_per_chunk(params.row_bytes));
constexpr std::size_t kSmemBytes = sizeof(Smem);
static const bool attr_set = [] {
for (auto fn : {kernel<int32_t>, kernel<int64_t>}) {
RuntimeDeviceCheck(
cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kSmemBytes)));
}
return true;
}();
(void)attr_set;
LaunchKernel(std::min(chunks, kBlockQuota), kThreads, device, kSmemBytes)(
use_int32 ? kernel<int32_t> : kernel<int64_t>, params);
}
// Cache operand viewed as [-1, D] rows; binds row dim, stride and dtype.
static void verify_cache(
const tvm::ffi::TensorView& t, host::SymbolicSize& D, host::SymbolicSize& stride, host::SymbolicDType& dtype) {
using namespace host;
TensorMatcher({-1, D}) //
.with_strides({stride, 1})
.with_dtype(dtype)
.with_device<kDLGPU, kDLGPUHost, kDLCPU>()
.verify(t);
}
static void verify_indices(
const tvm::ffi::TensorView& a,
const tvm::ffi::TensorView& b,
host::SymbolicSize& L,
host::SymbolicDType& dtype,
host::SymbolicDevice& device) {
using namespace host;
TensorMatcher({L}) //
.with_dtype<int32_t, int64_t>(dtype)
.with_device<kDLGPU>(device)
.verify(a)
.verify(b);
}
// One layer, direct pointers. `v_*` are ignored when `has_v == false`.
static void run_one_impl(
const tvm::ffi::TensorView k_cache_dst,
const tvm::ffi::TensorView v_cache_dst,
const tvm::ffi::TensorView indices_dst,
const tvm::ffi::TensorView k_cache_src,
const tvm::ffi::TensorView v_cache_src,
const tvm::ffi::TensorView indices_src,
bool has_v) {
using namespace host;
auto D = SymbolicSize{"row dim"};
auto N = SymbolicSize{"src stride"};
auto M = SymbolicSize{"dst stride"};
auto L = SymbolicSize{"indices length"};
auto cache_dtype = SymbolicDType{};
auto indices_dtype = SymbolicDType{};
auto indices_device = SymbolicDevice{};
verify_cache(k_cache_src, D, N, cache_dtype);
verify_cache(k_cache_dst, D, M, cache_dtype);
if (has_v) {
verify_cache(v_cache_src, D, N, cache_dtype);
verify_cache(v_cache_dst, D, M, cache_dtype);
}
verify_indices(indices_src, indices_dst, L, indices_dtype, indices_device);
const auto dtype_size = dtype_bytes(cache_dtype.unwrap());
const auto row_bytes = static_cast<uint32_t>(D.unwrap() * dtype_size);
const auto src_stride = static_cast<int64_t>(N.unwrap() * dtype_size);
HicacheTmaParams params{
.k_src = k_cache_src.data_ptr(),
.v_src = has_v ? v_cache_src.data_ptr() : nullptr,
.k_dst = k_cache_dst.data_ptr(),
.v_dst = has_v ? v_cache_dst.data_ptr() : nullptr,
.indices_src = indices_src.data_ptr(),
.indices_dst = indices_dst.data_ptr(),
.src_stride = src_stride,
.dst_stride = static_cast<int64_t>(M.unwrap() * dtype_size),
.row_bytes = row_bytes,
.length = static_cast<uint32_t>(L.unwrap()),
.num_layers = 1,
.units_per_row_magic = units_per_row_magic(row_bytes),
.src_is_table = false,
.dst_is_table = false,
.has_v = has_v,
.has_src_map = false,
};
if (src_stride != row_bytes && row_bytes <= kHicacheTmaMaxMapRowBytes) {
params.has_src_map = true;
encode_src_map(&params.src_map[0], params.k_src, k_cache_src.shape()[0], row_bytes, src_stride);
if (has_v) encode_src_map(&params.src_map[1], params.v_src, v_cache_src.shape()[0], row_bytes, src_stride);
}
launch(params, indices_dtype.unwrap().bits == 32, indices_device.unwrap());
}
// All layers through device-side pointer tables; strides and row bytes explicit.
static void run_all_impl(
const tvm::ffi::TensorView k_ptr_dst,
const tvm::ffi::TensorView v_ptr_dst,
const tvm::ffi::TensorView indices_dst,
const tvm::ffi::TensorView k_ptr_src,
const tvm::ffi::TensorView v_ptr_src,
const tvm::ffi::TensorView indices_src,
int64_t src_stride_bytes,
int64_t dst_stride_bytes,
int64_t row_bytes,
bool has_v) {
using namespace host;
auto N = SymbolicSize{"num_layers"};
auto L = SymbolicSize{"indices length"};
auto indices_dtype = SymbolicDType{};
auto device_ = SymbolicDevice{};
auto verify_table = [&](const tvm::ffi::TensorView& t) {
TensorMatcher({N}).with_dtype<uint64_t>().with_device<kDLGPU>(device_).verify(t);
};
verify_table(k_ptr_src);
verify_table(k_ptr_dst);
if (has_v) {
verify_table(v_ptr_src);
verify_table(v_ptr_dst);
}
verify_indices(indices_src, indices_dst, L, indices_dtype, device_);
const HicacheTmaParams params{
.k_src = k_ptr_src.data_ptr(),
.v_src = has_v ? v_ptr_src.data_ptr() : nullptr,
.k_dst = k_ptr_dst.data_ptr(),
.v_dst = has_v ? v_ptr_dst.data_ptr() : nullptr,
.indices_src = indices_src.data_ptr(),
.indices_dst = indices_dst.data_ptr(),
.src_stride = src_stride_bytes,
.dst_stride = dst_stride_bytes,
.row_bytes = static_cast<uint32_t>(row_bytes),
.length = static_cast<uint32_t>(L.unwrap()),
.num_layers = static_cast<uint32_t>(N.unwrap()),
.units_per_row_magic = units_per_row_magic(static_cast<uint32_t>(row_bytes)),
.src_is_table = true,
.dst_is_table = true,
.has_v = has_v,
.has_src_map = false,
};
launch(params, indices_dtype.unwrap().bits == 32, device_.unwrap());
}
static void run_one(
const tvm::ffi::TensorView k_cache_dst,
const tvm::ffi::TensorView v_cache_dst,
const tvm::ffi::TensorView indices_dst,
const tvm::ffi::TensorView k_cache_src,
const tvm::ffi::TensorView v_cache_src,
const tvm::ffi::TensorView indices_src) {
run_one_impl(k_cache_dst, v_cache_dst, indices_dst, k_cache_src, v_cache_src, indices_src, true);
}
static void run_one_mla(
const tvm::ffi::TensorView cache_dst,
const tvm::ffi::TensorView indices_dst,
const tvm::ffi::TensorView cache_src,
const tvm::ffi::TensorView indices_src) {
run_one_impl(cache_dst, cache_dst, indices_dst, cache_src, cache_src, indices_src, false);
}
static void run_all(
const tvm::ffi::TensorView k_ptr_dst,
const tvm::ffi::TensorView v_ptr_dst,
const tvm::ffi::TensorView indices_dst,
const tvm::ffi::TensorView k_ptr_src,
const tvm::ffi::TensorView v_ptr_src,
const tvm::ffi::TensorView indices_src,
const int64_t src_stride_bytes,
const int64_t dst_stride_bytes,
const int64_t row_bytes) {
run_all_impl(
k_ptr_dst,
v_ptr_dst,
indices_dst,
k_ptr_src,
v_ptr_src,
indices_src,
src_stride_bytes,
dst_stride_bytes,
row_bytes,
true);
}
static void run_all_mla(
const tvm::ffi::TensorView ptr_dst,
const tvm::ffi::TensorView indices_dst,
const tvm::ffi::TensorView ptr_src,
const tvm::ffi::TensorView indices_src,
const int64_t src_stride_bytes,
const int64_t dst_stride_bytes,
const int64_t row_bytes) {
run_all_impl(
ptr_dst,
ptr_dst,
indices_dst,
ptr_src,
ptr_src,
indices_src,
src_stride_bytes,
dst_stride_bytes,
row_bytes,
false);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh
+73 -30
View File
@@ -303,13 +303,26 @@ template <int NUM_TOP_K, int HOT_BUFFER_SIZE>
struct SmemLayout {
static constexpr int HASH_SIZE = NUM_TOP_K * 2;
static constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + WARP_SIZE - 1) / WARP_SIZE;
// int32_t region: top_k_tokens + chunk_offset + evict_chunk_offset + hash_keys + total_hits + newest_hit
// int32_t region: top_k_tokens + chunk offsets + hash keys + hit counters
static constexpr int TOTAL_INT32 = NUM_TOP_K + (NUM_BUFFER_CHUNKS + 1) + (NUM_BUFFER_CHUNKS + 1) + HASH_SIZE + 2;
// int16_t region: lru_slots_out + hash_vals
static constexpr int TOTAL_INT16 = HOT_BUFFER_SIZE + HASH_SIZE;
static constexpr size_t BYTES = TOTAL_INT32 * sizeof(int32_t) + TOTAL_INT16 * sizeof(int16_t);
};
template <int SPARSE_BLOCK_SIZE, bool TopKIsBlocks>
__device__ __forceinline__ int32_t resolve_selected_token(const int32_t* top_k, int32_t token_index) {
if constexpr (TopKIsBlocks) {
const int32_t block_index = top_k[token_index / SPARSE_BLOCK_SIZE];
if (block_index < 0) {
return -1;
}
return block_index * SPARSE_BLOCK_SIZE + token_index % SPARSE_BLOCK_SIZE;
} else {
return top_k[token_index];
}
}
// Each block processes one request
// req_pool_indices and seq_lens can each be int32_t or int64_t
// Layout: [HOT_BUFFER_SIZE slots for LRU] + [page_size slots for newest token]
@@ -319,23 +332,28 @@ struct SmemLayout {
// false -> generic byte-stride: device + host both linear, stride = item_size_bytes
// true -> DSv4 page-padded device + page-padded host (kvcacheio.cuh constants)
//
// TopKIsBlocks makes the kernel consume block ids directly. It resolves token
// positions in registers and writes the flattened token-slot table expected by
// sparse attention without materializing an intermediate token-index tensor.
// RecordMissPlan records this step's miss plan (miss_src/dst = host/device loc
// per miss, miss_count per request) for shared-index skip layers to replay via
// copy_cache_planned_kernel. SkipIO elides only the KV byte movement (timing
// probe; output is garbage). Both are compile-time flags so the production
// (false, false) instantiation stays byte-identical.
// probe; output is garbage). These are compile-time flags, so inactive paths
// are removed from each specialization.
template <
int BLOCK_SIZE,
int NUM_TOP_K,
int HOT_BUFFER_SIZE,
bool IsMLA,
bool IsDsv4Layout,
int SPARSE_BLOCK_SIZE,
bool TopKIsBlocks,
bool RecordMissPlan,
bool SkipIO,
typename SeqLensT,
typename ReqPoolIndicesT>
__global__ void load_cache_to_device_buffer_kernel(
const int32_t* __restrict__ top_k_tokens,
const int32_t* __restrict__ top_k,
int32_t* __restrict__ device_buffer_tokens,
const int64_t* __restrict__ host_cache_locs,
const int32_t* __restrict__ device_buffer_locs,
@@ -351,7 +369,7 @@ __global__ void load_cache_to_device_buffer_kernel(
int64_t buffer_stride_0,
int64_t host_stride,
int64_t lru_slot_stride_0,
int64_t top_k_tokens_stride,
int64_t top_k_stride,
int64_t top_k_device_locs_stride,
int64_t page_size,
int64_t item_size_bytes,
@@ -360,9 +378,12 @@ __global__ void load_cache_to_device_buffer_kernel(
int32_t* __restrict__ miss_count_out,
int64_t plan_stride) {
static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA).");
// todo hisparse: support page wise sparsity
static_assert(SPARSE_BLOCK_SIZE > 0, "SPARSE_BLOCK_SIZE must be positive.");
// Cache residency and LRU replacement remain token-granular even when the
// sparse-attention selection arrives as block ids.
constexpr int NUM_TOP_K_TOKENS = NUM_TOP_K * (TopKIsBlocks ? SPARSE_BLOCK_SIZE : 1);
constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K + WARP_SIZE - 1) / WARP_SIZE;
constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K_TOKENS + WARP_SIZE - 1) / WARP_SIZE;
constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + WARP_SIZE - 1) / WARP_SIZE;
const int bid = blockIdx.x;
@@ -372,7 +393,7 @@ __global__ void load_cache_to_device_buffer_kernel(
// CUDA graph pads the batch to a captured size. Keep padded output rows
// invalid without a separate fill kernel.
if (bid >= num_real_reqs[0]) {
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
req_top_k_device_locs[i] = -1;
}
return;
@@ -386,7 +407,7 @@ __global__ void load_cache_to_device_buffer_kernel(
const int64_t seq_len = seq_lens[bid];
// Calculate offsets for this request
const int32_t* req_top_k_tokens = top_k_tokens + bid * top_k_tokens_stride;
const int32_t* req_top_k = top_k + bid * top_k_stride;
const int64_t buffer_offset = rid * buffer_stride_0;
int32_t* req_device_buffer_tokens = device_buffer_tokens + buffer_offset;
@@ -396,14 +417,16 @@ __global__ void load_cache_to_device_buffer_kernel(
// Fast path: short sequences have all tokens in the device buffer in order.
if (seq_len <= HOT_BUFFER_SIZE) {
const int count = (seq_len < NUM_TOP_K) ? static_cast<int>(seq_len) : NUM_TOP_K;
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
const int count = (seq_len < NUM_TOP_K_TOKENS) ? static_cast<int>(seq_len) : NUM_TOP_K_TOKENS;
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
int32_t device_loc = -1;
if (i < count) {
int32_t token_pos = req_top_k_tokens[i];
if (token_pos >= 0) {
const int32_t token_pos = resolve_selected_token<SPARSE_BLOCK_SIZE, TopKIsBlocks>(req_top_k, i);
if constexpr (TopKIsBlocks) {
if (token_pos >= 0 && token_pos < seq_len) {
device_loc = req_device_buffer_locs[token_pos];
}
} else if (i < count && token_pos >= 0) {
device_loc = req_device_buffer_locs[token_pos];
}
req_top_k_device_locs[i] = device_loc;
}
@@ -418,21 +441,21 @@ __global__ void load_cache_to_device_buffer_kernel(
// Dynamic shared memory layout: int32_t arrays first, then int16_t arrays.
extern __shared__ char smem_raw[];
using Layout = SmemLayout<NUM_TOP_K, HOT_BUFFER_SIZE>;
using Layout = SmemLayout<NUM_TOP_K_TOKENS, HOT_BUFFER_SIZE>;
constexpr int HASH_SIZE = Layout::HASH_SIZE;
int32_t* smem_i32 = reinterpret_cast<int32_t*>(smem_raw);
// Top-k token positions; reused as miss-token scratch in the copy phase
int32_t* s_top_k_tokens = smem_i32;
// Prefix-sum offsets for hit counting and miss counting
int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K;
int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K_TOKENS;
// Prefix-sum offsets for evictable counting
int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1);
// Open-addressing hash table: top-k token_id -> top-k index (keys)
int32_t* s_hash_keys = s_evict_chunk_offset + (NUM_BUFFER_CHUNKS + 1);
// Scalar counters
int32_t& s_total_hits = s_hash_keys[HASH_SIZE];
int32_t& s_newest_hit = s_hash_keys[HASH_SIZE + 1];
int32_t& s_total_misses = s_hash_keys[HASH_SIZE + 1];
int16_t* smem_i16 = reinterpret_cast<int16_t*>(smem_i32 + Layout::TOTAL_INT32);
// Compacted slot ordering: [hits fwd-> ... <-evictables bwd]
@@ -443,7 +466,7 @@ __global__ void load_cache_to_device_buffer_kernel(
// Initialize shared memory: counters, hash table, prefix-sum offsets.
if (tid == 0) {
s_total_hits = 0;
s_newest_hit = 0;
s_total_misses = 0;
}
for (int i = tid; i < HASH_SIZE; i += BLOCK_SIZE) {
s_hash_keys[i] = HASH_EMPTY;
@@ -458,14 +481,20 @@ __global__ void load_cache_to_device_buffer_kernel(
const int32_t newest_token = seq_len - 1;
// Insert top-k tokens into shared-memory hash table.
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
int32_t token_idx = req_top_k_tokens[i];
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
const int32_t token_idx = resolve_selected_token<SPARSE_BLOCK_SIZE, TopKIsBlocks>(req_top_k, i);
if constexpr (TopKIsBlocks) {
if (token_idx < 0 || token_idx >= seq_len) {
s_top_k_tokens[i] = TOKEN_HIT;
req_top_k_device_locs[i] = -1;
continue;
}
}
if (token_idx == newest_token) {
// If topk includes the latest token, bind its canonical occurrence to newest_slot (at HOT_BUFFER_SIZE) and mark
// it as a hit. newest_slot is at the first position of the extra page, excluded from LRU tracking.
s_top_k_tokens[i] = TOKEN_HIT;
req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot];
s_newest_hit = 1;
} else {
int slot = hash_slot(token_idx, HASH_SIZE);
while (true) {
@@ -580,7 +609,7 @@ __global__ void load_cache_to_device_buffer_kernel(
const int chunk_token_start = chunk_idx * WARP_SIZE;
const int my_token_idx = chunk_token_start + lane_id;
const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K);
const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K_TOKENS);
int32_t my_token = 0;
bool is_miss = false;
@@ -611,6 +640,9 @@ __global__ void load_cache_to_device_buffer_kernel(
#else
total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses);
#endif
if (tid == 0) {
s_total_misses = total_misses;
}
}
__syncthreads();
@@ -632,7 +664,7 @@ __global__ void load_cache_to_device_buffer_kernel(
}
__syncthreads();
total_misses = NUM_TOP_K - s_total_hits - s_newest_hit;
total_misses = s_total_misses;
if constexpr (RecordMissPlan) {
if (tid == 0) {
miss_count_out[bid] = total_misses;
@@ -695,10 +727,12 @@ template <
int HOT_BUFFER_SIZE,
bool IsMLA,
bool IsDsv4Layout,
int SPARSE_BLOCK_SIZE,
bool TopKIsBlocks,
bool RecordMissPlan,
bool SkipIO>
void load_cache_to_device_buffer(
tvm::ffi::TensorView top_k_tokens,
tvm::ffi::TensorView top_k,
tvm::ffi::TensorView device_buffer_tokens,
tvm::ffi::TensorView host_cache_locs,
tvm::ffi::TensorView device_buffer_locs,
@@ -718,7 +752,8 @@ void load_cache_to_device_buffer(
tvm::ffi::TensorView miss_count_out) {
using namespace host;
const int64_t bs = top_k_tokens.shape()[0];
constexpr int NUM_TOP_K_TOKENS = NUM_TOP_K * (TopKIsBlocks ? SPARSE_BLOCK_SIZE : 1);
const int64_t bs = top_k.shape()[0];
const int64_t host_stride = host_cache_locs.shape()[1];
// Miss-plan side outputs; 0-dim sentinels when RecordMissPlan is false.
int64_t* const miss_src_ptr = RecordMissPlan ? static_cast<int64_t*>(miss_src_out.data_ptr()) : nullptr;
@@ -730,9 +765,9 @@ void load_cache_to_device_buffer(
}
const int64_t buffer_stride_0 = device_buffer_tokens.strides()[0];
const int64_t lru_slot_stride_0 = lru_slots.strides()[0];
const int64_t top_k_tokens_stride = top_k_tokens.strides()[0];
const int64_t top_k_stride = top_k.strides()[0];
const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0];
const auto kernel_device = top_k_tokens.device();
const auto kernel_device = top_k.device();
const auto device = LaunchKernel::resolve_device(kernel_device);
const void* const host_cache_k_ptr = runtime::get_device_accessible_ptr(host_cache_k);
const void* const host_cache_v_ptr =
@@ -741,7 +776,7 @@ void load_cache_to_device_buffer(
// Generic lambda: int32/int64 kernel variants are compiled for both
// seq_lens and req_pool_indices; the correct combo is selected at runtime.
auto launch = [&](auto kernel_fn, const auto* seq_lens_ptr, const auto* req_pool_indices_ptr) {
constexpr size_t smem_bytes = SmemLayout<NUM_TOP_K, HOT_BUFFER_SIZE>::BYTES;
constexpr size_t smem_bytes = SmemLayout<NUM_TOP_K_TOKENS, HOT_BUFFER_SIZE>::BYTES;
#ifndef USE_ROCM
if constexpr (smem_bytes > 48u * 1024u) {
cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
@@ -749,7 +784,7 @@ void load_cache_to_device_buffer(
#endif
LaunchKernel(bs, BLOCK_SIZE, device, smem_bytes)(
kernel_fn,
static_cast<const int32_t*>(top_k_tokens.data_ptr()),
static_cast<const int32_t*>(top_k.data_ptr()),
static_cast<int32_t*>(device_buffer_tokens.data_ptr()),
static_cast<const int64_t*>(host_cache_locs.data_ptr()),
static_cast<const int32_t*>(device_buffer_locs.data_ptr()),
@@ -765,7 +800,7 @@ void load_cache_to_device_buffer(
buffer_stride_0,
host_stride,
lru_slot_stride_0,
top_k_tokens_stride,
top_k_stride,
top_k_device_locs_stride,
page_size,
item_size_bytes,
@@ -788,6 +823,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int64_t,
@@ -802,6 +839,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int64_t,
@@ -816,6 +855,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int32_t,
@@ -830,6 +871,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int32_t,
python/sglang/kernels/jit/include/sgl_kernel/utils.cuh
+2
View File
@@ -53,6 +53,8 @@ inline constexpr auto cudaSuccess = hipSuccess;
#define cudaDeviceGetAttribute hipDeviceGetAttribute
#define cudaDevAttrComputeCapabilityMajor hipDeviceAttributeComputeCapabilityMajor
#define cudaDevAttrComputeCapabilityMinor hipDeviceAttributeComputeCapabilityMinor
#define cudaFuncSetAttribute hipFuncSetAttribute
#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
#endif
namespace sglang {
python/sglang/kernels/jit/utils/occupancy.py
+9 -3
View File
@@ -10,7 +10,11 @@ from sglang.kernels.jit.utils.compile import load_jit
if TYPE_CHECKING:
from tvm_ffi.module import Module
__all__ = ["get_max_active_clusters"]
__all__ = ["NoSchedulableClustersError", "get_max_active_clusters"]
class NoSchedulableClustersError(ValueError):
"""The occupancy query succeeded, but no cluster fits the requested shape."""
@cache_once
@@ -36,11 +40,13 @@ def get_max_active_clusters(cluster_size: int, occupancy: int) -> int:
dividing a GPC evenly. The probe kernel is pinned to ``occupancy`` blocks per
SM, so pass the occupancy the real kernel reaches (its second
``__launch_bounds__`` argument). Raises ``RuntimeError`` before sm90, which
has no clusters, and ``ValueError`` when nothing is schedulable.
has no clusters, and ``NoSchedulableClustersError`` (a ``ValueError``)
when the query succeeds but nothing is schedulable. Other probe errors
propagate to the caller.
"""
result = _get_max_active_clusters(cluster_size, occupancy)
if result == 0:
raise ValueError(
raise NoSchedulableClustersError(
f"no cluster of {cluster_size} fits at occupancy {occupancy}; "
"the cluster width is likely beyond what this device supports"
)
python/sglang/kernels/ops/attention/dsa/transform_index.py
+71 -7
View File
@@ -116,6 +116,50 @@ def transform_index_page_table_decode_kernel(
tl.store(result_ptr + offset, -1, mask=~mask)
@triton.jit
def transform_index_page_table_decode_tiled_kernel(
page_table_ptr: torch.Tensor,
topk_indices_ptr: torch.Tensor,
result_ptr: torch.Tensor,
page_table_row_stride: tl.constexpr,
topk_indices_stride_0: tl.constexpr,
topk_indices_stride_1: tl.constexpr,
result_stride_0: tl.constexpr,
result_stride_1: tl.constexpr,
TOPK: tl.constexpr,
BLOCK_TOPK: tl.constexpr,
):
"""Width-generic form of the kernel above.
The 2048 variant folds the row stride into a compile-time TOPK and covers a
whole row with one unmasked `tl.arange`, which needs TOPK to be a power of
two. k-pool widths are not: `index_topk + index_kpool - 1` is 2051 for
GLM-5.3-Flash. Tile the row instead and carry the strides explicitly.
"""
req_id = tl.program_id(0)
topk_offsets = tl.program_id(1) * BLOCK_TOPK + tl.arange(0, BLOCK_TOPK)
in_row = topk_offsets < TOPK
loaded_topk_indices = tl.load(
topk_indices_ptr
+ req_id * topk_indices_stride_0
+ topk_offsets * topk_indices_stride_1,
mask=in_row,
other=-1,
)
selected = in_row & (loaded_topk_indices >= 0)
loaded_kv_indices = tl.load(
page_table_ptr + req_id * page_table_row_stride + loaded_topk_indices,
mask=selected,
other=-1,
)
tl.store(
result_ptr + req_id * result_stride_0 + topk_offsets * result_stride_1,
loaded_kv_indices,
mask=in_row,
)
# Expanded EAGLE page tables are contiguous, so their row stride changes with
# the exact context length. Treating it as constexpr creates one cubin per
# observed length and grows the loaded-module set in long-lived processes.
@@ -194,18 +238,37 @@ def transform_index_page_table_decode_fast(
"""
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
assert topk_indices.shape[1] == 2048
qo_len = topk_indices.shape[0]
topk = topk_indices.shape[1]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
# Launch triton kernel
grid = (qo_len,)
transform_index_page_table_decode_kernel[grid](
if topk == 2048:
# Keep the single-program path for the unpooled width, which covers a
# whole row per program with no masking.
transform_index_page_table_decode_kernel[(qo_len,)](
page_table,
topk_indices,
result,
page_size,
page_table_row_stride=page_table.stride(0),
)
return result
block_topk = 256
transform_index_page_table_decode_tiled_kernel[
(qo_len, triton.cdiv(topk, block_topk))
](
page_table,
topk_indices,
result,
page_size,
page_table_row_stride=page_table.stride(0),
page_table.stride(0),
topk_indices.stride(0),
topk_indices.stride(1),
result.stride(0),
result.stride(1),
TOPK=topk,
BLOCK_TOPK=block_topk,
num_warps=4,
)
return result
@@ -220,7 +283,8 @@ def transform_index_page_table_prefill_fast(
cu_seqlens_q: Optional[torch.Tensor] = None,
) -> torch.Tensor:
assert page_size == 1
assert topk_indices.shape[1] == 2048
assert topk_indices.ndim == 2
assert topk_indices.shape[1] > 0
real_num_tokens = sum(extend_lens_cpu)
result = _allocate_prefill_result(topk_indices, real_num_tokens, output_num_tokens)
if real_num_tokens == 0:
python/sglang/kernels/ops/attention/dsv4/topk.py
+74 -22
View File
@@ -29,38 +29,40 @@ def _jit_topk_v1_module():
@cache_once
def _jit_topk_v2_module():
from sglang.kernels.jit.utils.occupancy import get_max_active_clusters
from sglang.kernels.jit.utils.occupancy import (
NoSchedulableClustersError,
get_max_active_clusters,
)
args = make_cpp_args(is_arch_support_pdl())
# Leave these undefined if the probe fails: topk_v2.cuh carries per-arch
# defaults, and a 0 would size the persistent pool to an empty grid.
# Enable each cluster path only when its occupancy probe reports capacity.
extra_cuda_cflags = []
if is_arch_support_pdl(): # set the persistent cluster size after hopper
try:
occ_8_2 = get_max_active_clusters(8, occupancy=2)
except Exception:
pass
else:
if occ_8_2 > 0:
extra_cuda_cflags.append(f"-DSGL_TOPK_V2_MAX_C8_OCC2={occ_8_2}")
try:
occ_16_1 = get_max_active_clusters(16, occupancy=1)
except Exception:
pass
else:
if occ_16_1 > 0:
extra_cuda_cflags.append(f"-DSGL_TOPK_V2_MAX_C16_OCC1={occ_16_1}")
for cluster_size, occupancy in ((8, 2), (16, 1)):
try:
max_active_clusters = get_max_active_clusters(
cluster_size, occupancy=occupancy
)
except NoSchedulableClustersError:
max_active_clusters = 0
extra_cuda_cflags.append(
f"-DSGL_TOPK_V2_MAX_C{cluster_size}_OCC{occupancy}={max_active_clusters}"
)
kernel = f"TopKKernel<{args}>"
wrappers = [
("topk_transform_paged", f"{kernel}::transform_paged"),
("topk_transform_ragged", f"{kernel}::transform_ragged"),
("topk_plan", f"{kernel}::plan"),
]
if is_hip_runtime():
# transform_packed only exists under USE_ROCM, see topk_v2.cuh
wrappers.append(("topk_transform_packed", f"{kernel}::transform_packed"))
return load_jit(
make_name("topk_v2"),
*args,
extra_cuda_cflags=extra_cuda_cflags,
cuda_files=["deepseek_v4/topk_v2.cuh"],
cuda_wrappers=[
("topk_transform_paged", f"{kernel}::transform_paged"),
("topk_transform_ragged", f"{kernel}::transform_ragged"),
("topk_plan", f"{kernel}::plan"),
],
cuda_wrappers=wrappers,
)
@@ -214,6 +216,9 @@ def topk_transform_paged_v2(
* Both outputs given -- ``out_page_indices`` receives the page-table
transform and ``out_raw_indices`` receives the selected raw indices.
For the packed (DSA extend prefill) layout see
:func:`topk_transform_packed_v2`.
NOTE: every entry of `seq_lens` must be NON-NEGATIVE, and `metadata` must
come from :func:`plan_topk_v2` over the same `seq_lens` values.
A length of 0 is the valid way to express "no tokens": the row takes the
@@ -249,3 +254,50 @@ def topk_transform_paged_v2(
metadata,
out_raw_indices,
)
def topk_transform_packed_v2(
scores: torch.Tensor,
seq_lens: torch.Tensor,
page_tables: torch.Tensor,
out_page_indices: torch.Tensor,
page_size: int,
*,
row_starts: torch.Tensor,
row_to_batch: Optional[torch.Tensor] = None,
) -> None:
"""Packed (DSA extend prefill) fused top-k + page-table transform.
Row ``i`` selects the top-k of ``scores[i, ks : ks + seq_lens[i]]``
(``ks = row_starts[i]``) and writes the page-table transform of the selected
row-local positions into ``out_page_indices``, ``-1`` padded. Prefill expands
one request into many query-token rows, so ``row_to_batch[i]`` (optional,
``(rows,)`` int32) names the ``page_tables`` row of the request row ``i``
belongs to; omitting it indexes the table by score row. ``row_to_batch`` is
not range-checked.
This is :func:`topk_transform_ragged_v2` with a page-table output instead of
an additive offset. Like ragged, it dispatches the implementation per row at
runtime, so it needs no plan and no :func:`plan_topk_v2` metadata.
NOTE: ``scores`` is MODIFIED IN PLACE -- the <= 3 columns ahead of each row's
window that the 16-byte-aligned read base pulls in are masked out. They are
invalid for that row and the buffer must have no other consumer, so do not
pass a view with overlapping rows.
``seq_lens`` entries must be NON-NEGATIVE, as for the paged entry point.
ROCm only: the kernel is compiled under ``USE_ROCM`` so that CUDA and XPU
builds are untouched. Nothing in it is AMD-specific -- no non-ROCm caller
produces this layout today.
"""
assert is_hip_runtime(), "topk_transform_packed_v2 is compiled under USE_ROCM only"
module = _jit_topk_v2_module()
module.topk_transform_packed(
scores,
seq_lens,
row_starts,
page_tables,
out_page_indices,
page_size,
row_to_batch,
)
python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py
+11 -23
View File
@@ -12,9 +12,10 @@ two transpose copies the unfused path needs to feed the conv kernel.
Scope (v1): chain speculation only (``speculative_eagle_topk == 1``, i.e.
``retrieve_next_token is None``). The tree path keeps the unfused reference
kernels. Requires ``T >= kernel_width - 1`` (the rolled conv state is then
exactly the last ``kernel_width - 1`` input tokens, matching the reference
kernel's store).
kernels. Requires ``T >= kernel_width - 1``.
State: conv_state and the SSM state are read-only. Verify is speculative, and
the commit scatter advances them from the selected intermediate window.
ReplaySSM (``cache_ring``): instead of per-step [HV, V, K] fp32 state
snapshots, stash each step's raw inputs (pre-l2norm k, pre-delta v, gate,
@@ -383,19 +384,8 @@ def fused_kda_conv_gating_verify_kernel(
)
tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
# Rolled conv state after consuming T >= W-1 tokens is exactly the last
# W-1 input tokens — which are the current window registers. The verify
# pass never writes the ssm state back (rollback happens at commit).
if is_qk_owner:
tl.store(cs_base + q_ch + 0 * stride_cs_tok, q_c0, mask=mask_k)
tl.store(cs_base + q_ch + 1 * stride_cs_tok, q_c1, mask=mask_k)
tl.store(cs_base + q_ch + 2 * stride_cs_tok, q_c2, mask=mask_k)
tl.store(cs_base + k_ch + 0 * stride_cs_tok, k_c0, mask=mask_k)
tl.store(cs_base + k_ch + 1 * stride_cs_tok, k_c1, mask=mask_k)
tl.store(cs_base + k_ch + 2 * stride_cs_tok, k_c2, mask=mask_k)
tl.store(cs_base + v_ch + 0 * stride_cs_tok, v_c0, mask=mask_v)
tl.store(cs_base + v_ch + 1 * stride_cs_tok, v_c1, mask=mask_v)
tl.store(cs_base + v_ch + 2 * stride_cs_tok, v_c2, mask=mask_v)
# No conv-state writeback: every V tile reads the same Q/K history, so a
# tile in a later wave would read what i_v == 0 had overwritten.
def fused_kda_conv_gating_verify(
@@ -423,13 +413,11 @@ def fused_kda_conv_gating_verify(
softplus_beta: float = 1.0,
softplus_threshold: float = 20.0,
use_qk_l2norm_in_kernel: bool = True,
# num_warps=4 is ~1.3x faster than the unfused pair in-graph; conv_state
# and the conv-window cache stay bit-identical to the reference, the bf16
# output within one ulp (the BV=4 tile reduces K in a different order).
# The fp32 intermediate-ssm rollback cache carries that ~1 ulp/step delta
# through the delta-rule recurrence — measured ~6e-8 at T=4 standard gate
# (the production MTP shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe
# gate. num_warps=1 is ~2.4x slower in-graph — numerics debugging only.
# num_warps=4 is ~1.3x faster than the unfused pair in-graph; 1 restores the
# reference reduction order but is ~2.4x slower, for numerics debugging only.
# The fp32 intermediate-ssm rollback cache carries the reduction-order delta
# furthest: ~6e-8 at T=4 standard gate (the production MTP shape), ~2e-3 at
# T=8 safe gate. conv_state is not comparable to the reference at all.
# The ReplaySSM ring values are bit-exact at any num_warps: they are
# elementwise (conv FMA chain, gate, sigmoid), upstream of every tl.sum.
num_warps: int = 4,
python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py
+29 -2
View File
@@ -5,6 +5,33 @@ import triton
import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.srt.utils import is_gfx95_supported, is_hip
_is_hip = is_hip()
_is_gfx95 = is_gfx95_supported()
def _select_recurrent_launch_config(
n: int,
h: int,
hv: int,
k: int,
v: int,
is_kda: bool,
) -> tuple[int, int]:
"""Select the value tile and warp count for recurrent GDN."""
if (
_is_hip
and _is_gfx95
and not is_kda
and 0 < n <= 32
and h == 4
and hv == 16
and k == 128
and v == 128
):
return (8, 4) if n == 1 else (16, 2)
return min(triton.next_power_of_2(v), 32), 1
@triton.jit(do_not_specialize=["T"])
@@ -401,11 +428,11 @@ def fused_sigmoid_gating_delta_rule_update(
stride_a = a.stride()[1] if a.ndim == 4 else a.stride()[-2]
HV = v.shape[2]
N = B if cu_seqlens is None else len(cu_seqlens) - 1
BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
BV, num_warps = _select_recurrent_launch_config(N, H, HV, K, V, is_kda)
BK = triton.next_power_of_2(K)
NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
assert NK == 1, "NK > 1 is not supported yet"
num_stages = 3
num_warps = 1
if scale is None:
scale = k.shape[-1] ** -0.5
python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py
+27 -6
View File
@@ -23,6 +23,7 @@ from ..common.utils import (
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
"HAS_SINK": lambda args: args["sink_ptr"] is not None,
"BATCH_SIZE_BUCKET": lambda args: triton.next_power_of_2(args["batch_size"]),
"HAS_HISPARSE_SLOTS": lambda args: args["hisparse_slots_ptr"] is not None,
}
)
@triton.autotune(
@@ -43,6 +44,7 @@ def _gqa_share_sparse_decode_kernel(
idx_ptr, # topk index: qh x b x topk
o_ptr, # O partial: c x b x qh x d
lse_ptr, # lse partial: c x b x qh
hisparse_slots_ptr, # pre-resolved device slots: kh x b x (topk * block)
seq_lens,
slot_ids,
# shape
@@ -52,6 +54,8 @@ def _gqa_share_sparse_decode_kernel(
head_dim,
max_topk,
max_kv_len,
hisparse_slots_stride_h,
hisparse_slots_stride_b,
# sm_scale
sm_scale,
# per-tensor KV dequant scales (1.0 when the cache is unit-scaled)
@@ -89,6 +93,7 @@ def _gqa_share_sparse_decode_kernel(
NUM_TOPK_CHUNKS: tl.constexpr,
HAS_SINK: tl.constexpr,
IS_FP8: tl.constexpr,
HAS_HISPARSE_SLOTS: tl.constexpr,
):
# decode program ids: split-K over the topk dimension to give every SM
# something to do at small batch. pid(0) folds (batch, chunk) together so
@@ -161,18 +166,30 @@ def _gqa_share_sparse_decode_kernel(
# only iterate over this chunk's topk slice. the load must respect the
# per-chunk start offset.
cur_idx_ptr = idx_base + chunk_start_topk * stride_ti_t
hisparse_topk_counter = chunk_start_topk
for _ in tl.range(chunk_start_topk, chunk_end_topk):
# load index
c = tl.load(cur_idx_ptr).to(tl.int32) * BLOCK_SIZE_N
cur_idx_ptr = cur_idx_ptr + stride_ti_t
# resolve slots for this block via req_to_token
pos = c + off_n
pos_mask = pos < seq_len
slots = tl.load(
req_to_token_ptr + sid * stride_r2t_b + pos,
mask=pos_mask,
other=0,
).to(tl.int64)
if HAS_HISPARSE_SLOTS:
slots = tl.load(
hisparse_slots_ptr
+ pid_kh * hisparse_slots_stride_h
+ pid_b * hisparse_slots_stride_b
+ hisparse_topk_counter * BLOCK_SIZE_N
+ off_n,
mask=off_n < BLOCK_SIZE_N,
other=0,
).to(tl.int64)
hisparse_topk_counter = hisparse_topk_counter + 1
else:
slots = tl.load(
req_to_token_ptr + sid * stride_r2t_b + pos,
mask=pos_mask,
other=0,
).to(tl.int64)
slots = (slots + max_slots) % max_slots # safety against negative
# load K as (head_dim, BLOCK_SIZE_N) via indirect addressing
k_off = (
@@ -321,6 +338,7 @@ def flash_decode_with_gqa_share_sparse(
q_scale: Optional[float] = None,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
hisparse_slots: Optional[torch.Tensor] = None,
) -> torch.Tensor:
triton.set_allocator(robust_allocator)
is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="decode")
@@ -384,6 +402,7 @@ def flash_decode_with_gqa_share_sparse(
topk_idx,
o_partial,
lse_partial,
hisparse_slots,
seq_lens,
slot_ids,
max_slots,
@@ -392,6 +411,8 @@ def flash_decode_with_gqa_share_sparse(
head_dim,
max_topk,
max_kv_len,
hisparse_slots.stride(0) if hisparse_slots is not None else 0,
hisparse_slots.stride(1) if hisparse_slots is not None else 0,
sm_scale,
k_scale,
v_scale,
python/sglang/kernels/ops/attention/minimax_sparse/prefill/topk_sparse.py
+11
View File
@@ -28,6 +28,7 @@ from ..common.utils import (
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
"BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"] * args["BLOCK_SIZE_H"],
"HAS_SINK": lambda args: args["sink_ptr"] is not None,
"HAS_LOC_MAPPING": lambda args: args["loc_mapping_ptr"] is not None,
}
)
@triton.autotune(
@@ -55,6 +56,7 @@ def _gqa_share_sparse_fwd_kernel(
t_ptr, # topk_idx: kh x n x k
o_ptr, # O: n x h x d
req_to_token_ptr, # req_to_token: max_reqs x max_kv_len
loc_mapping_ptr, # logical slot to HiSparse device slot
# seqlens
cu_seqlens_q,
cu_seqblocks_q,
@@ -106,6 +108,7 @@ def _gqa_share_sparse_fwd_kernel(
HAS_SINK: tl.constexpr,
USE_TMA: tl.constexpr,
IS_FP8: tl.constexpr,
HAS_LOC_MAPPING: tl.constexpr,
):
sm_scale_log2e = sm_scale * 1.4426950409
# get batch id and head id
@@ -199,6 +202,12 @@ def _gqa_share_sparse_fwd_kernel(
mask=pos_mask,
other=0,
).to(tl.int64)
if HAS_LOC_MAPPING:
slots = tl.load(
loc_mapping_ptr + slots,
mask=pos_mask,
other=0,
).to(tl.int64)
slots = (slots + max_slots) % max_slots # safety against negative
# k shape: [BLOCK_SIZE_KD, BLOCK_SIZE_K] (transposed for tl.dot)
k = tl.load(
@@ -289,6 +298,7 @@ def flash_prefill_with_gqa_share_sparse(
q_scale: Optional[float] = None,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
loc_mapping: Optional[torch.Tensor] = None,
) -> torch.Tensor:
triton.set_allocator(robust_allocator)
is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="prefill")
@@ -340,6 +350,7 @@ def flash_prefill_with_gqa_share_sparse(
topk_idx,
o,
req_to_token,
loc_mapping,
cu_seqlens,
cu_seqblocks_q,
seq_lens,
python/sglang/kernels/ops/attention/mrope.py
+5 -2
View File
@@ -7,11 +7,14 @@ import triton
import triton.language as tl
@triton.jit
# Sequence length varies across requests. Specializing S would create one
# compiled kernel variant per observed length, while the existing mask handles
# the final partial block.
@triton.jit(do_not_specialize=["S"])
def apply_interleaved_rope_kernel(
x_ptr,
out_ptr,
S: tl.constexpr,
S,
D: tl.constexpr,
stride_x_m,
stride_x_s,
python/sglang/kernels/ops/diffusion/__init__.py
+1
View File
@@ -736,6 +736,7 @@ _EXPORTS: dict[str, str] = {
"interpolate": "ext.hunyuan3d_rasterizer",
"rasterize": "ext.hunyuan3d_rasterizer",
"meshVerticeInpaint": "ext.mesh_processor",
"load_mesh_processor": "ext.mesh_processor",
}
python/sglang/kernels/ops/diffusion/ext/mesh_processor/__init__.py
+3 -3
View File
@@ -19,7 +19,7 @@ _abs_path = os.path.dirname(os.path.abspath(__file__))
_mesh_processor_kernel = None
def _load_mesh_processor():
def load_mesh_processor():
"""JIT compile and load the mesh processor kernel."""
global _mesh_processor_kernel
@@ -47,7 +47,7 @@ def meshVerticeInpaint(
method: str = "smooth",
) -> Tuple[np.ndarray, np.ndarray]:
"""Inpaint texture using mesh vertex connectivity."""
kernel = _load_mesh_processor()
kernel = load_mesh_processor()
texture = np.ascontiguousarray(texture, dtype=np.float32)
mask = np.ascontiguousarray(mask, dtype=np.uint8)
@@ -61,4 +61,4 @@ def meshVerticeInpaint(
)
__all__ = ["meshVerticeInpaint"]
__all__ = ["load_mesh_processor", "meshVerticeInpaint"]
python/sglang/kernels/ops/kv_canary/_dispatch.py
+12
View File
@@ -0,0 +1,12 @@
from __future__ import annotations
import torch
def use_torch_reference(device: torch.device) -> bool:
"""Whether a canary launcher must fall back to its byte-equal torch reference.
The write / verify / plan-entries kernels are CUDA-JIT only; HIP keeps them
since torch reports it as ``"cuda"``. XPU / CPU / anything else falls back.
"""
return device.type != "cuda"
python/sglang/kernels/ops/kv_canary/plan/api.py
+32 -5
View File
@@ -4,6 +4,7 @@ from typing import Optional
import torch
from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference
from sglang.kernels.ops.kv_canary.plan.entries_kernel import (
launch_plan_entries_kernel,
)
@@ -99,6 +100,7 @@ def launch_canary_plan_kernels(
Calling contract:
- Pure side-effect; no host work, no D2H.
- Safe in cuda-graph capture; caller refills all input tensors in-place before replay.
The reference path is not (host work, D2H) and must not be launched under capture.
- The wrapper launches the plan sub-kernels needed to fill both plans end-to-end.
- Padding rows contribute zero entries.
@@ -106,17 +108,42 @@ def launch_canary_plan_kernels(
:func:`sglang.kernels.ops.kv_canary.plan_ref.launch_canary_plan_kernels_torch_reference`; both the Triton
offsets kernel and the CUDA JIT entries kernel must match byte-for-byte.
"""
# SWA plans are meaningless without the full->swa LUT (entries would carry
# untranslated full-pool slots), so this is a cross-backend contract, not a
# CUDA-only guard. Enforce it before dispatching: the torch reference does not
# re-check, so leaving it below the early-return would silently skip it.
if swa_window_size > 0 and full_to_swa_index_mapping is None:
raise ValueError(
"kv-canary: launch_canary_plan_kernels requires full_to_swa_index_mapping when swa_window_size > 0"
)
if use_torch_reference(verify_plan_out.verify_slot_indices.device):
from sglang.kernels.ops.kv_canary.plan_ref import (
launch_canary_plan_kernels_torch_reference,
)
launch_canary_plan_kernels_torch_reference(
verify_plan_out=verify_plan_out,
write_plan_out=write_plan_out,
req_pool_indices=req_pool_indices,
prefix_lens=prefix_lens,
extend_seq_lens=extend_seq_lens,
req_to_token=req_to_token,
swa_window_size=swa_window_size,
full_to_swa_index_mapping=full_to_swa_index_mapping,
verify_capacity=verify_capacity,
req_to_verify_expected_tokens=req_to_verify_expected_tokens,
req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
)
return
bs = int(req_pool_indices.shape[0])
if bs > _PLAN_BS_BLOCK_SIZE:
raise ValueError(
f"kv-canary: launch_canary_plan_kernels supports at most bs={_PLAN_BS_BLOCK_SIZE} reqs per launch, "
f"got bs={bs}. Bump _PLAN_BS_BLOCK_SIZE if real workloads need this."
)
if swa_window_size > 0 and full_to_swa_index_mapping is None:
raise ValueError(
"kv-canary: launch_canary_plan_kernels requires full_to_swa_index_mapping when swa_window_size > 0"
)
device = verify_plan_out.verify_slot_indices.device
verify_offsets_scratch = torch.empty(
_PLAN_BS_BLOCK_SIZE + 1, dtype=torch.int64, device=device
python/sglang/kernels/ops/kv_canary/verify.py
+29 -1
View File
@@ -8,6 +8,7 @@ import torch
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.ops.kv_canary import consts
from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference
if TYPE_CHECKING:
from tvm_ffi.module import Module
@@ -120,6 +121,16 @@ class RealKvSource:
f"got {row_stride_bytes} bytes (shape={tuple(self.tensor.shape)}, "
f"dtype={self.tensor.dtype})"
)
# A row is addressed as page_size slots of num_bytes_per_token, unchecked at fold time;
# a narrower row hashes fewer bytes than asked and still reports the chain clean.
min_row_bytes = self.page_size * self.num_bytes_per_token
if row_stride_bytes < min_row_bytes:
raise ValueError(
f"kv-canary: RealKvSource.tensor dim-1 is {row_stride_bytes} bytes but "
f"page_size={self.page_size} x num_bytes_per_token={self.num_bytes_per_token} "
f"needs {min_row_bytes} (shape={tuple(self.tensor.shape)}, "
f"dtype={self.tensor.dtype})"
)
@dataclass(frozen=True, slots=True, kw_only=True)
@@ -301,7 +312,8 @@ def launch_canary_verify_kernel(
- Pure side-effect; never raises. Host polls violation_write_index[0] > 0 for is_errored and
violation_ring[0] for the first violation.
- kernel_run_counter is bumped every call (canary-ran health signal).
- Safe in cuda-graph capture; caller refills plan in-place before replay.
- Safe in cuda-graph capture; caller refills plan in-place before replay. The reference
path is not (host work, D2H) and must not be launched under capture.
Pinned by torch reference
:func:`sglang.kernels.ops.kv_canary.verify_ref.launch_canary_verify_kernel_torch_reference`; CUDA must match
@@ -309,12 +321,28 @@ def launch_canary_verify_kernel(
"""
canary_buf = context.canary_buf
real_kv_sources = context.real_kv_sources
# Enforce the source-count cap before dispatching: the torch reference is
# pinned to match the CUDA ABI byte-for-byte, so the limit is a cross-backend
# contract, not a CUDA-only guard. Checking after the reference early-return
# (XPU / CPU path) would silently skip it.
if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES:
raise ValueError(
f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
f"got {len(real_kv_sources)}"
)
if use_torch_reference(canary_buf.device):
from sglang.kernels.ops.kv_canary.verify_ref import (
launch_canary_verify_kernel_torch_reference,
)
launch_canary_verify_kernel_torch_reference(
context=context,
plan=plan,
check_verify_expected_token=check_verify_expected_token,
)
return
_assert_contiguous(canary_buf, "canary_buf")
_assert_contiguous(plan.verify_slot_indices, "plan.verify_slot_indices")
_assert_contiguous(plan.verify_expected_tokens, "plan.verify_expected_tokens")
python/sglang/kernels/ops/kv_canary/verify_ref.py
+79 -25
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
from typing import NamedTuple, Sequence
import torch
from sglang.kernels.ops.kv_canary import consts
@@ -93,6 +95,13 @@ def launch_canary_verify_kernel_torch_reference(
f"kv-canary: canary_buf slot stride must hold at least 4 int64 fields, got {slot_stride_i64}"
)
host_real_kv_sources = materialize_real_kv_sources(
real_kv_sources=real_kv_sources,
real_kv_hash_mode=real_kv_hash_mode,
slot_indices=slot_indices_list,
work_device=work_device,
)
violation_rows: list[list[int]] = []
for k in range(active):
@@ -118,9 +127,7 @@ def launch_canary_verify_kernel_torch_reference(
expected_real_kv_hash_u64 = _compute_real_kv_hash_scalar(
slot_idx=slot_idx,
real_kv_sources=real_kv_sources,
real_kv_hash_mode=real_kv_hash_mode,
work_device=work_device,
host_sources=host_real_kv_sources,
)
expected_real_kv_hash = _to_signed_int64(expected_real_kv_hash_u64)
@@ -190,37 +197,84 @@ def compute_slot_hash(buf_i64: torch.Tensor, source_slot_idx: int) -> int:
return splitmix64_mix3(prev_hash, token, position)
class _MaterializedRealKvSource(NamedTuple):
"""A ``RealKvSource`` narrowed to the rows one launch reads, on ``work_device``.
``row_lookup`` maps a source row (``slot_idx // page_size``) to its index in
``tensor_u8``, which holds only the gathered rows.
"""
tensor_u8: torch.Tensor
row_lookup: dict[int, int]
page_size: int
num_bytes_per_token: int
effective_read_bytes: int
def materialize_real_kv_sources(
*,
real_kv_sources: tuple[RealKvSource, ...],
real_kv_hash_mode: consts.RealKvHashMode,
slot_indices: Sequence[int],
work_device: torch.device,
) -> tuple[_MaterializedRealKvSource, ...]:
"""Gather each source's read rows onto ``work_device`` once per launch.
An empty tuple means nothing to hash; callers skip the per-slot fold."""
mode = int(real_kv_hash_mode)
if (
mode == int(consts.RealKvHashMode.NONE)
or len(real_kv_sources) == 0
or len(slot_indices) == 0
):
return ()
materialized: list[_MaterializedRealKvSource] = []
for source in real_kv_sources:
# Gather on device first: copying the whole source is a KV-layer-sized transfer
# (one row per token of the pool) on every launch.
rows = sorted({slot_idx // source.page_size for slot_idx in slot_indices})
row_index = torch.tensor(rows, dtype=torch.int64, device=source.tensor.device)
tensor_u8 = (
source.tensor.detach()
.index_select(0, row_index)
.to(device=work_device)
.contiguous()
.view(torch.uint8)
)
effective_read_bytes = (
16 if mode == int(consts.RealKvHashMode.PARTIAL) else source.read_bytes
)
materialized.append(
_MaterializedRealKvSource(
tensor_u8=tensor_u8,
row_lookup={row: i for i, row in enumerate(rows)},
page_size=source.page_size,
num_bytes_per_token=source.num_bytes_per_token,
effective_read_bytes=effective_read_bytes,
)
)
return tuple(materialized)
def _compute_real_kv_hash_scalar(
*,
slot_idx: int,
real_kv_sources: tuple[RealKvSource, ...],
real_kv_hash_mode: consts.RealKvHashMode,
work_device: torch.device,
host_sources: tuple[_MaterializedRealKvSource, ...],
) -> int:
mode = int(real_kv_hash_mode)
if mode == int(consts.RealKvHashMode.NONE) or len(real_kv_sources) == 0:
if len(host_sources) == 0:
return 0
acc: int = 0
for source in real_kv_sources:
page_size = source.page_size
num_bytes_per_token = source.num_bytes_per_token
read_bytes = source.read_bytes
tensor_u8 = (
source.tensor.detach().to(device=work_device).contiguous().view(torch.uint8)
)
for source in host_sources:
row = source.row_lookup[slot_idx // source.page_size]
col_within_page = slot_idx % source.page_size
col_start = col_within_page * source.num_bytes_per_token
row = slot_idx // page_size
col_within_page = slot_idx % page_size
col_start = col_within_page * num_bytes_per_token
effective_read_bytes = (
16 if mode == int(consts.RealKvHashMode.PARTIAL) else read_bytes
)
raw_bytes: list[int] = []
for b in range(effective_read_bytes):
raw_bytes.append(int(tensor_u8[row, col_start + b].item()))
raw_bytes = source.tensor_u8[
row, col_start : col_start + source.effective_read_bytes
].tolist()
source_hash = _splitmix64_fold_bytes_scalar(raw_bytes=raw_bytes)
python/sglang/kernels/ops/kv_canary/write.py
+24 -1
View File
@@ -7,6 +7,7 @@ import torch
from sglang.kernels.jit.utils import cache_once, load_jit
from sglang.kernels.ops.kv_canary import consts
from sglang.kernels.ops.kv_canary._dispatch import use_torch_reference
from sglang.kernels.ops.kv_canary.verify import (
VerifyOrWriteContext,
_assert_contiguous,
@@ -183,7 +184,8 @@ def launch_canary_write_kernel(
- Input-verification mismatch records violations but does NOT abort the chain.
- kernel_run_counter is bumped every call.
- Safe in cuda-graph capture; caller refills input_ids / positions / out_cache_loc / plan
in-place before replay.
in-place before replay. The reference path is not (host work, D2H) and must not be
launched under capture.
Pinned by torch reference
:func:`sglang.kernels.ops.kv_canary.write_ref.launch_canary_write_kernel_torch_reference`; CUDA must match
@@ -191,12 +193,33 @@ def launch_canary_write_kernel(
"""
canary_buf = context.canary_buf
real_kv_sources = context.real_kv_sources
# Enforce the source-count cap before dispatching: the torch reference is
# pinned to match the CUDA ABI byte-for-byte, so the limit is a cross-backend
# contract, not a CUDA-only guard. Checking after the reference early-return
# (XPU / CPU path) would silently skip it.
if len(real_kv_sources) > consts.MAX_REAL_KV_SOURCES:
raise ValueError(
f"kv-canary: at most {consts.MAX_REAL_KV_SOURCES} RealKvSource entries supported by the CUDA ABI, "
f"got {len(real_kv_sources)}"
)
if use_torch_reference(canary_buf.device):
from sglang.kernels.ops.kv_canary.write_ref import (
launch_canary_write_kernel_torch_reference,
)
launch_canary_write_kernel_torch_reference(
context=context,
plan=plan,
input_ids=input_ids,
positions=positions,
out_cache_loc=out_cache_loc,
enable_write_input_assert=enable_write_input_assert,
expected_input_tokens=expected_input_tokens,
expected_input_positions=expected_input_positions,
)
return
_assert_contiguous(canary_buf, "canary_buf")
_assert_contiguous(plan.write_offsets, "plan.write_offsets")
_assert_contiguous(plan.write_seed_slot_indices, "plan.write_seed_slot_indices")
python/sglang/kernels/ops/kv_canary/write_ref.py
+14 -3
View File
@@ -10,6 +10,7 @@ from sglang.kernels.ops.kv_canary.verify_ref import (
_compute_real_kv_hash_scalar,
_to_signed_int64,
compute_slot_hash,
materialize_real_kv_sources,
splitmix64_mix3,
)
from sglang.kernels.ops.kv_canary.write import WritePlan
@@ -96,6 +97,18 @@ def launch_canary_write_kernel_torch_reference(
expected_input_tokens_host = None
expected_input_positions_host = None
# A superset of the slots the loop below folds: the per-req entry ranges all lie
# inside [0, total_entries), and gathering a spare row is harmless.
write_slot_indices = [
slot for slot in out_cache_loc_host[:total_entries].tolist() if slot >= 0
]
host_real_kv_sources = materialize_real_kv_sources(
real_kv_sources=real_kv_sources,
real_kv_hash_mode=real_kv_hash_mode,
slot_indices=write_slot_indices,
work_device=work_device,
)
violation_rows: list[list[int]] = []
total_slots_written = 0
@@ -129,9 +142,7 @@ def launch_canary_write_kernel_torch_reference(
real_kv_hash_u64 = _compute_real_kv_hash_scalar(
slot_idx=slot,
real_kv_sources=real_kv_sources,
real_kv_hash_mode=real_kv_hash_mode,
work_device=work_device,
host_sources=host_real_kv_sources,
)
if enable_write_input_assert:
python/sglang/kernels/ops/kvcache/hicache.py
+131 -15
View File
@@ -3,6 +3,8 @@ from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import torch
from sglang.kernels.jit.utils import (
cache_once,
is_hip_runtime,
@@ -10,9 +12,9 @@ from sglang.kernels.jit.utils import (
make_cpp_args,
)
from sglang.kernels.kernel_api_logging import debug_kernel_api
from sglang.srt.environ import envs
if TYPE_CHECKING:
import torch
from tvm_ffi.module import Module
_is_hip = is_hip_runtime()
@@ -81,13 +83,81 @@ def _jit_hicache_staged_module(
)
# TMA staging ring per CTA: 32 KB stages x 6 keeps the host loads in flight
# under a 4-block launch; smaller stages make the loader's per-chunk cost the
# limit. 4 store warps drain a strided-destination stage faster than it fills.
TMA_STAGE_BYTES = 32 * 1024
TMA_NUM_STAGES = 6
TMA_STORE_WARPS = 4
# Each CTA is capped by its SM's write port, so the host link needs four of
# them; the register kernel keeps DEFAULT_BLOCK_QUOTA.
TMA_BLOCK_QUOTA = 4
@cache_once
def _jit_hicache_tma_module(*, block_quota: int) -> Module:
args = make_cpp_args(TMA_STAGE_BYTES, TMA_NUM_STAGES, TMA_STORE_WARPS, block_quota)
return load_jit(
"hicache_tma",
*args,
cuda_files=["kvcacheio/hicache_tma.cuh"],
cuda_wrappers=[
("launch_one", f"&HiCacheTmaKernel<{args}>::run_one"),
("launch_all", f"&HiCacheTmaKernel<{args}>::run_all"),
("launch_one_mla", f"&HiCacheTmaKernel<{args}>::run_one_mla"),
("launch_all_mla", f"&HiCacheTmaKernel<{args}>::run_all_mla"),
("fits_device", f"&HiCacheTmaKernel<{args}>::fits_device"),
],
)
def hicache_tma_rows_per_chunk(element_size: int) -> int:
"""Rows the TMA kernel moves per stage (mirrors hicache_tma_rows_per_chunk)."""
rows = 1
while rows * 2 <= TMA_STAGE_BYTES // element_size and rows * 2 <= 128:
rows *= 2
return rows
@cache_once
def use_hicache_tma_kernel(
*, element_size: int, block_quota: int, page_size: int | None = None
) -> bool:
"""Whether transfers of `element_size`-byte rows go through the TMA kernel.
A chunk that straddles pages degrades to one bulk op per row, so pages must
tile the chunk; `page_size=None` (caller unaware of paging) trusts the indices.
"""
if _is_hip or not envs.SGLANG_HICACHE_TMA_TRANSFER.get():
return False
if element_size % 16 != 0 or torch.cuda.get_device_capability()[0] < 9:
return False
if page_size is not None and page_size % hicache_tma_rows_per_chunk(element_size):
return False
try:
module = _jit_hicache_tma_module(block_quota=block_quota)
except Exception as e:
logging.getLogger(__name__).warning(
f"Failed to load the TMA HiCache kernel, using the register kernel: {e}"
)
return False
return bool(module.fits_device(torch.cuda.current_device()))
def can_use_hicache_jit_kernel(
*,
element_size: int,
unroll: int | None = None, # can be tuned for performance
block_quota: int | None = None, # can be tuned for less interference
page_size: int | None = None,
) -> bool:
logger = logging.getLogger(__name__)
if use_hicache_tma_kernel(
element_size=element_size,
block_quota=block_quota or TMA_BLOCK_QUOTA,
page_size=page_size,
):
return True
unroll = unroll or _default_unroll(element_size)
if not _tiles_across_lanes(element_size, unroll):
logger.warning(f"Unsupported {element_size = } for JIT HiCache kernel")
@@ -166,6 +236,7 @@ def transfer_hicache_one_layer(
element_dim: int | None = None,
unroll: int | None = None, # can be tuned for performance
block_quota: int | None = None, # can be tuned for less interference
page_size: int | None = None,
) -> None:
element_dim = element_dim or k_cache_dst.size(-1)
k_cache_src = k_cache_src.view(-1, element_dim)
@@ -173,13 +244,19 @@ def transfer_hicache_one_layer(
k_cache_dst = k_cache_dst.view(-1, element_dim)
v_cache_dst = v_cache_dst.view(-1, element_dim)
element_size = element_dim * k_cache_dst.element_size()
block_quota = block_quota or DEFAULT_BLOCK_QUOTA
unroll = unroll or _default_unroll(element_size)
module = _jit_hicache_module(
element_size=element_size,
unroll=unroll,
block_quota=block_quota,
)
tma_quota = block_quota or TMA_BLOCK_QUOTA
if use_hicache_tma_kernel(
element_size=element_size, block_quota=tma_quota, page_size=page_size
):
module = _jit_hicache_tma_module(block_quota=tma_quota)
else:
block_quota = block_quota or DEFAULT_BLOCK_QUOTA
unroll = unroll or _default_unroll(element_size)
module = _jit_hicache_module(
element_size=element_size,
unroll=unroll,
block_quota=block_quota,
)
module.launch_one(
k_cache_dst,
v_cache_dst,
@@ -204,11 +281,28 @@ def transfer_hicache_all_layer(
element_size: int | None = None,
unroll: int | None = None, # can be tuned for performance
block_quota: int | None = None, # can be tuned for less interference
page_size: int | None = None,
) -> None:
if element_size is None: # assume both contiguous
assert kv_cache_dst_stride_bytes == kv_cache_src_stride_bytes
element_size = kv_cache_dst_stride_bytes
tma_quota = block_quota or TMA_BLOCK_QUOTA
if use_hicache_tma_kernel(
element_size=element_size, block_quota=tma_quota, page_size=page_size
):
_jit_hicache_tma_module(block_quota=tma_quota).launch_all(
k_ptr_dst,
v_ptr_dst,
indices_dst,
k_ptr_src,
v_ptr_src,
indices_src,
kv_cache_src_stride_bytes,
kv_cache_dst_stride_bytes,
element_size,
)
return
block_quota = block_quota or DEFAULT_BLOCK_QUOTA
unroll = unroll or _default_unroll(element_size)
module = _jit_hicache_module(
@@ -237,18 +331,25 @@ def transfer_hicache_one_layer_mla(
element_dim: int | None = None,
unroll: int | None = None,
block_quota: int | None = None,
page_size: int | None = None,
) -> None:
element_dim = element_dim or cache_dst.size(-1)
cache_src = cache_src.view(-1, element_dim)
cache_dst = cache_dst.view(-1, element_dim)
element_size = element_dim * cache_dst.element_size()
block_quota = block_quota or DEFAULT_BLOCK_QUOTA
unroll = unroll or _default_unroll(element_size)
module = _jit_hicache_module(
element_size=element_size,
unroll=unroll,
block_quota=block_quota,
)
tma_quota = block_quota or TMA_BLOCK_QUOTA
if use_hicache_tma_kernel(
element_size=element_size, block_quota=tma_quota, page_size=page_size
):
module = _jit_hicache_tma_module(block_quota=tma_quota)
else:
block_quota = block_quota or DEFAULT_BLOCK_QUOTA
unroll = unroll or _default_unroll(element_size)
module = _jit_hicache_module(
element_size=element_size,
unroll=unroll,
block_quota=block_quota,
)
module.launch_one_mla(
cache_dst,
indices_dst,
@@ -268,11 +369,26 @@ def transfer_hicache_all_layer_mla(
element_size: int | None = None,
unroll: int | None = None,
block_quota: int | None = None,
page_size: int | None = None,
) -> None:
if element_size is None:
assert cache_dst_stride_bytes == cache_src_stride_bytes
element_size = cache_dst_stride_bytes
tma_quota = block_quota or TMA_BLOCK_QUOTA
if use_hicache_tma_kernel(
element_size=element_size, block_quota=tma_quota, page_size=page_size
):
_jit_hicache_tma_module(block_quota=tma_quota).launch_all_mla(
ptr_dst,
indices_dst,
ptr_src,
indices_src,
cache_src_stride_bytes,
cache_dst_stride_bytes,
element_size,
)
return
block_quota = block_quota or DEFAULT_BLOCK_QUOTA
unroll = unroll or _default_unroll(element_size)
module = _jit_hicache_module(
python/sglang/kernels/ops/kvcache/hisparse.py
+84 -1
View File
@@ -174,6 +174,8 @@ def _jit_sparse_module(
hot_buffer_size: int,
is_mla: bool = False,
is_dsv4_layout: bool = False,
top_k_block_size: int = 1,
top_k_is_blocks: bool = False,
record_miss_plan: bool = False,
skip_io: bool = False,
) -> Module:
@@ -185,6 +187,8 @@ def _jit_sparse_module(
hot_buffer_size,
is_mla,
is_dsv4_layout,
top_k_block_size,
top_k_is_blocks,
record_miss_plan,
skip_io,
)
@@ -195,6 +199,8 @@ def _jit_sparse_module(
hot_buffer_size,
is_mla,
is_dsv4_layout,
top_k_block_size,
top_k_is_blocks,
record_miss_plan,
skip_io,
)
@@ -308,7 +314,7 @@ def _load_cache_to_device_buffer_mla(
skip_io=skip_io,
)
empty = torch.empty(0)
empty = torch.empty(0, device=top_k_tokens.device)
if num_real_reqs is None:
num_real_reqs = torch.tensor(
@@ -399,6 +405,83 @@ def load_cache_to_device_buffer_mla(
)
def load_blocks_to_device_buffer_mha(
top_k_blocks: torch.Tensor,
device_buffer_tokens: torch.Tensor,
host_cache_locs: torch.Tensor,
device_buffer_locs: torch.Tensor,
host_cache_k: torch.Tensor,
host_cache_v: torch.Tensor,
device_buffer_k: torch.Tensor,
device_buffer_v: torch.Tensor,
top_k_device_locs: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
lru_slots: torch.Tensor,
item_size_bytes: int,
hot_buffer_size: int,
sparse_block_size: int,
page_size: int = 1,
block_size: int = 256,
num_real_reqs: torch.Tensor | None = None,
skip_io: bool = False,
) -> None:
"""Swap block-selected MHA K/V into the HiSparse device pool."""
num_top_k_blocks = top_k_blocks.size(1)
num_top_k_tokens = num_top_k_blocks * sparse_block_size
assert hot_buffer_size >= num_top_k_tokens, (
f"hot_buffer_size ({hot_buffer_size}) must be >= selected tokens "
f"({num_top_k_tokens})"
)
assert top_k_device_locs.size(1) >= num_top_k_tokens
k_stride = host_cache_k.stride(0) * host_cache_k.element_size()
v_stride = host_cache_v.stride(0) * host_cache_v.element_size()
assert k_stride == v_stride == item_size_bytes, (
"K/V token strides must equal item_size_bytes: "
f"k_stride={k_stride}, v_stride={v_stride}, "
f"item_size_bytes={item_size_bytes}"
)
module = _jit_sparse_module(
item_size_bytes,
block_size,
num_top_k_blocks,
hot_buffer_size,
is_mla=False,
is_dsv4_layout=False,
top_k_block_size=sparse_block_size,
top_k_is_blocks=True,
record_miss_plan=False,
skip_io=skip_io,
)
empty = torch.empty(0, device=top_k_blocks.device)
if num_real_reqs is None:
num_real_reqs = torch.tensor(
[top_k_blocks.size(0)], dtype=torch.int32, device=top_k_blocks.device
)
module.load_cache_to_device_buffer(
top_k_blocks,
device_buffer_tokens,
host_cache_locs,
device_buffer_locs,
host_cache_k,
host_cache_v,
device_buffer_k,
device_buffer_v,
top_k_device_locs,
req_pool_indices,
seq_lens,
lru_slots,
num_real_reqs,
page_size,
item_size_bytes,
empty,
empty,
empty,
)
def copy_cache_planned_mla(
*,
miss_src: torch.Tensor,
python/sglang/kernels/ops/kvcache/pd_dcp_gather.py
+14 -8
View File
@@ -1,4 +1,4 @@
from typing import Sequence
from typing import Optional, Sequence
import torch
import triton
@@ -15,10 +15,11 @@ def _copy_mla_rows_into_pack_kernel(
):
layer_id = tl.program_id(0)
block_id = tl.program_id(1)
metadata_offset = layer_id * 3
metadata_offset = layer_id * 4
src = tl.load(src_metadata + metadata_offset).to(pack.dtype)
row_nbytes = tl.load(src_metadata + metadata_offset + 1)
pack_offset = tl.load(src_metadata + metadata_offset + 2)
src_row_stride = tl.load(src_metadata + metadata_offset + 3)
offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
layer_nbytes = num_rows * row_nbytes
@@ -26,7 +27,7 @@ def _copy_mla_rows_into_pack_kernel(
row = offsets // row_nbytes
byte = offsets % row_nbytes
src_row = tl.load(row_indices + row, mask=mask, other=0)
values = tl.load(src + src_row * row_nbytes + byte, mask=mask)
values = tl.load(src + src_row * src_row_stride + byte, mask=mask)
tl.store(pack + pack_offset + offsets, values, mask=mask)
@@ -35,11 +36,14 @@ def copy_mla_rows_into_pack(
row_indices: torch.Tensor,
pack: torch.Tensor,
token_item_lens: Sequence[int],
src_token_item_lens: Optional[Sequence[int]] = None,
) -> None:
if len(kv_data_ptrs) != len(token_item_lens):
if src_token_item_lens is None:
src_token_item_lens = token_item_lens
if not (len(kv_data_ptrs) == len(token_item_lens) == len(src_token_item_lens)):
raise ValueError(
"kv_data_ptrs and token_item_lens length mismatch: "
f"{len(kv_data_ptrs)} vs {len(token_item_lens)}"
"KV pointers, copy widths, and source strides length mismatch: "
f"{len(kv_data_ptrs)}, {len(token_item_lens)}, {len(src_token_item_lens)}"
)
if not kv_data_ptrs:
return
@@ -47,11 +51,13 @@ def copy_mla_rows_into_pack(
n = int(row_indices.numel())
metadata = []
offset = 0
for ptr, item_len in zip(kv_data_ptrs, token_item_lens):
for ptr, item_len, src_item_len in zip(
kv_data_ptrs, token_item_lens, src_token_item_lens
):
item_len = int(item_len)
if item_len <= 0:
raise ValueError(f"MLA token item length must be positive, got {item_len}")
metadata.extend((int(ptr), item_len, offset))
metadata.extend((int(ptr), item_len, offset, int(src_item_len)))
offset += n * item_len
src_metadata = torch.tensor(metadata, dtype=torch.int64, device=pack.device)
python/sglang/kernels/ops/layernorm/mhc.py
+139 -4
View File
@@ -18,10 +18,99 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_interleave
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.utils.common import strict_contiguous
from sglang.srt.runtime_context import get_parallel, get_platform
from sglang.srt.utils import is_gfx95_supported, is_hip
from sglang.srt.utils.common import is_gfx1250_supported
logger = logging.getLogger(__name__)
_AITER_MHC_RUNTIME_DISABLED = False
_AITER_MHC_ACTIVE_LOGGED = False
def _use_aiter_mhc() -> bool:
return (
not _AITER_MHC_RUNTIME_DISABLED
and is_gfx95_supported()
and envs.SGLANG_USE_AITER.get()
)
def _try_aiter_mhc_pre(
residual: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
norm_weight: torch.Tensor | None,
norm_eps: float | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
global _AITER_MHC_RUNTIME_DISABLED, _AITER_MHC_ACTIVE_LOGGED
try:
from aiter.ops.mhc import mhc_pre as aiter_mhc_pre
except Exception as err:
logger.warning("AITER mHC pre is unavailable, falling back: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
kwargs = {}
if norm_weight is not None:
kwargs["norm_weight"] = norm_weight
kwargs["norm_eps"] = norm_eps if norm_eps is not None else rms_eps
try:
result = aiter_mhc_pre(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
**kwargs,
)
except Exception as err:
logger.warning("AITER mHC pre failed, disabling fast path: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
if not _AITER_MHC_ACTIVE_LOGGED:
logger.info("Using AITER gfx950 mHC pre/post kernels")
_AITER_MHC_ACTIVE_LOGGED = True
return result
def _try_aiter_mhc_post(
x: torch.Tensor,
residual: torch.Tensor,
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
) -> torch.Tensor | None:
global _AITER_MHC_RUNTIME_DISABLED
try:
from aiter.ops.mhc import mhc_post as aiter_mhc_post
except Exception as err:
logger.warning("AITER mHC post is unavailable, falling back: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
out = torch.empty_like(residual)
try:
aiter_mhc_post(out, x, residual, post_layer_mix, comb_res_mix)
except Exception as err:
logger.warning("AITER mHC post failed, disabling fast path: %s", err)
_AITER_MHC_RUNTIME_DISABLED = True
return None
return out
# This module is imported during model-registry discovery. Do not import the real
# TileLang package here: it loads native CUDA stubs. The proxy below lets
# module-level @tilelang.jit declarations parse, then imports and applies real
@@ -119,6 +208,24 @@ pass_configs = {
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
}
def _use_deep_gemm_hc_prenorm() -> bool:
if is_hip() or not envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
return False
from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM
return ENABLE_JIT_DEEPGEMM
def _use_tilelang_mhc_pre() -> bool:
return envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get() and not is_hip()
def _use_tilelang_mhc_post() -> bool:
return envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get() and not is_hip()
FP8 = "float8_e4m3"
BF16 = "bfloat16"
FP32 = "float32"
@@ -1041,7 +1148,7 @@ def mhc_pre(
num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
)
if envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
if _use_deep_gemm_hc_prenorm():
n_splits = _compute_num_split_for_mhc_pre(num_tokens, hc_hidden_size)
gemm_out_mul = torch.empty(
@@ -1653,7 +1760,7 @@ def mhc_fused_post_pre(
hidden_size,
)
if envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
if _use_deep_gemm_hc_prenorm():
import deep_gemm
deep_gemm.tf32_hc_prenorm_gemm(
@@ -1847,7 +1954,25 @@ def _mhc_pre_dispatch(
norm_eps: float | None = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
assert residual.dim() == 3, f"residual must be (s, n, h); got {residual.shape}"
if not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
if _use_aiter_mhc():
result = _try_aiter_mhc_pre(
residual=residual,
fn=fn,
hc_scale=hc_scale,
hc_base=hc_base,
rms_eps=rms_eps,
hc_pre_eps=hc_pre_eps,
hc_sinkhorn_eps=hc_sinkhorn_eps,
hc_post_mult_value=hc_post_mult_value,
sinkhorn_repeat=sinkhorn_repeat,
norm_weight=norm_weight,
norm_eps=norm_eps,
)
if result is not None:
post_mix, comb_mix, layer_input = result
return post_mix, comb_mix, layer_input, norm_weight is not None
if not _use_tilelang_mhc_pre():
post_mix, comb_mix, layer_input = _mhc_pre_torch(
residual=residual,
fn=fn,
@@ -1886,7 +2011,17 @@ def _mhc_post_dispatch(
) -> torch.Tensor:
assert x.dim() == 2 and residual.dim() == 3
assert post_layer_mix.dim() == 3 and comb_res_mix.dim() == 3
if not envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get():
if _use_aiter_mhc():
result = _try_aiter_mhc_post(
x=x,
residual=residual,
post_layer_mix=post_layer_mix,
comb_res_mix=comb_res_mix,
)
if result is not None:
return result
if not _use_tilelang_mhc_post():
return _mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
return mhc_post(x, residual, post_layer_mix, comb_res_mix)
python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py
+2 -1
View File
@@ -40,8 +40,9 @@ def _store_sf_interleaved_kernel(
tok_offsets = tok_start + tl.arange(0, BLOCK_T)
mask = tok_offsets < num_tokens
# Load slot indices
# Slot 0 is the reserved CUDA-graph padding sink; skip writes to it.
slots = tl.load(loc_ptr + tok_offsets, mask=mask, other=0)
mask = mask & (slots != 0)
page_offsets = slots % page_size
page_idxs = slots // page_size
python/sglang/kernels/ops/quantization/mxfp8_quant.py
+3
View File
@@ -143,6 +143,9 @@ def _mxfp8_quant_store_qkv_kernel(
tl.store(sfq_ptr + (t * NQ + r) * SF + blk, sf)
else:
myloc = tl.load(loc_ptr + t).to(tl.int64)
# Slot 0 is the reserved CUDA-graph padding sink; skip writes to it.
if myloc == 0:
return
if r < NQ + NKV:
h = r - NQ
cache = kc_ptr
python/sglang/kernels/ops/speculative/cache_locs.py
+43 -39
View File
@@ -77,16 +77,12 @@ def generate_draft_decode_kv_indices(
iter_upper: tl.constexpr,
num_tokens_upper: tl.constexpr,
page_size: tl.constexpr,
window_size: tl.constexpr = 0,
sink_size: tl.constexpr = 0,
NUM_STEPS: tl.constexpr = 0,
):
# Optional token-block parallelism (NUM_STEPS > 0): the first grid axis
# packs (draft step, token block) as ``step + NUM_STEPS * block``,
# spreading the per-request index copy below over many programs instead
# of one program crawling the whole context serially (which bottlenecks
# long-context spec decode, where this kernel runs every iteration).
# NUM_STEPS == 0 (default) is the historical one-program-per-step kernel:
# the same 128-wide copy loop, in the same order, with the token-block
# branches folded away at compile time.
# window_size > 0 restricts the draft (not the target) to sink_size prefix
# tokens + the most-recent window_size; window_size == 0 is the identity.
BLOCK_SIZE: tl.constexpr = 128 if NUM_STEPS == 0 else 512
pid0 = tl.program_id(axis=0)
bid = tl.program_id(axis=1)
@@ -108,45 +104,52 @@ def generate_draft_decode_kv_indices(
kv_indptr += kv_indptr_stride * iters
iters += 1
if NUM_STEPS == 0:
load_offset = tl.arange(0, bs_upper)
seq_lens = tl.load(
paged_kernel_lens + load_offset, mask=load_offset < bid, other=0
)
seq_len = tl.load(paged_kernel_lens + bid)
cum_seq_len = tl.sum(seq_lens)
load_offset = tl.arange(0, bs_upper)
seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0)
seq_len = tl.load(paged_kernel_lens + bid)
if window_size > 0:
cap = window_size + sink_size
seq_lens = tl.minimum(seq_lens, cap)
seq_len_w = tl.minimum(seq_len, cap)
s_eff = tl.minimum(sink_size, seq_len)
recent_start = seq_len - (seq_len_w - s_eff)
else:
seq_len = tl.load(paged_kernel_lens + bid)
num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
# Blocks with no copy work exit before the O(bs) prefix-sum below;
# block 0 always continues (it owns the extension and kv_indptr).
if blk >= num_loop and blk > 0:
return
load_offset = tl.arange(0, bs_upper)
seq_lens = tl.load(
paged_kernel_lens + load_offset, mask=load_offset < bid, other=0
)
cum_seq_len = tl.sum(seq_lens)
seq_len_w = seq_len
s_eff = 0
recent_start = 0
cum_seq_len = tl.sum(seq_lens)
# Update kv_indices
kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters)
kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len_w + iters)
kv_ptr = kv_indices + kv_offset
token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len
num_loop = tl.cdiv(seq_len_w, BLOCK_SIZE)
if NUM_STEPS != 0 and blk >= num_loop and blk > 0:
return
if NUM_STEPS == 0:
kv_offset = tl.arange(0, BLOCK_SIZE)
num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
copy_offset = tl.arange(0, BLOCK_SIZE)
for _ in range(num_loop):
mask = kv_offset < seq_len
data = tl.load(token_pool_ptr + kv_offset, mask=mask)
tl.store(kv_ptr + kv_offset, data, mask=mask)
kv_offset += BLOCK_SIZE
mask = copy_offset < seq_len_w
src = tl.where(
copy_offset < s_eff,
copy_offset,
recent_start + copy_offset - s_eff,
)
data = tl.load(token_pool_ptr + src, mask=mask)
tl.store(kv_ptr + copy_offset, data, mask=mask)
copy_offset += BLOCK_SIZE
else:
for i in range(blk, num_loop, num_blk):
tok_off = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = tok_off < seq_len
data = tl.load(token_pool_ptr + tok_off, mask=mask)
tl.store(kv_ptr + tok_off, data, mask=mask)
copy_offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = copy_offset < seq_len_w
src = tl.where(
copy_offset < s_eff,
copy_offset,
recent_start + copy_offset - s_eff,
)
data = tl.load(token_pool_ptr + src, mask=mask)
tl.store(kv_ptr + copy_offset, data, mask=mask)
# Extension entries and kv_indptr belong to token block 0 alone; other
# blocks neither compute nor store them.
@@ -178,18 +181,19 @@ def generate_draft_decode_kv_indices(
)
tl.store(
kv_ptr + seq_len + extend_offset,
kv_ptr + seq_len_w + extend_offset,
extend_data,
mask=extend_offset < iters,
)
# Update kv_indptr
bs_offset = tl.arange(0, num_tokens_upper)
zid = bid * topk + topk_id
if zid == 0:
zid = num_seqs * topk
pos_vals = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0)
if window_size > 0:
pos_vals = tl.minimum(pos_vals, window_size + sink_size)
base = tl.sum(pos_vals)
tl.store(kv_indptr + zid, base + zid * iters)
python/sglang/multimodal_gen/runtime/disaggregation/orchestrator.py
+1
View File
@@ -63,6 +63,7 @@ def _deserialize_request_metrics(data: dict | None) -> RequestMetrics | None:
metrics = RequestMetrics(request_id=data["request_id"])
metrics.stages = data.get("stages", {})
metrics.denoising_stages = set(data.get("denoising_stages", ()))
metrics.steps = data.get("steps", [])
metrics.total_duration_ms = data.get("total_duration_ms", 0.0)
for name, snapshot in data.get("memory_snapshots", {}).items():
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
+22 -11
View File
@@ -97,9 +97,12 @@ def init_world_group(
def _sync_srt_world_group() -> None:
import sglang.srt.distributed.parallel_state as srt_parallel_state
from sglang.srt.runtime_context import get_parallel
if srt_parallel_state._WORLD is None:
srt_parallel_state._WORLD = _WORLD
if srt_parallel_state._WORLD is _WORLD:
get_parallel().override_permanently(world_group=_WORLD)
def _clear_srt_world_group() -> None:
@@ -110,18 +113,11 @@ def _clear_srt_world_group() -> None:
def _sync_srt_tp_group() -> None:
"""Lend this package's TP group to `srt`, and state the widths it implies.
"""Expose this package's TP group, widths, and ranks to shared SRT layers.
Shared `srt` layers run in this package and ask `get_parallel()` for how to
shard -- `srt/layers/attention/vision.py` reads `attn_tp_size`. The
published `srt` config cannot answer: `gpu_worker.py` publishes a dummy
carrying *this* package's `tp_size`, which a sequence-parallel launch sets
to 1 while the group lent here is as wide as the world. So the widths are
permanently overridden alongside the group -- this runs with no `srt`
config published at all, which is exactly why it cannot go through
`RuntimeContext.override` (it requires one).
Only tensor parallelism folds this way, so every other dimension is one.
Use the group's actual width: sequence parallelism can make it wider than
the TP size in the dummy SRT configuration. Other parallel dimensions are
one. Overrides also work before SRT configuration is published.
"""
import sglang.srt.distributed.parallel_state as srt_parallel_state
from sglang.srt.runtime_context import derive_parallel_widths, get_parallel
@@ -132,6 +128,15 @@ def _sync_srt_tp_group() -> None:
srt_parallel_state._ATTN_TP = _TP
if srt_parallel_state._ATTN_TP is _TP:
get_parallel().override_permanently(
tp_group=_TP,
attn_tp_group=_TP,
tp_size=_TP.world_size,
tp_rank=_TP.rank_in_group,
attn_tp_rank=_TP.rank_in_group,
moe_tp_rank=_TP.rank_in_group,
attn_cp_rank=0,
pp_rank=0,
moe_ep_rank=0,
**derive_parallel_widths(
tp_size=_TP.world_size,
attn_cp_size=1,
@@ -151,6 +156,9 @@ def _clear_srt_tp_group() -> None:
if srt_parallel_state._ATTN_TP is _TP:
srt_parallel_state._ATTN_TP = None
get_parallel().clear_stamp()
if srt_parallel_state._WORLD is not None:
# Restore the still-active WORLD handle after clearing TP overrides.
get_parallel().override_permanently(world_group=srt_parallel_state._WORLD)
if srt_parallel_state._TP is _TP:
srt_parallel_state._TP = None
@@ -691,6 +699,9 @@ def use_tensor_parallel_group(tp_group: GroupCoordinator):
tp_size=tp_group.world_size,
tp_rank=tp_group.rank_in_group,
tp_group=tp_group,
attn_tp_group=tp_group,
attn_tp_rank=tp_group.rank_in_group,
moe_tp_rank=tp_group.rank_in_group,
# Only tensor parallelism folds here, so every other dimension is
# one and the quotients come out of the shared derivation.
**derive_parallel_widths(
python/sglang/multimodal_gen/runtime/launch_server.py
+1
View File
@@ -211,6 +211,7 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True):
reader.close()
logger.debug("All workers are ready")
logger.info("[server-load] workers_ready_monotonic_ns=%d", time.monotonic_ns())
if node_rank != 0:
# The TokenizerManager / HTTP surface lives on the node that owns
python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py
+3 -1
View File
@@ -279,7 +279,9 @@ class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
if not isinstance(module, LayerwiseOffloadableModuleMixin):
return
for manager in module.layerwise_offload_managers:
manager.release_all()
# Not release_all: this is a use ending, not a reset. The default
# still drops the resident set, so behaviour is unchanged here.
manager.release_after_use()
# The layers are gone; the rest of this component is dead weight on the
# device until it is used again, and the stage that follows may be the
# one that needs the room.
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
+28 -5
View File
@@ -2121,10 +2121,30 @@ class LayerwiseOffloadManager:
torch.mps.synchronize()
torch.mps.empty_cache()
@torch.compiler.disable
def release_after_use(self, *, keep_resident: bool = False) -> None:
"""This component's use has ended; release what that use was streaming.
Distinct from `release_all`, which is the literal operation and stays
that way for a full reset. A use ending asks a narrower question: the
streamed window is certainly dead, but the resident set only is if
nothing will want it before something else needs the room.
The two were the same call, and that is why `resident_layers` does
nothing for any component whose use is a single forward pass rather
than a denoise loop -- the set is prefetched at the start of the use
and dropped at the end of it, every request. `keep_resident` is how a
caller that knows the memory picture says otherwise; it defaults to the
long-standing behaviour, so nothing moves until someone asks.
"""
self._release_layers(drop_resident=not keep_resident)
@torch.compiler.disable
def release_all(self) -> None:
"""Release every layer, including the resident ones: this ends the
denoise stage that the resident set is scoped to."""
"""Release every layer, resident ones included. A full reset."""
self._release_layers(drop_resident=True)
def _release_layers(self, *, drop_resident: bool) -> None:
self._log_direct_read_summary()
self._log_debug_timing()
if self._mapped_populator is not None:
@@ -2140,10 +2160,13 @@ class LayerwiseOffloadManager:
self._collect_mapped_layer(layer_idx)
for layer_idx in list(self._gpu_layers):
self.release_layer(layer_idx, force=True)
# `force` is what overrides release_layer's own skip of the resident
# set, so not forcing is all it takes to leave that set alone.
self.release_layer(layer_idx, force=drop_resident)
# The next use starts a new request; its first pass over the layers may
# find their pages evicted and is the one worth faulting in sequentially.
self._first_pass = True
# find their pages evicted and is the one worth faulting in
# sequentially. Layers still on the device were never evicted.
self._first_pass = drop_resident
@torch.compiler.disable
def load_all_layers(self) -> None:
python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py
+15 -6
View File
@@ -381,8 +381,9 @@ class QwenImage21Attention(nn.Module):
class QwenImage21TransformerBlock(nn.Module):
def __init__(self, ac, quant_config, prefix):
def __init__(self, ac, quant_config, prefix, layer_id):
super().__init__()
self._layer_id = layer_id
self.img_norm1 = nn.LayerNorm(
ac.hidden_size, eps=ac.eps, elementwise_affine=False
)
@@ -404,6 +405,9 @@ class QwenImage21TransformerBlock(nn.Module):
ropes,
caches,
):
# Cache-DiT's UnifiedBlocks forwards the same args to every layer.
# Slice here so prefix KV stays per-layer after that wrap.
caches = [cache[self._layer_id] for cache in caches]
scale1, gate1, scale2, gate2 = modulation
prefixes = [
apply_modulation(
@@ -482,9 +486,12 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
self.modulation = nn.Sequential(
nn.SiLU(), nn.Linear(ac.hidden_size, ac.hidden_size * 4, bias=False)
)
self.num_layers = ac.num_layers
self.transformer_blocks = nn.ModuleList(
[
QwenImage21TransformerBlock(ac, quant_config, f"transformer_blocks.{i}")
QwenImage21TransformerBlock(
ac, quant_config, f"transformer_blocks.{i}", i
)
for i in range(ac.num_layers)
]
)
@@ -528,7 +535,7 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
)
prefix_modulation = self.prepare_modulation(zero_temb)
if prefix_caches is None:
prefix_caches = [[None] * len(self.transformer_blocks) for _ in layouts]
prefix_caches = [[None] * self.num_layers for _ in layouts]
prefix_states, ropes = [], []
for sample, layout in enumerate(layouts):
prefix = None
@@ -544,8 +551,10 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
)
prefix_states.append({"hidden_states": prefix})
ropes.append(layout["target_rope"][start:end])
# visit each block once so layerwise offload transfers weights once per batch
for i, block in enumerate(self.transformer_blocks):
# Same extras for every block so Cache-DiT's UnifiedBlocks wrap is valid.
# Each block slices prefix_caches by _layer_id. Visit once per layer for
# layerwise offload.
for block in self.transformer_blocks:
images = block(
images,
modulation,
@@ -553,7 +562,7 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
prefix_modulation,
layouts,
ropes,
[cache[i] for cache in prefix_caches],
prefix_caches,
)
output = self.proj_out(self.norm_out(images, temb))
if sp > 1:
python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_qwen3.py
-2
View File
@@ -29,7 +29,6 @@ from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, can_return_tuple
from transformers.utils.deprecation import deprecate_kwarg
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.srt.layers.layernorm import RMSNorm
from .transformers_compat import (
@@ -347,7 +346,6 @@ def make_qwen3_rms_norm(hidden_size: int, eps: float) -> RMSNorm:
hidden_size,
eps=eps,
cast_x_before_out_mul=True,
force_native=not current_platform.is_npu(),
)
python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage21.py
+3
View File
@@ -197,6 +197,9 @@ class QwenImage21RMS_norm(nn.Module):
class QwenImage21Upsample(nn.Upsample):
def forward(self, x):
# Nearest interpolation copies values; no FP32 arithmetic is needed.
if self.mode == "nearest-exact" and x.dtype in (torch.float16, torch.bfloat16):
return super().forward(x)
return super().forward(x.float()).type_as(x)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py
+2
View File
@@ -417,6 +417,8 @@ class PipelineStage(StageDedupMixin, ABC):
# Execute the actual stage logic with unified profiling.
previous_batch_is_warmup = self._current_batch_is_warmup
metrics = batch.metrics
if metrics is not None and self.role_affinity == RoleType.DENOISER:
metrics.denoising_stages.add(stage_name)
warmup_metrics = metrics if batch.is_warmup else None
previous_active_stage = (
warmup_metrics.active_stage_name if warmup_metrics is not None else None
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py
+4
View File
@@ -19,6 +19,7 @@ from PIL import Image
from torch import nn
from transformers import PreTrainedTokenizerBase
from sglang.kernels.ops.diffusion import load_mesh_processor
from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
Hunyuan3D2PipelineConfig,
@@ -790,6 +791,9 @@ class Hunyuan3DPaintPostprocessStage(PipelineStage):
def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
del server_args
if batch.is_warmup:
# compile without exporting warmup meshes or textures
load_mesh_processor()
if batch.is_warmup or batch.extra.get("_mesh_failed"):
return OutputBatch(output_file_paths=[], metrics=batch.metrics)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image21.py
+7
View File
@@ -234,6 +234,13 @@ def prepare_qwen21_mu(batch, server_args):
class QwenImage21DenoisingStage(DenoisingStage):
def _bcg_pad_prompt_kwargs(
self, call_kwargs, current_model=None, force_bucket=None
):
# Prefill runs eagerly. Later steps use exact-length prefix KV, so text
# padding only creates duplicate graphs without enabling more replay.
return call_kwargs
def _predict_noise(
self,
current_model,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py
+9 -2
View File
@@ -8,6 +8,7 @@ This module contains implementations of timestep preparation stages for diffusio
"""
import inspect
import logging
from dataclasses import dataclass
from typing import Any, Callable, Tuple
@@ -161,8 +162,14 @@ class TimestepPreparationStage(PipelineStage):
# Update batch with prepared timesteps
batch.timesteps = timesteps
batch.scheduler = scheduler
if not batch.is_warmup:
self.log_debug("timesteps: %s", timesteps)
if not batch.is_warmup and logger.isEnabledFor(logging.DEBUG):
# format on cpu to avoid first-use cuda kernels in tensor repr
logger.debug(
"[%s] timesteps (%s): %s",
self.__class__.__name__,
timesteps.device,
timesteps.detach().cpu(),
)
return batch
def build_dedup_fingerprint(
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
+7 -1
View File
@@ -58,6 +58,7 @@ class RequestMetrics:
def __init__(self, request_id: str):
self.request_id = request_id
self.stages: Dict[str, float] = {}
self.denoising_stages: set[str] = set()
self.steps: list[float] = []
self.steps_by_stage: Dict[str, list[float]] = {}
self.stage_iterations: Dict[str, tuple[int, int]] = {}
@@ -112,6 +113,7 @@ class RequestMetrics:
return {
"request_id": self.request_id,
"stages": self.stages,
"denoising_stages": sorted(self.denoising_stages),
"steps": self.steps,
"total_duration_ms": self.total_duration_ms,
"memory_snapshots": {
@@ -461,7 +463,11 @@ class PerformanceLogger:
Note that this accords to the time spent internally in server, postprocess is not included
"""
formatted_stages = [
{"name": name, "execution_time_ms": duration_ms}
{
"name": name,
"execution_time_ms": duration_ms,
"is_denoising": name in metrics.denoising_stages,
}
for name, duration_ms in metrics.stages.items()
]
python/sglang/multimodal_gen/runtime/warmup_request_builder.py
+3 -4
View File
@@ -365,12 +365,11 @@ def _resolve_warmup_num_frames(
if num_frames is None:
return num_frames
# Breakable CUDA graph replays only exact latent shapes: the warmup
# request must run the full serving frame count so its captured graphs
# match serving signatures (mirrors the uncapped-steps rule in
# _resolve_warmup_steps).
# explicit frame counts and breakable CUDA graphs must keep the requested
# latent shape; only default server warmup applies the bounded frame cap
if (
not server_based_warmup
or isinstance(explicit_num_frames, int)
or getattr(server_args, "enable_breakable_cuda_graph", False) is True
):
warmup_num_frames = num_frames
python/sglang/multimodal_gen/test/runner/PERFORMANCE.md
+196
View File
@@ -0,0 +1,196 @@
# CI performance guards and baselines
## Metric and failure contracts
Every generated testcase must report finite, positive E2E, including cases with
`run_perf_check=False`. Missing request records, absent performance logs and
missing/invalid E2E fail CI. `run_perf_check=False` disables stage/step and
memory checks, not the request's E2E threshold guard. Baseline generation still
requires valid E2E but skips baseline comparisons. Explicit GT generation skips
validation.
A performance failure stops the testcase's remaining repeated requests and
subsequent checks for that attempt. Performance failures use the existing
pytest retry budget (at most six retries), rerunning only failed items.
Exhausting the budget still fails CI; missing metrics and exceeded thresholds
are never treated as passing. Consistency failures remain non-retryable.
Standalone infrastructure failures retain their existing retry policy.
Valid failed measurements are recorded
before threshold validation; realtime chunk and memory guards remain enabled
according to their existing configuration.
## B200 runner baselines
`b200.json` keeps the existing Verda references as its defaults. Its
`runner_overrides` map applies metric overrides by the GitHub `RUNNER_NAME`
prefix. DeepInfra runners (`b200-di*`) use separate E2E references for the two
cases below; unknown runners keep the defaults. Loading, stage/step and memory
references, other cases, and the 25% E2E tolerance are unchanged.
| Case | Default E2E (ms) | DeepInfra E2E (ms) |
| --- | ---: | ---: |
| `flux1_modelopt_nvfp4_t2i` | 836.71 | 1334.16 |
| `qwen_image_2512_modelopt_nvfp4_t2i` | 9650.06 | 16126.87 |
The pool mismatch was observed in [B200 CI job 103856482654](https://github.com/sgl-project/sglang/actions/runs/34805428031/job/103856482654).
The DeepInfra references are the medians of three unprofiled, warmed requests
using the CI case configuration at commit
`9b7e11f32b88d4c1bfd9cf44550a30a702dd9df8`:
Flux: 1367.47, 1327.48, 1334.16 ms; Qwen: 16126.87, 16806.36, 15240.80 ms.
The matching Verda measurements were 811.26, 771.97, 783.31 ms and
9169.16, 9045.49, 9071.78 ms, respectively. These calibrate the runner pools;
they do not establish a model-level root cause for the difference.
When refreshing a pool-specific reference, update its `runner_overrides` entry,
not the shared `scenarios` entry. The baseline generation script writes shared
scenarios; use a separate `--out` file when collecting pool-specific candidates.
### Cirrascale historical CI reference
`b200-cirrascale1-0123` has separate E2E references of 1574.32 ms for
`flux1_modelopt_nvfp4_t2i` and 17742.04 ms for
`qwen_image_2512_modelopt_nvfp4_t2i`. The measured Cirrascale 3 runners below also
have separate references. Other Cirrascale runners retain the defaults until
calibrated.
Historical jobs on this runner, all using driver 580.126.20, already recorded
the slower timings before this PR's changes, with unchanged B200 case definitions:
| PR / CI job | Flux1 E2E (ms) | Qwen2512 E2E (ms) |
| --- | ---: | ---: |
| [#39021](https://github.com/sgl-project/sglang/actions/runs/34594936520/job/103248571918) | 1772.69 | 17836.93 |
| [#38782](https://github.com/sgl-project/sglang/actions/runs/34595934993/job/103251774792) | 2617.81 | 17742.04 |
| [#39022](https://github.com/sgl-project/sglang/actions/runs/34670636885/job/103506632867) | 1574.32 | 38392.34 |
| [#39291](https://github.com/sgl-project/sglang/actions/runs/34750152864/job/103707842024) | 3942.31 | 19846.93 |
Use the minimum observed E2E for each case, not the noisy maximum or a median
inflated by slow runs. The existing 25% tolerance yields limits of 1967.90 ms
and 22177.55 ms. Large transient slowdowns must still fail and use the bounded
failed-item retry policy. Historical green jobs did not enforce the new E2E
guard; their recorded timings, not their green status, support these references.
This does not identify the underlying host/GPU contention mechanism. Loading,
other metrics, other cases, and other runner references remain unchanged.
### Cirrascale 3 historical CI references
Match `b200-cirrascale3-0123` and `b200-cirrascale3-4567` separately, without
extending the override to unmeasured runners. Their two NVFP4 case definitions,
sampling configuration and model implementations are unchanged in the historical
comparisons below. These are warmed request timings, not model download or load
times. The independent PRs did not include this PR's E2E guard changes.
| PR / CI job | Runner suffix | Flux1 E2E (ms) | Qwen2512 E2E (ms) |
| --- | --- | ---: | ---: |
| [#40265](https://github.com/sgl-project/sglang/actions/runs/35437052621/job/105881489389) | `3-0123` | 1470.24 | 35377.30 |
| [#39206, earlier head](https://github.com/sgl-project/sglang/actions/runs/35433818034/job/105873136034) | `3-0123` | 1471.59 | 17894.49 |
| [#40293](https://github.com/sgl-project/sglang/actions/runs/35433914839/job/105879952157) | `3-4567` | 1547.77 | 17346.65 |
| [#39983](https://github.com/sgl-project/sglang/actions/runs/35448201646/job/105999591501) | `3-4567` | 1471.19 | 18294.78 |
| [#40374](https://github.com/sgl-project/sglang/actions/runs/35465121078/job/105956031846) | `3-4567` | 1500.43 | 18433.93 |
The earlier #39206 row uses the minimum of its seven attempts, not their noisy
maximum. Apply the same minimum-observed rule per runner across these records:
1470.24 / 17894.49 ms for `3-0123`, and 1471.19 / 17346.65 ms for `3-4567`.
Keep the 25% tolerance and all other metrics unchanged. In particular, the
35-second historical Qwen outlier still fails; it is not a new reference.
These records establish a pre-existing runner-specific mismatch with the Verda
reference, not the underlying cause of contention or a claim that all runs pass.
## Initial loading references
The initial H100 loading references cover 28 cases with at least three distinct
CI runs whose maximum/minimum startup-time ratio is at most 1.25. Each reference
is the minimum measured `load_time_ms`, rounded to two decimals; repeated requests
sharing one server do not count as separate startups. These are process-start to
all-workers-ready measurements, excluding warmup, not checkpoint-I/O-only times.
Sources are PR Test Base runs [34750203901](https://github.com/sgl-project/sglang/actions/runs/34750203901),
[34751664379](https://github.com/sgl-project/sglang/actions/runs/34751664379),
[34753876622](https://github.com/sgl-project/sglang/actions/runs/34753876622), and
[34755864886](https://github.com/sgl-project/sglang/actions/runs/34755864886).
Only saved valid loading measurements are used; this does not claim those runs
passed all other checks. No E2E reference or tolerance is raised. Cross-run
stability is a conservative selection criterion, not proof of an optimal load
time. The first batches left variable cases uncalibrated; the best-observed
references below now give those cases a loading guard without claiming stable
runtime. B200 and 5090 references are not inferred from H100 or development H200
measurements.
The six RTX 5090 loading references use the same selection rule, from runs
[34753876622](https://github.com/sgl-project/sglang/actions/runs/34753876622),
[34755864886](https://github.com/sgl-project/sglang/actions/runs/34755864886), and
[34764411082](https://github.com/sgl-project/sglang/actions/runs/34764411082).
Their per-case maximum/minimum ratios range from 1.048 to 1.203. The last run's
six recorded requests passed E2E validation and failed the then-missing loading
baseline check; subsequent checks and MiniMax's second request did not run.
The existing MiniMax wall-clock tolerance override is unchanged.
The same three runs also provide initial two-H100 loading references for
`ltx_2.3_one_stage_ti2v`, `ltx_2.3_two_stage_t2v_2gpus`,
`wan2_1_t2v_1.3b_cfg_parallel`, and `zimage_image_t2i_2_gpus`.
Seven partition-0 cases also meet this criterion: `flux2_modelopt_fp8_tp2_t2i`,
`flux_image_t2i_2_gpus`, `ideogram4_fp8_tp2_t2i`, `qwen_image_t2i_2_gpus`,
`wan2_2_i2v_a14b_2gpu`, `wan2_2_t2v_a14b_lora_2gpu`, and
`wan2_2_t2v_a14b_teacache_2gpu`. All eleven cases meet the same three-run
stability criterion. Existing loading references are not raised when a later
run exceeds their limits.
Five more two-H100 references use runs
[34755864886](https://github.com/sgl-project/sglang/actions/runs/34755864886),
[34764411082](https://github.com/sgl-project/sglang/actions/runs/34764411082), and
[34766506722](https://github.com/sgl-project/sglang/actions/runs/34766506722):
`flux_2_image_t2i_2_gpus`, `fsdp-inference`, `mova_360p_tp2`,
`wan2_1_i2v_14b_720P_2gpu`, and `zimage_image_t2i_2_gpus_non_square`.
Their maximum/minimum loading ratios range from 1.093 to 1.205. Each reference
is the minimum observed loading time rounded to two decimal places; the same
existing tolerances apply. These are initial references, not claims that the
full testcases or all later checks passed.
Seven further H100 references use that same recent three-run window:
`flux_image_t2i`, `flux_2_ti2i`, `joyai_image_edit_ti2i`,
`qwen_image_edit_2509_ti2i`, `qwen_image_layered_i2i`,
`minimax_h3_t2va_2gpu_h100`, and `qwen_image_edit_modelopt_fp8_ti2i`.
Their maximum/minimum ratios in this window range from 1.006 to 1.250
(the largest unrounded ratio is 1.249565). Earlier historical runs vary more;
these references do not claim stability across the entire history. Each value
is the minimum in the stated window, rounded to two decimals. Existing loading
references, E2E references, and tolerances are unchanged.
Four additional single-H100 references follow the same minimum-of-three rule.
`flux_2_image_t2i_upscaling_4x` uses runs 34764411082, 34766506722, and
[34770008768](https://github.com/sgl-project/sglang/actions/runs/34770008768).
`flux_2_t2i_customized_vae_path`, `flux_2_ti2i_multi_image_cache_dit`, and
`zimage_image_t2i` use runs 34766506722, 34770008768, and
[34771349145](https://github.com/sgl-project/sglang/actions/runs/34771349145).
Their maximum/minimum ratios range from 1.090 to 1.224. These are initial
loading references only; earlier variable runs remain diagnostic evidence,
and no existing performance reference or tolerance is increased.
`qwen_image_edit_ti2i` and `lingbot_world_realtime_plastic_beach` use runs
34764411082, 34766506722, and 34771349145. Their minimum loading times are
36793.38 ms and 31600.49 ms, respectively, with maximum/minimum ratios of
1.1745 and 1.1256. The realtime case uses the same process-start-to-ready
loading boundary, excluding warmup; its request E2E reference is unchanged.
### Best-observed references for variable cases
The remaining 16 H100 and five B200 cases use the fastest valid startup in
the saved reports from runs 34755864886, 34764411082, 34766506722,
34770008768, and 34771349145, where available. Each value is rounded to two
decimals. H100 cases have three to five distinct startups; B200 cases have
two. These are initial measured references, not claims of stability or proof
that every historical run passed. Requiring all noisy samples to converge
before establishing a guard would leave these cases without a quantified limit.
No existing reference or tolerance is increased. Samples above the resulting
limit still fail, and their infrastructure/code diagnosis remains separate.
In particular, LTX HQ retains the observed 116.531 s startup as its reference,
not the later 187-208 s startups. Downloads in H200 development measurements
are not used to establish either GPU pool's loading references.
| Source run | GPU | Cases supplying the minimum |
| --- | --- | --- |
| 34755864886 | H100 | `fast_hunyuan_video`, `joy_echo_t2v_2gpu`, `wan2_1_i2v_14b_480P_2gpu`, `wan2_1_t2v_14b_2gpu`, `wan2_2_t2v_a14b_2gpu` |
| 34764411082 | H100 | `lingbot_video_moe_t2v`, `ltx_2_3_hq_pipeline`, `qwen_image_t2i_2_gpus_extra_high`, `wan22_modelopt_fp8_t2v` |
| 34766506722 | H100 | `ltx_2_3_two_stage_ti2v_2gpus`, `ltx_2_5_diffusion_decoder_2gpus`, `ltx_2_two_stage_t2v`, `minimax_h3_ref2va_video_audio_2gpu_h100`, `sana_wm_ti2v`, `wan2_1_i2v_14b_lora_2gpu`, `wan2_1_t2v_1_3b_cache_dit_sp_only_2gpu` |
| 34755864886 | B200 | `flux1_modelopt_nvfp4_t2i`, `flux2_modelopt_nvfp4_t2i` |
| 34764411082 | B200 | `ideogram4_nvfp4_t2i`, `qwen_image_2512_modelopt_nvfp4_t2i`, `wan22_modelopt_nvfp4_t2v` |
python/sglang/multimodal_gen/test/runner/pytest_runner.py
+38 -13
View File
@@ -1,7 +1,9 @@
from __future__ import annotations
import os
import subprocess
import sys
import time
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Sequence
@@ -104,7 +106,6 @@ def _run_pytest_attempt(cmd: list[str]) -> tuple[int, str]:
stderr=subprocess.STDOUT,
bufsize=0,
)
output_bytes = bytearray()
while True:
chunk = process.stdout.read(4096)
@@ -118,6 +119,18 @@ def _run_pytest_attempt(cmd: list[str]) -> tuple[int, str]:
return process.returncode, output_bytes.decode("utf-8", errors="replace")
def _estimate_failed_test_time(xml_path: str | None, attempt_time: float) -> float:
if xml_path is None or not Path(xml_path).exists():
return attempt_time
failed_time = sum(
float(testcase.get("time", "0"))
for testcase in ET.parse(xml_path).getroot().iter("testcase")
if testcase.find("failure") is not None or testcase.find("error") is not None
)
return failed_time if failed_time > 0 else attempt_time
def _extract_collection_line(full_output: str) -> str | None:
for line in full_output.splitlines():
stripped = line.strip()
@@ -157,8 +170,7 @@ def _summary_has_retryable_failure(summary_lines: list[str]) -> bool:
for line in summary_lines:
lowered = line.lower()
if (
"[performance]" in line
or "SafetensorError" in line
"SafetensorError" in line
or "FileNotFoundError" in line
or "TimeoutError" in line
or "out of memory" in lowered
@@ -185,11 +197,10 @@ def _is_retryable_failure(full_output: str) -> bool:
if _is_consistency_failure(full_output):
return False
if "[performance]" in full_output:
return True
summary_lines = _extract_short_test_summary(full_output)
is_perf_assertion = (
"multimodal_gen/test/server/test_server_utils.py" in full_output
and "AssertionError" in full_output
)
is_aggregated_retryable_failure = _summary_has_retryable_failure(summary_lines)
is_flaky_ci_assertion = (
@@ -202,12 +213,7 @@ def _is_retryable_failure(full_output: str) -> bool:
"out of memory" in full_output.lower() or "oom killer" in full_output.lower()
)
return (
is_perf_assertion
or is_aggregated_retryable_failure
or is_flaky_ci_assertion
or is_oom_error
)
return is_aggregated_retryable_failure or is_flaky_ci_assertion or is_oom_error
def _print_attempt_tail_summary(
@@ -284,6 +290,8 @@ def run_pytest(
base_cmd.extend(["-k", filter_expr])
max_retries = 6
retry_deadline = os.environ.get("SGLANG_DIFFUSION_RETRY_DEADLINE")
retry_deadline = float(retry_deadline) if retry_deadline else None
attempt_reports = []
for i in range(max_retries + 1):
is_retry = i > 0
@@ -298,7 +306,9 @@ def run_pytest(
f"for {len(files)} assigned item(s)"
)
attempt_start = time.monotonic()
returncode, full_output = _run_pytest_attempt(cmd)
attempt_time = time.monotonic() - attempt_start
retryable = returncode not in (0, 5) and _is_retryable_failure(full_output)
attempt_reports.append(
{
@@ -343,6 +353,21 @@ def run_pytest(
_print_attempt_tail_summary(attempt_reports, len(files))
return (returncode, list(all_executed_cases), all_case_results)
if retry_deadline is not None:
remaining = retry_deadline - time.time()
retry_estimate = _estimate_failed_test_time(junit_xml_path, attempt_time)
# leave headroom for pytest startup and variation in the failed cases
required = retry_estimate * 1.1 + 30
if remaining < required:
print(
f"Retry budget exhausted: {remaining:.1f}s remaining, "
f"next failed-item retry needs approximately {required:.1f}s. "
"Preserving the failing result instead of starting another attempt.",
flush=True,
)
_print_attempt_tail_summary(attempt_reports, len(files))
return (returncode, list(all_executed_cases), all_case_results)
print(
f"Retryable failure detected on attempt {i + 1}. "
"Retrying only previously failed items."
python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py
+5
View File
@@ -1,6 +1,7 @@
import argparse
import inspect
import json
import math
import os
import re
import sys
@@ -137,6 +138,9 @@ def _run_case(case: DiffusionTestCase) -> dict:
perf = PerformanceSummary.from_req_perf_record(
rec, BASELINE_CONFIG.step_fractions
)
for name, value in (("load", ctx.load_time_ms), ("E2E", perf.e2e_ms)):
if value is None or not (math.isfinite(value) and value > 0):
raise ValueError(f"{case.id}: {name} duration missing or invalid")
if case.server_args.modality == "video" and sp.num_frames and sp.num_frames > 0:
if "per_frame_generation" not in perf.stage_metrics:
perf.stage_metrics["per_frame_generation"] = perf.e2e_ms / sp.num_frames
@@ -147,6 +151,7 @@ def _run_case(case: DiffusionTestCase) -> dict:
str(k): round(v, 2) for k, v in perf.all_denoise_steps.items()
},
"expected_e2e_ms": round(perf.e2e_ms, 2),
"expected_load_ms": round(ctx.load_time_ms, 2),
"expected_avg_denoise_ms": round(perf.avg_denoise_ms, 2),
"expected_median_denoise_ms": round(perf.median_denoise_ms, 2),
}
python/sglang/multimodal_gen/test/server/ascend/conftest.py
+14 -1
View File
@@ -1,4 +1,7 @@
"""Ascend NPU conftest: evict stale model page cache before each test case.
"""Ascend NPU fixtures for performance validation and model page cache.
NPU performance guards cover inference latency, not model loading latency.
Loading times are still collected and reported by the shared test harness.
Memory-capped CI runners (e.g. a 128 GiB cgroup on the 4-NPU A3 pool) count
reclaimable page cache from previously loaded models in cgroup
@@ -24,6 +27,16 @@ _CGROUP_V2_CURRENT = "/sys/fs/cgroup/memory.current"
_CGROUP_V1_USAGE = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
@pytest.fixture(autouse=True)
def _disable_load_latency_validation(monkeypatch):
"""Disable loading-latency comparisons only for tests in this directory."""
from sglang.multimodal_gen.test.server.test_server_utils import PerformanceValidator
monkeypatch.setattr(
PerformanceValidator, "validate_load", lambda self, summary: None
)
def _read_cgroup_memory_current() -> str:
"""Best-effort read of the container's cgroup memory usage in bytes."""
for path in (_CGROUP_V2_CURRENT, _CGROUP_V1_USAGE):
python/sglang/multimodal_gen/test/server/gpu_cases.py
+31 -22
View File
@@ -49,6 +49,7 @@ from sglang.multimodal_gen.test.test_utils import (
DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST,
DEFAULT_JOYAI_IMAGE_EDIT_MODEL_NAME_FOR_TEST,
DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_21_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST,
@@ -115,6 +116,7 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
),
PI05_ACTION_CI_sampling_params,
run_perf_check=False,
perf_warmup_requests=1,
run_component_accuracy_check=False,
run_t2v_input_reference_check=False,
),
@@ -181,6 +183,10 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
DiffusionServerArgs(
model_path=DEFAULT_COSMOS3_NANO_MODEL_NAME_FOR_TEST,
modality="image",
extras=[
"--warmup-num-frames 1",
"--component-residency transformer=resident",
],
),
COSMOS3_NANO_CI_sampling_params,
run_perf_check=False,
@@ -254,6 +260,8 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
DiffusionServerArgs(
model_path=DEFAULT_COSMOS3_NANO_MODEL_NAME_FOR_TEST,
modality="video",
# the latency baseline measures the warmed, resident transformer
extras=["--component-residency transformer=resident"],
env_vars={"SGLANG_DISABLE_COSMOS3_GUARDRAILS": "1"},
),
DiffusionSamplingParams(
@@ -398,6 +406,7 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
"sana_wm_ti2v",
DiffusionServerArgs(
model_path=DEFAULT_SANA_WM_STREAMING_MODEL_NAME_FOR_TEST,
extras=["--warmup-resolutions 384x640"],
),
SANA_WM_TI2V_CI_sampling_params,
run_perf_check=False,
@@ -1091,6 +1100,9 @@ TWO_GPU_CASES = [
# decoder headroom on 80 GB GPUs.
extras=[
"--load-diffusion-decoder",
"--warmup-resolutions 768x448",
"--warmup-num-frames 49",
"""--warmup-sampling-params '{"use_diffusion_decoder":true}'""",
"--component-residency "
"transformer=component-offload,text_encoder=component-offload",
],
@@ -1102,7 +1114,6 @@ TWO_GPU_CASES = [
expect_audio_output=True,
extras={"seed": 42, "use_diffusion_decoder": True},
),
run_perf_check=False,
run_component_accuracy_check=False,
),
# I2V LoRA test case
@@ -1131,26 +1142,25 @@ TWO_GPU_CASES = [
ring_degree=2,
),
),
# TODO: re-enable when the checkpoint is accessible to fork PR CI
# DiffusionTestCase(
# "qwen_image21_t2i_tp2",
# DiffusionServerArgs(
# model_path="Qwen/Qwen-Image-2.1",
# tp_size=2,
# ulysses_degree=1,
# ring_degree=1,
# ),
# replace(
# T2I_sampling_params,
# output_size="1024x1024",
# output_format="png",
# extras={"num_inference_steps": 40, "guidance_scale": 1, "seed": 42},
# ),
# perf_repeat_requests=2,
# run_perf_check=False,
# run_component_accuracy_check=False,
# run_t2v_input_reference_check=False,
# ),
DiffusionTestCase(
"qwen_image21_t2i_tp2",
DiffusionServerArgs(
model_path=DEFAULT_QWEN_IMAGE_21_MODEL_NAME_FOR_TEST,
tp_size=2,
ulysses_degree=1,
ring_degree=1,
),
replace(
T2I_sampling_params,
output_size="1024x1024",
output_format="png",
extras={"num_inference_steps": 40, "guidance_scale": 1, "seed": 42},
),
perf_repeat_requests=2,
run_perf_check=False,
run_component_accuracy_check=False,
run_t2v_input_reference_check=False,
),
DiffusionTestCase(
"qwen_image_t2i_2_gpus_extra_high",
DiffusionServerArgs(
@@ -1160,7 +1170,6 @@ TWO_GPU_CASES = [
ring_degree=2,
),
replace(T2I_sampling_params, extras={"quality": "extra-high"}),
run_perf_check=False,
run_component_accuracy_check=False,
run_models_api_check=False,
run_t2v_input_reference_check=False,
python/sglang/multimodal_gen/test/server/perf_baselines/5090.json
+6
View File
@@ -40,6 +40,7 @@
},
"scenarios": {
"flux_2_klein_base_image_t2i": {
"expected_load_ms": 30018.87,
"stages_ms": {
"DecodingStage": 18.19,
"DenoisingStage": 17612.9,
@@ -67,6 +68,7 @@
"estimated_full_test_time_s": 94.0
},
"wan2_1_t2v_1.3b": {
"expected_load_ms": 46026.44,
"stages_ms": {
"DecodingStage": 495.01,
"DenoisingStage": 21452.59,
@@ -93,6 +95,7 @@
"estimated_full_test_time_s": 160.9
},
"turbo_wan2_1_t2v_1.3b": {
"expected_load_ms": 43084.93,
"stages_ms": {
"InputValidationStage": 0.11,
"TextEncodingStage": 634.42,
@@ -120,6 +123,7 @@
"estimated_full_test_time_s": 200.3
},
"zimage_image_t2i": {
"expected_load_ms": 43930.51,
"stages_ms": {
"DecodingStage": 7.11,
"DenoisingStage": 2178.8,
@@ -146,6 +150,7 @@
"estimated_full_test_time_s": 329.8
},
"flux_image_t2i_layerwise_cpu_offload_5090": {
"expected_load_ms": 47725.91,
"stages_ms": {
"InputValidationStage": 0.11,
"TextEncodingStage": 204.63,
@@ -194,6 +199,7 @@
"estimated_full_test_time_s": 90.0
},
"minimax_h3_t2va_consumer_budget_1gpu_5090": {
"expected_load_ms": 47708.93,
"stages_ms": {
"InputValidationStage": 0.06,
"MiniMaxH3PartitionAdmissionStage": 0.03,
python/sglang/multimodal_gen/test/server/perf_baselines/b200.json
+39
View File
@@ -5,6 +5,40 @@
"description": "Reference estimates for B200-only diffusion cases, split out from the shared diffusion baseline file.",
"last_updated": "2026-07-01"
},
"runner_overrides": {
"b200-cirrascale1-0123": {
"flux1_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 1574.32
},
"qwen_image_2512_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 17742.04
}
},
"b200-cirrascale3-0123": {
"flux1_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 1470.24
},
"qwen_image_2512_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 17894.49
}
},
"b200-cirrascale3-4567": {
"flux1_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 1471.19
},
"qwen_image_2512_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 17346.65
}
},
"b200-di": {
"flux1_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 1334.16
},
"qwen_image_2512_modelopt_nvfp4_t2i": {
"expected_e2e_ms": 16126.87
}
}
},
"tolerances": {
"long_term": {
"e2e": 0.15,
@@ -40,6 +74,7 @@
},
"scenarios": {
"flux1_modelopt_nvfp4_t2i": {
"expected_load_ms": 38841.5,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 16.55,
@@ -68,6 +103,7 @@
"estimated_full_test_time_s": 71.2
},
"flux2_modelopt_nvfp4_t2i": {
"expected_load_ms": 71333.99,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 330.43,
@@ -97,6 +133,7 @@
"estimated_full_test_time_s": 592.3
},
"qwen_image_2512_modelopt_nvfp4_t2i": {
"expected_load_ms": 46619.98,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 221.4,
@@ -163,6 +200,7 @@
"estimated_full_test_time_s": 120.0
},
"wan22_modelopt_nvfp4_t2v": {
"expected_load_ms": 79653.94,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 316.23,
@@ -191,6 +229,7 @@
"estimated_full_test_time_s": 181.8
},
"ideogram4_nvfp4_t2i": {
"expected_load_ms": 33439.28,
"stages_ms": {
"InputValidationStage": 0.04,
"Ideogram4TextEncodingStage": 129.35,
python/sglang/multimodal_gen/test/server/perf_baselines/h100.json
+208 -23
View File
@@ -39,7 +39,32 @@
]
},
"scenarios": {
"sana_wm_ti2v": {
"expected_load_ms": 32625.62,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 7239.81,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0
},
"sana_video_2b_t2v": {
"expected_load_ms": 25382.51,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 2472.55,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0
},
"pi05_action_http": {
"expected_load_ms": 25954.41,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 21.37,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0
},
"qwen_image_t2i": {
"expected_load_ms": 40799.72,
"stages_ms": {
"TextEncodingStage": 232.03,
"DenoisingStage": 12402.12,
@@ -110,6 +135,7 @@
"estimated_full_test_time_s": 133.1
},
"qwen_image_t2i_2_gpus": {
"expected_load_ms": 41606.84,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 457.66,
@@ -180,41 +206,111 @@
"estimated_full_test_time_s": 65.6
},
"qwen_image21_t2i_tp2": {
"expected_load_ms": 45596.77,
"notes": "H100 CI run 35549625911 attempt 3, job 106188864502, PR head 044e71f5ecb507ffdf66729f934b51c3dc68b59f; request 1 after server warmup, complete cached public model snapshot b3179ad355be050328e483a9dfdd9e60cd62adfa. Existing tolerances are unchanged.",
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 2914.82,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 300.0
},
"qwen_image_t2i_2_gpus_extra_high": {
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"expected_load_ms": 46120.38,
"notes": "H100 CI run 34695613376, job 103558583416, head 5ae9847fc7092a7c75a275556d672f7507a1f933; request 1 after server warmup. E2E also agrees with runs 34693958419 and 34687251777. Existing tolerances are unchanged.",
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 455.77,
"LatentPreparationStage": 0.16,
"TimestepPreparationStage": 18.03,
"DenoisingStage": 9335.7,
"DecodingStage": 51.15
},
"denoise_step_ms": {
"0": 203.41,
"1": 191.47,
"2": 188.81,
"3": 186.65,
"4": 184.78,
"5": 185.17,
"6": 190.42,
"7": 191.36,
"8": 190.9,
"9": 183.73,
"10": 187.01,
"11": 184.09,
"12": 183.56,
"13": 182.9,
"14": 182.7,
"15": 197.53,
"16": 191.23,
"17": 188.21,
"18": 186.65,
"19": 184.42,
"20": 186.86,
"21": 185.31,
"22": 185.24,
"23": 184.98,
"24": 184.09,
"25": 183.76,
"26": 183.12,
"27": 182.28,
"28": 183.34,
"29": 183.79,
"30": 182.83,
"31": 182.17,
"32": 181.26,
"33": 182.59,
"34": 182.78,
"35": 182.32,
"36": 182.42,
"37": 182.26,
"38": 182.04,
"39": 183.17,
"40": 182.72,
"41": 181.68,
"42": 187.57,
"43": 190.77,
"44": 185.43,
"45": 185.76,
"46": 190.02,
"47": 182.43,
"48": 181.89,
"49": 183.58
},
"expected_e2e_ms": 9878.54,
"expected_avg_denoise_ms": 185.67,
"expected_median_denoise_ms": 184.09,
"load_peak_vram_mb": 42026.0,
"runtime_peak_vram_mb": 46696.0,
"warmup_peak_vram_mb": 43424.0,
"load_peak_allocated_mb": 41798.14,
"runtime_peak_allocated_mb": 44854.98,
"estimated_full_test_time_s": 54.4
},
"flux2_modelopt_fp8_tp2_t2i": {
"expected_load_ms": 50693.24,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 2089.62,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 75.4
},
"joy_echo_t2v_2gpu": {
"expected_load_ms": 46333.83,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 2469.07,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 62.8
},
"ideogram4_fp8_tp2_t2i": {
"expected_load_ms": 33515.44,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 8975.57,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 60.8
@@ -283,6 +379,7 @@
"estimated_full_test_time_s": 120.0
},
"flux_image_t2i": {
"expected_load_ms": 29629.92,
"stages_ms": {
"TimestepPreparationStage": 32.58,
"DenoisingStage": 6447.42,
@@ -353,6 +450,7 @@
"estimated_full_test_time_s": 127.4
},
"flux_2_image_t2i": {
"expected_load_ms": 55239.13,
"stages_ms": {
"TimestepPreparationStage": 15.77,
"DenoisingStage": 22276.29,
@@ -424,6 +522,7 @@
"estimated_full_test_time_s": 145.2
},
"flux_2_klein_image_t2i": {
"expected_load_ms": 28510.14,
"stages_ms": {
"DecodingStage": 6.41,
"InputValidationStage": 0.03,
@@ -449,6 +548,7 @@
"estimated_full_test_time_s": 120.5
},
"flux_2_klein_base_image_t2i": {
"expected_load_ms": 26405.08,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 53.19,
@@ -520,6 +620,7 @@
"estimated_full_test_time_s": 124.4
},
"flux_2_ti2i": {
"expected_load_ms": 54947.9,
"stages_ms": {
"TextEncodingStage": 364.98,
"DenoisingStage": 44465.25,
@@ -591,6 +692,7 @@
"estimated_full_test_time_s": 168.9
},
"flux_2_ti2i_multi_image_cache_dit": {
"expected_load_ms": 54569.41,
"stages_ms": {
"ImageVAEEncodingStage": 156.01,
"DenoisingStage": 23101.61,
@@ -662,6 +764,7 @@
"estimated_full_test_time_s": 148.6
},
"flux_image_t2i_2_gpus": {
"expected_load_ms": 31766.33,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 28.24,
@@ -732,6 +835,7 @@
"estimated_full_test_time_s": 44.9
},
"zimage_image_t2i": {
"expected_load_ms": 28221.57,
"stages_ms": {
"DecodingStage": 6.23,
"InputValidationStage": 0.03,
@@ -761,6 +865,7 @@
"estimated_full_test_time_s": 116.3
},
"zimage_image_t2i_fp8": {
"expected_load_ms": 27623.5,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 131.15,
@@ -790,6 +895,7 @@
"estimated_full_test_time_s": 123.7
},
"zimage_image_t2i_multi_lora": {
"expected_load_ms": 36426.15,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 129.63,
@@ -819,6 +925,7 @@
"estimated_full_test_time_s": 162.1
},
"zimage_image_t2i_2_gpus": {
"expected_load_ms": 31619.48,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 255.84,
@@ -848,6 +955,7 @@
"estimated_full_test_time_s": 39.3
},
"qwen_image_edit_ti2i": {
"expected_load_ms": 36793.38,
"stages_ms": {
"InputValidationStage": 24.78,
"ImageEncodingStage": 714.55,
@@ -919,6 +1027,7 @@
"estimated_full_test_time_s": 153.6
},
"joyai_image_edit_ti2i": {
"expected_load_ms": 34589.89,
"stages_ms": {
"InputValidationStage": 16.43,
"ImageEncodingStage": 602.8,
@@ -980,6 +1089,7 @@
"estimated_full_test_time_s": 117.6
},
"qwen_image_t2i_cache_dit_enabled": {
"expected_load_ms": 34703.95,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 231.07,
@@ -1050,6 +1160,7 @@
"estimated_full_test_time_s": 124.9
},
"wan2_1_t2v_1.3b_teacache_enabled": {
"expected_load_ms": 32869.61,
"stages_ms": {
"TextEncodingStage": 534.94,
"DenoisingStage": 3853.51,
@@ -1120,6 +1231,7 @@
"estimated_full_test_time_s": 126.0
},
"wan2_1_t2v_1.3b": {
"expected_load_ms": 30234.05,
"stages_ms": {
"DecodingStage": 361.6,
"InputValidationStage": 0.03,
@@ -1190,6 +1302,7 @@
"estimated_full_test_time_s": 129.3
},
"wan2_1_t2v_1.3b_cfg_parallel": {
"expected_load_ms": 30166.76,
"stages_ms": {
"LatentPreparationStage": 0.1,
"InputValidationStage": 0.05,
@@ -1260,6 +1373,7 @@
"estimated_full_test_time_s": 45.5
},
"turbo_wan2_1_t2v_1.3b": {
"expected_load_ms": 27199.67,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 364.61,
@@ -1285,6 +1399,7 @@
"estimated_full_test_time_s": 124.7
},
"ltx_2_two_stage_t2v": {
"expected_load_ms": 47470.86,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 403.5,
@@ -1356,6 +1471,7 @@
"estimated_full_test_time_s": 100.4
},
"wan2_2_ti2v_5b": {
"expected_load_ms": 37605.69,
"stages_ms": {
"InputValidationStage": 706.07,
"TextEncodingStage": 327.65,
@@ -1426,6 +1542,7 @@
"estimated_full_test_time_s": 141.7
},
"qwen_image_edit_2509_ti2i": {
"expected_load_ms": 37508.7,
"stages_ms": {
"ImageEncodingStage": 587.89,
"DenoisingStage": 37539.33,
@@ -1487,6 +1604,7 @@
"estimated_full_test_time_s": 160.2
},
"qwen_image_layered_i2i": {
"expected_load_ms": 37439.25,
"stages_ms": {
"QwenImageLayeredBeforeDenoisingStage": 144.26,
"TimestepPreparationStage": 0.0,
@@ -1555,6 +1673,7 @@
"estimated_full_test_time_s": 161.5
},
"fastwan2_2_ti2v_5b": {
"expected_load_ms": 31379.65,
"stages_ms": {
"InputValidationStage": 400.0,
"TextEncodingStage": 327.82,
@@ -1578,6 +1697,7 @@
"estimated_full_test_time_s": 125.2
},
"fast_hunyuan_video": {
"expected_load_ms": 32322.52,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 252.98,
@@ -1606,6 +1726,7 @@
"estimated_full_test_time_s": 77.0
},
"wan2_2_i2v_a14b_2gpu": {
"expected_load_ms": 78667.8,
"stages_ms": {
"InputValidationStage": 15.59,
"ImageVAEEncodingStage": 1484.2,
@@ -1667,6 +1788,7 @@
"estimated_full_test_time_s": 168.7
},
"wan2_1_i2v_14b_480P_2gpu": {
"expected_load_ms": 46402.22,
"stages_ms": {
"InputValidationStage": 7.88,
"LatentPreparationStage": 0.1,
@@ -1740,6 +1862,7 @@
"estimated_full_test_time_s": 128.8
},
"wan2_1_i2v_14b_720P_2gpu": {
"expected_load_ms": 37549.49,
"stages_ms": {
"InputValidationStage": 11.58,
"TextEncodingStage": 327.33,
@@ -1812,6 +1935,7 @@
"estimated_full_test_time_s": 182.2
},
"wan2_2_t2v_a14b_2gpu": {
"expected_load_ms": 61832.44,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 314.11,
@@ -1873,6 +1997,7 @@
"estimated_full_test_time_s": 188.0
},
"wan2_1_t2v_14b_2gpu": {
"expected_load_ms": 45201.26,
"stages_ms": {
"TextEncodingStage": 325.79,
"DecodingStage": 637.43,
@@ -1943,6 +2068,7 @@
"estimated_full_test_time_s": 96.3
},
"wan2_2_t2v_a14b_lora_2gpu": {
"expected_load_ms": 61440.19,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 325.03,
@@ -2003,6 +2129,7 @@
"estimated_full_test_time_s": 576.3
},
"wan2_1_t2v_1_3b_lora_1gpu": {
"expected_load_ms": 30139.26,
"stages_ms": {
"InputValidationStage": 0.05,
"TextEncodingStage": 329.6,
@@ -2073,6 +2200,7 @@
"estimated_full_test_time_s": 129.6
},
"wan2_1_i2v_14b_lora_2gpu": {
"expected_load_ms": 66554.82,
"stages_ms": {
"InputValidationStage": 12.06,
"TextEncodingStage": 325.79,
@@ -2145,6 +2273,7 @@
"estimated_full_test_time_s": 509.3
},
"flux_2_image_t2i_2_gpus": {
"expected_load_ms": 53086.36,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 372.55,
@@ -2216,6 +2345,7 @@
"estimated_full_test_time_s": 75.0
},
"qwen_image_edit_2511_ti2i": {
"expected_load_ms": 38727.15,
"stages_ms": {
"DecodingStage": 18.69,
"InputValidationStage": 48.98,
@@ -2277,6 +2407,7 @@
"estimated_full_test_time_s": 143.7
},
"fsdp-inference": {
"expected_load_ms": 31380.54,
"stages_ms": {
"InputValidationStage": 0.03,
"LatentPreparationStage": 0.11,
@@ -2306,6 +2437,7 @@
"estimated_full_test_time_s": 61.3
},
"hunyuan3d_shape_gen": {
"expected_load_ms": 31925.92,
"stages_ms": {
"Hunyuan3DShapeBeforeDenoisingStage": 54.28,
"Hunyuan3DShapeDenoisingStage": 1698.03,
@@ -2377,6 +2509,7 @@
"estimated_full_test_time_s": 420.1
},
"wan2_1_t2v_1.3b_frame_interp_2x": {
"expected_load_ms": 30176.3,
"stages_ms": {
"InputValidationStage": 0.05,
"TextEncodingStage": 534.56,
@@ -2447,6 +2580,7 @@
"estimated_full_test_time_s": 129.3
},
"flux_2_image_t2i_upscaling_4x": {
"expected_load_ms": 52489.37,
"stages_ms": {
"TextEncodingStage": 373.42,
"DenoisingStage": 21888.49,
@@ -2518,6 +2652,7 @@
"estimated_full_test_time_s": 145.1
},
"wan2_1_t2v_1.3b_upscaling_4x": {
"expected_load_ms": 31719.84,
"stages_ms": {
"DecodingStage": 362.62,
"InputValidationStage": 0.04,
@@ -2588,6 +2723,7 @@
"estimated_full_test_time_s": 129.3
},
"wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": {
"expected_load_ms": 30083.06,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 532.52,
@@ -2658,6 +2794,7 @@
"estimated_full_test_time_s": 129.4
},
"ltx_2.3_one_stage_ti2v": {
"expected_load_ms": 39264.11,
"stages_ms": {
"InputValidationStage": 3.49,
"TextEncodingStage": 397.91,
@@ -2711,6 +2848,7 @@
"estimated_full_test_time_s": 167.9
},
"ltx_2.3_two_stage_t2v_2gpus": {
"expected_load_ms": 45206.77,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 399.08,
@@ -2772,14 +2910,41 @@
"estimated_full_test_time_s": 186.0
},
"ltx_2_5_diffusion_decoder_2gpus": {
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"expected_load_ms": 28133.24,
"stages_ms": {
"InputValidationStage": 0.05,
"TextEncodingStage": 634.13,
"LTX2TextConnectorStage": 317.43,
"LTX2DurationStage": 0.01,
"LTX2SigmaPreparationStage": 0.03,
"TimestepPreparationStage": 246.11,
"LTX2AVLatentPreparationStage": 0.31,
"LTX2ImageEncodingStage": 0.02,
"LTX2AVDenoisingStage": 3750.7,
"LTX2AVDecodingStage": 3382.33
},
"denoise_step_ms": {
"0": 292.65,
"1": 259.94,
"2": 256.47,
"3": 256.26,
"4": 259.62,
"5": 255.44,
"6": 258.02,
"7": 255.91
},
"expected_e2e_ms": 10141.6,
"expected_avg_denoise_ms": 261.79,
"expected_median_denoise_ms": 257.25,
"load_peak_vram_mb": 1906,
"runtime_peak_vram_mb": 69546,
"warmup_peak_vram_mb": 28396,
"load_peak_allocated_mb": 1876.14,
"runtime_peak_allocated_mb": 48399,
"estimated_full_test_time_s": 545.1
},
"ltx_2_3_two_stage_ti2v_2gpus": {
"expected_load_ms": 50003.37,
"stages_ms": {
"InputValidationStage": 3.21,
"TextEncodingStage": 411.15,
@@ -2841,6 +3006,7 @@
"estimated_full_test_time_s": 101.4
},
"longlive2_t2v": {
"expected_load_ms": 29328.4,
"stages_ms": {
"InputValidationStage": 0.04,
"LongLive2TextEncodingStage": 328.05,
@@ -2861,6 +3027,7 @@
"estimated_full_test_time_s": 153.1
},
"longlive2_i2v": {
"expected_load_ms": 27107.25,
"stages_ms": {
"InputValidationStage": 23.02,
"LongLive2TextEncodingStage": 327.98,
@@ -2881,6 +3048,7 @@
"estimated_full_test_time_s": 149.4
},
"lingbot_video_moe_t2v": {
"expected_load_ms": 22654.46,
"stages_ms": {
"InputValidationStage": 0.06,
"LingBotVideoTextEncodingStage": 89.03,
@@ -2913,6 +3081,7 @@
"estimated_full_test_time_s": 600.0
},
"lingbot_world_realtime_plastic_beach": {
"expected_load_ms": 31600.49,
"stages_ms": {
"RealtimeInputValidationStage": 0.09,
"RealtimeTextEncodingStage": 0.04,
@@ -2924,7 +3093,7 @@
"CausalVaeDecodingStage": 217.75
},
"denoise_step_ms": {},
"expected_e2e_ms": 1912.29,
"expected_e2e_ms": 18687.61,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"load_peak_vram_mb": 42814.0,
@@ -2934,6 +3103,7 @@
"estimated_full_test_time_s": 126.0
},
"ltx_2_3_hq_pipeline": {
"expected_load_ms": 116530.53,
"stages_ms": {
"InputValidationStage": 4.63,
"TextEncodingStage": 401.75,
@@ -2982,6 +3152,7 @@
"estimated_full_test_time_s": 363.2
},
"qwen_image_t2i_cache_dit_scm_config_diffusers_1gpu": {
"expected_load_ms": 35585.13,
"stages_ms": {
"DiffusersExecutionStage": 1075.16
},
@@ -2996,6 +3167,7 @@
"estimated_full_test_time_s": 98.3
},
"cosmos3_nano_t2i": {
"expected_load_ms": 28691.77,
"stages_ms": {
"Cosmos3ImagePreprocessStage": 0.01,
"Cosmos3TokenizationStage": 155.89,
@@ -3047,6 +3219,7 @@
"estimated_full_test_time_s": 65.0
},
"cosmos3_nano_t2v": {
"expected_load_ms": 30688.3,
"stages_ms": {
"Cosmos3ImagePreprocessStage": 0.01,
"Cosmos3TokenizationStage": 143.82,
@@ -3068,30 +3241,34 @@
"estimated_full_test_time_s": 65.0
},
"flux_2_t2i_customized_vae_path": {
"expected_load_ms": 52103.72,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 23073.55,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 574.4
},
"wan2_2_t2v_a14b_teacache_2gpu": {
"expected_load_ms": 63938.03,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 82478.23,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 186.4
},
"wan2_1_t2v_1_3b_cache_dit_sp_only_2gpu": {
"expected_load_ms": 34366.19,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 1077.99,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 52.2
},
"minimax_h3_ref2va_video_audio_2gpu_h100": {
"expected_load_ms": 68100.42,
"stages_ms": {
"InputValidationStage": 0.08,
"MiniMaxH3PartitionAdmissionStage": 0.04,
@@ -3120,6 +3297,7 @@
"estimated_full_test_time_s": 340.0
},
"minimax_h3_t2va_2gpu_h100": {
"expected_load_ms": 79373.57,
"stages_ms": {
"InputValidationStage": 0.05,
"MiniMaxH3PartitionAdmissionStage": 0.03,
@@ -3151,14 +3329,16 @@
"estimated_full_test_time_s": 235.0
},
"mova_360p_tp2": {
"expected_load_ms": 60480.85,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 58242.09,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 141.9
},
"zimage_image_t2i_2_gpus_non_square": {
"expected_load_ms": 30675.16,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 259.45,
@@ -3188,9 +3368,10 @@
"estimated_full_test_time_s": 40.5
},
"flux1_modelopt_fp8_t2i": {
"expected_load_ms": 33073.38,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 1138.33,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 50.3
@@ -3204,14 +3385,16 @@
"estimated_full_test_time_s": 498.1
},
"wan22_modelopt_fp8_t2v": {
"expected_load_ms": 63380.1,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 7908.23,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 89.5
},
"hunyuanvideo_modelopt_fp8_t2v": {
"expected_load_ms": 25315.9,
"stages_ms": {
"InputValidationStage": 0.05,
"TextEncodingStage": 33.67,
@@ -3244,17 +3427,19 @@
"estimated_full_test_time_s": 64.3
},
"qwen_image_modelopt_fp8_t2i": {
"expected_load_ms": 47967.61,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 3038.04,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 65.5
},
"qwen_image_edit_modelopt_fp8_ti2i": {
"expected_load_ms": 49282.42,
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_e2e_ms": 2964.8,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 73.7
python/sglang/multimodal_gen/test/server/perf_baselines/xpu_b60.json
+158 -153
View File
@@ -2,8 +2,8 @@
"metadata": {
"model": "Diffusion Server",
"hardware": "CI Intel Arc Pro B60 (24 GiB) pool: bmg-multigen-models",
"description": "Reference numbers seeded from XPU multimodal_gen CI run 32736878259 job 97461821112 (PR #36100).",
"last_updated": "2026-08-24"
"description": "All scenarios re-seeded from XPU multimodal_gen CI run 35563559050 job 106220948594 (PR #39956), with SGLANG_DIFFUSION_SYNC_STAGE_PROFILING=1 so step timings record device time rather than host enqueue. Each value is the max over the 7 attempts of that job. Denoise numbers are stable across attempts (<=2% spread); expected_load_ms is host I/O bound and ranged 21.4-48.0s across attempts, so it is guarded only by the wide e2e tolerance.",
"last_updated": "2026-09-21"
},
"tolerances": {
"long_term": {
@@ -42,181 +42,186 @@
"zimage_image_t2i": {
"estimated_full_test_time_s": 65.1,
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 87.64,
"LatentPreparationStage": 0.21,
"TimestepPreparationStage": 136.9,
"DenoisingStage": 8361.1,
"DecodingStage": 8.39
"InputValidationStage": 0.1,
"TextEncodingStage": 89.87,
"LatentPreparationStage": 0.28,
"TimestepPreparationStage": 2.44,
"DenoisingStage": 8614.3,
"DecodingStage": 293.13
},
"denoise_step_ms": {
"0": 99.57,
"1": 927.61,
"2": 926.72,
"3": 928.32,
"4": 929.29,
"5": 928.8,
"6": 925.2,
"7": 927.22,
"8": 931.28
"0": 862.67,
"1": 859.76,
"2": 861.38,
"3": 860.12,
"4": 859.95,
"5": 859.53,
"6": 859.63,
"7": 860.1,
"8": 859.77
},
"expected_e2e_ms": 9312.75,
"expected_avg_denoise_ms": 836.0,
"expected_median_denoise_ms": 927.61
"expected_e2e_ms": 9799.49,
"expected_load_ms": 33283.41,
"expected_avg_denoise_ms": 860.1,
"expected_median_denoise_ms": 859.76
},
"flux_2_klein_image_t2i": {
"estimated_full_test_time_s": 106.8,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 37.83,
"ImageVAEEncodingStage": 0.0,
"LatentPreparationStage": 0.33,
"TimestepPreparationStage": 351.26,
"DenoisingStage": 2518.52,
"DecodingStage": 9.13
"InputValidationStage": 0.11,
"TextEncodingStage": 88.26,
"ImageVAEEncodingStage": 0.06,
"LatentPreparationStage": 0.47,
"TimestepPreparationStage": 2.37,
"DenoisingStage": 2803.24,
"DecodingStage": 292.44
},
"denoise_step_ms": {
"0": 83.44,
"1": 16.73,
"2": 18.04,
"3": 17.11
"0": 591.79,
"1": 536.64,
"2": 537.71,
"3": 537.7
},
"expected_e2e_ms": 3509.48,
"expected_avg_denoise_ms": 33.83,
"expected_median_denoise_ms": 17.58
"expected_e2e_ms": 3754.98,
"expected_load_ms": 26566.83,
"expected_avg_denoise_ms": 550.96,
"expected_median_denoise_ms": 537.7
},
"flux_2_klein_base_image_t2i": {
"stages_ms": {
"InputValidationStage": 0.03,
"TextEncodingStage": 36.3,
"ImageVAEEncodingStage": 0.0,
"LatentPreparationStage": 0.28,
"TimestepPreparationStage": 355.5,
"DenoisingStage": 55237.43,
"DecodingStage": 11.09
"InputValidationStage": 0.11,
"TextEncodingStage": 87.34,
"ImageVAEEncodingStage": 0.06,
"LatentPreparationStage": 0.44,
"TimestepPreparationStage": 2.49,
"DenoisingStage": 54080.81,
"DecodingStage": 287.66
},
"denoise_step_ms": {
"0": 106.73,
"1": 34.25,
"2": 34.59,
"3": 34.59,
"4": 35.08,
"5": 34.24,
"6": 34.03,
"7": 34.47,
"8": 34.17,
"9": 44.04,
"10": 34.83,
"11": 34.64,
"12": 35.38,
"13": 34.32,
"14": 35.51,
"15": 36.16,
"16": 35.32,
"17": 34.75,
"18": 35.28,
"19": 50.88,
"20": 35.41,
"21": 35.36,
"22": 34.53,
"23": 35.42,
"24": 34.75,
"25": 34.78,
"26": 34.95,
"27": 34.84,
"28": 35.38,
"29": 34.7,
"30": 34.42,
"31": 35.01,
"32": 34.98,
"33": 34.61,
"34": 34.78,
"35": 34.76,
"36": 34.61,
"37": 37.59,
"38": 35.0,
"39": 34.6,
"40": 35.52,
"41": 35.14,
"42": 34.84,
"43": 35.16,
"44": 33.86,
"45": 34.52,
"46": 34.94,
"47": 33.84,
"48": 35.14,
"49": 34.75
"0": 1123.84,
"1": 1068.88,
"2": 1068.05,
"3": 1067.09,
"4": 1068.74,
"5": 1068.64,
"6": 1067.83,
"7": 1068.84,
"8": 1068.23,
"9": 1068.88,
"10": 1067.99,
"11": 1068.93,
"12": 1069.2,
"13": 1067.82,
"14": 1068.2,
"15": 1068.84,
"16": 1068.15,
"17": 1067.99,
"18": 1068.51,
"19": 1068.77,
"20": 1068.34,
"21": 1068.12,
"22": 1069.19,
"23": 1068.35,
"24": 1069.71,
"25": 1068.26,
"26": 1068.36,
"27": 1069.13,
"28": 1068.15,
"29": 1068.11,
"30": 1069.23,
"31": 1069.6,
"32": 1069.44,
"33": 1068.34,
"34": 1068.38,
"35": 1069.22,
"36": 1068.4,
"37": 1067.92,
"38": 1068.22,
"39": 1068.85,
"40": 1068.95,
"41": 1069.67,
"42": 1068.07,
"43": 1068.93,
"44": 1069.75,
"45": 1067.77,
"46": 1068.09,
"47": 1068.11,
"48": 1068.67,
"49": 1068.8
},
"expected_e2e_ms": 56248.35,
"expected_avg_denoise_ms": 36.83,
"expected_median_denoise_ms": 34.84
"expected_e2e_ms": 55028.04,
"expected_load_ms": 31539.71,
"expected_avg_denoise_ms": 1069.67,
"expected_median_denoise_ms": 1068.45,
"estimated_full_test_time_s": 103.3
},
"wan2_1_t2v_1.3b": {
"estimated_full_test_time_s": 331.5,
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 704.97,
"LatentPreparationStage": 0.89,
"TimestepPreparationStage": 2.71,
"DenoisingStage": 114601.91,
"DecodingStage": 73.74,
"InputValidationStage": 0.11,
"TextEncodingStage": 764.42,
"LatentPreparationStage": 1.26,
"TimestepPreparationStage": 0.97,
"DenoisingStage": 70658.77,
"DecodingStage": 3102.93,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 208.34,
"1": 80.0,
"2": 84.33,
"3": 85.37,
"4": 92.98,
"5": 85.51,
"6": 84.71,
"7": 85.0,
"8": 84.82,
"9": 85.45,
"10": 85.89,
"11": 86.33,
"12": 85.68,
"13": 86.51,
"14": 94.06,
"15": 86.65,
"16": 86.39,
"17": 86.15,
"18": 86.54,
"19": 86.88,
"20": 87.36,
"21": 85.66,
"22": 86.22,
"23": 85.45,
"24": 94.66,
"25": 86.44,
"26": 86.3,
"27": 86.35,
"28": 85.29,
"29": 86.31,
"30": 85.8,
"31": 86.12,
"32": 86.03,
"33": 86.42,
"34": 91.27,
"35": 85.84,
"36": 86.3,
"37": 85.04,
"38": 84.99,
"39": 85.46,
"40": 85.72,
"41": 84.76,
"42": 88.92,
"43": 85.58,
"44": 86.57,
"45": 85.42,
"46": 87.03,
"47": 85.65,
"48": 86.77,
"49": 84.42
"0": 1677.96,
"1": 1401.03,
"2": 1400.7,
"3": 1401.14,
"4": 1400.29,
"5": 1401.15,
"6": 1400.76,
"7": 1400.83,
"8": 1399.92,
"9": 1401.37,
"10": 1400.79,
"11": 1401.01,
"12": 1400.59,
"13": 1400.05,
"14": 1400.17,
"15": 1400.44,
"16": 1400.18,
"17": 1400.41,
"18": 1399.86,
"19": 1400.23,
"20": 1400.58,
"21": 1400.54,
"22": 1400.3,
"23": 1400.26,
"24": 1400.32,
"25": 1400.66,
"26": 1400.31,
"27": 1400.01,
"28": 1400.15,
"29": 1400.61,
"30": 1400.58,
"31": 1400.79,
"32": 1400.66,
"33": 1400.92,
"34": 1401.43,
"35": 1400.94,
"36": 1400.91,
"37": 1400.93,
"38": 1401.66,
"39": 1400.7,
"40": 1400.75,
"41": 1401.24,
"42": 1400.85,
"43": 1400.86,
"44": 1400.98,
"45": 1400.81,
"46": 1400.27,
"47": 1400.9,
"48": 1400.3,
"49": 1400.39
},
"expected_e2e_ms": 116299.31,
"expected_avg_denoise_ms": 88.83,
"expected_median_denoise_ms": 86.07
"expected_e2e_ms": 75427.76,
"expected_load_ms": 48009.32,
"expected_avg_denoise_ms": 1406.13,
"expected_median_denoise_ms": 1400.58
}
}
}
python/sglang/multimodal_gen/test/server/realtime_consistency.py
+15 -2
View File
@@ -6,6 +6,7 @@ import asyncio
import os
import statistics
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any
@@ -50,9 +51,11 @@ class RealtimeChunkStats:
class RealtimeCollectionResult:
frames: list[np.ndarray]
chunk_stats: list[RealtimeChunkStats]
e2e_ms: float
_REALTIME_CHUNK_STATS_BY_CASE: dict[str, list[RealtimeChunkStats]] = {}
_REALTIME_E2E_MS_BY_CASE: dict[str, float] = {}
_REALTIME_KEY_FRAMES_BY_CASE: dict[str, list[np.ndarray]] = {}
@@ -217,15 +220,20 @@ def validate_realtime_perf_stats(
def record_realtime_perf_stats(
case_id: str, chunk_stats: list[RealtimeChunkStats]
case_id: str, chunk_stats: list[RealtimeChunkStats], e2e_ms: float
) -> None:
_REALTIME_CHUNK_STATS_BY_CASE[case_id] = list(chunk_stats)
_REALTIME_E2E_MS_BY_CASE[case_id] = e2e_ms
def pop_realtime_perf_stats(case_id: str) -> list[RealtimeChunkStats]:
return _REALTIME_CHUNK_STATS_BY_CASE.pop(case_id, [])
def pop_realtime_e2e_ms(case_id: str) -> float | None:
return _REALTIME_E2E_MS_BY_CASE.pop(case_id, None)
def select_realtime_key_frames(frames: list[np.ndarray]) -> list[np.ndarray]:
if not frames:
return []
@@ -357,6 +365,8 @@ async def collect_realtime_output(
sent_event_indices.add(event_idx)
async with websockets.connect(ws_url, max_size=None, ping_interval=None) as ws:
# exclude server startup, warmup and the later mp4 consistency encoding
request_start = time.perf_counter()
await ws.send(msgspec.msgpack.encode(init_payload))
await send_events_for_boundary(ws, -1)
@@ -399,5 +409,8 @@ async def collect_realtime_output(
if header.get("is_final_frame_batch", True):
received_chunks.add(chunk_index)
await send_events_for_boundary(ws, chunk_index)
e2e_ms = (time.perf_counter() - request_start) * 1000
return RealtimeCollectionResult(frames=frames, chunk_stats=chunk_stats)
return RealtimeCollectionResult(
frames=frames, chunk_stats=chunk_stats, e2e_ms=e2e_ms
)
python/sglang/multimodal_gen/test/server/test_server_common.py
+138 -61
View File
@@ -8,6 +8,7 @@ Each collected request prints a performance log before validation.
from __future__ import annotations
import json
import math
import os
import queue
import threading
@@ -26,6 +27,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.server.realtime_consistency import (
RealtimeChunkStats,
pop_realtime_e2e_ms,
pop_realtime_key_frames,
pop_realtime_perf_stats,
validate_realtime_perf_stats,
@@ -101,6 +103,10 @@ _SERVER_FATAL_LOG_PATTERNS = (
_CASE_LOG_SEPARATOR = "=" * 88
class PerformanceValidationError(AssertionError):
"""A terminal performance failure, including across repeated requests."""
def _print_case_log_separator(case_id: str, state: str) -> None:
print(
f"\n{_CASE_LOG_SEPARATOR}\n"
@@ -371,11 +377,14 @@ class DiffusionServerBase:
log_path = ctx.perf_log_path
log_wait_timeout = 30
req_perf_record = wait_for_req_perf_record(
rid,
log_path,
timeout=log_wait_timeout,
)
try:
req_perf_record = wait_for_req_perf_record(
rid,
log_path,
timeout=log_wait_timeout,
)
except AssertionError as exc:
raise PerformanceValidationError(f"[performance] {case_id}: {exc}") from exc
return (req_perf_record, content)
@@ -384,8 +393,13 @@ class DiffusionServerBase:
case: DiffusionTestCase,
perf_record: RequestPerfRecord,
request_index: int = 1,
load_time_ms: float | None = None,
) -> None:
"""Validate metrics and record results."""
if perf_record is None:
raise PerformanceValidationError(
f"[performance] {case.id}: request performance record is missing"
)
is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
scenario = BASELINE_CONFIG.scenarios.get(case.id)
@@ -412,22 +426,25 @@ class DiffusionServerBase:
)
summary = validator.collect_metrics(perf_record)
self._print_performance_log(case, summary, scenario)
summary.load_time_ms = load_time_ms
self._record_performance_result(case, summary, request_index)
self._print_performance_log(case, summary, scenario)
if is_baseline_generation_mode:
_PENDING_BASELINE_DUMPS.setdefault(case.id, []).append(summary)
return
if missing_scenario:
self._dump_baseline_for_testcase(case, summary, missing_scenario)
pytest.fail(
f"Testcase '{case.id}' not found in {get_perf_baseline_update_path()}"
)
# disabling stage checks must not disable the request's e2e guard
validator.validate_e2e(summary)
validator.validate_load(summary)
if case.run_perf_check:
if is_baseline_generation_mode:
_PENDING_BASELINE_DUMPS.setdefault(case.id, []).append(summary)
return
if missing_scenario:
self._dump_baseline_for_testcase(case, summary, missing_scenario)
if missing_scenario:
pytest.fail(
f"Testcase '{case.id}' not found in {get_perf_baseline_update_path()}"
)
return
if current_platform.is_cuda():
expected_load_peak_vram_mb = scenario.load_peak_vram_mb
expected_runtime_peak_vram_mb = scenario.runtime_peak_vram_mb
@@ -476,6 +493,57 @@ class DiffusionServerBase:
chunk_stats: list[RealtimeChunkStats],
request_index: int = 1,
) -> None:
e2e_ms = pop_realtime_e2e_ms(case.id)
scenario = BASELINE_CONFIG.scenarios.get(case.id)
summary = PerformanceSummary(e2e_ms, 0, 0, {}, [], {}, {})
check_memory = case.run_perf_check and current_platform.is_cuda()
if check_memory:
request_id = next(
(stat.request_id for stat in reversed(chunk_stats) if stat.request_id),
None,
)
if request_id is None:
pytest.fail(f"{case.id}: realtime chunk stats are missing request IDs")
perf_record = wait_for_req_perf_record(
request_id, ctx.perf_log_path, timeout=30
)
if perf_record is None:
pytest.fail(
f"{case.id}: realtime request performance record is missing"
)
if scenario is None:
pytest.fail(
f"Testcase '{case.id}' not found in {get_perf_baseline_update_path()}"
)
validator = PerformanceValidator(
scenario=scenario,
tolerances=BASELINE_CONFIG.tolerances,
step_fractions=BASELINE_CONFIG.step_fractions,
)
summary = validator.collect_metrics(perf_record)
# the last chunk's record supplies memory peaks, not the session's e2e
summary.e2e_ms = e2e_ms
summary.load_time_ms = ctx.load_time_ms
self._record_performance_result(case, summary, request_index)
self._print_performance_log(case, summary, scenario)
if os.environ.get("SGLANG_GEN_BASELINE", "0") == "1":
_PENDING_BASELINE_DUMPS.setdefault(case.id, []).append(summary)
return
if scenario is None:
pytest.fail(
f"Testcase '{case.id}' not found in {get_perf_baseline_update_path()}"
)
if not check_memory:
validator = PerformanceValidator(
scenario=scenario,
tolerances=BASELINE_CONFIG.tolerances,
step_fractions=BASELINE_CONFIG.step_fractions,
)
validator.validate_e2e(summary)
validator.validate_load(summary)
validate_realtime_perf_stats(
case.id,
chunk_stats,
@@ -484,48 +552,7 @@ class DiffusionServerBase:
case.sampling_params.realtime_perf_ignore_initial_chunks
),
)
if not case.run_perf_check or not current_platform.is_cuda():
return
request_id = next(
(stat.request_id for stat in reversed(chunk_stats) if stat.request_id),
None,
)
if request_id is None:
pytest.fail(f"{case.id}: realtime chunk stats are missing request IDs")
perf_record = wait_for_req_perf_record(
request_id,
ctx.perf_log_path,
timeout=30,
)
if perf_record is None:
pytest.fail(f"{case.id}: realtime request performance record is missing")
scenario = BASELINE_CONFIG.scenarios.get(case.id)
if scenario is None:
pytest.fail(
f"Testcase '{case.id}' not found in {get_perf_baseline_update_path()}"
)
validator = PerformanceValidator(
scenario=scenario,
tolerances=BASELINE_CONFIG.tolerances,
step_fractions=BASELINE_CONFIG.step_fractions,
)
summary = validator.collect_metrics(perf_record)
self._print_performance_log(case, summary, scenario)
self._record_performance_result(case, summary, request_index)
if os.environ.get("SGLANG_GEN_BASELINE", "0") == "1":
logger.info(
"%s realtime peak VRAM baseline: load=%.0fMiB, runtime=%.0fMiB, "
"warmup=%.0fMiB",
case.id,
summary.load_peak_vram_mb,
summary.runtime_peak_vram_mb,
summary.warmup_peak_vram_mb,
)
if not check_memory:
return
if scenario.load_peak_vram_mb is None or scenario.runtime_peak_vram_mb is None:
@@ -555,12 +582,28 @@ class DiffusionServerBase:
summary: PerformanceSummary,
request_index: int = 1,
) -> None:
if not isinstance(summary.e2e_ms, (int, float)) or not (
math.isfinite(summary.e2e_ms) and summary.e2e_ms > 0
):
raise PerformanceValidationError(
f"[performance] {case.id}: E2E duration missing or invalid: "
f"{summary.e2e_ms!r}"
)
if summary.load_time_ms is None or not (
math.isfinite(summary.load_time_ms) and summary.load_time_ms > 0
):
raise PerformanceValidationError(
f"[performance] {case.id}: Load duration missing or invalid: "
f"{summary.load_time_ms!r}"
)
result = {
"class_name": type(self).__name__,
"test_name": case.id,
"request_index": request_index,
"modality": case.server_args.modality,
"e2e_ms": summary.e2e_ms,
"load_time_ms": summary.load_time_ms,
"load_inclusive_e2e_ms": summary.load_time_ms + summary.e2e_ms,
"avg_denoise_ms": summary.avg_denoise_ms,
"median_denoise_ms": summary.median_denoise_ms,
"load_peak_vram_mb": summary.load_peak_vram_mb,
@@ -595,6 +638,8 @@ class DiffusionServerBase:
f"--- Performance Log: {case.id} ---",
(
f" e2e={summary.e2e_ms:.2f}ms, "
f"load={summary.load_time_ms:.2f}ms, "
f"load_inclusive_e2e={summary.load_time_ms + summary.e2e_ms:.2f}ms, "
f"avg_denoise={summary.avg_denoise_ms:.2f}ms, "
f"median_denoise={summary.median_denoise_ms:.2f}ms, "
f"load_peak_vram={summary.load_peak_vram_mb:.0f}MiB, "
@@ -662,6 +707,7 @@ class DiffusionServerBase:
"stages_ms": stages_formatted,
"denoise_step_ms": denoise_steps_formatted,
"expected_e2e_ms": round(max(s.e2e_ms for s in summaries), 2),
"expected_load_ms": round(max(s.load_time_ms for s in summaries), 2),
"expected_avg_denoise_ms": round(
max(s.avg_denoise_ms for s in summaries), 2
),
@@ -1573,6 +1619,28 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION}
if case.run_lora_dynamic_load_check:
self._test_dynamic_lora_loading(diffusion_server, case)
for warmup_index in range(case.perf_warmup_requests):
label = f"request warmup {warmup_index + 1}/{case.perf_warmup_requests}"
_print_case_log_separator(case.id, f"BEGIN {label}")
generate_fn = get_generate_fn(
model_path=case.server_args.model_path,
modality=case.server_args.modality,
sampling_params=case.sampling_params,
)
record, _ = self.run_and_collect(
diffusion_server, case.id, generate_fn, collect_perf=True
)
if record is None or not (
math.isfinite(record.total_duration_ms) and record.total_duration_ms > 0
):
raise PerformanceValidationError(
f"[performance] {case.id}: {label} E2E duration missing or invalid"
)
print(
f"[server-test] {case.id}: {label} e2e={record.total_duration_ms:.4f}ms"
)
_print_case_log_separator(case.id, f"END {label}")
failures = []
for request_index in range(1, case.perf_repeat_requests + 1):
label = f"request {request_index}/{case.perf_repeat_requests}"
@@ -1586,6 +1654,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION}
str(Path(artifact_dir) / f"request-{request_index}"),
)
self._test_diffusion_request(case, diffusion_server, request_index)
except PerformanceValidationError as exc:
_print_case_log_separator(case.id, f"FAILED {label}")
raise PerformanceValidationError(f"[{label}] {exc}") from exc
except pytest.skip.Exception as exc:
if request_index == 1:
raise
@@ -1635,6 +1706,10 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION}
except BaseException as exc:
if isinstance(exc, (KeyboardInterrupt, SystemExit)):
raise
if name == "performance" and isinstance(
exc, (AssertionError, pytest.fail.Exception)
):
raise PerformanceValidationError(f"[performance] {exc}") from exc
failures.append((name, str(exc)))
if is_realtime_case:
@@ -1651,7 +1726,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION}
else:
run_case_check(
"performance",
lambda: self._validate_and_record(case, perf_record, request_index),
lambda: self._validate_and_record(
case, perf_record, request_index, diffusion_server.load_time_ms
),
)
if case.server_args.custom_validator == "mesh":
python/sglang/multimodal_gen/test/server/test_server_utils.py
+50 -6
View File
@@ -6,7 +6,9 @@ from __future__ import annotations
import asyncio
import base64
import math
import os
import re
import shlex
import subprocess
import sys
@@ -187,6 +189,7 @@ class ServerContext:
log_dir: Path
_stdout_fh: Any = field(repr=False)
_log_thread: threading.Thread | None = field(default=None, repr=False)
load_time_ms: float | None = None
def log_tail(self, lines: int = 200) -> str:
"""Return recent server output for failure diagnostics."""
@@ -422,6 +425,9 @@ class ServerManager:
# regardless of log-level configuration.
print(f"[server-test] Running command: {cmd_str}", flush=True)
load_started_ns = time.monotonic_ns()
load_finished_ns = None
load_ready = threading.Event()
process = subprocess.Popen(
command,
stdout=subprocess.PIPE,
@@ -437,9 +443,17 @@ class ServerManager:
def _log_pipe(pipe: Any, file: Any) -> None:
"""Read from pipe and write to file and stdout."""
nonlocal load_finished_ns
try:
with pipe:
for line in iter(pipe.readline, ""):
match = re.search(
r"\[server-load\] workers_ready_monotonic_ns=(\d+)",
line,
)
if match and load_finished_ns is None:
load_finished_ns = int(match.group(1))
load_ready.set()
sys.stdout.write(line)
sys.stdout.flush()
file.write(line)
@@ -474,6 +488,10 @@ class ServerManager:
)
try:
self._wait_for_ready(process, stdout_path)
# health includes warmup; the worker marker's clock excludes it
load_ready.wait(timeout=5)
if load_finished_ns is not None:
context.load_time_ms = (load_finished_ns - load_started_ns) / 1e6
except BaseException:
context.cleanup()
raise
@@ -713,7 +731,7 @@ class PerformanceValidator:
if self.is_baseline_generation_mode:
return summary
self._validate_e2e(summary)
self.validate_e2e(summary)
self._validate_denoise_agg(summary)
self._validate_denoise_steps(summary)
self._validate_stages(summary)
@@ -738,9 +756,15 @@ class PerformanceValidator:
return profile_tolerance
return max(profile_tolerance, override)
def _validate_e2e(self, summary: PerformanceSummary) -> None:
def validate_e2e(self, summary: PerformanceSummary) -> None:
"""Validate end-to-end performance."""
assert summary.e2e_ms > 0, "E2E duration missing"
assert math.isfinite(summary.e2e_ms) and summary.e2e_ms > 0, (
"E2E duration missing or invalid"
)
expected = self.scenario.expected_e2e_ms
assert math.isfinite(expected) and expected > 0, (
"E2E baseline missing or invalid"
)
self._assert_le(
"E2E Latency",
summary.e2e_ms,
@@ -748,6 +772,24 @@ class PerformanceValidator:
self._timing_tol(self.tolerances.e2e),
)
def validate_load(self, summary: PerformanceSummary) -> None:
load_ms = summary.load_time_ms
expected_load_ms = self.scenario.expected_load_ms
assert load_ms is not None and math.isfinite(load_ms) and load_ms > 0, (
"Load duration missing or invalid"
)
assert (
expected_load_ms is not None
and math.isfinite(expected_load_ms)
and expected_load_ms > 0
), "Load baseline missing or invalid"
self._assert_le(
"Load Latency (excluding warmup)",
load_ms,
expected_load_ms,
self._timing_tol(self.tolerances.e2e),
)
def _validate_denoise_agg(self, summary: PerformanceSummary) -> None:
"""Validate aggregate denoising metrics."""
assert summary.avg_denoise_ms > 0, "Denoising step timings missing"
@@ -801,7 +843,7 @@ class PerformanceValidator:
assert actual is not None, f"Stage {stage} timing missing"
tolerance = self._timing_tol(
self.tolerances.denoise_stage
if stage == "DenoisingStage"
if stage in summary.denoising_stages
else self.tolerances.non_denoise_stage
)
if stage.endswith("DecodingStage"):
@@ -1479,9 +1521,9 @@ def get_generate_fn(
size=sampling_params.output_size,
seconds=video_seconds,
extra_body={
"reference_url": sampling_params.image_path,
"fps": sampling_params.fps,
"num_frames": sampling_params.num_frames,
**extra_body,
},
)
@@ -1541,7 +1583,9 @@ def get_generate_fn(
require_chunk_stats=True,
)
)
record_realtime_perf_stats(case_id, realtime_output.chunk_stats)
record_realtime_perf_stats(
case_id, realtime_output.chunk_stats, realtime_output.e2e_ms
)
record_realtime_key_frames(case_id, realtime_output.frames)
fps = int(sampling_params.fps or 24)
video_bytes = encode_realtime_frames_to_mp4(realtime_output.frames, fps=fps)
python/sglang/multimodal_gen/test/server/testcase_configs.py
+25
View File
@@ -115,6 +115,7 @@ class ScenarioConfig:
expected_avg_denoise_ms: float
expected_median_denoise_ms: float
estimated_full_test_time_s: float | None = None
expected_load_ms: float | None = None
load_peak_vram_mb: float | None = None
runtime_peak_vram_mb: float | None = None
# Peak of the warmup calibration probe (the default workload's full shape
@@ -146,6 +147,7 @@ class ScenarioConfig:
expected_avg_denoise_ms=float(cfg["expected_avg_denoise_ms"]),
expected_median_denoise_ms=float(cfg["expected_median_denoise_ms"]),
estimated_full_test_time_s=optional_float("estimated_full_test_time_s"),
expected_load_ms=optional_float("expected_load_ms"),
load_peak_vram_mb=optional_float("load_peak_vram_mb"),
runtime_peak_vram_mb=optional_float("runtime_peak_vram_mb"),
warmup_peak_vram_mb=optional_float("warmup_peak_vram_mb"),
@@ -172,6 +174,15 @@ class BaselineConfig:
with path.open("r", encoding="utf-8") as fh:
data = json.load(fh)
# runner pools with the same gpu can have different host-side latency
runner_name = os.environ.get("RUNNER_NAME", "")
for prefix, overrides in data.get("runner_overrides", {}).items():
if runner_name.startswith(prefix):
for name, metrics in overrides.items():
data["scenarios"][name].update(metrics)
print(f"--- Performance Runner Baseline: {prefix} ---")
break
# Get tolerance profile, defaulting to 'pr_test'
profile_name = "pr_test"
tolerances = ToleranceConfig.load_profile(
@@ -320,6 +331,7 @@ class DiffusionTestCase:
run_perf_check: bool = True
# Validate every repetition against the same baseline and GT.
perf_repeat_requests: int = 1
perf_warmup_requests: int = 0
run_consistency_check: bool = True
run_component_accuracy_check: bool = True
run_models_api_check: bool = True
@@ -333,12 +345,19 @@ class DiffusionTestCase:
def __post_init__(self) -> None:
if self.perf_repeat_requests < 1:
raise ValueError(f"{self.id}: perf_repeat_requests must be positive")
if self.perf_warmup_requests < 0:
raise ValueError(f"{self.id}: perf_warmup_requests must be non-negative")
if self.sampling_params is None:
object.__setattr__(
self,
"sampling_params",
get_default_sampling_params_for_server_args(self.server_args),
)
if (
self.perf_warmup_requests
and self.sampling_params.realtime_num_chunks is not None
):
raise ValueError(f"{self.id}: request warmup requires non-realtime metrics")
has_startup_lora = self.server_args.lora_path is not None
has_dynamic_lora = self.server_args.dynamic_lora_path is not None
@@ -468,6 +487,8 @@ class PerformanceSummary:
frames_per_second: float | None = None
total_frames: int | None = None
avg_frame_time_ms: float | None = None
denoising_stages: set[str] = field(default_factory=set)
load_time_ms: float | None = None
@staticmethod
def from_req_perf_record(
@@ -489,10 +510,13 @@ class PerformanceSummary:
# convert from list to dict
stage_metrics = {}
denoising_stages = set()
for item in record.stages:
if isinstance(item, dict) and "name" in item:
val = item.get("execution_time_ms", 0.0)
stage_metrics[item["name"]] = val
if item.get("is_denoising", item["name"] == "DenoisingStage"):
denoising_stages.add(item["name"])
load_peak_vram_mb = float(
record.memory_snapshots.get("load_peak", {}).get("peak_reserved_mb", 0.0)
@@ -528,6 +552,7 @@ class PerformanceSummary:
step_metrics=step_durations,
sampled_steps=sampled_steps,
all_denoise_steps=per_step,
denoising_stages=denoising_stages,
load_peak_vram_mb=load_peak_vram_mb,
runtime_peak_vram_mb=runtime_peak_vram_mb,
warmup_peak_vram_mb=warmup_peak_vram_mb,
python/sglang/multimodal_gen/test/test_utils.py
+1 -1
View File
@@ -40,7 +40,7 @@ logger = init_logger(__name__)
# NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows
# publish.
SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion"
SGL_TEST_FILES_CI_DATA_REVISION = "90a87cce5cdef73a9cd461f6d611ac66becef835"
SGL_TEST_FILES_CI_DATA_REVISION = "38ba32bd812b2dfb0eccc83ef063096c089e3389"
# The NPU pin is kept as a separate branch so ascend GT can be bumped independently
# when it's regenerated on its own cadence.
python/sglang/multimodal_gen/test/unit/realtime/test_realtime_consistency_harness.py
+140
View File
@@ -14,6 +14,8 @@ from sglang.multimodal_gen.runtime.realtime.video import (
RAW_RGBA_DELTA_GZIP_CONTENT_TYPE,
build_delta_gzip_raw_rgb_payload,
)
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.server import test_server_common
from sglang.multimodal_gen.test.server.realtime_consistency import (
build_realtime_event_payload,
build_realtime_init_payload,
@@ -24,16 +26,24 @@ from sglang.multimodal_gen.test.server.realtime_consistency import (
prepare_realtime_first_frame,
realtime_ws_url,
record_realtime_key_frames,
record_realtime_perf_stats,
select_realtime_key_frames,
summarize_realtime_perf_stats,
validate_realtime_perf_stats,
)
from sglang.multimodal_gen.test.server.test_server_common import (
DiffusionServerBase,
PerformanceValidationError,
)
from sglang.multimodal_gen.test.server.test_server_utils import get_generate_fn
from sglang.multimodal_gen.test.server.testcase_configs import (
DiffusionSamplingParams,
DiffusionServerArgs,
DiffusionTestCase,
LONGLIVE2_I2V_CI_sampling_params,
LONGLIVE2_T2V_CI_sampling_params,
REALTIME_MODEL_sampling_params,
ScenarioConfig,
)
# Request construction
@@ -340,6 +350,26 @@ def test_collect_realtime_output_skips_and_records_chunk_stats(monkeypatch):
np.testing.assert_array_equal(result.frames[1], second)
assert [stat.chunk_index for stat in result.chunk_stats] == [0, 1]
assert [stat.chunk_total_ms for stat in result.chunk_stats] == [31.0, 32.0]
assert result.e2e_ms > 0
record_realtime_perf_stats("stream-e2e", result.chunk_stats, result.e2e_ms)
case = DiffusionTestCase(
"stream-e2e",
DiffusionServerArgs(model_path="test", modality="video"),
DiffusionSamplingParams(prompt="test"),
run_perf_check=False,
)
runner = DiffusionServerBase()
runner._perf_results = []
monkeypatch.setenv("SGLANG_GEN_BASELINE", "0")
monkeypatch.setitem(
test_server_common.BASELINE_CONFIG.scenarios,
case.id,
ScenarioConfig({}, {}, 1000, 0, 0, expected_load_ms=100),
)
runner._validate_realtime_performance(
SimpleNamespace(load_time_ms=100), case, result.chunk_stats
)
assert runner._perf_results[0]["e2e_ms"] == result.e2e_ms
assert websocket.sent == [
{"type": "init", "prompt": "test"},
{
@@ -350,6 +380,116 @@ def test_collect_realtime_output_skips_and_records_chunk_stats(monkeypatch):
]
@pytest.mark.parametrize("e2e_ms", [None, 0, -1, float("nan"), float("inf")])
def test_realtime_requires_e2e_without_threshold_checks(e2e_ms):
case = DiffusionTestCase(
"missing-stream-e2e",
DiffusionServerArgs(model_path="test", modality="video"),
DiffusionSamplingParams(prompt="test"),
run_perf_check=False,
)
runner = DiffusionServerBase()
runner._perf_results = []
if e2e_ms is not None:
record_realtime_perf_stats(case.id, [], e2e_ms)
with pytest.raises(
PerformanceValidationError, match="E2E duration missing or invalid"
):
runner._validate_realtime_performance(
SimpleNamespace(load_time_ms=100), case, []
)
@pytest.mark.parametrize("baseline", [None, 0, float("nan"), 1000])
def test_realtime_e2e_guard_without_chunk_thresholds(monkeypatch, baseline):
monkeypatch.setenv("SGLANG_GEN_BASELINE", "0")
monkeypatch.setattr(test_server_common.current_platform, "is_hip", lambda: False)
case = DiffusionTestCase(
"stream-e2e-threshold",
DiffusionServerArgs(model_path="test", modality="video"),
DiffusionSamplingParams(prompt="test"),
run_perf_check=False,
)
if baseline is None:
monkeypatch.delitem(
test_server_common.BASELINE_CONFIG.scenarios, case.id, raising=False
)
error, message = pytest.fail.Exception, "not found"
else:
monkeypatch.setitem(
test_server_common.BASELINE_CONFIG.scenarios,
case.id,
ScenarioConfig({}, {}, baseline, 0, 0),
)
error = AssertionError
message = (
"E2E Latency" if baseline == 1000 else "E2E baseline missing or invalid"
)
runner = DiffusionServerBase()
runner._perf_results = []
record_realtime_perf_stats(case.id, [], 2000)
with pytest.raises(error, match=message):
runner._validate_realtime_performance(
SimpleNamespace(load_time_ms=100), case, []
)
assert runner._perf_results[0]["e2e_ms"] == 2000
@pytest.mark.parametrize("peak_mb", [1000, 2000])
def test_realtime_memory_guard_retains_session_e2e(monkeypatch, peak_mb):
case = DiffusionTestCase(
"stream-memory-e2e",
DiffusionServerArgs(model_path="test", modality="video"),
DiffusionSamplingParams(prompt="test"),
run_perf_check=True,
)
scenario = ScenarioConfig(
{},
{},
2000,
0,
0,
load_peak_vram_mb=1000,
runtime_peak_vram_mb=1000,
expected_load_ms=100,
)
monkeypatch.setitem(test_server_common.BASELINE_CONFIG.scenarios, case.id, scenario)
monkeypatch.setattr(test_server_common.current_platform, "is_cuda", lambda: True)
monkeypatch.setattr(test_server_common.current_platform, "is_hip", lambda: False)
monkeypatch.setenv("SGLANG_GEN_BASELINE", "0")
record = RequestPerfRecord(
request_id="last-chunk",
commit_hash="test",
tag="test",
stages=[],
steps=[],
total_duration_ms=20,
memory_snapshots={
"load_peak": {"peak_reserved_mb": 1000},
"runtime_peak": {"peak_reserved_mb": peak_mb},
},
)
monkeypatch.setattr(
test_server_common, "wait_for_req_perf_record", lambda *a, **k: record
)
stats = [
parse_realtime_chunk_stats(
msgspec.msgpack.decode(_packed_realtime_chunk_stats(0))
)
]
record_realtime_perf_stats(case.id, stats, 2000)
runner = DiffusionServerBase()
runner._perf_results = []
ctx = SimpleNamespace(perf_log_path="unused", load_time_ms=100)
if peak_mb > 1000:
with pytest.raises(AssertionError, match="Runtime Peak VRAM"):
runner._validate_realtime_performance(ctx, case, stats)
else:
runner._validate_realtime_performance(ctx, case, stats)
assert runner._perf_results[0]["e2e_ms"] == 2000
assert runner._perf_results[0]["runtime_peak_vram_mb"] == peak_mb
def test_collect_realtime_output_accepts_combined_frame_batch(monkeypatch):
frame = np.arange(12, dtype=np.uint8).reshape(2, 2, 3)
websocket = _FakeRealtimeWebSocket(
python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py
+99
View File
@@ -27,9 +27,13 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import (
from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import (
LongLive2T2VConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2_5 import LTX25PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig
from sglang.multimodal_gen.configs.sample.longlive2 import LongLive2SamplingParams
from sglang.multimodal_gen.configs.sample.ltx_2_5 import LTX25SamplingParams
from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams
from sglang.multimodal_gen.runtime.entrypoints.control_requests import (
SetLoraReq,
UnmergeLoraWeightsReq,
@@ -60,6 +64,8 @@ from sglang.multimodal_gen.runtime.warmup_request_builder import (
should_include_warmup_image,
supports_synthetic_warmup,
)
from sglang.multimodal_gen.test.server.gpu_cases import ONE_GPU_CASES, TWO_GPU_CASES
from sglang.multimodal_gen.test.server.testcase_configs import _get_extra_arg_value
def _make_bare_scheduler(enable_cfg_parallel: bool) -> Scheduler:
@@ -653,6 +659,99 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
self.assertEqual(num_frames, 17)
pipeline_config.adjust_num_frames.assert_called_once_with(17)
def test_server_warmup_preserves_explicit_frames_without_cuda_graphs(self):
server_args = SimpleNamespace(
pipeline_config=LTX25PipelineConfig(),
enable_breakable_cuda_graph=False,
pipeline_class_name="LTX2Pipeline",
num_gpus=2,
warmup_num_frames=49,
)
num_frames = _resolve_warmup_num_frames(
server_args, LTX25SamplingParams(), server_based_warmup=True
)
self.assertEqual(num_frames, 57)
def test_sana_ci_warmup_matches_formal_shape(self):
case = next(case for case in ONE_GPU_CASES if case.id == "sana_wm_ti2v")
resolution = _get_extra_arg_value(
case.server_args.extras, "--warmup-resolutions"
)
server_args = SimpleNamespace(
pipeline_config=SanaWMPipelineConfig(),
pipeline_class_name=None,
model_path=case.server_args.model_path,
model_id=None,
backend="sglang",
num_gpus=1,
warmup_steps=1,
warmup_num_frames=None,
warmup_sampling_params=None,
enable_breakable_cuda_graph=False,
enable_torch_compile=False,
enable_cfg_parallel=False,
)
with patch.object(
SamplingParams, "from_pretrained", return_value=SanaWMSamplingParams()
):
reqs = build_warmup_reqs(
server_args,
warmup_resolutions=[resolution],
warmup_input_path="synthetic-warmup.png",
server_based_warmup=True,
)
self.assertEqual(resolution, case.sampling_params.output_size)
self.assertEqual(len(reqs), 1)
self.assertEqual((reqs[0].width, reqs[0].height), (384, 640))
self.assertEqual(reqs[0].num_frames, case.sampling_params.num_frames)
def test_ltx25_ci_warmup_matches_formal_decoder_and_shape(self):
case = next(
case
for case in TWO_GPU_CASES
if case.id == "ltx_2_5_diffusion_decoder_2gpus"
)
extras = case.server_args.extras
server_args = SimpleNamespace(
pipeline_config=LTX25PipelineConfig(),
pipeline_class_name=None,
model_path=case.server_args.model_path,
model_id=None,
backend="sglang",
num_gpus=2,
warmup_steps=1,
warmup_num_frames=int(_get_extra_arg_value(extras, "--warmup-num-frames")),
warmup_sampling_params=_get_extra_arg_value(
extras, "--warmup-sampling-params"
),
enable_breakable_cuda_graph=False,
enable_torch_compile=False,
enable_cfg_parallel=False,
)
resolution = _get_extra_arg_value(extras, "--warmup-resolutions")
with patch.object(
SamplingParams, "from_pretrained", return_value=LTX25SamplingParams()
):
reqs = build_warmup_reqs(
server_args,
warmup_resolutions=[resolution],
warmup_input_path="synthetic-warmup.png",
server_based_warmup=True,
)
self.assertEqual(len(reqs), 1)
req = reqs[0]
self.assertEqual(resolution, case.sampling_params.output_size)
self.assertEqual(server_args.warmup_num_frames, case.sampling_params.num_frames)
self.assertEqual((req.width, req.height, req.num_frames), (768, 448, 57))
self.assertEqual(
req.sampling_params.use_diffusion_decoder,
case.sampling_params.extras["use_diffusion_decoder"],
)
self.assertEqual(req.num_inference_steps, 2)
def test_server_based_warmup_uses_video_supported_resolution_budget(self):
server_args = MagicMock()
server_args.warmup_steps = 1
python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py
+59 -27
View File
@@ -2,6 +2,7 @@ from contextlib import ExitStack
from types import SimpleNamespace
from unittest.mock import call, patch
import pytest
import torch
from torch import nn
@@ -17,7 +18,7 @@ from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import
initialize_parallel_runtime,
)
from sglang.srt.distributed import parallel_state as srt_parallel_state
from sglang.srt.runtime_context import get_parallel
from sglang.srt.runtime_context import ParallelContext, get_parallel
_UTILS = "sglang.multimodal_gen.test.single_test_file.component_accuracy.utils"
@@ -135,7 +136,7 @@ def test_srt_attention_tp_group_tracks_diffusion_tp_group():
# `world_size`, because lending the group also states the parallel widths it
# implies -- the shared `srt` vision layers ask for `attn_tp_size`, and this
# package publishes no `srt` config for that read to resolve against.
tp_group = SimpleNamespace(world_size=2)
tp_group = SimpleNamespace(world_size=2, rank_in_group=1)
with (
patch.object(parallel_state, "_TP", tp_group),
@@ -147,11 +148,17 @@ def test_srt_attention_tp_group_tracks_diffusion_tp_group():
assert srt_parallel_state._TP is tp_group
assert srt_parallel_state._ATTN_TP is tp_group
assert get_parallel().attn_tp_size == 2
assert get_parallel().tp_group is tp_group
assert get_parallel().attn_tp_group is tp_group
assert get_parallel().tp_rank == 1
assert get_parallel().attn_tp_rank == 1
parallel_state._clear_srt_tp_group()
assert srt_parallel_state._TP is None
assert srt_parallel_state._ATTN_TP is None
with pytest.raises(RuntimeError):
get_parallel().tp_group
def test_srt_owned_groups_are_not_overwritten_or_cleared():
@@ -171,62 +178,87 @@ def test_srt_owned_groups_are_not_overwritten_or_cleared():
assert srt_parallel_state._ATTN_TP is srt_attention_tp_group
def test_srt_tp_groups_follow_encoder_folding_context():
original_diffusion_tp_group = object()
original_srt_tp_group = object()
original_srt_attention_tp_group = object()
folding_tp_group = _tp_group(world_size=2, rank_in_group=1)
@pytest.mark.parametrize("rank", [0, 1])
def test_srt_tp_groups_follow_encoder_folding_context(rank):
original_tp_group = _tp_group()
folding_tp_group = _tp_group(world_size=2, rank_in_group=rank)
with (
patch.object(parallel_state, "_TP", original_diffusion_tp_group),
patch.object(srt_parallel_state, "_TP", original_srt_tp_group),
patch.object(
srt_parallel_state,
"_ATTN_TP",
original_srt_attention_tp_group,
),
patch.object(parallel_state, "_TP", original_tp_group),
patch.object(srt_parallel_state, "_TP", None),
patch.object(srt_parallel_state, "_ATTN_TP", None),
patch("sglang.srt.runtime_context._PARALLEL", ParallelContext()),
):
# Diffusion initialization stamps the original TP=1 group before an
# encoder temporarily folds the sequence-parallel group into TP=2.
parallel_state._sync_srt_tp_group()
with parallel_state.use_tensor_parallel_group(folding_tp_group):
assert parallel_state._TP is folding_tp_group
assert srt_parallel_state._TP is folding_tp_group
assert srt_parallel_state._ATTN_TP is folding_tp_group
assert get_parallel().tp_size == 2
assert get_parallel().tp_rank == 1
assert get_parallel().tp_rank == rank
assert get_parallel().tp_group is folding_tp_group
assert get_parallel().attn_tp_size == 2
assert get_parallel().attn_tp_rank == rank
assert get_parallel().attn_tp_group is folding_tp_group
assert get_parallel().moe_tp_rank == rank
assert parallel_state._TP is original_diffusion_tp_group
assert srt_parallel_state._TP is original_srt_tp_group
assert srt_parallel_state._ATTN_TP is original_srt_attention_tp_group
assert parallel_state._TP is original_tp_group
assert srt_parallel_state._TP is original_tp_group
assert srt_parallel_state._ATTN_TP is original_tp_group
assert get_parallel().tp_size == 1
assert get_parallel().tp_rank == 0
assert get_parallel().tp_group is original_tp_group
assert get_parallel().attn_tp_size == 1
assert get_parallel().attn_tp_rank == 0
assert get_parallel().attn_tp_group is original_tp_group
assert get_parallel().moe_tp_rank == 0
def test_encoder_folding_context_is_nested_and_restores_each_group():
original_tp_group = object()
original_tp_group = _tp_group()
outer_tp_group = _tp_group(world_size=4, rank_in_group=3)
inner_tp_group = _tp_group(world_size=2, rank_in_group=1)
with (
patch.object(parallel_state, "_TP", original_tp_group),
patch.object(srt_parallel_state, "_TP", original_tp_group),
patch.object(srt_parallel_state, "_ATTN_TP", original_tp_group),
patch.object(srt_parallel_state, "_TP", None),
patch.object(srt_parallel_state, "_ATTN_TP", None),
patch("sglang.srt.runtime_context._PARALLEL", ParallelContext()),
):
parallel_state._sync_srt_tp_group()
with parallel_state.use_tensor_parallel_group(outer_tp_group):
assert get_parallel().tp_size == 4
with parallel_state.use_tensor_parallel_group(inner_tp_group):
assert parallel_state._TP is inner_tp_group
assert srt_parallel_state._TP is inner_tp_group
assert srt_parallel_state._ATTN_TP is inner_tp_group
assert get_parallel().tp_size == 2
assert get_parallel().tp_rank == 1
with pytest.raises(RuntimeError, match="encoder load failed"):
with parallel_state.use_tensor_parallel_group(inner_tp_group):
assert parallel_state._TP is inner_tp_group
assert srt_parallel_state._TP is inner_tp_group
assert srt_parallel_state._ATTN_TP is inner_tp_group
assert get_parallel().tp_size == 2
assert get_parallel().tp_rank == 1
assert get_parallel().attn_tp_group is inner_tp_group
assert get_parallel().attn_tp_rank == 1
assert get_parallel().moe_tp_rank == 1
raise RuntimeError("encoder load failed")
assert parallel_state._TP is outer_tp_group
assert srt_parallel_state._TP is outer_tp_group
assert srt_parallel_state._ATTN_TP is outer_tp_group
assert get_parallel().tp_size == 4
assert get_parallel().tp_rank == 3
assert get_parallel().attn_tp_group is outer_tp_group
assert get_parallel().attn_tp_rank == 3
assert get_parallel().moe_tp_rank == 3
assert parallel_state._TP is original_tp_group
assert srt_parallel_state._TP is original_tp_group
assert srt_parallel_state._ATTN_TP is original_tp_group
assert get_parallel().tp_group is original_tp_group
assert get_parallel().tp_rank == 0
assert get_parallel().attn_tp_group is original_tp_group
assert get_parallel().attn_tp_rank == 0
assert get_parallel().moe_tp_rank == 0
def test_weight_transfer_uses_loader_for_implicit_srt_shard():
python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py
+16 -3
View File
@@ -2,11 +2,14 @@
import unittest
from types import SimpleNamespace
from unittest.mock import Mock, patch
import numpy as np
import torch
from diffusers import AutoencoderKL as DiffusersAutoencoderKL
from diffusers import LCMScheduler, UNet2DConditionModel
from sglang.kernels.ops.diffusion.ext import mesh_processor
from sglang.multimodal_gen.configs.models.vaes.stable_diffusion import (
StableDiffusionVAEConfig,
)
@@ -185,9 +188,19 @@ class TestHunyuan3DWarmupOutput(unittest.TestCase):
def test_paint_postprocess_skips_export_during_warmup(self):
stage = Hunyuan3DPaintPostprocessStage(Hunyuan3D2PipelineConfig())
output = stage.forward(self._batch(), SimpleNamespace())
kernel = Mock()
with (
patch.object(mesh_processor, "_mesh_processor_kernel", None),
patch.object(
mesh_processor, "load_extension_with_recovery", return_value=kernel
) as build,
):
output = stage.forward(self._batch(), SimpleNamespace())
build.assert_called_once()
array = np.zeros((1, 3), dtype=np.float32)
mesh_processor.meshVerticeInpaint(array, array, array, array, array, array)
build.assert_called_once()
kernel.meshVerticeInpaint.assert_called_once()
self.assertEqual(output.output_file_paths, [])
python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
+75
View File
@@ -1058,6 +1058,81 @@ def test_prepare_for_next_req_repins_residents(monkeypatch):
assert {0, 1, 2} <= manager._gpu_layers
def test_release_after_use_defaults_to_the_old_release_all(monkeypatch):
"""The rename must not move anything: `release_after_use()` == the previous call.
`finish_use` used to call `release_all()` unconditionally. It now says
`release_after_use()`, and with the default argument that has to clear exactly the
same layers, or this refactor is a behaviour change wearing a new name.
"""
_patch_fake_device(monkeypatch)
manager = _resident_manager(
_MultiBlockModel(6), num_layers=6, prefetch_size=1, resident_layers=3
)
_arm_residency(manager)
manager.prepare_for_next_req(non_blocking=False)
assert manager._gpu_layers
manager.release_after_use()
assert not manager._gpu_layers
assert manager._first_pass is True
def test_release_all_still_drops_everything(monkeypatch):
"""`release_all` keeps its literal contract for the full-reset callers.
`enable_offload` syncs to CPU and expects nothing left on the device; it
must not inherit the resident-set exemption.
"""
_patch_fake_device(monkeypatch)
manager = _resident_manager(
_MultiBlockModel(6), num_layers=6, prefetch_size=1, resident_layers=3
)
_arm_residency(manager)
manager.prepare_for_next_req(non_blocking=False)
manager.release_all()
assert not manager._gpu_layers
assert manager._first_pass is True
def test_release_after_use_can_keep_the_resident_set(monkeypatch):
"""`keep_resident` is the whole point of naming the two calls apart.
A component whose use is one forward pass has its resident set prefetched
at the start of the use and dropped at the end, so `resident_layers` buys
it nothing. Measured on Qwen-Image-2.1 / RTX 5090:
`--layerwise-resident-layers text_encoder=0.8` logs `resident=53/66` and
moves neither memory nor latency.
"""
_patch_fake_device(monkeypatch)
manager = _resident_manager(
_MultiBlockModel(6), num_layers=6, prefetch_size=1, resident_layers=3
)
_arm_residency(manager)
manager.prepare_for_next_req(non_blocking=False)
manager.release_after_use(keep_resident=True)
assert set(manager._gpu_layers) == set(manager._retained_set)
# Those layers never left the device, so the next use must not re-do the
# sequential first pass that exists for evicted pages.
assert manager._first_pass is False
def test_release_after_use_keeps_nothing_when_no_residents_are_configured(monkeypatch):
"""`keep_resident` with an empty resident set is still a full release."""
_patch_fake_device(monkeypatch)
manager = _resident_manager(
_MultiBlockModel(6), num_layers=6, prefetch_size=1, resident_layers=0
)
manager.prefetch_layer(0, non_blocking=False)
manager.prefetch_layer(1, non_blocking=False)
assert manager._gpu_layers
manager.release_after_use(keep_resident=True)
assert not manager._gpu_layers
def _record_prepare(manager, monkeypatch):
"""Log the order of prefetches and stream waits inside prepare_for_next_req.
python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py
+194
View File
@@ -0,0 +1,194 @@
import io
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
from sglang.multimodal_gen.runtime import launch_server as launcher
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.scripts import gen_perf_baselines
from sglang.multimodal_gen.test.server import test_server_utils as utils
from sglang.multimodal_gen.test.server.test_server_utils import PerformanceValidator
from sglang.multimodal_gen.test.server.testcase_configs import (
DiffusionSamplingParams,
DiffusionServerArgs,
DiffusionTestCase,
PerformanceSummary,
ScenarioConfig,
ToleranceConfig,
)
@pytest.fixture
def validator(monkeypatch):
monkeypatch.setenv("SGLANG_GEN_BASELINE", "0")
scenario = ScenarioConfig.from_dict(
{
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 1000,
"expected_avg_denoise_ms": 0,
"expected_median_denoise_ms": 0,
"expected_load_ms": 4000,
}
)
return PerformanceValidator(scenario, ToleranceConfig(0.25, 0, 0, 0, 0), [])
def test_slow_loading_fails_even_when_inference_passes(validator):
summary = PerformanceSummary(1000, 0, 0, {}, [], {}, {}, load_time_ms=6000)
validator.validate_e2e(summary)
with pytest.raises(AssertionError, match="Load Latency"):
validator.validate_load(summary)
def test_fast_loading_cannot_hide_inference_regression(validator):
summary = PerformanceSummary(2000, 0, 0, {}, [], {}, {}, load_time_ms=1000)
validator.validate_load(summary)
with pytest.raises(AssertionError, match="E2E Latency"):
validator.validate_e2e(summary)
@pytest.mark.parametrize("load_time_ms", [None, 0, float("nan"), 1234.5])
def test_baseline_script_preserves_required_load_measurement(monkeypatch, load_time_ms):
case = DiffusionTestCase(
"load-baseline",
DiffusionServerArgs("test", modality="image"),
DiffusionSamplingParams(prompt="test"),
)
context = SimpleNamespace(
port=1234, load_time_ms=load_time_ms, perf_log_path="unused", cleanup=Mock()
)
monkeypatch.setattr(
gen_perf_baselines,
"ServerManager",
Mock(return_value=Mock(start=lambda: context)),
)
monkeypatch.setattr(gen_perf_baselines, "get_dynamic_server_port", lambda: 1234)
monkeypatch.setattr(gen_perf_baselines, "_build_server_extra_args", lambda case: "")
monkeypatch.setattr(gen_perf_baselines, "_openai_client", Mock())
monkeypatch.setattr(
gen_perf_baselines,
"get_generate_fn",
lambda **kwargs: lambda *args: ("request", b"output"),
)
monkeypatch.setattr(gen_perf_baselines.current_platform, "is_cuda", lambda: False)
record = RequestPerfRecord(
request_id="request",
commit_hash="test",
tag="test",
stages=[],
steps=[],
total_duration_ms=1000,
)
monkeypatch.setattr(
gen_perf_baselines, "wait_for_req_perf_record", lambda *args, **kwargs: record
)
if load_time_ms == 1234.5:
scenario = ScenarioConfig.from_dict(gen_perf_baselines._run_case(case))
assert scenario.expected_load_ms == 1234.5
assert scenario.expected_e2e_ms == 1000
else:
with pytest.raises(ValueError, match="load duration missing or invalid"):
gen_perf_baselines._run_case(case)
context.cleanup.assert_called_once()
@pytest.mark.parametrize("duration", [None, 0, -1, float("nan"), float("inf")])
def test_missing_or_invalid_load_duration_fails(validator, duration):
summary = PerformanceSummary(1000, 0, 0, {}, [], {}, {}, load_time_ms=duration)
with pytest.raises(AssertionError, match="Load duration missing or invalid"):
validator.validate_load(summary)
@pytest.mark.parametrize("duration", [None, 0, -1, float("nan"), float("inf")])
def test_missing_or_invalid_load_baseline_fails(validator, duration):
validator.scenario.expected_load_ms = duration
summary = PerformanceSummary(1000, 0, 0, {}, [], {}, {}, load_time_ms=4000)
with pytest.raises(AssertionError, match="Load baseline missing or invalid"):
validator.validate_load(summary)
def test_fast_inference_cannot_hide_loading_regression(validator):
summary = PerformanceSummary(1, 0, 0, {}, [], {}, {}, load_time_ms=5500)
validator.validate_e2e(summary)
with pytest.raises(AssertionError, match="Load Latency"):
validator.validate_load(summary)
def test_repeated_requests_use_same_load_measurement(validator):
for inference_ms in (1000, 900, 950):
summary = PerformanceSummary(
inference_ms, 0, 0, {}, [], {}, {}, load_time_ms=4000
)
validator.validate_e2e(summary)
validator.validate_load(summary)
@pytest.mark.parametrize("workers", [1, 2])
@pytest.mark.parametrize("warmup_seconds", [0, 120])
def test_server_load_clock_excludes_warmup(
monkeypatch, tmp_path, workers, warmup_seconds, validator
):
clock = [1_000_000_000]
output = io.StringIO()
monkeypatch.setattr(utils.time, "monotonic_ns", lambda: clock[0])
monkeypatch.setattr(utils.current_platform, "is_hip", lambda: False)
monkeypatch.setattr(utils.tempfile, "gettempdir", lambda: str(tmp_path))
monkeypatch.setattr(
utils, "prepare_perf_log", lambda: (tmp_path, tmp_path / "perf.jsonl")
)
monkeypatch.setattr(launcher, "configure_logger", Mock())
monkeypatch.setattr(launcher, "logger", Mock())
launcher.logger.info.side_effect = lambda message, *args: output.write(
(message % args) + "\n"
)
def ready():
clock[0] += 1_000_000_000
return {"status": "ready"}
worker_context = Mock()
worker_context.Pipe.side_effect = lambda **kwargs: (Mock(recv=ready), Mock())
monkeypatch.setattr(launcher.mp, "get_context", Mock(return_value=worker_context))
monkeypatch.setattr(launcher, "shutdown_scheduler_processes", Mock())
def warmup(args):
clock[0] += warmup_seconds * 1_000_000_000
monkeypatch.setattr(launcher, "launch_http_server_only", warmup)
monkeypatch.setattr(ServerArgs, "__post_init__", lambda self: None)
args = ServerArgs(
model_path="test",
num_gpus=workers,
nnodes=1,
node_rank=0,
master_port=1234,
webui=False,
pipeline_config=PipelineConfig(),
)
def spawn(*unused_args, **unused_kwargs):
launcher.launch_server(args)
return SimpleNamespace(pid=1234, stdout=io.StringIO(output.getvalue()))
monkeypatch.setattr(utils.subprocess, "Popen", spawn)
manager = utils.ServerManager("test", 1234)
monkeypatch.setattr(manager, "_wait_for_ready", Mock())
context = manager.start()
launcher.mp.get_context.assert_called_once_with("spawn")
assert worker_context.Process.call_count == workers
for call in worker_context.Process.call_args_list:
restored_args = call.kwargs["args"][0].server_args.materialize()
assert isinstance(restored_args, ServerArgs)
assert restored_args.num_gpus == workers
context._log_thread.join(timeout=5)
assert not context._log_thread.is_alive()
assert context.load_time_ms == workers * 1000
summary = PerformanceSummary(
1000, 0, 0, {}, [], {}, {}, load_time_ms=context.load_time_ms
)
validator.validate_e2e(summary)
validator.validate_load(summary)
python/sglang/multimodal_gen/test/unit/test_performance_failure_policy.py
+268
View File
@@ -0,0 +1,268 @@
import os
import subprocess
import sys
import textwrap
import time
import pytest
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.runner.pytest_runner import (
_estimate_failed_test_time,
_is_retryable_failure,
run_pytest,
)
from sglang.multimodal_gen.test.server import test_server_common as common
from sglang.multimodal_gen.test.server.testcase_configs import (
DiffusionSamplingParams,
DiffusionServerArgs,
DiffusionTestCase,
ScenarioConfig,
)
@pytest.mark.parametrize("generate_baseline", [False, True])
def test_e2e_only_does_not_require_stage_metrics(monkeypatch, generate_baseline):
monkeypatch.setenv("SGLANG_GEN_BASELINE", str(int(generate_baseline)))
case = DiffusionTestCase(
"e2e_only",
DiffusionServerArgs(model_path="test", modality="image"),
DiffusionSamplingParams(prompt="test"),
run_perf_check=False,
)
scenario = ScenarioConfig({}, {}, 1000, 0, 0, expected_load_ms=100)
monkeypatch.setitem(common.BASELINE_CONFIG.scenarios, case.id, scenario)
monkeypatch.setattr(common, "_PENDING_BASELINE_DUMPS", {})
server = common.DiffusionServerBase()
server._perf_results = []
record = RequestPerfRecord(
request_id="guard",
commit_hash="test",
tag="guard",
stages=[],
steps=[],
total_duration_ms=2000 if generate_baseline else 1000,
)
server._validate_and_record(case, record, load_time_ms=100)
assert len(server._perf_results) == 1
assert bool(common._PENDING_BASELINE_DUMPS) == generate_baseline
@pytest.mark.parametrize(
"output",
[
"multimodal_gen/test/server/test_server_utils.py: AssertionError",
"Consistency check failed for example\nTimeoutError",
"[performance] Validation failed\nConsistency check failed for example",
],
)
def test_validation_failures_are_not_retryable(output):
assert not _is_retryable_failure(output)
@pytest.mark.parametrize(
"output",
[
"[performance] Validation failed for 'E2E Latency'",
"[performance] Validation failed for 'Load Latency (excluding warmup)'",
"[performance] Validation failed for 'Average Denoise Step'\nTimeoutError",
"[performance] E2E missing or invalid\nCUDA out of memory",
],
)
def test_performance_failures_are_retryable(output):
assert _is_retryable_failure(output)
@pytest.mark.parametrize(
"output", ["TimeoutError", "SafetensorError", "CUDA out of memory"]
)
def test_infrastructure_failure_policy_is_unchanged(output):
assert _is_retryable_failure(output)
@pytest.mark.parametrize("with_deadline", [False, True])
def test_performance_retry_recovers_only_failed_items(
tmp_path, monkeypatch, with_deadline
):
monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
if with_deadline:
monkeypatch.setenv("SGLANG_DIFFUSION_RETRY_DEADLINE", str(time.time() + 600))
else:
monkeypatch.delenv("SGLANG_DIFFUSION_RETRY_DEADLINE", raising=False)
test_file = tmp_path / "test_retry.py"
test_file.write_text(
"from pathlib import Path\n"
"def test_slow():\n"
" marker = Path(__file__).with_suffix('.attempt')\n"
" if not marker.exists():\n"
" marker.touch()\n"
" assert False, '[performance] Validation failed for E2E Latency'\n"
"def test_fast():\n"
" marker = Path(__file__).with_suffix('.passed')\n"
" assert not marker.exists(), 'passing case must not rerun'\n"
" marker.touch()\n"
)
code, _, _ = run_pytest([str(test_file)])
assert code == 0
assert test_file.with_suffix(".attempt").exists()
assert test_file.with_suffix(".passed").exists()
def test_retry_budget_preserves_failure_and_report(tmp_path, monkeypatch, capfd):
monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
monkeypatch.setenv("SGLANG_DIFFUSION_RETRY_DEADLINE", str(time.time() - 1))
test_file = tmp_path / "test_budget.py"
test_file.write_text(
"import pytest\n"
"@pytest.mark.parametrize('case_id', ['slow_case'])\n"
"def test_slow(case_id):\n"
" assert False, '[performance] Validation failed for E2E Latency'\n"
)
report = tmp_path / "junit.xml"
code, executed, results = run_pytest([str(test_file)], junit_xml_path=str(report))
output = capfd.readouterr().out
assert code == 1
assert executed == ["slow_case"]
assert results == {"slow_case": "fail"}
assert output.count("Starting pytest attempt") == 1
assert "Retry budget exhausted" in output
assert "Pytest Tail Summary" in output
def test_retry_estimate_excludes_successful_cases(tmp_path):
report = tmp_path / "junit.xml"
report.write_text(
"<testsuites><testsuite>"
'<testcase name="passed" time="100" />'
'<testcase name="failed" time="10"><failure /></testcase>'
'<testcase name="error" time="20"><error /></testcase>'
"</testsuite></testsuites>"
)
assert _estimate_failed_test_time(str(report), 130) == 30
assert _estimate_failed_test_time(None, 130) == 130
@pytest.mark.parametrize(
"problem",
[
"regression",
"missing_baseline",
"missing_record",
"missing_e2e",
"missing_log",
"e2e_only_regression",
"e2e_only_missing_baseline",
"e2e_only_zero_baseline",
"e2e_only_nan_baseline",
],
)
def test_performance_failure_survives_real_pytest_runner(tmp_path, problem):
# exercise the validator, request loop, pytest output and retry classifier together
test_file = tmp_path / "test_guard.py"
test_file.write_text(
textwrap.dedent(
"""
import pytest
from types import SimpleNamespace
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.server import test_server_common as common
from sglang.multimodal_gen.test.test_utils import wait_for_req_perf_record
from sglang.multimodal_gen.test.server.testcase_configs import (
DiffusionSamplingParams, DiffusionServerArgs, DiffusionTestCase, ScenarioConfig,
)
@pytest.mark.parametrize("case_id", ["threshold_guard"])
def test_guard(case_id, monkeypatch, tmp_path):
server = common.DiffusionServerBase()
server._perf_results = []
case = DiffusionTestCase(
case_id,
DiffusionServerArgs(model_path="test", modality="image", lora_path="test-lora"),
DiffusionSamplingParams(prompt="test"),
run_lora_basic_api_check=True, perf_repeat_requests=2,
run_consistency_check=False, run_models_api_check=False,
run_perf_check=not PROBLEM.startswith("e2e_only_") and PROBLEM not in ("missing_record", "missing_e2e", "missing_log"),
)
scenario = ScenarioConfig({}, {}, 1000, 100, 100, expected_load_ms=100)
if PROBLEM == "e2e_only_zero_baseline":
scenario.expected_e2e_ms = 0
if PROBLEM == "e2e_only_nan_baseline":
scenario.expected_e2e_ms = float("nan")
if PROBLEM in ("missing_baseline", "e2e_only_missing_baseline"):
monkeypatch.delitem(common.BASELINE_CONFIG.scenarios, case_id, raising=False)
else:
monkeypatch.setitem(common.BASELINE_CONFIG.scenarios, case_id, scenario)
monkeypatch.setattr(common.current_platform, "is_cuda", lambda: False)
monkeypatch.setattr(common.current_platform, "is_hip", lambda: False)
monkeypatch.setattr(common, "get_generate_fn", lambda **kwargs: None)
requests = []
lora_checks = []
def collect(*args, **kwargs):
requests.append(1)
if PROBLEM == "missing_record":
return None, b""
return RequestPerfRecord(
request_id="guard", commit_hash="test", tag="guard",
stages=[], steps=[100],
total_duration_ms=None if PROBLEM == "missing_e2e" else 2000,
), b""
context = SimpleNamespace(load_time_ms=100)
if PROBLEM == "missing_log":
log_path = tmp_path / "empty-perf.jsonl"
log_path.write_text("")
context = SimpleNamespace(perf_log_path=log_path, load_time_ms=100)
monkeypatch.setattr(server, "_client", lambda ctx: None)
def generate(*args):
requests.append(1)
return "guard", b""
monkeypatch.setattr(server, "_run_generation_with_server_watchdog", generate)
monkeypatch.setattr(
common, "wait_for_req_perf_record",
lambda rid, path, timeout: wait_for_req_perf_record(rid, path, timeout=0.01),
)
else:
monkeypatch.setattr(server, "run_and_collect", collect)
monkeypatch.setattr(
server, "_test_lora_api_functionality",
lambda *args: lora_checks.append(1),
)
try:
server._test_diffusion_generation_impl(case, context)
finally:
print(f"GUARD_REQUESTS={len(requests)} LORA_CHECKS={len(lora_checks)}")
print(f"RETAINED_METRICS={len(server._perf_results)}")
def test_unrelated_timeout():
raise TimeoutError("independent infrastructure failure")
"""
).replace("PROBLEM", repr(problem))
)
report = tmp_path / "junit.xml"
env = os.environ.copy()
env.update(
PYTEST_DISABLE_PLUGIN_AUTOLOAD="1",
SGLANG_GEN_BASELINE="0",
SGLANG_GEN_GT="0",
)
command = (
"from sglang.multimodal_gen.test.runner.pytest_runner import run_pytest; "
f"result = run_pytest([{str(test_file)!r}], junit_xml_path={str(report)!r}); "
"print('GUARD_RESULT', result); raise SystemExit(result[0])"
)
result = subprocess.run(
[sys.executable, "-c", command],
cwd=tmp_path,
env=env,
capture_output=True,
text=True,
timeout=600,
)
output = result.stdout + result.stderr
assert result.returncode == 1, output
assert "[performance]" in output, output
assert "GUARD_REQUESTS=1 LORA_CHECKS=0" in output, output
retained = 0 if problem in {"missing_record", "missing_e2e", "missing_log"} else 1
assert f"RETAINED_METRICS={retained}" in output, output
assert output.count("Starting pytest attempt") == 7, output
assert "Max retry exceeded (6)" in output, output
assert "'threshold_guard': 'fail'" in output, output
python/sglang/multimodal_gen/test/unit/test_performance_metrics.py
+83
View File
@@ -7,14 +7,30 @@ import torch
import sglang.multimodal_gen.runtime.managers.gpu_worker as gpu_worker_module
import sglang.multimodal_gen.runtime.managers.memory_managers.component_manager as component_manager_module
import sglang.multimodal_gen.runtime.utils.perf_logger as perf_logger_module
from sglang.multimodal_gen.runtime.disaggregation.orchestrator import (
_deserialize_request_metrics,
)
from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
WarmupPhasePeak,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising_av import (
LTX2RefinementStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.stages.denoising import (
MiniMaxH3DenoisingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
TextEncodingStage,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.perf_logger import (
MemorySnapshot,
PerformanceLogger,
RequestMetrics,
RequestPerfRecord,
)
@@ -26,6 +42,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
ScenarioConfig,
ToleranceConfig,
)
from sglang.multimodal_gen.test.test_utils import read_perf_logs
@pytest.fixture(autouse=True)
@@ -71,6 +88,72 @@ def test_request_metrics_attributes_steps_and_iterations_to_active_stage():
}
@pytest.mark.parametrize("roundtrip", [False, True])
@pytest.mark.parametrize(
"stage_class,profile_name,is_denoising",
[
(DenoisingStage, "DenoisingStage", True),
(MiniMaxH3DenoisingStage, "MiniMaxH3DenoisingStage", True),
(LTX2RefinementStage, "LTX2RefinementStage", True),
(LTX2RefinementStage, "custom_refinement", True),
(DenoisingStage, "BeforeDenoisingStage", True),
(TextEncodingStage, "BeforeDenoisingStage", False),
(TextEncodingStage, "TextEncodingStage", False),
],
)
def test_stage_role_reaches_performance_guard(
stage_class, profile_name, is_denoising, roundtrip, monkeypatch, tmp_path
):
# skip model construction and kernels, retaining the real stage role,
# call boundary, profiler, log writer/reader and threshold validator
stage = stage_class.__new__(stage_class)
stage.server_args = SimpleNamespace(
enable_layerwise_nvtx_marker=False, comfyui_mode=False
)
stage.set_profile_stage_name(profile_name)
monkeypatch.setattr(stage, "forward", lambda batch, args: batch)
monkeypatch.setattr(
stage, "verify_input", PipelineStage.verify_input.__get__(stage)
)
monkeypatch.setattr(
stage, "verify_output", PipelineStage.verify_output.__get__(stage)
)
monkeypatch.setattr(current_platform, "get_available_gpu_memory", lambda **_: 100)
monkeypatch.setattr(current_platform, "is_hip", lambda: False)
monkeypatch.setenv("SGLANG_PERF_LOG_DIR", str(tmp_path))
monkeypatch.setattr(perf_logger_module, "get_is_main_process", lambda: True)
monkeypatch.setattr(perf_logger_module, "get_git_commit_hash", lambda: "test")
metrics = RequestMetrics("stage-role")
batch = SimpleNamespace(is_warmup=False, metrics=metrics, perf_dump_path="metrics")
with patch.object(perf_logger_module.time, "perf_counter", side_effect=[10, 11.5]):
assert stage(batch, stage.server_args) is batch
metrics.total_duration_ms = 1500
if roundtrip:
metrics = _deserialize_request_metrics(
json.loads(json.dumps(metrics.to_dict()))
)
PerformanceLogger.log_request_summary(metrics)
(record,) = read_perf_logs(tmp_path / "performance.log")
assert record.stages == [
{
"name": profile_name,
"execution_time_ms": 1500.0,
"is_denoising": is_denoising,
}
]
validator = PerformanceValidator(
ScenarioConfig({profile_name: 1000}, {}, 1500, 1, 1),
ToleranceConfig(0.25, 0.25, 0.8, 0.3, 0.2),
(),
)
summary = validator.collect_metrics(record)
if is_denoising:
with pytest.raises(AssertionError, match="Stage '"):
validator._validate_stages(summary)
else:
validator._validate_stages(summary)
def test_performance_summary_separates_load_and_runtime_peaks():
summary = PerformanceSummary.from_req_perf_record(
_perf_record(
python/sglang/multimodal_gen/test/unit/test_qwen_image21.py
+92
View File
@@ -36,6 +36,7 @@ from sglang.multimodal_gen.runtime.models.encoders.qwen3vl_vision import (
from sglang.multimodal_gen.runtime.models.vaes.autoencoder_kl_qwenimage21 import (
AutoencoderKLQwenImage21,
QwenImage21RMS_norm,
QwenImage21Upsample,
_patchify,
_unpatchify,
)
@@ -50,6 +51,28 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.q
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("layout", ["contiguous", "channels_last", "transposed"])
@pytest.mark.parametrize("device", ["cpu", "cuda"])
def test_nearest_upsample_preserves_every_finite_low_precision_value(
dtype, layout, device
):
if device == "cuda" and not torch.cuda.is_available():
pytest.skip("CUDA required")
values = torch.arange(65536, dtype=torch.int32).to(torch.int16).view(dtype)
values = values[torch.isfinite(values)].reshape(1, 2, -1, 128).to(device)
if layout == "channels_last":
values = values.contiguous(memory_format=torch.channels_last)
elif layout == "transposed":
values = values.transpose(2, 3)
upsample = QwenImage21Upsample(scale_factor=2, mode="nearest-exact")
expected = torch.nn.functional.interpolate(
values.float(), scale_factor=2, mode="nearest-exact"
).to(dtype)
actual = upsample(values)
assert torch.equal(actual.view(torch.int16), expected.view(torch.int16))
@pytest.mark.parametrize("prompt", ["edit", ""])
@pytest.mark.parametrize("image_count", [0, 1, 2])
def test_prompt_conditioning_uses_training_template_and_pre_norm(prompt, image_count):
@@ -155,6 +178,75 @@ def test_condition_slots_expand_to_actual_latent_grid():
torch.testing.assert_close(collapsed[slots][0], hidden[2])
class _RecordingBlock(torch.nn.Module):
def __init__(self, layer_id):
super().__init__()
self._layer_id = layer_id
self.seen = None
def forward(self, hidden_states, *args):
caches = [cache[self._layer_id] for cache in args[-1]]
self.seen = caches[0]
return hidden_states
class _UnifiedBlocks(torch.nn.Module):
def __init__(self, blocks):
super().__init__()
self.transformer_blocks = torch.nn.ModuleList(blocks)
def forward(self, hidden_states, *args):
x = hidden_states
for block in self.transformer_blocks:
x = block(x, *args)
return x
def _run_blocks(blocks, prefix_caches):
x = torch.zeros(1)
for block in blocks:
x = block(x, prefix_caches)
return x
def test_cache_dit_wrapper_keeps_per_layer_prefix_kv():
inner = [_RecordingBlock(0), _RecordingBlock(1), _RecordingBlock(2)]
wrapped = torch.nn.ModuleList([_UnifiedBlocks(inner)])
prefix_caches = [[{"layer": 0}, {"layer": 1}, {"layer": 2}]]
_run_blocks(wrapped, prefix_caches)
assert [block.seen for block in inner] == prefix_caches[0]
def test_plain_blocks_still_get_per_layer_prefix_kv():
blocks = torch.nn.ModuleList([_RecordingBlock(0), _RecordingBlock(1)])
prefix_caches = [[{"layer": 0}, {"layer": 1}]]
_run_blocks(blocks, prefix_caches)
assert [block.seen for block in blocks] == prefix_caches[0]
class _FirstSlotBlock(torch.nn.Module):
"""Old loop body: take caches[0] and broadcast it to every layer."""
def __init__(self):
super().__init__()
self.seen = None
def forward(self, hidden_states, *args):
self.seen = args[-1][0]
return hidden_states
def test_first_slot_only_caches_are_shared_across_layers():
inner = [_FirstSlotBlock(), _FirstSlotBlock()]
unified = _UnifiedBlocks(inner)
prefix_caches = [[{"layer": 0}, {"layer": 1}]]
unified(torch.zeros(1), [cache[0] for cache in prefix_caches])
assert [block.seen for block in inner] == [prefix_caches[0][0], prefix_caches[0][0]]
def test_adjacent_image_slots_stay_distinct():
layout = build_layout(
[False, True, True, False], [(1, 2, 2), (1, 4, 2), (1, 2, 2)], (4, 6, 6), "cpu"
python/sglang/multimodal_gen/test/unit/test_qwen_image21_cuda.py
+29 -3
View File
@@ -16,6 +16,7 @@ from sglang.multimodal_gen.configs.models.dits.qwenimage21 import (
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image21 import (
QwenImage21PipelineConfig,
)
from sglang.multimodal_gen.configs.sample.qwenimage21 import QwenImage21SamplingParams
from sglang.multimodal_gen.runtime.breakable_cuda_graph.runner import (
DiffusionBreakableCudaGraphRunner,
)
@@ -33,6 +34,10 @@ from sglang.multimodal_gen.runtime.pipelines.qwen_image21 import QwenImage21Pipe
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
ComposedPipelineBase,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.qwen_image21 import (
QwenImage21DenoisingStage,
)
from sglang.multimodal_gen.runtime.server_args import (
ServerArgs,
set_global_server_args,
@@ -261,14 +266,35 @@ def test_cached_prefix_matches_full_recomputation(model, edit):
def test_graph_replay_uses_new_request_prefix(model, edit, sample_count):
first = batched_inputs([inputs(5 + i, edit) for i in range(sample_count)])
second = batched_inputs([inputs(9 + i, edit) for i in range(sample_count)])
for kwargs in (first, second):
kwargs["encoder_hidden_states_mask"] = torch.ones(
kwargs["encoder_hidden_states"].shape[:2], device="cuda", dtype=torch.bool
)
stage = object.__new__(QwenImage21DenoisingStage)
runner = DiffusionBreakableCudaGraphRunner(model, torch.device("cuda"))
try:
with torch.no_grad(), set_forward_context(None, None):
with (
torch.no_grad(),
set_forward_context(
None,
None,
Req(sampling_params=QwenImage21SamplingParams(), is_warmup=True),
),
):
model(**first)
assert runner.capture(**first)
stage._bcg_run(runner, first, model)
assert len(runner.entries) == 1
with (
torch.no_grad(),
set_forward_context(
None,
None,
Req(sampling_params=QwenImage21SamplingParams()),
),
):
model(**second)
expected = model(**second)
actual = runner(**second)
actual = stage._bcg_run(runner, second, model)
assert len(runner.entries) == 1
torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6)
finally:
python/sglang/multimodal_gen/test/unit/test_runner_perf_baselines.py
+78
View File
@@ -0,0 +1,78 @@
from dataclasses import replace
import pytest
from sglang.multimodal_gen.test.server.test_server_utils import PerformanceValidator
from sglang.multimodal_gen.test.server.testcase_configs import (
BaselineConfig,
PerformanceSummary,
get_perf_baseline_path,
)
@pytest.mark.parametrize(
"runner",
["b200-fin03-4-4567", "b200-cirrascale2", "b200-cirrascale4-0123", "unknown", ""],
)
def test_default_runner_baseline(monkeypatch, runner):
monkeypatch.setenv("RUNNER_NAME", runner)
config = BaselineConfig.load(get_perf_baseline_path("b200"))
assert config.scenarios["flux1_modelopt_nvfp4_t2i"].expected_e2e_ms == 836.71
assert (
config.scenarios["qwen_image_2512_modelopt_nvfp4_t2i"].expected_e2e_ms
== 9650.06
)
@pytest.mark.parametrize(
"runner,flux,qwen",
[
("b200-di01-4567", 1334.16, 16126.87),
("b200-cirrascale1-0123", 1574.32, 17742.04),
("b200-cirrascale3-0123", 1470.24, 17894.49),
("b200-cirrascale3-4567", 1471.19, 17346.65),
],
)
def test_runner_override_preserves_other_metrics(monkeypatch, runner, flux, qwen):
monkeypatch.delenv("RUNNER_NAME", raising=False)
default = BaselineConfig.load(get_perf_baseline_path("b200"))
h100_default = BaselineConfig.load(get_perf_baseline_path("h100"))
monkeypatch.setenv("RUNNER_NAME", runner)
pool = BaselineConfig.load(get_perf_baseline_path("b200"))
expected = {
"flux1_modelopt_nvfp4_t2i": flux,
"qwen_image_2512_modelopt_nvfp4_t2i": qwen,
}
for name, scenario in default.scenarios.items():
assert pool.scenarios[name] == replace(
scenario, expected_e2e_ms=expected.get(name, scenario.expected_e2e_ms)
)
assert pool.tolerances == default.tolerances
assert pool.step_fractions == default.step_fractions
assert BaselineConfig.load(get_perf_baseline_path("h100")) == h100_default
@pytest.mark.parametrize(
"runner",
[
"b200-di01-4567",
"b200-fin03-4-4567",
"b200-cirrascale1-0123",
"b200-cirrascale3-0123",
"b200-cirrascale3-4567",
],
)
def test_runner_baseline_enforces_e2e_boundary(monkeypatch, runner):
monkeypatch.setenv("RUNNER_NAME", runner)
monkeypatch.setenv("SGLANG_GEN_BASELINE", "0")
monkeypatch.delenv("SGLANG_E2E_TOLERANCE", raising=False)
config = BaselineConfig.load(get_perf_baseline_path("b200"))
for name in ("flux1_modelopt_nvfp4_t2i", "qwen_image_2512_modelopt_nvfp4_t2i"):
scenario = config.scenarios[name]
validator = PerformanceValidator(
scenario, config.tolerances, config.step_fractions
)
limit = scenario.expected_e2e_ms * (1 + config.tolerances.e2e)
validator.validate_e2e(PerformanceSummary(limit - 1, 0, 0, {}, [], {}, {}))
with pytest.raises(AssertionError, match="E2E Latency"):
validator.validate_e2e(PerformanceSummary(limit + 1, 0, 0, {}, [], {}, {}))
python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py
+2 -8
View File
@@ -380,18 +380,12 @@ def test_sensenova_u1_npu_fia_checks_operator_availability(monkeypatch, availabl
assert npu_fia_available() is available
@pytest.mark.parametrize(
("is_npu", "uses_native"),
[(False, True), (True, False)],
)
def test_sensenova_u1_shared_rmsnorm_dispatch(monkeypatch, is_npu, uses_native):
monkeypatch.setattr(current_platform, "is_npu", lambda: is_npu)
def test_sensenova_u1_shared_rmsnorm_uses_framework_dispatch():
norm = make_qwen3_rms_norm(64, eps=1e-6)
assert isinstance(norm, RMSNorm)
assert norm.cast_x_before_out_mul
assert (norm._forward_method == norm.forward_native) is uses_native
assert norm._forward_method != norm.forward_native
@torch.no_grad()
python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py
+141 -14
View File
@@ -1,6 +1,7 @@
import json
import os
from dataclasses import replace
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
@@ -20,6 +21,88 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
pytest_plugins = ["pytester"]
@pytest.mark.parametrize("load_time_ms", [None, 0, float("nan"), 1000])
def test_load_guard_is_terminal_without_stage_checks(
harness, monkeypatch, load_time_ms
):
runner, case = harness
case = replace(case, run_perf_check=False)
monkeypatch.setattr(
runner, "run_and_collect", Mock(return_value=(_perf_record(), b"output"))
)
with pytest.raises(test_server_common.PerformanceValidationError, match="Load"):
runner.test_diffusion_generation(
case, SimpleNamespace(load_time_ms=load_time_ms)
)
assert runner.run_and_collect.call_count == 1
@pytest.mark.parametrize("load_time_ms", [None, 0, float("nan")])
def test_baseline_generation_requires_loading_measurement(
harness, monkeypatch, load_time_ms
):
runner, case = harness
monkeypatch.setenv("SGLANG_GEN_BASELINE", "1")
monkeypatch.setattr(
runner, "run_and_collect", Mock(return_value=(_perf_record(), b"output"))
)
with pytest.raises(test_server_common.PerformanceValidationError, match="Load"):
runner.test_diffusion_generation(
case, SimpleNamespace(load_time_ms=load_time_ms)
)
assert test_server_common._PENDING_BASELINE_DUMPS == {}
def test_request_warmup_is_separate_from_guarded_requests(harness, monkeypatch, capsys):
runner, case = harness
case = replace(case, perf_warmup_requests=1)
cold = _perf_record()
cold.total_duration_ms = 3000
generate = Mock(
side_effect=[
(cold, b"output"),
(_perf_record(), b"output"),
(_perf_record(), b"output"),
]
)
monkeypatch.setattr(runner, "run_and_collect", generate)
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert generate.call_count == 3
assert len(runner._perf_results) == 2
assert runner._validate_consistency.call_count == 2
assert "request warmup 1/1 e2e=3000.0000ms" in capsys.readouterr().out
@pytest.mark.parametrize("duration", [None, 0, float("nan"), float("inf")])
def test_request_warmup_requires_e2e(harness, monkeypatch, duration):
runner, case = harness
case = replace(case, perf_warmup_requests=1)
record = _perf_record()
if duration is None:
record = None
else:
record.total_duration_ms = duration
generate = Mock(return_value=(record, b"output"))
monkeypatch.setattr(runner, "run_and_collect", generate)
with pytest.raises(
test_server_common.PerformanceValidationError, match="warmup.*E2E"
):
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert generate.call_count == 1
def test_request_after_warmup_still_enforces_e2e(harness, monkeypatch):
runner, case = harness
case = replace(case, perf_warmup_requests=1, run_perf_check=False)
slow = _perf_record()
slow.total_duration_ms = 3000
generate = Mock(side_effect=[(_perf_record(), b"output"), (slow, b"output")])
monkeypatch.setattr(runner, "run_and_collect", generate)
with pytest.raises(test_server_common.PerformanceValidationError, match="E2E"):
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert generate.call_count == 2
def _perf_record():
return RequestPerfRecord(
request_id="request",
@@ -48,6 +131,7 @@ def harness(monkeypatch):
expected_avg_denoise_ms=5,
expected_median_denoise_ms=5,
estimated_full_test_time_s=1,
expected_load_ms=100,
load_peak_vram_mb=1000,
runtime_peak_vram_mb=2000,
)
@@ -109,18 +193,30 @@ def test_each_request_failure_fails_case(harness, monkeypatch, bad_request, fail
outputs[bad_request] = RuntimeError("server request failed")
generate = Mock(side_effect=outputs)
monkeypatch.setattr(runner, "run_and_collect", generate)
ctx = object()
ctx = SimpleNamespace(load_time_ms=100)
with pytest.raises(pytest.fail.Exception, match=f"request {bad_request + 1}/2"):
terminal = failure in {"performance", "load_peak", "runtime_peak", "missing_memory"}
error_type = (
test_server_common.PerformanceValidationError
if terminal
else pytest.fail.Exception
)
with pytest.raises(error_type, match=f"request {bad_request + 1}/2"):
runner.test_diffusion_generation(case, ctx)
assert generate.call_count == 2
expected_requests = bad_request + 1 if terminal else 2
assert generate.call_count == expected_requests
assert all(call.args[0] is ctx for call in generate.call_args_list)
assert runner._validate_consistency.call_count == (
1 if failure == "generation" else 2
expected_consistency = (
bad_request if terminal else (1 if failure == "generation" else 2)
)
assert runner._validate_consistency.call_count == expected_consistency
# Even failed performance measurements must survive in the report.
expected = [i + 1 for i in range(2) if failure != "generation" or i != bad_request]
expected = [
i + 1
for i in range(expected_requests)
if failure != "generation" or i != bad_request
]
assert [r["request_index"] for r in runner._perf_results] == expected
@@ -131,7 +227,7 @@ def test_both_requests_pass(harness, monkeypatch):
"run_and_collect",
Mock(side_effect=[(_perf_record(), b"first"), (_perf_record(), b"second")]),
)
runner.test_diffusion_generation(case, object())
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert runner._validate_consistency.call_count == 2
assert [call.args[1] for call in runner._validate_consistency.call_args_list] == [
b"first",
@@ -140,6 +236,37 @@ def test_both_requests_pass(harness, monkeypatch):
assert [r["request_index"] for r in runner._perf_results] == [1, 2]
@pytest.mark.parametrize("run_perf_check", [False, True])
@pytest.mark.parametrize("e2e_ms", [None, 0, -1, float("nan"), float("inf")])
def test_e2e_is_required_even_without_threshold_checks(
harness, monkeypatch, run_perf_check, e2e_ms
):
runner, case = harness
case = replace(case, run_perf_check=run_perf_check)
record = _perf_record()
record.total_duration_ms = e2e_ms
generate = Mock(return_value=(record, b"output"))
monkeypatch.setattr(runner, "run_and_collect", generate)
with pytest.raises(
test_server_common.PerformanceValidationError,
match="E2E duration missing or invalid",
):
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert generate.call_count == 1
assert not runner._perf_results
def test_disabled_threshold_checks_still_record_e2e(harness, monkeypatch):
runner, case = harness
case = replace(case, run_perf_check=False)
monkeypatch.setattr(
runner, "run_and_collect", Mock(return_value=(_perf_record(), b"output"))
)
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert [r["e2e_ms"] for r in runner._perf_results] == [100, 100]
def test_request_artifacts_do_not_overwrite_each_other(harness, monkeypatch, tmp_path):
runner, case = harness
monkeypatch.setenv("SGLANG_DIFFUSION_ARTIFACT_DIR", str(tmp_path))
@@ -150,7 +277,7 @@ def test_request_artifacts_do_not_overwrite_each_other(harness, monkeypatch, tmp
return _perf_record(), b"output"
monkeypatch.setattr(runner, "run_and_collect", generate)
runner.test_diffusion_generation(case, object())
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert artifact_dirs == [str(tmp_path / f"request-{i}") for i in (1, 2)]
assert os.environ["SGLANG_DIFFUSION_ARTIFACT_DIR"] == str(tmp_path)
@@ -168,7 +295,7 @@ def test_later_skip_cannot_hide_earlier_failure(harness, monkeypatch):
),
)
with pytest.raises(pytest.fail.Exception, match="failed first"):
runner.test_diffusion_generation(case, object())
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
def test_second_request_cannot_be_skipped_after_first_passes(harness, monkeypatch):
@@ -179,7 +306,7 @@ def test_second_request_cannot_be_skipped_after_first_passes(harness, monkeypatc
Mock(side_effect=[(_perf_record(), b"output"), pytest.skip.Exception("skip")]),
)
with pytest.raises(pytest.fail.Exception, match="Required request skipped"):
runner.test_diffusion_generation(case, object())
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
def test_empty_content_is_not_a_consistency_pass(harness):
@@ -206,7 +333,7 @@ def test_audio_checked_even_when_video_consistency_fails(
runner, "run_and_collect", Mock(return_value=(_perf_record(), b"output"))
)
with pytest.raises(pytest.fail.Exception, match="audio consistency.*wrong audio"):
runner.test_diffusion_generation(case, object())
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert runner._validate_consistency.call_count == 2
assert runner._validate_audio_consistency.call_count == 2
@@ -252,7 +379,7 @@ def test_perf_fixture_retains_failed_case_results(pytester, monkeypatch):
DiffusionServerArgs("test", modality="image"),
DiffusionSamplingParams(prompt="test"),
)
summary = PerformanceSummary(100, 5, 5, {}, [], {}, {})
summary = PerformanceSummary(100, 5, 5, {}, [], {}, {}, load_time_ms=100)
for index in (1, 2):
self._record_performance_result(case, summary, index)
if case_id == "failed":
@@ -281,7 +408,7 @@ def test_gt_generation_runs_both_requests(harness, monkeypatch):
monkeypatch.setattr(runner, "run_and_collect", generate)
save = Mock()
monkeypatch.setattr(runner, "_save_gt_output", save)
runner.test_diffusion_generation(case, object())
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
assert [(call.args[0].id, call.args[1]) for call in save.call_args_list] == [
("first", b"first"),
("first", b"second"),
@@ -303,7 +430,7 @@ def test_baseline_generation_keeps_worst_of_both_requests(harness, monkeypatch):
monkeypatch.setattr(
runner, "run_and_collect", Mock(side_effect=[(r, b"output") for r in records])
)
runner.test_diffusion_generation(case, object())
runner.test_diffusion_generation(case, SimpleNamespace(load_time_ms=100))
summaries = test_server_common._PENDING_BASELINE_DUMPS[case.id]
assert len(summaries) == 2
log = Mock()
python/sglang/multimodal_gen/test/unit/test_suite_partitioning.py
+5
View File
@@ -140,3 +140,8 @@ def test_qwen_quality_variants_use_the_same_generation_request():
assert extra_high.prompt == lossless.prompt
assert extra_high.output_size == lossless.output_size
assert extra_high.extras == {"quality": "extra-high"}
scenarios = json.loads(_H100_BASELINE_PATH.read_text())["scenarios"]
for case_id in ("qwen_image_t2i_2_gpus", "qwen_image_t2i_2_gpus_extra_high"):
assert cases[case_id].run_perf_check
assert scenarios[case_id]["expected_e2e_ms"] > 0
python/sglang/multimodal_gen/test/unit/test_timestep_preparation_logging.py
+65
View File
@@ -0,0 +1,65 @@
import logging
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from sglang.multimodal_gen.configs.pipeline_configs.flux import (
Flux2KleinBasePipelineConfig,
)
from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
FlowMatchEulerDiscreteScheduler,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages import (
base,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages import (
timestep_preparation as module,
)
@pytest.mark.parametrize("warmup", [False, True])
@pytest.mark.parametrize("debug", [False, True])
def test_timestep_logging_preserves_scheduler_and_skips_unused_copy(warmup, debug):
args = SimpleNamespace(pipeline_config=Flux2KleinBasePipelineConfig())
scheduler = FlowMatchEulerDiscreteScheduler()
with patch.object(base, "get_global_server_args", return_value=args):
stage = module.TimestepPreparationStage(scheduler)
batch = Req(sampling_params=SamplingParams(num_inference_steps=4))
batch.is_warmup = warmup
records = []
class Capture(logging.Handler):
def emit(self, record):
records.append(record)
self.format(record)
test_logger = logging.Logger(
"timestep-test", logging.DEBUG if debug else logging.INFO
)
test_logger.addHandler(Capture())
with (
patch.object(module, "logger", test_logger),
patch.object(
module, "get_local_torch_device", return_value=torch.device("cpu")
),
):
if debug and not warmup:
result = stage.forward(batch, args)
else:
with patch.object(
torch.Tensor, "detach", side_effect=AssertionError("unused log copy")
):
result = stage.forward(batch, args)
assert result is batch
assert batch.scheduler is scheduler
assert batch.timesteps is scheduler.timesteps
assert len(records) == int(debug and not warmup)
if records:
value = records[0].args[-1]
assert value.device.type == "cpu"
assert torch.equal(value, batch.timesteps)
assert "TimestepPreparationStage" in records[0].getMessage()
python/sglang/multimodal_gen/test/unit/test_video_url_request_extras.py
+36
View File
@@ -0,0 +1,36 @@
from unittest.mock import Mock
import pytest
from sglang.multimodal_gen.test.server.test_server_utils import get_generate_fn
from sglang.multimodal_gen.test.server.testcase_configs import DiffusionSamplingParams
def test_url_video_request_preserves_sampling_extras():
extras = {
"profile": True,
"num_profiled_timesteps": 5,
"num_inference_steps": 12,
"seed": 0,
}
original_extras = extras.copy()
params = DiffusionSamplingParams(
prompt="test",
image_path="https://example.com/input.png",
direct_url_test=True,
fps=24,
num_frames=25,
extras=extras,
)
client = Mock()
client.videos.create.side_effect = ConnectionError("stop at transport boundary")
generate = get_generate_fn("test-model", "video", params)
with pytest.raises(ConnectionError, match="stop at transport boundary"):
generate("url-video", client)
assert client.videos.create.call_args.kwargs["extra_body"] == {
"reference_url": params.image_path,
"fps": 24,
"num_frames": 25,
**original_extras,
}
assert params.extras == original_extras
python/sglang/srt/arg_groups/arg_utils.py
+11 -19
View File
@@ -92,31 +92,23 @@ class Arg(msgspec.Struct, frozen=True):
fallback: Any = None
_NO_DEFAULT = object()
class Derived(msgspec.Struct, frozen=True):
"""Metadata for a field the configuration implies, not one anyone types.
"""Metadata for namespace fields that are not CLI inputs or record fields.
The other half of a namespace. An ``Arg`` field is the operator's input and
is collected into ``ServerArgs``; a ``Derived`` field carries no annotation,
so it is not a dataclass field and never reaches the record -- which is
right, because it has no input to preserve and the record is what crosses a
process boundary.
``fn`` names what computes it, as a dotted path resolved lazily so that a
declaration module stays free of runtime imports. Such a field is a pure
function of the published configuration, so it is computed once at
``publish`` and stored as an ordinary bag leaf -- a plain attribute load,
which is what a read inside compiled model code needs.
Every declaration carries ``fn`` today, the parallel quotients included:
they are a function of the configured leaves, so they are computed at
publish like the rest. What is special about them is not how they are
computed but that a stamp can move one afterwards -- an elastic scale-up
restamps ``attn_dp_size`` -- which ``ParallelContext`` answers above the
published leaf.
``fn`` is a lazily resolved dotted function path. It computes a value from
resolved configuration once at publication. Fields without ``fn``, such as
ranks and group handles, are set at runtime. Parallel overrides take
precedence over published values.
"""
doc: str = ""
fn: str = ""
# Default for runtime-only fields, e.g. ``gpu_id=None`` without a device.
# Fields without a default raise if read before initialization.
default: Any = _NO_DEFAULT
class NS(msgspec.Struct, frozen=True):
python/sglang/srt/arg_groups/field_order.py
+1
View File
@@ -205,6 +205,7 @@ POSITIONAL_FIELD_ORDER = (
"stat_loggers",
"constrained_json_whitespace_pattern",
"constrained_json_disable_any_whitespace",
"constrained_json_max_whitespace_cnt",
"attention_backend",
"decode_attention_backend",
"enable_lean_attention",
python/sglang/srt/arg_groups/fields/device.py
+12 -1
View File
@@ -17,7 +17,7 @@ from typing import (
import msgspec
from sglang.srt.arg_groups.arg_utils import A
from sglang.srt.arg_groups.arg_utils import A, Derived
class Device(msgspec.Struct):
@@ -40,6 +40,17 @@ class Device(msgspec.Struct):
int,
"The delta between consecutive GPU IDs that are used. For example, setting it to 2 will use GPU 0,2,4,...",
] = 1
gpu_id = Derived(
doc=(
"Which device this process runs on. Nobody types it and nothing "
"computes it from the configuration: the parent decides -- "
"reindexing narrows the visible devices before the spawn, and Ray "
"allocates from its own pool -- so the entry states it in the "
"bundle it hands `publish`. `None` is an answer rather than an "
"absence: most roles run on no device at all."
),
default=None,
)
random_seed: A[Optional[int], "The random seed."] = None
mlx_enable_sampling: A[
bool,
python/sglang/srt/arg_groups/fields/memory.py
+2 -1
View File
@@ -153,7 +153,7 @@ class Memory(msgspec.Struct):
hicache_storage_backend: A[
Optional[str],
Arg(
help="The storage backend for hierarchical KV cache. Built-in backends: file, mooncake, npu_memcache, hf3fs, nixl, aibrix. For dynamic backend, use --hicache-storage-backend-extra-config to specify: backend_name (custom name), module_path (Python module path), class_name (backend class name).",
help="The storage backend for hierarchical KV cache. Built-in backends: file, mooncake, npu_memcache, hf3fs, nixl, aibrix, tensorcast. For dynamic backend, use --hicache-storage-backend-extra-config to specify: backend_name (custom name), module_path (Python module path), class_name (backend class name).",
choices=[
"file",
"sim",
@@ -167,6 +167,7 @@ class Memory(msgspec.Struct):
"simm",
"mori",
"shm",
"tensorcast",
],
),
] = None
python/sglang/srt/arg_groups/fields/parallel.py
+65 -12
View File
@@ -280,18 +280,7 @@ class Parallel(msgspec.Struct):
"Maximum EP size the server can scale to at runtime. Pre-allocates active-rank state and backend buffers to this size. Defaults to the launch-time world size.",
] = None
# ---- derived: the quotients of the leaves above -------------------------
#
# Declared here, beside what they are computed from, because a namespace is
# one file and one class. They are not annotated, so they are not dataclass
# fields and `collect_input_fields` does not put them on the record -- which
# is right: a quotient has no operator input to preserve, and the record is
# what crosses a process boundary, so a width put there would be a stale
# copy the moment an elastic scale-up restamps one. Every input is a leaf
# above, so all six are fixed once the configuration is: `publish` computes
# them through `parallel_widths_of` and stores them as ordinary bag leaves,
# and `ParallelContext` answers with the stamp when a scale-up has moved
# one.
# Derived fields are computed at publication and are not stored in ServerArgs.
attn_tp_size = Derived(
fn="sglang.srt.runtime_context.attn_tp_size_of",
doc="Attention tensor-parallel width: `tp_size` divided by the "
@@ -320,3 +309,67 @@ class Parallel(msgspec.Struct):
doc="Whether decode context parallelism is in play: `dcp_size` is "
"wider than one rank, which is exactly when the group gets built.",
)
# Runtime fields: publish sets ranks; distributed initialization sets groups.
tp_rank = Derived(doc="This process's place in the tensor-parallel group.")
pp_rank = Derived(doc="This process's place in the pipeline group.")
moe_ep_rank = Derived(doc="This process's place in the expert-parallel group.")
moe_dp_rank = Derived(doc=("This process's place in the MoE data-parallel group."))
moe_tp_rank = Derived(
doc=("This process's place in the MoE tensor-parallel group.")
)
attn_tp_rank = Derived(
doc=("This process's place in the attention tensor-parallel group.")
)
attn_cp_rank = Derived(
doc=("This process's place in the attention context-parallel group.")
)
dcp_rank = Derived(
doc=("This process's place in the decode context-parallel group.")
)
attn_dcp_rank = Derived(
doc=(
"Decode context-parallel rank inside the attention TP group, "
"zero where decode context parallelism is off."
)
)
attn_dp_rank = Derived(
doc=(
"This process's index in the attention-DP group, computed from "
"`tp_rank` when `initialize_dp_attention` runs."
)
)
dp_rank = Derived(
doc=(
"Which data-parallel replica this process serves, as the data "
"parallel controller numbered them at spawn. `None` when there "
"is no controller: unlike the other ranks it is a position in "
"no group, which is why the spawn states it."
)
)
launch_world_rank = Derived(
doc=(
"This process's rank in the WORLD group as built. A scale-up "
"does not renumber it."
)
)
launch_world_size = Derived(
fn="sglang.srt.runtime_context.launch_world_size_of",
doc="Width the WORLD group was built at -- what a scale-up leaves "
"behind rather than updates.",
)
max_world_size = Derived(
fn="sglang.srt.runtime_context.max_world_size_of",
doc="Ranks the WORLD group has room for: `--max-ep-size` when set, "
"otherwise the launch width.",
)
world_group = Derived(doc="The WORLD group.")
tp_group = Derived(doc="The tensor-parallel group.")
pp_group = Derived(doc="The pipeline group.")
moe_ep_group = Derived(doc="The expert-parallel group.")
moe_dp_group = Derived(doc="The MoE data-parallel group.")
moe_tp_group = Derived(doc="The MoE tensor-parallel group.")
attn_tp_group = Derived(doc="The attention tensor-parallel group.")
attn_cp_group = Derived(doc="The attention context-parallel group.")
shared_experts_tp_group = Derived(doc=("The shared-expert tensor-parallel group."))
dcp_group = Derived(doc="The decode context-parallel group.")
python/sglang/srt/arg_groups/fields/serving.py
+4
View File
@@ -281,6 +281,10 @@ class Serving(msgspec.Struct):
bool,
"(xgrammar and llguidance backends only) Enforce compact representation in JSON constrained output.",
] = False
constrained_json_max_whitespace_cnt: A[
Optional[int],
"(xgrammar backend only) Max consecutive whitespace chars allowed in JSON constrained output. None means unbounded.",
] = None
# -------------------------------------------------------------------------
# Dynamic batch tokenizer
python/sglang/srt/arg_groups/fields/spec.py
+5 -1
View File
@@ -160,7 +160,11 @@ class Spec(msgspec.Struct):
] = None
speculative_draft_window_size: A[
Optional[int],
"Sliding window size for the draft model. Honored by Llama EAGLE-3 (`LlamaForCausalLMEagle3`) and DFLASH only; other EAGLE-3 backends (e.g. MLA-based drafters) silently ignore it. For Llama EAGLE-3, the drafter only attends to the most recent N keys (verifier hidden states + its own outputs); the verifier is unaffected. For DFLASH, the draft worker keeps a recent target-token window in its local KV cache (paged backends may retain up to one extra page on the left for alignment). Default is full attention/context.",
"Sliding window size for the draft model. Honored by Llama EAGLE-3 (`LlamaForCausalLMEagle3`), DFLASH, and the built-in EAGLE/MTP draft-decode path on the Triton and FlashInfer draft backends; other EAGLE-3 backends (e.g. MLA-based drafters) silently ignore it. For Llama EAGLE-3, the drafter only attends to the most recent N keys (verifier hidden states + its own outputs); the verifier is unaffected. For DFLASH, the draft worker keeps a recent target-token window in its local KV cache (paged backends may retain up to one extra page on the left for alignment). For the built-in EAGLE/MTP draft, each draft-decode step attends to a --speculative-draft-sink-size sink plus the most recent N tokens, leaving the target verify pass unchanged; it is ignored (with a warning) if the draft model has a native sliding window of its own. Default is full attention/context.",
] = None
speculative_draft_sink_size: A[
Optional[int],
"Number of leading 'attention sink' tokens the draft always attends to, in addition to the --speculative-draft-window-size recent window (StreamingLLM-style). Honored only by the built-in EAGLE/MTP draft-decode path on the Triton and FlashInfer draft backends; the Llama EAGLE-3 and DFLASH windows ignore it. 0/unset => pure recent window. Requires --speculative-draft-window-size.",
] = None
speculative_moe_runner_backend: A[
Optional[str],
python/sglang/srt/arg_groups/hisparse_hook.py
+8 -2
View File
@@ -86,14 +86,16 @@ def validate_hisparse(server_args: ServerArgs) -> None:
from sglang.srt.configs.model_config import (
is_deepseek_dsa,
is_deepseek_v4,
is_minimax_sparse,
)
hf_config = model_config_of(server_args).hf_config
is_v4_hisparse = is_deepseek_v4(hf_config)
is_m3_hisparse = is_minimax_sparse(hf_config)
is_hip = get_platform().is_hip
assert is_deepseek_dsa(hf_config) or is_v4_hisparse, (
assert is_deepseek_dsa(hf_config) or is_v4_hisparse or is_m3_hisparse, (
"--enable-hisparse is only supported for DSA (DeepSeek Sparse Attention) "
"models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. "
"models (e.g., DeepSeek V3.2, GLM-5), DeepSeek V4, and MiniMax M3 now. "
)
assert cfg.disable_radix_cache, (
@@ -121,6 +123,10 @@ def validate_hisparse(server_args: ServerArgs) -> None:
)
return
# MiniMax M3 uses its own Triton sparse kernels.
if is_m3_hisparse:
return
if resolved_view(server_args).kv_cache_dtype not in (
"bfloat16",
"auto",
python/sglang/srt/arg_groups/kv_cache_hook.py
+4 -6
View File
@@ -507,12 +507,10 @@ def handle_unified_memory_pool(server_args: Any) -> None:
"write loc, so a captured decode replay raises. "
"TODO(ch-wan): carry out_cache_loc_virtual into the child view."
)
assert not (cfg.enable_hierarchical_cache or cfg.enable_lmcache), (
"--enable-unified-memory is not yet compatible with hierarchical / "
"host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): "
"the unified-memory-pool init wires up no host pools, and its device mamba / "
"full-attention slots are VIRTUAL — the host-offload path does not "
"translate them to physical."
assert not cfg.enable_lmcache, (
"--enable-unified-memory is not yet compatible with --enable-lmcache: "
"the LMCache offload path indexes the device buffers with the ids it "
"is handed, and under the unified pool those are VIRTUAL."
)
if cfg.dcp_size > 1:
_validate_unified_memory_dcp(server_args)
python/sglang/srt/arg_groups/model_overrides/__init__.py
+1
View File
@@ -16,6 +16,7 @@ from sglang.srt.arg_groups.model_overrides import exaone # noqa: F401
from sglang.srt.arg_groups.model_overrides import falcon_h1 # noqa: F401
from sglang.srt.arg_groups.model_overrides import gemma2_gemma3 # noqa: F401
from sglang.srt.arg_groups.model_overrides import gemma4 # noqa: F401
from sglang.srt.arg_groups.model_overrides import gigachat35 # noqa: F401
from sglang.srt.arg_groups.model_overrides import glm4_moe # noqa: F401
from sglang.srt.arg_groups.model_overrides import gpt_oss # noqa: F401
from sglang.srt.arg_groups.model_overrides import granitemoehybrid # noqa: F401
python/sglang/srt/arg_groups/model_overrides/gigachat35.py
+26
View File
@@ -0,0 +1,26 @@
"""Config-time override declarations for gigachat35.
Architectures: GigaChat35ForCausalLM, GigaChat35ForCausalLMNextN.
"""
import logging
from typing import Any, Dict
from sglang.srt.arg_groups.model_override_base import (
_register_for,
resolving_view,
)
logger = logging.getLogger(__name__)
@_register_for("GigaChat35ForCausalLM", "GigaChat35ForCausalLMNextN")
def _gigachat35_overrides(server_args: Any, hf_config: Any) -> dict:
cfg = resolving_view(server_args)
overrides: Dict[str, Any] = {"disable_shared_experts_fusion": True}
if cfg.speculative_algorithm == "EAGLE":
logger.info(
"Enable multi-layer EAGLE speculative decoding for GigaChat 3.5 model."
)
overrides["enable_multi_layer_eagle"] = True
return overrides
python/sglang/srt/arg_groups/speculative_hook.py
+32 -4
View File
@@ -16,6 +16,7 @@ from sglang.srt.arg_groups.overrides import (
resolving_view,
run_post_process_pass,
)
from sglang.srt.environ import envs
from sglang.srt.platforms import current_platform
from sglang.srt.runtime_context import get_platform
@@ -35,6 +36,7 @@ def _should_auto_enable_hip_rejection_sampling(
accept_threshold_single: float,
accept_threshold_acc: float,
enable_deterministic_inference: bool,
simulate_acc_len: float = -1.0,
) -> bool:
"""Whether HIP may default ``speculative_use_rejection_sampling`` on.
@@ -43,6 +45,10 @@ def _should_auto_enable_hip_rejection_sampling(
would crash configs that previously ran greedy on HIP, including EAGLE3
stage-a ``test_basic_sanity_eagle3`` (draft 32000 vs target 128256). Skip
EAGLE3 and any EAGLE run that already has a token map.
Also skip when ``SGLANG_SIMULATE_ACC_LEN`` is on: AgentX throughput still
runs the real EAGLE verify then overwrites accept length, so the Triton
chain sampler is paid for and thrown away.
"""
return (
is_hip
@@ -53,6 +59,7 @@ def _should_auto_enable_hip_rejection_sampling(
and accept_threshold_single == 1.0
and accept_threshold_acc == 1.0
and not enable_deterministic_inference
and simulate_acc_len <= 0
)
@@ -161,8 +168,11 @@ def handle_speculative_decoding(server_args: ServerArgs) -> None:
),
)
# Validate --speculative-draft-window-size once, regardless of algorithm.
# Consumed by DFLASH (compact draft KV cache) and Llama EAGLE-3 (drafter attention SWA).
# Validate --speculative-draft-window-size / --speculative-draft-sink-size once,
# regardless of algorithm. Consumed by DFLASH (compact draft KV cache), Llama
# EAGLE-3 (drafter attention SWA), and the built-in MTP/NEXTN + EAGLE draft-decode
# path on the Triton and FlashInfer draft attention backends (StreamingLLM sink +
# recent window).
if cfg.speculative_draft_window_size is not None:
window_size = int(cfg.speculative_draft_window_size)
if window_size <= 0:
@@ -174,13 +184,30 @@ def handle_speculative_decoding(server_args: ServerArgs) -> None:
"handle_speculative_decoding",
speculative_draft_window_size=window_size,
)
if cfg.speculative_algorithm not in ("EAGLE3", "DFLASH"):
if cfg.speculative_algorithm not in ("EAGLE", "EAGLE3", "DFLASH"):
logger.warning(
"--speculative-draft-window-size has no effect with "
"speculative_algorithm=%s (honored by Llama EAGLE-3 and DFLASH only).",
"speculative_algorithm=%s (honored by DFLASH, Llama EAGLE-3, and the "
"EAGLE/MTP/NEXTN draft-decode path on the Triton/FlashInfer draft backends).",
cfg.speculative_algorithm,
)
if cfg.speculative_draft_sink_size is not None:
sink_size = int(cfg.speculative_draft_sink_size)
if sink_size < 0:
raise ValueError(
f"--speculative-draft-sink-size must be non-negative, got {sink_size}."
)
if cfg.speculative_draft_window_size is None:
raise ValueError(
"--speculative-draft-sink-size requires --speculative-draft-window-size."
)
declare_resolution(
server_args,
"handle_speculative_decoding",
speculative_draft_sink_size=sink_size,
)
algo = None
if cfg.speculative_algorithm is not None:
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -1005,6 +1032,7 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
accept_threshold_single=cfg.speculative_accept_threshold_single,
accept_threshold_acc=cfg.speculative_accept_threshold_acc,
enable_deterministic_inference=cfg.enable_deterministic_inference,
simulate_acc_len=float(envs.SGLANG_SIMULATE_ACC_LEN.get()),
):
declare_resolution(
server_args,
python/sglang/srt/configs/__init__.py
+2
View File
@@ -17,6 +17,7 @@ from sglang.srt.configs.dots_ocr import DotsOCRConfig
from sglang.srt.configs.dots_vlm import DotsVLMConfig
from sglang.srt.configs.exaone import ExaoneConfig
from sglang.srt.configs.falcon_h1 import FalconH1Config
from sglang.srt.configs.gigachat35 import GigaChat35Config
from sglang.srt.configs.glm5_next import Glm5NextConfig, Glm5NextTextConfig
from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
from sglang.srt.configs.hy_v4 import HYV4Config
@@ -132,6 +133,7 @@ __all__ = [
"Dots3Config",
"FalconH1Config",
"FalconMambaConfig",
"GigaChat35Config",
"GraniteMoeHybridConfig",
"HYV4Config",
"MambaConfig",
python/sglang/srt/configs/gigachat35.py
+252
View File
@@ -0,0 +1,252 @@
from __future__ import annotations
from typing import Optional
import torch
from transformers.configuration_utils import PretrainedConfig
from sglang.srt.configs.mamba_utils import (
Mamba2CacheParams,
Mamba2StateShape,
mamba2_state_dtype,
)
FULL_ATTENTION = "attention"
LINEAR_ATTENTION = "linear_attention"
_FULL_ATTENTION_ALIASES = {FULL_ATTENTION, "full_attention"}
_REQUIRED_LINEAR_ATTRS = (
"linear_conv_kernel_dim",
"linear_key_head_dim",
"linear_value_head_dim",
"linear_num_key_heads",
"linear_num_value_heads",
)
class GigaChat35Config(PretrainedConfig):
model_type = "gigachat3_5"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size: int = 128256,
hidden_size: int = 1280,
intermediate_size: int = 896,
moe_intermediate_size: int = 896,
num_hidden_layers: int = 12,
num_attention_heads: int = 16,
num_key_value_heads: int = 16,
hidden_act: str = "silu",
max_position_embeddings: int = 4096,
initializer_range: float = 0.006,
rms_norm_eps: float = 1e-6,
use_cache: bool = True,
rope_theta: float = 100000.0,
rope_scaling: Optional[dict] = None,
rope_interleave: bool = True,
attention_bias: bool = False,
tie_word_embeddings: bool = False,
q_lora_rank: Optional[int] = 1536,
kv_lora_rank: int = 512,
qk_nope_head_dim: int = 128,
qk_rope_head_dim: int = 64,
v_head_dim: int = 128,
head_dim: int = 64,
n_routed_experts: int = 64,
n_shared_experts: int = 2,
num_experts_per_tok: int = 6,
moe_layer_freq: int = 1,
first_k_dense_replace: int = 1,
routed_scaling_factor: float = 2.5,
n_group: int = 1,
topk_group: int = 1,
topk_method: str = "noaux_tc",
scoring_func: str = "sigmoid",
norm_topk_prob: bool = True,
use_shared_expert_sigmoid: bool = False,
linear_attention_type: str = "Qwen3NextGatedDeltaNet",
layer_types: Optional[list[str]] = None,
full_attention_layers: Optional[list[int]] = None,
linear_conv_kernel_dim: int = 4,
linear_key_head_dim: int = 128,
linear_value_head_dim: int = 128,
linear_num_key_heads: int = 8,
linear_num_value_heads: int = 16,
linear_sigmoid_gate_scale: float = 2.0,
output_gate_type: str = "sigmoid",
norm_type: str = "ZeroCenteredGatedNorm",
layernorm_type: str = "pre_post",
layernorm_gating_weight: float = 2.0,
gated_attention: bool = True,
use_mla_scaling_factor: bool = True,
swiglu_limit: Optional[float] = None,
num_nextn_predict_layers: int = 0,
nextn_is_sparse: bool = False,
**kwargs,
):
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.moe_intermediate_size = moe_intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.max_position_embeddings = max_position_embeddings
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
self.rope_theta = rope_theta
self.rope_scaling = rope_scaling
self.rope_interleave = rope_interleave
self.attention_bias = attention_bias
self.q_lora_rank = q_lora_rank
self.kv_lora_rank = kv_lora_rank
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
self.v_head_dim = v_head_dim
self.head_dim = head_dim
self.n_routed_experts = n_routed_experts
self.n_shared_experts = n_shared_experts
self.num_experts_per_tok = num_experts_per_tok
self.moe_layer_freq = moe_layer_freq
self.first_k_dense_replace = first_k_dense_replace
self.routed_scaling_factor = routed_scaling_factor
self.n_group = n_group
self.topk_group = topk_group
self.topk_method = topk_method
self.scoring_func = scoring_func
self.norm_topk_prob = norm_topk_prob
self.use_shared_expert_sigmoid = use_shared_expert_sigmoid
self.linear_attention_type = linear_attention_type
self.layer_types = layer_types
self.full_attention_layers = full_attention_layers
self.linear_conv_kernel_dim = linear_conv_kernel_dim
self.linear_key_head_dim = linear_key_head_dim
self.linear_value_head_dim = linear_value_head_dim
self.linear_num_key_heads = linear_num_key_heads
self.linear_num_value_heads = linear_num_value_heads
self.linear_sigmoid_gate_scale = linear_sigmoid_gate_scale
self.output_gate_type = output_gate_type
self.linear_num_key_heads_cpu = linear_num_key_heads
self.linear_num_value_heads_cpu = linear_num_value_heads
self.norm_type = norm_type
self.layernorm_type = layernorm_type
self.layernorm_gating_weight = layernorm_gating_weight
self.gated_attention = gated_attention
self.use_mla_scaling_factor = use_mla_scaling_factor
self.swiglu_limit = (
swiglu_limit if (swiglu_limit is not None and swiglu_limit > 0) else None
)
self.num_nextn_predict_layers = num_nextn_predict_layers
self.nextn_is_sparse = nextn_is_sparse
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
_dtype = getattr(self, "torch_dtype", None) or getattr(self, "dtype", None)
if isinstance(_dtype, str):
_dtype = getattr(torch, _dtype, torch.bfloat16)
self.torch_dtype = _dtype or torch.bfloat16
def _resolve_layer_types(self) -> list[str]:
"""Return a normalized per-layer list of FULL_ATTENTION / LINEAR_ATTENTION.
Resolution order: explicit ``layer_types`` -> ``full_attention_layers``.
Defaults to all-linear if nothing is specified (degenerate, but
well-defined).
"""
n = self.num_hidden_layers
if self.layer_types is not None:
if len(self.layer_types) != n:
raise ValueError(
f"layer_types must have length num_hidden_layers ({n}), "
f"got {len(self.layer_types)}."
)
resolved = []
for idx, lt in enumerate(self.layer_types):
if lt == LINEAR_ATTENTION:
resolved.append(LINEAR_ATTENTION)
elif lt in _FULL_ATTENTION_ALIASES:
resolved.append(FULL_ATTENTION)
else:
raise ValueError(f"Unsupported layer type {lt!r} at index {idx}.")
return resolved
resolved = [LINEAR_ATTENTION] * n
if self.full_attention_layers is not None:
for lid in self.full_attention_layers:
resolved[lid] = FULL_ATTENTION
return resolved
@property
def layers_block_type(self) -> list[str]:
return self._resolve_layer_types()
@property
def linear_layer_ids(self) -> list[int]:
return [
i for i, lt in enumerate(self.layers_block_type) if lt == LINEAR_ATTENTION
]
@property
def full_attention_layer_ids(self) -> list[int]:
return [
i for i, lt in enumerate(self.layers_block_type) if lt == FULL_ATTENTION
]
def is_linear_attention_layer(self, layer_id: int) -> bool:
if layer_id >= self.num_hidden_layers:
return False
return self.layers_block_type[layer_id] == LINEAR_ATTENTION
@property
def mamba2_cache_params(self) -> Mamba2CacheParams:
from sglang.srt.runtime_context import get_parallel
missing = [a for a in _REQUIRED_LINEAR_ATTRS if getattr(self, a, None) is None]
if missing:
raise ValueError(
"GigaChat35 hybrid GDN config is missing required linear-attention "
"fields: " + ", ".join(missing)
)
shape = Mamba2StateShape.create(
tp_world_size=get_parallel().attn_tp_size,
intermediate_size=self.linear_value_head_dim * self.linear_num_value_heads,
n_groups=self.linear_num_key_heads,
num_heads=self.linear_num_value_heads,
head_dim=self.linear_value_head_dim,
state_size=self.linear_key_head_dim,
conv_kernel=self.linear_conv_kernel_dim,
)
return Mamba2CacheParams(
shape=shape,
layers=self.linear_layer_ids,
dtype=mamba2_state_dtype(self),
)
from sglang.srt.configs.linear_attn_model_registry import ( # noqa: E402
LinearAttnModelSpec,
register_linear_attn_model,
)
register_linear_attn_model(
LinearAttnModelSpec(
config_class=GigaChat35Config,
backend_class_name="sglang.srt.layers.attention.linear.gdn_backend.GDNAttnBackend",
arch_names=[
"GigaChat35ForCausalLM",
"GigaChat35ForCausalLMNextN",
],
uses_mamba_radix_cache=True,
support_mamba_cache=True,
)
)
python/sglang/srt/configs/hybrid_arch.py
+5 -2
View File
@@ -133,6 +133,10 @@ def glm5_next_config(model_config: ModelConfig):
return None
def hybrid_kda_config(model_config: ModelConfig):
return kimi_linear_config(model_config) or glm5_next_config(model_config)
def linear_attn_model_spec(model_config: ModelConfig):
result = _get_linear_attn_registry_result(model_config)
return result[0] if result else None
@@ -142,8 +146,7 @@ def mambaish_config(model_config: ModelConfig):
existing = (
mamba2_config(model_config)
or hybrid_gdn_config(model_config)
or kimi_linear_config(model_config)
or glm5_next_config(model_config)
or hybrid_kda_config(model_config)
or hybrid_lightning_config(model_config)
)
if existing:
python/sglang/srt/configs/model_config.py
+7
View File
@@ -927,6 +927,11 @@ class ModelConfig:
and self.hf_config.architectures[0] == "InklingForConditionalGeneration"
):
self.hf_config.architectures[0] = "InklingForConditionalGenerationMTP"
if (
is_draft_model
and self.hf_config.architectures[0] == "GigaChat35ForCausalLM"
):
self.hf_config.architectures[0] = "GigaChat35ForCausalLMNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "Step3p7ForConditionalGeneration"
@@ -1191,6 +1196,8 @@ class ModelConfig:
or "MistralLarge3ForCausalLMEagle" in self.hf_config.architectures
or "KimiK25ForConditionalGeneration" in self.hf_config.architectures
or "Eagle3DeepseekV2ForCausalLM" in self.hf_config.architectures
or "GigaChat35ForCausalLM" in self.hf_config.architectures
or "GigaChat35ForCausalLMNextN" in self.hf_config.architectures
):
self.head_dim = 256
self.attention_arch = AttentionArch.MLA
python/sglang/srt/constrained/base_grammar_backend.py
+1
View File
@@ -385,6 +385,7 @@ def create_grammar_backend(
vocab_size=vocab_size,
model_eos_token_ids=eos_list,
any_whitespace=not get_serving().constrained_json_disable_any_whitespace,
max_whitespace_cnt=get_serving().constrained_json_max_whitespace_cnt,
)
except TokenizerNotSupportedError as e:
if get_serving().enable_strict_thinking:
python/sglang/srt/constrained/xgrammar_backend.py
+10 -2
View File
@@ -215,6 +215,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
vocab_size: int,
model_eos_token_ids: Optional[List[int]] = None,
any_whitespace: bool = True,
max_whitespace_cnt: Optional[int] = None,
):
super().__init__()
@@ -244,6 +245,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
self.vocab_size = vocab_size
self.override_stop_tokens = override_stop_tokens
self.any_whitespace = any_whitespace
self.max_whitespace_cnt = max_whitespace_cnt
@property
def is_support_token_filter(self):
@@ -259,11 +261,15 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
@staticmethod
def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
if logits.device.type in {"cuda", "npu", "xpu", "musa"}:
if logits.device.type in {"cuda", "xpu", "musa"}:
if _is_hip:
apply_token_bitmask_inplace_cuda(logits, vocab_mask)
else:
apply_token_bitmask_inplace_triton(logits, vocab_mask)
elif logits.device.type == "npu":
import sgl_kernel_npu # noqa: F401
torch.ops.npu.apply_token_bitmask(logits, vocab_mask)
else:
raise RuntimeError(f"Unsupported device: {logits.device.type}")
@@ -348,7 +354,9 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
schema = json.loads(key_string)
validate_xgrammar_json_schema(schema)
ctx = self.grammar_compiler.compile_json_schema(
schema=key_string, any_whitespace=self.any_whitespace
schema=key_string,
any_whitespace=self.any_whitespace,
max_whitespace_cnt=self.max_whitespace_cnt,
)
except (
python/sglang/srt/disaggregation/base/conn.py
+5
View File
@@ -24,6 +24,7 @@ class StateType(str, enum.Enum):
# only the live subrange of that row for the current open pool.
DSA_TAIL = "dsa_tail"
MINIMAX_INDEX_K = "minimax_index_k"
MINIMAX_DENSE_KV = "minimax_dense_kv"
# DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
# (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.
SWA_RING = "swa_ring"
@@ -171,6 +172,10 @@ class BaseKVSender(ABC):
def pop_decode_prefix_len(self) -> int:
return 0
def get_max_transfer_tokens(self) -> Optional[int]:
"""Optional page-aligned limit for one scheduler KV send."""
return None
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
return num_pages > 0
python/sglang/srt/disaggregation/common/conn.py
+25 -1
View File
@@ -35,9 +35,12 @@ from sglang.srt.environ import envs
from sglang.srt.runtime_context import (
get_disagg,
get_parallel,
get_schedule,
get_serving,
max_prefill_buffer_tokens,
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils.common import ceil_align
from sglang.srt.utils.network import (
NetworkAddress,
get_local_ip_auto,
@@ -181,6 +184,7 @@ class CommonKVManager(BaseKVManager):
envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get()
)
self._dcp_pack_buffers = None
self._dcp_pack_max_tokens: Optional[int] = None
# for p/d multi node infer
self.bootstrap_host = get_serving().host
self.bootstrap_port = get_disagg().disaggregation_bootstrap_port
@@ -378,19 +382,26 @@ class CommonKVManager(BaseKVManager):
f"{type(self).__name__} does not support staging memory registration"
)
def _init_dcp_pack_buffers_once(self, dcp_size: int) -> None:
def _init_dcp_pack_buffers_once(
self, dcp_size: int, *, include_draft: bool = False
) -> None:
if self._dcp_pack_buffers is not None:
return
if not self.kv_args.kv_item_lens:
return
from sglang.srt.disaggregation.common.dcp_pack import init_dcp_pack_buffers
max_tokens = max_prefill_buffer_tokens() or get_schedule().max_prefill_tokens
max_tokens = ceil_align(max_tokens, self.kv_args.page_size)
self._dcp_pack_buffers = init_dcp_pack_buffers(
self._register_staging_memory,
self.kv_args,
len(self.transfer_queues),
dcp_size,
max_tokens,
include_draft=include_draft,
)
self._dcp_pack_max_tokens = max_tokens
def check_status(self, bootstrap_room: int) -> KVPoll:
return self.request_status[bootstrap_room]
@@ -1565,6 +1576,19 @@ class CommonKVSender(BaseKVSender):
def pop_decode_prefix_len(self) -> int:
return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
def get_max_transfer_tokens(self) -> Optional[int]:
if self.kv_mgr._dcp_pack_max_tokens is None:
return None
for peer, info in self.kv_mgr.transfer_infos.get(
self.bootstrap_room, {}
).items():
if (
not info.is_dummy
and self.kv_mgr.decode_kv_args_table[peer].requires_dcp_relayout
):
return self.kv_mgr._dcp_pack_max_tokens
return None
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
return num_pages > 0 or last_chunk

Some files were not shown because too many files have changed in this diff Show More