diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 40823cf28..6d39630ef 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -514,6 +514,21 @@ jobs: rust_ext_artifact: '' secrets: inherit + base-c-test-8-gpu-b300: + needs: [check-changes, call-gate, wait-for-base-b, sgl-kernel-build-wheels, rust-ext-build] + if: ${{ !failure() && !cancelled() }} + uses: ./.github/workflows/_pr-test-stage.yml + with: + self_name: base-c-test-8-gpu-b300 + runner_config: 8-gpu-b300 + check_changes: ${{ toJson(needs.check-changes.outputs) }} + caller_inputs: ${{ toJson(inputs) }} + partitions: ${{ needs.check-changes.outputs.partitions }} + run_timeout_minutes: '60' + timeout_per_file: '3600' + rust_ext_artifact: ${{ needs.rust-ext-build.outputs.artifact_name }} + secrets: inherit + # List every build and test job: `skipped` passes here, so an omission turns that # job's failure into a green run with no tests. pr-test-finish: @@ -549,6 +564,7 @@ jobs: base-c-test-deepep-8-gpu-h200, base-c-test-4-gpu-b200, base-c-test-4-gpu-gb300, + base-c-test-8-gpu-b300, ] if: always() runs-on: ubuntu-latest diff --git a/docker/Dockerfile b/docker/Dockerfile index 89b76b1eb..710e65570 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -12,7 +12,7 @@ ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee ARG BUILD_AND_DOWNLOAD_PARALLEL=8 ARG SGL_KERNEL_VERSION=0.4.5 ARG SGL_VERSION -ARG SGL_DEEP_GEMM_VERSION=0.1.5 +ARG SGL_DEEP_GEMM_VERSION=0.1.5.post1 ARG USE_LATEST_SGLANG=0 ARG GDRCOPY_VERSION=2.5.1 ARG PIP_DEFAULT_INDEX diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 6fe524634..619e5af52 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -1966,8 +1966,8 @@ SGLang supports various environment variables that can be used to configure its SGLANG_USE_IPC_POOL_HANDLE_CACHE - Cache CUDA IPC pool handles. - false + When CUDA IPC multimodal feature transport is selected, reuse mappings to its existing bounded pool. This does not enable CUDA IPC transport or reserve another pool. + true SGLANG_MM_FEATURE_CACHE_MB diff --git a/python/pyproject.toml b/python/pyproject.toml index 2650861a8..8f6cc75f9 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -67,7 +67,7 @@ dependencies = [ "scipy", "sentencepiece", "setproctitle", - "sgl-deep-gemm==0.1.5", + "sgl-deep-gemm==0.1.5.post1", "sglang-kernel==0.4.5", "smg-grpc-servicer>=0.5.0", "soundfile==0.13.1", diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh index b2eb01106..e306f5442 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh @@ -18,9 +18,11 @@ #include #include -#include #include #include +#ifndef USE_ROCM +#include +#endif namespace sglang { diff --git a/python/sglang/kernels/jit/csrc/moe/align_single_token.cuh b/python/sglang/kernels/jit/csrc/moe/align_single_token.cuh new file mode 100644 index 000000000..d13a03c26 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/moe/align_single_token.cuh @@ -0,0 +1,107 @@ +// Tiny moe_align_block_size for M == 1 decode: one warp replaces the +// moe_align_block_size + count_and_sort_expert_tokens kernel pair (~4.4us) +// with a single ~1.5us launch. +// +// For a single token the top-k expert ids are distinct, so the aligned +// layout is exactly: experts sorted ascending, one block per expert, +// slot i of block b = flat topk index for b's expert, remaining block +// slots padded with numel (= topk). num_tokens_post_padded = topk * block. + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include // For device::cast +#include // For LaunchKernel + +#include + +#include + +namespace { + +struct AlignSingleTokenParams { + const int32_t* __restrict__ topk_ids; // [1, topk] + int32_t* __restrict__ sorted_ids; // [topk * block_size] + int32_t* __restrict__ expert_ids; // [topk] + int32_t* __restrict__ num_post; // [1] + uint32_t topk; + uint32_t block_size; +}; + +template +__global__ void align_single_token_kernel(const AlignSingleTokenParams __grid_constant__ params) { + using namespace device; + const uint32_t lane = threadIdx.x; // one warp + const uint32_t topk = params.topk; + const uint32_t bs = params.block_size; + + PDLWaitPrimary(); + + int32_t my_id = (lane < topk) ? params.topk_ids[lane] : INT32_MAX; + + // Rank of my expert id among the topk (ids are distinct for one token; + // tie-break on lane keeps this robust anyway). + uint32_t rank = 0; + for (uint32_t j = 0; j < topk; ++j) { + int32_t other = __shfl_sync(0xffffffff, my_id, j); + if (other < my_id || (other == my_id && j < lane)) { + rank++; + } + } + + if (lane < topk) { + params.expert_ids[rank] = my_id; + // block `rank`: first slot is my flat index (token 0, slot `lane`), + // rest padded with numel (= topk). + params.sorted_ids[rank * bs] = static_cast(lane); + } + // Fill padding cooperatively: positions not equal to a block start. + for (uint32_t p = lane; p < topk * bs; p += 32) { + if (p % bs != 0) { + params.sorted_ids[p] = static_cast(topk); + } + } + if (lane == 0) { + params.num_post[0] = static_cast(topk * bs); + } + + PDLTriggerSecondary(); +} + +template +struct AlignSingleTokenKernel { + static constexpr auto kernel = align_single_token_kernel; + + static void + run(const tvm::ffi::TensorView topk_ids, + const tvm::ffi::TensorView sorted_ids, + const tvm::ffi::TensorView expert_ids, + const tvm::ffi::TensorView num_post, + int64_t block_size) { + using namespace host; + + auto One_ = SymbolicSize{"one"}; + auto K_ = SymbolicSize{"topk"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({One_, K_}).with_dtype().with_device(device).verify(topk_ids); + + const auto topk = static_cast(K_.unwrap()); + RuntimeCheck(One_.unwrap() == 1, "moe_align_single_token requires M == 1"); + RuntimeCheck(topk <= 32, "moe_align_single_token requires topk <= 32"); + + const auto params = AlignSingleTokenParams{ + .topk_ids = static_cast(topk_ids.data_ptr()), + .sorted_ids = static_cast(sorted_ids.data_ptr()), + .expert_ids = static_cast(expert_ids.data_ptr()), + .num_post = static_cast(num_post.data_ptr()), + .topk = topk, + .block_size = static_cast(block_size), + }; + + LaunchKernel(dim3(1), 32, device.unwrap()).enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/kernels/ops/moe/moe_align_single_token.py b/python/sglang/kernels/ops/moe/moe_align_single_token.py new file mode 100644 index 000000000..68677f325 --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_align_single_token.py @@ -0,0 +1,49 @@ +"""CUDA JIT single-warp moe_align_block_size for M == 1 decode batches.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_align_single_token_module() -> Module: + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + "moe_align_single_token", + *args, + cuda_files=["moe/align_single_token.cuh"], + cuda_wrappers=[("run", f"AlignSingleTokenKernel<{args}>::run")], + extra_cuda_cflags=["-O3"], + ) + + +def moe_align_single_token( + topk_ids: torch.Tensor, block_size: int +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """moe_align_block_size for a single token (distinct expert ids). + + Returns (sorted_token_ids, expert_ids, num_tokens_post_padded) with the + same layout as moe_align_block_size: experts ascending, one block per + expert, padding value = numel. + """ + topk = topk_ids.shape[1] + device = topk_ids.device + sorted_ids = torch.empty((topk * block_size,), dtype=torch.int32, device=device) + expert_ids = torch.empty((topk,), dtype=torch.int32, device=device) + num_post = torch.empty((1,), dtype=torch.int32, device=device) + _jit_align_single_token_module().run( + topk_ids, sorted_ids, expert_ids, num_post, block_size + ) + return sorted_ids, expert_ids, num_post diff --git a/python/sglang/kernels/ops/sampling/renorm_triton.py b/python/sglang/kernels/ops/sampling/renorm_triton.py new file mode 100644 index 000000000..a50956688 --- /dev/null +++ b/python/sglang/kernels/ops/sampling/renorm_triton.py @@ -0,0 +1,172 @@ +"""ROCm-compatible top-k / top-p probability renormalization fallbacks.""" + +from __future__ import annotations + +from typing import Union + +import torch +import triton +import triton.language as tl + +_BLOCK_SIZE = 1024 + + +@triton.jit +def _mask_and_partial_sum_kernel( + probs_ptr, + pivots_ptr, + out_ptr, + partial_sums_ptr, + vocab_size: tl.constexpr, + num_chunks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + chunk = tl.program_id(1) + offsets = chunk * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < vocab_size + row_offsets = row * vocab_size + offsets + + probs = tl.load(probs_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + pivot = tl.load(pivots_ptr + row) + kept = tl.where(mask & (probs >= pivot), probs, 0.0) + + tl.store(out_ptr + row_offsets, kept, mask=mask) + tl.store(partial_sums_ptr + row * num_chunks + chunk, tl.sum(kept, axis=0)) + + +@triton.jit +def _normalize_kernel( + out_ptr, + row_sums_ptr, + numel, + vocab_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + row = offsets // vocab_size + values = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + denominator = tl.load(row_sums_ptr + row, mask=mask, other=1.0) + tl.store(out_ptr + offsets, values / denominator, mask=mask) + + +def _prepare_probs(probs: torch.Tensor) -> torch.Tensor: + if probs.ndim != 2: + raise ValueError(f"probs must be 2D, got shape={tuple(probs.shape)}") + if not probs.is_cuda: + raise ValueError("renorm kernels require a CUDA/HIP tensor") + return probs.float().contiguous() + + +def _renorm_from_pivots(probs_fp32: torch.Tensor, pivots: torch.Tensor) -> torch.Tensor: + batch_size, vocab_size = probs_fp32.shape + num_chunks = triton.cdiv(vocab_size, _BLOCK_SIZE) + out = torch.empty_like(probs_fp32) + partial_sums = torch.empty( + (batch_size, num_chunks), device=probs_fp32.device, dtype=torch.float32 + ) + _mask_and_partial_sum_kernel[(batch_size, num_chunks)]( + probs_fp32, + pivots, + out, + partial_sums, + vocab_size=vocab_size, + num_chunks=num_chunks, + BLOCK_SIZE=_BLOCK_SIZE, + num_warps=8, + ) + + row_sums = partial_sums.sum(dim=1) + _normalize_kernel[(triton.cdiv(out.numel(), _BLOCK_SIZE),)]( + out, + row_sums, + out.numel(), + vocab_size=vocab_size, + BLOCK_SIZE=_BLOCK_SIZE, + num_warps=8, + ) + return out + + +def top_p_renorm_probs_triton( + probs: torch.Tensor, top_p: Union[torch.Tensor, float] +) -> torch.Tensor: + """Apply exact top-p thresholding and renormalize each probability row. + + Sorting and prefix sums use PyTorch's device kernels because a vocabulary-sized + in-register Triton sort does not scale to 100K+ vocabularies. Triton performs + the bandwidth-heavy masking, partial reduction, and normalization. + """ + probs_fp32 = _prepare_probs(probs) + batch_size, vocab_size = probs_fp32.shape + if batch_size == 0 or vocab_size == 0: + return probs_fp32 + + if isinstance(top_p, torch.Tensor): + top_ps = top_p.to(device=probs.device, dtype=torch.float32).reshape(-1) + if top_ps.numel() == 1: + top_ps = top_ps.expand(batch_size) + elif top_ps.numel() != batch_size: + raise ValueError( + f"top_p must be scalar or have one value per row, got " + f"{top_ps.numel()} values for {batch_size} rows" + ) + else: + if not 0.0 < float(top_p) <= 1.0: + raise ValueError("top_p values must be in (0, 1]") + top_ps = torch.full( + (batch_size,), float(top_p), device=probs.device, dtype=torch.float32 + ) + + # Match FlashInfer's threshold semantics: sort ascending, discard the prefix + # whose cumulative mass is below 1 - p, and retain all ties at the pivot. + sorted_probs = torch.sort(probs_fp32, dim=-1).values + cdf = torch.cumsum(sorted_probs, dim=-1) + cutoff = torch.searchsorted(cdf, (1.0 - top_ps).unsqueeze(1), right=False).squeeze( + 1 + ) + cutoff.clamp_(max=vocab_size - 1) + pivots = sorted_probs.gather(1, cutoff.unsqueeze(1)).squeeze(1).contiguous() + + return _renorm_from_pivots(probs_fp32, pivots) + + +def top_k_renorm_probs_triton( + probs: torch.Tensor, top_k: Union[torch.Tensor, int] +) -> torch.Tensor: + """Apply exact top-k thresholding and renormalize each probability row. + + Sorting uses PyTorch's device kernels because a vocabulary-sized in-register + Triton sort does not scale to 100K+ vocabularies. Triton performs the + bandwidth-heavy masking, partial reduction, and normalization. + """ + probs_fp32 = _prepare_probs(probs) + batch_size, vocab_size = probs_fp32.shape + if batch_size == 0 or vocab_size == 0: + return probs_fp32 + + if isinstance(top_k, torch.Tensor): + top_ks = top_k.to(device=probs.device, dtype=torch.int64).reshape(-1) + if top_ks.numel() == 1: + top_ks = top_ks.expand(batch_size) + elif top_ks.numel() != batch_size: + raise ValueError( + f"top_k must be scalar or have one value per row, got " + f"{top_ks.numel()} values for {batch_size} rows" + ) + else: + top_ks = torch.full( + (batch_size,), int(top_k), device=probs.device, dtype=torch.int64 + ) + + # Match FlashInfer's threshold semantics: sort descending, keep the k highest + # probabilities, and retain all ties at the pivot. + sorted_probs = torch.sort(probs_fp32, dim=-1, descending=True).values + cutoff = (top_ks - 1).clamp_(min=0, max=vocab_size - 1) + pivots = sorted_probs.gather(1, cutoff.unsqueeze(1)).squeeze(1).contiguous() + + return _renorm_from_pivots(probs_fp32, pivots) + + +__all__ = ["top_k_renorm_probs_triton", "top_p_renorm_probs_triton"] diff --git a/python/sglang/srt/arg_groups/kimi_k3_hook.py b/python/sglang/srt/arg_groups/kimi_k3_hook.py new file mode 100644 index 000000000..7fe32a90c --- /dev/null +++ b/python/sglang/srt/arg_groups/kimi_k3_hook.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +def apply_kimi_k3_spec_backend_defaults(server_args: ServerArgs) -> None: + """Apply speculative backend defaults for Kimi hybrid models.""" + from sglang.srt.utils import is_sm100_supported + + if server_args.speculative_algorithm is None: + return + + # Use the fused Kimi-K3/DSPARK CuTeDSL kernel for KDA target verification. + # Decode is left free (its bf16-ssm SM100+ flashinfer default is fine -- the + # target only verifies under spec); the verify backend is pinned directly. + if server_args.linear_attn_verify_backend is None: + server_args.linear_attn_verify_backend = "nv_cutedsl" + logger.info( + "Kimi hybrid model with speculative decoding: pinning " + "--linear-attn-verify-backend to nv_cutedsl (uses the fused " + "Kimi-K3/DSPARK CuTeDSL kernel)." + ) + + # dspark's draft is dense MQA; trtllm_mha avoids flashinfer's blocking + # per-step host plan. DSPARK-only: other spec algos use MLA-family drafts. + if ( + server_args.speculative_algorithm == "DSPARK" + and server_args.speculative_draft_attention_backend is None + and is_sm100_supported() + ): + server_args.speculative_draft_attention_backend = "trtllm_mha" + logger.info( + "Kimi hybrid DSPARK: defaulting " + "--speculative-draft-attention-backend to trtllm_mha." + ) + + +def disable_kimi_k3_symm_mem(server_args: ServerArgs) -> None: + """Turn `--enable-symm-mem` back off unless every phase runs eager. + + Symm-mem allocations are per-forward, so an address captured into a graph is + neither reserved for its lifetime nor at the same offset on every rank. Under + capture that corrupts spec decode: accept collapses to 1.000, or the server + silently emits garbage with accept pinned at the ceiling. Prefill counts too -- + the same allocation sits in any captured RowParallelLinear. + + Gates on the arch itself: this runs from cuda-graph resolution, which is earlier + than the model-specific hook block. + """ + from sglang.srt.connector import ConnectorType + from sglang.srt.model_executor.cuda_graph_config import Backend + from sglang.srt.utils import parse_connector_type + + if not server_args.enable_symm_mem: + return + if parse_connector_type(server_args.model_path) == ConnectorType.INSTANCE: + return + if server_args.get_model_config().hf_config.architectures[0] not in ( + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", + ): + return + graph = server_args.cuda_graph_config + if ( + graph.decode.backend == Backend.DISABLED + and graph.prefill.backend == Backend.DISABLED + ): + return + server_args.enable_symm_mem = False + logger.warning( + "Kimi hybrid model: ignoring --enable-symm-mem because CUDA graphs are on. " + "The symmetric-memory pool's per-forward allocations are not valid for the " + "lifetime of a captured graph, which corrupts speculative decoding and can " + "silently produce wrong output. The auto-probed K3 fused all-reduce is faster " + "anyway. Disable capture on every phase " + "(--cuda-graph-backend-decode=disabled --cuda-graph-backend-prefill=disabled) " + "if you genuinely need symmetric memory." + ) + + +def apply_kimi_k3_linear_attn_defaults(server_args: ServerArgs) -> None: + """KDA decode-fallback default for Kimi hybrid models (spec-independent).""" + from sglang.srt.utils import is_sm100_supported + + # Preempts the generic SM100+bf16 flashinfer switch (a GDN default): on + # KDA shapes the triton packed decode measures ~35% faster than + # recurrent_kda across bs 1-256, and ReplaySSM requires triton. + if ( + server_args.linear_attn_decode_backend is None + and server_args.mamba_ssm_dtype == "bfloat16" + and is_sm100_supported() + ): + server_args.linear_attn_decode_backend = "triton" + logger.info( + "Kimi hybrid model with bf16 SSM state: defaulting " + "--linear-attn-decode-backend to triton." + ) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 3ab05bcc4..41f3550be 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -30,6 +30,7 @@ Two declaration forms, keyed on ``hf_config.architectures[0]``: from __future__ import annotations import dataclasses +import inspect import logging from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple @@ -39,6 +40,7 @@ from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.utils.common import ( cpu_has_amx_support, get_device_capability, + get_device_name, get_device_sm, get_nvidia_driver_version, get_quantization_config, @@ -48,6 +50,7 @@ from sglang.srt.utils.common import ( is_flashinfer_available, is_gfx95_supported, is_hip, + is_mnnvl_fabric_device, is_musa, is_npu, is_sm90_supported, @@ -315,6 +318,234 @@ def _register_for(*architectures: str): return decorator +def _dspark_verify_on_decode_backend( + backend: Optional[str], q_len: int, kv_cache_dtype: Optional[str] +) -> bool: + """Whether the MLA decode backend can serve a q_len-wide target verify.""" + if backend == "trtllm_mla": + return True + if backend == "tokenspeed_mla": + return kv_cache_dtype == "fp8_e4m3" and q_len <= 8 + if backend == "cutedsl_mla": + # The cute-dsl kernel rejects q_len >= 5 with no fallback. + return q_len <= 4 + return False + + +_KIMI_K3_DCP_PATCH_URL = ( + "https://github.com/sgl-project/sglang/blob/" + "b701464720ca22aa1851d5dda7144e84a410f2c7/" + "docker/kimi_k3/kimi_k3_cu13.Dockerfile#L116-L123" +) + + +def _require_kimi_k3_cutedsl_dcp_support() -> None: + try: + from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla + + parameters = inspect.signature(trtllm_batch_decode_with_kv_cache_mla).parameters + except (ImportError, TypeError, ValueError) as exc: + raise RuntimeError( + "Kimi-K3 DCP with decode_attention_backend='cutedsl_mla' requires " + "a DCP-patched FlashInfer " + "trtllm_batch_decode_with_kv_cache_mla exposing enable_dcp in its " + f"signature. Apply the patch as shown in {_KIMI_K3_DCP_PATCH_URL}." + ) from exc + + if "enable_dcp" not in parameters: + raise RuntimeError( + "Kimi-K3 DCP with decode_attention_backend='cutedsl_mla' requires " + "enable_dcp in the signature of " + "flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla. Apply " + f"the FlashInfer DCP patch as shown in {_KIMI_K3_DCP_PATCH_URL}." + ) + + +@_register_for("KimiK3ForConditionalGeneration") +def _kimi_k3_overrides(server_args: Any, hf_config: Any) -> dict: + if server_args.dcp_size > 1: + overrides = {} + if server_args.enable_symm_mem: + logger.warning( + "Kimi-K3 DCP disables --enable-symm-mem due to decode CUDA " + "graph correctness issues." + ) + overrides["enable_symm_mem"] = False + + if server_args.speculative_algorithm == "DSPARK": + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) + + ragged_mode = read_ragged_verify_mode() + if ragged_mode is not RaggedVerifyMode.STATIC: + raise ValueError( + "Kimi-K3 DCP + DSPARK currently requires " + "SGLANG_RAGGED_VERIFY_MODE=static; compact/cap-accept are " + f"not validated under DCP (got {ragged_mode.value!r})." + ) + + # DSPARK target-verify + draft-extend must run on the decode + # (cutedsl_mla) backend, whose _run_decode_kernel implements the DCP + # signature (causal_seqs / cp_world / cp_rank). The default + # "prefill" routes verify to trtllm_mla, whose base _run_decode_kernel + # lacks that DCP path (TypeError: unexpected kwarg 'causal_seqs'). + overrides["speculative_attention_mode"] = "decode" + + prefill_backend, decode_backend = attention_backends_of(server_args) + if decode_backend == "cutedsl_mla" or decode_backend is None: + _require_kimi_k3_cutedsl_dcp_support() + logger.info( + "Kimi-K3 DCP keeps decode attention backend 'cutedsl_mla' " + f"(prefill={prefill_backend!r} -> 'trtllm_mla')." + ) + overrides.update( + prefill_attention_backend="trtllm_mla", + decode_attention_backend="cutedsl_mla", + ) + elif decode_backend == "tokenspeed_mla": + logger.info( + "Kimi-K3 DCP overrides attention backends: " + f"prefill={prefill_backend!r}, decode={decode_backend!r} -> " + "'tokenspeed_mla'." + ) + logger.info( + "Kimi-K3 DCP with tokenspeed mla backend overrides KV cache dtype: " + f"{server_args.kv_cache_dtype!r} -> 'fp8_e4m3'." + ) + overrides.update( + prefill_attention_backend="tokenspeed_mla", + decode_attention_backend="tokenspeed_mla", + kv_cache_dtype="fp8_e4m3", + ) + else: + raise AssertionError( + f"Decode attention backend for Kimi-K3 DCP must be 'cutedsl_mla' or 'tokenspeed_mla', got {decode_backend!r}." + ) + + if server_args.dcp_replicate_q_proj is None: + logger.info("Kimi-K3 DCP enables replicated Q projection by default.") + overrides["dcp_replicate_q_proj"] = True + + device_name = get_device_name() + dcp_comm_backend = "fi_a2a" if is_mnnvl_fabric_device() else "a2a" + logger.info( + "Kimi-K3 DCP selects communication backend on " + f"{device_name!r}: {server_args.dcp_comm_backend!r} -> " + f"{dcp_comm_backend!r}." + ) + overrides["dcp_comm_backend"] = dcp_comm_backend + return overrides + + if not (is_sm100_supported() and get_device_sm() in (100, 103)): + return {} + backends_unset = server_args.is_attention_backend_not_set() + if server_args.speculative_algorithm != "DSPARK": + if not backends_unset: + return {} + logger.info( + "Use trtllm_mla as the default prefill and decode attention " + "backend for Kimi-K3 on SM100/SM103." + ) + return { + "decode_attention_backend": "trtllm_mla", + "prefill_attention_backend": "trtllm_mla", + } + # DSPARK: verify runs on the decode backend (mode=decode below), so this + # picks the verify kernel -- mode=prefill routes it to flashinfer, which is + # slow and syncs, while plain decode is cold under dspark. + q_len = server_args.speculative_num_draft_tokens or ( + server_args.speculative_dspark_block_size + 1 + if server_args.speculative_dspark_block_size is not None + # Checkpoint auto-infer happens after overrides; K3 draft uses block 7. + else 8 + ) + overrides = {} + if backends_unset: + backend = "trtllm_mla" + overrides["decode_attention_backend"] = backend + overrides["prefill_attention_backend"] = "trtllm_mla" + else: + # Explicit backend knobs keep priority, but the mode is a separate knob + # that still needs declaring -- else verify stays on the prefill backend, + # whose host-side plan (flashinfer by default) forces a per-step D2H. + _, backend = attention_backends_of(server_args) + if _dspark_verify_on_decode_backend(backend, q_len, server_args.kv_cache_dtype): + overrides["speculative_attention_mode"] = "decode" + logger.info( + "Kimi-K3 DSPARK on SM100/SM103: decode/verify attention backend " + f"{backend} (speculative_attention_mode=decode)." + ) + else: + logger.warning( + f"Kimi-K3 DSPARK: decode attention backend {backend!r} cannot serve " + f"target verify at q_len={q_len}, so verify runs on the prefill " + "backend (speculative_attention_mode=prefill). A host-plan prefill " + "backend costs a per-step seq_lens D2H sync; leave the attention " + "backend knobs unset for the sync-free default." + ) + return overrides + + +def _is_mxfp4_pack_quantized(hf_config: Any) -> bool: + qc = getattr( + getattr(hf_config, "text_config", hf_config), "quantization_config", None + ) + if not isinstance(qc, dict): + return False + groups = qc.get("config_groups") or {} + return any( + "mxfp4" in str(g.get("format", "")) + for g in groups.values() + if isinstance(g, dict) + ) + + +@_register_for("KimiK3ForConditionalGeneration") +def _kimi_k3_moe_runner_overrides(server_args: Any, hf_config: Any) -> dict: + # MoE runner default, independent of the attention-backend gate above. + # trtllm-gen fused MoE (flashinfer_mxfp4) beats marlin on both the decode + # (M=bs) and the target-verify (M=bs*(gamma+1)) regimes on SM100/SM103; + # it hard-requires the SiTU cubin pool on the box (K3's SiTU activation has + # no public cubins). Do not silently trade W4A8 for Marlin W4A16 when the + # default cannot start; explicit non-FlashInfer runner choices still win. + if server_args.moe_runner_backend not in ("auto", "flashinfer_mxfp4"): + return {} + if not (is_sm100_supported() and get_device_sm() in (100, 103)): + return {} + if not _is_mxfp4_pack_quantized(hf_config): + return {} + from sglang.kernels.ops.moe.trtllm_gen_moe import available as _trtllm_gen_moe_ok + + if not _trtllm_gen_moe_ok(): + raise RuntimeError( + "Kimi-K3 on Blackwell with moe_runner_backend='auto' or " + "'flashinfer_mxfp4' requires a valid " + "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL. Install it with:\n" + "wget https://github.com/sgl-project/whl/releases/download/" + "trtllm_gen_moe_cubin_20260617/" + "trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip\n" + "sudo mkdir -p /opt/trtllm_gen_moe_cubin_pool\n" + "sudo unzip -q " + "trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip -d " + "/opt/trtllm_gen_moe_cubin_pool\n" + "export " + "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool/" + "trtllm_gen_moe_cubin_pool_20260617_v0613rc1\n" + "To use Marlin " + "instead, set --moe-runner-backend marlin explicitly." + ) + + if server_args.moe_runner_backend == "auto": + logger.info( + "Kimi-K3 on SM100/SM103: moe_runner_backend=flashinfer_mxfp4 " + "(trtllm-gen SiTU cubin pool found)." + ) + return {"moe_runner_backend": "flashinfer_mxfp4"} + return {} + + @_register_for( "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", @@ -1175,6 +1406,7 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: _MAMBA_RADIX_CACHE_ARCHS = frozenset( { "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", "BailingMoeV2_5ForCausalLM", "Qwen3NextForCausalLM", "Qwen3_5MoeForConditionalGeneration", @@ -1208,6 +1440,10 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset( "GraniteMoeHybridForCausalLM", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", + # KDA-based: same MambaPool ping-pong machinery as GDN; requires the + # KDA backend's track-snapshot writes (decode + extend) so donated + # slots hold real states for prefix-cache restores. + "KimiK3ForConditionalGeneration", } ) diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index cf7f9bed5..c1748a4a1 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -19,6 +19,7 @@ from sglang.srt.configs.interns2preview import InternS2PreviewConfig from sglang.srt.configs.janus_pro import MultiModalityConfig from sglang.srt.configs.jet_nemotron import JetNemotronConfig from sglang.srt.configs.jet_vlm import JetVLMConfig +from sglang.srt.configs.kimi_k3 import KimiK3Config from sglang.srt.configs.kimi_k25 import KimiK25Config from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.configs.kimi_vl import KimiVLConfig @@ -71,6 +72,7 @@ __all__ = [ "Step3VisionEncoderConfig", "Olmo3Config", "KimiLinearConfig", + "KimiK3Config", "KimiK25Config", "LagunaConfig", "Qwen3NextConfig", diff --git a/python/sglang/srt/configs/hybrid_arch.py b/python/sglang/srt/configs/hybrid_arch.py index e3e05b28a..31c284053 100644 --- a/python/sglang/srt/configs/hybrid_arch.py +++ b/python/sglang/srt/configs/hybrid_arch.py @@ -102,6 +102,9 @@ def kimi_linear_config(model_config: ModelConfig): config = model_config.hf_config if isinstance(config, KimiLinearConfig): return config + text_config = getattr(config, "text_config", None) + if isinstance(text_config, KimiLinearConfig): + return text_config return None diff --git a/python/sglang/srt/configs/kimi_k3.py b/python/sglang/srt/configs/kimi_k3.py new file mode 100644 index 000000000..fa82cd15a --- /dev/null +++ b/python/sglang/srt/configs/kimi_k3.py @@ -0,0 +1,124 @@ +from transformers.configuration_utils import PretrainedConfig + +from sglang.srt.configs.kimi_linear import KimiLinearConfig + + +class KimiK3VisionConfig(PretrainedConfig): + model_type = "kimi_k3_vision" + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + vt_num_attention_heads: int = 12, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1024, + vt_intermediate_size: int = 4096, + merge_kernel_size: tuple[int, int] = (2, 2), + video_attn_type: str = "spatial_temporal", + merge_type: str = "sd2_tpool", + _attn_implementation: str = "flash_attention_2", + mm_projector_type: str = "patchmergerv2", + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + qkv_hidden_size: int = 1536, + norm_type: str = "rmsnorm", + attn_bias: bool = False, + patch_embed_proj_bias: bool = False, + mlp_type: str = "mlp2", + linear_bias: bool = False, + activation_func: str = "gelu_pytorch_tanh", + pos_emb_interpolation_mode: str = "bilinear", + text_hidden_size: int = 2304, + **kwargs, + ): + super().__init__(**kwargs) + + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = tuple(merge_kernel_size) + self.video_attn_type = video_attn_type + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = ( + mm_hidden_size if mm_hidden_size is not None else vt_hidden_size + ) + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + self.qkv_hidden_size = qkv_hidden_size + self.norm_type = norm_type + self.attn_bias = attn_bias + self.patch_embed_proj_bias = patch_embed_proj_bias + self.mlp_type = mlp_type + self.linear_bias = linear_bias + self.activation_func = activation_func + self.pos_emb_interpolation_mode = pos_emb_interpolation_mode + + # Aliases consumed by the K2.5 vision implementation. + self.num_attention_heads = vt_num_attention_heads + self.num_hidden_layers = vt_num_hidden_layers + self.hidden_size = vt_hidden_size + self.intermediate_size = vt_intermediate_size + + +class KimiK3Config(PretrainedConfig): + model_type = "kimi_k3" + + def __init__( + self, + text_config: dict | KimiLinearConfig | None = None, + vision_config: dict | KimiK3VisionConfig | None = None, + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + image_placeholder: str = "<|kimi_image_placeholder|>", + **kwargs, + ): + if text_config is None: + self.text_config = KimiLinearConfig() + elif isinstance(text_config, dict): + self.text_config = KimiLinearConfig(**text_config) + else: + self.text_config = text_config + + if vision_config is None: + self.vision_config = KimiK3VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = KimiK3VisionConfig(**vision_config) + else: + self.vision_config = vision_config + + if self.vision_config.text_hidden_size != self.text_config.hidden_size: + self.vision_config.text_hidden_size = self.text_config.hidden_size + + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + self.image_placeholder = image_placeholder + + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + @property + def hidden_size(self) -> int: + return self.text_config.hidden_size + + @property + def vocab_size(self) -> int: + return self.text_config.vocab_size diff --git a/python/sglang/srt/configs/kimi_linear.py b/python/sglang/srt/configs/kimi_linear.py index 181dca1d3..45e4838c4 100644 --- a/python/sglang/srt/configs/kimi_linear.py +++ b/python/sglang/srt/configs/kimi_linear.py @@ -43,6 +43,7 @@ class KimiLinearConfig(PretrainedConfig): use_grouped_topk: bool = True, num_expert_group: int = 1, topk_group: int = 1, + topk_method: str = "noaux_tc", q_lora_rank: int | None = None, kv_lora_rank: int | None = None, qk_nope_head_dim: int | None = None, @@ -51,6 +52,13 @@ class KimiLinearConfig(PretrainedConfig): mla_use_nope: bool | None = False, num_nextn_predict_layers: int = 0, linear_attn_config: dict | None = None, + attn_res_block_size: int | None = None, + routed_expert_hidden_size: int | None = None, + latent_moe_use_norm: bool = False, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + mla_use_output_gate: bool = False, + max_position_embeddings: int = 4096, **kwargs, ): self.model_type = model_type @@ -83,6 +91,7 @@ class KimiLinearConfig(PretrainedConfig): self.mla_use_nope = mla_use_nope # moe config self.n_routed_experts = self.num_experts = num_experts + self.topk_method = topk_method self.num_experts_per_token = num_experts_per_token self.moe_renormalize = moe_renormalize self.num_shared_experts = num_shared_experts @@ -97,6 +106,14 @@ class KimiLinearConfig(PretrainedConfig): self.topk_group = topk_group self.num_nextn_predict_layers = num_nextn_predict_layers + self.attn_res_block_size = attn_res_block_size + self.routed_expert_hidden_size = routed_expert_hidden_size + self.latent_moe_use_norm = latent_moe_use_norm + self.activation_situ_beta = activation_situ_beta + self.activation_situ_linear_beta = activation_situ_linear_beta + self.mla_use_output_gate = mla_use_output_gate + self.max_position_embeddings = max_position_embeddings + if linear_attn_config is not None: assert linear_attn_config["kda_layers"] is not None assert linear_attn_config["full_attn_layers"] is not None diff --git a/python/sglang/srt/configs/mamba_utils.py b/python/sglang/srt/configs/mamba_utils.py index 11eefe174..4a9591d54 100644 --- a/python/sglang/srt/configs/mamba_utils.py +++ b/python/sglang/srt/configs/mamba_utils.py @@ -128,18 +128,26 @@ class BaseLinearStateParams(ABC): """Per-slot bytes of the ReplaySSM spec-verify fold window (all layers). Not part of ``mamba_cache_per_req``, so the memory solver must charge it separately. MUST mirror the ``MambaPool`` allocation: - raw v/k in the conv dtype + fp32 g and beta. GDN scalar-g layout - only.""" - assert not self.is_kda, "replayssm ring accounting supports GDN only" + raw v/k in the conv dtype + fp32 beta, plus the fp32 gate ring + (per-head scalar for GDN, per-K vector for KDA). KDA additionally + keeps the chunked d/k rings under spec (its forward_decode routes on + their presence), also in the conv dtype.""" hv, v_dim, k_dim = self.shape.temporal h_k = self.shape.num_k_heads_per_tp conv_b = self.dtype.conv.itemsize fp32_b = 4 per_layer = ( - hv * record_len * v_dim * conv_b - + h_k * record_len * k_dim * conv_b - + 2 * hv * record_len * fp32_b + hv * record_len * v_dim * conv_b # rawv + + h_k * record_len * k_dim * conv_b # rawk + + hv * record_len * fp32_b # beta (fp32) + # g (fp32): GDN per-head scalar, KDA per-K vector + + hv * record_len * (k_dim if self.is_kda else 1) * fp32_b ) + if self.is_kda: + per_layer += ( + hv * record_len * v_dim * conv_b # d + + h_k * record_len * k_dim * conv_b # k + ) return per_layer * len(self.layers) @property diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 7d5c25f21..15b84ed23 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -121,6 +121,10 @@ def is_deepseek_dsa(config) -> bool: ) +def is_kimi_k3(config) -> bool: + return _hf_arch(config) == "KimiK3ForConditionalGeneration" + + def is_deepseek_v4(config) -> bool: return _hf_arch(config) in ( "DeepseekV4ForCausalLM", @@ -484,6 +488,9 @@ class ModelConfig: self.is_multimodal_breakable_cuda_graph_supported = enable_multimodal and ( is_multimodal_breakable_cuda_graph_supported(self.hf_config.architectures) ) + self.is_mla_breakable_cuda_graph_supported = ( + is_mla_breakable_cuda_graph_supported(self.hf_config.architectures) + ) self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) # Derive context length and model shapes @@ -912,18 +919,20 @@ class ModelConfig: self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.v_head_dim = self.hf_text_config.v_head_dim self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim - elif "KimiLinearForCausalLM" in self.hf_config.architectures: + elif ( + "KimiLinearForCausalLM" in self.hf_config.architectures + or "KimiK3ForConditionalGeneration" in self.hf_config.architectures + ): + tc = self.hf_text_config self.head_dim = 72 self.attention_arch = AttentionArch.MLA - self.kv_lora_rank = self.hf_config.kv_lora_rank - self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim - self.v_head_dim = self.hf_config.v_head_dim - self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim + self.kv_lora_rank = tc.kv_lora_rank + self.qk_rope_head_dim = tc.qk_rope_head_dim + self.v_head_dim = tc.v_head_dim + self.qk_nope_head_dim = tc.qk_nope_head_dim self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - if self.hf_config.rope_scaling: - self.scaling = compute_mla_mscale_scaling( - self.hf_config.rope_scaling, self.scaling - ) + if getattr(tc, "rope_scaling", None): + self.scaling = compute_mla_mscale_scaling(tc.rope_scaling, self.scaling) elif ( "BailingMoeV2_5ForCausalLM" in self.hf_config.architectures or "BailingMoeForCausalLMNextN" in self.hf_config.architectures @@ -1841,6 +1850,15 @@ multimodal_breakable_cuda_graph_supported_model_archs = [ "Qwen3_5MoeForConditionalGeneration", ] +# MLA archs validated to run breakable CUDA graph when it is explicitly +# requested (--cuda-graph-backend-prefill=breakable bypasses the ServerArgs +# disable rules). Dispatch pins the absorbed MLA path inside capture/replay +# for these archs, so the prefill runner's MHA-companion prefix restrictions +# do not apply (see PrefillCudaGraphRunner.mla_pinned_under_bcg). +mla_breakable_cuda_graph_supported_model_archs = [ + "KimiK3ForConditionalGeneration", +] + if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): multimodal_model_archs.append(external_mm_model_arch) @@ -1915,6 +1933,14 @@ def is_multimodal_breakable_cuda_graph_supported(model_architectures: List[str]) ) +def is_mla_breakable_cuda_graph_supported(model_architectures: List[str]): + """Whether an MLA arch may keep prefill breakable CUDA graph enabled.""" + return any( + arch in mla_breakable_cuda_graph_supported_model_archs + for arch in model_architectures + ) + + # SequenceClassification models that use CrossEncodingPooler _cross_encoding_pooler_archs = [ "BertForSequenceClassification", @@ -1936,8 +1962,13 @@ def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float """Compute MLA attention scaling factor from rope_scaling with mscale. Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models. - Warns if 'factor' is missing from rope_scaling (common in v5 configs). + Transformers v5 also exposes the default RoPE parameters through + ``rope_scaling``. Those parameters do not request any scaling. """ + rope_type = rope_scaling.get("rope_type") or rope_scaling.get("type") + if rope_type == "default": + return base_scaling + if not rope_scaling.get("apply_yarn_scaling", True) or not rope_scaling.get( "apply_scale", True ): diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 8b29a73ce..c25d4755a 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -206,6 +206,7 @@ class CommonKVManager(BaseKVManager): self.request_status: Dict[int, KVPoll] = {} self._socket_cache: Dict[str, zmq.Socket] = {} self._monitor_cache: Dict[str, zmq.Socket] = {} + self._socket_send_locks: Dict[str, threading.Lock] = {} self._socket_lock = threading.Lock() self.failure_records: Dict[int, str] = {} self.failure_lock = threading.Lock() @@ -776,8 +777,18 @@ class CommonKVManager(BaseKVManager): self._monitor_cache[endpoint] = sock.get_monitor_socket( zmq.EVENT_DISCONNECTED ) + self._socket_send_locks.setdefault(endpoint, threading.Lock()) return sock + def _send_multipart_locked( + self, endpoint: str, parts: List[bytes], is_ipv6: bool = False + ): + # Cached sockets are shared across sender threads and zmq sockets are + # not thread-safe; serialize sends per endpoint. + sock = self._connect(endpoint, is_ipv6=is_ipv6) + with self._socket_send_locks[endpoint]: + sock.send_multipart(parts) + def get_mha_kv_ptrs_with_pp( self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int] ) -> Tuple[List[int], List[int], List[int], List[int], int]: diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 7ec83b677..240e20def 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -211,7 +211,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): speculative_eagle_topk: Optional[int] = None, linear_replayssm_cache_len: int = 16, mamba_envelope_layout: bool = False, - enable_gdn_replayssm_spec: bool = False, + enable_linear_replayssm_spec: bool = False, ): DecodeReqToTokenPool.__init__( self, @@ -258,7 +258,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): speculative_eagle_topk=speculative_eagle_topk, linear_replayssm_cache_len=linear_replayssm_cache_len, mamba_envelope_layout=mamba_envelope_layout, - enable_gdn_replayssm_spec=enable_gdn_replayssm_spec, + enable_linear_replayssm_spec=enable_linear_replayssm_spec, ) def clear(self): diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 861149645..07961389d 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -481,6 +481,7 @@ _GENERAL_VIDEO_META_ATTRS = ( "video_timestamps", "second_per_grid_ts", ) +_GENERAL_IMAGE_META_ATTRS = ("original_image_sizes",) # MiMo-VL audio-in-video fields; appended only when model_type is MiMo. _MIMO_VIDEO_AUDIO_META_ATTRS = ( "video_audio_feature_lens", @@ -551,6 +552,8 @@ class MultiModalEmbeddingData(EmbeddingData): self.img_grid_thw = [None] * num_parts self.video_grid_thw = [None] * num_parts self.audio_feature_lens = [None] * num_parts + for attr in _GENERAL_IMAGE_META_ATTRS: + setattr(self, attr, [None] * num_parts) self.modality_list = [ modality if part_idx == i else None for i in range(num_parts) ] @@ -567,6 +570,18 @@ class MultiModalEmbeddingData(EmbeddingData): self._set_part_grid(part_idx, modality, self.get_grid()) if modality == Modality.VIDEO: self._set_video_meta_for_part(part_idx, kwargs) + if modality == Modality.IMAGE: + self._set_image_meta_for_part(part_idx, kwargs) + + def _set_image_meta_for_part(self, part_idx, source): + for attr_name in _GENERAL_IMAGE_META_ATTRS: + val = ( + source.get(attr_name) + if isinstance(source, dict) + else getattr(source, attr_name, None) + ) + if val is not None: + getattr(self, attr_name)[part_idx] = val def _set_part_grid(self, part_idx, modality, grid): """Set the grid for one part according to modality (IMAGE/VIDEO/AUDIO).""" @@ -601,6 +616,10 @@ class MultiModalEmbeddingData(EmbeddingData): val = getattr(embedding_data, attr, None) if val is not None: extra[attr] = val + for attr in _GENERAL_IMAGE_META_ATTRS: + val = getattr(embedding_data, attr, None) + if val is not None: + extra[attr] = val mm_data = cls( part_idx=embedding_data.part_idx, num_parts=embedding_data.num_parts, @@ -651,6 +670,10 @@ class MultiModalEmbeddingData(EmbeddingData): kwargs[attr] = torch.cat(valid, dim=0) else: kwargs[attr] = list(itertools.chain(*valid)) + for attr in _GENERAL_IMAGE_META_ATTRS: + valid = [value for value in getattr(self, attr) if value is not None] + if valid: + kwargs[attr] = list(itertools.chain(*valid)) return kwargs def add(self, embedding_data: EmbeddingData): @@ -669,6 +692,8 @@ class MultiModalEmbeddingData(EmbeddingData): self._set_part_grid(pid, embedding_data.modality, embedding_data.get_grid()) if embedding_data.modality == Modality.VIDEO: self._set_video_meta_for_part(pid, embedding_data) + if embedding_data.modality == Modality.IMAGE: + self._set_image_meta_for_part(pid, embedding_data) class WaitingImageRequestStatus(IntEnum): @@ -678,6 +703,13 @@ class WaitingImageRequestStatus(IntEnum): TIMEOUT = -2 +def _select_mm_processor_prompt(recv_req, mm_processor): + """Mirror tokenizer-side prompt selection for scheduler-side EPD rebuilds.""" + if mm_processor.prefer_tokenized_input and recv_req.input_ids is not None: + return list(recv_req.input_ids) + return recv_req.input_text or recv_req.input_ids + + def create_part_req_id(original_req_id: str, part_idx: int) -> str: """Create a unique part request ID by appending part index suffix.""" return f"{original_req_id}_local_part_{part_idx}" @@ -724,6 +756,8 @@ class WaitingImageRequest: model_type, host_name, receive_count, + zmq_context=None, + embedding_port=None, ): self.rid = rid self.recv_req = recv_req @@ -736,9 +770,16 @@ class WaitingImageRequest: self.host_name = host_name self.receive_count = receive_count self.num_items_assigned = recv_req.num_items_assigned - self.embedding_port, self.recv_socket = get_zmq_socket_on_host( - zmq.Context(), zmq.PULL, host=host_name - ) + self.zmq_context = zmq_context + if embedding_port is None: + if self.zmq_context is None: + raise ValueError("zmq_context is required for a per-request socket") + self.embedding_port, self.recv_socket = get_zmq_socket_on_host( + self.zmq_context, zmq.PULL, host=host_name + ) + else: + self.embedding_port = embedding_port + self.recv_socket = None logger.info(f"Waiting for input {self.embedding_port = }") self.recv_embedding_data = None # ok=1 pending=0 fail=-1 @@ -835,71 +876,78 @@ class WaitingImageRequest: def _try_recv_mm_data(self): if self.status != WaitingImageRequestStatus.PENDING: return - while self.recv_embedding_data is None or not self.recv_embedding_data.ready: + if self.recv_socket is None: + return + while self.status == WaitingImageRequestStatus.PENDING: try: parts = self.recv_socket.recv_multipart(flags=zmq.NOBLOCK, copy=False) except zmq.Again: # No data available yet, wait a bit and retry return - try: - recv_obj: EmbeddingData = safe_pickle_loads(parts[0]) - if getattr(recv_obj, "error_msg", None) is not None: - logger.warning( - f"Received error signal from encoder for {self.rid}: {recv_obj.error_msg} {recv_obj.error_code = }" - ) - self.error_msg = recv_obj.error_msg - self.error_code = recv_obj.error_code - self.status = WaitingImageRequestStatus.FAIL - self.recv_socket.close() - return + self.consume_parts(parts) - # Extract original req_id from part_req_id and drop stale payloads - # that may arrive on a reused ZMQ port after a prior request aborted. - original_req_id = extract_original_req_id(recv_obj.req_id) - if original_req_id != self.recv_req.rid: - logger.warning( - f"Dropping stale embedding data: expected rid={self.recv_req.rid}, " - f"got rid={recv_obj.req_id} (likely from ZMQ port reuse)" - ) - continue - recv_obj.req_id = original_req_id + def consume_parts(self, parts): + if self.status != WaitingImageRequestStatus.PENDING: + return - buffer = parts[1].buffer if hasattr(parts[1], "buffer") else parts[1] - recv_obj.embedding = ( - torch.frombuffer(buffer, dtype=recv_obj.dtype) - .reshape(recv_obj.shape) - .clone() + try: + recv_obj: EmbeddingData = safe_pickle_loads(parts[0]) + if getattr(recv_obj, "error_msg", None) is not None: + logger.warning( + f"Received error signal from encoder for {self.rid}: {recv_obj.error_msg} {recv_obj.error_code = }" ) - - if self.recv_embedding_data is None: - self.recv_embedding_data = ( - MultiModalEmbeddingData.from_embedding_data( - recv_obj, model_type=self.model_type - ) - ) - else: - self.recv_embedding_data.add(recv_obj) - except Exception as e: - # A message the scheduler cannot decode (blocked unpickle, - # bad shape/dtype, ...) must fail this request, not crash the - # scheduler event loop; FAIL still reaches the TP-wide status - # all-reduce in _process_waiting_requests. - logger.exception( - "Failed to decode embedding message for rid=%s", self.rid - ) - self.error_msg = f"Failed to decode embedding message: {e}" + self.error_msg = recv_obj.error_msg + self.error_code = recv_obj.error_code self.status = WaitingImageRequestStatus.FAIL - self._cleanup_gpu_buffer() - self.recv_socket.close() + self.close_recv_socket() return + # Extract original req_id from part_req_id and drop stale payloads + # that may arrive on a reused ZMQ port after a prior request aborted. + original_req_id = extract_original_req_id(recv_obj.req_id) + if original_req_id != self.recv_req.rid: + logger.warning( + f"Dropping stale embedding data: expected rid={self.recv_req.rid}, " + f"got rid={recv_obj.req_id} (likely from ZMQ port reuse)" + ) + return + recv_obj.req_id = original_req_id + + buffer = parts[1].buffer if hasattr(parts[1], "buffer") else parts[1] + recv_obj.embedding = ( + torch.frombuffer(buffer, dtype=recv_obj.dtype) + .reshape(recv_obj.shape) + .clone() + ) + + if self.recv_embedding_data is None: + self.recv_embedding_data = MultiModalEmbeddingData.from_embedding_data( + recv_obj, model_type=self.model_type + ) + else: + self.recv_embedding_data.add(recv_obj) + except Exception as e: + # A message the scheduler cannot decode (blocked unpickle, + # bad shape/dtype, ...) must fail this request, not crash the + # scheduler event loop; FAIL still reaches the TP-wide status + # all-reduce in _process_waiting_requests. + logger.exception("Failed to decode embedding message for rid=%s", self.rid) + self.error_msg = f"Failed to decode embedding message: {e}" + self.status = WaitingImageRequestStatus.FAIL + self._cleanup_gpu_buffer() + self.close_recv_socket() + return + + if not self.recv_embedding_data.ready: + return + # Assemble mm_inputs. Wrapped so an assembly failure still reaches the # TP-wide status all-reduce in _process_waiting_requests instead of # raising past it. try: recv_embedding = self.recv_embedding_data.get_embedding(is_concat=True) mm_inputs = self.mm_processor.get_mm_data( - self.recv_req.input_text, + _select_mm_processor_prompt(self.recv_req, self.mm_processor), recv_embedding, **self.recv_embedding_data.get_mm_extra_meta(), ) @@ -913,7 +961,12 @@ class WaitingImageRequest: self.status = WaitingImageRequestStatus.FAIL self.error_msg = f"Failed to assemble multimodal inputs: {e}" self._cleanup_gpu_buffer() - self.recv_socket.close() + self.close_recv_socket() + + def close_recv_socket(self): + if self.recv_socket is not None: + self.recv_socket.close() + self.recv_socket = None def _cleanup_gpu_buffer(self): pass @@ -975,11 +1028,13 @@ class WaitingImageRDMARequest(WaitingImageRequest): encoder_urls, host_name, receive_count, + zmq_context, embeddings_engine, dtype, gpu_id=0, model_type: Optional[str] = None, embedding_pool=None, + embedding_port=None, ): super().__init__( rid=rid, @@ -989,6 +1044,8 @@ class WaitingImageRDMARequest(WaitingImageRequest): model_type=model_type, host_name=host_name, receive_count=receive_count, + zmq_context=zmq_context, + embedding_port=embedding_port, ) self.embeddings_engine = embeddings_engine self.dtype = dtype @@ -1256,7 +1313,7 @@ class WaitingImageRDMARequest(WaitingImageRequest): else: recv_embedding = self.recv_embedding_data.get_embedding(is_concat=True) mm_inputs = self.mm_processor.get_mm_data( - self.recv_req.input_text, + _select_mm_processor_prompt(self.recv_req, self.mm_processor), recv_embedding, **self.recv_embedding_data.get_mm_extra_meta(), ) @@ -1511,6 +1568,10 @@ class MMReceiverBase(ABC): encode_urls: Optional[List[str]] = None, ): self.context = zmq.asyncio.Context(20) + # Scheduler-side receive is polled synchronously. Keep one regular ZMQ + # context alive for the process instead of creating a temporary context + # whose destruction also closes its per-request socket. + self.scheduler_context = zmq.Context() self.encoder_transfer_backend = server_args.encoder_transfer_backend # When ``encode_urls`` is shared with an :class:`EncoderBootstrapServer` # (tokenizer manager process), it grows / shrinks in place as encoders @@ -1529,6 +1590,24 @@ class MMReceiverBase(ABC): self.nnodes = server_args.nnodes self.hostname = get_local_ip_auto() self.waiting_list: List[WaitingImageRequest] = [] + self.waiting_by_rid: Dict[str, WaitingImageRequest] = {} + self.scheduler_embedding_port = None + self.scheduler_recv_socket = None + if ( + self.encoder_transfer_backend == "zmq_to_scheduler" + and scheduler is not None + ): + ( + self.scheduler_embedding_port, + self.scheduler_recv_socket, + ) = get_zmq_socket_on_host( + self.scheduler_context, zmq.PULL, host=self.hostname + ) + logger.info( + "Scheduler TP rank %s reuses ZMQ embedding port %s", + self.tp_rank, + self.scheduler_embedding_port, + ) self.scheduler = scheduler self.gpu_id = scheduler.ps.gpu_id if scheduler is not None else 0 self.wait_timeout = envs.SGLANG_ENCODER_RECV_TIMEOUT.get() @@ -1864,6 +1943,28 @@ class MMReceiverBase(ABC): waiting_req.error_code = best_code # For zmq_to_scheduler + def _drain_scheduler_embeddings(self): + if self.scheduler_recv_socket is None: + return + + while True: + try: + parts = self.scheduler_recv_socket.recv_multipart( + flags=zmq.NOBLOCK, copy=False + ) + except zmq.Again: + return + + recv_obj: EmbeddingData = safe_pickle_loads(parts[0]) + rid = extract_original_req_id(recv_obj.req_id) + waiting_req = self.waiting_by_rid.get(rid) + if waiting_req is None: + logger.warning( + "Dropping embedding data for inactive request %s", recv_obj.req_id + ) + continue + waiting_req.consume_parts(parts) + def _process_waiting_requests(self, recv_reqs, waiting_cls, **extra_kwargs): new_recv_reqs = [] for recv_req in recv_reqs: @@ -1886,8 +1987,16 @@ class MMReceiverBase(ABC): model_type=self.model_type, host_name=self.hostname, receive_count=self.tp_size, + zmq_context=( + None + if self.scheduler_recv_socket is not None + else self.scheduler_context + ), + embedding_port=self.scheduler_embedding_port, **extra_kwargs, ) + if self.scheduler_recv_socket is not None: + self.waiting_by_rid[waiting_req.rid] = waiting_req waiting_req.send_encode_request() self.waiting_list.append(waiting_req) else: @@ -1896,14 +2005,16 @@ class MMReceiverBase(ABC): if len(self.waiting_list) == 0: return new_recv_reqs, [] + self._drain_scheduler_embeddings() current_time = time.time() local_status = [] for waiting_req in self.waiting_list: - waiting_req._try_recv_mm_data() + if self.scheduler_recv_socket is None: + waiting_req._try_recv_mm_data() if current_time - waiting_req.start_time > self.wait_timeout: waiting_req.status = WaitingImageRequestStatus.TIMEOUT waiting_req._cleanup_gpu_buffer() - waiting_req.recv_socket.close() + waiting_req.close_recv_socket() local_status.append(waiting_req.status) local_status = torch.tensor(local_status, device="cpu", dtype=torch.int32) @@ -1945,6 +2056,8 @@ class MMReceiverBase(ABC): ) else: # status_value == WaitingImageRequestStatus.PENDING new_waiting.append(waiting_req) + continue + self.waiting_by_rid.pop(waiting_req.rid, None) self.waiting_list = new_waiting return new_recv_reqs, abort_reqs diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 76c434db5..84732b837 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -55,6 +55,9 @@ from sglang.srt.managers.io_struct import ( ) from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache +from sglang.srt.model_executor.model_runner_components.load_model_utils import ( + maybe_precompile_model_kernels_after_loading, +) from sglang.srt.model_loader import get_model from sglang.srt.multimodal.processors.qwen_vl import preprocess_video from sglang.srt.observability.metrics_collector import EncoderMetricsCollector @@ -114,6 +117,7 @@ mooncake_send_done_count: Dict[str, int] = dict() use_image_processor_gpu = envs.SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU.get() ENCODER_MAX_BATCH_SIZE = envs.SGLANG_ENCODER_MAX_BATCH_SIZE.get() +ENCODER_MAX_BATCH_SIZE_EXPLICIT = envs.SGLANG_ENCODER_MAX_BATCH_SIZE.is_set() # Watchdog: max time to wait for a batched /encode result. Bounds HTTP latency # if the batch worker stalls (NCCL hang, dead worker proc, etc.). ENCODER_REQ_TIMEOUT = envs.SGLANG_ENCODER_REQ_TIMEOUT.get() @@ -186,7 +190,7 @@ def _convert(data): _mm_grid_attrs = { - # Kimi K2.5 HF processor uses grid_thws (see base_processor.ATTR_NAME_TO_MODALITY). + # Kimi K2.5/K3 HF processors use grid_thws (see base_processor.ATTR_NAME_TO_MODALITY). Modality.IMAGE: ["image_grid_thw", "image_grid_hws", "grid_thws"], Modality.VIDEO: ["video_grid_thw"], Modality.AUDIO: ["audio_feature_lens_raw"], @@ -200,14 +204,18 @@ _mm_feature_attrs = { def _get_mm_grid_dim(mm_inputs, modality, model_type: Optional[str] = None): + # Kimi K2.5/K3 vision processors only emit `grid_thws`; prefer it over generic keys + # so we never pick a mis-typed or stale `image_grid_hws` field from kwargs. attrs = _mm_grid_attrs[modality] model_type = (model_type or "").lower() if modality == Modality.IMAGE: - # Kimi K2.5 emits grid_thws, while Kimi-VL emits image_grid_hws. - if model_type == "kimi_k25": + # Kimi K2.5/K3 emit grid_thws, while Kimi-VL emits image_grid_hws. + # Other model types keep the generic attr order above. + if model_type in ("kimi_k25", "kimi_k3"): attrs = ("grid_thws", "image_grid_thw", "image_grid_hws") elif model_type == "kimi_vl": attrs = ("image_grid_hws", "image_grid_thw", "grid_thws") + for attr in attrs: if attr in mm_inputs and mm_inputs[attr] is not None: return _convert(mm_inputs[attr]) @@ -249,10 +257,29 @@ def _normalize_aux_value(val): def _build_mm_aux_data(mm_inputs, model_type=None): # Video aux metadata, scoped to model_type's video-meta attrs. - return { + aux = { attr: _normalize_aux_value(mm_inputs.get(attr)) for attr in video_meta_attrs_for(model_type) } + if model_type == "kimi_k3": + aux["original_image_sizes"] = _normalize_aux_value( + mm_inputs.get("original_image_sizes") + ) + return aux + + +def _get_original_image_size(image): + """Return an image's original (width, height) before encoder preprocessing.""" + if isinstance(image, dict): + image = image.get("image") + if isinstance(image, torch.Tensor): + if image.ndim < 2: + raise ValueError(f"Invalid image tensor shape: {tuple(image.shape)}") + return [int(image.shape[-1]), int(image.shape[-2])] + if hasattr(image, "size"): + width, height = image.size + return [int(width), int(height)] + raise TypeError(f"Cannot determine original image size from {type(image)}") class MMEncoder: @@ -320,9 +347,12 @@ class MMEncoder: load_config=self.load_config, device_config=self.device_config, ) + maybe_precompile_model_kernels_after_loading(self.model, self.device) self.context = zmq.asyncio.Context(2) self.sync_context = zmq.Context() # Reuse sync context for thread pool + self.scheduler_send_sockets = {} + self.scheduler_send_locks = {} self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=10) # Dedicated executor for image preprocessing (resize/normalize). # Separate from self.executor (ZMQ sends) to avoid contention under high concurrency. @@ -799,7 +829,7 @@ class MMEncoder: return self._get_feat_extract_output_lengths(input_length) else: if ( - self.model_type in ["kimi_k25", "kimi_vl"] + self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"] and modality == Modality.IMAGE ): return self._kimi_tokens_from_patch_grid(grid) @@ -855,6 +885,9 @@ class MMEncoder: """ if grid_thw is None: grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) + split_kimi_k3_images = ( + self.model_type == "kimi_k3" and modality == Modality.IMAGE + ) # Audio features are per-item (list of mels for mimo_v2, or batched # N x n_mels x T_max for qwen2_audio); slice by item index and keep @@ -874,31 +907,50 @@ class MMEncoder: offsets.append(curr) for idx in indices: sub_feature_list.append(mm_feature[offsets[idx] : offsets[idx + 1]]) - sub_feature = torch.cat(sub_feature_list, dim=0) + if not split_kimi_k3_images: + sub_feature = torch.cat(sub_feature_list, dim=0) - mm_item = MultimodalDataItem.from_dict( - { - "modality": modality, - "feature": ( - sub_feature - if isinstance(sub_feature, list) - else _convert(sub_feature) - ), - } - ) + if split_kimi_k3_images: + mm_items = [ + MultimodalDataItem.from_dict( + { + "modality": modality, + "feature": _convert(feature), + } + ) + for feature in sub_feature_list + ] + else: + mm_items = [ + MultimodalDataItem.from_dict( + { + "modality": modality, + "feature": ( + sub_feature + if isinstance(sub_feature, list) + else _convert(sub_feature) + ), + } + ) + ] for k, v in mm_inputs.items(): if k in _mm_feature_attrs.get(modality, []): continue val = _convert(v) if k in _mm_grid_attrs.get(modality, []): - mm_item.set(k, val[indices]) + if split_kimi_k3_images: + for mm_item, idx in zip(mm_items, indices): + mm_item.set(k, val[idx : idx + 1]) + else: + mm_items[0].set(k, val[indices]) else: - mm_item.set(k, val) + for mm_item in mm_items: + mm_item.set(k, val) forward_start = time.perf_counter() with torch.inference_mode(): - new_embeddings = get_feature_fn([mm_item]) + new_embeddings = get_feature_fn(mm_items) if not keep_on_gpu: new_embeddings = new_embeddings.cpu() if new_embeddings.ndim != 2: @@ -1473,12 +1525,20 @@ class MMEncoder: def _grid_count_per_leaf(self, leaves: List, modality: Modality) -> List[int]: """Number of grid entries each leaf produces under the model's processor. - Most processors map 1 leaf → 1 grid. Kimi-VL/K25 image processors expand + Most processors map 1 leaf → 1 grid. Kimi-VL/K2.5/K3 image processors expand a leaf shaped {"type": "image", "image": [pil1, pil2, ...]} into N grids (see _normalize_kimi_encoder_images). Cross-request batching needs these counts to keep per-request boundaries aligned with grid_dim. """ - if self.model_type not in ("kimi_k25", "kimi_vl") or modality != Modality.IMAGE: + if ( + self.model_type + not in ( + "kimi_k25", + "kimi_k3", + "kimi_vl", + ) + or modality != Modality.IMAGE + ): return [1] * len(leaves) def count(leaf): @@ -1527,7 +1587,7 @@ class MMEncoder: normalized.append(img) return normalized - # Kimi-K2.5 vision processor expects media dicts. + # Kimi-K2.5/K3 vision processors expect media dicts. normalized = [] for img in images: wrapped = wrap_one(img) @@ -1580,12 +1640,16 @@ class MMEncoder: if model_preprocessor: return model_preprocessor(images, Modality.IMAGE, self.vision_config) image_config = self.vision_config.get("image", {}) - if self.model_type in ["kimi_k25", "kimi_vl"]: + original_image_sizes = [_get_original_image_size(item) for item in images] + if self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]: images = self._normalize_kimi_encoder_images(images) - return await asyncio.get_running_loop().run_in_executor( + processor_input = await asyncio.get_running_loop().run_in_executor( self.preproc_executor, functools.partial(self.image_processor, images=images, **image_config), ) + if self.model_type == "kimi_k3": + processor_input["original_image_sizes"] = original_image_sizes + return processor_input async def _process_video_items(self, mm_items, model_preprocessor): if model_preprocessor: @@ -1893,20 +1957,79 @@ class MMEncoder: serialized_data = pickle.dumps(new_mm_data) buffer = embedding_tensor.__buffer__() - # Use thread pool executor for parallel ZMQ send operations + _zmq_xfer_start = time.perf_counter() + if ( + self.server_args.encoder_transfer_backend == "zmq_to_scheduler" + and url is not None + ): + lock = self.scheduler_send_locks.get(endpoint) + if lock is None: + lock = asyncio.Lock() + self.scheduler_send_locks[endpoint] = lock + + async with lock: + sock = self.scheduler_send_sockets.get(endpoint) + if sock is None: + sock = self.context.socket(zmq.PUSH) + config_socket(sock, zmq.PUSH) + sock.setsockopt(zmq.IMMEDIATE, 1) + sock.setsockopt(zmq.SNDTIMEO, int(self.send_timeout * 1000)) + sock.connect(endpoint) + self.scheduler_send_sockets[endpoint] = sock + try: + frames = ( + [serialized_data, buffer] + if buffer is not None + else [serialized_data] + ) + tracker = await sock.send_multipart(frames, copy=False, track=True) + except Exception: + if self.scheduler_send_sockets.get(endpoint) is sock: + self.scheduler_send_sockets.pop(endpoint, None) + sock.close(linger=0) + raise + + # MessageTracker.wait() protects the zero-copy source buffer; it + # is not a receiver acknowledgement. Waiting under the per-peer + # lock serialized every large embedding on that TCP connection. + # Queue sends in order under the lock, then wait for buffer + # ownership independently so libzmq can pipeline the connection. + try: + await asyncio.to_thread(tracker.wait, self.send_timeout) + except Exception: + if self.scheduler_send_sockets.get(endpoint) is sock: + self.scheduler_send_sockets.pop(endpoint, None) + sock.close(linger=0) + raise + + if encoder_metrics_collector is not None: + encoder_metrics_collector.observe_transfer( + time.perf_counter() - _zmq_xfer_start, + backend=self.server_args.encoder_transfer_backend, + ) + return + + # Per-request sockets remain for zmq_to_tokenizer and legacy direct + # scheduler sends. Scheduler URL sends use persistent sockets above. def send_with_socket(): sock = self.sync_context.socket(zmq.PUSH) config_socket(sock, zmq.PUSH) + sock.setsockopt(zmq.IMMEDIATE, 1) + sock.setsockopt(zmq.SNDTIMEO, int(self.send_timeout * 1000)) try: sock.connect(endpoint) if buffer is not None: - sock.send_multipart([serialized_data, buffer], copy=False) + tracker = sock.send_multipart( + [serialized_data, buffer], copy=False, track=True + ) else: - sock.send_multipart([serialized_data], copy=False) + tracker = sock.send_multipart( + [serialized_data], copy=False, track=True + ) + tracker.wait(timeout=self.send_timeout) finally: sock.close(linger=5000) - _zmq_xfer_start = time.perf_counter() await asyncio.get_event_loop().run_in_executor(self.executor, send_with_socket) if ( encoder_metrics_collector is not None @@ -2077,7 +2200,7 @@ class MMEncoder: try: mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality) grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) - aux_data = _build_mm_aux_data(mm_inputs) + aux_data = _build_mm_aux_data(mm_inputs, self.model_type) # Setup metadata and event management nbytes, total_tokens, embedding_dim, event = ( @@ -2178,7 +2301,7 @@ class MMEncoder: """Cross-request encoder fusion (image/audio). No cache path.""" # items_per_req counts grid entries (post-expansion) so per-request # slicing of grid_dim/final_slices stays aligned for processors that - # expand one leaf into multiple grids (e.g. Kimi-VL/K25 dict-of-images). + # expand one leaf into multiple grids (e.g. Kimi-VL/K2.5/K3 dict-of-images). flat_items, items_per_req = [], [] for req in requests: leaves = MMEncoder._flatten_nested_items(req["mm_items"]) @@ -2235,14 +2358,17 @@ class MMEncoder: if self.profiler is not None: for _ in requests: self.profiler.step() - # No aux_data here: batch_encode only handles IMAGE/AUDIO - # (_BATCHABLE_MODALITIES), and _build_mm_aux_data only extracts - # video-meta fields — which never appear in image/audio mm_inputs. + aux_data = _build_mm_aux_data(mm_inputs, self.model_type) results = [] offset = 0 for req, n in zip(requests, items_per_req): slices = final_slices[offset : offset + n] emb = slices[0] if n == 1 else torch.cat(slices, dim=0) + req_aux_data = {} + if aux_data.get("original_image_sizes") is not None: + req_aux_data["original_image_sizes"] = aux_data[ + "original_image_sizes" + ][offset : offset + n] if self.rank == 0: self.embedding_to_send[req["req_id"]] = EmbeddingData( req["req_id"], @@ -2251,6 +2377,7 @@ class MMEncoder: grid_dim[offset : offset + n], modality, emb, + **req_aux_data, ) results.append((emb.nbytes, emb.shape[0], emb.shape[1], None, None)) offset += n @@ -2462,6 +2589,20 @@ class PendingRequest: # VIDEO excluded: per-video preprocess kwargs (do_sample_frames, video_metadata) # vary per request and can't merge into one HF processor call. _BATCHABLE_MODALITIES = {Modality.IMAGE, Modality.AUDIO} +_KIMI_K3_DEFAULT_ENCODER_MAX_BATCH_SIZE = 2 + + +def _resolve_encoder_batch_policy( + model_type: str, + configured_max_batch_size: int, + max_batch_size_is_explicit: bool, +) -> Tuple[int, bool]: + """Return effective batch size and same-turn coalescing policy.""" + max_batch_size = max(1, int(configured_max_batch_size)) + coalesce_same_turn = model_type == "kimi_k3" + if coalesce_same_turn and not max_batch_size_is_explicit: + max_batch_size = min(max_batch_size, _KIMI_K3_DEFAULT_ENCODER_MAX_BATCH_SIZE) + return max_batch_size, coalesce_same_turn class EncoderScheduler: @@ -2472,11 +2613,13 @@ class EncoderScheduler: encoder: "MMEncoder", send_sockets: List[zmq.Socket], max_batch_size: int, + coalesce_same_turn: bool = False, request_timeout: float = ENCODER_REQ_TIMEOUT, ): self.encoder = encoder self.send_sockets = send_sockets self.max_batch_size = max(1, int(max_batch_size)) + self.coalesce_same_turn = bool(coalesce_same_turn) self.request_timeout = max(1.0, float(request_timeout)) self.pending_queue: asyncio.Queue[PendingRequest] = asyncio.Queue() self._worker_task: Optional[asyncio.Task] = None @@ -2485,7 +2628,9 @@ class EncoderScheduler: if self._worker_task is None: self._worker_task = asyncio.create_task(self._batch_worker()) logger.info( - f"EncoderScheduler started with max_batch_size={self.max_batch_size}" + "EncoderScheduler started with " + f"max_batch_size={self.max_batch_size}, " + f"coalesce_same_turn={self.coalesce_same_turn}" ) async def stop(self) -> None: @@ -2520,6 +2665,17 @@ class EncoderScheduler: async def _collect_batch(self) -> List[PendingRequest]: batch = [await self.pending_queue.get()] + first_modality = Modality.from_str(batch[0].request.get("modality", "image")) + should_yield = ( + self.coalesce_same_turn + and self.max_batch_size > 1 + and first_modality in _BATCHABLE_MODALITIES + ) + if should_yield: + # Let HTTP handlers that arrived in the same event-loop turn enqueue + # before dispatch. Unlike a fixed sleep, this adds no millisecond-scale + # tax to an isolated request. + await asyncio.sleep(0) while len(batch) < self.max_batch_size: try: batch.append(self.pending_queue.get_nowait()) @@ -2603,23 +2759,29 @@ class EncoderScheduler: encoder_metrics_collector.observe_queue_wait( max(0.0, start - p.submit_time), modality=modality_str ) - for sock in self.send_sockets: - sock_send( - sock, - wrap_as_pickle( - { - "type": "batch_encode", - "modality": modality.name, - "requests": requests, - "enter_time": start, - } - ), - ) - - logger.info(f"Dispatching batch of {len(group)} {modality.name} requests") - try: - results = await self.encoder.batch_encode(requests, modality) + # The scheduler is the sole owner of batched dispatch order. Keep + # the collective broadcast and rank-0 execution under the same + # lock, while allowing concurrent HTTP handlers to enqueue before + # waiting on their individual futures. + async with self.encoder.encode_dispatch_lock: + for sock in self.send_sockets: + sock_send( + sock, + wrap_as_pickle( + { + "type": "batch_encode", + "modality": modality.name, + "requests": requests, + "enter_time": start, + } + ), + ) + + logger.info( + f"Dispatching batch of {len(group)} {modality.name} requests" + ) + results = await self.encoder.batch_encode(requests, modality) if len(group) > 1: logger.info( f"Batch of {len(group)} {modality.name} requests completed in " @@ -3376,8 +3538,16 @@ async def run_dp_worker( encoder_metrics_collector = EncoderMetricsCollector(labels) enc.dp_rank = dp_rank + max_batch_size, coalesce_same_turn = _resolve_encoder_batch_policy( + enc.model_type, + ENCODER_MAX_BATCH_SIZE, + ENCODER_MAX_BATCH_SIZE_EXPLICIT, + ) sched = EncoderScheduler( - encoder=enc, send_sockets=[], max_batch_size=ENCODER_MAX_BATCH_SIZE + encoder=enc, + send_sockets=[], + max_batch_size=max_batch_size, + coalesce_same_turn=coalesce_same_turn, ) ctx = zmq.asyncio.Context(2) @@ -3386,12 +3556,12 @@ async def run_dp_worker( send_lock = asyncio.Lock() inflight: Set[asyncio.Task] = set() # Acquire-before-recv → back-pressure propagates to the dispatcher - # PUSH buffer. Must be ≥ ENCODER_MAX_BATCH_SIZE or batching degrades. + # PUSH buffer. Must be at least max_batch_size or batching degrades. max_inflight = envs.SGLANG_ENCODER_DP_WORKER_MAX_INFLIGHT.get() - if max_inflight < ENCODER_MAX_BATCH_SIZE: + if max_inflight < max_batch_size: logger.warning( f"SGLANG_ENCODER_DP_WORKER_MAX_INFLIGHT={max_inflight} is below " - f"ENCODER_MAX_BATCH_SIZE={ENCODER_MAX_BATCH_SIZE}; the encoder " + f"the effective encoder max_batch_size={max_batch_size}; the encoder " f"will never assemble a full batch." ) inflight_sem = asyncio.Semaphore(max_inflight) @@ -3472,8 +3642,16 @@ async def _lifespan(app: FastAPI): yield return if encoder is not None: + max_batch_size, coalesce_same_turn = _resolve_encoder_batch_policy( + encoder.model_type, + ENCODER_MAX_BATCH_SIZE, + ENCODER_MAX_BATCH_SIZE_EXPLICIT, + ) encoder_scheduler = EncoderScheduler( - encoder, send_sockets, max_batch_size=ENCODER_MAX_BATCH_SIZE + encoder, + send_sockets, + max_batch_size=max_batch_size, + coalesce_same_turn=coalesce_same_turn, ) encoder_scheduler.start() try: @@ -3937,8 +4115,10 @@ async def handle_encode_request(request: dict): }, ) else: - # Lock direct dispatch together with rank0 await so its NCCL launch - # order matches the ZMQ dispatch order rank>0 sees. + # Non-batched requests still own their collective dispatch order + # directly; batched requests take this lock in _dispatch_group. + # Locking direct dispatch together with the rank0 await keeps its + # NCCL launch order matching the ZMQ dispatch order rank>0 sees. async with encoder.encode_dispatch_lock: for socket in send_sockets: sock_send(socket, wrap_as_pickle(request)) @@ -4157,12 +4337,6 @@ async def health_generate(): if encoder is None: return Response(status_code=503) - # Skip the dummy encode when real requests are already in flight — the - # ongoing traffic already proves liveness, matching the scheduler's - # `is_fully_idle`-based health-check skip pattern. - if encoder.embedding_to_send: - return Response(status_code=200) - # Pick the first available modality for the dummy encode if encoder.image_processor is not None: mm_items = [f"data:image/png;base64,{MINIMUM_PNG_PICTURE_BASE64}"] @@ -4186,21 +4360,26 @@ async def health_generate(): "part_idx": 0, } - # Broadcast to other TP ranks so distributed ops stay in sync - for socket in send_sockets: - sock_send(socket, wrap_as_pickle(dummy_request)) + # A health encode participates in the same TP collectives as a real + # request. Serialize its broadcast and rank-0 forward with every other + # collective dispatch, then recheck whether traffic made the probe + # unnecessary while it waited for the lock. + async with encoder.encode_dispatch_lock: + if encoder.embedding_to_send: + return Response(status_code=200) + for socket in send_sockets: + sock_send(socket, wrap_as_pickle(dummy_request)) - # Run encode on rank 0 with timeout - _, _, _, error_msg, _ = await asyncio.wait_for( - encoder.encode( - mm_items=mm_items, - modality=modality, - req_id=req_id, - num_parts=1, - part_idx=0, - ), - timeout=HEALTH_CHECK_TIMEOUT, - ) + _, _, _, error_msg, _ = await asyncio.wait_for( + encoder.encode( + mm_items=mm_items, + modality=modality, + req_id=req_id, + num_parts=1, + part_idx=0, + ), + timeout=HEALTH_CHECK_TIMEOUT, + ) # Clean up stored embedding encoder.embedding_to_send.pop(req_id, None) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 782661658..1907e0ee6 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -404,10 +404,8 @@ class MooncakeKVManager(CommonKVManager): def _send_chunk_ready(self, req, chunk_idx, kv_chunk, prefill_unique_rank): """Notify decode that a non-last staging chunk RDMA is complete.""" na = NetworkAddress(req.endpoint, req.dst_port) - self._connect( + self._send_multipart_locked( na.to_tcp(), - is_ipv6=na.is_ipv6, - ).send_multipart( [ b"CHUNK_READY", str(req.room).encode("ascii"), @@ -416,7 +414,8 @@ class MooncakeKVManager(CommonKVManager): str(len(kv_chunk.prefill_kv_indices)).encode("ascii"), req.mooncake_session_id.encode("ascii"), str(prefill_unique_rank).encode("ascii"), - ] + ], + is_ipv6=na.is_ipv6, ) def _do_staging_transfer( @@ -1091,9 +1090,8 @@ class MooncakeKVManager(CommonKVManager): data: bytes, ): na = NetworkAddress(remote, dst_port) - socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6) - - socket.send_multipart( + self._send_multipart_locked( + na.to_tcp(), [ MooncakeKVManager.AUX_DATA_HEADER, str(room).encode("ascii"), @@ -1101,7 +1099,8 @@ class MooncakeKVManager(CommonKVManager): str(aux_index).encode("ascii"), struct.pack(">I", len(data)), data, - ] + ], + is_ipv6=na.is_ipv6, ) def _handle_aux_data(self, msg: List[bytes]): @@ -1502,12 +1501,14 @@ class MooncakeKVManager(CommonKVManager): self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int ): na = NetworkAddress(remote, dst_port) - self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart( + self._send_multipart_locked( + na.to_tcp(), [ str(room).encode("ascii"), str(status).encode("ascii"), str(prefill_rank).encode("ascii"), - ] + ], + is_ipv6=na.is_ipv6, ) def transfer_worker( @@ -1736,12 +1737,34 @@ class MooncakeKVManager(CommonKVManager): if kv_chunk.is_last_chunk: if kv_chunk.state_indices and not skip_state: - self.maybe_send_extra( + state_rc = self.maybe_send_extra( req, kv_chunk.state_indices, executor, target_rank_registration_info, ) + if state_rc != 0: + with self.session_lock: + self.session_failures[ + req.mooncake_session_id + ] += 1 + self.failed_sessions.add( + req.mooncake_session_id + ) + self.record_failure( + kv_chunk.room, + f"Failed to send state components of {kv_chunk.room} to " + f"{NetworkAddress(req.endpoint, req.dst_port).to_host_port_str()}", + ) + self.update_status(kv_chunk.room, KVPoll.Failed) + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + KVPoll.Failed, + prefill_unique_rank, + ) + break # Only the last chunk we need to send the aux data ret = self.send_aux( @@ -1856,11 +1879,13 @@ class MooncakeKVManager(CommonKVManager): # Send ACK back to decode endpoint try: na = NetworkAddress(decode_ip, decode_port) - self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart( + self._send_multipart_locked( + na.to_tcp(), [ b"ABORT_ACK", str(room_to_be_aborted).encode("ascii"), - ] + ], + is_ipv6=na.is_ipv6, ) logger.debug( f"Sent ABORT_ACK for room {room_to_be_aborted} to " diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 7935e9a7b..a519d65cf 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -58,6 +58,10 @@ _SEMAPHORE_BYTES = 128 _MAX_GRAPH_INPUTS = 131072 # resolved once at import time; explicit constructor sizes take precedence _DEFAULT_MAX_SIZE = envs.SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB.get() * 1024 +# forced per-direction workspace sizes; highest priority, override both the +# tuned config and explicit constructor sizes (None = not forced) +_FORCE_PULL_SIZE_KB = envs.SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB.get() +_FORCE_PUSH_SIZE_KB = envs.SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB.get() class _PullMode(enum.Enum): @@ -109,6 +113,10 @@ class CustomAllReduceV2: the tuned size and ``max_size``. :param max_push_size: explicit per-buffer push workspace size; overrides both the tuned size and ``max_size``. + + ``SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB`` / + ``SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB`` take the highest + priority and override all of the size parameters above. """ self.disabled = True if not can_use_custom_all_reduce_v2(group=group, device=device): @@ -123,6 +131,25 @@ class CustomAllReduceV2: max_pull_size = min(base_config.max_pull_bytes, max_size) if max_push_size is None: max_push_size = min(base_config.max_push_bytes, max_size) + if _FORCE_PULL_SIZE_KB is not None: + max_pull_size = int(_FORCE_PULL_SIZE_KB) * 1024 + if _FORCE_PUSH_SIZE_KB is not None: + max_push_size = int(_FORCE_PUSH_SIZE_KB) * 1024 + + def force_thresholds(heuristic): + # forced sizes bypass the tuned NCCL-crossover heuristics: lift + # each direction's ceiling to the forced workspace capacity (the + # clip() below only ever lowers thresholds) + if _FORCE_PULL_SIZE_KB is not None: + heuristic = heuristic._replace(two_shot_pull_threshold=max_pull_size) + if _FORCE_PUSH_SIZE_KB is not None: + heuristic = heuristic._replace(one_shot_push_threshold=max_push_size) + return heuristic + + base_config = base_config._replace( + graph=force_thresholds(base_config.graph), + eager=force_thresholds(base_config.eager), + ) # a minimal workspace keeps the Communicator valid even when a caller # only uses one direction (e.g. push-only fused qk-norm instances) self.max_pull_size = _ceil_align(max(max_pull_size, _ALIGN_BYTES), _ALIGN_BYTES) @@ -138,8 +165,12 @@ class CustomAllReduceV2: max_push_bytes=self.max_push_size, max_pull_bytes=self.max_pull_size )._replace(num_pull_blocks=num_pull_blocks, num_push_blocks=num_push_blocks) self.override_algo: Optional[AllReduceAlgo] = None - self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() - + # On a multi-node (MNNVL) group the symm-mem workspace plane works + # across nodes, but graph zero-copy input registration (cudaIpc / + # node-local VMM remap) does not — force eager pull inside graphs. + is_multinode = not all(in_the_same_node_as(group, source_rank=0)) + tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() + self._is_graph_mode_supported = not tms_cudagraph and not is_multinode # device-side pointer table: one row of world_size pointers per # graph-captured all-reduce input self.graph_params = torch.zeros( @@ -342,7 +373,7 @@ class CustomAllReduceV2: yield return try: - self._graph_mode_allowed = not self.tms_cudagraph + self._graph_mode_allowed = self._is_graph_mode_supported yield finally: self._graph_mode_allowed = False @@ -426,6 +457,22 @@ def can_use_custom_all_reduce_v2( group: ProcessGroup, device: torch.device, ) -> bool: + # Multi-node (MNNVL): the node-local NVLink/P2P topology checks below + # are meaningless across nodes; the torch symm-mem rendezvous (fabric + # handles) is the real capability gate there. + if envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get() and not all( + in_the_same_node_as(group, source_rank=0) + ): + world_size = dist.get_world_size(group=group) + if world_size in range(2, 9): + logger.warning( + "CustomAllReduceV2 enabled on a multi-node group " + "(world_size=%d); graph zero-copy is disabled.", + world_size, + ) + return True + return False + supported = list(range(2, 17)) if dist.get_world_size(group=group) not in supported: return False diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index c90dc0a96..6fb820082 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -695,10 +695,22 @@ class GroupCoordinator: self.pymscclpp_comm is not None and self.pymscclpp_comm.should_mscclpp_allreduce(input_) ) + # With the MNNVL opt-in, let CustomAllReduceV2 take eligible (small) + # inputs ahead of the symm-mem pynccl fast path; otherwise pynccl + # would absorb every all-reduce whenever --enable-symm-mem is on and + # v2 never runs. Large inputs fail should_custom_ar and still go to + # the symm-mem path below. + _ca_takes_input = ( + _CA_V2_MULTINODE + and self.ca_comm is not None + and not self.ca_comm.disabled + and self.ca_comm.should_custom_ar(input_) + ) if ( self.pynccl_comm is not None and self.is_symmetric_memory_enabled() and not should_use_pymscclpp_allreduce + and not _ca_takes_input ): self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce") with self.pynccl_comm.change_state(enable=True): @@ -1984,6 +1996,9 @@ logger = logging.getLogger(__name__) _ENABLE_CUSTOM_ALL_REDUCE = True _ENABLE_MSCCLPP_ALL_REDUCE = False _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = False +# Read once at import: whether CustomAllReduceV2 is opted in on a multi-node +# (MNNVL) group. Used on the all_reduce hot path (see GroupCoordinator). +_CA_V2_MULTINODE = envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get() def set_custom_all_reduce(enable: bool): diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index c55ea80ff..6eb0ad1af 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -107,6 +107,7 @@ from sglang.srt.utils import ( configure_logger, get_bool_env_var, is_cuda, + is_mnnvl_fabric_device, kill_process_tree, launch_dummy_health_check_server, maybe_reindex_device_id, @@ -1572,6 +1573,12 @@ class Engine(EngineScoreMixin, EngineBase): def _set_envs_and_config(server_args: ServerArgs): # Set global environments + # MNNVL fabric (GB200/GB300) multi-node: cross-node NVLink needs NCCL's + # cuMem-based buffers and MNNVL transport. Default them on (user-set + # values win; the symm-mem override below only fires when unset). + if server_args.nnodes > 1 and is_mnnvl_fabric_device(): + os.environ.setdefault("NCCL_CUMEM_ENABLE", "1") + os.environ.setdefault("NCCL_MNNVL_ENABLE", "1") if "NCCL_CUMEM_ENABLE" not in os.environ or server_args.enable_symm_mem: os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem)) if ( diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index cb0256d0d..ad1e56cb8 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -194,6 +194,7 @@ class OpenAIServingChat(OpenAIServingBase): """Handler for /v1/chat/completions requests""" _default_sampling_params_logged = False + _KIMI_K3_GENERATION_STUB_TOKENS = 3 def __init__( self, @@ -650,6 +651,14 @@ class OpenAIServingChat(OpenAIServingBase): content["meta_info"].get("cached_tokens", 0) ) + def _reported_prompt_tokens(self, meta_info: Dict[str, Any]) -> int: + prompt_tokens = meta_info.get("prompt_tokens", 0) + if self.chat_encoding_spec == "kimi_k3": + # K3's three-token assistant generation stub is model input, but the + # reference API excludes it from billed/reported prompt tokens. + prompt_tokens = max(0, prompt_tokens - self._KIMI_K3_GENERATION_STUB_TOKENS) + return prompt_tokens + async def _generate_stream_content( self, content: Dict[str, Any], @@ -1493,7 +1502,9 @@ class OpenAIServingChat(OpenAIServingBase): ): index = content.get("index", 0) - prompt_tokens[index] = content["meta_info"].get("prompt_tokens", 0) + prompt_tokens[index] = self._reported_prompt_tokens( + content["meta_info"] + ) completion_tokens[index] = content["meta_info"].get( "completion_tokens", 0 ) @@ -1720,6 +1731,20 @@ class OpenAIServingChat(OpenAIServingBase): created: int, ) -> Union[ChatCompletionResponse, ORJSONResponse]: """Build chat completion response from generation results""" + if self.chat_encoding_spec == "kimi_k3": + ret = [ + { + **item, + "meta_info": { + **item["meta_info"], + "prompt_tokens": self._reported_prompt_tokens( + item["meta_info"] + ), + }, + } + for item in ret + ] + choices = [] # Build sglext at response level (from first ret_item, as these are per-request) @@ -2378,7 +2403,7 @@ class OpenAIServingChat(OpenAIServingBase): # Add usage stats if continuous_usage_stats is enabled if continuous_usage_stats: - prompt_tokens = content["meta_info"].get("prompt_tokens", 0) + prompt_tokens = self._reported_prompt_tokens(content["meta_info"]) completion_tokens = content["meta_info"].get("completion_tokens", 0) reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0) chunk.usage = UsageProcessor.calculate_token_usage( @@ -2431,7 +2456,7 @@ class OpenAIServingChat(OpenAIServingBase): # Add usage stats if continuous_usage_stats is enabled if continuous_usage_stats: - prompt_tokens = content["meta_info"].get("prompt_tokens", 0) + prompt_tokens = self._reported_prompt_tokens(content["meta_info"]) completion_tokens = content["meta_info"].get("completion_tokens", 0) reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0) chunk.usage = UsageProcessor.calculate_token_usage( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 30c796851..820d0b362 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -808,6 +808,39 @@ class Envs: # Triton two_dot variant, 1.16-1.38x faster across GLM/DS shapes). SGLANG_OPT_Q8KV8_QPREP_VARIANT = EnvStr("auto") + # TRT-LLM-gen fused MoE (SiTU) via sglang JIT: path to an unpacked SiTU + # cubin pool (cubins + flat ABI headers + overlay/; distributed as a + # single downloadable archive). Needs the public flashinfer package + # installed for the unmodified JIT sources. Unset = feature off. + SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL = EnvStr(None) + + # MNNVL fused all-reduce (bf16, TP8): zero-copy 1shot multicast-push for + # small messages and in-place NVLS 2shot on symmetric-memory tensors for + # large ones, with an optional fused residual add. Covers the KDA o_proj + # output and the latent|shared MoE reduce; everything else falls back to + # the regular all-reduce path. Auto-enabled on SM100/SM103 when + # CustomAllReduceV2 with multicast is available; set 0/1 to override in + # either direction. See srt/layers/k3_ar_fusion.py. + SGLANG_K3_AR_FUSION = EnvBool(False) + # K3 SP-MoE fused residual + reduce-scatter and matching all-gather over + # CustomAllReduceV2's MNNVL push workspace. Auto-probed for the validated + # TP8 GB300 configuration; set 0/1 to override. See + # srt/layers/k3_sp_collective.py. + SGLANG_K3_SP_COLLECTIVE = EnvBool(False) + # Keep K3's post-MoE residual stream token-sharded between consecutive + # SP-MoE layers. The next attention-residual aggregation and snapshot + # bank write run on the local shard, then only the normalized attention + # input is all-gathered. Requires SGLANG_K3_SP_COLLECTIVE. + SGLANG_K3_SP_ATTN_RES = EnvBool(False) + # Fused o_proj GEMM + all-reduce (bf16, TP 2..8, SM100+): one + # kernel computes the TP-local o_proj partial and the cross-rank sum over + # a P2P comm region, replacing the GEMM + NCCL AR pair at M <= 512. + SGLANG_K3_GEMM_AR = EnvBool(False) + # Merge the router gate and routed_expert_down_proj weights so the K3 MoE + # front reads hidden_states once, and run the top-k plus the bf16 cast in one + # epilogue kernel. See kernels/ops/moe/moe_front.py. Default on. + SGLANG_K3_FUSED_FRONT = EnvBool(True) + # sgl-kernel SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False) @@ -848,6 +881,16 @@ class Envs: # Default per-direction workspace cap for CustomAllReduceV2; explicit # constructor sizes take precedence over this. SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB = EnvInt(16 * 1024) + SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB = EnvInt(None) + SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB = EnvInt(None) + # Allow CustomAllReduceV2 on a process group that spans nodes (MNNVL + # fabric). Requires torch symmetric memory to rendezvous across nodes + # (fabric handles + IMEX). Graph zero-copy input registration is not + # supported in this mode and is disabled; all-reduce inside CUDA graphs + # falls back to eager pull from the symm workspace. Auto-enabled on + # MNNVL-fabric devices (GB200/GB300) when nnodes > 1; set 0/1 to + # override in either direction. + SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE = EnvBool(False) SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096) SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048) SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096) @@ -903,6 +946,9 @@ class Envs: SGLANG_MM_BUFFER_SIZE_MB = EnvInt(0) SGLANG_MM_PRECOMPUTE_HASH = EnvBool(False) SGLANG_VIT_ENABLE_CUDA_GRAPH = EnvBool(False) + SGLANG_KIMI_K3_VIT_CUDA_GRAPH_CACHE_CAPACITY = EnvInt(2) + SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MIN_HITS = EnvInt(2) + SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MAX_SEQLEN = EnvInt(6144) # Use the fully-vectorized ViT position-embedding interpolation (no per-image # Python loop / CPU<->GPU sync). Bit-exact with the legacy implementation; # set False to fall back to the per-image loop. @@ -914,7 +960,9 @@ class Envs: # VLM Item CUDA IPC Transport SGLANG_USE_CUDA_IPC_TRANSPORT = EnvBool(False) - SGLANG_USE_IPC_POOL_HANDLE_CACHE = EnvBool(False) + # Reuse the mapping for the already-allocated bounded CUDA IPC pool. This + # has no effect unless CUDA IPC feature transport is explicitly selected. + SGLANG_USE_IPC_POOL_HANDLE_CACHE = EnvBool(True) SGLANG_MM_FEATURE_CACHE_MB = EnvInt(1 * 1024) SGLANG_MM_ITEM_MEM_POOL_RECYCLE_INTERVAL_SEC = EnvFloat(0.05) @@ -929,7 +977,6 @@ class Envs: # mamba pool ratio accordingly. Frees one resident slot per running request, # raising max_running_requests. Off = original locking + ratio (escape hatch). SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK = EnvBool(False) - # Unified Radix Tree SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False) # Registered TreeCore backend serving the unified radix cache. diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py index ab30b9b5e..083718ef9 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py @@ -85,6 +85,7 @@ def fused_sigmoid_gating_delta_rule_update( dt_bias=dt_bias, softplus_beta=softplus_beta, softplus_threshold=softplus_threshold, + lower_bound=0.0, q=q, k=k, v=v, @@ -118,6 +119,7 @@ def fused_sigmoid_gating_delta_rule_update( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_VARLEN=cu_seqlens is not None, IS_KDA=is_kda, + USE_LOWER_BOUND=False, DISABLE_STATE_UPDATE=disable_state_update, CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None, HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None, diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 8062eb465..837b766e7 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -181,6 +181,37 @@ class SiluAndMul(MultiPlatformOp): return self._musa_swish_glu(x) +class SituAndMul(MultiPlatformOp): + """SituGLU activation used by Kimi K3. + + Computes beta * tanh(gate / beta) * sigmoid(gate) * up. + When linear_beta is set, up is softly clipped: + up = linear_beta * tanh(up / linear_beta). + """ + + def __init__(self, beta: float = 1.0, linear_beta: float | None = None): + super().__init__() + self.beta = float(beta) + self.linear_beta = None if linear_beta is None else float(linear_beta) + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].float() + up = x[..., d:].float() + gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (gate * up).to(x.dtype) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + from sglang.kernels.ops.kimi_k3.activation import situ_and_mul + + return situ_and_mul(x, None, self.beta, self.linear_beta) + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + class GeluAndMul(MultiPlatformOp): def __init__(self, approximate="tanh"): super().__init__() diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 39e9aa51d..9af985d73 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -70,6 +70,18 @@ def create_flashinfer_backend(runner): def create_trtllm_mla_backend(runner): if not runner.use_mla_backend: raise ValueError("trtllm_mla backend can only be used with MLA models.") + if ( + runner.server_args.dcp_size > 1 + and runner.server_args.speculative_algorithm is not None + ): + _, decode_backend = runner.server_args.get_attention_backends() + if decode_backend == "trtllm_mla": + raise ValueError( + "trtllm_mla cannot serve decode context parallelism with speculative " + "decoding: it does not forward the cyclic DCP metadata to its decode " + "kernel and returns no rank-local LSE for the cross-rank merge. " + "Select cutedsl_mla or tokenspeed_mla." + ) from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend return TRTLLMMLABackend(runner) diff --git a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py index 772b69465..f9207b4fe 100644 --- a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py @@ -178,15 +178,18 @@ class CuteDslMLABackend(TRTLLMMLABackend): bs, num_tokens, forward_mode, seq_lens, device ) if get_parallel().dcp_enabled: - if forward_mode.is_target_verify(): - self.forward_decode_metadata.global_seq_lens_k = torch.zeros_like( - self.forward_decode_metadata.seq_lens_k - ) - self.forward_decode_metadata.max_seq_len_k = ( - self._get_dcp_local_max_seq_len( - self.max_context_len - + (self.num_draft_tokens if forward_mode.is_target_verify() else 0) + metadata = self.forward_decode_metadata + if metadata.global_seq_lens_k is None: + # Plain decode under DCP also keeps the int32 GLOBAL lens in a + # capture-stable buffer (super allocates it only for verify): + # the DCP kernel consumes both the rank-local and the global + # lens every MLA layer, so both are maintained once per step. + metadata.global_seq_lens_k = torch.zeros( + (bs,), dtype=torch.int32, device=device ) + metadata.max_seq_len_k = self._get_dcp_local_max_seq_len( + self.max_context_len + + (self.num_draft_tokens if forward_mode.is_target_verify() else 0) ) def _apply_cuda_graph_metadata( @@ -224,7 +227,13 @@ class CuteDslMLABackend(TRTLLMMLABackend): local_seq_lens = self._get_dcp_local_seq_lens(seq_lens) else: seq_lens = seq_lens[:bs] - local_seq_lens = self._get_dcp_local_seq_lens(seq_lens) + # Hoist: refresh the int32 global + rank-local lens once per step + # into the capture-stable buffers; forward_decode reads them + # instead of recomputing get_dcp_lens + two int32 casts per MLA + # layer. + metadata.global_seq_lens_k.copy_(seq_lens) + metadata.seq_lens_k.copy_(self._get_dcp_local_seq_lens(seq_lens)) + local_seq_lens = metadata.seq_lens_k self._fill_dcp_block_kv_indices( metadata.block_kv_indices, @@ -249,6 +258,19 @@ class CuteDslMLABackend(TRTLLMMLABackend): metadata.seq_lens_k = self._get_dcp_local_seq_lens( metadata.global_seq_lens_k ) + elif ( + forward_batch.forward_mode.is_decode_or_idle() + and self.forward_decode_metadata.seq_lens_k is not None + ): + # Same hoist as verify: the parent stored the int32 GLOBAL + # lens in seq_lens_k; keep it as global_seq_lens_k and derive + # the rank-local view once per step (forward_decode consumes + # both every MLA layer). + metadata = self.forward_decode_metadata + metadata.global_seq_lens_k = metadata.seq_lens_k + metadata.seq_lens_k = self._get_dcp_local_seq_lens( + metadata.global_seq_lens_k + ) self.forward_decode_metadata.max_seq_len_k = ( self._get_dcp_local_max_seq_len( self.forward_decode_metadata.max_seq_len_k @@ -349,12 +371,30 @@ class CuteDslMLABackend(TRTLLMMLABackend): # Query / KV preparation mirrors the base cute-dsl decode (both FP16 and # FP8 KV), then swaps to the DCP kernel call + rank-local return. merge_query = q_rope is not None + query = None if self.data_type == torch.float8_e4m3fn: assert q_rope is not None and k_rope is not None if cos_sin_cache is None: - q, k, k_rope = mla_quantize_without_rope_for_fp8( - q, q_rope, k.squeeze(1), k_rope.squeeze(1) - ) + if ( + save_kv_cache + and self._fused_set_kv_concat_q_fp8 + and not self._unified_mla + ): + # Static pool: out_cache_loc is already the physical loc. + # Fused: bf16->fp8 quantize + KV scatter + q concat in one + # launch; None when not covered. + query = self._set_kv_and_concat_q_fp8_fused( + layer=layer, + loc=forward_batch.out_cache_loc, + q=q, + q_rope=q_rope, + k=k, + k_rope=k_rope, + ) + if query is None: + q, k, k_rope = mla_quantize_without_rope_for_fp8( + q, q_rope, k.squeeze(1), k_rope.squeeze(1) + ) else: q, k, k_rope = mla_quantize_and_rope_for_fp8( q, @@ -369,13 +409,15 @@ class CuteDslMLABackend(TRTLLMMLABackend): ) merge_query = False - if save_kv_cache: + if query is None and save_kv_cache: assert k is not None and k_rope is not None self.token_to_kv_pool.set_mla_kv_buffer( layer, forward_batch.out_cache_loc, k, k_rope ) - if merge_query: + if query is not None: + pass # fused fp8 path already built the query and wrote KV + elif merge_query: q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_rope_reshaped = q_rope.view( -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim @@ -404,8 +446,14 @@ class CuteDslMLABackend(TRTLLMMLABackend): self.init_forward_metadata(forward_batch) metadata = forward_batch.decode_trtllm_mla_metadata - global_seq_lens = forward_batch.seq_lens[: forward_batch.batch_size] - local_seq_lens = self._get_dcp_local_seq_lens(global_seq_lens) + if metadata.seq_lens_k is not None and metadata.global_seq_lens_k is not None: + # Hoisted path: int32 rank-local + global lens maintained once per + # step by metadata init / graph replay-prep. + local_seq_lens = metadata.seq_lens_k[: forward_batch.batch_size] + global_seq_lens = metadata.global_seq_lens_k[: forward_batch.batch_size] + else: + global_seq_lens = forward_batch.seq_lens[: forward_batch.batch_size] + local_seq_lens = self._get_dcp_local_seq_lens(global_seq_lens) raw_out, lse = self._run_decode_kernel( query=query, kv_cache=kv_cache, diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 62c43eb7c..96518ded3 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -30,6 +30,9 @@ from sglang.srt.layers.dcp import ( ) from sglang.srt.layers.dcp.planner import plan_dcp_decode_metadata from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + is_in_breakable_cuda_graph, +) from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) @@ -205,6 +208,10 @@ class FlashInferMhaChunkKVRunner: class FlashInferMLAAttnBackend(AttentionBackend): """Flashinfer attention kernels.""" + # Verify metadata is ragged-layout aware via generate_attn_arg_prefill; + # graphs key their wrappers by token tier (_verify_graph_key). + supports_ragged_verify_graph: bool = True + def __init__( self, model_runner: ModelRunner, @@ -352,7 +359,9 @@ class FlashInferMLAAttnBackend(AttentionBackend): kv_len_arr=self.cuda_graph_kv_lens[:bs], backend="auto", ) - self.prefill_cuda_graph_metadata[bs] = prefill_wrapper + self.prefill_cuda_graph_metadata[ + self._verify_graph_key(bs, spec_info) + ] = prefill_wrapper self.forward_metadata = PrefillMetadata(prefill_wrapper, False) else: raise ValueError(f"Invalid mode: {forward_mode=}") @@ -404,8 +413,11 @@ class FlashInferMLAAttnBackend(AttentionBackend): use_ragged = ( not get_exec().kernel.flashinfer_mla_disable_ragged and extend_no_prefix - # Piecewise cuda graph should use paged prefill to be compatible with prefix cache + # Captured prefill (tc_piecewise or breakable) must use paged + # prefill: it stays compatible with prefix cache, and the ragged + # wrapper rejects the absorbed-MLA head dims (qk=576, vo=512). and not is_in_tc_piecewise_cuda_graph() + and not is_in_breakable_cuda_graph() ) self.indices_updater_prefill.update( @@ -452,6 +464,15 @@ class FlashInferMLAAttnBackend(AttentionBackend): "kv_indices": self.cuda_graph_kv_indices, } + @staticmethod + def _verify_graph_key(bs: int, spec_info: Optional[SpecInput]): + """bs for uniform graphs; token tier for ragged (tiers share slot + counts but each graph must replay its own recorded plan buffers).""" + layout = spec_info.ragged_verify_layout if spec_info is not None else None + if layout is None: + return bs + return ("ragged", layout.graph_num_tokens) + def _apply_cuda_graph_metadata( self, bs: int, @@ -496,7 +517,9 @@ class FlashInferMLAAttnBackend(AttentionBackend): seq_lens[:bs], seq_lens_sum, prefix_lens=None, - prefill_wrapper_paged=self.prefill_cuda_graph_metadata[bs], + prefill_wrapper_paged=self.prefill_cuda_graph_metadata[ + self._verify_graph_key(bs, spec_info) + ], use_ragged=False, spec_info=spec_info, ) diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index cd350b6d8..984a0edd3 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_attn_backend.py @@ -36,14 +36,23 @@ class HybridAttnBackend(AttentionBackend): self.spec_attn_is_prefill = ( model_runner.server_args.speculative_attention_mode == "prefill" ) - # decide_needs_cpu_seq_lens ORs this flag across backends; without the - # delegation the base-class default (True) forces a per-step seq_lens - # D2H + host sync even when both sub-backends opted out. - self.needs_cpu_seq_lens = ( - prefill_backend.needs_cpu_seq_lens or decode_backend.needs_cpu_seq_lens + # Gates the FutureMap's per-step seq_lens D2H (decide_needs_cpu_seq_lens + # ORs it across backends). Count only what runs in the spec decode loop: + # decode always, prefill only when mode=prefill routes verify to it -- + # else a cpu-lens prefill backend forces the D2H on steps it never serves. + self.needs_cpu_seq_lens = decode_backend.needs_cpu_seq_lens or ( + self.spec_attn_is_prefill and prefill_backend.needs_cpu_seq_lens ) self.max_context_len = model_runner.model_config.context_len + @property + def supports_ragged_verify_graph(self) -> bool: + # Ragged verify is TARGET_VERIFY-only; delegate to its executor. + backend = ( + self.decode_backend if self.spec_attn_is_decode else self.prefill_backend + ) + return backend.supports_ragged_verify_graph + def _select_backend(self, forward_mode: ForwardMode) -> AttentionBackend: """ Select the appropriate attention backend based on the forward mode. @@ -113,6 +122,19 @@ class HybridAttnBackend(AttentionBackend): def get_cuda_graph_seq_len_fill_value(self): return self.decode_backend.get_cuda_graph_seq_len_fill_value() + def init_mha_chunk_metadata( + self, forward_batch: ForwardBatch, disable_flashinfer_ragged: bool = False + ): + # Chunked-prefix / one-shot MHA metadata is a prefill concern. Without + # this delegation the MLA MHA path silently skips (re)planning its + # ragged wrappers when the full-attn backend is this prefill/decode + # split (e.g. --decode-attention-backend trtllm_mla), and any + # prefix-cache-hit extend batch then runs against a stale plan: + # ValueError: q.shape[0] (...) does not match qo_indptr[-1] (...) + init = getattr(self.prefill_backend, "init_mha_chunk_metadata", None) + if init is not None: + init(forward_batch, disable_flashinfer_ragged) + @property def verify_mask(self) -> Optional[VerifyMask]: return self._select_backend(ForwardMode.TARGET_VERIFY).verify_mask diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 00226405b..f62abfe89 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -11,6 +11,7 @@ from sglang.kernels.ops.mamba.mamba_state_indices_triton import ( ) from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( scatter_mamba_states_after_mtp_verify, + track_mamba_states_all_layers, track_mamba_states_if_needed, ) from sglang.srt.configs.hybrid_arch import mamba2_config @@ -118,7 +119,7 @@ class MambaAttnBackendBase(AttentionBackend): # The ring cursor is a per-slot decode counter shared by all GDN layers; # manage it once here (snapshot, hand to layers, advance mod L), not per-layer. # Gate on the linear_replayssm FLAG, not on cursor-tensor presence: the - # spec-verify ring (--enable-gdn-replayssm-spec) shares the write_pos + # spec-verify ring (--enable-linear-replayssm-spec) shares the write_pos # allocation but owns it exclusively via commit_gdn_replayssm_spec # (advance-by-accept-count once per verify step). Advancing it here as # well inserts one phantom/stale ring entry per step and cumulatively @@ -178,13 +179,18 @@ class MambaAttnBackendBase(AttentionBackend): # skip mamba metadata. query_start_loc = None elif forward_batch.forward_mode.is_target_verify(): - query_start_loc = torch.arange( - 0, - forward_batch.input_ids.shape[0] + 1, - step=forward_batch.spec_info.draft_token_num, - dtype=torch.int32, - device=forward_batch.input_ids.device, - ) + ragged_layout = forward_batch.spec_info.ragged_verify_layout + if ragged_layout is not None: + # Compact ragged verify: variable per-request verify lens. + query_start_loc = ragged_layout.qo_indptr_device + else: + query_start_loc = torch.arange( + 0, + forward_batch.input_ids.shape[0] + 1, + step=forward_batch.spec_info.draft_token_num, + dtype=torch.int32, + device=forward_batch.input_ids.device, + ) if self.topk > 1: retrieve_next_token = forward_batch.spec_info.retrieve_next_token @@ -371,7 +377,7 @@ class MambaAttnBackendBase(AttentionBackend): """True iff --enable-linear-replayssm is on for this pool. Gate on the FLAG, not on ``replayssm_write_pos is not None``: the - spec-verify ring (--enable-gdn-replayssm-spec) also allocates the + spec-verify ring (--enable-linear-replayssm-spec) also allocates the cursor tensor but owns it exclusively via commit_gdn_replayssm_spec. The decode-ring metadata machinery gated here (per-bs static cursor buffers, the per-replay snapshot + advance-by-one in _replay_metadata, @@ -485,9 +491,16 @@ class MambaAttnBackendBase(AttentionBackend): self.cached_cuda_graph_decode_query_start_loc[: bs + 1] ) elif forward_mode.is_target_verify(): - self.query_start_loc_list[bs - 1].copy_( - self.cached_cuda_graph_verify_query_start_loc[: bs + 1] + ragged_layout = ( + spec_info.ragged_verify_layout if spec_info is not None else None ) + if ragged_layout is not None: + # Ragged capture: qsl from the runner's synthetic layout. + self.query_start_loc_list[bs - 1].copy_(ragged_layout.qo_indptr_device) + else: + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_verify_query_start_loc[: bs + 1] + ) else: raise ValueError(f"Invalid forward mode: {forward_mode=}") mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) @@ -654,7 +667,19 @@ class MambaAttnBackendBase(AttentionBackend): bs - num_padding ) elif forward_mode.is_target_verify(): - if num_padding == 0: + ragged_layout = ( + spec_info.ragged_verify_layout if spec_info is not None else None + ) + if ragged_layout is not None: + # Mamba kernels index dense [bs, N] scratch, so they need the + # capped variant (see padded_to_bucket). Padding rows carry + # mamba slot -1 and are skipped. + if ragged_layout.bs != bs or ragged_layout.cap is None: + ragged_layout = ragged_layout.padded_to_bucket( + padded_bs=bs, cap=spec_info.draft_token_num + ) + self.query_start_loc_list[bs - 1].copy_(ragged_layout.qo_indptr_device) + elif num_padding == 0: self.query_start_loc_list[bs - 1].copy_( self.cached_cuda_graph_verify_query_start_loc[: bs + 1] ) @@ -705,26 +730,73 @@ class MambaAttnBackendBase(AttentionBackend): def get_cpu_graph_seq_len_fill_value(self): return 1 + def _track_pools(self): + """Full [num_layers, pool_size, ...] conv/ssm pools plus the pool index + of the last mamba layer, for the fused all-layers track launch. None if + the pool shape is not the expected layout.""" + cached = getattr(self, "_track_pools_cache", False) + if cached is False: + pools = None + try: + mamba_cache = self.req_to_token_pool.mamba_pool.mamba_cache + conv_pool = mamba_cache.conv[0] + ssm_pool = mamba_cache.temporal + last_pool_idx = max(self.req_to_token_pool.mamba_map.values()) + if ( + conv_pool.dim() >= 3 + and ssm_pool.dim() >= 3 + and conv_pool.shape[0] == ssm_pool.shape[0] == last_pool_idx + 1 + ): + pools = (conv_pool, ssm_pool, last_pool_idx) + except (AttributeError, IndexError): + pools = None + self._track_pools_cache = pools + cached = pools + return cached + def _track_mamba_state_decode( self, forward_batch: ForwardBatch, conv_states: torch.Tensor, ssm_states: torch.Tensor, cache_indices: torch.Tensor, + layer_id: Optional[int] = None, ): """Copy decode conv/SSM states to track slots for prefix caching. Track dests come from the metadata (under cuda-graph: the static buffer), so the - InputBuffer registry slot is never mutated.""" - if forward_batch.mamba_track_mask is not None: - track_mamba_states_if_needed( - conv_states, - ssm_states, - cache_indices, - forward_batch.mamba_track_mask, - self.forward_metadata.mamba_track_indices, - forward_batch.batch_size, - check_freed_slots=self.enable_unified_memory, - ) + InputBuffer registry slot is never mutated. + + With a known layer_id, the per-layer launches collapse into ONE + all-layers launch fired at the last mamba layer: the mask/src/dst + indices are shared across layers and every layer's state is final by + then, so the result is identical.""" + if forward_batch.mamba_track_mask is None: + return + if layer_id is not None: + pools = self._track_pools() + if pools is not None: + conv_pool, ssm_pool, last_pool_idx = pools + if self.req_to_token_pool.mamba_map.get(layer_id) != last_pool_idx: + return + track_mamba_states_all_layers( + conv_pool, + ssm_pool, + cache_indices, + forward_batch.mamba_track_mask, + self.forward_metadata.mamba_track_indices, + forward_batch.batch_size, + check_freed_slots=self.enable_unified_memory, + ) + return + track_mamba_states_if_needed( + conv_states, + ssm_states, + cache_indices, + forward_batch.mamba_track_mask, + self.forward_metadata.mamba_track_indices, + forward_batch.batch_size, + check_freed_slots=self.enable_unified_memory, + ) def _track_mamba_state_extend( self, @@ -893,8 +965,18 @@ class HybridLinearAttnBackend(AttentionBackend): @property def data_type(self): + # KV-cache dtype readers (e.g. the trtllm_mla fused-rope check) reach the + # wrapper since split backends are wrapped once (#31439); the full-attn + # side owns the KV cache, so its dtype is authoritative. return self.full_attn_backend.data_type + @property + def supports_ragged_verify_graph(self) -> bool: + return ( + self.full_attn_backend.supports_ragged_verify_graph + and self.linear_attn_backend.supports_ragged_verify_graph + ) + def _is_full_attn( self, layer: Optional[RadixAttention], layer_id: Optional[int] = None ) -> bool: @@ -1121,9 +1203,32 @@ class HybridLinearAttnBackend(AttentionBackend): ] ) - mamba_caches = ( - self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() - ) + req_pool = self.linear_attn_backend.req_to_token_pool + mamba_caches = req_pool.get_speculative_mamba2_params_all_layers() + + # ReplaySSM-KDA: the accepted drafts live in the per-slot ring (written + # during verify); no intermediate_ssm is allocated. Replay the accepted + # prefix into `temporal` instead of scattering an intermediate state. + # dspark/dflash call this method directly; the generic spec_utils commit + # handles replayssm before reaching here (returns early), so this branch is + # only hit by the direct callers. Chain layout only (topk <= 1), so + # accept_lens == last_correct_step_indices + 1. + mamba_pool = req_pool.mamba_pool + if getattr(mamba_pool, "replayssm_is_kda", False): + from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( + commit_kda_replayssm_after_verify, + ) + + commit_kda_replayssm_after_verify( + spec_state=mamba_caches, + state_batch_indices=state_indices_tensor, + accept_lens=last_correct_step_indices + 1, + last_correct_step_indices=last_correct_step_indices, + mamba_track_indices=mamba_track_indices, + mamba_steps_to_track=mamba_steps_to_track, + null_block_id=-1, + ) + return scatter_mamba_states_after_mtp_verify( mamba_caches, diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 7143b4124..7539cc1cd 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -12,6 +12,7 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBack from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, + build_verify_intermediate_state_indices, get_linear_attn_decode_backend, get_linear_attn_prefill_backend, ) @@ -346,8 +347,13 @@ class GDNAttnBackend(MambaAttnBackendBase): decode_backend = get_linear_attn_decode_backend() prefill_backend = get_linear_attn_prefill_backend() self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend) - self.verify_intermediate_state_indices = torch.arange( - self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device + # Sized past the pool for attn_tp-padded warmup/MLP-sync batches (see helper). + self.verify_intermediate_state_indices = ( + build_verify_intermediate_state_indices( + self.req_to_token_pool.size, + model_runner.server_args, + model_runner.device, + ) ) def init_forward_metadata(self, forward_batch: ForwardBatch): @@ -427,7 +433,7 @@ class GDNAttnBackend(MambaAttnBackendBase): replayssm_force_flush=replayssm_force_flush, ) self._track_mamba_state_decode( - forward_batch, conv_states, ssm_states, cache_indices + forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id ) return core_attn_out @@ -456,7 +462,7 @@ class GDNAttnBackend(MambaAttnBackendBase): ) self._track_mamba_state_decode( - forward_batch, conv_states, ssm_states, cache_indices + forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id ) return core_attn_out @@ -633,13 +639,13 @@ class GDNAttnBackend(MambaAttnBackendBase): ) else: # The recurrent fallback needs the per-draft snapshots, which - # the pool gates OFF under --enable-gdn-replayssm-spec (the + # the pool gates OFF under --enable-linear-replayssm-spec (the # same flag that makes `use_replayssm_spec` true above), so # this branch is unreachable with a None buffer by # construction -- keep it loud rather than silently frozen. assert intermediate_state_cache is not None, ( "recurrent target_verify fallback requires intermediate_ssm, " - "which is not allocated under --enable-gdn-replayssm-spec" + "which is not allocated under --enable-linear-replayssm-spec" ) core_attn_out = self.kernel_dispatcher.target_verify( A_log=layer.A_log, diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 86bc22d69..06fd7dd78 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -1,7 +1,9 @@ +import importlib.util from typing import Optional, Tuple, Union import torch +from sglang.kernels.ops.attention import kda_fused_decode from sglang.kernels.ops.mamba.causal_conv1d_triton import ( causal_conv1d_fn, causal_conv1d_update, @@ -10,8 +12,10 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBack from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, + build_verify_intermediate_state_indices, get_linear_attn_decode_backend, get_linear_attn_prefill_backend, + get_linear_attn_verify_backend, ) from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.utils import is_cpu, is_cuda, is_npu @@ -39,7 +43,9 @@ class KDAKernelDispatcher: self, decode_backend: LinearAttnKernelBackend, prefill_backend: LinearAttnKernelBackend, + verify_backend: LinearAttnKernelBackend, ): + self.verify_backend = verify_backend triton_kernel = TritonKDAKernel() if decode_backend.is_triton(): @@ -68,15 +74,39 @@ class KDAKernelDispatcher: "KDA supports 'triton', 'cutedsl', or 'flashinfer'." ) - # target_verify (MTP / speculative decode) kernel: each decode backend - # verifies with its own kernel. FlashInfer decode uses recurrent_kda (SM100, - # chain only); Triton -- and CuTe DSL, which has no verify of its own -- use - # the Triton fused KDA verify, which handles chain + tree - # (retrieve_parent_token) and per-step checkpointing and is the reference the - # KDA backend correctness tests assert against. - self.verify_kernel = ( - self.decode_kernel if decode_backend.is_flashinfer() else triton_kernel - ) + # target_verify kernel, selected via --linear-attn-verify-backend (defaults + # to follow decode: flashinfer -> recurrent_kda, else triton). + # triton: fused chain + tree (retrieve_parent_token) verify; the reference + # the KDA correctness tests assert against. + # flashinfer: recurrent_kda (SM100, chain only); reuses the decode kernel + # when decode is also flashinfer. + # nv_cutedsl: fused Kimi-K3/DSpARK dense verify. + if verify_backend.is_triton(): + self.verify_kernel = triton_kernel + elif verify_backend.is_flashinfer(): + if decode_backend.is_flashinfer(): + self.verify_kernel = self.decode_kernel + else: + if not is_cuda(): + raise ValueError("KDA FlashInfer verify backend requires CUDA") + from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import ( + FlashInferKDAKernel, + ) + + self.verify_kernel = FlashInferKDAKernel() + elif verify_backend.is_nv_cutedsl(): + self.verify_kernel = triton_kernel + elif verify_backend.is_custom(): + # Future custom KDA verify kernel plugs in here. + raise NotImplementedError( + "--linear-attn-verify-backend custom: no custom KDA verify kernel " + "is registered yet." + ) + else: + raise ValueError( + f"Unsupported KDA verify backend: {verify_backend}. " + "KDA verify supports 'triton', 'nv_cutedsl', or 'flashinfer'." + ) if prefill_backend.is_triton(): self.extend_kernel = triton_kernel @@ -104,11 +134,40 @@ class KDAKernelDispatcher: rank0_log( "KDA cutedsl prefill needs SM100; falling back to Triton extend." ) + elif prefill_backend.is_ptx_kda(): + if not is_cuda(): + raise ValueError("PTX KDA prefill backend requires CUDA") + from sglang.srt.layers.attention.linear.kernels.kda_ptx import PtxKDAKernel + + ptx_kda_kernel = PtxKDAKernel() + if ptx_kda_kernel.supports_prefill: + self.extend_kernel = ptx_kda_kernel + else: + self.extend_kernel = triton_kernel + rank0_log( + "PTX KDA prefill needs SM103 (GB300); falling back to Triton " + "extend." + ) + elif prefill_backend.is_nvidia_kda(): + if not is_cuda(): + raise ValueError("NVIDIA KDA prefill backend requires CUDA") + from sglang.srt.layers.attention.linear.kernels.kda_nvidia import ( + NvidiaKDAKernel, + ) + + nvidia_kda_kernel = NvidiaKDAKernel() + if nvidia_kda_kernel.supports_prefill: + self.extend_kernel = nvidia_kda_kernel + else: + self.extend_kernel = triton_kernel + rank0_log( + "NVIDIA KDA prefill needs SM100; falling back to Triton extend." + ) else: raise ValueError( f"Unsupported KDA prefill backend: {prefill_backend}. " - "KDA supports 'triton', 'flashkda', or 'cutedsl' " - "(cutedsl prefill needs SM100)." + "KDA supports 'triton', 'flashkda', 'cutedsl', 'nvidia_kda', or " + "'ptx_kda' (cutedsl/nvidia_kda prefill need SM100, ptx_kda SM103)." ) self.supports_packed_decode = getattr( @@ -201,6 +260,7 @@ class KDAKernelDispatcher: intermediate_state_indices: torch.Tensor, cache_steps: int, retrieve_parent_token: torch.Tensor, + lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: """MTP / speculative-decode verify, routed to ``self.verify_kernel`` @@ -221,6 +281,9 @@ class KDAKernelDispatcher: intermediate_state_indices=intermediate_state_indices, cache_steps=cache_steps, retrieve_parent_token=retrieve_parent_token, + lower_bound=lower_bound, + # Forward extras (e.g. the fused ring-write cache_ring/replayssm_*). + **kwargs, ) def extend( @@ -249,19 +312,48 @@ class KDAKernelDispatcher: ) +def ragged_verify_dense_scatter_indices( + *, + query_start_loc: torch.Tensor, + seq_len: int, + draft_token_num: int, +) -> torch.Tensor: + """Dense [bs, draft_token_num] slot index per packed ragged-verify token. + + Rows never exceed draft_token_num under either layout variant (cap for + graph replay, planner construction for eager -- see + RaggedVerifyLayout.padded_to_bucket), so in-row offsets stay in-row; + tokens past the layout's coverage collapse into one ghost row at index + bs * draft_token_num. + """ + batch_size = query_start_loc.shape[0] - 1 + token_pos = torch.arange(seq_len, device=query_start_loc.device, dtype=torch.int32) + token_slots = torch.searchsorted(query_start_loc[1:], token_pos, right=True) + return ( + token_slots * draft_token_num + + (token_pos - query_start_loc[token_slots]).to(torch.int64) + ).clamp_(max=batch_size * draft_token_num) + + class KDAAttnBackend(MambaAttnBackendBase): """Attention backend for KDA (Kimi Delta Attention) linear attention.""" - # Same GPU-only contract as GDNAttnBackend / Mamba2AttnBackend: KDA metadata - # never reads the spec-v2 seq_lens_cpu mirror (replay padding comes from - # forward_batch.num_padding, and the replayssm track-flush mask paths are - # gated `not is_kda`), so don't force FutureMap's blocking per-step - # seq_lens D2H (~0.5 ms/step host stall in bs=1 MTP decode). + # The verify kernel is varlen and the conv path scatters ragged tokens + # to its dense layout, so ragged verify graphs are supported. + supports_ragged_verify_graph: bool = True + + # Read by decide_needs_cpu_seq_lens. Decode/verify metadata is GPU-only + # (graph replay already passes seq_lens_cpu=None), extend reads + # extend_seq_lens_cpu from schedule, mamba track indices rebuild from req + # objects, and the replayssm seq_lens_cpu force-flush is GDN-only. needs_cpu_seq_lens: bool = False def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) - # mamba_cache.conv is [..., kernel-1, dim] while conv_states_shape expects the window length (kernel-1) at shape[-1], hence the transpose. + # Needed by the extra_buffer track path: _init_track_conv_indices reads + # conv_states_shape[-1] as the conv window length (kernel_size - 1). + # The KDA pool stores conv states as [kernel-1, dim] — transposed vs + # Mamba2/GDN's [dim, kernel-1] — so expose the transposed shape here. self.conv_states_shape = ( model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0] .transpose(-1, -2) @@ -269,22 +361,35 @@ class KDAAttnBackend(MambaAttnBackendBase): ) decode_backend = get_linear_attn_decode_backend() prefill_backend = get_linear_attn_prefill_backend() - # KDA FlashInfer speculative decode (target_verify) is linear-chain only -- - # recurrent_kda has no tree-ancestor traversal. Reject EAGLE tree verify - # (topk > 1) early at setup instead of deep in the per-step verify call. - # (The kernel keeps a per-call retrieve_parent_token guard as a backstop; it - # also covers ngram tree, which this topk field does not.) + verify_backend = get_linear_attn_verify_backend() + # KDA FlashInfer target_verify (recurrent_kda) is chain-only (no tree-ancestor + # traversal). Reject EAGLE tree verify (topk > 1) early at setup, keyed on the + # verify backend (not decode). The kernel keeps a per-call + # retrieve_parent_token backstop that also covers ngram tree. speculative_topk = model_runner.server_args.speculative_eagle_topk or 1 - if decode_backend.is_flashinfer() and speculative_topk > 1: + if verify_backend.is_flashinfer() and speculative_topk > 1: raise ValueError( "KDA FlashInfer speculative decoding only supports topk=1 " "(EAGLE tree verify / retrieve_parent_token is unsupported)." ) - self.kernel_dispatcher = KDAKernelDispatcher(decode_backend, prefill_backend) + self.kernel_dispatcher = KDAKernelDispatcher( + decode_backend, prefill_backend, verify_backend + ) + # One-shot; emitted at the first fused-decode interception below. + self._fused_override_notice = ( + "K3 fused KDA decode engaged: --linear-attn-decode-backend " + f"{decode_backend} only picks the fallback kernel for shapes " + "the fused kernel does not cover." + ) # Per-request row index into the speculative `intermediate_ssm` scratch, - # used by the MTP / target_verify path (mirrors GDNAttnBackend). - self.verify_intermediate_state_indices = torch.arange( - self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device + # used by the MTP / target_verify path (mirrors GDNAttnBackend). Sized + # past the pool for attn_tp-padded warmup/MLP-sync batches (see helper). + self.verify_intermediate_state_indices = ( + build_verify_intermediate_state_indices( + self.req_to_token_pool.size, + model_runner.server_args, + model_runner.device, + ) ) def init_forward_metadata(self, forward_batch: ForwardBatch): @@ -326,6 +431,77 @@ class KDAAttnBackend(MambaAttnBackendBase): replayssm_k = layer_cache.replayssm_k replayssm_g = layer_cache.replayssm_g + # Fully fused decode step: conv1d update + delta-rule recurrence + + # gated RMSNorm in one kernel. Engages only when the model handed off + # the output-norm gate for this forward (attempt-and-verify stash, + # see kimi_k3.py) and the shapes are covered; the model applies the + # norm itself whenever the stash is left unconsumed. + if replayssm_d is None: + fused_static = getattr(layer, "_k3_fused_decode_args", None) + onorm_gate = getattr(layer, "_k3_onorm_gate", None) + if ( + fused_static is not None + and onorm_gate is not None + and mixed_qkv.shape[0] == cache_indices.shape[0] + and b.ndim == 3 + and kda_fused_decode.covered( + mixed_qkv, + a, + b[0], + conv_states, + ssm_states, + cache_indices, + onorm_gate, + ) + ): + if self._fused_override_notice is not None: + rank0_log(self._fused_override_notice) + self._fused_override_notice = None + w_q_t, w_k_t, w_v_t, conv_bias, a_log, onorm_w, onorm_eps = fused_static + core_attn_out = kda_fused_decode.kda_fused_decode( + mixed_qkv, + a, + b[0], + conv_states, + w_q_t, + w_k_t, + w_v_t, + conv_bias, + a_log, + layer.dt_bias, + onorm_gate, + onorm_w, + ssm_states, + cache_indices, + scale=layer.head_k_dim**-0.5, + onorm_eps=onorm_eps, + lower_bound=layer.lower_bound, + ) + layer._k3_onorm_consumed = True + self._track_mamba_state_decode( + forward_batch, + conv_states, + ssm_states, + cache_indices, + layer.layer_id, + ) + return core_attn_out + elif fused_static is not None and onorm_gate is not None: + # One-shot diagnostics: the model offered the handoff but the + # runtime shapes were rejected (decode stays on the unfused + # chain, which is correct but slower). + if not getattr(KDAAttnBackend, "_fused_reject_logged", False): + KDAAttnBackend._fused_reject_logged = True + rank0_log( + "KDA fused decode rejected by covered(): " + f"mixed_qkv {tuple(mixed_qkv.shape)}/{mixed_qkv.dtype} " + f"stride {mixed_qkv.stride()}, " + f"conv_states {tuple(conv_states.shape)} " + f"stride {conv_states.stride()}, " + f"ssm_states {tuple(ssm_states.shape)}/{ssm_states.dtype}, " + f"b {tuple(b.shape)}, indices {cache_indices.dtype}" + ) + qkv = causal_conv1d_update( mixed_qkv, conv_states.transpose(-1, -2), @@ -353,6 +529,7 @@ class KDAAttnBackend(MambaAttnBackendBase): cache_indices=cache_indices, num_v_heads=layer.num_v_heads, head_v_dim=layer.head_v_dim, + lower_bound=layer.lower_bound, replayssm_d=replayssm_d, replayssm_k=replayssm_k, replayssm_g=replayssm_g, @@ -360,7 +537,7 @@ class KDAAttnBackend(MambaAttnBackendBase): replayssm_force_flush=replayssm_force_flush, ) self._track_mamba_state_decode( - forward_batch, conv_states, ssm_states, cache_indices + forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id ) return core_attn_out @@ -380,10 +557,11 @@ class KDAAttnBackend(MambaAttnBackendBase): ssm_states=ssm_states, cache_indices=cache_indices, query_start_loc=query_start_loc, + lower_bound=layer.lower_bound, ) self._track_mamba_state_decode( - forward_batch, conv_states, ssm_states, cache_indices + forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id ) return core_attn_out @@ -418,6 +596,11 @@ class KDAAttnBackend(MambaAttnBackendBase): has_initial_state = forward_batch.extend_prefix_lens > 0 if self.forward_metadata.has_mamba_track_mask: + # Snapshot the conv sliding window at the last track-aligned chunk + # boundary into the ping-pong track slots (the prefix-cache restore + # source). The KDA pool stores conv states as [kernel-1, dim], so + # rows of the raw [tokens, dim] pre-conv input index in directly + # (GDN needs a transpose here; KDA does not). mamba_cache_params.conv[0][ self.forward_metadata.conv_states_mask_indices ] = mixed_qkv[self.forward_metadata.track_conv_indices] @@ -483,14 +666,23 @@ class KDAAttnBackend(MambaAttnBackendBase): query_start_loc=query_start_loc, A_log=layer.A_log, dt_bias=layer.dt_bias, - lower_bound=getattr(layer, "lower_bound", None), + lower_bound=layer.lower_bound, extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu, # draft_extend_v2 must stay rollback-able, so kernels that commit state # in place (e.g. FlashKDA) must not run for it. is_spec_decode=forward_batch.forward_mode.is_draft_extend_v2(), return_intermediate_states=track_ssm, + # Which global chunk rows of h the track snapshot will read; lets + # kernels that cannot materialize per-chunk states (NVIDIA KDA) take the + # fast path when the snapshot only needs the final state. + track_ssm_h_src=( + self.forward_metadata.track_ssm_h_src if track_ssm else None + ), ) if track_ssm: + # Snapshot the SSM state at the last track-aligned chunk boundary + # from the kernel's per-chunk states (h) / final states into the + # ping-pong track slots (see _init_track_ssm_indices). core_attn_out, h = core_attn_out self._track_mamba_state_extend( forward_batch, h, ssm_states, self.forward_metadata @@ -525,7 +717,12 @@ class KDAAttnBackend(MambaAttnBackendBase): conv_states = mamba_cache_params.conv[0] ssm_states = mamba_cache_params.temporal intermediate_state_cache = getattr(mamba_cache_params, "intermediate_ssm", None) - if intermediate_state_cache is None: + # ReplaySSM: intermediate_ssm is intentionally None (the ring + commit-time + # fold replace the per-step snapshots). Pass None to the verify kernel so it + # skips the write (CACHE_INTERMEDIATE_STATES=False); the output is unaffected. + replayssm_rawv = getattr(mamba_cache_params, "replayssm_rawv", None) + replayssm_on = replayssm_rawv is not None + if intermediate_state_cache is None and not replayssm_on: raise RuntimeError( "KDA target_verify requires a speculative mamba cache " "(MambaPool.SpeculativeState); none found." @@ -534,13 +731,66 @@ class KDAAttnBackend(MambaAttnBackendBase): intermediate_state_indices = self.verify_intermediate_state_indices draft_token_num = forward_batch.spec_info.draft_token_num - batch_size = seq_len // draft_token_num + ragged_layout = forward_batch.spec_info.ragged_verify_layout + if self._can_run_dspark_cutedsl_mtp( + layer=layer, + mixed_qkv=mixed_qkv, + a=a, + b=b, + draft_token_num=draft_token_num, + ragged_layout=ragged_layout, + conv_states=conv_states, + ssm_states=ssm_states, + intermediate_state_cache=intermediate_state_cache, + intermediate_conv_window_cache=intermediate_conv_window_cache, + retrieve_parent_token=retrieve_parent_token, + replayssm_rawv=replayssm_rawv, + ): + return self._run_dspark_cutedsl_mtp( + layer=layer, + mixed_qkv=mixed_qkv, + a=a, + b=b, + conv_states=conv_states, + ssm_states=ssm_states, + intermediate_state_cache=intermediate_state_cache, + intermediate_conv_window_cache=intermediate_conv_window_cache, + intermediate_state_indices=intermediate_state_indices, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=mamba_cache_params.replayssm_rawk, + replayssm_g=mamba_cache_params.replayssm_g, + replayssm_beta=mamba_cache_params.replayssm_beta, + ) + if ragged_layout is None: + batch_size = seq_len // draft_token_num + dense_token_indices = None + mixed_qkv_dense = mixed_qkv.view(batch_size, draft_token_num, -1) + else: + # Conv update and its per-step scratch want the dense + # [bs, draft_token_num] layout: scatter ragged tokens to their + # step slots, gather back after (pad steps are never committed). + # Uncovered tier-leftover tokens land in the ghost row (see + # ragged_verify_dense_scatter_indices); their values are pad + # garbage, discarded downstream (ghost collisions are + # value-irrelevant). + batch_size = query_start_loc.shape[0] - 1 + num_dense_tokens = batch_size * draft_token_num + dense_token_indices = ragged_verify_dense_scatter_indices( + query_start_loc=query_start_loc, + seq_len=seq_len, + draft_token_num=draft_token_num, + ) + dense = mixed_qkv.new_zeros(num_dense_tokens + 1, mixed_qkv.shape[-1]) + dense.index_copy_(0, dense_token_indices, mixed_qkv) + mixed_qkv_dense = dense[:num_dense_tokens].view( + batch_size, draft_token_num, -1 + ) # causal_conv1d_update expects [.., dim, width]. KDA keeps dense conv-window # scratch because the deduplicated overlapping layout cannot be transposed. - mixed_qkv_reshaped = mixed_qkv.view(batch_size, draft_token_num, -1).transpose( - 1, 2 - ) + mixed_qkv_reshaped = mixed_qkv_dense.transpose(1, 2) mixed_qkv_processed = causal_conv1d_update( mixed_qkv_reshaped, conv_states.transpose(-1, -2), @@ -554,14 +804,41 @@ class KDAAttnBackend(MambaAttnBackendBase): retrieve_next_sibling=retrieve_next_sibling, retrieve_parent_token=retrieve_parent_token, ) - mixed_qkv = mixed_qkv_processed.transpose(1, 2).reshape(seq_len, -1) + mixed_qkv_flat = mixed_qkv_processed.transpose(1, 2).reshape( + batch_size * draft_token_num, -1 + ) + if dense_token_indices is None: + mixed_qkv = mixed_qkv_flat + else: + # Ghost row (zeros) so uncovered tail tokens gather finite values. + padded_flat = mixed_qkv_flat.new_zeros( + batch_size * draft_token_num + 1, mixed_qkv_flat.shape[-1] + ) + padded_flat[: batch_size * draft_token_num] = mixed_qkv_flat + mixed_qkv = padded_flat[dense_token_indices] q, k, v = mixed_qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0) # n (h d) -> 1 n h d k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) - return self.kernel_dispatcher.target_verify( + # ReplaySSM: the ring-write is fused into the triton verify kernel + # (CACHE_RING). Ragged layouts work natively -- step_idx is the + # within-row step under varlen, so row i writes + # ring[slot][0..verify_lens[i]) and commit folds at most commit_lens + # of them (absorb overflow is bounded in-kernel). ring_kwargs stays + # empty for non-triton verify kernels, which never see replayssm. + ring_kwargs = {} + if replayssm_rawv is not None: + ring_kwargs = dict( + cache_ring=True, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=mamba_cache_params.replayssm_rawk, + replayssm_g=mamba_cache_params.replayssm_g, + replayssm_beta=mamba_cache_params.replayssm_beta, + ) + + core_attn_out = self.kernel_dispatcher.target_verify( A_log=layer.A_log, dt_bias=layer.dt_bias, q=q, @@ -576,4 +853,193 @@ class KDAAttnBackend(MambaAttnBackendBase): intermediate_state_indices=intermediate_state_indices, cache_steps=draft_token_num, retrieve_parent_token=retrieve_parent_token, + lower_bound=layer.lower_bound, + **ring_kwargs, ) + if dense_token_indices is not None: + # Kernel output is empty-allocated and the capped qsl skips the + # uncovered tail rows; zero them so discarded pad hidden states + # stay finite. Uncovered == clamped-to-ghost. + covered = dense_token_indices < (batch_size * draft_token_num) + core_attn_out = torch.where(covered.view(1, -1, 1, 1), core_attn_out, 0.0) + return core_attn_out + + def _can_run_dspark_cutedsl_mtp( + self, + *, + layer: RadixLinearAttention, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + draft_token_num: int, + ragged_layout, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + intermediate_state_cache: torch.Tensor, + intermediate_conv_window_cache: torch.Tensor, + retrieve_parent_token: Optional[torch.Tensor], + replayssm_rawv: Optional[torch.Tensor], + ) -> bool: + """Return whether the fixed Kimi-K3/DSpARK CuTe contract is satisfied.""" + if not self.kernel_dispatcher.verify_backend.is_nv_cutedsl() or not is_cuda(): + return False + if importlib.util.find_spec("cutlass") is None: + return False + if torch.cuda.get_device_capability()[0] != 10: + return False + if ragged_layout is not None or retrieve_parent_token is not None: + return False + # draft_token_num = 1 bonus + dspark block size; the CuTe kernel is + # specialized per block size and capped at 8 by shared-memory growth. + if not 2 <= draft_token_num <= 8: + return False + if layer.bias is not None or layer.lower_bound is None: + return False + if ( + layer.head_q_dim != 128 + or layer.head_k_dim != 128 + or layer.head_v_dim != 128 + ): + return False + if ( + layer.num_q_heads != layer.num_k_heads + or layer.num_k_heads != layer.num_v_heads + ): + return False + if ( + mixed_qkv.dtype != torch.bfloat16 + or a.dtype != torch.bfloat16 + or b.dtype != torch.bfloat16 + ): + return False + if layer.conv_weights is None or tuple(layer.conv_weights.shape) != ( + layer.q_dim + layer.k_dim + layer.v_dim, + 4, + ): + return False + if layer.conv_weights.dtype != torch.float32: + return False + if ( + ssm_states.dtype != torch.float32 + or tuple(ssm_states.shape[-3:]) + != (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim) + or tuple(ssm_states.stride()[-3:]) + != ( + layer.head_v_dim * layer.head_k_dim, + layer.head_k_dim, + 1, + ) + or ssm_states.stride(0) % 4 != 0 + or ssm_states.storage_offset() % 4 != 0 + ): + return False + if conv_states.shape[-2] != 3 or intermediate_conv_window_cache.shape[-2] != 3: + return False + if intermediate_state_cache is None: + # ReplaySSM: the ring replaces the per-step snapshots. The kernel + # wrapper validates ring layout/dtypes and raises loudly (there is + # no recurrent fallback without intermediate_ssm). + return replayssm_rawv is not None + if intermediate_state_cache.shape[1] < draft_token_num: + return False + if tuple(intermediate_state_cache.shape[-3:]) != ( + layer.num_v_heads, + layer.head_v_dim, + layer.head_k_dim, + ): + return False + return True + + @staticmethod + def _run_dspark_cutedsl_mtp( + *, + layer: RadixLinearAttention, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + intermediate_state_cache: torch.Tensor, + intermediate_conv_window_cache: torch.Tensor, + intermediate_state_indices: torch.Tensor, + cache_indices: torch.Tensor, + query_start_loc: torch.Tensor, + replayssm_rawv: Optional[torch.Tensor] = None, + replayssm_rawk: Optional[torch.Tensor] = None, + replayssm_g: Optional[torch.Tensor] = None, + replayssm_beta: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + from sglang.kernels.ops.kimi_k3.kda_decode_mtp import ( + fused_kda_decode_mtp_dspark, + ) + + seq_len = mixed_qkv.shape[0] + h = layer.num_v_heads + x_q, x_k, x_v = mixed_qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) + x_q = x_q.reshape(1, seq_len, h, layer.head_q_dim) + x_k = x_k.reshape(1, seq_len, h, layer.head_k_dim) + x_v = x_v.reshape(1, seq_len, h, layer.head_v_dim) + w_q, w_k, w_v = layer.conv_weights.split( + [layer.q_dim, layer.k_dim, layer.v_dim], dim=0 + ) + cs_q, cs_k, cs_v = conv_states.split( + [layer.q_dim, layer.k_dim, layer.v_dim], dim=-1 + ) + cs_q = cs_q.transpose(-1, -2) + cs_k = cs_k.transpose(-1, -2) + cs_v = cs_v.transpose(-1, -2) + intermediate_conv_window_cache = intermediate_conv_window_cache.transpose( + -1, -2 + ) + ic_q, ic_k, ic_v = intermediate_conv_window_cache.split( + [layer.q_dim, layer.k_dim, layer.v_dim], dim=-2 + ) + # The kernel owns all 128 output channels and can fold gated RMSNorm + # into the recurrence. + onorm_gate = getattr(layer, "_k3_onorm_gate", None) + fused_static = getattr(layer, "_k3_fused_decode_args", None) + apply_onorm = onorm_gate is not None and fused_static is not None + if apply_onorm: + onorm_weight = fused_static[5] + onorm_eps = fused_static[6] + onorm_gate = onorm_gate.reshape(1, seq_len, h, layer.head_v_dim) + else: + onorm_weight = None + onorm_eps = None + onorm_gate = None + + out = fused_kda_decode_mtp_dspark( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=a, + beta=b, + A_log=layer.A_log.reshape(-1), + dt_bias=layer.dt_bias.reshape(-1), + recurrent_state=ssm_states, + intermediate_ssm=intermediate_state_cache, + intermediate_state_indices=intermediate_state_indices, + intermediate_conv_q=ic_q, + intermediate_conv_k=ic_k, + intermediate_conv_v=ic_v, + ssm_state_indices=cache_indices.to(torch.int32), + cu_seqlens=query_start_loc.to(torch.int32), + lower_bound=float(layer.lower_bound), + scale=layer.head_q_dim**-0.5, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_g=replayssm_g, + replayssm_beta=replayssm_beta, + onorm_gate=onorm_gate, + onorm_weight=onorm_weight, + onorm_eps=onorm_eps, + ) + if apply_onorm: + layer._k3_onorm_consumed = True + return out diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py b/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py index 74891c91d..e2323b6ae 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py @@ -75,6 +75,11 @@ class CuteDSLKDAKernel(LinearAttnKernelBase): query_start_loc: torch.Tensor, **kwargs, ) -> torch.Tensor: + if kwargs.get("lower_bound") is not None: + raise NotImplementedError( + "KDA safe gate (lower_bound) is not implemented in the CuTe DSL " + "decode kernel; use --linear-attn-decode-backend triton." + ) return cutedsl_fused_sigmoid_gating_kda_update( A_log=A_log, dt_bias=dt_bias, @@ -108,6 +113,10 @@ class CuteDSLKDAKernel(LinearAttnKernelBase): **kwargs, ) -> torch.Tensor: if kwargs.get("return_intermediate_states"): + # The mamba radix extra_buffer track path needs per-chunk states + # (h), which chunk_kda_cutedsl does not expose. Refuse instead of + # silently skipping the snapshot (that corrupts prefix-cache + # restores). Use the Triton prefill backend with extra_buffer. raise NotImplementedError( "CuteDSLKDAKernel.extend cannot return intermediate chunk " "states required by mamba_radix_cache_strategy=extra_buffer; " diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py index cd566f4e6..561c4b71e 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py @@ -78,8 +78,51 @@ class FlashInferKDAKernel(LinearAttnKernelBase): # Cache the constant per-(row-map, batch, T) verify scatter indices # (ssm_state_indices), which never change across verify calls. self._verify_idx_cache: dict = {} + # State pools whose stride layout has been validated against the + # recurrent_kda contract (per-layer views are pool-stable, so id() is + # a stable key — same lifetime argument as _gate_cache). + self._state_contract_ok: set = set() logger.info("Using FlashInfer KDA kernel") + def _check_state_stride_contract(self, ssm_states: torch.Tensor) -> None: + """One-time (per pool view) check that ``ssm_states`` matches the + layout ``recurrent_kda`` was compiled for. + + The kernel's state argument is a CuTe fake tensor of shape + ``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)`` + and ``assumed_align=32`` (flashinfer ``kda_kernels/recurrent_kda.py``): + the slot stride is free — which is what lets the envelope-strided pools + (unified memory / page-major layout, slot stride = per-slot envelope + pitch) be passed in and updated IN PLACE on the cu_seqlens path — but + the inner strides are compiled-in constants and the divisibility / + alignment are hard assumptions. A pool violating them would mis-address + state in-kernel without any error; fail loudly here instead. + """ + key = id(ssm_states) + if key in self._state_contract_ok: + return + if ssm_states.dim() != 4: + raise ValueError( + f"recurrent_kda needs a [N, HV, V, K] state pool; got " + f"shape {tuple(ssm_states.shape)}" + ) + _, hv, v, k = ssm_states.shape + if ssm_states.stride()[1:] != (v * k, k, 1): + raise ValueError( + "recurrent_kda state inner strides must be compact " + f"(V*K, K, 1)=({v * k}, {k}, 1); got {ssm_states.stride()[1:]} " + "(only the slot stride may be non-compact)" + ) + base_bytes = ssm_states.storage_offset() * ssm_states.element_size() + if ssm_states.stride(0) % 16 != 0 or base_bytes % 32 != 0: + raise ValueError( + "recurrent_kda state pool breaks the compiled stride contract: " + f"slot stride {ssm_states.stride(0)} elements must be a multiple " + f"of 16 and the base byte offset {base_bytes} a multiple of 32 " + "(sym_int64(divisibility=16) / assumed_align=32)" + ) + self._state_contract_ok.add(key) + # ---- gate / beta normalization (shared by decode + verify) ---- def _prep_gate_params(self, A_log: torch.Tensor, dt_bias: torch.Tensor): @@ -118,6 +161,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase): ssm_states: torch.Tensor, cache_indices: torch.Tensor, query_start_loc: torch.Tensor, + lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: batch_size = cache_indices.shape[0] @@ -126,6 +170,11 @@ class FlashInferKDAKernel(LinearAttnKernelBase): num_v_heads = v.shape[2] head_v_dim = v.shape[3] + # The committed pool goes into the kernel as-is (in-place update); under + # unified memory / page-major it is an envelope-strided view, which the + # cu_seqlens path supports — verify the compiled contract once per pool. + self._check_state_stride_contract(ssm_states) + # Pack each request as a length-1 sequence ([1, B, ...] + cu_seqlens) so # recurrent_kda indexes the committed pool IN-KERNEL via ssm_state_indices. # The plain [B, 1, ...] path (no cu_seqlens) instead python-gathers @@ -144,9 +193,8 @@ class FlashInferKDAKernel(LinearAttnKernelBase): A_log_fi, dt_bias_fi = self._prep_gate_params(A_log, dt_bias) - # Softplus gate (lower_bound=None) to match the Triton KDA decode path; - # in-place state update into the committed pool (no rollback for decode). - # query_start_loc is the decode cu_seqlens (one token per request). + # Gate contract matches the Triton decode path (safe gate when + # lower_bound set); in-place state update, no rollback for decode. output_fi, _ = self._recurrent_kda( q=query_fi, k=key_fi, @@ -160,7 +208,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase): output_final_state=False, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, - lower_bound=None, + lower_bound=lower_bound, cu_seqlens=query_start_loc.to(torch.int32), ssm_state_indices=cache_indices.to(torch.int32), ) @@ -186,6 +234,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase): intermediate_state_indices: torch.Tensor, cache_steps: int, retrieve_parent_token: torch.Tensor, + lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: if retrieve_parent_token is not None: @@ -272,7 +321,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase): output_final_state=False, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, - lower_bound=None, + lower_bound=lower_bound, cu_seqlens=query_start_loc.to(torch.int32), ssm_state_indices=ssm_state_indices, num_spec_tokens=num_spec_tokens, diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py index 9420c3f57..8dd3b273f 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py @@ -117,6 +117,10 @@ class FlashKDAKernel(LinearAttnKernelBase): return_intermediate_states: bool = False, **kwargs, ) -> torch.Tensor: + # The fused kernel cannot expose per-chunk states (h), which the mamba + # radix extra_buffer track path needs; route tracked batches through + # the Triton chunk_kda fallback instead of silently skipping the + # snapshot (that would corrupt prefix-cache restores). if return_intermediate_states or self._should_fall_back( lower_bound, is_spec_decode, query_start_loc, extend_seq_lens_cpu ): diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py b/python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py new file mode 100644 index 000000000..7ef05a2dc --- /dev/null +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py @@ -0,0 +1,495 @@ +# SPDX-License-Identifier: Apache-2.0 +"""NVIDIA split KDA prefill backend (K1-K4 CuTe/Triton/cuTile pipeline). + +Wraps the vendored ``kda_nvidia_prefill`` package through its FLA-compatible +``chunk_kda_fwd`` interface. SGLang owns the boundary conversion from its +V-major ``[B,H,V,K]`` cache to the vendor's K-major ``[B,H,K,V]`` state and +back; the vendored split kernels keep their original internal layouts. + +The serving wrapper repacks ordinary packed prefill batches into bounded +equal-length groups (B <= 8) and pads every sequence to one of +2k/4k/8k/16k. This makes short multi-sequence prefill use the same NVIDIA KDA +pipeline while bounding compile variants and transient workspaces. Everything +that is not an ordinary state-committing prefill stays on its dedicated path: + +- single-sequence prefill, including chunks of one long request, where Triton + retains the neutral BS=1 performance; +- track batches needing an *interior* chunk snapshot (mamba extra_buffer + with a non-boundary track point): the pipeline never materializes + per-chunk states. Boundary-aligned track batches reuse the final state; +- spec-decode extends, which must stay rollback-able. + +Inputs are staged into per-bucket persistent buffers (2k/4k/8k/16k): +stable pointers keep the vendored pipeline's pointer-keyed wrapper and +fast-path caches hot, and bound both the workspace set and the compile +variants. Pad rows are state-neutral (k/v/beta zero => no rank-1 update; +g sentinel -1000 => decay exactly 1). + +Any unexpected kernel failure falls back to Triton for that batch with a +loud log (attempt-and-verify, same stance as the fused decode kernel). +""" + +import logging +from typing import Optional + +import torch + +from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel +from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( + LinearAttnKernelBase, +) + +logger = logging.getLogger(__name__) + +_BUCKETS = (2048, 4096, 8192, 16384) +_MAX_NVIDIA_KDA_BATCH = 8 + + +def _to_nvidia_kda_state_layout( + state: torch.Tensor, *, head_k_dim: int, head_v_dim: int +) -> torch.Tensor: + """Materialize SGLang [B,H,V,K] state as vendor [B,H,K,V].""" + expected = (head_v_dim, head_k_dim) + if state.ndim != 4 or tuple(state.shape[-2:]) != expected: + raise ValueError( + "SGLang KDA state must be [B,H,V,K] with " + f"(V,K)={expected}, got {tuple(state.shape)}" + ) + return state.transpose(-1, -2).float().contiguous() + + +def _from_nvidia_kda_state_layout( + state: torch.Tensor, + *, + head_k_dim: int, + head_v_dim: int, + dtype: torch.dtype, +) -> torch.Tensor: + """Materialize vendor [B,H,K,V] state as SGLang [B,H,V,K].""" + expected = (head_k_dim, head_v_dim) + if state.ndim != 4 or tuple(state.shape[-2:]) != expected: + raise ValueError( + "NVIDIA KDA state must be [B,H,K,V] with " + f"(K,V)={expected}, got {tuple(state.shape)}" + ) + return state.transpose(-1, -2).to(dtype=dtype).contiguous() + + +class NvidiaKDAKernel(LinearAttnKernelBase): + def __init__(self): + # This kernel uses tcgen05 + TMEM, which are available on datacenter + # Blackwell (SM100/SM103, reported as capability major 10), but not on + # Blackwell desktop SM120 even though its capability number is larger. + self.supports_prefill = torch.cuda.is_available() and ( + torch.cuda.get_device_capability()[0] == 10 + ) + self._fwd = None + self._l2norm = None + self._triton = TritonKDAKernel() + # Stable detached fp32 views of the (frozen) gate params: nn.Parameters + # trip cute's from_dlpack (grad-tracking) and per-call views defeat + # the pointer-keyed wrapper caches. Keyed per source tensor -- this + # kernel instance is shared by every KDA layer (one dispatcher per + # attention backend) and each layer has its own A_log/dt_bias. + self._param_flat = {} + # One-shot engagement log per (batch size, bucket) (observability). + self._engaged_logged = set() + self._unsupported_logged = set() + # (batch size, bucket, shape, device) -> staging dict. + self._staging = {} + + def _ensure_loaded(self): + if self._fwd is None: + from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd + from sglang.kernels.ops.attention.linear.kda_nvidia_prefill import ( + chunk_kda_fwd, + ) + + self._fwd = chunk_kda_fwd + self._l2norm = l2norm_fwd + logger.info("Using NVIDIA chunked KDA prefill (Blackwell)") + + def decode(self, *args, **kwargs): + raise NotImplementedError("NvidiaKDAKernel is prefill-only") + + def target_verify(self, *args, **kwargs): + raise NotImplementedError("NvidiaKDAKernel does not support target_verify") + + def _triton_extend( + self, + q, + k, + v, + g, + beta, + ssm_states, + cache_indices, + query_start_loc, + A_log, + dt_bias, + lower_bound, + return_intermediate_states, + kwargs, + ): + return self._triton.extend( + q, + k, + v, + g, + beta, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + return_intermediate_states=return_intermediate_states, + **kwargs, + ) + + def _flat_param(self, t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + if t is None: + return None + key = (t.data_ptr(), t.dtype, tuple(t.shape)) + v = self._param_flat.get(key) + if v is None: + v = t.detach().reshape(-1).float().contiguous() + self._param_flat[key] = v + return v + + def _get_staging(self, batch_size, bucket, num_heads, head_k_dim, head_v_dim, dev): + key = ( + batch_size, + bucket, + num_heads, + head_k_dim, + head_v_dim, + dev.index, + ) + st = self._staging.get(key) + if st is None: + st = { + "q": torch.zeros( + batch_size, + bucket, + num_heads, + head_k_dim, + dtype=torch.bfloat16, + device=dev, + ), + "k": torch.zeros( + batch_size, + bucket, + num_heads, + head_k_dim, + dtype=torch.bfloat16, + device=dev, + ), + "v": torch.zeros( + batch_size, + bucket, + num_heads, + head_v_dim, + dtype=torch.bfloat16, + device=dev, + ), + "g": torch.full( + (batch_size, bucket, num_heads, head_k_dim), + -1000.0, + dtype=torch.bfloat16, + device=dev, + ), + "beta": torch.zeros( + batch_size, + bucket, + num_heads, + dtype=torch.bfloat16, + device=dev, + ), + "s0": torch.zeros( + batch_size, + num_heads, + head_k_dim, + head_v_dim, + dtype=torch.float32, + device=dev, + ), + } + self._staging[key] = st + return st + + def extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + query_start_loc: torch.Tensor, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + lower_bound: Optional[float] = None, + return_intermediate_states: bool = False, + **kwargs, + ) -> torch.Tensor: + num_tokens = q.shape[1] + seq_lens_cpu = kwargs.get("extend_seq_lens_cpu") + track_h_src = kwargs.get("track_ssm_h_src") + supported_shape = ( + q.ndim == 4 + and k.ndim == 4 + and v.ndim == 4 + and q.shape[-1] == 128 + and k.shape[-1] == 128 + and v.shape[-1] == 128 + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and v.dtype == torch.bfloat16 + and g.dtype == torch.bfloat16 + and beta.dtype == torch.float32 + and A_log is not None + ) + needs_interior_snapshot = ( + return_intermediate_states + and track_h_src is not None + and track_h_src.numel() > 0 + ) + eligible = ( + not needs_interior_snapshot + and not kwargs.get("is_spec_decode") + and seq_lens_cpu is not None + and len(seq_lens_cpu) > 1 + and supported_shape + ) + if not eligible: + if ( + seq_lens_cpu is not None + and len(seq_lens_cpu) > 1 + and not kwargs.get("is_spec_decode") + and not needs_interior_snapshot + and not supported_shape + ): + unsupported_key = ( + q.dtype, + k.dtype, + v.dtype, + g.dtype, + beta.dtype, + q.shape[-1], + k.shape[-1], + v.shape[-1], + A_log is not None, + ) + if unsupported_key not in self._unsupported_logged: + self._unsupported_logged.add(unsupported_key) + logger.warning( + "NVIDIA KDA prefill only supports BF16 q/k/v/g, FP32 beta, " + "K=V=128 with A_log; " + "got q/k/v/g/beta=%s/%s/%s/%s/%s, Kq/Kk/V=%d/%d/%d, " + "A_log=%s. Falling back to Triton.", + q.dtype, + k.dtype, + v.dtype, + g.dtype, + beta.dtype, + q.shape[-1], + k.shape[-1], + v.shape[-1], + A_log is not None, + ) + return self._triton_extend( + q, + k, + v, + g, + beta, + ssm_states, + cache_indices, + query_start_loc, + A_log, + dt_bias, + lower_bound, + return_intermediate_states, + kwargs, + ) + + self._ensure_loaded() + + seq_lens = [int(length) for length in seq_lens_cpu] + if ( + any(length <= 0 or length > _BUCKETS[-1] for length in seq_lens) + or sum(seq_lens) != num_tokens + or len(seq_lens) != cache_indices.numel() + ): + logger.warning( + "NVIDIA KDA cannot repack prefill shape T=%d seq_lens=%s; falling " + "back to Triton", + num_tokens, + seq_lens, + ) + return self._triton_extend( + q, + k, + v, + g, + beta, + ssm_states, + cache_indices, + query_start_loc, + A_log, + dt_bias, + lower_bound, + return_intermediate_states, + kwargs, + ) + + num_heads = q.shape[2] + head_k_dim = q.shape[-1] + head_v_dim = v.shape[-1] + + alog_flat = self._flat_param(A_log) + dtb_flat = self._flat_param(dt_bias) + + all_slot_indices = torch.where( + cache_indices >= 0, cache_indices, ssm_states.shape[0] - 1 + ).to(torch.int64) + state_backup = ssm_states.index_select(0, all_slot_indices).clone() + packed_output = torch.empty_like(v) + + try: + seq_start = 0 + token_start = 0 + while seq_start < len(seq_lens): + group_lens = seq_lens[seq_start : seq_start + _MAX_NVIDIA_KDA_BATCH] + group_size = len(group_lens) + group_tokens = sum(group_lens) + bucket = next(b for b in _BUCKETS if b >= max(group_lens)) + engage_key = (group_size, bucket) + if engage_key not in self._engaged_logged: + self._engaged_logged.add(engage_key) + logger.info( + "NVIDIA KDA prefill engaged: sequences=%d max_T=%d -> B=%d " + "bucket=%d", + len(seq_lens), + max(group_lens), + group_size, + bucket, + ) + + st = self._get_staging( + group_size, + bucket, + num_heads, + head_k_dim, + head_v_dim, + q.device, + ) + st["q"].zero_() + st["k"].zero_() + st["v"].zero_() + st["g"].fill_(-1000.0) + st["beta"].zero_() + + row_token_start = token_start + for row, length in enumerate(group_lens): + row_token_end = row_token_start + length + st["q"][row, :length].copy_( + self._l2norm(q[0, row_token_start:row_token_end].contiguous()) + ) + st["k"][row, :length].copy_( + self._l2norm(k[0, row_token_start:row_token_end].contiguous()) + ) + st["v"][row, :length].copy_(v[0, row_token_start:row_token_end]) + g_in = g[0, row_token_start:row_token_end] + if g_in.dim() == 2: + g_in = g_in.view(length, num_heads, head_k_dim) + st["g"][row, :length].copy_(g_in) + st["beta"][row, :length].copy_( + beta[0, row_token_start:row_token_end] + ) + row_token_start = row_token_end + + group_slots = all_slot_indices[seq_start : seq_start + group_size] + st["s0"].copy_( + _to_nvidia_kda_state_layout( + ssm_states.index_select(0, group_slots), + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + ) + ) + res = self._fwd( + st["q"], + st["k"], + st["v"], + st["g"], + st["beta"], + scale=head_k_dim**-0.5, + initial_state=st["s0"], + output_final_state=True, + cu_seqlens=None, + safe_gate=lower_bound is not None, + lower_bound=lower_bound, + use_gate_in_kernel=True, + A_log=alog_flat, + dt_bias=dtb_flat, + ) + group_output, final_state = res[0], res[1] + + row_token_start = token_start + for row, length in enumerate(group_lens): + row_token_end = row_token_start + length + packed_output[0, row_token_start:row_token_end].copy_( + group_output[row, :length] + ) + row_token_start = row_token_end + + ssm_states.index_copy_( + 0, + group_slots, + _from_nvidia_kda_state_layout( + final_state, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + dtype=ssm_states.dtype, + ), + ) + seq_start += group_size + token_start += group_tokens + except Exception: + ssm_states.index_copy_(0, all_slot_indices, state_backup) + logger.warning( + "NVIDIA KDA prefill failed (T=%d sequences=%d); restored states " + "and fell back to Triton for this batch", + num_tokens, + len(seq_lens), + exc_info=True, + ) + return self._triton_extend( + q, + k, + v, + g, + beta, + ssm_states, + cache_indices, + query_start_loc, + A_log, + dt_bias, + lower_bound, + return_intermediate_states, + kwargs, + ) + + if return_intermediate_states: + # Boundary-aligned track: the snapshot comes from the final state + # already scattered into the pool (track_ssm_final_src); h is + # never indexed (track_ssm_h_src is empty), a zero-row stand-in + # keeps the (out, h) contract. + h_empty = q.new_empty( + (1, 0) + tuple(ssm_states.shape[1:]), dtype=torch.float32 + ) + return packed_output, h_empty + return packed_output diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py b/python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py new file mode 100644 index 000000000..2ea19c465 --- /dev/null +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PTX/tcgen05 KDA chunked-prefill backend (``--linear-attn-prefill-backend ptx_kda``). + +Wraps the vendored hand-CUDA ``kda_ptx_prefill`` kernel (GB300 / sm_103a) through +its FLA-compatible ``chunk_kda_fwd`` interface. The kernel selects a fused +long-sequence route or a two-launch high-head/many-sequence route at the KDA +serving shape: K = V = 128, chunk 64. + +Scope: ordinary extend batches satisfying the kernel's fixed tensor contract. +Correctness-sensitive cases stay on Triton: + +- track batches receive dense intermediate SSM states directly from the kernel + when the cache checkpoint stride is also 64 tokens. Other interior snapshots + stay on Triton; boundary-only tracking can still use the final state; +- spec-decode extends, which must stay rollback-able. + +Single-sequence token counts that are not a multiple of the kernel's 64-token +chunk are padded up to a bucket (1k/2k/4k/8k/16k/32k) in a persistent staging +buffer, which bounds the resident workspace set. Pad rows are state-neutral: +k/v/beta zero => no rank-1 update; raw gate -1000 => transformed decay of +exactly 1. Multi-sequence batches go through the kernel's own varlen grid +(real cu_seqlens, no padding), so their shapes are whatever the scheduler +produces and each distinct shape can retain another workspace. + +The kernel caches one workspace per distinct shape for the process lifetime; +callers should account for that resident-memory behavior when sending highly +variable packed batches. Any unexpected kernel failure restores the state and +falls back to Triton for that batch with a loud log. +""" + +import logging +from typing import Optional + +import torch + +from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel +from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( + LinearAttnKernelBase, +) + +logger = logging.getLogger(__name__) + +_CHUNK = 64 +_PAD_BUCKETS = (1024, 2048, 4096, 8192, 16384, 32768) +# Raw-gate sentinel for pad rows: softplus(-1000) ~ 0 and +# sigmoid(-1000) ~ 0, so either gate transform yields glog 0 => decay 1. +_PAD_GATE = -1000.0 + + +class PtxKDAKernel(LinearAttnKernelBase): + def __init__(self): + # tcgen05 + TMEM with sm_103a-only encodings: GB300 (SM103) only. + self.supports_prefill = torch.cuda.is_available() and ( + torch.cuda.get_device_capability() == (10, 3) + ) + self._fwd = None + self._triton = TritonKDAKernel() + # Stable detached fp32 views of the (frozen) gate params, keyed per + # source tensor: this kernel instance is shared by every KDA layer. + self._param_flat = {} + self._unsupported_logged = False + # (bucket, H, K, V, device) -> staging dict for ragged token counts. + self._staging = {} + + def _ensure_loaded(self): + if self._fwd is None: + from sglang.kernels.ops.attention.linear.kda_ptx_prefill import ( + chunk_kda_fwd, + load_ext, + ) + + logger.info("Building the PTX KDA prefill extension (first use, ~1-2 min)") + load_ext() + self._fwd = chunk_kda_fwd + logger.info("Using PTX KDA chunked prefill (GB300 / sm_103a)") + + def decode(self, *args, **kwargs): + raise NotImplementedError("PtxKDAKernel is prefill-only") + + def target_verify(self, *args, **kwargs): + raise NotImplementedError("PtxKDAKernel does not support target_verify") + + def _flat_param(self, t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + if t is None: + return None + key = (t.data_ptr(), t.dtype, tuple(t.shape)) + flat = self._param_flat.get(key) + if flat is None: + flat = t.detach().reshape(-1).float().contiguous() + self._param_flat[key] = flat + return flat + + def _get_staging(self, bucket, num_heads, head_k_dim, head_v_dim, dev): + key = (bucket, num_heads, head_k_dim, head_v_dim, dev.index) + st = self._staging.get(key) + if st is None: + st = { + "q": torch.zeros( + 1, bucket, num_heads, head_k_dim, dtype=torch.bfloat16, device=dev + ), + "k": torch.zeros( + 1, bucket, num_heads, head_k_dim, dtype=torch.bfloat16, device=dev + ), + "v": torch.zeros( + 1, bucket, num_heads, head_v_dim, dtype=torch.bfloat16, device=dev + ), + "g": torch.full( + (1, bucket, num_heads, head_k_dim), + _PAD_GATE, + dtype=torch.bfloat16, + device=dev, + ), + "beta": torch.zeros( + 1, bucket, num_heads, dtype=torch.bfloat16, device=dev + ), + } + self._staging[key] = st + return st + + def _triton_extend( + self, + q, + k, + v, + g, + beta, + ssm_states, + cache_indices, + query_start_loc, + A_log, + dt_bias, + lower_bound, + return_intermediate_states, + kwargs, + ): + return self._triton.extend( + q, + k, + v, + g, + beta, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + return_intermediate_states=return_intermediate_states, + **kwargs, + ) + + def extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + query_start_loc: torch.Tensor, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + lower_bound: Optional[float] = None, + return_intermediate_states: bool = False, + **kwargs, + ) -> torch.Tensor: + num_tokens = q.shape[1] + seq_lens_cpu = kwargs.get("extend_seq_lens_cpu") + track_h_src = kwargs.get("track_ssm_h_src") + supported_shape = ( + q.ndim == 4 + and k.ndim == 4 + and v.ndim == 4 + and q.shape[-1] == 128 + and k.shape[-1] == 128 + and v.shape[-1] == 128 + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and v.dtype == torch.bfloat16 + and g.dtype == torch.bfloat16 + and A_log is not None + ) + needs_interior_snapshot = ( + return_intermediate_states + and track_h_src is not None + and track_h_src.numel() > 0 + ) + intermediate_stride_supported = True + if needs_interior_snapshot: + # The CUDA kernel hardcodes BT=64 and h[c] is the state before + # global chunk c. SGLang indexes h using mamba_cache_chunk_size, + # which may be larger (for example 256 for Nemotron-H). Returning + # the raw 64-token rows in that case would silently select the + # wrong boundary and compute wrong offsets for later sequences. + from sglang.srt.runtime_context import get_server_args + + intermediate_stride_supported = ( + get_server_args().mamba_cache_chunk_size == _CHUNK + ) + seq_lens = ( + [int(length) for length in seq_lens_cpu] if seq_lens_cpu is not None else [] + ) + shape_known = ( + len(seq_lens) >= 1 + and len(seq_lens) == cache_indices.numel() + and sum(seq_lens) == num_tokens + and min(seq_lens, default=0) > 0 + ) + # Only the single-sequence path pads (staging buffers keep the resident + # workspace set bounded); varlen batches go in as-is. A ragged single + # sequence beyond the largest bucket uses the kernel's ragged grid. + padded_tokens = num_tokens + if len(seq_lens) == 1 and num_tokens % _CHUNK != 0: + padded_tokens = next( + (bucket for bucket in _PAD_BUCKETS if bucket >= num_tokens), + num_tokens, + ) + eligible = ( + not kwargs.get("is_spec_decode") + and intermediate_stride_supported + and shape_known + and supported_shape + ) + if not eligible: + if shape_known and not supported_shape and not self._unsupported_logged: + self._unsupported_logged = True + logger.warning( + "PTX KDA prefill only supports BF16 q/k/v/g with K=V=128 and " + "A_log; got q/k/v/g=%s/%s/%s/%s, Kq/Kk/V=%d/%d/%d, A_log=%s. " + "Falling back to Triton.", + q.dtype, + k.dtype, + v.dtype, + g.dtype, + q.shape[-1], + k.shape[-1], + v.shape[-1], + A_log is not None, + ) + return self._triton_extend( + q, + k, + v, + g, + beta, + ssm_states, + cache_indices, + query_start_loc, + A_log, + dt_bias, + lower_bound, + return_intermediate_states, + kwargs, + ) + + self._ensure_loaded() + + num_heads = q.shape[2] + head_k_dim = q.shape[-1] + head_v_dim = v.shape[-1] + + slot = torch.where( + cache_indices >= 0, cache_indices, ssm_states.shape[0] - 1 + ).to(torch.int64) + state_backup = ssm_states.index_select(0, slot).clone() + + try: + qn = q + kn = k + g4 = g.reshape(1, num_tokens, num_heads, head_k_dim) + b3 = beta.reshape(1, num_tokens, num_heads) + if padded_tokens != num_tokens: + st = self._get_staging( + padded_tokens, num_heads, head_k_dim, head_v_dim, q.device + ) + for name, src in (("q", qn), ("k", kn), ("v", v), ("g", g4)): + st[name][0, :num_tokens].copy_(src[0]) + st["beta"][0, :num_tokens].copy_(b3[0]) + st["q"][0, num_tokens:].zero_() + st["k"][0, num_tokens:].zero_() + st["v"][0, num_tokens:].zero_() + st["g"][0, num_tokens:].fill_(_PAD_GATE) + st["beta"][0, num_tokens:].zero_() + qn, kn, v_in, g4, b3 = st["q"], st["k"], st["v"], st["g"], st["beta"] + else: + v_in = v + + # Multi-sequence packed batch, or a single sequence beyond the + # padding buckets -> the kernel's varlen grid. Host + # cu_seqlens avoids the D2H sync its piece table would otherwise do. + cu_kwargs = {} + if len(seq_lens) > 1 or padded_tokens % _CHUNK != 0: + cu_cpu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32) + cu_cpu[1:] = torch.tensor(seq_lens, dtype=torch.int32).cumsum(0) + cu_kwargs = dict(cu_seqlens=query_start_loc, cu_seqlens_cpu=cu_cpu) + + result = self._fwd( + qn, + kn, + v_in, + g4, + b3, + scale=head_k_dim**-0.5, + initial_state=ssm_states.index_select(0, slot), + output_final_state=True, + **cu_kwargs, + # SGLang keeps the recurrent state V-major ([slots, H, V, K]); + # the kernel is K-major, so let the wrapper transpose both ways. + state_v_first=True, + safe_gate=lower_bound is not None, + lower_bound=lower_bound, + use_gate_in_kernel=True, + A_log=self._flat_param(A_log), + dt_bias=self._flat_param(dt_bias), + return_intermediate_states=return_intermediate_states, + use_qk_l2norm_in_kernel=True, + ) + out, final_state, h = result[0], result[1], result[10] + ssm_states.index_copy_(0, slot, final_state.to(ssm_states.dtype)) + except Exception: + ssm_states.index_copy_(0, slot, state_backup) + logger.warning( + "PTX KDA prefill failed (T=%d); restored state and fell back to " + "Triton for this batch", + num_tokens, + exc_info=True, + ) + return self._triton_extend( + q, + k, + v, + g, + beta, + ssm_states, + cache_indices, + query_start_loc, + A_log, + dt_bias, + lower_bound, + return_intermediate_states, + kwargs, + ) + + packed_output = out[:, :num_tokens].contiguous() + if return_intermediate_states: + return packed_output, h + return packed_output diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py index c2e8e0ed8..ebf06d49f 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py @@ -38,6 +38,7 @@ class TritonKDAKernel(LinearAttnKernelBase): cache_indices: torch.Tensor, num_v_heads: int, head_v_dim: int, + lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: """Packed decode fast path: feed the conv-1d output ``mixed_qkv`` @@ -67,6 +68,11 @@ class TritonKDAKernel(LinearAttnKernelBase): and replayssm_g is not None and replayssm_write_pos is not None ): + if lower_bound is not None: + raise NotImplementedError( + "KDA safe gate (lower_bound) is not implemented in the " + "ReplaySSM decode kernel; disable --enable-linear-replayssm." + ) K = ssm_states.shape[-1] # ssm_states: [num_slots, HV, V, K] fused_recurrent_linear_replayssm_decode( mixed_qkv=mixed_qkv, @@ -105,6 +111,7 @@ class TritonKDAKernel(LinearAttnKernelBase): out=out, ssm_state_indices=cache_indices, use_qk_l2norm_in_kernel=True, + lower_bound=lower_bound, ) # [B, 1, HV, V] -> [1, B, HV, V] view to match existing decode layout. return out.transpose(0, 1) @@ -122,6 +129,7 @@ class TritonKDAKernel(LinearAttnKernelBase): ssm_states: torch.Tensor, cache_indices: torch.Tensor, query_start_loc: torch.Tensor, + lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: return fused_sigmoid_gating_delta_rule_update( @@ -139,6 +147,7 @@ class TritonKDAKernel(LinearAttnKernelBase): softplus_beta=1.0, softplus_threshold=20.0, is_kda=True, + lower_bound=lower_bound, ) def target_verify( @@ -158,6 +167,13 @@ class TritonKDAKernel(LinearAttnKernelBase): intermediate_state_indices: torch.Tensor, cache_steps: int, retrieve_parent_token: torch.Tensor, + lower_bound: Optional[float] = None, + # fused ReplaySSM ring-write (dense verify only; off elsewhere). + cache_ring: bool = False, + replayssm_rawv: Optional[torch.Tensor] = None, + replayssm_rawk: Optional[torch.Tensor] = None, + replayssm_g: Optional[torch.Tensor] = None, + replayssm_beta: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: # KDA MTP / speculative-decode verify via the fused KDA kernel (IS_KDA=True), @@ -186,6 +202,12 @@ class TritonKDAKernel(LinearAttnKernelBase): intermediate_state_indices=intermediate_state_indices, cache_steps=cache_steps, retrieve_parent_token=retrieve_parent_token, + lower_bound=lower_bound, + cache_ring=cache_ring, + replayssm_rawv=replayssm_rawv, + replayssm_rawk=replayssm_rawk, + replayssm_g=replayssm_g, + replayssm_beta=replayssm_beta, ) def extend( diff --git a/python/sglang/srt/layers/attention/linear/utils.py b/python/sglang/srt/layers/attention/linear/utils.py index 520c10c78..99bdc07f1 100644 --- a/python/sglang/srt/layers/attention/linear/utils.py +++ b/python/sglang/srt/layers/attention/linear/utils.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from enum import Enum -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Dict, Optional from sglang.srt.utils.common import rank0_log @@ -15,8 +15,11 @@ logger = logging.getLogger(__name__) class LinearAttnKernelBackend(Enum): TRITON = "triton" CUTEDSL = "cutedsl" + NV_CUTEDSL = "nv_cutedsl" FLASHINFER = "flashinfer" FLASHKDA = "flashkda" + NVIDIA_KDA = "nvidia_kda" + PTX_KDA = "ptx_kda" CUSTOM = "custom" @classmethod @@ -29,51 +32,109 @@ class LinearAttnKernelBackend(Enum): def is_cutedsl(self): return self == LinearAttnKernelBackend.CUTEDSL + def is_nv_cutedsl(self): + return self == LinearAttnKernelBackend.NV_CUTEDSL + def is_flashinfer(self): return self == LinearAttnKernelBackend.FLASHINFER def is_flashkda(self): return self == LinearAttnKernelBackend.FLASHKDA + def is_nvidia_kda(self): + return self == LinearAttnKernelBackend.NVIDIA_KDA + + def is_ptx_kda(self): + return self == LinearAttnKernelBackend.PTX_KDA + def is_custom(self): return self == LinearAttnKernelBackend.CUSTOM -LINEAR_ATTN_DECODE_BACKEND: Optional[LinearAttnKernelBackend] = None -LINEAR_ATTN_PREFILL_BACKEND: Optional[LinearAttnKernelBackend] = None +_BACKENDS: Dict[str, Optional[LinearAttnKernelBackend]] = { + "decode": None, + "prefill": None, + "verify": None, +} def initialize_linear_attn_config( server_args: ServerArgs, prefill_default: Optional[str] = None ): - global LINEAR_ATTN_DECODE_BACKEND - global LINEAR_ATTN_PREFILL_BACKEND - base = server_args.linear_attn_backend decode = server_args.linear_attn_decode_backend or base prefill = server_args.linear_attn_prefill_backend or prefill_default or base - LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend(decode) - LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend(prefill) + _BACKENDS["decode"] = LinearAttnKernelBackend(decode) + _BACKENDS["prefill"] = LinearAttnKernelBackend(prefill) - rank0_log(f"Linear attention kernel backend: decode={decode}, prefill={prefill}") + # Verify backend. Unset -> follow decode (flashinfer -> its recurrent kernel, + # else triton), preserving historical behavior. + verify = server_args.linear_attn_verify_backend + if verify is None: + verify = decode if _BACKENDS["decode"].is_flashinfer() else "triton" + _BACKENDS["verify"] = LinearAttnKernelBackend(verify) + + rank0_log( + f"Linear attention kernel backend: decode={decode}, prefill={prefill}, " + f"verify={verify}" + ) + + +def _get_backend(phase: str) -> LinearAttnKernelBackend: + backend = _BACKENDS[phase] + if backend is None: + logger.warning( + "linear-attn %s backend is not initialized, using triton backend", phase + ) + backend = _BACKENDS[phase] = LinearAttnKernelBackend.TRITON + return backend def get_linear_attn_decode_backend() -> LinearAttnKernelBackend: - global LINEAR_ATTN_DECODE_BACKEND - if LINEAR_ATTN_DECODE_BACKEND is None: - logger.warning( - "LINEAR_ATTN_DECODE_BACKEND is not initialized, using triton backend" - ) - LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend.TRITON - return LINEAR_ATTN_DECODE_BACKEND + return _get_backend("decode") def get_linear_attn_prefill_backend() -> LinearAttnKernelBackend: - global LINEAR_ATTN_PREFILL_BACKEND - if LINEAR_ATTN_PREFILL_BACKEND is None: - logger.warning( - "LINEAR_ATTN_PREFILL_BACKEND is not initialized, using triton backend" + return _get_backend("prefill") + + +def get_linear_attn_verify_backend() -> LinearAttnKernelBackend: + return _get_backend("verify") + + +def build_verify_intermediate_state_indices( + pool_size: int, server_args: ServerArgs, device +): + """Per-request row index into the speculative intermediate scratch + (`intermediate_ssm` / `intermediate_conv_window`) for the MTP / + target_verify path: request slot i owns scratch row i. + + The scratch is allocated with one extra padding row (the `+1` in + MambaPool.SpeculativeState, index `pool_size`). Warmup and MLP-sync + batches can be padded past the pool capacity — under DP attention + `get_eager_max_batch_size` ceil-aligns the eager warmup bs to attn_tp — + and the verify kernels index this table positionally up to that padded + bs. Size the table to the padded maximum and clamp every out-of-pool row + onto the padding row: pad rows race onto one discard row, which is + value-irrelevant (same convention as the ragged-verify ghost row). + """ + import torch + + from sglang.srt.utils.common import get_eager_max_batch_size + + padded_bs = max(get_eager_max_batch_size(server_args, pool_size), pool_size) + indices = torch.arange(pool_size, dtype=torch.int32, device=device) + if padded_bs > pool_size: + indices = torch.cat( + [ + indices, + torch.full( + (padded_bs - pool_size,), + pool_size, + dtype=torch.int32, + device=device, + ), + ] ) - LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend.TRITON - return LINEAR_ATTN_PREFILL_BACKEND + return indices diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index cef26b2a2..17ee1b76c 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -55,6 +55,7 @@ from sglang.srt.layers.attention.trtllm_mla_backend import ( TRTLLMMLAMultiStepDraftBackend, ) from sglang.srt.layers.dcp.layout import get_dcp_lens +from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import is_flashinfer_available, is_tokenspeed_mla_available @@ -75,7 +76,8 @@ logger = logging.getLogger(__name__) # Workspace upper bound for tokenspeed_mla_decode: # num_sms * num_heads * max_q_len * (kv_lora_rank + 1) * sizeof(float32) -# MAX_Q_LEN=8 covers EAGLE3 num_draft_tokens=4 plus headroom. +# MAX_Q_LEN=8 covers EAGLE3 num_draft_tokens=4 plus headroom. Larger +# speculative widths grow the bound at backend initialization. _TOKENSPEED_MAX_Q_LEN = 8 @@ -401,7 +403,10 @@ class TokenspeedMLABackend(TRTLLMMLABackend): ): if not get_parallel().dcp_enabled: return super()._apply_cuda_graph_metadata( - bs, req_pool_indices, seq_lens, forward_mode + bs, + req_pool_indices, + seq_lens, + forward_mode, ) metadata = self.decode_cuda_graph_metadata[bs] @@ -511,6 +516,25 @@ class TokenspeedMLABackend(TRTLLMMLABackend): llama_4_scaling: Optional[torch.Tensor] = None, ): parallel = get_parallel() + # FlashInfer autotunes MoE kernels with a synthetic full-model decode + # and discards the attention/logits result. On multi-node GB300, the + # synthetic full-head DCP metadata can make both the TokenSpeed and + # TRTLLM decode kernels surface cudaErrorNvlinkUncorrectable. Skip + # attention only inside that explicitly scoped dummy pass. Real + # requests and CUDA graph capture continue through TokenSpeed below. + if parallel.dcp_enabled and get_in_autotune_dummy_run(): + output = torch.zeros( + (q.shape[0], layer.tp_q_head_num * layer.v_head_dim), + dtype=self.q_data_type, + device=q.device, + ) + lse = torch.zeros( + (q.shape[0], layer.tp_q_head_num), + dtype=torch.float32, + device=q.device, + ) + return output, lse + if not parallel.dcp_enabled: return super().forward_decode( q, @@ -615,7 +639,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend): # (prepare_prefill_qkv / pack_prefix_chunk_kv); no quantize here. # Hybrid MLA models resolve the model-side hook through the outer # HybridLinearAttnBackend, so their fallback MHA path can pass V as a - # last-dimension slice of kv_b_proj (stride(-2) > size(-1)). The + # last-dimension slice of kv_b_proj (stride(-2) > size(-1)). The # TokenSpeed prefill kernel requires dense Q/K/V layouts even though # the public wrapper accepts arbitrary torch tensors. q = q.contiguous() diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 57e90f2b4..c378c80ab 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -465,11 +465,19 @@ class TritonAttnBackend(AttentionBackend): spec_info, ): """Fill all cuda-graph buffers for target_verify mode.""" + # Prefer the spec_info's per-request query length (DSpark draft propose + # uses gamma < verify window); fall back to the configured verify window. + num_draft_tokens = self.num_draft_tokens + if ( + spec_info is not None + and getattr(spec_info, "draft_token_num", None) is not None + ): + num_draft_tokens = int(spec_info.draft_token_num) qo_indptr = self.qo_indptr[: bs + 1] qo_indptr[: bs + 1] = torch.arange( 0, - (1 + bs) * self.num_draft_tokens, - step=self.num_draft_tokens, + (1 + bs) * num_draft_tokens, + step=num_draft_tokens, dtype=torch.int32, device=self.device, ) @@ -506,7 +514,7 @@ class TritonAttnBackend(AttentionBackend): custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask else: custom_mask = None - seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens) + seq_mask_len = num_draft_tokens * (seq_lens + num_draft_tokens) mask_indptr = self.mask_indptr[: bs + 1] mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0) return ( @@ -799,10 +807,18 @@ class TritonAttnBackend(AttentionBackend): max_extend_len = None elif forward_batch.forward_mode.is_target_verify(): bs = len(forward_batch.req_pool_indices) + # self.num_draft_tokens is the verify window (gamma + 1), while + # DSpark draft propose runs a gamma-token TARGET_VERIFY forward. + num_draft_tokens = self.num_draft_tokens + if ( + spec_info is not None + and getattr(spec_info, "draft_token_num", None) is not None + ): + num_draft_tokens = int(spec_info.draft_token_num) qo_indptr = torch.arange( 0, - (1 + bs) * self.num_draft_tokens, - step=self.num_draft_tokens, + (1 + bs) * num_draft_tokens, + step=num_draft_tokens, dtype=torch.int32, device=self.device, ) @@ -839,13 +855,13 @@ class TritonAttnBackend(AttentionBackend): ) custom_mask = spec_info.custom_mask - seq_mask_len = self.num_draft_tokens * ( - forward_batch.seq_lens + self.num_draft_tokens + seq_mask_len = num_draft_tokens * ( + forward_batch.seq_lens + num_draft_tokens ) mask_indptr = self.mask_indptr mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0) mask_indptr = mask_indptr[: bs + 1] - max_extend_len = self.num_draft_tokens + max_extend_len = num_draft_tokens num_kv_splits = None attn_logits = None attn_lse = None @@ -1093,10 +1109,16 @@ class TritonAttnBackend(AttentionBackend): and getattr(spec_info, "custom_mask", None) is not None else None ) + max_extend_len = self.num_draft_tokens + if ( + spec_info is not None + and getattr(spec_info, "draft_token_num", None) is not None + ): + max_extend_len = int(spec_info.draft_token_num) return ForwardMetadata( attn_logits=None, attn_lse=None, - max_extend_len=self.num_draft_tokens, + max_extend_len=max_extend_len, num_kv_splits=None, kv_indptr=self.kv_indptr[: bs + 1], kv_indices=self.cuda_graph_kv_indices, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 5c3a76a2e..2fbafa3e5 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -19,6 +19,20 @@ from sglang.kernels.ops.attention.pad import ( from sglang.kernels.ops.attention.pad import ( unpad_draft_extend_output as unpad_draft_extend_output_triton, ) +from sglang.kernels.ops.attention.set_mla_kv_concat_q import ( + can_use_set_mla_kv_concat_q, + can_use_set_mla_kv_concat_q_fp8, +) +from sglang.kernels.ops.attention.set_mla_kv_concat_q import ( + covered as set_mla_kv_concat_q_covered, +) +from sglang.kernels.ops.attention.set_mla_kv_concat_q import ( + covered_fp8 as set_mla_kv_concat_q_fp8_covered, +) +from sglang.kernels.ops.attention.set_mla_kv_concat_q import ( + set_mla_kv_concat_q, + set_mla_kv_concat_q_fp8, +) from sglang.kernels.ops.attention.utils import ( concat_mla_absorb_q_general, mla_quantize_and_rope_for_fp8, @@ -146,6 +160,13 @@ class TRTLLMMLAPrefillMetadata: fallback_to_flashinfer_impl: bool = False +from sglang.kernels.jit.utils import is_arch_support_pdl + +# Arm PDL on the trtllm-gen decode launch so its prolog overlaps the tail of +# the query-prep kernels (which already trigger their PDL secondary). +_ENABLE_PDL = is_arch_support_pdl() + + @dataclass class TRTLLMMLADecodeMetadata: """Metadata for TRTLLM MLA decode operations.""" @@ -167,6 +188,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): # read seq_lens_cpu / seq_lens_sum; opt out of the D2H sync. needs_cpu_seq_lens: bool = False + # Ragged verify: the packed query is front-aligned into the dense + # [bs, draft_token_num] layout in forward_extend; metadata stays uniform. + supports_ragged_verify_graph: bool = True + def __init__( self, model_runner: ModelRunner, @@ -266,6 +291,27 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): # write translates through the pool's _full_translate hook instead). self._decode_dense_loc: Optional[torch.Tensor] = None self.cuda_graph_out_cache_loc_dense: Optional[torch.Tensor] = None + # Fused KV-scatter + q-concat on the decode dense-loc path (one launch + # instead of set_mla_kv_buffer + concat_mla_absorb_q). Disabled under + # async asserts: the fused path writes the pool directly and would + # skip the pool's OOB probe. + self._fused_set_kv_concat_q = ( + self.data_type == torch.bfloat16 + and not envs.SGLANG_ENABLE_ASYNC_ASSERT.get() + and can_use_set_mla_kv_concat_q( + self.kv_lora_rank * 2, self.qk_rope_head_dim * 2 + ) + ) + # fp8 sibling: quantize + KV scatter + q concat in one launch + # (replaces mla_quantize_without_rope_for_fp8's concat + three aten + # casts plus the KV-row write on the fp8 decode path). + self._fused_set_kv_concat_q_fp8 = ( + self.data_type == torch.float8_e4m3fn + and not envs.SGLANG_ENABLE_ASYNC_ASSERT.get() + and self.kv_lora_rank == 512 + and self.qk_rope_head_dim == 64 + and can_use_set_mla_kv_concat_q_fp8() + ) def _calc_padded_blocks(self, max_seq_len: int) -> int: """ @@ -415,6 +461,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): device: torch.device, ): """Allocate persistent metadata buffers for CUDA graph capture.""" + if forward_mode.is_target_verify() and bs in self.decode_cuda_graph_metadata: + # Token tiers at the same slot count must share one per-bs buffer + # set (each graph bakes in the tensors it captured). + self.forward_decode_metadata = self.decode_cuda_graph_metadata[bs] + return metadata = TRTLLMMLADecodeMetadata() if forward_mode.is_target_verify(): @@ -438,6 +489,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ) metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device) + if metadata.seq_lens_k is None: + # Plain decode: static int32 seq_lens buffer, refreshed by the + # capture+replay body below (same pattern as target-verify). + metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device) + # Capture with full width so future longer sequences are safe during replay. max_blocks_per_seq = self._calc_padded_blocks(self.max_context_len) block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_blocks_per_seq] @@ -476,6 +532,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): metadata.sum_seq_lens_q = num_tokens_per_req * bs seq_lens = seq_lens[:bs] metadata.seq_lens_k.copy_(seq_lens) + elif metadata.seq_lens_k is not None: + # Plain decode: int64 -> int32 downcast copy into the static + # buffer (once per step, replacing the per-layer conversion). + metadata.seq_lens_k.copy_(seq_lens[:bs]) # Update block indices for new sequences. create_flashmla_kv_indices_triton[ @@ -643,6 +703,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): self.forward_decode_metadata.global_seq_lens_k = ( self.forward_decode_metadata.seq_lens_k ) + elif forward_batch.forward_mode.is_decode_or_idle(): + # One int32 conversion per step; forward_decode reads it back + # so the per-layer .to(int32) in _run_decode_kernel stays a + # no-op (24 elementwise copies/step otherwise). + self.forward_decode_metadata.seq_lens_k = seq_lens.to(torch.int32) elif forward_batch.forward_mode.is_draft_extend_v2(): sum_seq_lens_q = sum(forward_batch.extend_seq_lens_cpu) max_seq_len_q = max(forward_batch.extend_seq_lens_cpu) @@ -699,14 +764,32 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): cu_seqlens_q: torch.Tensor, seq_lens_q: torch.Tensor, sum_seq_lens_q: int, + zero_uncovered: bool = False, ) -> torch.Tensor: - """Unpad draft extended output using Triton kernel.""" + """Unpad draft extended output using Triton kernel. + + zero_uncovered: ragged verify's clamped rows leave output positions + unwritten; zero the destination so those discarded rows stay finite + (draft_extend writes every position and does not need this). + """ + output_buffer = self.unpad_output_buffer + if zero_uncovered: + if output_buffer is not None: + output_buffer[:sum_seq_lens_q].zero_() + else: + # No persistent buffer without cuda graph state; the triton + # wrapper's dynamic fallback is torch.empty. + output_buffer = torch.zeros( + (sum_seq_lens_q, raw_out.shape[2], raw_out.shape[3]), + dtype=raw_out.dtype, + device=raw_out.device, + ) return unpad_draft_extend_output_triton( raw_out, cu_seqlens_q, seq_lens_q, sum_seq_lens_q, - self.unpad_output_buffer, + output_buffer, ) def _compute_decode_bmm1_scale(self, layer: RadixAttention) -> float: @@ -738,8 +821,25 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): seq_lens: torch.Tensor, max_seq_len: int, layer: RadixAttention, + *, + causal_seqs: Optional[torch.Tensor] = None, + cp_world: int = 1, + cp_rank: int = 0, + return_lse: bool = False, ) -> torch.Tensor: - """Hook for subclasses to swap the decode/spec-verify kernel.""" + """Hook for subclasses to swap the decode/spec-verify kernel. + + The DCP arguments belong to the hook contract because forward_extend + passes them on the DCP target-verify path. This implementation does not + forward them to the kernel and returns no LSE, so only the DCP-capable + subclasses serve them.""" + if cp_world > 1 or return_lse: + raise NotImplementedError( + "trtllm_mla does not forward the cyclic DCP metadata to its " + "decode kernel and returns no rank-local LSE for the cross-rank " + "merge; select cutedsl_mla or tokenspeed_mla for a DCP " + "target-verify run" + ) # Scale computation for TRTLLM MLA kernel BMM1 operation: # The final BMM1 scale is computed as: q_scale * k_scale * softmax_scale @@ -761,6 +861,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): return flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( query=query, kv_cache=kv_cache, + enable_pdl=_ENABLE_PDL, workspace_buffer=self.workspace_buffer, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, @@ -818,6 +919,101 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(), ) + def _set_kv_and_concat_q_fused( + self, + layer: RadixAttention, + loc: torch.Tensor, + k: torch.Tensor, + k_rope: torch.Tensor, + q: torch.Tensor, + q_rope: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Decode: scatter the KV row at ``loc`` (already physical — the + dense-loc buffer on the unified pool, or out_cache_loc on the static + pool where ``_full_translate`` is identity) and build the + [q_nope | q_rope] fmha query in one kernel launch (saves one launch + per MLA layer and keeps the PDL chain intact). + + Returns the concatenated query, or None when the fused kernel does + not cover the inputs (caller falls back to the two-kernel path). + """ + k_nope_2d = k.view(k.shape[0], -1) + k_rope_2d = k_rope.view(k_rope.shape[0], -1) + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope_3d = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + # Same raw per-layer buffer the decode kernel reads below (bf16-only + # gate means store_dtype == dtype, so no view); get_key_buffer applies + # the hybrid pool's full-attention layer-id mapping. + kv_raw = self.token_to_kv_pool.get_key_buffer(layer.layer_id) + kv_2d = kv_raw.view(kv_raw.shape[0], -1) if kv_raw.dim() != 2 else kv_raw + if not set_mla_kv_concat_q_covered( + kv_buffer=kv_2d, + loc=loc, + k_nope=k_nope_2d, + k_rope=k_rope_2d, + q_nope=q_nope, + q_rope=q_rope_3d, + ): + return None + return set_mla_kv_concat_q( + kv_buffer=kv_2d, + loc=loc, + cache_k_nope=k_nope_2d, + cache_k_rope=k_rope_2d, + q_nope=q_nope, + q_rope=q_rope_3d, + ) + + def _set_kv_and_concat_q_fp8_fused( + self, + layer: RadixAttention, + loc: torch.Tensor, + q: torch.Tensor, + q_rope: torch.Tensor, + k: torch.Tensor, + k_rope: torch.Tensor, + ) -> Optional[torch.Tensor]: + """fp8-KV decode: quantize + scatter the KV row at ``loc`` (already + physical) and build the fp8 [q_nope | q_rope] query in one launch. + + Returns the fp8 query, or None when the fused kernel does not cover + the inputs (caller falls back to the aten quantize chain). + """ + k_nope_2d = k.view(k.shape[0], -1) + k_rope_2d = k_rope.view(k_rope.shape[0], -1) + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope_3d = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + # fp8 view of the pool's uint8 store; same buffer the decode kernel + # reads below. + kv_raw = self.token_to_kv_pool.get_key_buffer(layer.layer_id) + kv_2d = kv_raw.view(kv_raw.shape[0], -1) if kv_raw.dim() != 2 else kv_raw + if not set_mla_kv_concat_q_fp8_covered( + kv_buffer=kv_2d, + loc=loc, + k_nope=k_nope_2d, + k_rope=k_rope_2d, + q_nope=q_nope, + q_rope=q_rope_3d, + ): + return None + parallel = get_parallel() + return set_mla_kv_concat_q_fp8( + kv_buffer=kv_2d, + loc=loc, + cache_k_nope=k_nope_2d, + cache_k_rope=k_rope_2d, + q_nope=q_nope, + q_rope=q_rope_3d, + # DCP cyclic KV sharding: virtual loc -> owner mask + loc//world + # (identity when attn_dcp_size == 1). + dcp_world_size=parallel.attn_dcp_size, + dcp_rank=parallel.attn_dcp_rank, + ) + def forward_decode( self, q: torch.Tensor, # q_nope @@ -834,12 +1030,33 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ) -> torch.Tensor: """Run forward for decode using TRTLLM MLA kernel.""" merge_query = q_rope is not None + fused_fp8_query = None if self.data_type == torch.float8_e4m3fn: assert q_rope is not None and k_rope is not None if cos_sin_cache is None: - q, k, k_rope = mla_quantize_without_rope_for_fp8( - q, q_rope, k.squeeze(1), k_rope.squeeze(1) - ) + if save_kv_cache and self._fused_set_kv_concat_q_fp8: + loc = ( + self._decode_dense_loc + if self._decode_dense_loc is not None + else ( + None if self._unified_mla else forward_batch.out_cache_loc + ) + ) + if loc is not None: + # Fused: bf16->fp8 quantize + KV scatter + q concat + # in one launch; None when not covered. + fused_fp8_query = self._set_kv_and_concat_q_fp8_fused( + layer=layer, + loc=loc, + q=q, + q_rope=q_rope, + k=k, + k_rope=k_rope, + ) + if fused_fp8_query is None: + q, k, k_rope = mla_quantize_without_rope_for_fp8( + q, q_rope, k.squeeze(1), k_rope.squeeze(1) + ) else: q, k, k_rope = mla_quantize_and_rope_for_fp8( q, @@ -854,34 +1071,65 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ) merge_query = False - # Save KV cache if requested - if save_kv_cache: + # Save KV cache if requested (the fused fp8 path already wrote it) + query = fused_fp8_query + if query is None and save_kv_cache: assert ( k is not None and k_rope is not None ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." if self._decode_dense_loc is not None: # cuda-graph path: dense write loc precomputed out-of-graph, so # the in-graph write captures no translate allocation. - self.token_to_kv_pool.set_mla_kv_buffer( - layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True - ) + if merge_query and self._fused_set_kv_concat_q: + # Fused: KV scatter + [q_nope | q_rope] concat in one + # launch; None when the inputs are not covered. + query = self._set_kv_and_concat_q_fused( + layer=layer, + loc=self._decode_dense_loc, + k=k, + k_rope=k_rope, + q=q, + q_rope=q_rope, + ) + if query is None: + self.token_to_kv_pool.set_mla_kv_buffer( + layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True + ) else: # eager (or static pool): the pool's _full_translate handles it. - self.token_to_kv_pool.set_mla_kv_buffer( - layer, forward_batch.out_cache_loc, k, k_rope - ) + if ( + merge_query + and self._fused_set_kv_concat_q + and not self._unified_mla + ): + # Static pool: _full_translate is identity, so + # out_cache_loc is already the physical write loc. + query = self._set_kv_and_concat_q_fused( + layer=layer, + loc=forward_batch.out_cache_loc, + k=k, + k_rope=k_rope, + q=q, + q_rope=q_rope, + ) + if query is None: + self.token_to_kv_pool.set_mla_kv_buffer( + layer, forward_batch.out_cache_loc, k, k_rope + ) - # Prepare query tensor inline - if merge_query: - # For FP16 path, we merge the query and rope parts into a single tensor - q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) - q_rope_reshaped = q_rope.view( - -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim - ) - query = concat_mla_absorb_q_general(q_nope, q_rope_reshaped) - else: - # For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function - query = q.view(-1, layer.tp_q_head_num, layer.head_dim) + # Prepare query tensor inline (already built when the fused save-KV + # path ran) + if query is None: + if merge_query: + # For FP16 path, we merge the query and rope parts into a single tensor + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope_reshaped = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + query = concat_mla_absorb_q_general(q_nope, q_rope_reshaped) + else: + # For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function + query = q.view(-1, layer.tp_q_head_num, layer.head_dim) # Apply llama 4 scaling if provided if llama_4_scaling is not None: @@ -915,7 +1163,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): query=query, kv_cache=kv_cache, block_tables=metadata.block_kv_indices, - seq_lens=forward_batch.seq_lens, + seq_lens=( + metadata.seq_lens_k + if metadata.seq_lens_k is not None + else forward_batch.seq_lens + ), max_seq_len=metadata.max_seq_len_k, layer=layer, ) @@ -1038,9 +1290,38 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): max_seq_len = metadata.max_seq_len_k + ( 0 if dcp_enabled else draft_token_num ) - # For target_verify, all sequences have the same number of draft tokens - q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) - needs_unpad = False + ragged_layout = forward_batch.spec_info.ragged_verify_layout + if ragged_layout is None or dcp_enabled: + q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) + needs_unpad = False + else: + if ragged_layout.bs != bs: + # Capped variant: the dense [bs, draft_token_num] q + # buffer below cannot take a row the full-coverage pad + # may inflate past the verify window, and it keeps + # qo_indptr consistent with the clamped lens (same + # contract as the KDA dense path). + ragged_layout = ragged_layout.padded_to_bucket( + padded_bs=bs, cap=draft_token_num + ) + total_tokens = q.shape[0] + seq_lens_q = torch.clamp( + ragged_layout.verify_lens, max=draft_token_num + ) + cu_seqlens_q = ragged_layout.qo_indptr_device + padded_q = torch.zeros( + (bs, draft_token_num, layer.tp_q_head_num, layer.head_dim), + dtype=q.dtype, + device=q.device, + ) + q = self.pad_draft_extend_query( + q, padded_q, seq_lens_q, cu_seqlens_q + ) + needs_unpad = True + unpad_zero_uncovered = True + unpad_seq_lens_q = seq_lens_q + unpad_cu_seqlens_q = cu_seqlens_q + unpad_sum_seq_lens_q = total_tokens else: # draft_extend: handle varying num_correct_drafts_per_req. If total_tokens % bs == 0, # we can directly reshape q; otherwise, pad to max_seq_len_q. @@ -1084,6 +1365,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): q, padded_q, actual_seq_lens_q, actual_cu_seqlens_q ) needs_unpad = True + unpad_zero_uncovered = False unpad_seq_lens_q = actual_seq_lens_q unpad_cu_seqlens_q = actual_cu_seqlens_q unpad_sum_seq_lens_q = total_tokens @@ -1145,6 +1427,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): unpad_cu_seqlens_q, unpad_seq_lens_q, unpad_sum_seq_lens_q, + zero_uncovered=unpad_zero_uncovered, ) output = output.view(-1, layer.tp_q_head_num * layer.v_head_dim) else: diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 7f50be156..0d69dea7d 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -169,6 +169,43 @@ def prepare_vision_attention_metadata( ) +def prepare_flashinfer_cudnn_vision_attention_metadata( + cu_seqlens: torch.Tensor, + device: torch.device, + *, + elem_per_token: int, +) -> VisionAttentionMetadata: + cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32, non_blocking=True) + seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] + batch_size = int(seq_lens.numel()) + padded_batch_size = next( + (size for size in BATCH_BUCKETS if size >= batch_size), + math.ceil(batch_size / BATCH_BUCKETS[0]) * BATCH_BUCKETS[0], + ) + if padded_batch_size != batch_size: + pad_size = padded_batch_size - batch_size + padded_indptrs = torch.cat([cu_seqlens, cu_seqlens[-1].expand(pad_size)]) + padded_seq_lens = torch.cat([seq_lens, seq_lens.new_zeros(pad_size)]) + else: + padded_indptrs = cu_seqlens + padded_seq_lens = seq_lens + + elem_indptrs = padded_indptrs * elem_per_token + real_max_seqlen = int(seq_lens.max().item()) + max_seqlen = next( + (size for size in FLASHINFER_MAX_SEQLEN_BUCKETS if size >= real_max_seqlen), + math.ceil(real_max_seqlen / FLASHINFER_MAX_SEQLEN_BUCKETS[-1]) + * FLASHINFER_MAX_SEQLEN_BUCKETS[-1], + ) + return prepare_vision_attention_metadata( + cu_seqlens, + device=device, + packed_indptrs=torch.cat([elem_indptrs] * 3), + sequence_lengths=padded_seq_lens.view(-1, 1, 1, 1), + flashinfer_max_seqlen=max_seqlen, + ) + + # TODO: requires real seqlens from images @functools.lru_cache(maxsize=128) def _get_cu_seqlens_for_shape(batch_size: int, seqlen: int, device) -> torch.Tensor: diff --git a/python/sglang/srt/layers/attn_residual.py b/python/sglang/srt/layers/attn_residual.py new file mode 100644 index 000000000..9db79ca8b --- /dev/null +++ b/python/sglang/srt/layers/attn_residual.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +# Kimi-K3 Attention Residual: snapshot bank + aggregation. +# +# The public API is the AttnResidual class (constructed once per forward pass). +# It owns the frozen snapshot bank [T, NB, H] and the valid-row counter, and +# dispatches each aggregation point (score rows → softmax → weighted sum → +# RMSNorm) by hardware capability: +# fast — warp-specialized TMA kernel: cp.async.bulk producer + +# online-softmax consumers over a double-buffered chunk ring, out +# norm fused, per-nvb tuned launch config, one persistent CTA per +# SM. Taken on SM100+ with H=7168. +# fused — Triton 2-kernel pipeline with full H-parallelism; the fallback +# everywhere the fast kernel does not apply. +# aggregate_stream_torch is the eager reference (tests and the +# H % _BLOCK_H != 0 shape fallback of aggregate_stream). + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ReplicatedLinear + +_BLOCK_H: int = 1024 # H = 7168 = 7 x 1024 +_MAX_ROWS: int = 16 # next_pow2(8 + 1), K3 has <= 8 snapshots + +_FAST_SUPPORTED = None + + +def _use_fast(hidden_size: int) -> bool: + """The TMA kernel needs SM100+ (tcgen05, cp.async.bulk) and its H=7168 + template instantiation; everything else takes the triton pipeline.""" + global _FAST_SUPPORTED + if _FAST_SUPPORTED is None: + major, _ = torch.cuda.get_device_capability() + _FAST_SUPPORTED = major >= 10 + return _FAST_SUPPORTED and hidden_size == 7168 + + +def get_cw( + proj: ReplicatedLinear, + norm: RMSNorm, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Cached product norm_weight ⊙ proj_weight (both [H]) in `dtype`. + + Cached per dtype: the fast kernel consumes bf16 while the triton path + consumes fp32, and a shared slot would hand one path the other's dtype.""" + cache = getattr(proj, "_attn_res_cw_cache", None) + if cache is None: + cache = {} + proj._attn_res_cw_cache = cache + cw = cache.get(dtype) + if cw is None: + cw = (norm.weight.float() * proj.weight.squeeze().float()).contiguous() + cw = cache[dtype] = cw.to(dtype) + return cw + + +def _aggregate_fast( + prefix_sum: torch.Tensor, + bank: torch.Tensor, + nvb: int, + score_proj: ReplicatedLinear, + score_norm: RMSNorm, + out_norm: RMSNorm, + write_bank_row: bool = False, +) -> torch.Tensor: + """Warp-specialized TMA kernel: online softmax over row chunks with the + output RMSNorm fused, one persistent CTA per SM, per-nvb tuned launch + config (GB300 benchmark winner across nvb). With write_bank_row the kernel + also snapshots the prefix row into bank[:, nvb, :] (bit-exact, zero extra + reads — the row streams through the score pass anyway).""" + from sglang.kernels.ops.kimi_k3.attn_res import attn_res_fused_tma + + # The kernel applies one eps to both the score norm and the output norm. + assert score_norm.variance_epsilon == out_norm.variance_epsilon + + cw = get_cw(score_proj, score_norm, dtype=torch.bfloat16) + out = torch.empty_like(prefix_sum) + attn_res_fused_tma( + prefix_sum, + bank, + cw, + out_norm.weight, + out, + nvb, + score_norm.variance_epsilon, + write_prefix=write_bank_row, + ) + return out + + +@triton.jit +def _score_kernel( + prefix_ptr, # [T, H] + bank_ptr, # [T, NB_total, H] + cw_ptr, # [H] fp32 + scores_ptr, # [T, MAX_ROWS] fp32 + NVB, + eps, + stride_pm, + stride_bm, + stride_bb, + stride_sm, + H: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """One CTA per (token, row): scan H, output one scalar score.""" + pid_t = tl.program_id(0) + j = tl.program_id(1) + if j > NVB: + return + sumsq = 0.0 + dotv = 0.0 + for h0 in tl.static_range(0, H, BLOCK_H): + offs_h = h0 + tl.arange(0, BLOCK_H) + if j < NVB: + v = tl.load(bank_ptr + pid_t * stride_bm + j * stride_bb + offs_h).to( + tl.float32 + ) + else: + v = tl.load(prefix_ptr + pid_t * stride_pm + offs_h).to(tl.float32) + cw = tl.load(cw_ptr + offs_h) + sumsq += tl.sum(v * v) + dotv += tl.sum(v * cw) + rrms = 1.0 / tl.sqrt(sumsq / H + eps) + tl.store(scores_ptr + pid_t * stride_sm + j, dotv * rrms) + + +@triton.jit +def _combine_kernel( + prefix_ptr, + bank_ptr, + scores_ptr, # [T, MAX_ROWS] fp32 + out_ptr, # [T, H] + NVB, + stride_pm, + stride_bm, + stride_bb, + stride_sm, + stride_om, + BLOCK_H: tl.constexpr, + MAX_ROWS: tl.constexpr, +): + """One CTA per (token, H-chunk): softmax(scores) → weighted sum → write chunk. + + Softmax is redundantly computed by each H-chunk CTA (≤16 elements, trivial). + This gives full H-parallelism: 7 CTAs for H=7168/1024. + """ + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + h0 = pid_h * BLOCK_H + + # Softmax (redundant per chunk, 16 fp32 ops) + offs_b = tl.arange(0, MAX_ROWS) + mask_b = offs_b <= NVB + raw = tl.load( + scores_ptr + pid_t * stride_sm + offs_b, mask=mask_b, other=float("-inf") + ) + m = tl.max(raw, axis=0) + e = tl.where(mask_b, tl.exp(raw - m), 0.0) + p = e / tl.sum(e, axis=0) + + # Weighted sum for this H chunk + offs_h = h0 + tl.arange(0, BLOCK_H) + acc = tl.zeros([BLOCK_H], tl.float32) + for j in range(0, NVB + 1): + if j < NVB: + v = tl.load(bank_ptr + pid_t * stride_bm + j * stride_bb + offs_h).to( + tl.float32 + ) + else: + v = tl.load(prefix_ptr + pid_t * stride_pm + offs_h).to(tl.float32) + p_j = tl.sum(tl.where(offs_b == j, p, 0.0), axis=0) + acc += p_j * v + tl.store( + out_ptr + pid_t * stride_om + offs_h, + acc.to(out_ptr.dtype.element_ty), + ) + + +def _mix_fused( + prefix_sum: torch.Tensor, + bank: torch.Tensor, + nvb: int, + score_proj: ReplicatedLinear, + score_norm: RMSNorm, +) -> torch.Tensor: + """Triton score + combine pair: returns the pre-norm mixture.""" + T, H = prefix_sum.shape + cw = get_cw(score_proj, score_norm) + n_h_blocks = H // _BLOCK_H + + # Step 1: score each row (2D grid, full row-parallelism) + scores = torch.empty((T, _MAX_ROWS), dtype=torch.float32, device=prefix_sum.device) + _score_kernel[(T, nvb + 1)]( + prefix_sum, + bank, + cw, + scores, + nvb, + score_norm.variance_epsilon, + prefix_sum.stride(0), + bank.stride(0), + bank.stride(1), + scores.stride(0), + H=H, + BLOCK_H=_BLOCK_H, + num_warps=8, + ) + + # Step 2: softmax + weighted sum (2D grid, full H-parallelism) + out = torch.empty_like(prefix_sum) + _combine_kernel[(T, n_h_blocks)]( + prefix_sum, + bank, + scores, + out, + nvb, + prefix_sum.stride(0), + bank.stride(0), + bank.stride(1), + scores.stride(0), + out.stride(0), + BLOCK_H=_BLOCK_H, + MAX_ROWS=_MAX_ROWS, + num_warps=4, + ) + return out + + +def _aggregate_fused( + prefix_sum: torch.Tensor, + bank: torch.Tensor, + nvb: int, + score_proj: ReplicatedLinear, + score_norm: RMSNorm, + out_norm: RMSNorm, +) -> torch.Tensor: + # Step 3: standard RMSNorm (sglang's optimized kernel) + return out_norm(_mix_fused(prefix_sum, bank, nvb, score_proj, score_norm)) + + +def aggregate_stream_torch( + prefix_sum: torch.Tensor, + bank: torch.Tensor, + nvb: int, + score_proj: ReplicatedLinear, + score_norm: RMSNorm, +) -> torch.Tensor: + """Eager reference for aggregate_stream (materializes [T, R, H]).""" + if nvb == 0: + return prefix_sum + T, H = prefix_sum.shape + # rows = [bank[0..nvb-1], prefix_sum] shape [T, nvb+1, H] + rows = torch.cat([bank[:, :nvb, :], prefix_sum.unsqueeze(1)], dim=1) + R = nvb + 1 + normed = score_norm(rows.reshape(T * R, H)) + scores = score_proj(normed)[0].reshape(T, R) + probs = torch.softmax(scores.float(), dim=-1) + mixed = (probs.unsqueeze(-1) * rows.float()).sum(dim=1) + return mixed.to(prefix_sum.dtype) + + +def aggregate_stream( + prefix_sum: torch.Tensor, + bank: torch.Tensor, + nvb: int, + score_proj: ReplicatedLinear, + score_norm: RMSNorm, +) -> torch.Tensor: + """Pre-norm aggregated stream value (softmax mixture, no output norm): + the K3 analogue of the residual stream, for dspark aux capture -- the + raw wire only carries the current block's running prefix.""" + if nvb == 0: + return prefix_sum + if prefix_sum.shape[1] % _BLOCK_H != 0: + return aggregate_stream_torch(prefix_sum, bank, nvb, score_proj, score_norm) + return _mix_fused(prefix_sum, bank, nvb, score_proj, score_norm) + + +def _aggregate_fused_add( + prefix_a: torch.Tensor, + prefix_b: torch.Tensor, + bank: torch.Tensor, + nvb: int, + score_proj: ReplicatedLinear, + score_norm: RMSNorm, + out_norm: RMSNorm, + write_bank_row: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Aggregation point with a pending upstream residual add: materialize + prefix = prefix_a + prefix_b, then aggregate. Returns (normed, prefix). + write_bank_row rides _aggregate (fast path only).""" + prefix = prefix_a + prefix_b + return ( + _aggregate( + prefix, + bank, + nvb, + score_proj, + score_norm, + out_norm, + write_bank_row=write_bank_row, + ), + prefix, + ) + + +def _aggregate( + prefix_sum: torch.Tensor, + bank: torch.Tensor, + nvb: int, + score_proj: ReplicatedLinear, + score_norm: RMSNorm, + out_norm: RMSNorm, + write_bank_row: bool = False, +) -> torch.Tensor: + """Single aggregation point: score → softmax → mix → norm. + + Caller handles nvb == 0 (layer 0 attn side: just out_norm(prefix_sum)). + write_bank_row is fast-path only (in-kernel snapshot of the prefix row + into bank[:, nvb, :]); the triton path keeps the standalone .write() copy — + the caller (AttnResidual.forward) owns that fallback. + """ + if _use_fast(prefix_sum.shape[1]): + return _aggregate_fast( + prefix_sum, + bank, + nvb, + score_proj, + score_norm, + out_norm, + write_bank_row=write_bank_row, + ) + assert not write_bank_row, "fused bank write is fast-path only" + return _aggregate_fused(prefix_sum, bank, nvb, score_proj, score_norm, out_norm) + + +class AttnResidual: + """Snapshot bank + aggregation of one K3 attention-residual stream, + backed by the capability-dispatched kernels above. + + One instance lives for one model forward pass. + """ + + def __init__( + self, + hidden_states: torch.Tensor, + block_num: int, + block_residual: Optional[torch.Tensor] = None, + ) -> None: + num_tokens, hidden_size = hidden_states.shape + # Frozen snapshot rows [T, NB, H]; raw tensor for PP transfer and the + # legacy kernel path. + self.block_residual = hidden_states.new_empty( + (num_tokens, block_num, hidden_size) + ) + self.num_valid_blocks = 0 + if block_residual is not None: # inherited from the previous PP rank + self.num_valid_blocks = block_residual.size(1) + self.block_residual[:, : self.num_valid_blocks, :].copy_(block_residual) + + def write(self, prefix_sum: torch.Tensor, rows: Optional[slice] = None) -> None: + """Snapshot the pre-attention prefix into the next bank row. + + Under SP attention-residual carry each rank owns a disjoint token + slice, so only that slice is written and subsequently read locally. + """ + bank = self.block_residual if rows is None else self.block_residual[rows] + bank[:, self.num_valid_blocks, :].copy_(prefix_sum) + self.num_valid_blocks += 1 + + def forward( + self, + hidden_states: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + score_proj: ReplicatedLinear, + score_norm: RMSNorm, + out_norm: RMSNorm, + rows: Optional[slice] = None, + write: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Aggregate; with write=True also snapshot the aggregated prefix + (the second return value) into the next bank row — fused into the + fast kernel (the row streams through its score pass anyway), a + standalone .write() copy on every other path.""" + nvb = self.num_valid_blocks + # Layer 0 attention side: nothing banked yet + if nvb == 0: + assert prefix_sum is None + if write: + self.write(hidden_states, rows) + return out_norm(hidden_states), hidden_states + + # SP-MoE: the caller holds only its token shard; align the banked + # residual rows to it (dim-0 slice of a contiguous buffer stays + # contiguous for the jit kernels). + block_residual = ( + self.block_residual if rows is None else self.block_residual[rows] + ) + + fused_write = write and _use_fast(hidden_states.shape[1]) + if prefix_sum is None: + # hidden_states already is the whole head (PP entry or a + # block-boundary restart). + normed = _aggregate( + hidden_states, + block_residual, + nvb, + score_proj, + score_norm, + out_norm, + write_bank_row=fused_write, + ) + prefix = hidden_states + else: + # Pending add: materialize the prefix, then aggregate. + normed, prefix = _aggregate_fused_add( + prefix_sum, + hidden_states, + block_residual, + nvb, + score_proj, + score_norm, + out_norm, + write_bank_row=fused_write, + ) + if fused_write: + self.num_valid_blocks += 1 # row nvb written in-kernel + elif write: + self.write(prefix, rows) + return normed, prefix + + def forward_sp_all_gather( + self, + hidden_states: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + score_proj: ReplicatedLinear, + score_norm: RMSNorm, + out_norm: RMSNorm, + rows: slice, + write: bool = False, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Fuse a local aggregation point and the following row all-gather.""" + nvb = self.num_valid_blocks + if nvb == 0: + return None + if prefix_sum is not None and prefix_sum.shape != hidden_states.shape: + prefix_sum = prefix_sum[rows] + prefix = hidden_states if prefix_sum is None else prefix_sum.add(hidden_states) + bank = self.block_residual[rows] + cw = get_cw(score_proj, score_norm, dtype=torch.bfloat16) + assert score_norm.variance_epsilon == out_norm.variance_epsilon + from sglang.srt.layers import k3_sp_collective + + normed = k3_sp_collective.attn_res_all_gather( + prefix, + bank, + cw, + out_norm.weight, + nvb, + score_norm.variance_epsilon, + write_prefix=write, + ) + if normed is None: + return None + if write: + self.num_valid_blocks += 1 + return normed, prefix + + def forward_sp_reduce_scatter( + self, + hidden_states: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + score_proj: ReplicatedLinear, + score_norm: RMSNorm, + out_norm: RMSNorm, + rows: slice, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Fuse o_proj RS, the pending local prefix add, and aggregation.""" + nvb = self.num_valid_blocks + if nvb == 0: + return None + local_tokens = rows.stop - rows.start + residual = prefix_sum + if residual is not None and residual.shape[0] != local_tokens: + residual = residual[rows] + bank = self.block_residual[rows] + cw = get_cw(score_proj, score_norm, dtype=torch.bfloat16) + assert score_norm.variance_epsilon == out_norm.variance_epsilon + from sglang.srt.layers import k3_sp_collective + + return k3_sp_collective.reduce_scatter_attn_res( + hidden_states, + residual, + bank, + cw, + out_norm.weight, + nvb, + score_norm.variance_epsilon, + ) diff --git a/python/sglang/srt/layers/k3_ar_fusion.py b/python/sglang/srt/layers/k3_ar_fusion.py new file mode 100644 index 000000000..c15eb6c11 --- /dev/null +++ b/python/sglang/srt/layers/k3_ar_fusion.py @@ -0,0 +1,386 @@ +"""K3 MNNVL fused all-reduce dispatch. + +Auto-enabled on SM100/SM103 when CustomAllReduceV2 with multicast is +available; ``SGLANG_K3_AR_FUSION`` overrides in either direction (0 = off, +1 = attempt anywhere and warn when unavailable). + +Glue between the model and the ``kernels.ops.kimi_k3.all_reduce`` kernels. Small +messages take the 1shot multicast-push through the group's CustomAllReduceV2 +workspace; large ones take the in-place NVLS 2shot, which needs its input to be a +:func:`symm_buffer` slice and otherwise falls back to the regular all-reduce. + +Call-site contract when the fusion is active: ``enabled()`` was +checked once (which initializes the state), ``x`` is bf16 and contiguous, +and the residual is identical on every rank (a fully reduced tensor such +as the attn-res prefix sum) or ``None``. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, List, NamedTuple, Optional + +import torch + +import sglang.srt.runtime_context as ctx +from sglang.kernels.jit.utils import cache_once +from sglang.srt.environ import envs + +if TYPE_CHECKING: + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + +logger = logging.getLogger(__name__) + +# push wins below ~0.5 MB on B200x8/GB300 (bs<=32 @ H=7168); NVLS 2shot wins +# above (measured: push 8.5us vs pull 11.0us at 448KB, 15.5 vs 11.0 at 896KB) +_PUSH_MAX_BYTES = 512 * 1024 + +# Latent width of the K3 latent|shared MoE buffer the fused-norm AR expects +# ([N, NORM_DIM] latent then [N, 2*NORM_DIM] shared). MUST match kNormDim in +# csrc/kimi_k3/comm/ar_fusion.cuh — the mod hardcodes this row width. +NORM_DIM = 3584 + +# :func:`symm_buffer` names: o_proj's [rows, hidden_size] output, and the MoE +# [latent | shared] pair flattened to rows x (moe_hidden_size + hidden_size). +ATTN_O_PROJ = "attn_o_proj" +MOE_LATENT_SHARED = "moe_latent_shared" + + +class _State(NamedTuple): + comm: CustomAllReduceV2 + world_size: int + group_name: str + + +@cache_once +def _get_state() -> Optional[_State]: + explicit = envs.SGLANG_K3_AR_FUSION.is_set() + if explicit: + if not envs.SGLANG_K3_AR_FUSION.get(): + return None + else: + # Auto-probe: only SM100/SM103, and not under --enable-symm-mem or EP + # a2a (their allocator contexts misroute the buffers the pull needs). + from sglang.srt.utils.common import get_device_sm + + if get_device_sm() not in (100, 103): + return None + server_args = ctx.get_server_args() + if server_args.enable_symm_mem or server_args.moe_a2a_backend != "none": + logger.info( + "K3 all-reduce fusion auto-probe: skipping " + "(enable_symm_mem=%s, moe_a2a_backend=%s; under symm-mem the " + "allocator contexts conflict, and under EP a2a the model's " + "symm-pool allocation contract does not hold on every AR " + "call-site. Set SGLANG_K3_AR_FUSION=1 to force.)", + server_args.enable_symm_mem, + server_args.moe_a2a_backend, + ) + return None + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + from sglang.srt.distributed.parallel_state import get_tp_group + from sglang.srt.runtime_context import get_parallel + + if get_parallel().tp_size <= 1: + return None + group = get_tp_group() + comm = group.ca_comm + if ( + not isinstance(comm, CustomAllReduceV2) + or comm.disabled + or comm.mc_base_ptr == 0 + ): + if explicit: + logger.warning( + "SGLANG_K3_AR_FUSION requested but CustomAllReduceV2 with multicast " + "is unavailable; falling back to the regular all-reduce path." + ) + else: + logger.info( + "K3 all-reduce fusion auto-probe: CustomAllReduceV2 with " + "multicast is unavailable; using the regular all-reduce path." + ) + return None + from sglang.kernels.ops.kimi_k3 import all_reduce as mod + + mod.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr) + logger.info("K3 all-reduce fusion enabled (world_size=%d)", comm.world_size) + return _State(comm, comm.world_size, group.cpu_group.group_name) + + +def enabled() -> bool: + return _get_state() is not None + + +class _Buffer(NamedTuple): + region: torch.Tensor # flat [max_rows * width] + base: int # local VA of region[0] + mc_base: int # multicast VA of region[0] + max_rows: int + width: int + + +# Reverse pointer -> multicast lookup for get_mc_ptr(). +_BUFS: List[_Buffer] = [] + + +@cache_once +def _max_buffer_rows() -> int: + """Rows to reserve per buffer: the largest batch the server args allow.""" + server_args = ctx.get_server_args() + chunked = server_args.chunked_prefill_size + if chunked is not None and chunked > 0: + return int(chunked) + return int(server_args.max_prefill_tokens or 0) + + +def _create_buffer( + name: str, width: int, dtype: torch.dtype, group_name: str +) -> _Buffer: + """One symmetric buffer, rendezvoused once. ``empty_strided_p2p`` skips the + caching allocator, so this is right for a persistent buffer only.""" + from torch._C._distributed_c10d import _SymmetricMemory + + max_rows = _max_buffer_rows() + # outside inference mode on purpose: created inside it the buffer would be an + # inference tensor, and the in-place write from a no_grad path (CUDA graph + # capture) is then rejected + with torch.inference_mode(False): + region = _SymmetricMemory.empty_strided_p2p( + (max_rows * width,), + [1], + dtype, + torch.device("cuda", torch.cuda.current_device()), + group_name, + ).view(max_rows, width) + handle = _SymmetricMemory.rendezvous(region) + assert handle is not None and handle.multicast_ptr != 0 + buf = _Buffer( + region=region, + base=int(region.data_ptr()), + mc_base=int(handle.multicast_ptr), + max_rows=max_rows, + width=width, + ) + _BUFS.append(buf) + logger.info( + "K3 symm buffer %r: %d rows x %d, %.1f MB (rendezvoused once on first " + "use; no per-forward symmetric allocation)", + name, + max_rows, + width, + region.numel() * region.element_size() / (1024 * 1024), + ) + return buf + + +def symm_buffer( + name: str, + rows: int, + width: int, + dtype: torch.dtype, + group_name: Optional[str] = None, +): + """``[rows, width]`` view of a named persistent symmetric buffer. + + The pull reduces in place on the input's multicast alias, so every rank must + resolve the same ``(buffer, offset)``; a per-forward symm-pool allocation does + not, because torch's caching allocator breaks ties on the absolute address and + ``rendezvous`` validates sizes only. + + ``group_name`` defaults to the TP group, the one the fused all-reduce reduces + over; a caller reducing over another group (SP-MoE's attention TP group) passes + its own. Each name belongs to one group, so the name alone identifies it. + """ + if group_name is None: + from sglang.srt.distributed.parallel_state import get_tp_group + + group_name = get_tp_group().cpu_group.group_name + buf: _Buffer = ctx.get_buffer( + f"k3_symm:{name}", lambda: _create_buffer(name, width, dtype, group_name) + ) + assert 0 < rows <= buf.max_rows and width == buf.width and dtype == buf.region.dtype + return buf.region[:rows] + + +def find_mc_ptr(x: torch.Tensor) -> Optional[int]: + """Multicast VA of ``x`` if it lies inside a named buffer, else None.""" + ptr = int(x.data_ptr()) + end = ptr + x.numel() * x.element_size() + for buf in _BUFS: + if buf.base <= ptr and end <= buf.base + buf.region.nbytes: + return buf.mc_base + (ptr - buf.base) + return None + + +def get_mc_ptr(x: torch.Tensor) -> int: + """Multicast VA of ``x``, which must be a :func:`symm_buffer` slice.""" + mc = find_mc_ptr(x) + assert mc is not None, "K3 fused pull input is not a symm_buffer slice" + return mc + + +def all_reduce( + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """In-place ``x = allreduce(x) [+ residual]``; returns ``x``.""" + from sglang.kernels.ops.kimi_k3 import all_reduce as mod + + state = _get_state() + assert state is not None + if x.shape[0] == 0: + return x if residual is None else x + residual + nbytes = x.numel() * 2 + if nbytes <= min(_PUSH_MAX_BYTES, state.comm.max_push_size): + return mod.all_reduce_push_res( + state.world_size, x, residual, ws_mc_base=state.comm.mc_base_ptr + ) + return mod.all_reduce_pull_res( + state.world_size, x, residual, input_mc_ptr=get_mc_ptr(x) + ) + + +def all_reduce_low_sm( + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + *, + num_blocks: Optional[int] = None, + unroll: Optional[int] = None, +) -> torch.Tensor: + """In-place ``x = allreduce(x) [+ residual]``, pinned to the low-SM NVLS pull. + + For side-stream use: push would fan out over many blocks and steal SMs from + the GEMMs it overlaps. ``num_blocks`` / ``unroll`` default to the tuned tables. + """ + from sglang.kernels.ops.kimi_k3 import all_reduce as mod + + state = _get_state() + assert state is not None + if x.shape[0] == 0: + return x if residual is None else x + residual + return mod.all_reduce_pull_res( + state.world_size, + x, + residual, + input_mc_ptr=get_mc_ptr(x), + num_blocks=num_blocks, + unroll=unroll, + ) + + +def all_reduce_norm( + x: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + *, + num_tokens: int, +) -> torch.Tensor: + """In-place ``x = allreduce(x)`` with a fused RMSNorm over the first + ``num_tokens`` rows of the 2D ``[rows, NORM_DIM]`` input.""" + from sglang.kernels.ops.kimi_k3 import all_reduce as mod + + state = _get_state() + assert state is not None + if x.shape[0] == 0: + return x + nbytes = x.numel() * 2 + if nbytes <= min(_PUSH_MAX_BYTES, state.comm.max_push_size): + return mod.all_reduce_push_norm( + state.world_size, + x, + weight, + eps, + num_norm_rows=num_tokens, + ws_mc_base=state.comm.mc_base_ptr, + ) + return mod.all_reduce_pull_norm( + state.world_size, + x, + weight, + eps, + num_norm_rows=num_tokens, + input_mc_ptr=get_mc_ptr(x), + ) + + +def finalize_push_fits(num_tokens: int) -> bool: + """Whether a [num_tokens, NORM_DIM] latent fits the push window; the + finalize-fused AR is push-only.""" + state = _get_state() + assert state is not None + nbytes = num_tokens * NORM_DIM * 2 + return nbytes <= min(_PUSH_MAX_BYTES, state.comm.max_push_size) + + +def gemm_ag_up_fits(num_tokens: int) -> bool: + """Whether :func:`gemm_ag_up_proj` covers this decode batch: TP8, within the + kernel's win range, and the staging slice fits a push slot.""" + from sglang.kernels.ops.kimi_k3 import gemm_ag as mod + + state = _get_state() + assert state is not None + comm = state.comm + return ( + 0 < num_tokens <= mod.MAX_TOKENS + and state.world_size == 8 + and num_tokens * (2 * NORM_DIM // 8) * 2 <= comm.max_push_size + # The GEMV producer reads a single phase counter, so its grid is not + # bound to the counter array; the spin consumer still needs one block + # per counter slot plus the cleanup block. + and comm.config.num_push_blocks >= 2 + ) + + +def gemm_ag_up_proj( + x: torch.Tensor, + weight: torch.Tensor, + b: torch.Tensor, + c: Optional[torch.Tensor], +) -> torch.Tensor: + """``up_proj(x) + b (+ c)`` via column-parallel GEMV + multicast all-gather + + fused add3. ``weight`` is the FULL replicated up_proj (the kernel slices this + rank's rows). Caller checked :func:`gemm_ag_up_fits`.""" + from sglang.kernels.ops.kimi_k3 import gemm_ag as mod + + state = _get_state() + assert state is not None + return mod.gemm_ag_up_proj( + state.world_size, + x, + weight, + b, + c, + torch.empty_like(b), + ws_mc_base=state.comm.mc_base_ptr, + ) + + +def finalize_all_reduce_push_norm( + out: torch.Tensor, + gemm2_out: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + expert_weights: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, +) -> torch.Tensor: + """Deferred MoE finalize fused into the 1shot push AR + RMSNorm; ``out`` is + output-only. Caller checked :func:`finalize_push_fits`.""" + from sglang.kernels.ops.kimi_k3 import all_reduce as mod + + state = _get_state() + assert state is not None + return mod.finalize_all_reduce_push_norm( + state.world_size, + out, + gemm2_out, + expanded_idx_to_permuted_idx, + expert_weights, + weight, + eps, + ws_mc_base=state.comm.mc_base_ptr, + ) diff --git a/python/sglang/srt/layers/k3_gemm_ar.py b/python/sglang/srt/layers/k3_gemm_ar.py new file mode 100644 index 000000000..ff902f66d --- /dev/null +++ b/python/sglang/srt/layers/k3_gemm_ar.py @@ -0,0 +1,95 @@ +"""K3 fused o_proj GEMM+AR dispatch (``SGLANG_K3_GEMM_AR``). + +Glue between the model and ``kernels.ops.kimi_k3.gemm_ar``: lazily allocates +the P2P comm region on the TP group and swaps the o_proj RowParallelLinear +forward for the single fused GEMM+all-reduce kernel at decode shapes +(falling through to the regular GEMM + AR path whenever the input doesn't +fit — prefill, capture, non-2D input). Eager-mode only; see +kernels/ops/kimi_k3/GEMM_AR_README.md. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.environ import envs + +if TYPE_CHECKING: + from sglang.srt.layers.linear import RowParallelLinear + +logger = logging.getLogger(__name__) + +_INITIALIZED = False +_ENABLED = False + + +def _init() -> bool: + global _INITIALIZED, _ENABLED + if _INITIALIZED: + return _ENABLED + _INITIALIZED = True + if not envs.SGLANG_K3_GEMM_AR.get(): + return False + from sglang.srt.runtime_context import get_parallel + + world_size = get_parallel().tp_size + if not (2 <= world_size <= 8): + logger.warning("SGLANG_K3_GEMM_AR requires 2 <= TP <= 8; disabled.") + return False + if torch.cuda.get_device_capability() < (10, 0): + logger.warning("SGLANG_K3_GEMM_AR requires SM100+; disabled.") + return False + _ENABLED = True + logger.info("K3 fused o_proj GEMM+AR enabled (world_size=%d)", world_size) + return True + + +def maybe_wrap_o_proj(o_proj: RowParallelLinear) -> None: + """SGLANG_K3_GEMM_AR: route decode-shaped o_proj calls through the fused + GEMM+all-reduce kernel; everything else falls through to the original + forward. The comm region + JIT compile happen HERE (model build, weight + shape known, well before CUDA-graph capture) — capture must only see the + ready-to-launch path.""" + if not _init(): + return + from sglang.kernels.ops.kimi_k3 import gemm_ar as mod + from sglang.srt.distributed.parallel_state import get_tp_group + from sglang.srt.runtime_context import get_parallel + + parallel = get_parallel() + world_size = parallel.tp_size + if not o_proj.reduce_results or o_proj.tp_size != world_size: + return + weight = o_proj.weight + if not ( + isinstance(weight, torch.Tensor) + and weight.shape[0] == mod.N + and weight.shape[1] % 128 == 0 + ): + return + + mod.init( + world_size=world_size, + rank=parallel.tp_rank, + group=get_tp_group().cpu_group, + k=weight.shape[1], + ) + # per-K compile + base-address stash, pre-capture + mod._module_with_bases(weight.shape[1], world_size) + + inner = o_proj.forward + + def _gemm_ar_forward(x, *args, **kwargs): + weight = o_proj.weight + if ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and mod.fits(x) + ): + return mod.o_proj_gemm_ar(x, weight), None + return inner(x, *args, **kwargs) + + o_proj.forward = _gemm_ar_forward diff --git a/python/sglang/srt/layers/k3_sp_collective.py b/python/sglang/srt/layers/k3_sp_collective.py new file mode 100644 index 000000000..e7ed3d386 --- /dev/null +++ b/python/sglang/srt/layers/k3_sp_collective.py @@ -0,0 +1,470 @@ +"""K3 SP-MoE MNNVL reduce-scatter/all-gather dispatch. + +The reduce-scatter folds the pending attention residual into its reduction +epilogue. The matching all-gather reassembles the token shards after MoE. +Both reuse CustomAllReduceV2's push workspace and fall back as a pair outside +the checked-in GB300 tuning envelope. +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers import k3_ar_fusion + +if TYPE_CHECKING: + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + from sglang.srt.distributed.parallel_state import GroupCoordinator + +logger = logging.getLogger(__name__) + +_HIDDEN_SIZE = 7168 +_SUPPORTED_WORLD_SIZES = (4, 8) + +# Named persistent symmetric buffers, one per NVLS-aliased tensor. Every rank +# must resolve the same (buffer, offset) for these, which a per-forward +# allocation does not give -- see k3_ar_fusion.symm_buffer. +_O_PROJ = "sp_o_proj" # TP-partial o_proj output, read by the pull RS +_ALL_GATHER = "sp_all_gather" # full-batch output of the direct AG +_ATTN_RES_AG = "sp_attn_res_ag" # ... and of its attention-residual fusion + + +class _State: + def __init__(self, group: GroupCoordinator, comm: CustomAllReduceV2): + self.group = group + self.comm = comm + + +_STATE: Optional[_State] = None +_INITIALIZED = False +_O_PROJ_RESULT_BUFFERS: dict[int, torch.Tensor] = {} +_O_PROJ_OUTPUT_ROWS: ContextVar[Optional[int]] = ContextVar( + "k3_sp_o_proj_output_rows", default=None +) + + +def _init_state() -> Optional[_State]: + global _STATE, _INITIALIZED + if _INITIALIZED: + return _STATE + _INITIALIZED = True + + explicit = envs.SGLANG_K3_SP_COLLECTIVE.is_set() + if explicit and not envs.SGLANG_K3_SP_COLLECTIVE.get(): + return None + + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + from sglang.srt.runtime_context import get_parallel, get_server_args + from sglang.srt.utils.common import get_device_sm + + server_args = get_server_args() + a2a = server_args.moe_a2a_backend + group = get_parallel().attn_tp_group + comm = group.ca_comm + if ( + get_device_sm() != 103 + or group.world_size not in _SUPPORTED_WORLD_SIZES + or a2a not in ("megamoe", "deepep") + or not isinstance(comm, CustomAllReduceV2) + or comm.disabled + or comm.mc_base_ptr == 0 + ): + message = ( + "K3 SP collective requires SM103, TP4/TP8, MegaMoE/DeepEP, and " + "CustomAllReduceV2 with multicast; using NCCL." + ) + (logger.warning if explicit else logger.info)(message) + return None + + from sglang.kernels.ops.kimi_k3 import attn_res, sp_collective + + # Refuse to enable without a checked-in table for this exact device. + if ( + sp_collective.get_dispatch( + "reduce_scatter", + group.world_size, + _HIDDEN_SIZE, + 1, + torch.device("cuda"), + ) + is None + ): + (logger.warning if explicit else logger.info)( + "K3 SP collective has no tuning table for %s; using NCCL.", + torch.cuda.get_device_name(), + ) + return None + + sp_collective.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr) + attn_res.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr) + _STATE = _State(group, comm) + logger.info( + "K3 SP collective enabled (TP%d, fused RS residual + AG)", + group.world_size, + ) + if envs.SGLANG_K3_SP_ATTN_RES.get(): + logger.info( + "K3 SP attention-residual carry enabled " + "(local agg/bank write + normalized AG)" + ) + return _STATE + + +def enabled() -> bool: + return _init_state() is not None + + +def _symm_buffer( + state: _State, name: str, rows: int, width: int, dtype: torch.dtype +) -> torch.Tensor: + """Named persistent symmetric buffer over the group this module reduces on.""" + return k3_ar_fusion.symm_buffer( + name, rows, width, dtype, group_name=state.group.cpu_group.group_name + ) + + +def requires_symmetric_rs(num_tokens: int, device: torch.device) -> bool: + """Whether standalone or fused RS reads o_proj through its NVLS alias.""" + state = _init_state() + if state is None: + return False + from sglang.kernels.ops.kimi_k3 import sp_collective + + dispatch = sp_collective.get_dispatch( + "reduce_scatter", + state.group.world_size, + _HIDDEN_SIZE, + num_tokens, + device, + ) + if dispatch is not None and dispatch.strategy == "pull": + return True + if not envs.SGLANG_K3_SP_ATTN_RES.get(): + return False + fusion = sp_collective.get_fusion_dispatch( + "reduce_scatter_attn_res", + state.group.world_size, + _HIDDEN_SIZE, + num_tokens, + device, + ) + return fusion is not None and fusion.strategy == "fused_pull" + + +def get_o_proj_output_buffer( + num_tokens: int, dtype: torch.dtype, hidden_size: int = _HIDDEN_SIZE +) -> torch.Tensor: + """Persistent symmetric storage for a TP-partial o_proj output.""" + state = _init_state() + assert state is not None, "K3 SP collective is not initialized" + return _symm_buffer(state, _O_PROJ, num_tokens, hidden_size, dtype) + + +@contextmanager +def o_proj_output_rows(num_rows: int): + """Tell o_proj to target a padded full-batch symmetric output.""" + token = _O_PROJ_OUTPUT_ROWS.set(num_rows) + try: + yield + finally: + _O_PROJ_OUTPUT_ROWS.reset(token) + + +def get_o_proj_output_rows(default: int) -> int: + num_rows = _O_PROJ_OUTPUT_ROWS.get() + return default if num_rows is None else num_rows + + +def register_o_proj_output(result: torch.Tensor, output: torch.Tensor) -> None: + """Associate a model-visible o_proj result/view with its symmetric backing.""" + if ( + result.ndim != output.ndim + or result.shape[1:] != output.shape[1:] + or result.shape[0] > output.shape[0] + or result.dtype != output.dtype + ): + raise RuntimeError( + "K3 o_proj result does not match its persistent symmetric output" + ) + _O_PROJ_RESULT_BUFFERS[result.data_ptr()] = output + + +def finish_padded_o_proj_output( + result: torch.Tensor, num_padded: int +) -> Optional[torch.Tensor]: + """Zero the padding tail and return the full persistent symmetric buffer.""" + output = _O_PROJ_RESULT_BUFFERS.get(result.data_ptr()) + if output is None or output.shape[0] != num_padded: + return None + output[result.shape[0] :].zero_() + _O_PROJ_RESULT_BUFFERS[output.data_ptr()] = output + return output + + +def _resolve_symmetric_o_proj_input(tensor: torch.Tensor) -> tuple[torch.Tensor, int]: + """Recover the persistent o_proj buffer if a model wrapper rebound its view. + + The second element is 0 when neither is symmetric, and the caller falls back. + """ + input_mc_ptr = k3_ar_fusion.find_mc_ptr(tensor) + if input_mc_ptr is not None: + return tensor, input_mc_ptr + output = _O_PROJ_RESULT_BUFFERS.get(tensor.data_ptr()) + if output is None: + return tensor, 0 + return output, k3_ar_fusion.find_mc_ptr(output) or 0 + + +def _eligible( + state: _State, + tensor: torch.Tensor, + residual: Optional[torch.Tensor] = None, +) -> bool: + if ( + tensor.dtype != torch.bfloat16 + or not tensor.is_contiguous() + or tensor.ndim != 2 + or tensor.shape[1] != _HIDDEN_SIZE + or tensor.shape[0] <= 0 + or tensor.shape[0] % state.group.world_size != 0 + ): + return False + if residual is not None: + local_shape = (tensor.shape[0] // state.group.world_size, tensor.shape[1]) + if ( + residual.shape not in (tensor.shape, local_shape) + or residual.dtype != tensor.dtype + or not residual.is_contiguous() + ): + return False + local_bytes = tensor.numel() * tensor.element_size() // state.group.world_size + return local_bytes <= state.comm.max_push_size + + +def reduce_scatter_res( + tensor: torch.Tensor, residual: Optional[torch.Tensor] +) -> Optional[torch.Tensor]: + """Return a fused local shard, or None when the NCCL fallback should run.""" + state = _init_state() + if state is None or not _eligible(state, tensor, residual): + return None + from sglang.kernels.ops.kimi_k3 import sp_collective + + dispatch = sp_collective.get_dispatch( + "reduce_scatter", + state.group.world_size, + tensor.shape[1], + tensor.shape[0], + tensor.device, + ) + if dispatch is None: + return None + output = torch.empty( + (tensor.shape[0] // state.group.world_size, tensor.shape[1]), + dtype=tensor.dtype, + device=tensor.device, + ) + if dispatch.strategy == "push": + return sp_collective.reduce_scatter_res( + state.group.world_size, + tensor, + output, + residual, + tuning=dispatch.tuning, + ) + if dispatch.strategy == "pull": + tensor, input_mc_ptr = _resolve_symmetric_o_proj_input(tensor) + if input_mc_ptr == 0: + logger.warning("K3 pull RS input is not symmetric; using NCCL.") + return None + return sp_collective.reduce_scatter_pull( + state.group.world_size, + tensor, + output, + residual, + input_mc_ptr=input_mc_ptr, + tuning=dispatch.tuning, + ) + raise AssertionError(f"unknown K3 reduce-scatter strategy: {dispatch.strategy}") + + +def reduce_scatter_attn_res( + tensor: torch.Tensor, + residual: Optional[torch.Tensor], + bank: torch.Tensor, + cw: torch.Tensor, + ow: torch.Tensor, + nvb: int, + eps: float, +) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Fused NVLS pull RS + local residual + attention-residual aggregation.""" + state = _init_state() + if ( + state is None + or not envs.SGLANG_K3_SP_ATTN_RES.get() + or not _eligible(state, tensor, residual) + ): + return None + local_tokens = tensor.shape[0] // state.group.world_size + if ( + bank.shape[0] != local_tokens + or bank.ndim != 3 + or bank.shape[2] != _HIDDEN_SIZE + or not bank.is_contiguous() + or cw.shape != (_HIDDEN_SIZE,) + or ow.shape != (_HIDDEN_SIZE,) + ): + return None + from sglang.kernels.ops.kimi_k3 import attn_res, sp_collective + + dispatch = sp_collective.get_fusion_dispatch( + "reduce_scatter_attn_res", + state.group.world_size, + tensor.shape[1], + tensor.shape[0], + tensor.device, + ) + if dispatch is None or dispatch.strategy != "fused_pull": + return None + tensor, input_mc_ptr = _resolve_symmetric_o_proj_input(tensor) + if input_mc_ptr == 0: + logger.warning("K3 fused pull RS input is not symmetric; using separate path.") + return None + normed = torch.empty( + (local_tokens, tensor.shape[1]), dtype=tensor.dtype, device=tensor.device + ) + prefix = torch.empty_like(normed) + attn_res.attn_res_fused_pull_rs( + state.group.world_size, + tensor, + residual, + bank, + cw, + ow, + normed, + prefix, + nvb, + eps, + input_mc_ptr=input_mc_ptr, + max_blocks=dispatch.max_blocks, + ) + return normed, prefix + + +def all_gather(tensor: torch.Tensor) -> Optional[torch.Tensor]: + """Return the reassembled full batch, or None for the NCCL fallback.""" + state = _init_state() + if state is None: + return None + global_tokens = tensor.shape[0] * state.group.world_size + if ( + tensor.dtype != torch.bfloat16 + or not tensor.is_contiguous() + or tensor.ndim != 2 + or tensor.shape[1] != _HIDDEN_SIZE + or tensor.shape[0] <= 0 + or tensor.numel() * tensor.element_size() > state.comm.max_push_size + ): + return None + from sglang.kernels.ops.kimi_k3 import sp_collective + + dispatch = sp_collective.get_dispatch( + "all_gather", + state.group.world_size, + tensor.shape[1], + global_tokens, + tensor.device, + ) + if dispatch is None: + return None + output_shape = (global_tokens, tensor.shape[1]) + if dispatch.strategy == "push": + output = torch.empty(output_shape, dtype=tensor.dtype, device=tensor.device) + return sp_collective.all_gather( + state.group.world_size, + tensor, + output, + ws_mc_base=state.comm.mc_base_ptr, + tuning=dispatch.tuning, + ) + if dispatch.strategy == "direct": + output = _symm_buffer( + state, _ALL_GATHER, global_tokens, tensor.shape[1], tensor.dtype + ) + return sp_collective.all_gather_direct( + state.group.world_size, + tensor, + output, + output_mc_ptr=k3_ar_fusion.get_mc_ptr(output), + tuning=dispatch.tuning, + ) + raise AssertionError(f"unknown K3 all-gather strategy: {dispatch.strategy}") + + +def attn_res_all_gather( + prefix: torch.Tensor, + bank: torch.Tensor, + cw: torch.Tensor, + ow: torch.Tensor, + nvb: int, + eps: float, + *, + write_prefix: bool = False, +) -> Optional[torch.Tensor]: + """Fused local attention-residual aggregation + direct multicast AG.""" + state = _init_state() + if ( + state is None + or not envs.SGLANG_K3_SP_ATTN_RES.get() + or prefix.dtype != torch.bfloat16 + or not prefix.is_contiguous() + or prefix.ndim != 2 + or prefix.shape[1] != _HIDDEN_SIZE + or prefix.shape[0] <= 0 + or bank.ndim != 3 + or bank.shape[0] != prefix.shape[0] + or bank.shape[2] != _HIDDEN_SIZE + or not bank.is_contiguous() + or cw.shape != (_HIDDEN_SIZE,) + or ow.shape != (_HIDDEN_SIZE,) + ): + return None + from sglang.kernels.ops.kimi_k3 import attn_res, sp_collective + + global_tokens = prefix.shape[0] * state.group.world_size + dispatch = sp_collective.get_fusion_dispatch( + "attn_res_all_gather", + state.group.world_size, + prefix.shape[1], + global_tokens, + prefix.device, + ) + if dispatch is None or dispatch.strategy != "fused_direct": + return None + output = _symm_buffer( + state, _ATTN_RES_AG, global_tokens, prefix.shape[1], prefix.dtype + ) + attn_res.attn_res_fused_direct_ag( + state.group.world_size, + prefix, + bank, + cw, + ow, + output, + nvb, + eps, + output_mc_ptr=k3_ar_fusion.get_mc_ptr(output), + max_blocks=dispatch.max_blocks, + write_prefix=write_prefix, + ) + return output diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 2886cb01e..67662c9a7 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -1560,7 +1560,13 @@ class RowParallelLinear(LinearBase): # Fallback for parameters that don't accept additional args param.load_row_parallel_weight(loaded_weight) - def forward(self, input_, skip_all_reduce=False, forward_batch=None): + def forward( + self, + input_, + skip_all_reduce=False, + forward_batch=None, + output_tensor=None, + ): if self.input_is_parallel: input_parallel = input_ else: @@ -1581,7 +1587,20 @@ class RowParallelLinear(LinearBase): get_tp_group(), disabled=not is_allocation_symmetric() ) with symm_ctx: - output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_) + if output_tensor is None: + output_parallel = self.quant_method.apply( + self, input_parallel, bias=bias_ + ) + else: + apply_into = getattr(self.quant_method, "apply_into", None) + if apply_into is None: + raise RuntimeError( + f"{type(self.quant_method).__name__} cannot write into " + "caller-owned linear output" + ) + output_parallel = apply_into( + self, input_parallel, output_tensor, bias=bias_ + ) # skip_all_reduce: explicit call-site override. Also honor # ForwardFlags (fuse_mlp_allreduce / mlp_reduce_scatter) published by diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index f6a445b42..a30b6bd73 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -125,6 +125,15 @@ class DeepEPMoE(FusedMoE): in ("modelopt_fp4", "modelopt_mixed", "nvfp4_online") ): self.deprecate_flag = True + elif ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and get_moe_runner_backend().is_deep_gemm() + and quant_config is not None + and quant_config.get_name() == "mxfp4" + ): + # MXFP4 experts (e.g. Kimi K3) on the DeepGEMM fp8_fp4 W4A8 path: + # route through the modern FusedMoE runner (Mxfp4MoEMethod.apply). + self.deprecate_flag = True elif ( quant_config is None and self.w13_weight.dtype == torch.bfloat16 @@ -322,7 +331,7 @@ class DeepEPMoE(FusedMoE): self, dispatch_output: DeepEPNormalDispatchOutput, ): - assert self.moe_runner_config.activation == "silu" + assert self.moe_runner_config.activation in ("silu", "situ") assert isinstance(self.quant_method, W4AFp8MoEMethod) return self.quant_method.apply_deepep_normal( layer=self, @@ -333,7 +342,7 @@ class DeepEPMoE(FusedMoE): self, dispatch_output: DeepEPLLDispatchOutput, ): - assert self.moe_runner_config.activation == "silu" + assert self.moe_runner_config.activation in ("silu", "situ") assert isinstance(self.quant_method, W4AFp8MoEMethod) return self.quant_method.apply_deepep_ll( layer=self, diff --git a/python/sglang/srt/layers/moe/fused_moe_native.py b/python/sglang/srt/layers/moe/fused_moe_native.py index ba9a51e82..d72f0e9a3 100644 --- a/python/sglang/srt/layers/moe/fused_moe_native.py +++ b/python/sglang/srt/layers/moe/fused_moe_native.py @@ -6,7 +6,7 @@ It is based on https://github.com/pytorch-labs/gpt-fast/blob/32971d3129541c5bfb4 import torch from torch.nn import functional as F -from sglang.srt.layers.activation import GeluAndMul, SiluAndMul +from sglang.srt.layers.activation import GeluAndMul, SiluAndMul, SituAndMul from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( swiglu_gpt_oss_sigmoid_alpha, @@ -40,6 +40,14 @@ def fused_moe_forward_native( x1 = F.silu(x1) elif moe_runner_config.activation == "gelu": x1 = F.gelu(x1) + elif moe_runner_config.activation == "situ": + beta = ( + moe_runner_config.gemm1_alpha + if moe_runner_config.gemm1_alpha is not None + else 4.0 + ) + x1 = beta * torch.tanh(x1.float() / beta) * torch.sigmoid(x1.float()) + x1 = x1.to(x.dtype) else: raise ValueError(f"Unsupported activation: {moe_runner_config.activation=}") x3 = torch.einsum("ti, taoi -> tao", x, w3_weights) @@ -77,6 +85,14 @@ def moe_forward_native( act = SiluAndMul() elif moe_runner_config.activation == "gelu": act = GeluAndMul() + elif moe_runner_config.activation == "situ": + situ_beta = ( + moe_runner_config.gemm1_alpha + if moe_runner_config.gemm1_alpha is not None + else 4.0 + ) + situ_linear_beta = moe_runner_config.gemm1_clamp_limit + act = SituAndMul(beta=situ_beta, linear_beta=situ_linear_beta) else: raise ValueError(f"Unsupported activation: {moe_runner_config.activation=}") diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index 9ab1dd29c..fd37548ac 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py @@ -2,7 +2,10 @@ from typing import Optional import torch import torch.nn.functional as F +import triton +import triton.language as tl +from sglang.srt.layers import zero_copy_context from sglang.srt.utils import is_cuda from sglang.srt.utils.custom_op import register_custom_op @@ -15,6 +18,69 @@ if _is_cuda: from sglang.kernels.ops.moe.moe_wna16_marlin import moe_wna16_marlin_gemm +@triton.jit +def _tl_tanh(x): + return 2.0 * tl.sigmoid(2.0 * x) - 1.0 + + +@triton.jit +def _situ_and_mul_kernel( + x_ptr, # [M, 2N] gate;up halves (non-interleaved) + out_ptr, # [M, N] + N, + situ_beta, + linear_beta, + stride_xm, + stride_om, + BLOCK_N: tl.constexpr, + HAS_LINEAR_BETA: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask = offs < N + base = x_ptr + pid_m * stride_xm + gate = tl.load(base + offs, mask=mask, other=0.0).to(tl.float32) + up = tl.load(base + N + offs, mask=mask, other=0.0).to(tl.float32) + gate = situ_beta * _tl_tanh(gate / situ_beta) * tl.sigmoid(gate) + if HAS_LINEAR_BETA: + up = linear_beta * _tl_tanh(up / linear_beta) + out = gate * up + tl.store( + out_ptr + pid_m * stride_om + offs, + out.to(out_ptr.dtype.element_ty), + mask=mask, + ) + + +def situ_and_mul( + output: torch.Tensor, + x: torch.Tensor, + situ_beta: float, + linear_beta: Optional[float], +) -> None: + """SiTU gated activation (Kimi K3), fused into one elementwise kernel: + out = situ_beta*tanh(gate/situ_beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) + where x = [gate; up] halves along the last dim. + """ + M, N2 = x.shape + N = N2 // 2 + assert output.shape == (M, N) + BLOCK_N = 1024 + grid = (M, triton.cdiv(N, BLOCK_N)) + _situ_and_mul_kernel[grid]( + x, + output, + N, + float(situ_beta), + float(linear_beta) if linear_beta is not None else 0.0, + x.stride(0), + output.stride(0), + BLOCK_N=BLOCK_N, + HAS_LINEAR_BETA=linear_beta is not None, + ) + + def get_scalar_type( num_bits: int, has_zp: bool, @@ -177,9 +243,25 @@ def fused_marlin_moe( if global_num_experts == -1: global_num_experts = E - sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( - topk_ids, block_size_m, global_num_experts - ) + if ( + M == 1 + and topk <= 32 + and expert_map is None + # The JIT kernel is int32-only; torch-native topk emits int64 -- let + # that (test-only) shape take the generic path instead of casting. + and topk_ids.dtype == torch.int32 + ): + # Single-token decode: top-k ids are distinct, so alignment is a + # single-warp sort instead of the align + count_and_sort kernel pair. + from sglang.kernels.ops.moe.moe_align_single_token import moe_align_single_token + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_single_token( + topk_ids, block_size_m + ) + else: + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, block_size_m, global_num_experts + ) if workspace is None: max_workspace_size = (max(2 * N, K) // 64) * ( @@ -266,6 +348,13 @@ def fused_marlin_moe( ) elif activation == "silu" and is_gated: silu_and_mul(intermediate_cache1.view(-1, gemm1_n), intermediate_cache2) + elif activation == "situ" and is_gated: + situ_and_mul( + intermediate_cache2, + intermediate_cache1.view(-1, gemm1_n), + situ_beta=gemm1_alpha if gemm1_alpha is not None else 4.0, + linear_beta=clamp_limit, + ) elif activation == "silu" and not is_gated: intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N)) elif activation == "relu2" and not is_gated: @@ -305,10 +394,28 @@ def fused_marlin_moe( is_zp_float=False, ).view(-1, topk, K) - output = hidden_states if inplace else torch.empty_like(hidden_states) + output = zero_copy_context.get_moe_output(hidden_states) + if output is None: + output = hidden_states if inplace else torch.empty_like(hidden_states) if is_mxfp4_marlin: - return torch.sum(intermediate_cache3, dim=1, out=output) + # Top-k weights (incl. routed scaling) are already applied above via + # mul_topk_weights, so this is a plain sum over the topk dim. The JIT + # vectorized pass (~1.5us at decode shapes) beats sgl_kernel's + # moe_sum_reduce_kernel_general (~5.7us) and the generic at::native + # reduce_kernel torch.sum dispatches to (~6.7us). + if ( + intermediate_cache3.dtype == torch.bfloat16 + and intermediate_cache3.is_contiguous() + and output.is_contiguous() + and intermediate_cache3.shape[-1] % 8 == 0 + ): + from sglang.kernels.ops.moe.moe_topk_sum import moe_topk_sum + + moe_topk_sum(intermediate_cache3, output) + else: + moe_sum_reduce(intermediate_cache3, output, 1.0) + return output else: if routed_scaling_factor is None: routed_scaling_factor = 1.0 diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index ee82f29bc..b3f1291f4 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -658,13 +658,32 @@ class FusedMoE(torch.nn.Module): if not is_bias and self.use_triton_kernels: # do not transpose for bias loaded_weight = loaded_weight.transpose(-2, -1) + # When the buffer is padded (e.g., MXFP4 SM100 rounds + # intermediate up to 256), shard_size from the buffer may + # exceed the checkpoint's per-TP slice. Derive the actual + # shard size from the loaded weight so we index correctly. + loaded_shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size loaded_weight = loaded_weight.narrow( - shard_dim, shard_size * tp_rank, shard_size + shard_dim, loaded_shard_size * tp_rank, loaded_shard_size ) expert_data = expert_data.narrow(shard_dim, start, shard_size) + loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight) - expert_data.copy_(loaded_weight) + # loaded_weight may be smaller than expert_data along shard_dim when + # the buffer is padded. Copy into the leading slice and leave the + # trailing padding as zeros. Rank-mismatched tensors (bias / scalar + # scales) don't carry shard_dim; they keep the plain broadcast copy. + if ( + loaded_weight.dim() == expert_data.dim() + and shard_dim < expert_data.dim() + and loaded_weight.shape[shard_dim] < expert_data.shape[shard_dim] + ): + expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim]).copy_( + loaded_weight + ) + else: + expert_data.copy_(loaded_weight) def _load_w2( self, @@ -729,13 +748,28 @@ class FusedMoE(torch.nn.Module): if not is_bias and not self.use_presharded_weights: if self.use_triton_kernels: loaded_weight = loaded_weight.transpose(-2, -1) + # Derive shard size from the loaded weight so padded buffers + # do not cause out-of-bounds indexing into the checkpoint. + loaded_shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size loaded_weight = loaded_weight.narrow( - shard_dim, shard_size * tp_rank, shard_size + shard_dim, loaded_shard_size * tp_rank, loaded_shard_size ) # w2, down_proj: Load into only logical weight of w2. loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight) - expert_data.copy_(loaded_weight) + # loaded_weight may be smaller than expert_data along shard_dim when + # the buffer is padded. Copy into the leading slice only. See the + # rank-mismatch note in _load_w13. + if ( + loaded_weight.dim() == expert_data.dim() + and shard_dim < expert_data.dim() + and loaded_weight.shape[shard_dim] < expert_data.shape[shard_dim] + ): + expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim]).copy_( + loaded_weight + ) + else: + expert_data.copy_(loaded_weight) def _maybe_load_fp8_shared_expert_as_fp4( self, @@ -873,7 +907,7 @@ class FusedMoE(torch.nn.Module): # if expert_id is None, then # all the experts are loaded at the same time if ( - not expert_id + expert_id is None and self.quant_config is not None and self.quant_config.get_name() == "mxfp4" and self.quant_config.is_static_cfg() diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index c67d0aa9e..8fa06487d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -91,7 +91,11 @@ class AiterRunnerOutput(RunnerOutput): return MoeRunnerBackend.AITER -_AITER_ACTIVATIONS = {"silu": "Silu", "swiglu": "Swiglu"} +_AITER_ACTIVATIONS = { + "silu": "Silu", + "swiglu": "Swiglu", + "situ": "Situv2", +} def _aiter_activation(activation: str): @@ -162,7 +166,15 @@ class AiterRunnerCore(MoeRunnerCore): extra["num_local_tokens"] = runner_input.num_local_tokens if runner_input.output_dtype is not None: extra["dtype"] = runner_input.output_dtype - if quant_info.swiglu_limit > 0: + if self.config.activation == "situ": + from aiter.ops.flydsl.moe_common import GateMode + + extra["gate_mode"] = GateMode.SEPARATED.value + if self.config.gemm1_alpha is not None: + extra["beta"] = float(self.config.gemm1_alpha) + if self.config.gemm1_clamp_limit is not None: + extra["linear_beta"] = float(self.config.gemm1_clamp_limit) + elif quant_info.swiglu_limit > 0: # GateMode is only needed for the gpt-oss MXFP4 swiglu_limit path. # Import lazily so models that don't use it (e.g. DeepSeek-V3 fp8, # swiglu_limit==0) still run on aiter builds where this module diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 52658f656..b769cce03 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple import einops import torch +import triton +import triton.language as tl from sglang.kernels.ops.attention.dsv4 import silu_and_mul_masked_post_quant from sglang.kernels.ops.quantization import per_token_group_quant @@ -165,7 +167,9 @@ class DeepGemmMoeQuantInfo(MoeQuantInfo): class DeepGemmRunnerCore(MoeRunnerCore): def __init__(self, config: MoeRunnerConfig): super().__init__(config) - assert self.config.activation == "silu" + # SiTU (Kimi K3) is applied outside the GEMMs in python, so it only + # needs the masked-gemm activation site to branch (see _run_masked_gemm). + assert self.config.activation in ("silu", "situ") assert self.config.is_gated self.swiglu_limit = self.config.swiglu_limit self.use_swizzle = False @@ -256,7 +260,61 @@ class DeepGemmRunnerCore(MoeRunnerCore): dispose_tensor(hidden_states) dispose_tensor(hidden_states_scale) - if envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get(): + if self.config.activation == "situ": + situ_beta = self.config.gemm1_alpha + situ_linear_beta = self.config.gemm1_clamp_limit + assert situ_beta is not None and situ_linear_beta is not None + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + # Fused SiTU + per-group fp8 quant over the compacted rows, + # then the proven round-up e8m0 cast (mn-major packed layout). + rows = gateup_output.shape[0] + half_n = N // 2 + kg = half_n // scale_block_size + down_input_fp8 = torch.empty( + (rows, half_n), + device=gateup_output.device, + dtype=torch.float8_e4m3fn, + ) + s = torch.empty( + (rows, kg), device=gateup_output.device, dtype=torch.float32 + ) + _situ_mul_quant_contig_kernel[(rows,)]( + gateup_output, + down_input_fp8, + s, + half_n, + kg, + situ_beta, + situ_linear_beta, + GROUP=scale_block_size, + KG_POW2=triton.next_power_of_2(kg), + num_warps=8, + ) + del gateup_output + down_input_scale = _cast_to_e8m0_with_rounding_up( + s.unsqueeze(0) + ).squeeze(0) + else: + from sglang.kernels.ops.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + + gate = gateup_output[:, : N // 2].float() + up = gateup_output[:, N // 2 :].float() + gate = situ_beta * torch.tanh(gate / situ_beta) * torch.sigmoid(gate) + up = situ_linear_beta * torch.tanh(up / situ_linear_beta) + down_input = (gate * up).to(torch.bfloat16) + del gateup_output + + down_input_fp8, down_input_scale = sglang_per_token_group_quant_fp8( + down_input, + scale_block_size, + column_major_scales=False, + scale_tma_aligned=False, + scale_ue8m0=False, + ) + del down_input + elif envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get(): swiglu_limit_arg: Optional[float] = self.swiglu_limit down_input_fp8 = torch.empty( @@ -519,28 +577,38 @@ class DeepGemmRunnerCore(MoeRunnerCore): ) # Act. - topk_ids_rs = running_state.get("topk_ids") - num_real_tokens = ( - topk_ids_rs.shape[0] - if ( - use_mxfp8 - and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 - and topk_ids_rs is not None - and "src2dst" in running_state + if self.config.activation == "situ": + down_input, down_input_scale = _varlen_deep_gemm_situ_mul_quant( + gateup_output, + masked_m, + group_size=128, + topk=self.config.top_k, + beta=self.config.gemm1_alpha, + linear_beta=self.config.gemm1_clamp_limit, + ) + else: + topk_ids_rs = running_state.get("topk_ids") + num_real_tokens = ( + topk_ids_rs.shape[0] + if ( + use_mxfp8 + and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + and topk_ids_rs is not None + and "src2dst" in running_state + ) + else None + ) + down_input, down_input_scale = _varlen_deep_gemm_silu_mul_quant( + gateup_output, + masked_m, + group_size=scale_block_size, + topk=self.config.top_k, + swiglu_limit=swiglu_limit_arg, + swizzle=self.use_swizzle, + gemm1_alpha=self.config.gemm1_alpha, + gemm1_clamp_limit=self.config.gemm1_clamp_limit, + num_real_tokens=num_real_tokens, ) - else None - ) - down_input, down_input_scale = _varlen_deep_gemm_silu_mul_quant( - gateup_output, - masked_m, - group_size=scale_block_size, - topk=self.config.top_k, - swiglu_limit=swiglu_limit_arg, - swizzle=self.use_swizzle, - gemm1_alpha=self.config.gemm1_alpha, - gemm1_clamp_limit=self.config.gemm1_clamp_limit, - num_real_tokens=num_real_tokens, - ) del gateup_output # Down activation is quantised locally at scale_block_size (never DeepEP-LL), @@ -885,10 +953,11 @@ def post_permute_deep_gemm_to_standard( hidden_states_shape = running_state["hidden_states_shape"] hidden_states_dtype = running_state["hidden_states_dtype"] hidden_states_device = running_state["hidden_states_device"] - src2dst = running_state["src2dst"] topk_ids = running_state["topk_ids"] topk_weights = running_state["topk_weights"] + src2dst = running_state["src2dst"] + with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): output = torch.empty( hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device @@ -975,7 +1044,7 @@ def pre_permute_deepep_normal_to_deep_gemm( topk_weights, num_recv_tokens_per_expert, ) = dispatch_output - assert runner_config.activation == "silu" + assert runner_config.activation in ("silu", "situ") all_tokens = sum(num_recv_tokens_per_expert) running_state["all_tokens"] = all_tokens @@ -1090,6 +1159,53 @@ def post_permute_deep_gemm_to_deepep_normal( ) +def _varlen_deep_gemm_situ_mul_quant( + gateup_output: torch.Tensor, + masked_m: torch.Tensor, + group_size: int, + topk: int, + beta: float, + linear_beta: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused SiTU activation + per-group fp8 quant via CUDA JIT kernel.""" + from sglang.kernels.ops.kimi_k3 import situ_and_mul_masked_post_quant + + E, N, D_2 = gateup_output.shape + D = D_2 // 2 + G = D // group_size + packed_ue8m0 = deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + + down_input = torch.empty( + (E, N, D), device=gateup_output.device, dtype=torch.float8_e4m3fn + ) + if packed_ue8m0: + down_input_scale = torch.empty( + (E, G // 4, N), device=gateup_output.device, dtype=torch.int32 + ) + else: + down_input_scale = torch.empty( + (E, N, G), device=gateup_output.device, dtype=torch.float32 + ) + + situ_and_mul_masked_post_quant( + gateup_output, + down_input, + down_input_scale, + group_size, + masked_m, + beta=beta, + linear_beta=linear_beta, + scale_ue8m0=packed_ue8m0, + topk=topk, + transposed=packed_ue8m0, + ) + + if packed_ue8m0: + down_input_scale = down_input_scale.transpose(-1, -2) + + return down_input, down_input_scale + + def _varlen_deep_gemm_silu_mul_quant( gateup_output: torch.Tensor, masked_m: Optional[torch.Tensor], @@ -1202,6 +1318,37 @@ def _varlen_deep_gemm_silu_mul_quant( ) +@triton.jit +def _situ_mul_quant_contig_kernel( + g_ptr, # [rows, 2N] bf16, non-interleaved [gate; up] halves + q_ptr, # [rows, N] fp8 out + s_ptr, # [rows, KG] fp32 scales out + N, + KG, + situ_beta, + situ_linear_beta, + GROUP: tl.constexpr, + KG_POW2: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + rows2d = tl.arange(0, KG_POW2)[:, None] + cols = tl.arange(0, GROUP)[None, :] + offs = rows2d * GROUP + cols + mask = rows2d < KG + gate = tl.load(g_ptr + row * 2 * N + offs, mask=mask, other=0.0).to(tl.float32) + up = tl.load(g_ptr + row * 2 * N + N + offs, mask=mask, other=0.0).to(tl.float32) + # tanh(x) == 2*sigmoid(2x) - 1 (avoids a libdevice dependency) + gate_t = 2.0 * tl.sigmoid(2.0 * gate / situ_beta) - 1.0 + gate = situ_beta * gate_t * tl.sigmoid(gate) + up_t = 2.0 * tl.sigmoid(2.0 * up / situ_linear_beta) - 1.0 + y = gate * situ_linear_beta * up_t + amax = tl.clamp(tl.max(tl.abs(y), axis=1), min=1e-10, max=float("inf")) + q = (y * (448.0 / amax)[:, None]).to(tl.float8e4nv) + tl.store(q_ptr + row * N + offs, q, mask=mask) + srow = tl.arange(0, KG_POW2) + tl.store(s_ptr + row * KG + srow, amax / 448.0, mask=srow < KG) + + def _apply_swiglu_limit( gateup_output: torch.Tensor, swiglu_limit: float ) -> torch.Tensor: diff --git a/python/sglang/srt/layers/moe/moe_runner/marlin.py b/python/sglang/srt/layers/moe/moe_runner/marlin.py index c595d6757..0c5ea7fae 100644 --- a/python/sglang/srt/layers/moe/moe_runner/marlin.py +++ b/python/sglang/srt/layers/moe/moe_runner/marlin.py @@ -140,7 +140,10 @@ def fused_experts_none_to_marlin( ) if runner_config.is_gated: - assert runner_config.activation == "silu", "Only gated SiLU is supported." + assert runner_config.activation in { + "silu", + "situ", + }, f"Only gated SiLU/SiTU is supported, got {runner_config.activation}." elif runner_config.activation not in {"silu", "relu2"}: raise ValueError( f"Unsupported Marlin MoE activation: {runner_config.activation}" diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index 964682ab8..64bfb8922 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -685,6 +685,17 @@ def _fused_moe_kernel_sequence( x = intermediate_cache1.view(-1, N) d = x.shape[-1] // 2 intermediate_cache2.copy_(F.silu(x[..., :d]) * x[..., d:]) + elif activation == "situ" and is_gated: + d = N // 2 + x = intermediate_cache1.view(-1, N) + gate = x[..., :d].float() + up = x[..., d:].float() + situ_beta = gemm1_alpha if gemm1_alpha is not None else 4.0 + gate = situ_beta * torch.tanh(gate / situ_beta) * torch.sigmoid(gate) + situ_linear_beta = gemm1_limit + if situ_linear_beta is not None: + up = situ_linear_beta * torch.tanh(up / situ_linear_beta) + intermediate_cache2.copy_((gate * up).to(intermediate_cache1.dtype)) elif activation == "gelu" and is_gated: assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu" assert gemm1_limit is None, "gemm1_limit is not supported for gelu" diff --git a/python/sglang/srt/layers/moe/route_quant_handoff.py b/python/sglang/srt/layers/moe/route_quant_handoff.py new file mode 100644 index 000000000..a8e490ae6 --- /dev/null +++ b/python/sglang/srt/layers/moe/route_quant_handoff.py @@ -0,0 +1,130 @@ +"""Attempt-and-verify handoff for the fused K3 MoE-front prep launch. + +At decode batch sizes the chain between the K3 fused-front GEMM and the +trtllm-gen SiTU MoE op is three tiny back-to-back kernels on the critical +path — route_radix (top-16 of 896), the triton ``(id << 16) | bf16(weight)`` +pack, and per_token_group_quant (mxfp8, ``[T, 3584]``) — about 7.5 us busy +plus two extra launches per MoE layer, each leaving the SMs near idle. The +fused kernel (kernels/ops/moe/moe_route_quant_fused.py) runs all three in one +launch: routing CTAs and one quant CTA per token, concurrently. + +The inputs live in different modules (router logits reach the router through +TopK, activations through the MoE runner), so the fusion is wired as a +consume-once stash instead of new signatures: + + KimiK3MoE._forward_routed* stage(x) before self.topk, clear() after + the experts call + biased_grouped_topk_gpu try_route_quant_fused() replaces the + moe_fused_gate call on a hit + Mxfp4MoEMethod.apply (situ) take(x) skips the quant and the pack + +Every step is fallback-safe: if the routing dispatch never consumes the staged +request (different model, uncovered shape, triton fallback) or the runner's +``take`` misses (activations re-viewed or copied), the unfused chain runs as +before. The stage/clear bracket in the model layer guarantees a published +entry can never leak into another layer whose allocator reused the same +activation address. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import msgspec +import torch + + +class _Handoff(msgspec.Struct): + # staged by the model layer: the activation rows the runner will quantize + request_x: Optional[torch.Tensor] = None + # published by the routing dispatch, keyed by the staged activations + produced_x: Optional[torch.Tensor] = None + packed: Optional[torch.Tensor] = None + x_q: Optional[torch.Tensor] = None + x_s: Optional[torch.Tensor] = None + + +_handoff = _Handoff() + + +def stage(x: torch.Tensor) -> None: + """Publish the routed activations for the upcoming topk call. Caller pairs + this with clear() after the experts call (try/finally).""" + _handoff.request_x = x + _handoff.produced_x = None + + +def clear() -> None: + _handoff.request_x = None + _handoff.produced_x = None + _handoff.packed = None + _handoff.x_q = None + _handoff.x_s = None + + +def try_route_quant_fused( + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + *, + num_fused_shared_experts: int, + renormalize: bool, + routed_scaling_factor: Optional[float], + apply_routed_scaling_factor_on_output: bool, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fused replacement for the ungrouped-sigmoid moe_fused_gate call when a + staged request covers it. Returns (weights, ids) on a hit, None otherwise + (caller falls through to the unfused router).""" + x = _handoff.request_x + if x is None or num_fused_shared_experts != 0: + return None + + from sglang.kernels.ops.moe import moe_route_quant_fused + + if ( + not moe_route_quant_fused.covered(gating_output, correction_bias, topk, x) + or not moe_route_quant_fused.available() + ): + return None + + weights, ids, packed, x_q, x_s = moe_route_quant_fused.route_quant_fused( + gating_output, + correction_bias, + x, + topk, + renormalize=renormalize, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + apply_scale=apply_routed_scaling_factor_on_output, + ) + _handoff.request_x = None + _handoff.produced_x = x + _handoff.packed = packed + _handoff.x_q = x_q + _handoff.x_s = x_s + return weights, ids + + +def take( + x: torch.Tensor, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Consume the published (packed_topk, x_q, x_s int32) for these exact + activation rows, or None. Storage identity is verified so a re-viewed or + copied tensor simply misses.""" + produced = _handoff.produced_x + if produced is None: + return None + if ( + produced.data_ptr() != x.data_ptr() + or produced.shape != x.shape + or produced.dtype != x.dtype + or produced.stride() != x.stride() + ): + return None + out = (_handoff.packed, _handoff.x_q, _handoff.x_s) + _handoff.produced_x = None + _handoff.packed = None + _handoff.x_q = None + _handoff.x_s = None + return out diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 768e03137..02b426466 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1574,7 +1574,22 @@ def biased_grouped_topk_gpu( return topk_weights, topk_ids else: num_experts = gating_output.shape[1] - if _is_cuda and num_experts == 384 and num_expert_group == 1: + # The JIT triton router (single fused kernel: scoring + bias + topk + + # renorm) handles the ungrouped case with arbitrary num_experts/topk. + # Original user: Kimi K2 (384 experts). Also dispatch shapes the other + # fused kernels cannot cover, e.g. Kimi K3 (896 experts, top-16): + # fused_topk_deepseek needs pow2 experts + topk<=8, jit_grouped_topk + # needs experts<=512 + topk<=8. + _jit_gate_ok = ( + _is_cuda + and num_expert_group == 1 + and (topk_group is None or topk_group == 1) + and ( + num_experts == 384 + or (num_experts <= 1024 and (num_experts > 512 or topk > 8)) + ) + ) + if _jit_gate_ok: # ===== TO BE REFACTORED ==== _use_jit_bf16_gate = False if _SGLANG_EXPERIMENTAL_LORA_OPTI: @@ -1601,9 +1616,36 @@ def biased_grouped_topk_gpu( # ===== END TO BE REFACTORED ==== from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate as jit_gate + # Pass BF16 logits through untouched for the automatic radix-select + # fast path. BF16 -> FP32 is exact, and unsupported shapes fall back + # inside moe_fused_gate. + _gating = ( + gating_output + if gating_output.dtype == torch.bfloat16 + else gating_output.to(dtype=torch.float32) + ) + # K3 staged fusion: when the model layer staged the routed + # activations (route_quant_handoff), the radix route, the trtllm id + # pack and the mxfp8 quant run as one launch. Bit-identical + # (weights, ids); a miss falls through to the unfused router. + from sglang.srt.layers.moe import route_quant_handoff + + fused = route_quant_handoff.try_route_quant_fused( + _gating, + correction_bias.to(dtype=torch.float32), + topk, + num_fused_shared_experts=num_fused_shared_experts, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + apply_routed_scaling_factor_on_output=bool( + apply_routed_scaling_factor_on_output + ), + ) + if fused is not None: + return fused return jit_gate( - gating_output.to(dtype=torch.float32), - correction_bias, + _gating, + correction_bias.to(dtype=torch.float32), topk=topk, scoring_func="sigmoid", num_fused_shared_experts=num_fused_shared_experts, @@ -2255,6 +2297,49 @@ def select_experts( return StandardTopKOutput(topk_weights, topk_ids, router_logits) +def precomputed_topk_postprocess_is_noop( + topk_config: TopKConfig, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, +) -> bool: + """Whether :func:`build_precomputed_topk_output` can stand in for + :func:`select_experts`' post-processing. + + A router that produces (weights, ids) itself -- e.g. K3's fused gate+top-k + kernel -- skips select_experts entirely, so it must not skip the work + select_experts does *after* the top-k: the EPLB logical->physical remap, the + padded-region mask, and the shared-expert append. This returns True only + when all of those reduce to no-ops, leaving just the capture hook and the + distribution recorder (which the builder below still runs). + """ + return ( + _is_cuda + and topk_config.num_fused_shared_experts == 0 + and num_token_non_padded is None + and expert_location_dispatch_info is None + and not envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get() + and not envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get() + ) + + +def build_precomputed_topk_output( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + topk_config: TopKConfig, + layer_id: int, +) -> StandardTopKOutput: + """Wrap a router's own (weights, ids) as a STANDARD top-k output, running the + capture hook and the expert-distribution recorder that select_experts would. + + Only valid when :func:`precomputed_topk_postprocess_is_noop` holds. + """ + capture_routed_experts_if_allowed(topk_config, layer_id, topk_ids) + get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids) + # router_logits is only read by the BYPASSED formats and by the + # shared-expert append (excluded above); STANDARD consumers take ids/weights. + return StandardTopKOutput(topk_weights, topk_ids, None) + + # NOTE: the AOT sgl_kernel::moe_fused_gate and sgl_kernel::kimi_k2_moe_fused_gate # ops (and their torch.compile fake impls) were retired here — both CUDA gate # paths now route through the unified Triton router (kernels/ops/moe/moe_fused_gate.py), diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 984d7cbc5..3490fc6a1 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -178,6 +178,17 @@ class CompressedTensorsConfig(QuantizationConfig): from sglang.srt.layers.moe.fused_moe_triton import FusedMoE if isinstance(layer, FusedMoE): + # Detect MXFP4 before the scheme-based path: MXFP4 uses a + # dedicated FusedMoEMethodBase (Mxfp4MoEMethod) that already + # handles all MoE backends, bypassing the scheme abstraction. + if self._is_mxfp4_moe(layer_name=prefix): + from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod + + logger.info_once( + "Using Mxfp4MoEMethod for MXFP4 compressed-tensors MoE" + ) + return Mxfp4MoEMethod(prefix=prefix) + layer.scheme = self.get_moe_scheme(layer=layer, layer_name=prefix) if layer.scheme is None: # ignored layer use_triton_kernels = get_moe_runner_backend().is_triton_kernels() @@ -559,6 +570,28 @@ class CompressedTensorsConfig(QuantizationConfig): return is_mxint4 and input_quant_none and is_symmetric and is_static + def _is_mxfp4_moe(self, layer_name: str) -> bool: + """Detect MXFP4-quantized MoE from global format or target scheme.""" + if "mxfp4" in (self.quant_format or ""): + return True + self._add_fused_moe_to_target_scheme_map() + for key in ["FusedMoE", "Linear"]: + scheme = self.target_scheme_map.get(key) + if scheme is None: + continue + wq = scheme.get("weights") + if wq is None: + continue + if ( + wq.num_bits == 4 + and wq.type == QuantizationType.FLOAT + and wq.strategy == QuantizationStrategy.GROUP.value + and wq.group_size == 32 + and wq.symmetric + ): + return True + return False + def _is_dynamic_token_w4( self, weight_quant: BaseModel, input_quant: BaseModel ) -> bool: diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index ffa210ff1..a475e3c93 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/python/sglang/srt/layers/quantization/marlin_utils.py @@ -250,7 +250,7 @@ def check_moe_marlin_supports_layer( # apply_router_weight_on_input is not supported for moe marlin supports_router_weight = not layer.moe_runner_config.apply_router_weight_on_input if layer.moe_runner_config.is_gated: - supports_activation = layer.moe_runner_config.activation == "silu" + supports_activation = layer.moe_runner_config.activation in {"silu", "situ"} else: supports_activation = layer.moe_runner_config.activation in { "silu", diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index a5bb043c5..c9c8b6ae5 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -33,13 +33,13 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) from sglang.srt.environ import envs +from sglang.srt.layers import zero_copy_context from sglang.srt.layers.amx_utils import ( CPUQuantMethod, _amx_process_weight_after_loading, ) from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig -from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend from sglang.srt.layers.quantization.base_config import ( @@ -145,6 +145,7 @@ if TYPE_CHECKING: _is_cpu = is_cpu() _is_hip = is_hip() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_aiter_k3_opt = _use_aiter and get_bool_env_var("SGLANG_AITER_K3_OPT") _is_shuffle_moe_mxfp4 = is_gfx95_supported() _is_cpu_amx_available = cpu_has_amx_support() @@ -332,6 +333,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.with_bias = False self.use_flashinfer = get_moe_runner_backend().is_flashinfer_mxfp4() self.use_marlin = get_moe_runner_backend().is_marlin() + # True W4A8: DeepGEMM fp8_fp4 grouped GEMM (SM100 MXF8F6F4 UMMA). + # Weights stay MXFP4 (e2m1 + ue8m0 g32, zero requantization); + # activations are quantized to fp8 per-token-group-128. + self.use_deep_gemm = get_moe_runner_backend().is_deep_gemm() self.flashinfer_mxfp4_moe_precision = ( get_exec().moe.flashinfer_mxfp4_moe_precision ) @@ -392,12 +397,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): intermediate_size_per_partition_after_pad - layer.intermediate_size_per_partition ) + elif self.use_deep_gemm: + # DeepGEMM fp8_fp4 grouped GEMM consumes the checkpoint layout + # directly (packed e2m1 K-major + ue8m0 g32 scales); no padding. + pass elif is_sm100_supported(): if self.use_flashinfer: + # FlashInfer trtllm-gen FP4 kernel actual alignment: + # intermediate: scale shuffle needs M%128==0 → intermediate%64==0 + # hidden: finalize kernel needs K%32==0, block quant needs K%32==0 + # Using 128 alignment (not 256) since 256 was overly conservative + # and causes 60GB/GPU waste for models like K3 (384 intermediate). intermediate_size_per_partition_after_pad = round_up( - intermediate_size_per_partition, 256 + intermediate_size_per_partition, 128 ) - hidden_size = round_up(hidden_size, 256) + hidden_size = round_up(hidden_size, 128) else: intermediate_size_per_partition_after_pad = round_up( intermediate_size_per_partition, triton_kernels_padding_alignment @@ -422,9 +436,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # naive-copy fast path is correct. intermediate_size_per_partition_after_pad = intermediate_size_per_partition elif _use_aiter: - + # Expert intermediate is padded to 128 or 256. On K3 (TP8) the 256 + # default costs +33% mem (3072/8=384 -> 512). K3 FlyDSL MoE works with + # 128 padding, so relaxing to 128 saves that memory and lets the TP8 + # decode CUDA graph range grow from 20 to 120. + _inter_align = 128 if _aiter_k3_opt else 256 intermediate_size_per_partition_after_pad = round_up( - intermediate_size_per_partition, 256 + intermediate_size_per_partition, _inter_align ) hidden_size = round_up(hidden_size, 256) @@ -465,17 +483,20 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w13_weight_scale", w13_weight_scale) set_weight_attrs(w13_weight_scale, extra_weight_attrs) + w13_weight_scale.quant_method = "group" - w13_weight_bias = torch.nn.Parameter( - torch.zeros( - layer.num_local_experts, - 2 * intermediate_size_per_partition_after_pad, - dtype=torch.bfloat16, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_bias", w13_weight_bias) - set_weight_attrs(w13_weight_bias, extra_weight_attrs) + create_bias = with_bias or not _is_hip + if create_bias: + w13_weight_bias = torch.nn.Parameter( + torch.zeros( + layer.num_local_experts, + 2 * intermediate_size_per_partition_after_pad, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_bias", w13_weight_bias) + set_weight_attrs(w13_weight_bias, extra_weight_attrs) # down_proj (row parallel) w2_weight = torch.nn.Parameter( @@ -501,13 +522,15 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_weight_scale", w2_weight_scale) set_weight_attrs(w2_weight_scale, extra_weight_attrs) + w2_weight_scale.quant_method = "group" - w2_weight_bias = torch.nn.Parameter( - torch.zeros(layer.num_local_experts, hidden_size, dtype=torch.bfloat16), - requires_grad=False, - ) - layer.register_parameter("w2_weight_bias", w2_weight_bias) - set_weight_attrs(w2_weight_bias, extra_weight_attrs) + if create_bias: + w2_weight_bias = torch.nn.Parameter( + torch.zeros(layer.num_local_experts, hidden_size, dtype=torch.bfloat16), + requires_grad=False, + ) + layer.register_parameter("w2_weight_bias", w2_weight_bias) + set_weight_attrs(w2_weight_bias, extra_weight_attrs) def process_weights_after_loading(self, layer): if self.use_marlin: @@ -519,19 +542,73 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): prepare_moe_mxfp4_layer_for_marlin, ) - if not is_sm90_supported() and not is_sm120_supported(): - raise RuntimeError("MXFP4 Marlin requires SM90 or SM120.") + if ( + not is_sm90_supported() + and not is_sm100_supported() + and not is_sm120_supported() + ): + raise RuntimeError("MXFP4 Marlin requires SM90+.") if not check_moe_marlin_supports_layer(layer, 32, allow_tile_padding=True): raise RuntimeError( "Current MXFP4 MoE layer is not supported by Marlin." ) - if self.moe_runner_config.gemm1_alpha is not None: + if self.moe_runner_config.gemm1_alpha is not None and getattr( + self.moe_runner_config, "gate_up_interleaved", True + ): deinterleave_moe_mxfp4_w13_for_marlin(layer) prepare_moe_mxfp4_layer_for_marlin(layer) layer._mxfp4_backend = "marlin" return + if self.use_deep_gemm: + from deep_gemm import transform_sf_into_required_layout + + # Packed fp4 (e2m1 x2 per byte) weights: DeepGEMM expects int8. + layer.w13_weight.data = layer.w13_weight.data.view(torch.int8) + layer.w2_weight.data = layer.w2_weight.data.view(torch.int8) + # Checkpoint scales are uint8 e8m0 (biased exponents). DeepGEMM + # SM100 needs them in packed-UE8M0 TMA-aligned MN-major layout. + # Round-trip through fp32 is exact (values are powers of two). + for scale_name, weight in ( + ("w13_weight_scale", layer.w13_weight), + ("w2_weight_scale", layer.w2_weight), + ): + scale = getattr(layer, scale_name) + num_experts, n, _ = scale.data.shape + k = weight.shape[2] * 2 + scale_f32 = scale.data.view(torch.float8_e8m0fnu).to(torch.float32) + scale.data = transform_sf_into_required_layout( + scale_f32, + mn=n, + k=k, + recipe=(1, 32), + num_groups=num_experts, + disable_ue8m0_cast=False, + ) + if get_moe_a2a_backend().is_megamoe(): + # MegaMoE consumes the same transformed sf, plus its own + # interleaved/UTCCP weight layout. K3 routes EVERY batch + # through mega (the megamoe backend has no a2a fallback), so + # the contig-layout originals are dead weight — repoint the + # params at the mega tensors to reclaim ~1.9GB/layer/rank + # (keeping both layouts OOMs: 92 layers double the experts). + from deep_gemm import transform_weights_for_mega_moe + + l1_pair, l2_pair = transform_weights_for_mega_moe( + (layer.w13_weight.data, layer.w13_weight_scale.data), + (layer.w2_weight.data, layer.w2_weight_scale.data), + ) + layer.mega_l1_weights = l1_pair + layer.mega_l2_weights = l2_pair + layer.w13_weight.data = l1_pair[0] + layer.w13_weight_scale.data = l1_pair[1] + layer.w2_weight.data = l2_pair[0] + layer.w2_weight_scale.data = l2_pair[1] + layer._mega_moe_weights_built = True + layer._mxfp4_backend = "deep_gemm" + return + if self._fi_kernel == "cutlass_sm90": self._process_weights_for_sm90_cutlass(layer) return @@ -539,38 +616,42 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self._process_weights_for_sm120_cutlass(layer) return if self.use_flashinfer: - # TODO: these values are hardcoded for now, we need to get them from the model + # Per-expert buffers are local (create_weights uses num_local_experts); + # the global self.num_experts here breaks EP>1. Mirrors the SM90 path. + E = layer.num_local_experts + _alpha = getattr(layer.moe_runner_config, "gemm1_alpha", None) or 1.702 + _limit = getattr(layer.moe_runner_config, "gemm1_clamp_limit", None) or 7.0 layer.gemm1_alpha = Parameter( - torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(), + torch.tensor([_alpha] * E, dtype=torch.float32).cuda(), requires_grad=False, ) layer.gemm1_beta = Parameter( - torch.tensor([1.0] * self.num_experts, dtype=torch.float32).cuda(), + torch.tensor([1.0] * E, dtype=torch.float32).cuda(), requires_grad=False, ) layer.gemm1_clamp_limit = Parameter( - torch.tensor([7.0] * self.num_experts, dtype=torch.float32).cuda(), + torch.tensor([_limit] * E, dtype=torch.float32).cuda(), requires_grad=False, ) sf_block_size = 32 # mxfp4 block size assert ( layer.w13_weight.dim() == 3 - and layer.w13_weight.shape[0] == self.num_experts + and layer.w13_weight.shape[0] == E and layer.w13_weight.shape[1] == self.intermediate_size_per_partition * 2 and layer.w13_weight.shape[2] == self.hidden_size // 2 ) assert ( layer.w13_weight_scale.dim() == 3 - and layer.w13_weight_scale.shape[0] == self.num_experts + and layer.w13_weight_scale.shape[0] == E and layer.w13_weight_scale.shape[1] == self.intermediate_size_per_partition * 2 and layer.w13_weight_scale.shape[2] == self.hidden_size // sf_block_size ) assert ( layer.w2_weight.dim() == 3 - and layer.w2_weight.shape[0] == self.num_experts + and layer.w2_weight.shape[0] == E and layer.w2_weight.shape[1] == self.hidden_size and layer.w2_weight.shape[2] == self.intermediate_size_per_partition // 2 @@ -583,13 +664,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) assert ( layer.w13_weight_bias.dim() == 2 - and layer.w13_weight_bias.shape[0] == self.num_experts + and layer.w13_weight_bias.shape[0] == E and layer.w13_weight_bias.shape[1] == self.intermediate_size_per_partition * 2 ) assert ( layer.w2_weight_bias.dim() == 2 - and layer.w2_weight_bias.shape[0] == self.num_experts + and layer.w2_weight_bias.shape[0] == E and layer.w2_weight_bias.shape[1] == self.hidden_size ) @@ -618,9 +699,26 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): new_shape = list(shape) return x.reshape(*new_shape) - w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2) - w13_weight = swap_every_two_rows(w13_weight, -2) - w13_bias = swap_every_two_rows(w13_bias, -1) + if getattr(layer.moe_runner_config, "gate_up_interleaved", True): + w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2) + w13_weight = swap_every_two_rows(w13_weight, -2) + w13_bias = swap_every_two_rows(w13_bias, -1) + else: + # Non-interleaved layout (e.g. K3 Latent MoE): first half = gate + # (w1), second half = up (w3). The trtllm-gen fused gated-act + # epilogue wants rows interleaved as (up_i, gate_i) PAIRS - + # the layout get_reorder_rows_for_gated_act_gemm_row_indices + # produces from [linear; gate] halves, and what the + # interleaved branch above arrives at via swap_every_two_rows. + # A plain halves swap keeps rows blocked and pairs up_i with + # up_{i+1}, which scrambles the gated activation. + half = w13_weight.shape[-2] // 2 + pair_idx = torch.empty(2 * half, dtype=torch.long) + pair_idx[0::2] = torch.arange(half) + half # up (w3) + pair_idx[1::2] = torch.arange(half) # gate (w1) + w13_weight = w13_weight[..., pair_idx, :].contiguous() + w13_weight_scale = w13_weight_scale[..., pair_idx, :].contiguous() + w13_bias = w13_bias[..., pair_idx].contiguous() # Shuffle weights and scaling factors for transposed mma output gemm1_weights_mxfp4_shuffled = [] @@ -658,7 +756,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): epilogue_tile_m, ) - for i in range(self.num_experts): + for i in range(E): gemm1_weights_mxfp4_shuffled.append( w13_weight[i] .view(torch.uint8)[w13_weight_permute_indices] @@ -699,7 +797,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w13_weight_scale = ( torch.stack(gemm1_scales_mxfp4_shuffled) .reshape( - self.num_experts, + E, 2 * self.intermediate_size_per_partition, self.hidden_size // sf_block_size, ) @@ -710,7 +808,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w2_weight_scale = ( torch.stack(gemm2_scales_mxfp4_shuffled) .reshape( - self.num_experts, + E, self.hidden_size, self.intermediate_size_per_partition // sf_block_size, ) @@ -722,44 +820,57 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): layer.w2_weight = Parameter(w2_weight, requires_grad=False) layer.w2_weight_scale = Parameter(w2_weight_scale, requires_grad=False) layer.w13_weight_bias = Parameter( - torch.stack(gemm1_bias_shuffled).reshape(self.num_experts, -1), + torch.stack(gemm1_bias_shuffled).reshape(E, -1), requires_grad=False, ) layer.w2_weight_bias = Parameter( - torch.stack(gemm2_bias_shuffled).reshape(self.num_experts, -1), + torch.stack(gemm2_bias_shuffled).reshape(E, -1), requires_grad=False, ) return if _use_aiter: - if layer.w13_weight_bias is not None: + if getattr(layer, "w13_weight_bias", None) is not None: layer.w13_weight_bias.data = layer.w13_weight_bias.data.to( torch.float32 ) - if layer.w2_weight_bias is not None: + if getattr(layer, "w2_weight_bias", None) is not None: layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32) e, n, k = layer.w13_weight.shape - layer.w13_weight.view(torch.uint8).copy_( - layer.w13_weight.data.view(torch.uint8) - .view(e, n // 2, 2, k) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, k) - ) - layer.w13_weight_scale.data = ( - layer.w13_weight_scale.data.view(e, n // 2, 2, -1) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, -1) - ) - layer.w13_weight_bias.data = ( - layer.w13_weight_bias.data.view(-1, n // 2, 2) - .permute(0, 2, 1) - .contiguous() - .view(-1, n) + gate_up_interleaved = getattr( + layer.moe_runner_config, "gate_up_interleaved", True ) - if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get(): + if gate_up_interleaved: + layer.w13_weight.view(torch.uint8).copy_( + layer.w13_weight.data.view(torch.uint8) + .view(e, n // 2, 2, k) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, k) + ) + layer.w13_weight_scale.data = ( + layer.w13_weight_scale.data.view(e, n // 2, 2, -1) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, -1) + ) + if getattr(layer, "w13_weight_bias", None) is not None: + layer.w13_weight_bias.data = ( + layer.w13_weight_bias.data.view(-1, n // 2, 2) + .permute(0, 2, 1) + .contiguous() + .view(-1, n) + ) + + k3_situ_a8w4 = ( + os.environ.get("AITER_SITUV2_A8W4", "0") == "1" + and getattr(layer.moe_runner_config, "activation", None) == "situ" + ) + use_aiter_gu_interleave = k3_situ_a8w4 or ( + envs.SGLANG_USE_AITER_MOE_GU_ITLV.get() and gate_up_interleaved + ) + if use_aiter_gu_interleave: layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True) shuffled_w13_scale = shuffle_scale_a16w4( layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]), @@ -925,11 +1036,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # half along its row dim (N) from N_un to N_pad with zeros, and along # its last dim (K) from K_un (or K_un / sf_block_size) to K_pad. + _interleaved = getattr(layer.moe_runner_config, "gate_up_interleaved", True) + def _stack_up_gate_w13(unpadded_w13, last_pad, last_un): # unpadded_w13: [E, 2*N_un, last_un] # Returns: [E, 2*N_pad, last_pad] in [up_padded; gate_padded] order. - gate_rows = unpadded_w13[:, 0::2, :] # [E, N_un, last_un] - up_rows = unpadded_w13[:, 1::2, :] # [E, N_un, last_un] + if _interleaved: + gate_rows = unpadded_w13[:, 0::2, :] # [E, N_un, last_un] + up_rows = unpadded_w13[:, 1::2, :] # [E, N_un, last_un] + else: + # Non-interleaved: first half=gate, second half=up + gate_rows = unpadded_w13[:, :N_un, :] + up_rows = unpadded_w13[:, N_un:, :] out = torch.zeros( E, 2 * N_pad, last_pad, dtype=unpadded_w13.dtype, device=device ) @@ -948,8 +1066,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): K_un // sf_block_size, ) # Bias: same de-interleave on dim=-1. - w13_bias_gate = layer.w13_weight_bias.data[:, 0::2] # [E, N_un] - w13_bias_up = layer.w13_weight_bias.data[:, 1::2] # [E, N_un] + if _interleaved: + w13_bias_gate = layer.w13_weight_bias.data[:, 0::2] # [E, N_un] + w13_bias_up = layer.w13_weight_bias.data[:, 1::2] # [E, N_un] + else: + w13_bias_gate = layer.w13_weight_bias.data[:, :N_un] + w13_bias_up = layer.w13_weight_bias.data[:, N_un:] w13_bias_padded = torch.zeros(E, 2 * N_pad, dtype=bias_dtype, device=device) w13_bias_padded[:, :N_un] = w13_bias_up w13_bias_padded[:, N_pad : N_pad + N_un] = w13_bias_gate @@ -971,9 +1093,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device) w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data - # ---- Per-expert SwiGLU scalars (GPT-OSS defaults) ------------------ + # ---- Per-expert SwiGLU scalars (read from runner config, fallback to GPT-OSS defaults) + _sm90_alpha = getattr(layer.moe_runner_config, "gemm1_alpha", None) or 1.702 + _sm90_limit = getattr(layer.moe_runner_config, "gemm1_clamp_limit", None) or 7.0 layer.swiglu_alpha = Parameter( - torch.full((E,), 1.702, dtype=torch.float32, device=device), + torch.full((E,), _sm90_alpha, dtype=torch.float32, device=device), requires_grad=False, ) layer.swiglu_beta = Parameter( @@ -981,7 +1105,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): requires_grad=False, ) layer.swiglu_limit = Parameter( - torch.full((E,), 7.0, dtype=torch.float32, device=device), + torch.full((E,), _sm90_limit, dtype=torch.float32, device=device), requires_grad=False, ) @@ -1121,14 +1245,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): moe_runner_backend = MoeRunnerBackend.TRITON if moe_runner_backend.is_aiter(): - # MXFP4 hard-codes Swiglu in the AITER kernel path. - self.runner = MoeRunner( - moe_runner_backend, replace(moe_runner_config, activation="swiglu") - ) + # MXFP4 hard-codes Swiglu in the AITER kernel path, so the + # checkpoint's "silu" has to be translated. K3's SiTU is the one + # activation the kernel selects on its own -- leave it alone. + aiter_config = moe_runner_config + if aiter_config.activation != "situ": + aiter_config = replace(aiter_config, activation="swiglu") + self.runner = MoeRunner(moe_runner_backend, aiter_config) elif ( moe_runner_backend.is_triton_kernels() or moe_runner_backend.is_triton() or moe_runner_backend.is_marlin() + or moe_runner_backend.is_deep_gemm() ): self.runner = MoeRunner(moe_runner_backend, moe_runner_config) elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in ( @@ -1172,6 +1300,28 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) return self.runner.run(dispatch_output, quant_info) + def _apply_marlin(self, layer, dispatch_output): + """MXFP4 x BF16 MoE via the Marlin runner. The quant_info (incl. the + dispatcher's EP mapping) is shared with + :class:`~sglang.srt.layers.quantization.mxfp4_marlin_moe.Mxfp4MarlinMoEMethod`; + only the input padding is local (``hidden_size`` is pre-rounded to the + Marlin-padded width at create_weights).""" + from sglang.srt.layers.quantization.mxfp4_marlin_moe import ( + build_marlin_moe_quant_info, + ) + + x = dispatch_output.hidden_states + if x.shape[-1] == self.hidden_size: + x_padded = x + else: + x_padded = torch.nn.functional.pad( + x, (0, self.hidden_pad), mode="constant", value=0.0 + ) + quant_info = build_marlin_moe_quant_info(layer) + return self.runner.run( + dispatch_output._replace(hidden_states=x_padded), quant_info + ) + def _apply_sm120_cutlass(self, layer, dispatch_output): """SM120 GPT-OSS MXFP8 x MXFP4 MoE via FlashInfer CUTLASS.""" from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import ( @@ -1206,6 +1356,24 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker + if self.use_deep_gemm: + # Handles standard AND deepep_ll/deepep_normal dispatch formats via + # the runner's registered pre/post-permute functions. Must run + # before the `.topk_output` unpack below: deepep dispatch outputs + # carry topk_ids/topk_weights directly and have no `.topk_output`. + from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo + + quant_info = DeepGemmMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + use_fp8=True, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + block_shape=[128, 128], + is_fp4_experts=True, + ) + return self.runner.run(dispatch_output, quant_info) + x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output if _is_cpu: @@ -1248,27 +1416,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): if self.use_marlin: assert TopKOutputChecker.format_is_standard(topk_output) - if x.shape[-1] == self.hidden_size: - x_padded = x - else: - x_padded = torch.nn.functional.pad( - x, (0, self.hidden_pad), mode="constant", value=0.0 - ) - quant_info = MarlinMoeQuantInfo( - w13_qweight=layer.w13_weight, - w2_qweight=layer.w2_weight, - w13_scales=layer.w13_weight_scale, - w2_scales=layer.w2_weight_scale, - w13_g_idx_sort_indices=None, - w2_g_idx_sort_indices=None, - weight_bits=4, - is_k_full=True, - w13_bias=getattr(layer, "w13_weight_bias", None), - w2_bias=getattr(layer, "w2_weight_bias", None), - ) - return self.runner.run( - dispatch_output._replace(hidden_states=x_padded), quant_info - ) + return self._apply_marlin(layer, dispatch_output) if self._fi_kernel == "cutlass_sm90": return self._apply_sm90_cutlass(layer, dispatch_output) @@ -1279,6 +1427,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations, # which can theoretically improve performance origin_hidden_states_dim = x.shape[-1] + # Filled by the staged K3 route+pack+quant fusion below; the pack + # site further down falls back to PackTopkIds when it is None. + prepared_packed_topk = None if self.flashinfer_mxfp4_moe_precision == "bf16": assert x.dtype == torch.bfloat16 x_quant = x @@ -1293,31 +1444,200 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): value=0.0, ) elif self.flashinfer_mxfp4_moe_precision == "default": - from sglang.srt.layers.quantization.fp8_utils import ( - flashinfer_mxfp8_quantize, - ) + if x.shape[-1] == self.hidden_size: + if x.dim() > 2: + x = x.view(-1, x.shape[-1]) + # K3 staged fusion (route_quant_handoff): the routing + # dispatch already quantized these rows and packed the + # topk ids in the fused route launch — consume both and + # skip the two standalone kernels. Identity-verified; + # a miss runs the unfused chain below. + from sglang.srt.layers.moe import route_quant_handoff - x_quant, x_scale = flashinfer_mxfp8_quantize( - x, False, alignment=self.hidden_size - ) - x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1) + prepared = route_quant_handoff.take(x) + if prepared is not None: + prepared_packed_topk, x_quant, x_scale = prepared + x_scale = x_scale.view(torch.float8_e4m3fn) + else: + from sglang.kernels.ops.quantization.per_token_group_quant import ( + per_token_group_quant, + ) + + x_quant, x_scale = per_token_group_quant( + x, group_size=32, scale_ue8m0=True + ) + x_scale = x_scale.view(torch.float8_e4m3fn) + else: + from sglang.srt.layers.quantization.fp8_utils import ( + flashinfer_mxfp8_quantize, + ) + + x_quant, x_scale = flashinfer_mxfp8_quantize( + x, False, alignment=self.hidden_size + ) + x_scale = x_scale.view(torch.float8_e4m3fn).reshape( + *x.shape[:-1], -1 + ) else: raise NotImplementedError() assert x_quant.shape[-1] == self.hidden_size - assert TopKOutputChecker.format_is_bypassed(topk_output) + is_standard = TopKOutputChecker.format_is_standard(topk_output) + # The situ path accepts precomputed (standard) routing; the + # public path below is logits-only. + assert is_standard or TopKOutputChecker.format_is_bypassed( + topk_output + ), f"unsupported topk format: {topk_output.format}" + if is_standard: + assert ( + self.moe_runner_config.activation == "situ" + ), "standard topk output only wired for the situ path" + top_k = topk_output.topk_ids.shape[1] + router_logits = None + else: + top_k = topk_output.topk_config.top_k + router_logits = topk_output.router_logits - top_k = topk_output.topk_config.top_k - router_logits = topk_output.router_logits + num_tokens = x_quant.shape[0] + hidden_size = origin_hidden_states_dim + # The K3 fused-front path publishes its [latent | shared] buffer + # slice as the output destination (zero_copy_context); writing + # the finalize output there directly skips this allocation and + # the copy_ back in _forward_fused. The slice lives in the same + # symmetric buffer the caller all-reduces. + symm_output = zero_copy_context.get_moe_output_spec( + torch.Size((num_tokens, hidden_size)), + torch.bfloat16, + x_quant.device, + ) + if symm_output is None: + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + symm_output = torch.empty( + num_tokens, + hidden_size, + dtype=torch.bfloat16, + device=x_quant.device, + ) - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): - num_tokens = x_quant.shape[0] - hidden_size = origin_hidden_states_dim - symm_output = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=x_quant.device + if self.moe_runner_config.activation == "situ": + # SiTU is only in the private trtllm-gen cubin pool (the + # public artifact bakes swiglu into the fused-act cubins and + # silently computes the wrong activation). Routing must also + # be noaux_tc (sigmoid + correction bias, DeepSeekV3 method), + # not the renormalize-softmax default below. + from sglang.kernels.ops.moe import trtllm_gen_moe as situ_moe + + if not situ_moe.available(): + raise RuntimeError( + "activation='situ' with the flashinfer_mxfp4 runner " + "needs the SiTU cubin pool: set " + "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL (see " + "sglang/kernels/ops/moe/trtllm_gen_moe.py)." + ) + # EP is cubin-internal: each rank computes its local expert slice + # [offset, +num_local) and the caller all-reduces. ep=1 -> TP path. + local_expert_offset = layer.moe_ep_rank * layer.num_local_experts + if TopKOutputChecker.format_is_standard(topk_output): + # Precomputed routing (radix router upstream): skip the + # in-op routing kernels entirely. At small T the in-op + # single-CTA routing costs ~22 us/layer vs ~6 us for the + # external radix router. + if prepared_packed_topk is not None: + packed_topk = prepared_packed_topk + else: + from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds + + packed_topk = PackTopkIds.execute( + topk_output.topk_ids, topk_output.topk_weights + ) + # Deferred finalize (K3 forward_deferred_finalize): return + # the finalize inputs instead of the finalized output. + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + _deferred_finalize_enabled, + ) + + defer_finalize = _deferred_finalize_enabled.get() + result = situ_moe.trtllm_fp4_block_scale_routed_moe( + packed_topk_ids=packed_topk, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=layer.w13_weight, + gemm1_weights_scale=layer.w13_weight_scale, + gemm1_alpha=layer.gemm1_alpha, + # SiTuGlu: gatedActBeta is the linear-half tanh + # clip; K3 stores it in gemm1_clamp_limit. + gemm1_beta=layer.gemm1_clamp_limit, + gemm2_weights=layer.w2_weight, + gemm2_weights_scale=layer.w2_weight_scale, + output1_scale_scalar=None, + output1_scale_gate_scalar=None, + output2_scale_scalar=None, + num_experts=layer.num_experts, + top_k=packed_topk.shape[1], + intermediate_size=self.intermediate_size_per_partition, + activation_type=situ_moe.ACTIVATION_SITU, + local_expert_offset=local_expert_offset, + local_num_experts=layer.num_local_experts, + output=symm_output, + do_finalize=not defer_finalize, + ) + if defer_finalize: + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + FlashInferTrtllmDeferredFinalizeOutput, + ) + + gemm2_out, topk_weights, expanded_idx = result + result = FlashInferTrtllmDeferredFinalizeOutput( + gemm2_out=gemm2_out, + expert_weights=topk_weights, + expanded_idx_to_permuted_idx=expanded_idx, + top_k=packed_topk.shape[1], + ) + return StandardCombineInput(hidden_states=result) + + # Bypassed topk: route from logits inside the op. + correction_bias = topk_output.topk_config.correction_bias + bias_bf16 = getattr(layer, "_situ_routing_bias_bf16", None) + if bias_bf16 is None and correction_bias is not None: + bias_bf16 = correction_bias.to(torch.bfloat16) + layer._situ_routing_bias_bf16 = bias_bf16 + situ_moe.trtllm_fp4_block_scale_moe( + # router_logits is a row-strided slice of the K3 fused + # front GEMM output; the FFI reads it as dense. + routing_logits=router_logits.to(torch.bfloat16).contiguous(), + routing_bias=bias_bf16, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=layer.w13_weight, + gemm1_weights_scale=layer.w13_weight_scale, + gemm1_alpha=layer.gemm1_alpha, + # SiTuGlu: gatedActBeta is the linear-half tanh clip; + # K3 stores it in gemm1_clamp_limit (situ_linear_beta). + gemm1_beta=layer.gemm1_clamp_limit, + gemm2_weights=layer.w2_weight, + gemm2_weights_scale=layer.w2_weight_scale, + output1_scale_scalar=None, + output1_scale_gate_scalar=None, + output2_scale_scalar=None, + num_experts=layer.num_experts, + top_k=top_k, + n_group=topk_output.topk_config.num_expert_group, + topk_group=topk_output.topk_config.topk_group, + intermediate_size=self.intermediate_size_per_partition, + routed_scaling_factor=( + topk_output.topk_config.routed_scaling_factor or 1.0 + ), + routing_method_type=situ_moe.ROUTING_DEEPSEEK_V3, + activation_type=situ_moe.ACTIVATION_SITU, + norm_topk_prob=topk_output.topk_config.renormalize, + local_expert_offset=local_expert_offset, + local_num_experts=layer.num_local_experts, + output=symm_output, ) + return StandardCombineInput(hidden_states=symm_output) + trtllm_gen_output = trtllm_fp4_block_scale_moe( router_logits.to(torch.bfloat16), None, # routing_bias @@ -1385,22 +1705,23 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): quant_type=AiterQuantType.PER_1X32, w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - b13=layer.w13_weight_bias, - b2=layer.w2_weight_bias, + b13=layer.w13_weight_bias if self.with_bias else None, + b2=layer.w2_weight_bias if self.with_bias else None, expert_mask=layer.dispatcher.expert_mask_gpu, doweight_stage1=self.moe_runner_config.apply_router_weight_on_input, hidden_pad=self.hidden_pad, intermediate_pad=self.intermediate_pad, - # Triggers aiter's INTERLEAVE gate_mode dispatch (required for our - # preshuffled gate/up-interleaved weight layout) and applies the - # model's swiglu clamp. Models populate the same scalar under - # different MoeRunnerConfig fields: gpt-oss uses `gemm1_clamp_limit` - # (renamed in `models/gpt_oss.py` from `config.swiglu_limit`); DSv4 - # / FP8 uses `swiglu_limit` directly. Accept either. + # Applies swiglu clamp for GPT-OSS-style activations. K3 SiTU + # uses gemm1_clamp_limit as linear_beta, which is forwarded by + # the AITER runner and must not be treated as swiglu_limit. swiglu_limit=( - self.moe_runner_config.gemm1_clamp_limit - or self.moe_runner_config.swiglu_limit - or 0.0 + 0.0 + if self.moe_runner_config.activation == "situ" + else ( + self.moe_runner_config.gemm1_clamp_limit + or self.moe_runner_config.swiglu_limit + or 0.0 + ) ), ) return self.runner.run( diff --git a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py index 3cdf7cbdc..7826032ef 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py @@ -17,6 +17,31 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def build_marlin_moe_quant_info(layer: Module) -> MarlinMoeQuantInfo: + """Build the Marlin quant_info for an MXFP4 MoE layer. + + Single source for the runner inputs shared by the marlin path of + ``Mxfp4MoEMethod.apply`` and :class:`Mxfp4MarlinMoEMethod`, including + the dispatcher's EP mapping (global -> local expert ids) when EP is on. + """ + expert_map = getattr(layer.dispatcher, "local_expert_mapping", None) + global_num_experts = layer.dispatcher.num_experts if expert_map is not None else -1 + return MarlinMoeQuantInfo( + w13_qweight=layer.w13_weight, + w2_qweight=layer.w2_weight, + w13_scales=layer.w13_weight_scale, + w2_scales=layer.w2_weight_scale, + w13_g_idx_sort_indices=None, + w2_g_idx_sort_indices=None, + weight_bits=4, + is_k_full=True, + w13_bias=getattr(layer, "w13_weight_bias", None), + w2_bias=getattr(layer, "w2_weight_bias", None), + expert_map=expert_map, + global_num_experts=global_num_experts, + ) + + class Mxfp4MarlinMoEMethod: """MXFP4 (E8M0 scales) MoE quantization method using the Marlin backend.""" @@ -162,18 +187,7 @@ class Mxfp4MarlinMoEMethod: value=0.0, ) - quant_info = MarlinMoeQuantInfo( - w13_qweight=layer.w13_weight, - w2_qweight=layer.w2_weight, - w13_scales=layer.w13_weight_scale, - w2_scales=layer.w2_weight_scale, - w13_g_idx_sort_indices=None, - w2_g_idx_sort_indices=None, - weight_bits=4, - is_k_full=True, - w13_bias=getattr(layer, "w13_weight_bias", None), - w2_bias=getattr(layer, "w2_weight_bias", None), - ) + quant_info = build_marlin_moe_quant_info(layer) runner_output = self.runner.run( dispatch_output._replace(hidden_states=hidden_states_padded), quant_info=quant_info, diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 15b9ee52d..5a66d91b2 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -254,6 +254,48 @@ class UnquantizedLinearMethod(LinearMethodBase): return F.linear(x, layer.weight, bias) + def apply_into( + self, + layer: torch.nn.Module, + x: torch.Tensor, + output: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Run an inference-only BF16 linear into caller-owned storage.""" + if ( + get_bf16_gemm_backend().is_cutedsl() + and x.is_cuda + and x.ndim == 2 + and x.dtype == torch.bfloat16 + and layer.weight.dtype == torch.bfloat16 + and output.dtype == torch.bfloat16 + and output.is_contiguous() + and output.shape == (x.shape[0], layer.weight.shape[0]) + and (bias is None or bias.dtype == torch.bfloat16) + and not layer.weight.requires_grad + and (bias is None or not bias.requires_grad) + and _use_cutedsl_bf16_gemm( + x.shape[0], layer.weight.shape[0], layer.weight.shape[1] + ) + ): + from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import ( + cutedsl_bf16_gemm_out, + ) + + return cutedsl_bf16_gemm_out(x, layer.weight, output, bias) + + if x.ndim != 2: + raise ValueError("caller-owned linear output currently requires a 2D input") + if output.shape != (x.shape[0], layer.weight.shape[0]): + raise ValueError( + f"linear output has shape {output.shape}, expected " + f"{(x.shape[0], layer.weight.shape[0])}" + ) + torch.mm(x, layer.weight.t(), out=output) + if bias is not None: + output.add_(bias) + return output + class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): """MoE method without quantization.""" diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 6696945d2..2dfc00d79 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -74,6 +74,7 @@ class RadixLinearAttention(nn.Module): self.A_log = A_log self.dt_bias = dt_bias + self.lower_bound = None def forward( self, diff --git a/python/sglang/srt/layers/zero_copy_context.py b/python/sglang/srt/layers/zero_copy_context.py new file mode 100644 index 000000000..ce2e32978 --- /dev/null +++ b/python/sglang/srt/layers/zero_copy_context.py @@ -0,0 +1,56 @@ +import contextlib +from typing import Iterator, Optional + +import msgspec +import torch + + +class ZeroCopyContext(msgspec.Struct, frozen=True): + moe_output: Optional[torch.Tensor] = None + + +ctx = ZeroCopyContext() + + +@contextlib.contextmanager +def set_moe_output(out: torch.Tensor) -> Iterator[None]: + """Publish `out` as the MoE runner's output destination for the block.""" + old_output = ctx.moe_output + msgspec.structs.force_setattr(ctx, "moe_output", out) + try: + yield + finally: + msgspec.structs.force_setattr(ctx, "moe_output", old_output) + + +def get_moe_output(ref: torch.Tensor) -> Optional[torch.Tensor]: + """The published destination iff it can stand in for empty_like(ref).""" + return get_moe_output_spec(ref.shape, ref.dtype, ref.device) + + +def get_moe_output_spec( + shape: torch.Size, dtype: torch.dtype, device: torch.device +) -> Optional[torch.Tensor]: + """Spec form of get_moe_output for callers that would otherwise have to + materialize a reference tensor just for the match (e.g. the trtllm-gen + runner, whose activation input is fp4-packed and shaped differently + from its output).""" + out = ctx.moe_output + if ( + out is not None + and out.shape == shape + and out.dtype == dtype + and out.device == device + and out.is_contiguous() + ): + return out + return None + + +__all__ = [ + "ZeroCopyContext", + "ctx", + "set_moe_output", + "get_moe_output", + "get_moe_output_spec", +] diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 125af0e00..5764a9f60 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -5,7 +5,7 @@ import logging import os import threading import time -from queue import Queue +from queue import Empty, Queue from typing import TYPE_CHECKING, Any, Callable, List, Optional import torch @@ -675,14 +675,62 @@ class HybridCacheController(BaseHiCacheController): operation.pool_transfers_done = True def _page_backup(self, operation): - # Backup extra pools - if operation.pool_transfers: + # MLA KV is replicated across TP ranks and should still be written only + # by TP0. On follower ranks, only the rank-sharded Mamba/KDA pool is + # owned by the rank and must be written here. Do not replicate other + # sidecar pools (for example SWA or indexer state) accidentally. + backup_transfers = operation.pool_transfers + if self.backup_skip: + backup_transfers = [ + transfer + for transfer in operation.pool_transfers or [] + if transfer.name == PoolName.MAMBA + ] + + if backup_transfers: self._resolve_sidecar_derived_pool_transfers(operation) - results = self.storage_backend.batch_set_v2(operation.pool_transfers) + results = self.storage_backend.batch_set_v2(backup_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) - # Backup kv pools - super()._page_backup(operation) + if not self.backup_skip: + super()._page_backup(operation) + else: + sidecar_ok = bool(backup_transfers) + if sidecar_ok: + for transfer in backup_transfers: + result = results.get(transfer.name) + if result is None: + result = results.get(transfer.name.value) + expected = len(transfer.keys or []) + if expected == 0 and transfer.host_indices is not None: + expected = int(transfer.host_indices.numel()) + if ( + not isinstance(result, (list, tuple)) + or len(result) != expected + or not all(bool(ok) for ok in result) + ): + sidecar_ok = False + break + operation.completed_tokens = ( + len(operation.hash_value) * self.page_size if sidecar_ok else 0 + ) + + def backup_thread_func(self): + """Back up rank-sharded sidecars on every TP rank. + + The base implementation skips the entire operation on non-zero MLA TP + ranks. That optimization is valid for replicated MLA KV, but not for + hybrid rank-sharded pools such as Kimi-K3 Mamba state. + """ + while not self.storage_stop_event.is_set(): + try: + operation = self.backup_queue.get(block=True, timeout=1) + if operation is None: + continue + self._page_backup(operation) + self.ack_backup_queue.put(operation) + except Empty: + continue def _resolve_sidecar_derived_pool_transfers(self, operation): for transfer in operation.pool_transfers: diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 1096cd958..21990fff5 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -27,7 +27,7 @@ from sglang.srt.mem_cache.pool_host.mha import ( get_mha_host_pool_cls, ) from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost -from sglang.srt.mem_cache.unified_cache.components import ComponentType +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.runtime_context import get_parallel if TYPE_CHECKING: diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index d93d38238..acd0f4137 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -8,7 +8,11 @@ from typing import TYPE_CHECKING, Any, Optional import msgspec import torch -from sglang.srt.configs.hybrid_arch import hybrid_gdn_config, mambaish_config +from sglang.srt.configs.hybrid_arch import ( + hybrid_gdn_config, + kimi_linear_config, + mambaish_config, +) from sglang.srt.configs.model_config import ( ModelConfig, get_dsa_index_head_dim, @@ -702,9 +706,16 @@ class KVCacheConfigurator: start_layer=self.layer_info.start_layer, linear_replayssm_cache_len=get_exec().mamba.linear_replayssm_cache_len, mamba_envelope_layout=get_memory().enable_page_major_kv_layout, - enable_gdn_replayssm_spec=( - get_exec().mamba.enable_gdn_replayssm_spec - and self.hybrid_gdn_config is not None + # ReplaySSM spec-verify is for linear-attn models (GDN fold or KDA + # fold); activate the pool machinery only for those, so any other + # mamba-ish model (Mamba2/Nemotron, lightning, ...) run with the + # flag set stays byte-identical to flag-off. + enable_linear_replayssm_spec=( + get_exec().mamba.enable_linear_replayssm_spec + and ( + self.hybrid_gdn_config is not None + or kimi_linear_config(self.model_config) is not None + ) ), ) return req_to_token_pool @@ -738,6 +749,18 @@ class KVCacheConfigurator: max_num_reqs: int, extra_max_context_len: int, ) -> ReqToTokenPool: + # DSPARK/DFLASH commit routes through the backend fold (KDA-only); a + # non-KDA model there would scatter a None intermediate_ssm and crash. + _algo = (self.server_args.speculative_algorithm or "").upper() + if ( + get_exec().mamba.enable_linear_replayssm_spec + and _algo in ("DSPARK", "DFLASH") + and kimi_linear_config(self.model_config) is None + ): + raise ValueError( + "--enable-linear-replayssm-spec with DSPARK/DFLASH requires a KDA " + "(kimi_linear) model; got a non-KDA model." + ) req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=get_schedule().max_mamba_cache_size, @@ -762,9 +785,16 @@ class KVCacheConfigurator: enable_linear_replayssm=get_exec().mamba.enable_linear_replayssm, linear_replayssm_cache_len=get_exec().mamba.linear_replayssm_cache_len, mamba_envelope_layout=get_memory().enable_page_major_kv_layout, - enable_gdn_replayssm_spec=( - get_exec().mamba.enable_gdn_replayssm_spec - and self.hybrid_gdn_config is not None + # ReplaySSM spec-verify is for linear-attn models (GDN fold or KDA + # fold); activate the pool machinery only for those, so any other + # mamba-ish model (Mamba2/Nemotron, lightning, ...) run with the + # flag set stays byte-identical to flag-off. + enable_linear_replayssm_spec=( + get_exec().mamba.enable_linear_replayssm_spec + and ( + self.hybrid_gdn_config is not None + or kimi_linear_config(self.model_config) is not None + ) ), ) return req_to_token_pool @@ -1688,11 +1718,23 @@ class KVCacheConfigurator: requested_per_worker = None max_num_reqs = min(estimated, token_capacity // 2) + capped_by_mamba = False if self.mambaish_config is not None: ratio = self._calculate_mamba_ratio() - max_num_reqs = min( - max_num_reqs, get_schedule().max_mamba_cache_size // ratio - ) + mamba_cap = get_schedule().max_mamba_cache_size // ratio + if mamba_cap < max_num_reqs: + capped_by_mamba = True + logger.warning( + "max_running_requests is capped to %d by the mamba state " + "cache (max_mamba_cache_size=%d, %d state slots per " + "request). To raise it: increase --mamba-full-memory-ratio " + "or --max-mamba-cache-size, or halve the state size with " + "--mamba-ssm-dtype bfloat16.", + mamba_cap, + get_schedule().max_mamba_cache_size, + ratio, + ) + max_num_reqs = min(max_num_reqs, mamba_cap) if max_num_reqs <= 0: raise RuntimeError( @@ -1703,7 +1745,11 @@ class KVCacheConfigurator: f"(2) increase --mem-fraction-static, or " f"(3) use GPUs with more memory." ) - if requested_per_worker is not None and max_num_reqs < requested_per_worker: + if ( + requested_per_worker is not None + and max_num_reqs < requested_per_worker + and not capped_by_mamba + ): logger.warning( "max_running_requests was reduced from the requested %d to %d " "(per dp worker) due to the available KV cache capacity.", @@ -1760,18 +1806,25 @@ class KVCacheConfigurator: assert config is not None has_spec_dec = not self.spec_algorithm.is_none() + # ReplaySSM drops the per-step intermediate_ssm scratch, so its mamba budget + # no longer reserves the (1 + D/ratio) intermediate factor -- the whole + # budget goes to persistent slots (K sized like non-spec), which is how the + # freed ~9GB turns into higher max_running. # The ring is allocated per slot but is not part of mamba_cache_per_req; # the solve must charge it too or num_slots is over-provisioned. - replayssm_active = ( - get_exec().mamba.enable_gdn_replayssm_spec - and self.hybrid_gdn_config is not None + replayssm_active = get_exec().mamba.enable_linear_replayssm_spec and ( + self.hybrid_gdn_config is not None + or kimi_linear_config(self.model_config) is not None ) if replayssm_active: - record_len = ( - server_args.max_speculative_num_draft_tokens - if server_args.max_speculative_num_draft_tokens is not None - else get_exec().mamba.linear_replayssm_cache_len - ) + # GDN sizes the fold window to the draft maximum; the KDA ring + # stays --linear-replayssm-cache-len long (mirrors MambaPool). + if kimi_linear_config(self.model_config) is not None: + record_len = get_exec().mamba.linear_replayssm_cache_len + elif server_args.max_speculative_num_draft_tokens is not None: + record_len = server_args.max_speculative_num_draft_tokens + else: + record_len = get_exec().mamba.linear_replayssm_cache_len replayssm_ring_per_req = ( config.mamba2_cache_params.replayssm_ring_bytes_per_req( record_len=record_len @@ -1790,7 +1843,9 @@ class KVCacheConfigurator: max_mamba_cache_size=get_schedule().max_mamba_cache_size // self.ps.attn_dp_size, ) - # Reserve intermediate memory based on capped max_num_reqs (+1 padding slot) + # Reserve intermediate memory based on capped max_num_reqs (+1: the + # pool's padding slot, see memory_pool.py). Skipped under replayssm + # (no intermediate_ssm allocated). if has_spec_dec and not replayssm_active: ratio = self._calculate_mamba_ratio() capped_reqs = min( @@ -1813,7 +1868,8 @@ class KVCacheConfigurator: max_mamba_cache_size=get_schedule().max_running_requests // self.ps.attn_dp_size, ) - # Reserve intermediate memory based on capped max_num_reqs (+1 padding slot) + # Reserve intermediate memory based on capped max_num_reqs (+1: the + # pool's padding slot). Skipped under replayssm. if has_spec_dec and not replayssm_active: intermediate_size = ( config.mamba2_cache_params.mamba_cache_per_req @@ -1880,7 +1936,9 @@ class KVCacheConfigurator: f"(4) use GPUs with more memory." ) - # +1: the pool's padding slot + # +1: the pool's padding slot is allocated alongside the request slots. + # ReplaySSM ring rides on every slot too (replayssm_ring_per_req is 0 when + # the ring is not allocated). mamba_state_memory = ( (get_schedule().max_mamba_cache_size + 1) * (config.mamba2_cache_params.mamba_cache_per_req + replayssm_ring_per_req) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index a57bd4c13..a9b2b8b99 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -344,7 +344,7 @@ class MambaPool: # replayssm_rawv: [num_layers, num_slots, HV, L, V] (conv/activation dtype) # replayssm_rawk: [num_layers, num_slots, H, L, K] (conv/activation dtype) # replayssm_beta: [num_layers, num_slots, HV, L] (fp32) - # The raw rings + beta exist only under --enable-gdn-replayssm-spec: the + # The raw rings + beta exist only under --enable-linear-replayssm-spec: the # closed-loop exact fold sequentially replays them through the recurrent # update at flush -- bit-identical to the recurrent baseline -- instead # of folding the chunked `d` records open-loop (which accumulates error @@ -380,7 +380,7 @@ class MambaPool: @dataclass(frozen=True, kw_only=True) class SpeculativeState(State): - # None under --enable-gdn-replayssm-spec: the spec ring owns rollback + # None under --enable-linear-replayssm-spec: the spec ring owns rollback # (verify writes ring records, commit moves cursors), so the per-draft # full-state snapshots are never produced or consumed. intermediate_ssm: Optional[torch.Tensor] @@ -468,7 +468,7 @@ class MambaPool: enable_linear_replayssm: bool = False, linear_replayssm_cache_len: int = 16, envelope_layout: bool = False, - enable_gdn_replayssm_spec: bool = False, + enable_linear_replayssm_spec: bool = False, ): conv_state_shape = cache_params.shape.conv temporal_state_shape = cache_params.shape.temporal @@ -487,12 +487,13 @@ class MambaPool: self.linear_replayssm_cache_len = linear_replayssm_cache_len # ReplaySSM: the decode ring (--enable-linear-replayssm) allocates the # chunked (d, k) records + write_pos; the spec-verify flag - # (--enable-gdn-replayssm-spec) always uses fold-every-commit and + # (--enable-linear-replayssm-spec) always uses fold-every-commit and # allocates only the raw (v, k, g, beta) window -- no chunked records, - # no cursors. The shared g allocation gates on `_replayssm_on`. - self.enable_gdn_replayssm_spec = enable_gdn_replayssm_spec - self.replayssm_spec_fold = bool(enable_gdn_replayssm_spec) - _replayssm_on = enable_linear_replayssm or enable_gdn_replayssm_spec + # no cursors (KDA additionally keeps d/k, see the allocation below). + # The shared g allocation gates on `_replayssm_on`. + self.enable_linear_replayssm_spec = enable_linear_replayssm_spec + self.replayssm_spec_fold = bool(enable_linear_replayssm_spec) + _replayssm_on = enable_linear_replayssm or enable_linear_replayssm_spec # for disagg with nvlink self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( @@ -570,7 +571,7 @@ class MambaPool: # flag is on; otherwise left as None so the legacy State is # byte-identical. temporal_state_shape == (HV, V, K). Either the decode # ring (--enable-linear-replayssm) or the spec-verify ring - # (--enable-gdn-replayssm-spec) shares this allocation. + # (--enable-linear-replayssm-spec) shares this allocation. replayssm_d = replayssm_k = replayssm_g = None replayssm_rawv = replayssm_rawk = replayssm_beta = None if _replayssm_on: @@ -580,16 +581,23 @@ class MambaPool: num_slots = size + 1 # Ring dtype. DECODE ring (--enable-linear-replayssm): records # follow the SSM dtype -- its flush folds `d` directly into the - # state. SPEC-verify ring (--enable-gdn-replayssm-spec): d/k feed + # state. SPEC-verify ring (--enable-linear-replayssm-spec): d/k feed # ONLY the one-shot output reconstruction (the closed-loop exact # fold replays the raw rings for state instead), so their # quantization noise stays below the bf16 output cast; keep them # in the conv/activation dtype instead of the (fp32-enforced) # SSM dtype to halve the ring traffic. g stays fp32 everywhere # (exact-fold input). The two flags are mutually exclusive. - ring_dtype = conv_dtype if enable_gdn_replayssm_spec else ssm_dtype - # Fold-every-commit: one verify window, no chunked (d, k) records. - if self.replayssm_spec_fold: + ring_dtype = conv_dtype if enable_linear_replayssm_spec else ssm_dtype + # Fold-every-commit: one verify window, no chunked (d, k) + # records. KDA is the exception on both counts: its window + # stays L-sized (the fused verify ring-write drops + # absorb-inflated rows past L), and d/k stay allocated -- + # forward_decode routes on `replayssm_d is None` (fused vs + # decode-ring), so skipping them would flip KDA decode to the + # fused path, a behavior change needing its own validation + # (memory follow-up). + if self.replayssm_spec_fold and not cache_params.is_kda: record_len = ( speculative_num_draft_tokens if speculative_num_draft_tokens is not None @@ -597,6 +605,7 @@ class MambaPool: ) else: record_len = L + if not self.replayssm_spec_fold or cache_params.is_kda: replayssm_d = torch.zeros( size=(num_mamba_layers, num_slots, hv, L, v_dim), dtype=ring_dtype, @@ -626,7 +635,22 @@ class MambaPool: # flush replays these through the recurrent update sequentially # (bit-identical to the recurrent baseline) instead of folding # the chunked `d` records open-loop. - if enable_gdn_replayssm_spec: + if enable_linear_replayssm_spec: + if cache_params.is_kda: + # Backstop for the KDA ring invariants; this pool is + # sized with the final adaptive-aware draft maximum. + if L & (L - 1) != 0: + raise ValueError( + f"spec-verify ring length must be a power of two, got {L}" + ) + if ( + speculative_num_draft_tokens is not None + and L < 2 * speculative_num_draft_tokens + ): + raise ValueError( + f"spec-verify ring too small: {L} < " + f"2 * {speculative_num_draft_tokens} (early-flush margin)" + ) replayssm_rawv = torch.zeros( size=(num_mamba_layers, num_slots, hv, record_len, v_dim), dtype=conv_dtype, @@ -662,7 +686,11 @@ class MambaPool: # The recurrent-verify fallback cannot be reached under the flag # (GDN + linear chain + triton enforced in server_args; the # backend asserts loudly if it ever is). - if enable_gdn_replayssm_spec: + # ReplaySSM skips this dominant scratch (~9GB @ K3 dspark γ=7): the + # KDA verify kernel takes intermediate_states_buffer=None (skips the + # per-step write, CACHE_INTERMEDIATE_STATES=False) and the commit + # replays the ring into the checkpoint instead. This is the memory win. + if enable_linear_replayssm_spec: intermediate_ssm_state_cache = None else: intermediate_ssm_state_cache = torch.zeros( @@ -795,7 +823,7 @@ class MambaPool: f"rawv={get_tensor_size_bytes(replayssm_rawv) / GB:.3f}GB, " f"rawk={get_tensor_size_bytes(replayssm_rawk) / GB:.3f}GB, " f"beta={get_tensor_size_bytes(replayssm_beta) / GB:.3f}GB " - if enable_gdn_replayssm_spec + if enable_linear_replayssm_spec else "" ) ) @@ -818,12 +846,12 @@ class MambaPool: # all GDN layers of one verify step; advanced by commit_gdn_replayssm_spec. self.replayssm_cache_base = ( torch.zeros((size + 1,), dtype=torch.int32, device=device) - if enable_gdn_replayssm_spec and not self.replayssm_spec_fold + if enable_linear_replayssm_spec and not self.replayssm_spec_fold else None ) self.replayssm_is_flush = ( torch.zeros((size + 1,), dtype=torch.int8, device=device) - if enable_gdn_replayssm_spec and not self.replayssm_spec_fold + if enable_linear_replayssm_spec and not self.replayssm_spec_fold else None ) mem_usage_bytes = self.mamba_cache.mem_usage_bytes() @@ -1130,7 +1158,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm: bool = False, linear_replayssm_cache_len: int = 16, mamba_envelope_layout: bool = False, - enable_gdn_replayssm_spec: bool = False, + enable_linear_replayssm_spec: bool = False, ): super().__init__( size=size, @@ -1157,7 +1185,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm=enable_linear_replayssm, linear_replayssm_cache_len=linear_replayssm_cache_len, mamba_envelope_layout=mamba_envelope_layout, - enable_gdn_replayssm_spec=enable_gdn_replayssm_spec, + enable_linear_replayssm_spec=enable_linear_replayssm_spec, ) def _init_mamba_pool( @@ -1173,7 +1201,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm: bool = False, linear_replayssm_cache_len: int = 16, mamba_envelope_layout: bool = False, - enable_gdn_replayssm_spec: bool = False, + enable_linear_replayssm_spec: bool = False, ): self.mamba_pool = self.mamba_pool_cls( size=mamba_size, @@ -1187,7 +1215,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_linear_replayssm=enable_linear_replayssm, linear_replayssm_cache_len=linear_replayssm_cache_len, envelope_layout=mamba_envelope_layout, - enable_gdn_replayssm_spec=enable_gdn_replayssm_spec, + enable_linear_replayssm_spec=enable_linear_replayssm_spec, ) self.mamba_allocator = MambaSlotAllocator( size=mamba_size, @@ -3663,7 +3691,7 @@ class HybridLinearKVPool(KVCache): def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]: # Hybrid layer ids are global model-layer ids, while the backing pool - # is dense over only full-attention layers. Shape discovery does not + # is dense over only full-attention layers. Shape discovery does not # need a global layer lookup, so delegate it to that backing pool. return self.full_kv_pool.get_kv_buffer_shape() diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 05725ac50..1575616ac 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -710,13 +710,14 @@ class UnifiedMambaPool(MambaPool): self.enable_custom_mem_pool = False self.custom_mem_pool = None self.num_mamba_layers = spec.layer_num - # GDN/KDA ReplaySSM unsupported; replicate parent's disabled-state attrs so - # paths guarded by `replayssm_write_pos is not None` don't AttributeError. + # GDN/KDA ReplaySSM / spec unsupported; replicate parent's disabled-state + # attrs so unconditional reads (e.g. `replayssm_cache_base is not None` in + # the req-slot alloc path) and `... is not None` guards don't AttributeError. self.enable_linear_replayssm = False self.linear_replayssm_cache_len = 16 self.replayssm_write_pos = None self.replayssm_is_kda = False - self.enable_gdn_replayssm_spec = False + self.enable_linear_replayssm_spec = False self.replayssm_spec_fold = False self.replayssm_cache_base = None self.replayssm_is_flush = None @@ -941,10 +942,10 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool): mamba_envelope_layout: bool = False, enable_linear_replayssm: bool = False, linear_replayssm_cache_len: int = 16, - enable_gdn_replayssm_spec: bool = False, + enable_linear_replayssm_spec: bool = False, ): # mamba_envelope_layout / speculative_eagle_topk / enable_linear_replayssm / - # linear_replayssm_cache_len / enable_gdn_replayssm_spec: accepted to match + # linear_replayssm_cache_len / enable_linear_replayssm_spec: accepted to match # the parent signature but NOT forwarded — the shared pool's conv/temporal # state are fixed-shape views (replayssm/spec are gated off under unified). assert mamba_size == self._shared_mamba_size, ( diff --git a/python/sglang/srt/model_executor/input_buffers.py b/python/sglang/srt/model_executor/input_buffers.py index 444522f21..1a6702126 100644 --- a/python/sglang/srt/model_executor/input_buffers.py +++ b/python/sglang/srt/model_executor/input_buffers.py @@ -66,12 +66,35 @@ def share_input_buffers_in(obj) -> None: setattr(obj, name, share_input_buffer(name, buffer)) +# Values that index the rope table, the KV pool, req_to_token, or the mamba +# state pool, so stale content is unsafe to execute. +_INDEX_SEMANTIC_BUFFERS = frozenset( + { + "positions", + "mrope_positions", + "out_cache_loc", + "req_pool_indices", + "mamba_track_indices", + "mamba_track_mask", + } +) + + @dataclass class ForwardInputBuffers: def _share_one_buffer(self, name: str, new_buffer: torch.Tensor) -> torch.Tensor: return share_input_buffer(name, new_buffer) + def reset_index_buffers(self) -> None: + """Zero the index-semantic buffers this set declares.""" + for f in fields(self): + if f.name not in _INDEX_SEMANTIC_BUFFERS: + continue + buffer = getattr(self, f.name) + if buffer is not None: + buffer.zero_() + def share_buffers(self): # disable share input buffer on npu due to accuracy issue if is_npu(): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ce9b04c43..8e957f661 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -134,6 +134,7 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import ( load_model_with_memory_saver, maybe_downgrade_dtype_for_legacy_gpu, maybe_enable_ipc_weight_cache, + maybe_precompile_model_kernels_after_loading, maybe_register_debug_tensor_dump_hook, maybe_trigger_remote_instance_nccl_send_group, report_online_quantization, @@ -736,6 +737,14 @@ class ModelRunner: start_layer=self.layer_info.start_layer, ) + def get_pp_proxy_residual_num_blocks(self) -> Optional[int]: + return misc_utils.resolve_pp_proxy_residual_num_blocks( + model_config=self.model_config, + pp_size=self.ps.pp_size, + pp_rank=self.ps.pp_rank, + start_layer=self.layer_info.start_layer, + ) + def decode_num_tokens_per_req( self, *, num_draft_tokens: Optional[int] = None ) -> int: @@ -1060,6 +1069,8 @@ class ModelRunner: if not self.is_draft_worker: get_offloader().post_init() + self.maybe_precompile_model_kernels_after_loading() + # Register model for layerwise NVTX profiling if enabled if get_exec().comm.enable_layerwise_nvtx_marker: pyt_hooks = PytHooks() @@ -1124,6 +1135,9 @@ class ModelRunner: is_ep_joiner=self.server_args.is_ep_joiner, ) + def maybe_precompile_model_kernels_after_loading(self) -> None: + maybe_precompile_model_kernels_after_loading(self.model, self.device) + def maybe_init_dwdp(self): if self.is_draft_worker: return @@ -1190,7 +1204,9 @@ class ModelRunner: model_dtype=getattr(self, "dtype", torch.bfloat16), is_draft_worker=getattr(self, "is_draft_worker", False), is_dflash=( - spec_algorithm.is_dflash() if spec_algorithm is not None else False + spec_algorithm.is_dflash_family() + if spec_algorithm is not None + else False ), speculative_draft_attention_backend=getattr( self.server_args, "speculative_draft_attention_backend", None diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index f91339d5e..f3ada403b 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -24,6 +24,7 @@ from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( RemoteInstanceWeightLoaderBackend, trigger_init_weights_send_group_for_remote_instance_request, ) +from sglang.srt.platforms import current_platform from sglang.srt.utils.common import is_npu from sglang.srt.utils.network import NetworkAddress @@ -40,6 +41,20 @@ _is_npu = is_npu() UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data processing +def maybe_precompile_model_kernels_after_loading(model, device: str) -> None: + precompile = getattr(model, "precompile_kernels_after_loading", None) + if precompile is None: + return + + if device == "cuda": + current_platform.synchronize() + current_platform.empty_cache() + precompile() + if device == "cuda": + current_platform.synchronize() + current_platform.empty_cache() + + class LoadedModel(msgspec.Struct, frozen=True, kw_only=True): loader: Any model: Any diff --git a/python/sglang/srt/model_executor/model_runner_components/misc_utils.py b/python/sglang/srt/model_executor/model_runner_components/misc_utils.py index d2308d123..653b3a18c 100644 --- a/python/sglang/srt/model_executor/model_runner_components/misc_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/misc_utils.py @@ -3,7 +3,11 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any, Optional -from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa +from sglang.srt.configs.model_config import ( + dsa_layer_skips_topk, + is_deepseek_dsa, + is_kimi_k3, +) from sglang.srt.runtime_context import get_context, get_exec, get_schedule from sglang.srt.server_args import CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS @@ -68,3 +72,16 @@ def resolve_pp_proxy_topk_size( ): return None return getattr(hf_config, "index_topk", None) + + +def resolve_pp_proxy_residual_num_blocks( + *, model_config: ModelConfig, pp_size: int, pp_rank: int, start_layer: int +) -> Optional[int]: + """Return the inherited Kimi K3 attention-residual bank width.""" + if pp_size <= 1 or pp_rank == 0 or not is_kimi_k3(model_config.hf_config): + return None + + block_size = getattr(model_config.hf_text_config, "attn_res_block_size", None) + if block_size is None: + return None + return (start_layer + block_size - 1) // block_size diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 7664b66d5..656b070d5 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -81,6 +81,7 @@ def _allocate_decode_buffers( ne_token_table: Optional[torch.Tensor] = None, hc_hidden_size: Optional[int] = None, pp_proxy_topk_size: Optional[int] = None, + pp_proxy_residual_num_blocks: Optional[int] = None, ) -> SimpleNamespace: """Allocate the FB-shared decode buffers.""" with torch.device(device): @@ -115,9 +116,14 @@ def _allocate_decode_buffers( "hidden_states": torch.zeros((max_bs, hs), dtype=dtype), } if not is_mhc: - pp_proxy_tensors["residual"] = torch.zeros( - (max_bs, hidden_size), dtype=dtype + # Only Kimi K3 supplies num_blocks: its PP bank is token-major + # [T, blocks, H]. Other models keep the legacy [max_bs, H]. + residual_shape = ( + (max_num_token, pp_proxy_residual_num_blocks, hidden_size) + if pp_proxy_residual_num_blocks is not None + else (max_bs, hidden_size) ) + pp_proxy_tensors["residual"] = torch.zeros(residual_shape, dtype=dtype) if pp_proxy_topk_size is not None: pp_proxy_tensors["topk_indices"] = torch.zeros( (max_num_token, pp_proxy_topk_size), dtype=torch.int32 @@ -338,6 +344,7 @@ class BaseRunner(ABC): ), hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None), pp_proxy_topk_size=mr.get_pp_proxy_topk_size(), + pp_proxy_residual_num_blocks=mr.get_pp_proxy_residual_num_blocks(), ) def _dummy_run( diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 7323effe8..c7a20c8ef 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -297,6 +297,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): else None ) self._ragged_graph_size = 0 + # Per-tier capture layouts; their verify_lens / qo_indptr tensors are + # baked into the captured graphs and refreshed in place each replay. + self._captured_ragged_layouts: dict[int, object] = {} if self.ragged_verify_mode and ( self.enable_two_batch_overlap or model_runner.server_args.enable_lora @@ -377,6 +380,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): self.model_runner.model_config, "hc_hidden_size", None ), pp_proxy_topk_size=self.model_runner.get_pp_proxy_topk_size(), + pp_proxy_residual_num_blocks=( + self.model_runner.get_pp_proxy_residual_num_blocks() + ), ) self.buffers.share_buffers() # FB-shared slot registry adopting DecodeInputBuffers storage (same @@ -495,11 +501,27 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): num_slots=self._ragged_capture_slots(num_tokens), num_draft_tokens=self.captured_req_width, ) - return RaggedVerifyLayout.from_verify_lens( + layout = RaggedVerifyLayout.from_verify_lens( verify_lens_cpu=verify_lens_cpu, device=self.device, grid=self.capture_num_tokens, ) + self._captured_ragged_layouts[num_tokens] = layout + return layout + + def _stage_ragged_verify_layout(self, ragged_layout, graph_size_key: int) -> None: + # Without this refresh every replay reuses the capture-time synthetic + # verify_lens / qo_indptr and mis-slices the packed q rows. + cap_layout = self._captured_ragged_layouts.get(graph_size_key) + if cap_layout is None: + return + live = ragged_layout + if live.bs != cap_layout.bs or live.cap is None: + live = live.padded_to_bucket( + padded_bs=cap_layout.bs, cap=self.captured_req_width + ) + cap_layout.verify_lens.copy_(live.verify_lens) + cap_layout.qo_indptr_device.copy_(live.qo_indptr_device) def can_run_graph(self, forward_batch: ForwardBatch): # Disable for token embedding overrides (dynamic per-request) @@ -823,6 +845,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): # captured graph needs before capturing. self.buffers.seq_lens.fill_(self.seq_len_fill_value) self.buffers.seq_lens_cpu.fill_(self.seq_len_fill_value) + # Capture runs real forwards, so a mid-serving recapture would index -- + # and write KV -- through the previous batch's live values. Replay is + # already covered by the registry's padding policy. + self.buffers.reset_index_buffers() # Trigger CUDA graph capture for specific shapes. # Capture the large shapes first so that the smaller shapes @@ -1036,6 +1062,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): f"stale ragged raw_num_token {self.raw_num_token} != " f"{ragged_layout.graph_num_tokens}" ) + self._stage_ragged_verify_layout(ragged_layout, graph_size_key) self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) if ( @@ -1072,6 +1099,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): f"{raw_bs}; the planner must reject this batch before replay" ) padded_num_tokens = graph_size_key + self._stage_ragged_verify_layout(ragged_layout, graph_size_key) else: raw_num_token = raw_bs * self.captured_req_width if self.require_mlp_tp_gather: diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index aa50b7255..20e9e83f2 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -337,6 +337,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): self.has_mha_companion_layers = any( layer is not None for layer in self.mha_companion_layers ) + # Archs on the MLA-BCG allowlist pin the absorbed MLA path inside + # capture/replay (attention_backend_handler), so the MHA companion is + # never captured and the MHA-prefix restrictions below don't apply. + self.mla_pinned_under_bcg = ( + self.model_runner.model_config.is_mla_breakable_cuda_graph_supported + ) self.moe_layers = self.model_runner.moe_layers self.moe_fusions = self.model_runner.moe_fusions self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None) @@ -1052,11 +1058,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): return False # A prefix forces the MHA companion path, whose captured state is # frozen prefix-free; DSA models are exempt (capture/replay force - # the sparse path, which takes any prefix via device metadata). + # the sparse path, which takes any prefix via device metadata), as + # are archs on the MLA-BCG allowlist (they pin the absorbed MLA path + # inside capture/replay, so the MHA companion is never captured). if ( self.prefill_backend_name == Backend.BREAKABLE and self.has_mha_companion_layers and not self.dsa_sparse_prefill_forced + and not self.mla_pinned_under_bcg and prefix_lens is not None and any(prefix_lens) ): @@ -1536,6 +1545,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): if ( isinstance(self.backend, BreakableCudaGraphBackend) and self.has_mha_companion_layers + and not self.mla_pinned_under_bcg ): self._restore_mha_capture_state(static_forward_batch) diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 71c506427..97c7702c7 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/python/sglang/srt/model_executor/runner_utils/buffers.py @@ -105,6 +105,7 @@ class DecodeInputBuffers(ForwardInputBuffers): ne_token_table: Optional[torch.Tensor] = None, hc_hidden_size: Optional[int] = None, pp_proxy_topk_size: Optional[int] = None, + pp_proxy_residual_num_blocks: Optional[int] = None, ) -> DecodeInputBuffers: with torch.device(device): input_ids = torch.zeros((max_num_token,), dtype=torch.int64) @@ -135,8 +136,15 @@ class DecodeInputBuffers(ForwardInputBuffers): "hidden_states": torch.zeros((max_bs, hs), dtype=dtype), } if not is_mhc: + # Only Kimi K3 supplies num_blocks: its PP bank is token-major + # [T, blocks, H]. Other models keep the legacy [max_bs, H]. + residual_shape = ( + (max_num_token, pp_proxy_residual_num_blocks, hidden_size) + if pp_proxy_residual_num_blocks is not None + else (max_bs, hidden_size) + ) pp_proxy_tensors["residual"] = torch.zeros( - (max_bs, hidden_size), dtype=dtype + residual_shape, dtype=dtype ) if pp_proxy_topk_size is not None: pp_proxy_tensors["topk_indices"] = torch.zeros( diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 9445213bb..92fab1bc3 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -80,7 +80,10 @@ def _support_mha_one_shot(attn, forward_batch, backend_name): def _handle_attention_backend(attn, forward_batch, backend_name): - if is_in_tc_piecewise_cuda_graph(): + # Captured prefill (tc_piecewise or breakable) must keep a single attention + # path: pin the absorbed MLA method — MHA one-shot/chunked shapes vary with + # kv-len and cannot be captured. + if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): return AttnForwardMethod.MLA # MLA prefill CP forces absorbed MLA regardless of prefix length: the @@ -145,7 +148,7 @@ def handle_attention_fa4(attn, forward_batch): def handle_attention_trtllm_mla(attn, forward_batch): - if is_in_tc_piecewise_cuda_graph(): + if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): return AttnForwardMethod.MLA sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch) @@ -190,7 +193,7 @@ def handle_attention_dsa(attn, forward_batch): def handle_attention_triton(attn, forward_batch): - if is_in_tc_piecewise_cuda_graph(): + if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): return AttnForwardMethod.MLA # when deterministic inference is enabled, use MLA diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 07d1502f1..49b495880 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -27,6 +27,7 @@ from sglang.srt.layers.dcp import ( cp_lse_ag_out_rs_mla, dcp_a2a_lse_reduce, ) +from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run from sglang.srt.layers.quantization.fp8_utils import ( materialize_bpreshuffle_fp8_scale_tuple, ) @@ -68,7 +69,7 @@ from sglang.srt.models.deepseek_common.utils import ( _use_aiter_bpreshuffle_gfx95, _use_aiter_gfx95, ) -from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args, get_spec +from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args from sglang.srt.state_capturer.indexer_topk import ( maybe_capture_indexer_topk, ) @@ -92,23 +93,12 @@ class MlaBmmFusionPlan: attn_output_buf: torch.Tensor -def _is_dcp_mla_decode_phase(forward_batch: ForwardBatch) -> bool: - if not get_parallel().dcp_enabled: - return False - if forward_batch.forward_mode.is_decode(): - return True - if not forward_batch.forward_mode.is_target_verify() or not _is_cuda: - return False - - server_args = get_server_args() - decode_backend = ( - server_args.decode_attention_backend or server_args.attention_backend - ) - return ( - get_spec().speculative_algorithm == "DSPARK" - and get_spec().speculative_attention_mode == "decode" - and decode_backend in ("tokenspeed_mla", "cutedsl_mla") - ) +def _select_local_dcp_heads_for_autotune( + attn_output: torch.Tensor, num_local_heads: int +) -> torch.Tensor: + """Select this rank's head shard without communicating dummy outputs.""" + rank = get_parallel().attn_dcp_rank + return attn_output.narrow(1, rank * num_local_heads, num_local_heads) def _is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool: @@ -366,7 +356,11 @@ class DeepseekMLAForwardMixin: # weights and skip the per-layer Q all-gather (bf16 decode absorb only). q_replicate_active = ( get_parallel().dcp_replicate_q_proj - and _is_dcp_mla_decode_phase(forward_batch) + and get_parallel().dcp_enabled + and ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + ) and not self.use_deep_gemm_bmm and self.w_kc_qrep is not None and self.q_b_proj_qrep_weight is not None @@ -802,7 +796,10 @@ class DeepseekMLAForwardMixin: # all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch. if get_parallel().dcp_enabled: - if _is_dcp_mla_decode_phase(forward_batch): + if ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + ): if not q_replicate_active: q_nope_out, q_pe = all_gather_q_for_mla_decode( q_nope_out=q_nope_out, @@ -955,7 +952,10 @@ class DeepseekMLAForwardMixin: topk_indices=topk_indices, ) attn_output = fusion_plan.attn_output_buf - elif _is_dcp_mla_decode_phase(forward_batch): + elif ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + ) and get_parallel().dcp_enabled: # set return_lse=True to correct attn_output attn_output, lse = self.attn_mqa_for_dcp_decode( q_nope_out, @@ -1029,31 +1029,44 @@ class DeepseekMLAForwardMixin: ) # correct attn_output with respect to lse from other ranks - if _is_dcp_mla_decode_phase(forward_batch): + if ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + ) and get_parallel().dcp_enabled: attn_output = attn_output.view( -1, self.num_local_heads * get_parallel().attn_dcp_size, self.kv_lora_rank, ) - dcp_comm_backend = get_parallel().dcp_comm_backend - is_lse_base_on_e = _is_mla_dcp_lse_base_on_e(self.current_attention_backend) - if dcp_comm_backend in ("a2a", "fi_a2a"): - # A2A exchange of head partials + LSE, then local Triton combine. - attn_output = dcp_a2a_lse_reduce( - attn_output.contiguous(), - lse.contiguous(), - get_parallel().dcp_group, - is_lse_base_on_e=is_lse_base_on_e, - comm_backend=dcp_comm_backend, + if get_in_autotune_dummy_run(): + # The synthetic FlashInfer MoE autotune pass discards model + # outputs. Avoid an unnecessary cross-node MNNVL exchange of + # zero attention partials. + attn_output = _select_local_dcp_heads_for_autotune( + attn_output, self.num_local_heads ) else: - attn_output = cp_lse_ag_out_rs_mla( - attn_output, - lse, - get_parallel().dcp_group, - is_lse_base_on_e=is_lse_base_on_e, + dcp_comm_backend = get_parallel().dcp_comm_backend + is_lse_base_on_e = _is_mla_dcp_lse_base_on_e( + self.current_attention_backend ) - attn_output = attn_output.transpose(0, 1) + if dcp_comm_backend in ("a2a", "fi_a2a"): + # A2A exchange of head partials + LSE, then local Triton combine. + attn_output = dcp_a2a_lse_reduce( + attn_output.contiguous(), + lse.contiguous(), + get_parallel().dcp_group, + is_lse_base_on_e=is_lse_base_on_e, + comm_backend=dcp_comm_backend, + ) + else: + attn_output = cp_lse_ag_out_rs_mla( + attn_output, + lse, + get_parallel().dcp_group, + is_lse_base_on_e=is_lse_base_on_e, + ) + attn_output = attn_output.transpose(0, 1) attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) _kvb_v = None diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index a1ca3f145..f9ff64733 100644 --- a/python/sglang/srt/models/dspark.py +++ b/python/sglang/srt/models/dspark.py @@ -596,11 +596,17 @@ class DSparkDraftMixin: class DSparkDraftModel(DSparkDraftMixin, DFlashDraftModel): - pass + def prune_to_ctx_kv_injection(self) -> None: + self.markov_head = None + self.confidence_head = None + for layer in self.layers: + layer.mlp = None + layer.self_attn.o_proj = None + torch.cuda.empty_cache() class Qwen3DSparkModel(DSparkDraftModel): pass -EntryClass = [Qwen3DSparkModel] +EntryClass = [Qwen3DSparkModel, DSparkDraftModel] diff --git a/python/sglang/srt/models/inkling_common/dense_mlp.py b/python/sglang/srt/models/inkling_common/dense_mlp.py index cafaafc60..66d3edadc 100644 --- a/python/sglang/srt/models/inkling_common/dense_mlp.py +++ b/python/sglang/srt/models/inkling_common/dense_mlp.py @@ -217,6 +217,7 @@ class InklingBatchDenseMLP(nn.Module, FusedMoELoadingMixin): self.quant_method, self.moe_runner_config, self.moe_tp_rank, + self.moe_tp_size, ) self.quant_method.create_weights( layer=self, diff --git a/python/sglang/srt/models/inkling_common/util.py b/python/sglang/srt/models/inkling_common/util.py index 19c27da0e..b7ec0a13f 100644 --- a/python/sglang/srt/models/inkling_common/util.py +++ b/python/sglang/srt/models/inkling_common/util.py @@ -88,6 +88,7 @@ class FusedMoELoadingMixin(abc.ABC): quant_method: UnquantizedFusedMoEMethod, moe_runner_config: MoeRunnerConfig, moe_tp_rank: int, + moe_tp_size: int, ) -> None: super().__init__() helper = FusedMoE.__new__(FusedMoE) @@ -97,6 +98,7 @@ class FusedMoELoadingMixin(abc.ABC): helper.moe_runner_config = moe_runner_config helper.use_triton_kernels = False helper.moe_tp_rank = moe_tp_rank + helper.moe_tp_size = moe_tp_size helper.use_presharded_weights = False helper.use_flashinfer_trtllm_moe = False # Keep this parameterless loading helper out of the module tree so diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py new file mode 100644 index 000000000..1d07b16a0 --- /dev/null +++ b/python/sglang/srt/models/kimi_k3.py @@ -0,0 +1,3203 @@ +# Kimi-K3 multimodal model: KimiLinear text backbone + MoonViT3d vision tower. +# Based on kimi_linear.py with K3-specific features: +# - Attention Residual (attn_res_block_size) +# - Latent MoE (routed_expert_hidden_size) +# - SiTU activation +# - MLA output gate (mla_use_output_gate) +# - Full-rank KDA gate (use_full_rank_gate) + +import logging +from collections.abc import Iterable +from functools import cached_property +from typing import TYPE_CHECKING, List, Optional, Tuple + +import torch +from torch import nn + +from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated +from sglang.srt.configs.kimi_k3 import KimiK3Config +from sglang.srt.configs.kimi_linear import KimiLinearConfig +from sglang.srt.distributed import ( + divide, + get_pp_group, + get_tp_group, + tensor_model_parallel_all_reduce, +) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers import ( + k3_ar_fusion, + k3_gemm_ar, + k3_sp_collective, + zero_copy_context, +) +from sglang.srt.layers.activation import SiluAndMul, SituAndMul +from sglang.srt.layers.attn_residual import AttnResidual, aggregate_stream, get_cw +from sglang.srt.layers.dcp.planner import prepare_decode_context_parallel_metadata +from sglang.srt.layers.dp_attention import ( + dp_gather_replicate, + dp_scatter, + get_global_dp_buffer, + get_local_dp_buffer, + is_allocation_symmetric, + is_dp_attention_enabled, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelBatchedLinear, + ColumnParallelLinear, + MergedColumnParallelLinear, + MergedColumnParallelRepeatedLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import route_quant_handoff +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import ( + TopK, + TopKOutputFormat, + build_precomputed_topk_output, + precomputed_topk_postprocess_is_noop, +) +from sglang.srt.layers.moe.utils import ( + RoutingMethodType, + get_moe_a2a_backend, + get_moe_runner_backend, +) +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, + get_embedding_tp_kwargs, +) +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, +) +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, + sharded_weight_loader, +) +from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, MoEGate +from sglang.srt.models.kimi_k3_vl import ( + KimiK3MultiModalProjector, + KimiK3VisionTower, +) +from sglang.srt.models.transformers import maybe_prefix +from sglang.srt.models.utils import WeightsMapper +from sglang.srt.multimodal.mm_utils import materialize_multimodal_features +from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args +from sglang.srt.utils import is_blackwell_supported, is_hip, make_layers +from sglang.srt.utils.common import ( + BumpAllocator, + add_prefix, + get_bool_env_var, + rank0_log, + require_mlp_sync, + set_weight_attrs, +) + +logger = logging.getLogger(__name__) +_is_hip = is_hip() +_aiter_k3_opt = get_bool_env_var("SGLANG_AITER_K3_OPT") + + +def _cdiv(a: int, b: int) -> int: + return (a + b - 1) // b + + +# MegaMoE SiTU sentinel: DeepGEMM 0.1.5.post1+ selects the K3 SiTU +# activation when activation_clamp == 0.03125 (2^-5: exactly representable and +# unused by any legitimate swiglu clamp; the host asserts clamp >= 0 so a +# negative sentinel is impossible). beta=4.0 / linear_beta=25.0 are baked into +# the DeepGEMM kernel. +_K3_MEGA_SITU_SENTINEL_CLAMP = 0.03125 + + +def _k3_bf16_gemm( + x: torch.Tensor, + weight: torch.Tensor, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """F.linear / torch.mm with the same TGV dispatch module-level GEMMs get + through UnquantizedLinearMethod. The fused MoE front and the deferred + shared down GEMM call torch directly on raw merged weights, so the + --bf16-gemm-backend cutedsl selection would silently skip them.""" + if x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16: + from sglang.srt.layers.quantization.unquant import get_bf16_gemm_backend + + if get_bf16_gemm_backend().is_cutedsl(): + from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import ( + cutedsl_bf16_gemm, + cutedsl_bf16_gemm_out, + use_cutedsl_bf16_gemm, + ) + + if use_cutedsl_bf16_gemm(x.shape[0], weight.shape[0], weight.shape[1]): + if out is None: + return cutedsl_bf16_gemm(x, weight) + if out.is_contiguous(): + # TGV stores straight into caller memory (same entry the + # UnquantizedLinearMethod out-buffer path uses); no + # staging tensor + copy. + return cutedsl_bf16_gemm_out(x, weight, out) + out.copy_(cutedsl_bf16_gemm(x, weight)) + return out + if out is None: + return torch.nn.functional.linear(x, weight) + return torch.mm(x, weight.t(), out=out) + + +# Fully fused KDA decode step (conv1d + delta rule + gated RMSNorm in one +# kernel, kernels/ops/attention/kda_fused_decode). The model hands the output-norm gate +# to the KDA backend via an attempt-and-verify stash on the attention layer; +# unconsumed stashes fall back to the unfused chain + o_norm here. + + +def _merge_weights_as_views( + mods: list, pad_rows_to: int = 1 +) -> tuple[torch.Tensor, list[int]]: + """Cat module weights along dim 0; re-point each module's weight to a view + of the merged buffer so the original storage is freed (net extra memory ~0). + + With pad_rows_to > 1 the merged buffer gets zero rows appended up to the + next multiple, so every row of the fused GEMM output stays 16-byte aligned + for vectorized consumers.""" + ws = [m.weight.data for m in mods] + sizes = [w.shape[0] for w in ws] + pad = (-sum(sizes)) % pad_rows_to + if pad: + ws = ws + [ws[0].new_zeros((pad, ws[0].shape[1]))] + merged = torch.cat(ws, dim=0).contiguous() + off = 0 + for m, n in zip(mods, sizes): + m.weight.data = merged[off : off + n] + off += n + return merged, sizes + + +# DP attention helpers. +# +# K3 cannot use LayerCommunicator: the attn-res aggregation kernels replace +# input_layernorm / post_attention_layernorm, which the communicator expects +# to own. Instead the MLP/MoE modules gather/scatter around their own body: +# attention and the attn-res buffers stay in local (per-DP-rank) token space, +# the MLP/MoE runs on the DP-gathered global batch with plain full-TP +# semantics (its internal all-reduces are unchanged and required — the latent +# reduce must happen in latent space before the norm), and the delayed +# prefix_sum add stays local, applied after the scatter back. + + +def _dp_local_buffer_group(): + """Symmetric-memory group for the local DP buffer (mirrors + CommunicateSummableTensorPairFn._scatter_hidden_states).""" + parallel = get_parallel() + if parallel.tp_size == parallel.attn_dp_size: + return get_tp_group() + return parallel.attn_tp_group + + +def _sp_all_gather_rows(hidden_states: torch.Tensor) -> torch.Tensor: + """Reassemble contiguous token shards, using the tuned K3 AG when covered.""" + group = get_parallel().attn_tp_group + hidden_states = hidden_states.contiguous() + full = k3_sp_collective.all_gather(hidden_states) + if full is None: + full = torch.empty( + (hidden_states.shape[0] * group.world_size, hidden_states.shape[1]), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + group.all_gather_into_tensor(full, hidden_states) + return full + + +def _sp_local_rows(hidden_states: torch.Tensor) -> slice: + """Full-batch row interval owned by this rank's contiguous token shard.""" + group = get_parallel().attn_tp_group + lo = group.rank_in_group * hidden_states.shape[0] + return slice(lo, lo + hidden_states.shape[0]) + + +class KimiK3MLP(nn.Module): + """K3 MLP; SiLU or SiTU activation.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + _tp_kwargs = ( + dict(tp_rank=tp_rank, tp_size=tp_size) if tp_size is not None else {} + ) + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + **_tp_kwargs, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + **_tp_kwargs, + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "situ": + self.act_fn = SituAndMul( + beta=activation_situ_beta or 1.0, + linear_beta=activation_situ_linear_beta, + ) + else: + raise ValueError(f"Unsupported activation: {hidden_act}") + self._dp_attention = is_dp_attention_enabled() + + def forward( + self, + hidden_states: torch.Tensor, + *, + prefix_sum: Optional[torch.Tensor] = None, + forward_batch: Optional[ForwardBatch] = None, + ) -> torch.Tensor: + # DP attention only when driven from the decoder layer (forward_batch + # given); the shared-experts instance inside KimiK3MoE passes None and + # runs on the already-gathered buffer. + use_dp = self._dp_attention and forward_batch is not None + if use_dp: + local_hidden_states = hidden_states + hidden_states = get_global_dp_buffer(get_tp_group()) + dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) + gate_up, _ = self.gate_up_proj(hidden_states) + hidden_states = self.act_fn(gate_up) + hidden_states, _ = self.down_proj(hidden_states) + if use_dp: + global_out = hidden_states + hidden_states = get_local_dp_buffer(_dp_local_buffer_group()) + dp_scatter(hidden_states, global_out, forward_batch) + # TODO(dark): maybe fuse residual with all reduce of down projection + if prefix_sum is not None: + hidden_states = hidden_states + prefix_sum + return hidden_states + + +def _add3( + a: torch.Tensor, + b: torch.Tensor, + c: Optional[torch.Tensor], + *, + prefetch_bc: bool = False, +) -> torch.Tensor: + """bf16(a + b) [+ c]. A pending c (the attn-res delayed +prefix_sum) + collapses the two elementwise adds into the 3-way JIT kernel — one + launch and one memory pass; its double rounding matches the unfused + pair bit-for-bit. prefetch_bc loads b/c before the PDL wait: only pass + True when their producers are at least two kernels back.""" + if c is None: + return a + b + from sglang.kernels.ops.elementwise import add3 + + return add3.add3(a, b, c, prefetch_bc=prefetch_bc) + + +# One-shot log guard: proves the merged front is live (see _ep_front). +_EP_FRONT_LOGGED = False + + +def _o_proj_takes_output(o_proj: RowParallelLinear) -> bool: + """Whether o_proj can write into caller-owned storage. ``apply_into`` is an + optional quant-method capability; only the unquantized method has it.""" + return getattr(o_proj.quant_method, "apply_into", None) is not None + + +def _k3_symm_o_proj_out(o_proj: RowParallelLinear, x: torch.Tensor) -> torch.Tensor: + """Symmetric storage for o_proj's TP-partial output; the fused attention + all-reduce reduces it in place.""" + return k3_ar_fusion.symm_buffer( + k3_ar_fusion.ATTN_O_PROJ, x.shape[0], o_proj.weight.shape[0], x.dtype + ) + + +class KimiK3MoE(nn.Module): + """K3 MoE with Latent MoE (experts run in moe_hidden_size space).""" + + def __init__( + self, + config: KimiLinearConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + layer_idx: int = 0, + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + hidden_size = config.hidden_size + moe_intermediate_size = config.moe_intermediate_size + moe_renormalize = config.moe_renormalize + self.tp_size = get_parallel().tp_size + self.routed_scaling_factor = config.routed_scaling_factor + self.num_shared_experts = config.num_shared_experts + self.layer_idx = layer_idx + self.alt_stream = alt_stream + self._dp_attention = is_dp_attention_enabled() + + self.use_latent_moe = config.routed_expert_hidden_size is not None + # Merged front weight ([H, gate_up + E + latent]), built after weight + # loading by _merge_front_weights(). + self._front_w: Optional[torch.Tensor] = None + self._front_sizes: Optional[List[int]] = None + # True when _front_w merges only [gate, routed_expert_down_proj] (the EP + # a2a pair) rather than the three-way fused-front weight. + self._front_is_ep_pair = False + self.moe_hidden_size = ( + config.routed_expert_hidden_size if self.use_latent_moe else hidden_size + ) + + # Gate — fp32 output so routing (sigmoid, bias add, top-k) runs in + # full precision (matches GateLinear in mke). codespell:ignore mke + self.gate = MoEGate(config, quant_config=None, prefix=f"{prefix}.gate") + + # For MXFP4 compressed-tensors, replace quant_config with Mxfp4Config + # so FusedMoE's weight_loader uses the MXFP4 fast path + moe_quant_config = quant_config + if quant_config is not None and getattr(quant_config, "quant_format", None): + if "mxfp4" in quant_config.quant_format: + from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config + + moe_quant_config = Mxfp4Config(is_checkpoint_mxfp4_serialized=True) + + # Routed experts (operate in moe_hidden_size space) + # gate_up_interleaved=False: K3 loads per-expert w1/w3 into non-interleaved layout + self.experts = get_moe_impl_class(moe_quant_config)( + num_experts=getattr(config, "n_routed_experts", config.num_experts), + top_k=config.num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + layer_id=self.layer_idx, + quant_config=moe_quant_config, + routed_scaling_factor=self.routed_scaling_factor, + activation=config.hidden_act, + gemm1_alpha=config.activation_situ_beta, + gemm1_clamp_limit=config.activation_situ_linear_beta, + gate_up_interleaved=False, + # trtllm fused-routing MoE backends (e.g. nvfp4 w4a4) route inside + # the kernel and require the routing method; K3 uses DSv3-style + # grouped topk with e_score_correction_bias. + routing_method_type=getattr( + config, "routing_method_type", RoutingMethodType.DeepSeekV3 + ), + prefix=add_prefix("experts", prefix), + ) + + self.topk = TopK( + top_k=config.num_experts_per_token, + renormalize=moe_renormalize, + use_grouped_topk=True, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + correction_bias=self.gate.e_score_correction_bias, + quant_config=quant_config, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk, + # flashinfer_mxfp4 + situ consumes precomputed routing + # (PackedPrecomputed): keep the radix router in the TopK layer + # and hand its ids/weights to the MoE op. Other quantized paths + # keep the runner-resolved format (marlin -> standard anyway, + # bypassed only for the public logits-routing path). + output_format=( + TopKOutputFormat.STANDARD + if quant_config is None + or ( + config.hidden_act == "situ" + and get_moe_runner_backend().is_flashinfer_mxfp4() + ) + # mega pre-dispatch consumes raw topk_ids/topk_weights + or get_moe_a2a_backend().is_megamoe() + else None + ), + ) + + # MegaMoE (deep_gemm fused a2a+GEMM over the EP symm buffer): a drop-in + # replacement for the routed experts call below. K3 routes ALL batches + # through it when enabled — the megamoe backend's non-mega fallback is + # a StandardDispatcher without a2a, which is wrong for scattered + # tokens — so SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK must + # cover the per-rank prefill chunk. SiTU is selected inside the + # DeepGEMM mega kernel via a sentinel activation_clamp with the K3 + # constants baked in. + self._use_mega_moe = get_moe_a2a_backend().is_megamoe() + self._mega_intermediate_size = moe_intermediate_size + self._mega_top_k = config.num_experts_per_token + if self._use_mega_moe: + assert self.use_latent_moe and config.hidden_act == "situ" + assert ( + config.activation_situ_beta, + config.activation_situ_linear_beta, + ) == (4.0, 25.0), ( + "mega SiTU kernel patch bakes beta=4.0/linear_beta=25.0; " + "got a checkpoint with different constants" + ) + + # EP a2a backends (megamoe / DeepEP) move each row to its experts + # directly, so the MoE region can consume whatever rows this rank + # holds — an SP-MoE token shard (attn_tp > 1) or the DP-local batch + # (DP attention) — with every global token dispatched exactly once. + # No DP gather and no TP reduce is needed anywhere in the region. + _a2a_backend = get_moe_a2a_backend() + self._ep_a2a = _a2a_backend.is_megamoe() or _a2a_backend.is_deepep() + + # The flashinfer_mxfp4 (trtllm-gen) runner quantizes routed_input with + # the strided-input JIT group quant (_use_jit_mxfp8_quant in mxfp4.py), + # so the fused-front split view can be consumed as is; other runners + # (e.g. marlin) require a dense buffer. + self._moe_front_needs_contiguous = ( + not get_moe_runner_backend().is_flashinfer_mxfp4() + ) + + # Defer the trtllm-gen finalize (top-k weighted unpermute) out of the + # MoE op and fuse it into the push all-reduce's staging pass + # (k3_ar_fusion.finalize_all_reduce_push_norm): the rank-local latent + # never materializes. Only the situ packed-routing trtllm-gen path + # serves the deferral; sizes beyond the push window fall back to the + # in-op finalize at runtime (finalize_push_fits). + self._defer_moe_finalize = ( + get_moe_runner_backend().is_flashinfer_mxfp4() + and config.hidden_act == "situ" + ) + + # Shared experts (operate in original hidden_size space). + # Replicate the shared-expert weights (tp1, DSv2 convention) under EP + # a2a: the block runs on partial batches (shard / DP-local rows), and + # a TP-sharded partial sum could never be reduced across ranks that + # hold different tokens. + self._shared_experts_tp1 = self._ep_a2a + if self.num_shared_experts is not None and self.num_shared_experts > 0: + shared_intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiK3MLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + **(dict(tp_rank=0, tp_size=1) if self._shared_experts_tp1 else {}), + ) + else: + self.shared_experts = None + + # SBO (single batch overlap): the shared experts read a fixed slab of + # weights the routed path never touches (bf16 — the checkpoint leaves + # shared_experts unquantized — and tp1-replicated under EP a2a, so + # ~264 MB per layer per rank), while the routed path is a2a-latency + # bound in decode with HBM mostly idle. Issue the shared experts on the + # side stream so the two run concurrently instead of back to back; the + # join happens right before the tail add. Measured on 2x4 GB300 + # (TP8/EP8 MegaMoE + SP-MoE): +4~5% output tok/s and −5% ITL over + # bs 1–32, GSM8K unchanged — so it is on whenever the shape allows, + # no flag. + # EP a2a only: with plain-TP experts the fused front already lands both + # partial sums in one collective (_forward_fused), a strictly better + # overlap than two streams. + self._sbo_shared_overlap = ( + self._ep_a2a + and self.shared_experts is not None + and self.alt_stream is not None + ) + + if self.use_latent_moe: + self.routed_expert_down_proj = ReplicatedLinear( + hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_down_proj", + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) + if config.latent_moe_use_norm + else None + ) + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + else: + self.routed_expert_down_proj = None + self.routed_expert_norm = None + self.routed_expert_up_proj = None + + # Static eligibility for fusing the fused-front latent all-reduce with + # the RMSNorm epilogue (SGLANG_K3_AR_FUSION). The kernel views the flat + # [latent | shared] buffer as [3N, NORM_DIM] rows and norms the first N, + # so it requires latent width == NORM_DIM and shared width == 2*NORM_DIM + # (K3: 3584 / 7168). Decided once here so the hot path only reads a bool + # and never re-validates dims per forward. + self.fuse_ar_norm = ( + self.routed_expert_norm is not None + and self.moe_hidden_size == k3_ar_fusion.NORM_DIM + and hidden_size == 2 * k3_ar_fusion.NORM_DIM + ) + # Static eligibility for the column-parallel up_proj tail (gemm_ag): + # per-rank 1/8-column GEMV -> multicast all-gather staged in the v2 + # push workspace -> spin-add3 with shared_output (+ prefix_sum), + # replacing the replicated [3584, 7168] GEMM + _add3 (~1.5-2x at + # decode sizes, 1/8 of the weight bytes read per rank). Kernel dims + # are fixed to fuse_ar_norm's (3584 -> 7168) over TP8; per-batch + # capacity checks live in k3_ar_fusion.gemm_ag_up_fits. + self._gemm_ag_up_eligible = ( + self.fuse_ar_norm + and self.tp_size == 8 + and self.routed_expert_up_proj is not None + and isinstance(self.routed_expert_up_proj.weight, torch.Tensor) + and self.routed_expert_up_proj.weight.dtype == torch.bfloat16 + and self.routed_expert_up_proj.weight.is_contiguous() + ) + + def _merge_front_weights(self) -> None: + """Merge shared gate_up + router gate + latent down_proj weights. + + All three GEMMs consume the same hidden_states; at decode each one is a + skinny memory-bound GEMV with its own splitK epilogue. One merged + [H, gu+E+latent] GEMM reads the input once and drops 2 GEMM launches + plus their splitK-reduce tails per MoE layer. + + Called once from load_weights (after all weights are loaded, before + cuda graph capture); only plain bf16/fp16 dense weights are merged — + quantized or mixed-dtype checkpoints keep the unfused path. + """ + if not self.use_latent_moe: + return + if self.shared_experts is not None and get_moe_a2a_backend().is_none(): + mods = [ + self.shared_experts.gate_up_proj, + self.gate, + self.routed_expert_down_proj, + ] + elif envs.SGLANG_K3_FUSED_FRONT.get(): + # EP a2a: the shared experts are tp1-replicated and run on the side + # stream, so they stay out of the merge -- but the router gate and the + # latent down-proj still read the same hidden_states, and merging just + # those two is what lets one GEMM read the activations once. The gate + # GEMM alone is only 896 rows, which is too few to use the machine + # well; folded into the 3584-row down-proj it comes almost free. + mods = [self.gate, self.routed_expert_down_proj] + else: + return + dtypes = {m.weight.dtype for m in mods} + if len(dtypes) != 1 or dtypes.pop() not in (torch.bfloat16, torch.float16): + return + self._front_w, self._front_sizes = _merge_weights_as_views(mods) + self._front_is_ep_pair = len(mods) == 2 + # Invalidate the cached properties. + for prop in ( + "_eligible_for_fused_front", + "_routing_contract_ok", + "_ep_front_eligible", + ): + self.__dict__.pop(prop, None) + + @cached_property + def _routed_needs_reduce(self): + return self.tp_size > 1 and get_moe_a2a_backend().is_none() + + @cached_property + def _eligible_for_fused_front(self) -> bool: + """The fused front commits to the single-collective tail (both + partial sums in one symmetric buffer), so beyond the merged front + weight it requires plain-TP routed sums (an a2a combine already + returns the complete sum — all-reducing it again would multiply by + tp_size) and a dense shared down weight for the direct out= GEMM.""" + return ( + self.use_latent_moe + and self.shared_experts is not None + and self._front_w is not None + and not self._front_is_ep_pair + and get_moe_a2a_backend().is_none() + and self.shared_experts.down_proj.weight.dtype + in (torch.bfloat16, torch.float16) + ) + + def _forward_mega_experts( + self, routed_input: torch.Tensor, topk_output + ) -> torch.Tensor: + """Routed experts via deep_gemm MegaMoE: fused a2a dispatch + grouped + GEMMs + SiTU + combine over the EP-group symmetric buffer. Semantically + equivalent to `self.experts(routed_input, topk_output)` on an a2a + backend (combine returns fully-summed rows; `_reduce_latent` then only + applies the norm).""" + import deep_gemm + + from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch + from sglang.srt.distributed.parallel_state import get_moe_ep_group + from sglang.srt.environ import envs + from sglang.srt.layers.moe.mega_moe import _get_mega_moe_symm_buffer + + # In SP-MoE mode (KimiK3DecoderLayer reduce-scatters the o_proj + # output) the incoming rows are already this rank's token shard, so + # the fused a2a below dispatches each token exactly once. On the + # non-scattered fallback path the rows are the full batch (redundant + # across ranks but correct). + num_tokens = routed_input.shape[0] + num_max_tokens_per_rank = ( + envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get() + ) + assert num_tokens <= num_max_tokens_per_rank, ( + f"mega MoE: num_tokens={num_tokens} exceeds " + f"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=" + f"{num_max_tokens_per_rank}; K3 has no non-mega fallback — raise " + f"the env var to cover the per-rank rows" + ) + buf = _get_mega_moe_symm_buffer( + get_moe_ep_group().device_group, + num_experts=self.experts.num_experts, + num_max_tokens_per_rank=num_max_tokens_per_rank, + num_topk=self._mega_top_k, + hidden=self.moe_hidden_size, + intermediate_hidden=self._mega_intermediate_size, + ) + + if num_tokens > 0: + topk_ids_in = topk_output.topk_ids.to(torch.int32) + topk_weights_in = topk_output.topk_weights.to(torch.float32) + else: + topk_ids_in = routed_input.new_empty( + (0, self._mega_top_k), dtype=torch.int32 + ) + topk_weights_in = routed_input.new_empty( + (0, self._mega_top_k), dtype=torch.float32 + ) + + mega_moe_pre_dispatch( + routed_input, + topk_ids_in, + topk_weights_in, + buf.x, + buf.x_sf, + buf.topk_idx, + buf.topk_weights, + quant_group_size=32, + ) + # At least one row so the tvm-ffi binding sees a non-null data_ptr. + y = torch.empty( + (max(num_tokens, 1), self.moe_hidden_size), + dtype=torch.bfloat16, + device=routed_input.device, + ) + deep_gemm.fp8_fp4_mega_moe( + y, + self.experts.mega_l1_weights, + self.experts.mega_l2_weights, + buf, + recipe=(1, 1, 32), + activation="swiglu", + # Sentinel: selects the K3 SiTU branch in the DeepGEMM mega kernel + # (beta=4.0 / linear_beta=25.0 baked in). + activation_clamp=_K3_MEGA_SITU_SENTINEL_CLAMP, + fast_math=True, + ) + y = y[:num_tokens] + if not self.experts.should_fuse_routed_scaling_factor_in_topk: + if ( + self.routed_scaling_factor is not None + and self.routed_scaling_factor != 1.0 + ): + y.mul_(self.routed_scaling_factor) + return y + + def _latent_norm(self, latent: torch.Tensor) -> torch.Tensor: + if self.routed_expert_norm is None: + return latent + return self.routed_expert_norm(latent) + + @cached_property + def _routing_contract_ok(self) -> bool: + """Whether a kernel may emit (weights, ids) itself and bypass + select_experts. Shared by the fused router and the merged front.""" + if self._eligible_for_fused_front: + return False + cfg = self.topk.topk_config + if cfg.output_format is not TopKOutputFormat.STANDARD: + return False + # The kernel implements sigmoid scoring with bias-ranked ungrouped top-k. + # Do NOT test cfg.scoring_func: it defaults to "softmax" and TopK + # documents it as unused. What actually selects sigmoid is the + # grouped-topk-with-correction-bias route (DSv3 noaux_tc), which calls + # biased_grouped_topk and hardwires scoring_func="sigmoid". + if not (cfg.use_grouped_topk and cfg.correction_bias is not None): + return False + if (cfg.num_expert_group or 1) > 1 or (cfg.topk_group or 1) > 1: + return False + # A waterfill balancer rewrites the routing after the top-k; leave it on + # the layer path that supports it. + if self.topk.waterfill_balancer is not None or self.topk.enable_waterfill: + return False + if self.gate.e_score_correction_bias is None: + return False + # K3 calls self.topk() without a padding mask or EPLB dispatch info, so + # select_experts' post-processing collapses to the capture hook and the + # recorder -- both of which build_precomputed_topk_output runs. Bail out + # if that ever stops holding rather than silently dropping the remap. + if not precomputed_topk_postprocess_is_noop(cfg): + return False + if get_exec().deterministic.enable_deterministic_inference: + return False + try: + from sglang.kernels.ops.moe import moe_front + except Exception: + return False + return moe_front.available() + + @cached_property + def _ep_front_eligible(self) -> bool: + """Static eligibility for the merged EP front (gate + latent down-proj in + one GEMM). Requires the two-module merge from _merge_front_weights and the + same routing contract the single-kernel router needs.""" + return ( + envs.SGLANG_K3_FUSED_FRONT.get() + and self._front_w is not None + and self._front_is_ep_pair + and self.use_latent_moe + and self.routed_expert_down_proj is not None + and self._routing_contract_ok + ) + + def _ep_front(self, hidden_states: torch.Tensor): + """Merged front: returns ``(topk_output, routed_input)``, or None when the + shape is not covered and the caller should run the unmerged path.""" + if not self._ep_front_eligible: + return None + from sglang.kernels.ops.moe import moe_front + + cfg = self.topk.topk_config + bias = self.gate.e_score_correction_bias + if ( + moe_front.get_front_strategy(hidden_states.shape[0], hidden_states.device) + != "merged_fp32" + ): + return None + if not moe_front.fused_front_covered( + hidden_states, self._front_w, bias, cfg.top_k, self.moe_hidden_size + ): + return None + + w, i, routed = moe_front.fused_front( + hidden_states, + self._front_w, + bias, + latent=self.moe_hidden_size, + topk=cfg.top_k, + renormalize=cfg.renormalize, + routed_scaling_factor=cfg.routed_scaling_factor, + apply_routed_scaling_factor_on_output=cfg.apply_routed_scaling_factor_on_output, + ) + + global _EP_FRONT_LOGGED + if not _EP_FRONT_LOGGED: + # An absence of fallback warnings does not prove a fast path ran. + _EP_FRONT_LOGGED = True + logger.info( + "K3 merged MoE front active (layer %d, %d tokens)", + self.layer_idx, + hidden_states.shape[0], + ) + return build_precomputed_topk_output(w, i, cfg, self.layer_idx), routed + + def _ep_front_overlap(self, hidden_states: torch.Tensor): + """Overlap the exact fp32 gate+top-k with the latent down projection. + + The side stream is joined before returning. It is then free for the + existing shared-expert overlap, which is deliberately issued later. + """ + if ( + not self._ep_front_eligible + or self.alt_stream is None + or hidden_states.shape[0] == 0 + ): + return None + from sglang.kernels.ops.moe import moe_front + + if ( + moe_front.get_front_strategy(hidden_states.shape[0], hidden_states.device) + != "overlap" + ): + return None + + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + with torch.cuda.stream(self.alt_stream): + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + + routed_input, _ = self.routed_expert_down_proj(hidden_states) + current_stream.wait_stream(self.alt_stream) + # Top-k tensors were allocated on alt_stream but are consumed by the + # routed experts on current_stream. Tell the caching allocator about + # that lifetime before alt_stream is reused for the shared experts. + for value in topk_output: + if isinstance(value, torch.Tensor): + value.record_stream(current_stream) + return topk_output, routed_input + + def _reduce_latent(self, latent: torch.Tensor) -> torch.Tensor: + """Unfused-front latent tail: TP-partial routed sums must be reduced + in latent space BEFORE the RMSNorm (sum(norm(x_i)) != norm(sum(x_i))).""" + if not self._routed_needs_reduce: + return self._latent_norm(latent) + return self._latent_norm(tensor_model_parallel_all_reduce(latent)) + + def _forward_unfused( + self, hidden_states: torch.Tensor, *, prefix_sum: Optional[torch.Tensor] + ) -> torch.Tensor: + """Front section with three separate GEMMs, each reading + hidden_states: shared-expert MLP, router gate, latent down-proj.""" + # Shared experts on original hidden_states. Under SBO they go to the + # side stream and are joined at the tail (see _sbo_shared_overlap). + # + # Issued *after* the front, deliberately: alt_stream.wait_stream() makes + # the side stream wait for whatever the main stream has enqueued so far, + # so issuing here means the shared experts overlap the routed a2a rather + # than the front GEMMs. The shared branch is the shorter of the two and + # does not need a head start; running it against the front only takes + # bandwidth away from the critical path. + shared_output = None + shared_event = None + + def issue_shared(): + nonlocal shared_output, shared_event + if self.shared_experts is None or hidden_states.shape[0] == 0: + return + if self._sbo_shared_overlap: + self.alt_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.alt_stream): + shared_output = self.shared_experts(hidden_states) + shared_event = self.alt_stream.record_event() + else: + shared_output = self.shared_experts(hidden_states) + + # Front: gate + TopK (+ latent down-proj when the merged front covers it). + # The gate and the latent down-proj read the same hidden_states, so the + # merged-weight strategies compute both in one GEMM; see + # kernels/ops/moe/moe_front.py for the strategy table. + routed_input = self._ep_front(hidden_states) + if routed_input is None: + routed_input = self._ep_front_overlap(hidden_states) + topk_output = None + if routed_input is not None: + topk_output, routed_input = routed_input + else: + # MoEGate produces fp32 router logits on CUDA (via linear_bf16_fp32 + # or dsv3_router_gemm); non-CUDA falls back to F.linear (bf16). The + # fp32 logits reach the radix router from moe_fused_gate. + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + + issue_shared() + + if not self.use_latent_moe: + expert_output = self.experts(hidden_states, topk_output) + if shared_event is not None: + torch.cuda.current_stream().wait_event(shared_event) + if shared_output is not None: + expert_output = expert_output + shared_output + if self.tp_size > 1: + expert_output = tensor_model_parallel_all_reduce(expert_output) + if prefix_sum is not None: + expert_output = expert_output + prefix_sum + return expert_output + + # Latent MoE: compress after routing, before experts + if TYPE_CHECKING: + assert ( + self.routed_expert_down_proj is not None + and self.routed_expert_up_proj is not None + ) + + if routed_input is None: + routed_input, _ = self.routed_expert_down_proj(hidden_states) + expert_output = ( + self._forward_mega_experts(routed_input, topk_output) + if self._use_mega_moe + else self.experts(routed_input, topk_output) + ) + latent = self._reduce_latent(expert_output) + # up_proj is replicated, so the routed output is now fully reduced. + out, _ = self.routed_expert_up_proj(latent) + if shared_event is not None: + # SBO join: as late as possible, so the side-stream shared experts + # get the whole routed a2a + latent tail to hide under. + torch.cuda.current_stream().wait_event(shared_event) + if shared_output is not None: + # tp1 shared experts (SP-MoE) are complete per-rank; TP-sharded + # ones need the partial-sum reduction. + if self.tp_size > 1 and not self._shared_experts_tp1: + shared_output = tensor_model_parallel_all_reduce(shared_output) + return _add3(out, shared_output, prefix_sum) + return out if prefix_sum is None else out + prefix_sum + + @cached_property + def _route_quant_fuse_eligible(self) -> bool: + """Whether to stage routed_input for the fused route+pack+quant launch + (route_quant_handoff). Only the trtllm-gen SiTU runner with mxfp8 + activations consumes the staged quant, so only that runner stages.""" + from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod + + method = self.experts.quant_method + return ( + isinstance(method, Mxfp4MoEMethod) + and method.use_flashinfer + and not method.use_marlin + and method.flashinfer_mxfp4_moe_precision == "default" + and self.experts.moe_runner_config.activation == "situ" + ) + + def _forward_routed(self, hidden_states, router_logits, routed_input, latent): + if self._route_quant_fuse_eligible: + route_quant_handoff.stage(routed_input) + try: + topk_output = self.topk(hidden_states, router_logits) + with zero_copy_context.set_moe_output(latent): + expert_output = self.experts(routed_input, topk_output) + finally: + route_quant_handoff.clear() + if expert_output.data_ptr() != latent.data_ptr(): + latent.copy_(expert_output) + + def _forward_routed_deferred(self, hidden_states, router_logits, routed_input): + """Routed experts with the in-op finalize skipped: returns the + FlashInferTrtllmDeferredFinalizeOutput triple (permuted gemm2 output, + expanded_idx_to_permuted_idx, expert_weights) for the finalize-fused + all-reduce.""" + if self._route_quant_fuse_eligible: + route_quant_handoff.stage(routed_input) + try: + topk_output = self.topk(hidden_states, router_logits) + return self.experts.forward_deferred_finalize(routed_input, topk_output) + finally: + route_quant_handoff.clear() + + def _forward_shared(self, gate_up, shared_output): + shared = self.shared_experts + if TYPE_CHECKING: + assert shared is not None and isinstance( + shared.down_proj.weight, torch.Tensor + ) + assert shared is not None + _k3_bf16_gemm( + shared.act_fn(gate_up), + shared.down_proj.weight, + out=shared_output, + ) + + def _get_fused_norm_params(self) -> tuple[torch.Tensor, float]: + norm = self.routed_expert_norm + assert self.fuse_ar_norm and norm is not None + return norm.weight, norm.variance_epsilon + + def _forward_fused( + self, hidden_states: torch.Tensor, *, prefix_sum: Optional[torch.Tensor] + ) -> torch.Tensor: + """Fused-front pipeline: read hidden_states once through the merged + [H, gate_up + E + latent] weight, then land both TP-partial sums in + one flat symmetric [latent | shared] buffer with zero copies — the + shared down GEMM writes its slice via out=, the MoE runner writes + its top-k sum via the zero-copy context — and all-reduce the pair + in a single collective (the symmetric mempool keeps the one-shot + allreduce path; same trick as RowParallelLinear).""" + if TYPE_CHECKING: # NOTE: precondition for this case + assert ( + self._front_w is not None + and self._front_sizes is not None + and self.moe_hidden_size is not None + and self.shared_experts is not None + and isinstance(self.shared_experts.down_proj.weight, torch.Tensor) + and self.routed_expert_up_proj is not None + ) + + num_tokens, hidden_size = hidden_states.shape + fused = _k3_bf16_gemm(hidden_states, self._front_w) + gate_up, router_logits, routed_input = torch.split( + fused, self._front_sizes, dim=-1 + ) + if num_tokens > 1 and _is_hip and not _aiter_k3_opt: + router_logits = router_logits.contiguous() + if num_tokens > 1 and self._moe_front_needs_contiguous: + routed_input = routed_input.contiguous() + latent_numel = num_tokens * self.moe_hidden_size + if k3_ar_fusion.enabled(): + # the shared-expert AR is pull-only, so its input must be a + # symm_buffer slice for every rank to resolve the same offset + buf = k3_ar_fusion.symm_buffer( + k3_ar_fusion.MOE_LATENT_SHARED, + num_tokens, + self.moe_hidden_size + hidden_size, + hidden_states.dtype, + ).view(-1) + else: + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + buf = hidden_states.new_empty(latent_numel + num_tokens * hidden_size) + + latent = buf[:latent_numel].view(num_tokens, self.moe_hidden_size) + shared_output = buf[latent_numel:].view(num_tokens, hidden_size) + fused_norm = False + if self.alt_stream is not None and k3_ar_fusion.enabled(): + defer_finalize = ( + self._defer_moe_finalize + and self.fuse_ar_norm + and k3_ar_fusion.finalize_push_fits(num_tokens) + ) + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + if defer_finalize: + deferred = self._forward_routed_deferred( + hidden_states, router_logits, routed_input + ) + else: + self._forward_routed(hidden_states, router_logits, routed_input, latent) + with torch.cuda.stream(self.alt_stream): + self._forward_shared(gate_up, shared_output) + # low-SM pull so the side-stream AR leaves the SMs to the + # routed GEMMs it overlaps (K3 dims are fixed; tuned here) + k3_ar_fusion.all_reduce_low_sm(shared_output, num_blocks=4, unroll=8) + current_stream.wait_stream(self.alt_stream) + # NOTE: the latent AR must stay serialized after the shared AR + # (both reuse the v2 pull semaphores; concurrent calls would + # corrupt each other's barrier windows) — the join above does it. + if defer_finalize: + # finalize folded into the push AR's staging pass; the norm + # covers every latent row + fused_norm = True + k3_ar_fusion.finalize_all_reduce_push_norm( + latent, + deferred.gemm2_out, + deferred.expanded_idx_to_permuted_idx, + deferred.expert_weights, + *self._get_fused_norm_params(), + ) + elif self.fuse_ar_norm: + fused_norm = True + k3_ar_fusion.all_reduce_norm( + latent.view(-1, self.moe_hidden_size), + *self._get_fused_norm_params(), + num_tokens=num_tokens, + ) + else: + k3_ar_fusion.all_reduce(latent) + # the gemm_ag tail wants the normed latent straight out of the + # fused-norm AR (its GEMV chains on it via PDL) + if ( + fused_norm + and self._gemm_ag_up_eligible + and k3_ar_fusion.gemm_ag_up_fits(num_tokens) + ): + return k3_ar_fusion.gemm_ag_up_proj( + latent, + self.routed_expert_up_proj.weight, # type: ignore + shared_output, + prefix_sum, + ) + else: # single collective over the flat [latent | shared] pair + self._forward_shared(gate_up, shared_output) + self._forward_routed(hidden_states, router_logits, routed_input, latent) + if self.fuse_ar_norm and k3_ar_fusion.enabled(): + fused_norm = True + k3_ar_fusion.all_reduce_norm( + buf.view(-1, k3_ar_fusion.NORM_DIM), + *self._get_fused_norm_params(), + num_tokens=num_tokens, + ) + elif k3_ar_fusion.enabled(): + k3_ar_fusion.all_reduce(buf) + else: + buf = tensor_model_parallel_all_reduce(buf) + + latent = buf[:latent_numel].view(num_tokens, self.moe_hidden_size) + shared_output = buf[latent_numel:].view(num_tokens, hidden_size) + if not fused_norm: + latent = self._latent_norm(latent) + out, _ = self.routed_expert_up_proj(latent) + + # prefetch_bc: b (shared_output) was produced by the all-reduce and + # c (prefix_sum) even earlier; the AR is a plain launch (full + # barrier), so both are complete once the norm / up_proj GEMM chain + # starts — only `a`'s producer can still be in flight at PDL entry. + return _add3(out, shared_output, prefix_sum, prefetch_bc=True) + + def forward( + self, + hidden_states: torch.Tensor, + *, + prefix_sum: Optional[torch.Tensor] = None, + forward_batch: Optional[ForwardBatch] = None, + ) -> torch.Tensor: + """A pending prefix_sum is always consumed here: folded into the + 3-way JIT tail add when covered, plain adds otherwise (bit-identical + either way). + + Under DP attention with TP-sharded experts (a2a none, forward_batch + given) the experts run on the DP-gathered global batch — the internal + reduces stay over the full TP group, which is exactly right in + gathered space — while prefix_sum stays in local token space, added + after the scatter back. EP a2a backends skip the gather entirely: + dispatching the DP-local rows (or the SP-MoE shard of them the + decoder layer already produced) covers every global token exactly + once, and prefix_sum is consumed in the tail add like the non-DP + path — gathering first would just replicate the whole batch onto + every rank (tp-fold redundant compute + a2a traffic).""" + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + use_dp = self._dp_attention and forward_batch is not None and not self._ep_a2a + if use_dp: + local_hidden_states = hidden_states + hidden_states = get_global_dp_buffer(get_tp_group()) + dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) + dp_prefix_sum, prefix_sum = prefix_sum, None + if hidden_states.shape[0] > 0 and self._eligible_for_fused_front: + out = self._forward_fused(hidden_states, prefix_sum=prefix_sum) + else: + out = self._forward_unfused(hidden_states, prefix_sum=prefix_sum) + if use_dp: + global_out = out + out = get_local_dp_buffer(_dp_local_buffer_group()) + dp_scatter(out, global_out, forward_batch) + if dp_prefix_sum is not None: + out = out + dp_prefix_sum + return out.view(num_tokens, hidden_size) + + +class KimiK3DeltaAttention(nn.Module): + """KDA attention; optional full-rank gate.""" + + def __init__( + self, + layer_idx: int, + hidden_size: int, + config: KimiLinearConfig, + quant_config: Optional[QuantizationConfig] = None, + rms_norm_eps: float = 1e-5, + prefix: str = "", + all_reduce_fusion: bool = False, + bfa_alt_stream: Optional[torch.cuda.Stream] = None, + **kwargs, + ) -> None: + super().__init__() + self.all_reduce_fusion = all_reduce_fusion + # Side stream for the [f_a|b] + f_b tiny GEMVs: they read only + # hidden_states, so they can run concurrently with the wide fused + # [q,k,v,g] GEMM on the main stream (graphed decode/verify only). + # Same SM bound rationale as the MLA gate stream. + self._bfa_alt_stream = bfa_alt_stream + self._bfa_bs_limit = ( + (128 if is_blackwell_supported() else 64) + if bfa_alt_stream is not None + else 0 + ) + self.tp_size = get_parallel().tp_size + # KDA is an attention layer: all head-sharded params must follow the + # attention-TP group (= tp under plain TP, = 1 under DP attention), + # matching the mamba state cache sizing (KimiLinearCacheParams uses + # get_attention_tp_size). Mirrors GLM5-next's head_shard_size pattern. + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank + self.hidden_size = hidden_size + self.config = config + self.head_dim = config.linear_attn_config["head_dim"] + self.num_heads = config.linear_attn_config["num_heads"] + self.num_k_heads = config.linear_attn_config["num_heads"] + self.num_v_heads = config.linear_attn_config["num_heads"] + self.head_k_dim = config.linear_attn_config["head_dim"] + self.head_v_dim = config.v_head_dim + self.layer_idx = layer_idx + self.prefix = prefix + assert self.num_heads % self.attn_tp_size == 0 + self.local_num_heads = divide(self.num_heads, self.attn_tp_size) + + projection_size = self.head_dim * self.num_heads + self.conv_size = config.linear_attn_config["short_conv_kernel_size"] + self.use_full_rank_gate = config.linear_attn_config.get( + "use_full_rank_gate", False + ) + + # The fused path hardcodes tp_size sharding, so require attn_tp == tp. + # For the full-rank gate (K3) the checkpoint quantizes only the MoE + # experts; attention linears resolve to UnquantizedLinearMethod, so a + # non-None quant_config is fine for the merged projection. + self.do_fuse_qkvbfg = self.attn_tp_size == self.tp_size and ( + quant_config is None or self.use_full_rank_gate + ) + + if self.do_fuse_qkvbfg and self.use_full_rank_gate: + # Fuse only the alignment-friendly wide projections [q, k, v, g] + # (6144/rank at TP8). Folding b (12/rank) and f_a (128, replicated) + # in as well skews the output dim to 6284 and measurably degrades + # the GEMM kernel selection; they stay as separate tiny GEMVs. + self.fused_qkvg_proj = MergedColumnParallelLinear( + self.hidden_size, + [ + projection_size, + projection_size, + projection_size, + projection_size, + ], + bias=False, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.fused_qkvg_proj", + ) + self.split_sizes = [ + 3 * projection_size // self.tp_size, + projection_size // self.tp_size, + ] + self.b_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads, + bias=False, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.b_proj", + ) + self.f_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.f_a_proj", + ) + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.f_b_proj", + ) + # Merged [f_a | b] weight, built after weight loading by + # _merge_bfa_weights(). + self._bfa_w: Optional[torch.Tensor] = None + elif self.do_fuse_qkvbfg: + self.qkvb_sizes = [ + projection_size, + projection_size, + projection_size, + self.num_heads, + ] + self.fg_sizes = [self.head_dim, self.head_dim] + + self.fused_qkvbfg_a_proj = MergedColumnParallelRepeatedLinear( + self.hidden_size, + self.qkvb_sizes, + self.fg_sizes, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkvbfg_a_proj", + ) + self.split_sizes = [ + 3 * projection_size // self.tp_size, + self.num_heads // self.tp_size, + 2 * self.head_dim, + ] + _dtype = config.dtype + if isinstance(_dtype, str): + _dtype = getattr(torch, _dtype, torch.bfloat16) + self.fused_fg_b_proj = ColumnParallelBatchedLinear( + 2, self.head_dim, projection_size, dtype=_dtype + ) + else: + attn_tp_rank = self.attn_tp_rank + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.num_heads, + self.num_k_heads, + bias=False, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=self.attn_tp_size, + v_head_size=self.head_v_dim, + prefix=f"{prefix}.qkv_proj", + ) + + self.f_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.f_a_proj", + ) + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.f_b_proj", + ) + self.b_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads, + bias=False, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.b_proj", + ) + + if self.use_full_rank_gate: + self.g_proj = ColumnParallelLinear( + self.hidden_size, + projection_size, + bias=False, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.g_proj", + ) + else: + self.g_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_a_proj", + ) + self.g_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.g_b_proj", + ) + + self.dt_bias = nn.Parameter( + torch.empty(divide(projection_size, self.attn_tp_size), dtype=torch.float32) + ) + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + self.qkv_conv1d = MergedColumnParallelLinear( + input_size=self.conv_size, + output_sizes=[projection_size, projection_size, projection_size], + bias=False, + params_dtype=torch.float32, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=f"{prefix}.qkv_conv1d", + ) + self.qkv_conv1d.weight.data = self.qkv_conv1d.weight.data.unsqueeze(1) + + # K3 checkpoint stores A_log as [head_dim] (128), but the FLA kernel + # expects exactly local_num_heads elements. We define the param as + # [1, 1, local_num_heads, 1] (matching the kimi_linear.py convention) + # and attach a custom weight_loader that handles both the old 4-D + # format and the K3 1-D [head_dim] format by narrowing to the first + # num_heads elements then TP-sharding. + self.A_log = nn.Parameter( + torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) + ) + + def _a_log_weight_loader( + param: torch.Tensor, loaded_weight: torch.Tensor + ) -> None: + tp_rank = get_parallel().attn_tp_rank + shard_size = param.data.shape[2] # local_num_heads + start_idx = tp_rank * shard_size + + # Handle old 4-D checkpoint format: [1, 1, H, 1] -> [H] + if loaded_weight.dim() == 4: + loaded_weight = loaded_weight.view(loaded_weight.shape[2]) + # Now loaded_weight is 1-D (either [num_heads] or [head_dim]). + # Narrow to the TP shard along the head dimension. + loaded_weight = loaded_weight.narrow(0, start_idx, shard_size) + # Reshape to match param shape [1, 1, local_num_heads, 1] + param.data.copy_(loaded_weight.view(param.data.shape)) + + set_weight_attrs(self.A_log, {"weight_loader": _a_log_weight_loader}) + + self.o_norm = FusedRMSNormGated( + self.head_dim, eps=rms_norm_eps, activation="sigmoid" + ) + self.o_proj = RowParallelLinear( + projection_size, + self.hidden_size, + bias=False, + # SGLANG_K3_AR_FUSION: keep the o_proj output TP-partial and + # complete the reduce at the decoder layer via the fused MNNVL + # all-reduce (which can fold the attn-res prefix add in). Only + # valid when the attn TP group is the full TP group (the fused + # comm lives there). + reduce_results=not self.all_reduce_fusion, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + # Reduce within the attn-TP group: the default reduce path uses + # the full-TP collective, which at attn_tp>1 is both the wrong + # group (sums across DP groups) and asymmetric vs idle DP ranks + # (deadlocks the per-layer DP gather). Off under all_reduce_fusion: + # the fused AR does the reduce itself (reduce_results=False) and the + # forward hands o_proj a slice of a persistent symmetric region — + # leaving this True would wrap the GEMM in + # use_symmetric_memory(attn_tp), which allocates its own output and + # so defeats the caller-owned buffer. At the fusion config + # attn_tp==tp so the fused full-TP reduce is the same group anyway. + use_dp_attention_reduce=not self.all_reduce_fusion, + prefix=f"{prefix}.o_proj", + ) + if self.all_reduce_fusion and not _o_proj_takes_output(self.o_proj): + # the fused AR reduces o_proj's output in place out of a symmetric + # buffer, which needs the GEMM to write into caller-owned storage + self.all_reduce_fusion = False + self.o_proj.reduce_results = True + self.o_proj.use_dp_attention_reduce = True + k3_gemm_ar.maybe_wrap_o_proj(self.o_proj) + conv_weights = self.qkv_conv1d.weight.squeeze(1) + bias = self.qkv_conv1d.bias + + self.attn = RadixLinearAttention( + layer_id=self.layer_idx, + num_q_heads=self.num_k_heads // self.attn_tp_size, + num_k_heads=self.num_k_heads // self.attn_tp_size, + num_v_heads=self.num_v_heads // self.attn_tp_size, + head_q_dim=self.head_k_dim, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_weights=conv_weights, + bias=bias, + A_log=self.A_log, + dt_bias=self.dt_bias, + ) + # KDA safe gate: checkpoint trained with gate_lower_bound=-5.0 + self.attn.lower_bound = config.linear_attn_config.get("gate_lower_bound", None) + # Set by _prepare_fused_decode() once weights are loaded. + self._kda_fused_decode_ready = False + + def forward_qkvbfg(self, hidden_states: torch.Tensor): + qkv, _ = self.qkv_proj(hidden_states) + beta = self.b_proj(hidden_states)[0] + forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0] + if self.use_full_rank_gate: + g_proj_states = self.g_proj(hidden_states)[0] + else: + g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0] + return qkv, beta, forget_gate, g_proj_states + + def _merge_bfa_weights(self) -> None: + """Merge f_a_proj (head_dim outputs) + b_proj (heads/tp outputs). + + Both are skinny same-input GEMVs at decode: b lands in a cublas dot + kernel pair, f_a in a splitK GEMM. One [H, head_dim + heads/tp (+pad)] + GEMV replaces both. f_a leads so its output slice starts at offset 0, + and the width is padded to a multiple of 8 so every fused-output row + stays 16-byte aligned for vectorized consumers (tiny-GEMM on f_b). + + Called once from load_weights (after all weights are loaded, before + cuda graph capture).""" + if not self.use_full_rank_gate: + return + self._bfa_w, sizes = _merge_weights_as_views( + [self.f_a_proj, self.b_proj], pad_rows_to=8 + ) + self._bfa_fa_size, self._bfa_b_size = sizes + + def _prepare_fused_decode(self) -> None: + """Static inputs for the fused KDA decode kernel + (kernels/ops/attention/kda_fused_decode): per-segment transposed fp32 conv + weights [4, seg], dense fp32 conv bias, fp32 output-norm weight. Stashed on the + attention layer for the KDA backend; when the shapes do not match + the compiled kernel the stash stays unset and decode keeps the + unfused chain. Called once from load_weights (after all weights are + loaded, before cuda graph capture).""" + if _is_hip: + # The fused KDA decode kernel is NVIDIA-only + return + layer = self.attn + w = layer.conv_weights + seg = 12 * 128 # compiled for H = HV = 12 heads of 128 (TP8) + if ( + w is None + or w.ndim != 2 + or w.shape != (3 * seg, 4) + or w.dtype != torch.float32 + or layer.A_log is None + or layer.A_log.numel() != 12 + or layer.A_log.dtype != torch.float32 + or layer.dt_bias is None + or tuple(layer.dt_bias.shape) != (seg,) + or layer.dt_bias.dtype != torch.float32 + ): + rank0_log( + "K3 fused KDA decode disabled: unexpected conv/A_log/dt_bias " + f"layout (conv {None if w is None else tuple(w.shape)}, " + f"A_log {None if layer.A_log is None else tuple(layer.A_log.shape)}, " + f"dt_bias {None if layer.dt_bias is None else tuple(layer.dt_bias.shape)})" + ) + return + # Conv weights/bias stay fp32 (checkpoint dtype; the kernel loads + # them as fp32, matching the triton chain's precision exactly). + wt = w.t().contiguous() # [4, 3*seg] + bias = layer.bias + conv_bias = ( + bias.float().contiguous() + if bias is not None + else torch.zeros(3 * seg, dtype=torch.float32, device=w.device) + ) + layer._k3_fused_decode_args = ( + wt[:, :seg].contiguous(), + wt[:, seg : 2 * seg].contiguous(), + wt[:, 2 * seg :].contiguous(), + conv_bias, + layer.A_log.detach().reshape(-1), # view; kernel wants [12] + self.o_norm.weight.data.float().contiguous(), + float(self.o_norm.eps), + ) + self._kda_fused_decode_ready = True + + def forward_qkvbfg_fused(self, hidden_states: torch.Tensor): + if self.use_full_rank_gate: + if self._bfa_w is not None: + w = self._bfa_w + n_fa, n_b = self._bfa_fa_size, self._bfa_b_size + from sglang.kernels.ops.kimi_k3 import kimi_k3_tiny_gemm as gemm + + if ( + self._bfa_alt_stream is not None + and get_is_capture_mode() + and 0 < hidden_states.shape[0] <= self._bfa_bs_limit + ): + # Issue the tiny [f_a|b] + f_b GEMVs on the side stream, + # then the wide [q,k,v,g] GEMM on the main stream; both + # read only hidden_states. Join before the split's + # consumers touch beta/forget_gate. + alt = self._bfa_alt_stream + cur = torch.cuda.current_stream() + alt.wait_stream(cur) + with torch.cuda.stream(alt): + bfa = gemm(hidden_states, w) + forget_gate = gemm(bfa[..., :n_fa], self.f_b_proj.weight) + beta = bfa[..., n_fa : n_fa + n_b] + fused_states, _ = self.fused_qkvg_proj(hidden_states) + qkv, g_proj_states = torch.split( + fused_states, self.split_sizes, dim=-1 + ) + cur.wait_stream(alt) + return qkv, beta, forget_gate, g_proj_states + + fused_states, _ = self.fused_qkvg_proj(hidden_states) + qkv, g_proj_states = torch.split(fused_states, self.split_sizes, dim=-1) + bfa = gemm(hidden_states, w) + forget_gate = gemm(bfa[..., :n_fa], self.f_b_proj.weight) + beta = bfa[..., n_fa : n_fa + n_b] + else: + fused_states, _ = self.fused_qkvg_proj(hidden_states) + qkv, g_proj_states = torch.split(fused_states, self.split_sizes, dim=-1) + beta = self.b_proj(hidden_states)[0] + forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0] + else: + fused_states = self.fused_qkvbfg_a_proj(hidden_states) + qkv, beta, fg_a_states = torch.split(fused_states, self.split_sizes, dim=-1) + forget_gate, g_proj_states = self.fused_fg_b_proj( + fg_a_states.view(-1, 2, self.head_dim).transpose(0, 1) + ) + return qkv, beta, forget_gate, g_proj_states + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + if self.do_fuse_qkvbfg: + mixed_qkv, beta, forget_gate, g_proj_states = self.forward_qkvbfg_fused( + hidden_states + ) + else: + mixed_qkv, beta, forget_gate, g_proj_states = self.forward_qkvbfg( + hidden_states + ) + + if not forward_batch.forward_mode.is_decode(): + forget_gate = forget_gate.unflatten(-1, (-1, self.head_dim)) + if not forward_batch.forward_mode.is_target_verify(): + # Only chunk_kda (extend) wants pre-activated beta; the verify + # kernel sigmoids it in-kernel like decode. + beta = beta.float().sigmoid() + forget_gate = forget_gate.unsqueeze(0) + beta = beta.unsqueeze(0) + + # Fused KDA handoff (attempt-and-verify): offer the output-norm gate + # so covered decode and target-verify kernels can fold gated RMSNorm + # into the recurrence kernel. If the backend leaves the stash + # unconsumed (env off or shape not covered), apply o_norm here as + # before. + fused_onorm = self._kda_fused_decode_ready and ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + ) + if fused_onorm: + self.attn._k3_onorm_gate = g_proj_states + self.attn._k3_onorm_consumed = False + + core_attn_out = self.attn( + forward_batch, + mixed_qkv=mixed_qkv, + a=forget_gate, + b=beta, + ) + + if fused_onorm: + self.attn._k3_onorm_gate = None + fused_onorm = self.attn._k3_onorm_consumed + if not fused_onorm: + norm_gate = g_proj_states.unflatten(-1, (-1, self.head_dim)) + core_attn_out = self.o_norm(core_attn_out, norm_gate) + core_attn_out = core_attn_out.squeeze(0).flatten(-2) + if self.all_reduce_fusion: + out = _k3_symm_o_proj_out(self.o_proj, core_attn_out) + partial, _ = self.o_proj(core_attn_out, output_tensor=out) + return partial + return self.o_proj(core_attn_out)[0] + + +class KimiK3MLAAttention(DeepseekV2AttentionMLA): + """MLA with output gate for K3. Gate is applied in TP-local space before o_proj.""" + + def __init__( + self, + config, + layer_idx: int, + quant_config: Optional[QuantizationConfig] = None, + all_reduce_fusion: bool = False, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + gate_alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + self.all_reduce_fusion = all_reduce_fusion + self.use_output_gate = getattr(config, "mla_use_output_gate", False) + super().__init__( + layer_id=layer_idx, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + quant_config=quant_config, + prefix=prefix, + config=config, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + skip_rope=True, + reduce_results=not self.all_reduce_fusion, + alt_stream=alt_stream, + ) + # Installed before the output-gate wrap below so the gate multiply is + # applied to x before the fused GEMM+AR sees it. + if self.all_reduce_fusion and not _o_proj_takes_output(self.o_proj): + # the fused AR reduces o_proj's output in place out of a symmetric + # buffer, which needs the GEMM to write into caller-owned storage + self.all_reduce_fusion = False + self.o_proj.reduce_results = True + self.o_proj.use_dp_attention_reduce = True + k3_gemm_ar.maybe_wrap_o_proj(self.o_proj) + if self.all_reduce_fusion: + # reduce_results=False was passed through super().__init__ above; + # the fused all-reduce does the reduce itself and reduces the o_proj + # output in place, so hand the GEMM a slice of the persistent + # symmetric buffer (k3_ar_fusion.symm_buffer) and do NOT set + # use_dp_attention_reduce — its inner attn_tp symm_ctx allocates its + # own output and would defeat the caller-owned buffer. At the fusion + # config (attn_tp==tp) the fused full-TP reduce is the same group as + # the attn_tp reduce. + # The wrap is installed before the output-gate wrap below so the + # gate multiply stays outside it and only the o_proj GEMM writes the + # region slice. NOTE: the captured name must differ from the + # gate block's `_orig_o_proj_forward` — closures capture the + # __init__ local by reference, and reusing the name would rebind it + # to this wrapper (infinite recursion + nested pool enter). + _symm_inner_o_proj_forward = self.o_proj.forward + _symm_o_proj = self.o_proj + + def _symm_o_proj_forward(x, *args, **kwargs): + return _symm_inner_o_proj_forward( + x, + *args, + output_tensor=_k3_symm_o_proj_out(_symm_o_proj, x), + **kwargs, + ) + + self.o_proj.forward = _symm_o_proj_forward + else: + # K3 has no LayerCommunicator, so o_proj (reduce_results=True by + # default here, unlike deepseek's communicator flow) must reduce + # within the attn-TP group itself — the default full-TP collective + # is the wrong group at attn_tp>1 and deadlocks against idle DP + # ranks. + self.o_proj.use_dp_attention_reduce = True + if self.use_output_gate: + projection_size = config.num_attention_heads * config.v_head_dim + # Shard by attn-TP to match the attention output (DSV2 MLA shards + # heads across the attention-TP group, not the global TP group). + self.g_proj = ColumnParallelLinear( + config.hidden_size, + projection_size, + bias=False, + quant_config=quant_config, + tp_rank=get_parallel().attn_tp_rank, + tp_size=get_parallel().attn_tp_size, + prefix=f"{prefix}.g_proj", + ) + # Output gate must multiply the TP-local attention output right + # before o_proj (vLLM: attn_out * sigmoid(g_proj(hidden_states))). + # o_proj is invoked deep inside DeepseekV2AttentionMLA forward + # cores, so wrap its forward at the instance level; the module + # itself (weights, reduce_results, loading path) is untouched. + self._gate_hidden_states = None + # (gate, producer stream) issued on the alt stream by forward(); + # None when the lazy path computes the gate here instead. + self._gate_precomputed = None + self._gate_alt_stream = gate_alt_stream + # Above this token count the attention-core kernels fill the SMs + # on their own and the overlap only adds sync overhead (same + # bound as deepseek_v4). + self._gate_bs_limit = ( + (128 if is_blackwell_supported() else 64) + if self._gate_alt_stream is not None + else 0 + ) + _orig_o_proj_forward = self.o_proj.forward + + def _gated_o_proj_forward(x, *args, **kwargs): + gate_input = self._gate_hidden_states + self._gate_hidden_states = None + precomputed = self._gate_precomputed + self._gate_precomputed = None + if precomputed is not None: + # Use wait_stream rather than an explicit event so the + # breakable-CUDA-graph runner can track the side-stream + # join across graph-segment boundaries. + torch.cuda.current_stream().wait_stream(precomputed[1]) + if gate_input is not None and not isinstance(x, tuple): + gate = ( + precomputed[0] + if precomputed is not None + else self.g_proj(gate_input)[0] + ) + from sglang.kernels.ops.kimi_k3 import mla_output_gate + + if mla_output_gate.covered(x, gate): + # One kernel for x * sigmoid(gate); double rounding + # matches the unfused pair bit-for-bit. + x = mla_output_gate.kimi_k3_mla_output_gate(x, gate) + else: + x = x * torch.sigmoid(gate) + return _orig_o_proj_forward(x, *args, **kwargs) + + self.o_proj.forward = _gated_o_proj_forward + + def _precompute_output_gate(self, hidden_states: torch.Tensor) -> None: + """Issue the output-gate GEMM on the alt stream so it overlaps the + attention core; the lazy path in the o_proj wrap otherwise computes + it on the critical path right before the gate multiply. The gate + tensor stays referenced via _gate_precomputed until the wrap joins, + so its memory cannot be reused while the alt stream still writes.""" + self._gate_precomputed = None + if ( + self._gate_alt_stream is not None + and get_is_capture_mode() + # The attention-core break ends the segment between the alt-stream + # event record and the o_proj-side wait, so under breakable capture + # the wait would cross graph segments; use the lazy path instead. + and not is_in_breakable_cuda_graph() + and (0 < hidden_states.shape[0] <= self._gate_bs_limit) + ): + alt = self._gate_alt_stream + alt.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(alt): + gate, _ = self.g_proj(hidden_states) + self._gate_precomputed = (gate, alt) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + **kwargs, + ): + if self.use_output_gate: + self._gate_hidden_states = hidden_states + self._precompute_output_gate(hidden_states) + return super().forward( + positions, hidden_states, forward_batch, zero_allocator, **kwargs + ) + + +class KimiK3DecoderLayer(nn.Module): + """Decoder layer carrying the K3 attention-residual stream.""" + + def __init__( + self, + config: KimiLinearConfig, + layer_idx: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_streams: Optional[List[torch.cuda.Stream]] = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.is_moe = config.is_moe + self.layer_idx = layer_idx + self._dp_attention = is_dp_attention_enabled() + # mlp-sync (DP attention OR MoE a2a/EP) pads extend batches to + # attn_tp multiples; attention must then run on the real rows only. + self._trim_padded_attn = require_mlp_sync(get_server_args()) + # A layer runs MoE (vs a plain dense MLP) iff it is past the dense + # prefix and on the MoE cadence — same predicate the mlp construction + # below uses. + self._is_moe_layer = ( + self.is_moe + and config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + # SP-MoE (EP a2a backend — megamoe or DeepEP): o_proj defers its + # attention-TP reduction; this layer completes it as a reduce-scatter + # so the whole MoE region (agg2, norms, gate, latent projs, tp1 + # shared experts, EP a2a dispatch) runs on 1/attn_tp of the rows, + # then all-gathers rows back after the MoE tail add. RS+AG moves the + # same bytes the o_proj all-reduce did, the shared-expert all-reduce + # disappears via tp1 weights, and each rank dispatches only its shard + # through the a2a (kills the attn_tp-fold dispatch redundancy) — + # strictly less communication + MoE-front compute /attn_tp. Works the + # same under DP attention: the attn_tp group is then the + # within-replica subgroup, rows are the DP-local batch, and + # KimiK3MoE skips the DP gather under EP a2a so the shard flows + # straight into the a2a. With attn_tp == 1 (full DP attention) there + # is no attention reduce to convert — the MoE-side gather skip alone + # removes the replication. Dense layers are excluded: their + # column-parallel MLP has no per-token decomposition that survives a + # token shard. + _a2a_backend = get_moe_a2a_backend() + self._sp_moe = ( + (_a2a_backend.is_megamoe() or _a2a_backend.is_deepep()) + and self._is_moe_layer + and get_parallel().attn_tp_group.world_size > 1 + ) + + # The fused all-reduce only serves the attn-res path (attn_res is + # config-static), so the standard path stays byte-for-byte untouched + # and always sees a reduced attention output. + # Mutually exclusive with SP-MoE: both complete o_proj's deferred + # reduction, but SP-MoE reduce-scatters to a shard whereas the fusion + # produces the full batch in a symm buffer — an SP-MoE layer builds + # o_proj in the plain deferred-reduce config (reduce_results forced off + # below) and reduce-scatters instead. + attn_tp_size = get_parallel().attn_tp_size + self.all_reduce_fusion = ( + not self._sp_moe + and attn_tp_size > 1 + and attn_tp_size == get_parallel().tp_size + and config.attn_res_block_size is not None + and k3_ar_fusion.enabled() + ) + + # Attention + if config.is_kda_layer(layer_idx): + self.self_attn = KimiK3DeltaAttention( + layer_idx=layer_idx, + hidden_size=config.hidden_size, + config=config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + all_reduce_fusion=self.all_reduce_fusion, + # Shared with the MLA gate stream: KDA and MLA layers never + # run concurrently within one forward, so the stream is free. + bfa_alt_stream=(alt_streams[2] if alt_streams is not None else None), + ) + else: + self.self_attn = KimiK3MLAAttention( + config=config, + layer_idx=layer_idx, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + all_reduce_fusion=self.all_reduce_fusion, + alt_stream=alt_streams[1] if alt_streams is not None else None, + gate_alt_stream=alt_streams[2] if alt_streams is not None else None, + ) + + # the attention drops the fusion when its o_proj cannot write into + # caller-owned storage; the layer's own AR call-site must agree + self.all_reduce_fusion = self.self_attn.all_reduce_fusion + + # MLP / MoE + if self._is_moe_layer: + self.mlp = KimiK3MoE( + config=config, + quant_config=quant_config, + layer_idx=layer_idx, + prefix=f"{prefix}.mlp", + alt_stream=alt_streams[0] if alt_streams is not None else None, + ) + else: + self.mlp = KimiK3MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + # Attention Residual + self.use_attn_residuals = config.attn_res_block_size is not None + if self.use_attn_residuals: + self.attn_res_block_size = config.attn_res_block_size + self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 + self.prev_valid_blocks = _cdiv(layer_idx, self.attn_res_block_size) + self.self_attention_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + if self._sp_moe: + # o_proj emits TP-partial sums; _finish_attn_reduce completes the + # reduction (RS on the clean attn-res path, AR on fallbacks). + o_proj = getattr(self.self_attn, "o_proj", None) + assert o_proj is not None, "SP-MoE requires attention exposing o_proj" + o_proj.reduce_results = False + if k3_sp_collective.enabled(): + # The table selects NVLS pull RS for larger token buckets. + # Only those o_proj outputs come from the persistent symmetric + # buffer; small push RS keeps the regular graph allocator. + _sp_inner_o_proj_forward = o_proj.forward + + def _sp_o_proj_forward(x, *args, **kwargs): + output_rows = k3_sp_collective.get_o_proj_output_rows(x.shape[0]) + if k3_sp_collective.requires_symmetric_rs(output_rows, x.device): + output = k3_sp_collective.get_o_proj_output_buffer( + output_rows, x.dtype, o_proj.output_size + ) + result = _sp_inner_o_proj_forward( + x, *args, output_tensor=output[: x.shape[0]], **kwargs + ) + k3_sp_collective.register_o_proj_output(result[0], output) + return result + return _sp_inner_o_proj_forward(x, *args, **kwargs) + + o_proj.forward = _sp_o_proj_forward + + def _finish_attn_reduce( + self, + attn_out: torch.Tensor, + allow_scatter: bool, + residual: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, int, bool]: + """Complete o_proj's deferred TP reduction under SP-MoE. + + Returns (reduced tensor, shard row offset, residual_fused); offset is + -1 when the result covers the full batch (non-SP mode, or fallback + all-reduce for row counts not divisible by attn_tp).""" + if not self._sp_moe: + return attn_out, -1, False + group = get_parallel().attn_tp_group + num_tokens = attn_out.shape[0] + if allow_scatter and num_tokens > 0 and num_tokens % group.world_size == 0: + shard = num_tokens // group.world_size + custom_out = k3_sp_collective.reduce_scatter_res(attn_out, residual) + if custom_out is not None: + return ( + custom_out, + group.rank_in_group * shard, + residual is not None, + ) + out = torch.empty( + (shard, attn_out.shape[1]), + dtype=attn_out.dtype, + device=attn_out.device, + ) + group.reduce_scatter_tensor(out, attn_out) + return out, group.rank_in_group * shard, False + return group.all_reduce(attn_out), -1, False + + def _run_self_attn( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + # DP attention: idle ranks (padded to the global shape) have no + # attention metadata; pass hidden_states through shape-preserving + # (same as the LayerCommunicator models' is_idle skip). + if forward_batch.forward_mode.is_idle(): + return hidden_states + + # mlp-sync (DP attention OR MoE a2a/EP — require_mlp_sync) pads + # extend batches to a multiple of attn_tp_size + # (prepare_mlp_sync_batch ceil_align), but the attention metadata + # (qo_indptr / query_start_loc) covers only the real tokens — the + # flashinfer ragged prefill rejects the row mismatch, and silent + # paths write the padded rows' garbage KV through the zero-padded + # out_cache_loc entries (clobbering pool slot 0 → cross-request + # corruption). Run attention on the real rows and zero-pad the + # output back; padded rows are discarded downstream. + num_padded = hidden_states.shape[0] + num_real = num_padded + if self._trim_padded_attn and forward_batch.forward_mode.is_extend(): + extend_lens = forward_batch.extend_seq_lens_cpu + if extend_lens is not None: + num_real = min(int(sum(extend_lens)), num_padded) + if num_real != num_padded: + with k3_sp_collective.o_proj_output_rows(num_padded): + attn_out = self._run_self_attn_inner( + hidden_states[:num_real], + positions[:num_real], + forward_batch, + zero_allocator, + ) + padded_o_proj = k3_sp_collective.finish_padded_o_proj_output( + attn_out, num_padded + ) + if padded_o_proj is not None: + return padded_o_proj + out = hidden_states.new_zeros(num_padded, attn_out.shape[-1]) + out[:num_real] = attn_out + return out + return self._run_self_attn_inner( + hidden_states, positions, forward_batch, zero_allocator + ) + + def _run_self_attn_inner( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + # For MLA layers with q_lora_rank, set up communicator attn_inputs + # before the forward call (normally done by LayerCommunicator). + from sglang.srt.layers.communicator import ( + AttentionInputs, + get_attn_tp_context, + ) + + qkv_latent_func = getattr(self.self_attn, "prepare_qkv_latent", None) + if qkv_latent_func is not None: + attn_inputs = AttentionInputs(hidden_states, forward_batch, qkv_latent_func) + get_attn_tp_context().set_attn_inputs(attn_inputs) + + result = self.self_attn( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + + if qkv_latent_func is not None: + get_attn_tp_context().clear_attn_inputs() + + return result + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + attn_res: Optional[AttnResidual], + zero_allocator: BumpAllocator, + input_sharded: bool = False, + keep_sharded: bool = False, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], bool]: + if attn_res is not None: + return self._forward_attn_residual( + positions, + hidden_states, + residual, + attn_res, + forward_batch, + zero_allocator, + input_sharded, + keep_sharded, + ) + + assert not input_sharded + # Standard residual path + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self._run_self_attn( + hidden_states, positions, forward_batch, zero_allocator + ) + # standard path returns a full-size residual to the next layer, so + # complete the deferred o_proj reduction as a plain all-reduce + hidden_states, _, _ = self._finish_attn_reduce( + hidden_states, allow_scatter=False + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states, forward_batch=forward_batch) + return hidden_states, residual, False + + def _forward_attn_residual( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + attn_res: AttnResidual, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + input_sharded: bool, + keep_sharded: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], bool]: + # Between attn-res layers hidden_states carries the previous layer's + # un-added MLP delta and prefix_sum the prefix it extends (None at + # stream start / PP entry, where hidden_states already is the head). + + # ---- Aggregation 1: attention side. Write layers snapshot the + # pre-attention prefix into the bank in the same call (fused into + # the fast kernel; standalone copy on other paths). ---- + if input_sharded: + assert self._sp_moe + input_rows = _sp_local_rows(hidden_states) + fused_ag = attn_res.forward_sp_all_gather( + hidden_states, + prefix_sum, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + rows=input_rows, + write=self.is_block_write_layer, + ) + if fused_ag is not None: + hidden_states, prefix_sum = fused_ag + else: + hidden_states, prefix_sum = attn_res.forward( + hidden_states, + prefix_sum, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + rows=input_rows, + write=self.is_block_write_layer, + ) + # Aggregate/norm and snapshot only this rank's rows, then + # gather the normalized tensor consumed by attention. + hidden_states = _sp_all_gather_rows(hidden_states) + else: + hidden_states, prefix_sum = attn_res.forward( + hidden_states, + prefix_sum, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + write=self.is_block_write_layer, + ) + if self.is_block_write_layer: + prefix_sum = None + + # ---- Attention ---- + hidden_states = self._run_self_attn( + hidden_states, positions, forward_batch, zero_allocator + ) + + # ---- Complete o_proj's deferred reduction ---- + # SP-MoE takes precedence (reduce-scatter to this rank's token shard); + # otherwise the fused all-reduce when enabled; otherwise o_proj already + # reduced itself (use_dp_attention_reduce, on when neither is active). + rows = None + shard_lo = -1 + agg2_fused = False + if self._sp_moe: + group = get_parallel().attn_tp_group + if ( + hidden_states.shape[0] > 0 + and hidden_states.shape[0] % group.world_size == 0 + ): + shard = hidden_states.shape[0] // group.world_size + fused_rows = slice( + group.rank_in_group * shard, + (group.rank_in_group + 1) * shard, + ) + fused_rs = attn_res.forward_sp_reduce_scatter( + hidden_states, + prefix_sum, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + rows=fused_rows, + ) + else: + fused_rs = None + if fused_rs is not None: + hidden_states, prefix_sum = fused_rs + rows = fused_rows + shard_lo = fused_rows.start + agg2_fused = True + else: + hidden_states, shard_lo, residual_fused = self._finish_attn_reduce( + hidden_states, allow_scatter=True, residual=prefix_sum + ) + if shard_lo >= 0: + rows = slice(shard_lo, shard_lo + hidden_states.shape[0]) + if residual_fused: + prefix_sum = None + elif prefix_sum is not None: + # Shard carry already holds the destination-local prefix; + # the first sharded layer still holds a full-batch prefix. + if prefix_sum.shape[0] != hidden_states.shape[0]: + prefix_sum = prefix_sum[rows] + elif self.all_reduce_fusion: + # Complete the o_proj reduce here, folding the pending prefix add + # into the fused all-reduce; attn_res then takes the pre-added + # tensor through its prefix_sum=None branch (same semantics: + # (normed, new_prefix) with new_prefix = prefix + attn_out). + hidden_states = k3_ar_fusion.all_reduce(hidden_states, prefix_sum) + prefix_sum = None + + # ---- Aggregation 2: MLP side (on the shard under SP-MoE) ---- + if not agg2_fused: + hidden_states, prefix_sum = attn_res.forward( + hidden_states, + prefix_sum, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + rows=rows, + ) + + # ---- MLP (consumes +prefix_sum: MoE folds it into the 3-way tail + # add, dense adds it after down_proj) ---- + out = self.mlp( + hidden_states, prefix_sum=prefix_sum, forward_batch=forward_batch + ) + if shard_lo >= 0: + if keep_sharded: + return out, None, True + out = _sp_all_gather_rows(out) + return out, None, False + + +class KimiK3LinearModel(nn.Module): + """K3 language-model backbone.""" + + def __init__( + self, + config: KimiLinearConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.pp_group = get_pp_group() + self.dspark_layers_to_capture: Optional[list[int]] = None + self._dp_attention = is_dp_attention_enabled() + self._trim_padded_attn = require_mlp_sync(get_server_args()) + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + # Under DP attention each rank embeds only its local tokens: + # reduce within the attention-TP group, not the full TP group. + **get_embedding_tp_kwargs(), + ) + else: + self.embed_tokens = PPMissingLayer() + + # Multi-stream pool (deepseek_v4 pattern): every alt stream is + # constructed here and threaded down to the layers. Slots: + # [0] MoE dual-stream shared-expert tail + # [1] DeepseekV2AttentionMLA base internals (forwarded; unused by K3) + # [2] MLA output-gate GEMM, overlaps the attention core + # (The attn-res bank write no longer needs a stream: it is fused + # into the agg1 fast kernel, see AttnResidual.forward(write=True).) + # Disable on HIP code path. + self.alt_streams = None if _is_hip else [torch.cuda.Stream() for _ in range(3)] + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: KimiK3DecoderLayer( + layer_idx=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + alt_streams=self.alt_streams, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=f"{prefix}.layers", + ) + + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.attn_res_block_size is not None: + self.output_attn_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + forward_batch: ForwardBatch, + inputs_embeds: torch.Tensor | None = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(input_ids) + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + if TYPE_CHECKING: + assert isinstance(hidden_states, torch.Tensor) + assert isinstance(residual, torch.Tensor | None) + + # mlp-sync (DP attention OR MoE a2a/EP) pads extend batches to a + # multiple of attn_tp_size; attention layers run on the real rows + # only (_run_self_attn trims), so the KV write locations must match + # the trimmed length. positions and hidden_states keep the padded + # length for the DP gather/scatter and the MoE. + if ( + self._trim_padded_attn + and forward_batch.forward_mode.is_extend() + and forward_batch.out_cache_loc is not None + and forward_batch.extend_seq_lens_cpu is not None + ): + num_real = int(sum(forward_batch.extend_seq_lens_cpu)) + if forward_batch.out_cache_loc.shape[0] > num_real: + forward_batch.out_cache_loc = forward_batch.out_cache_loc[:num_real] + + total_num_layers = self.end_layer - self.start_layer + device = hidden_states.device + zero_allocator = BumpAllocator( + buffer_size=total_num_layers * 2, + dtype=torch.float32, + device=device, + ) + + attn_res = None + if self.config.attn_res_block_size is not None: + attn_res_block_num = _cdiv(self.end_layer, self.config.attn_res_block_size) + attn_res = AttnResidual( + hidden_states, + attn_res_block_num, + block_residual=residual, + ) + residual = None + + # Carry the raw residual stream as a token shard across consecutive + # SP-MoE layers. PP transfer and dspark capture require full tensors, + # so those uncommon paths keep the established gather-per-layer flow. + sp_attn_res = ( + attn_res is not None + and envs.SGLANG_K3_SP_ATTN_RES.get() + and self.pp_group.world_size == 1 + and self.dspark_layers_to_capture is None + and k3_sp_collective.enabled() + ) + sp_sharded = False + aux_hidden_states = [] + for i in range(self.start_layer, self.end_layer): + if sp_sharded and not self.layers[i]._sp_moe: + hidden_states = _sp_all_gather_rows(hidden_states) + sp_sharded = False + with get_global_expert_distribution_recorder().with_current_layer(i): + hidden_states, residual, sp_sharded = self.layers[i]( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + residual=residual, + attn_res=attn_res, + zero_allocator=zero_allocator, + input_sharded=sp_sharded, + keep_sharded=sp_attn_res, + ) + if ( + self.dspark_layers_to_capture is not None + and i in self.dspark_layers_to_capture + ): + aux_hidden_states.append( + self._dspark_capture_stream(i, hidden_states, residual, attn_res) + ) + + if not self.pp_group.is_last_rank: + assert not sp_sharded + if attn_res is not None: + if residual is not None: + # Materialize the delayed MLP add: the wire carries the + # full stream head (bit-identical to the fused fold). + hidden_states = residual + hidden_states + residual = attn_res.block_residual # raw bank across ranks + return PPProxyTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + if hidden_states.shape[0] != 0: + if attn_res is not None: + # ---- Final aggregation (output side, folds delayed add) ---- + if sp_sharded: + output_rows = _sp_local_rows(hidden_states) + fused_output = attn_res.forward_sp_all_gather( + hidden_states, + residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + rows=output_rows, + ) + if fused_output is not None: + hidden_states, _ = fused_output + else: + hidden_states, _ = attn_res.forward( + hidden_states, + residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + rows=output_rows, + ) + hidden_states = _sp_all_gather_rows(hidden_states) + else: + hidden_states, _ = attn_res.forward( + hidden_states, + residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + else: + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + + if self.dspark_layers_to_capture is not None: + return hidden_states, aux_hidden_states + return hidden_states + + def _dspark_capture_stream( + self, + layer_idx: int, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + attn_res: Optional[AttnResidual], + ) -> torch.Tensor: + """Stream value after `layer_idx`: the pre-norm mixture its next + consumer would compute (next layer's attention side; output side + for the last layer).""" + if attn_res is None: + return hidden_states if residual is None else hidden_states + residual + if residual is not None: + # Materialize a delayed MLP add (mirrors the PP-wire fold). + hidden_states = residual + hidden_states + if layer_idx + 1 < self.end_layer: + next_layer = self.layers[layer_idx + 1] + score_proj = next_layer.self_attention_res_proj + score_norm = next_layer.self_attention_res_norm + nvb = next_layer.prev_valid_blocks + else: + # Last layer: the model's own output-side aggregation weights. + score_proj = self.output_attn_res_proj + score_norm = self.output_attn_res_norm + nvb = _cdiv(self.end_layer, self.config.attn_res_block_size) + return aggregate_stream( + hidden_states, attn_res.block_residual, nvb, score_proj, score_norm + ) + + +class KimiK3LinearForCausalLM(nn.Module): + """Text-only K3 causal LM.""" + + def __init__( + self, + config: KimiLinearConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = KimiK3LinearModel( + config, quant_config, prefix=maybe_prefix(prefix, "model") + ) + self.pp_group = get_pp_group() + if self.pp_group.is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + logit_scale = getattr(config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor(config=config, logit_scale=logit_scale) + self.capture_aux_hidden_states = False + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None: + if self.pp_group.world_size > 1: + # Capture layers living on non-last PP ranks would be silently + # skipped (the flag is only set on the last rank). + raise NotImplementedError("DSPARK aux hidden capture requires PP=1.") + if not self.pp_group.is_last_rank: + return + if layer_ids is None: + raise ValueError( + "DSPARK requires explicit layer_ids for aux hidden capture." + ) + self.capture_aux_hidden_states = True + self.model.dspark_layers_to_capture = list(layer_ids) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + embeds = input_embeds if input_embeds is not None else inputs_embeds + hidden_states = self.model( + input_ids, positions, forward_batch, embeds, pp_proxy_tensors + ) + if self.pp_group.is_last_rank: + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + aux_hidden_states, + ) + return hidden_states + + def prepare_context_parallel_metadata_for_dcp( + self, + seq_lens: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_prefix_lens_cpu: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + seq_lens_sum: int, + kv_buffer_shape: torch.Size, + kv_cache_dtype, + kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn, + ): + return prepare_decode_context_parallel_metadata( + seq_lens=seq_lens, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + seq_lens_sum=seq_lens_sum, + kv_buffer_shape=kv_buffer_shape, + kv_cache_dtype=kv_cache_dtype, + kv_cache_device=kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn=create_chunked_prefix_cache_kv_indices_fn, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + use_full_rank_gate = bool( + (self.config.linear_attn_config or {}).get("use_full_rank_gate", False) + ) + if use_full_rank_gate: + # Fused layout (K3): [q, k, v, g] column-parallel; b / f_a / f_b + # are standalone modules loaded by name. + fused_qkvbfg_mapping = [ + (".fused_qkvg_proj", ".q_proj", 0), + (".fused_qkvg_proj", ".k_proj", 1), + (".fused_qkvg_proj", ".v_proj", 2), + (".fused_qkvg_proj", ".g_proj", 3), + ] + else: + # Fused layout (low-rank gate): [q, k, v, b] + [f_a, g_a] + fused_qkvbfg_mapping = [ + (".fused_qkvbfg_a_proj", ".q_proj", 0), + (".fused_qkvbfg_a_proj", ".k_proj", 1), + (".fused_qkvbfg_a_proj", ".v_proj", 2), + (".fused_qkvbfg_a_proj", ".b_proj", 3), + (".fused_qkvbfg_a_proj", ".f_a_proj", 4), + (".fused_qkvbfg_a_proj", ".g_a_proj", 5), + (".fused_fg_b_proj", ".f_b_proj", 0), + (".fused_fg_b_proj", ".g_b_proj", 1), + ] + + stacked_params_mapping = [ + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + *fused_qkvbfg_mapping, + # Unfused QKV path + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + # Conv1d fusion + (".qkv_conv1d", ".q_conv1d", 0), + (".qkv_conv1d", ".k_conv1d", 1), + (".qkv_conv1d", ".v_conv1d", 2), + ] + + if self.config.is_moe: + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + else: + expert_params_mapping = [] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + num_hidden_layers = self.config.num_hidden_layers + for args in weights: + name, loaded_weight = args[:2] + kwargs = args[2] if len(args) > 2 else {} + + layer_id = get_layer_id(name) + if layer_id is not None and ( + layer_id < self.model.start_layer or layer_id >= self.model.end_layer + ): + continue + + # Skip weights of layers outside a truncated config (e.g. + # num_hidden_layers override for fast testing); the checkpoint may + # carry more layers than the instantiated model. + if ".layers." in name: + _lid = name.split(".layers.")[1].split(".")[0] + if _lid.isdigit() and int(_lid) >= num_hidden_layers: + continue + + # compressed-tensors MXFP4 stores as weight_packed; Mxfp4MoEMethod uses weight + if "weight_packed" in name: + name = name.replace("weight_packed", "weight") + + # MLA: fuse q_a_proj + kv_a_proj_with_mqa → fused_qkv_a_proj_with_mqa + if ".q_a_proj." in name or ".kv_a_proj_with_mqa." in name: + fused_name = name.replace(".q_a_proj.", ".fused_qkv_a_proj_with_mqa.") + fused_name = fused_name.replace( + ".kv_a_proj_with_mqa.", ".fused_qkv_a_proj_with_mqa." + ) + if fused_name in params_dict: + param = params_dict[fused_name] + if ".q_a_proj." in name: + param.data[: loaded_weight.shape[0]].copy_(loaded_weight) + else: + q_lora_rank = self.config.q_lora_rank or 0 + param.data[q_lora_rank:].copy_(loaded_weight) + loaded_params.add(fused_name) + continue + + if "rotary_emb.inv_freq" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if ("mlp.experts." in name) and name not in params_dict: + continue + # Fused projections only apply to KDA layers + if param_name in { + ".fused_qkvbfg_a_proj", + ".fused_fg_b_proj", + ".fused_qkvg_proj", + }: + layer_id = int(name.split(".")[2]) + if not self.config.is_kda_layer(layer_id): + continue + layer = self.model.layers[layer_id].self_attn + if not getattr(layer, "do_fuse_qkvbfg", False): + continue + if weight_name in {".q_proj", ".k_proj", ".v_proj"}: + layer_id = int(name.split(".")[2]) + if not self.config.is_kda_layer(layer_id): + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for idx, (param_name, weight_name, expert_id, shard_id) in enumerate( + expert_params_mapping + ): + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip experts of layers outside a truncated config (e.g. + # num_hidden_layers override), mirroring the non-expert + # `name not in params_dict` guard below. + if name not in params_dict: + break + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=shard_id, + ) + break + else: + if ( + name.endswith(".bias") + and name not in params_dict + and not self.config.is_linear_attn + ): + continue + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + + self.post_load_weights() + + def post_load_weights(self): + # Also invoked by loader post-load hooks (DummyModelLoader, + # ShardedStateLoader, remote-instance flows -- none of which call + # load_weights), so e.g. dummy-weight benchmarks get w_kc/w_vc and + # the fused buffers too. Same pattern as deepseek_v4. + # Post-load: absorb kv_b_proj into w_kc and w_vc for MLA layers + for layer_id in self.config.full_attention_layer_ids: + if layer_id >= len(self.model.layers): + continue # truncated config (e.g. num_hidden_layers override) + layer = self.model.layers[layer_id] + if isinstance(layer, PPMissingLayer): + continue + self_attn = layer.self_attn + w_kc, w_vc = self_attn.kv_b_proj.weight.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + self_attn.w_vc = w_vc.contiguous().transpose(1, 2) + if hasattr(self_attn.kv_b_proj, "weight_scale"): + self_attn.w_scale = self_attn.kv_b_proj.weight_scale + + # Post-load: precompute the attn-res combined score weights BEFORE + # cuda graph capture (a lazy first call inside get_cw would bake the + # multiply into every captured graph replay otherwise). Warm both + # dtypes: the fast kernel consumes bf16, the triton fallback fp32. + def _warm_cw(proj, norm): + get_cw(proj, norm, dtype=torch.bfloat16) + get_cw(proj, norm) + + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if layer.use_attn_residuals: + _warm_cw(layer.self_attention_res_proj, layer.self_attention_res_norm) + _warm_cw(layer.mlp_res_proj, layer.mlp_res_norm) + if hasattr(self.model, "output_attn_res_proj"): + _warm_cw(self.model.output_attn_res_proj, self.model.output_attn_res_norm) + + # Post-load: merge the horizontally-fused decode weights. Module + # weights are re-pointed to views of the merged buffers (net extra + # memory ~0), so this must run after all weights are loaded and + # before cuda graph capture. + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if isinstance(layer.mlp, KimiK3MoE): + layer.mlp._merge_front_weights() + # The router consumes the correction bias in fp32; convert the + # bf16 checkpoint values once (exact) so the per-call + # .to(float32) in topk becomes a no-op instead of one upcast + # kernel per MoE layer per step. + bias = layer.mlp.gate.e_score_correction_bias + if bias.dtype != torch.float32: + bias.data = bias.data.to(torch.float32) + if isinstance(layer.self_attn, KimiK3DeltaAttention): + layer.self_attn._merge_bfa_weights() + layer.self_attn._prepare_fused_decode() + + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer) or not isinstance( + layer.self_attn, KimiK3DeltaAttention + ): + continue + from sglang.kernels.ops.attention.fla.kda import ( + precompile_k3_recompute_w_u_kernel, + ) + + if precompile_k3_recompute_w_u_kernel( + num_heads=layer.self_attn.local_num_heads, + dtype=layer.self_attn.o_proj.weight.dtype, + device=layer.self_attn.dt_bias.device, + ): + rank0_log("Precompiled the Kimi-K3 KDA prefill kernel.") + break + + +class KimiK3ForConditionalGeneration(nn.Module): + """K3 multimodal wrapper: MoonViT3d tower + KimiK3LinearForCausalLM.""" + + # Raw HF checkpoint prefixes, before hf_to_sglang_mapper is applied. + encoder_only_safetensors_weight_prefixes = ( + "vision_tower.", + "mm_projector.", + ) + + hf_to_sglang_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + }, + orig_to_new_substr={ + "block_sparse_moe": "mlp", + }, + ) + + def __init__( + self, + config: KimiK3Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + **kwargs, + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + + # The dedicated K3 tower runs replicated (per-rank full weights); + # shard work across ranks image-wise via the DP runner. + self.use_data_parallel = True + + self.vision_tower = KimiK3VisionTower(config.vision_config) + self.mm_projector = KimiK3MultiModalProjector(config.vision_config) + + self.language_model = None + if not config.encoder_only: + self.language_model = KimiK3LinearForCausalLM( + config.text_config, + quant_config, + prefix="", + ) + + @property + def model(self): + return self.language_model + + def __setattr__(self, name, value): + if name == "model": + return + super().__setattr__(name, value) + + def post_load_weights(self): + # Delegate so DummyModelLoader's post-load hook reaches the LM tower. + if self.language_model is not None: + self.language_model.post_load_weights() + + def precompile_kernels_after_loading(self) -> None: + if self.config.language_only: + return + if self.vision_tower.precompile_fused_rope(): + logger.info("Precompiled dynamic-token fused K3 vision RoPE kernel") + if self.vision_tower.precompile_attention_backend(): + logger.info("Precompiled Kimi-K3 vision FA4 kernel") + + def get_input_embeddings(self): + if self.language_model is None: + raise AttributeError( + "get_input_embeddings() is not available in encoder-only mode" + ) + return self.language_model.model.embed_tokens + + @property + def lm_head(self): + if self.language_model is None: + raise AttributeError("lm_head is not available in encoder-only mode") + return self.language_model.lm_head + + def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None: + if self.language_model is None: + raise AttributeError( + "DSPARK layer capture is not available in encoder-only mode" + ) + self.language_model.set_dspark_layers_to_capture(layer_ids) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + device = self.vision_tower.device + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + image_grid_thws = [] + for item in items: + grid_thw = item.model_specific_data.get("image_grid_thw") + if grid_thw is None: + grid_thw = item.model_specific_data["grid_thws"] + if grid_thw.shape[0] != 1: + # One item must carry exactly one logical image so the DP + # owner assignment and the bounded CUDA-IPC lease accounting + # stay per-item; aggregated encoder inputs are split upstream + # (EPD encode server) before reaching this point. + raise ValueError( + "Kimi-K3 expects one vision grid per MultimodalDataItem; " + "split aggregated encoder inputs before get_image_feature()" + ) + image_grid_thws.append(grid_thw) + grid_thws_host = torch.concat(image_grid_thws, dim=0).cpu() + grid_thw_list = grid_thws_host.tolist() + + def materialize_item_features(image_indices: List[int]) -> torch.Tensor: + """Materialize features for the images assigned to this rank. + + K3 vision is image-wise data-parallel, so each image is consumed + by exactly one TP rank. Deferred CUDA-IPC proxies are + reconstructed here, after the assignment is known, so an image + crosses the tokenizer/scheduler boundary once instead of once + per rank; CPU-transport features likewise only pay their H2D + copy on the owner rank. The consumer count matches + MmItemMemoryPool.try_to_recycle(), which waits for the server TP + size rather than the attention subgroup size. + """ + parallel = get_parallel() + server_args = get_server_args() + ipc_consumer_count = max( + getattr(server_args, "tp_size", parallel.attn_tp_size), 1 + ) + device_index = device.index + if device.type == "cuda" and device_index is None: + device_index = torch.cuda.current_device() + + features = [] + for image_index in image_indices: + item = items[image_index] + if device.type == "cuda": + item.reconstruct( + device_index, ipc_consumer_count=ipc_consumer_count + ) + feature = item.feature + if not isinstance(feature, torch.Tensor): + raise TypeError( + "Kimi-K3 image feature must be a torch.Tensor, " + f"got {type(feature)}" + ) + features.append(feature) + return materialize_multimodal_features( + features, device=device, dtype=target_dtype + ) + + if self.use_data_parallel: + from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model + + image_embeds = run_dp_sharded_mrope_vision_model( + self.vision_tower, + None, + grid_thw_list, + rope_type="rope_2d", + # K3's tower pools the temporal dimension away: a t>1 grid + # still yields h*w/merge_area output embeddings, so the DP + # gather length must ignore t. + pool_temporal_dimension=True, + pass_grid_thw_list=True, + load_local_pixel_values=materialize_item_features, + pixel_values_device=device, + pixel_values_dtype=target_dtype, + ) + return self.mm_projector(image_embeds) + + pixel_values = materialize_item_features(list(range(len(items)))) + image_embeds = self.vision_tower(pixel_values, grid_thws_host.to(device)) + return self.mm_projector(image_embeds) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + pattern = MultiModalityDataPaddingPatternMultimodalTokens() + return pattern.pad_input_tokens(input_ids, mm_inputs) + + @property + def start_layer(self) -> int: + if self.language_model is None: + return 0 + return self.language_model.model.start_layer + + @property + def end_layer(self) -> int: + if self.language_model is None: + return self.config.text_config.num_hidden_layers + return self.language_model.model.end_layer + + def prepare_context_parallel_metadata_for_dcp( + self, + seq_lens: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_prefix_lens_cpu: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + seq_lens_sum: int, + kv_buffer_shape: torch.Size, + kv_cache_dtype, + kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn, + ): + return self.language_model.prepare_context_parallel_metadata_for_dcp( + seq_lens=seq_lens, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + seq_lens_sum=seq_lens_sum, + kv_buffer_shape=kv_buffer_shape, + kv_cache_dtype=kv_cache_dtype, + kv_cache_device=kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn=create_chunked_prefix_cache_kv_indices_fn, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = False, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ): + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.language_model, + data_embedding_funcs={ + Modality.IMAGE: self.get_image_feature, + }, + positions=positions, + pp_proxy_tensors=pp_proxy_tensors, + ) + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + mapper = getattr(self, "hf_to_sglang_mapper", None) + if mapper is not None: + weights = mapper.apply(weights) + + vision_params = ( + None + if self.config.language_only + else dict(self.named_parameters(remove_duplicate=False)) + ) + + def stream_language_weights(): + for name, loaded_weight in weights: + if "vision_tower" in name or "mm_projector" in name: + if vision_params is None: + continue + if name not in vision_params: + logger.warning("Unmapped vision weight: %s", name) + continue + param = vision_params[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + continue + yield name.replace("language_model.", ""), loaded_weight + + if self.language_model is not None: + self.language_model.load_weights(stream_language_weights()) + else: + # The vision weights are loaded as a side effect of advancing this + # streaming iterator. Encoder-only mode must therefore drain it + # even though it discards every language-model tensor. + for _ in stream_language_weights(): + pass + + +EntryClass = [KimiK3ForConditionalGeneration] diff --git a/python/sglang/srt/models/kimi_k3_vl.py b/python/sglang/srt/models/kimi_k3_vl.py new file mode 100644 index 000000000..a1c37bbf8 --- /dev/null +++ b/python/sglang/srt/models/kimi_k3_vl.py @@ -0,0 +1,937 @@ +"""Kimi K3 vision tower (MoonViT3d) and projector. + +Faithful port of the checkpoint reference implementation +(modeling_kimi_k3.py). Dedicated to K3 — do not share with Kimi K2.5: +K3 uses qkv_hidden_size != hidden_size (head_dim = qkv_hidden_size // +num_heads), RMSNorm encoder norms, bias-free linears, and the +PatchMergerMLPV2 projector (no pre-norm, post RMSNorm), all of which +differ from the K2.5 vision code. + +Weight names match the checkpoint exactly (wqkv/wo, mlp.fc0/fc1, +mm_projector.proj.0/proj.2, mm_projector.post_norm), so loading needs no +renames. +""" + +from dataclasses import dataclass, replace +from typing import List, Optional, Sequence, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.kernels.ops.attention.vision_rope import ( + apply_fused_qk_complex_rope, + precompile_fused_qk_complex_rope, +) +from sglang.srt.environ import envs +from sglang.srt.layers.attention.vision import ( + FLASHINFER_WORKSPACE_SIZE_BYTES, + QKV_BACKEND_IMPL, + VisionAttentionMetadata, + prepare_flashinfer_cudnn_vision_attention_metadata, + prepare_vision_attention_metadata, +) +from sglang.srt.models.kimi_vl_moonvit import tpool_patch_merger +from sglang.srt.multimodal.mm_utils import concat_or_single +from sglang.srt.runtime_context import get_server_args +from sglang.srt.utils import get_bool_env_var, is_hip, print_info_once + +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _use_aiter: + from aiter.ops.triton.conv.conv2d import conv2d as aiter_conv2d + +_SM103_TRITON_MAX_SEQLEN = 1536 +_SM103_FA4_MIN_ATTENTION_WORK = 3_000_000 + +GridTHW = Tuple[int, int, int] +SegmentBounds = Tuple[Tuple[int, int], ...] + + +def _resolve_grid_thw_list( + grid_thws: torch.Tensor, grid_thw_list: Optional[Sequence[Sequence[int]]] = None +) -> Tuple[GridTHW, ...]: + values = grid_thws.tolist() if grid_thw_list is None else grid_thw_list + return tuple((int(t), int(h), int(w)) for t, h, w in values) + + +def _get_mm_attention_backend() -> str: + try: + server_args = get_server_args() + except ValueError: + return "auto" + return server_args.mm_attention_backend or "auto" + + +def _is_fa4_available() -> bool: + try: + from sglang.kernels.ops.attention.flash_attention_v4 import ( + is_flash_attention_v4_available, + ) + except ImportError: + return False + return is_flash_attention_v4_available() + + +def _resolve_mm_attention_backend( + configured_backend: str, + *, + max_seqlen: int, + total_tokens: int, + device: torch.device, + fa4_available: Optional[bool] = None, +) -> str: + if configured_backend != "auto": + return configured_backend + if device.type != "cuda": + return "sdpa" + if torch.cuda.get_device_capability(device) != (10, 3): + return "sdpa" + + use_fa4 = ( + max_seqlen > _SM103_TRITON_MAX_SEQLEN + or max_seqlen * total_tokens >= _SM103_FA4_MIN_ATTENTION_WORK + ) + if use_fa4: + if fa4_available is None: + fa4_available = _is_fa4_available() + return "fa4" if fa4_available else "sdpa" + return "triton_attn" + + +def apply_rope( + xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + freqs_cis = freqs_cis.unsqueeze(-2) + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xk.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def _can_use_fused_rope(hidden_states: torch.Tensor, freqs_cis: torch.Tensor) -> bool: + return _can_use_fused_rope_for_shape( + dtype=hidden_states.dtype, + device=hidden_states.device, + freqs_cis=freqs_cis, + ) + + +def _can_use_fused_rope_for_shape( + *, + dtype: torch.dtype, + device: torch.device, + freqs_cis: torch.Tensor, +) -> bool: + if not ( + device.type == "cuda" + and freqs_cis.is_cuda + and device == freqs_cis.device + and dtype in (torch.bfloat16, torch.float16) + and freqs_cis.dtype == torch.complex64 + ): + return False + major, _ = torch.cuda.get_device_capability(device) + return major >= 9 + + +def sdpa_varlen_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: torch.Tensor, + *, + segment_bounds: Optional[SegmentBounds] = None, +) -> torch.Tensor: + is_single_segment = ( + len(segment_bounds) == 1 + if segment_bounds is not None + else cu_seqlens.numel() == 2 + ) + if is_single_segment: + out = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), + ) + return out.squeeze(0).transpose(0, 1).flatten(start_dim=-2) + outputs = [] + if segment_bounds is None: + bounds = cu_seqlens.tolist() + segment_bounds = tuple(zip(bounds[:-1], bounds[1:])) + for start, end in segment_bounds: + seg_q = q[start:end].transpose(0, 1).unsqueeze(0) + seg_k = k[start:end].transpose(0, 1).unsqueeze(0) + seg_v = v[start:end].transpose(0, 1).unsqueeze(0) + out = F.scaled_dot_product_attention(seg_q, seg_k, seg_v) + outputs.append(out.squeeze(0).transpose(0, 1)) + return torch.cat(outputs).flatten(start_dim=-2) + + +def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray: + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega + + pos = pos.reshape(-1) + out = np.einsum("m,d->md", pos, omega) + + emb_sin = np.sin(out) + emb_cos = np.cos(out) + return np.concatenate([emb_sin, emb_cos], axis=1) + + +def get_1d_sincos_pos_embed(embed_dim: int, t_size: int) -> np.ndarray: + grid_t = np.arange(t_size, dtype=np.float32) + return get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) + + +def interpolate_pos_emb( + weight: torch.Tensor, interpolation_mode: str, shape: Tuple[int, int] +) -> torch.Tensor: + return ( + F.interpolate( + weight.permute((2, 0, 1)).contiguous().unsqueeze(0), + size=shape, + mode=interpolation_mode, + ) + .squeeze(0) + .permute((1, 2, 0)) + .flatten(end_dim=1) + ) + + +class Learnable2DInterpPosEmbDividedFixed(nn.Module): + def __init__( + self, + height: int, + width: int, + num_frames: int, + dim: int, + interpolation_mode: str = "bicubic", + ) -> None: + super().__init__() + self.height = height + self.width = width + self.num_frames = num_frames + self.dim = dim + self.interpolation_mode = interpolation_mode + self.weight = nn.Parameter(torch.empty(height, width, dim)) + self.register_buffer( + "time_weight", + torch.from_numpy(get_1d_sincos_pos_embed(dim, num_frames)) + .float() + .unsqueeze(1), + persistent=False, + ) + + def position_embeddings( + self, + grid_thws: torch.Tensor, + *, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + ) -> torch.Tensor: + pos_embs = [] + for t, h, w in _resolve_grid_thw_list(grid_thws, grid_thw_list): + assert t <= self.num_frames, f"t:{t} > num_frames:{self.num_frames}" + if (h, w) == self.weight.shape[:-1]: + pos_emb_2d = self.weight.flatten(end_dim=1) + else: + pos_emb_2d = interpolate_pos_emb( + self.weight, self.interpolation_mode, (h, w) + ) + + if t == 1: + pos_emb_3d = pos_emb_2d + else: + pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[ + 0:t + ].to(pos_emb_2d.dtype) + + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) + + return torch.cat(pos_embs) + + def forward( + self, + x: torch.Tensor, + grid_thws: torch.Tensor, + *, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + position_embeddings: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if position_embeddings is None: + position_embeddings = self.position_embeddings( + grid_thws, grid_thw_list=grid_thw_list + ) + return x + position_embeddings + + +class MoonVision3dPatchEmbed(nn.Module): + def __init__( + self, + out_dim: int, + in_dim: int = 3, + patch_size: Union[int, Tuple[int, int]] = (14, 14), + pos_emb_height: int = 14, + pos_emb_width: int = 14, + pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + pos_emb_interpolation_mode: str = "bicubic", + patch_embed_proj_bias: bool = True, + ): + super().__init__() + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + self.patch_size = patch_size + + self.proj = nn.Conv2d( + in_dim, + out_dim, + kernel_size=patch_size, + stride=patch_size, + bias=patch_embed_proj_bias, + ) + + if pos_emb_type != "divided_fixed": + raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") + self.pos_emb = Learnable2DInterpPosEmbDividedFixed( + height=pos_emb_height, + width=pos_emb_width, + num_frames=pos_emb_time, + dim=out_dim, + interpolation_mode=pos_emb_interpolation_mode, + ) + + def forward( + self, + x: torch.Tensor, + grid_thws: torch.Tensor, + *, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + position_embeddings: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # MIOpen can overflow grid_size for some patch shapes. Prefer AITER's + # Triton convolution on AMD, with an equivalent linear fallback. + if _use_aiter: + x = aiter_conv2d( + x, + self.proj.weight, + self.proj.bias, + stride=self.patch_size, + padding=(0, 0), + dilation=(1, 1), + ).view(x.size(0), -1) + elif _is_hip: + x = F.linear(x.flatten(1), self.proj.weight.flatten(1), self.proj.bias) + else: + x = self.proj(x).view(x.size(0), -1) + return self.pos_emb( + x, + grid_thws, + grid_thw_list=grid_thw_list, + position_embeddings=position_embeddings, + ) + + +class Rope2DPosEmbRepeated(nn.Module): + def __init__( + self, dim: int, max_height: int, max_width: int, theta_base: float = 10000 + ): + super().__init__() + assert dim % 4 == 0, "dim must be divisible by 4" + self.dim = dim + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) + freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() + y_freqs = torch.outer(y_pos, freqs).float() + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) + freqs_cis = torch.cat( + [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 + ) + return freqs_cis.reshape(self.max_height, self.max_width, -1) + + def get_freqs_cis( + self, + grid_thws: torch.Tensor, + device: torch.device, + *, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + ) -> torch.Tensor: + if not hasattr(self, "freqs_cis"): + self.register_buffer( + "freqs_cis", self._precompute_freqs_cis(device), persistent=False + ) + + shapes = _resolve_grid_thw_list(grid_thws, grid_thw_list) + assert all( + 1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes + ), (shapes, self.max_height, self.max_width) + return torch.cat( + [ + self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) + for t, h, w in shapes + ], + dim=0, + ) + + +class MLP2(nn.Module): + def __init__(self, dims: List[int], activation, bias: bool = True): + super().__init__() + assert len(dims) == 3 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + self.activation = activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc1(self.activation(self.fc0(x))) + + +def _make_norm(norm_type: str, dim: int) -> nn.Module: + if norm_type == "layernorm": + return nn.LayerNorm(dim) + if norm_type == "rmsnorm": + return nn.RMSNorm(dim) + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + +class MoonViTEncoderLayer(nn.Module): + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + qkv_hidden_size: Optional[int] = None, + norm_type: str = "layernorm", + *, + activation=F.gelu, + attn_bias: bool = False, + linear_bias: bool = True, + attention_backend: str = "sdpa", + attention_workspace: Optional[torch.Tensor] = None, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + self.qkv_hidden_size = ( + hidden_dim if qkv_hidden_size is None else qkv_hidden_size + ) + self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads + + self.norm0 = _make_norm(norm_type, hidden_dim) + self.norm1 = _make_norm(norm_type, hidden_dim) + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation, bias=linear_bias) + self.wqkv = nn.Linear(hidden_dim, self.qkv_hidden_size * 3, bias=attn_bias) + self.wo = nn.Linear(self.qkv_hidden_size, hidden_dim, bias=attn_bias) + self.attention_backend = attention_backend + if attention_backend == "auto": + if not torch.cuda.is_available(): + implementation_backends = () + elif _is_hip: + implementation_backends = ("triton_attn",) + else: + implementation_backends = ("triton_attn", "fa4") + elif attention_backend == "sdpa": + implementation_backends = () + else: + implementation_backends = (attention_backend,) + self.attention_backend_impls = nn.ModuleDict( + { + backend: QKV_BACKEND_IMPL[backend]( + use_data_parallel=True, + workspace_buffer=attention_workspace, + ) + for backend in implementation_backends + } + ) + + def _attention( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + segment_bounds: SegmentBounds, + rope_freqs_cis: torch.Tensor, + forward_metadata: VisionAttentionMetadata, + use_fused_rope: bool, + selected_attention_backend: str, + ) -> torch.Tensor: + xqkv = self.wqkv(x) + qkv_shape = xqkv.size()[:-1] + ( + 3, + self.num_heads, + self.hidden_size_per_attention_head, + ) + xqkv = xqkv.view(*qkv_shape) + xq, xk, xv = torch.unbind(xqkv, dim=-3) + + if use_fused_rope: + xq, xk = apply_fused_qk_complex_rope(xq, xk, rope_freqs_cis) + else: + xq, xk = apply_rope(xq, xk, rope_freqs_cis) + + if selected_attention_backend == "sdpa": + attn_out = sdpa_varlen_attention( + xq, + xk, + xv, + cu_seqlens, + segment_bounds=segment_bounds, + ) + else: + if selected_attention_backend == "flashinfer_cudnn": + xv = xv.contiguous() + attn_out = self.attention_backend_impls[selected_attention_backend]( + xq, + xk, + xv, + cu_seqlens=cu_seqlens, + bsz=1, + seq_len=xq.shape[0], + forward_metadata=forward_metadata, + ).flatten(start_dim=-2) + return self.wo(attn_out) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + segment_bounds: SegmentBounds, + rope_freqs_cis: torch.Tensor, + forward_metadata: VisionAttentionMetadata, + use_fused_rope: bool, + selected_attention_backend: str, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm0(hidden_states) + hidden_states = self._attention( + hidden_states, + cu_seqlens, + segment_bounds, + rope_freqs_cis, + forward_metadata, + use_fused_rope, + selected_attention_backend, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +@dataclass(frozen=True) +class KimiK3VisionForwardMetadata: + grid_thw_list: Tuple[GridTHW, ...] + segment_bounds: SegmentBounds + rope_freqs_cis: torch.Tensor + attention: VisionAttentionMetadata + use_fused_rope: bool + selected_attention_backend: str + position_embeddings: Optional[torch.Tensor] = None + + +class MoonViT3dEncoder(nn.Module): + def __init__(self, hidden_dim: int, num_layers: int, block_cfg: dict) -> None: + super().__init__() + qkv_hidden_size = block_cfg.get("qkv_hidden_size") or block_cfg["hidden_dim"] + attention_backend = _get_mm_attention_backend() + if attention_backend != "auto" and attention_backend not in QKV_BACKEND_IMPL: + raise ValueError( + f"Unsupported Kimi-K3 vision attention backend: {attention_backend}" + ) + attention_workspace = None + if attention_backend == "flashinfer_cudnn" and torch.cuda.is_available(): + attention_workspace = torch.empty( + FLASHINFER_WORKSPACE_SIZE_BYTES, + dtype=torch.uint8, + device=torch.device("cuda", torch.cuda.current_device()), + ) + self.attention_backend = attention_backend + if attention_backend == "auto": + print_info_once( + "Kimi-K3 vision attention uses shape-aware auto selection on " + "B300/GB300 (Triton for small workloads, FA4 otherwise)." + ) + self.attention_width = qkv_hidden_size + self.rope_2d = Rope2DPosEmbRepeated( + qkv_hidden_size // block_cfg["num_heads"], 512, 512 + ) + self.blocks = nn.ModuleList( + [ + MoonViTEncoderLayer( + **block_cfg, + attention_backend=attention_backend, + attention_workspace=attention_workspace, + ) + for _ in range(num_layers) + ] + ) + self.final_layernorm = _make_norm( + block_cfg.get("norm_type", "layernorm"), hidden_dim + ) + + def precompile_fused_rope(self, dtype: torch.dtype, device: torch.device) -> bool: + if not self.blocks: + return False + return precompile_fused_qk_complex_rope( + num_heads=self.blocks[0].num_heads, + head_dim=self.blocks[0].hidden_size_per_attention_head, + dtype=dtype, + device=device, + ) + + def precompile_attention_backend( + self, dtype: torch.dtype, device: torch.device + ) -> bool: + if ( + self.attention_backend != "auto" + or not self.blocks + or device.type != "cuda" + or torch.cuda.get_device_capability(device) != (10, 3) + or not _is_fa4_available() + ): + return False + + block = self.blocks[0] + if "fa4" not in block.attention_backend_impls: + return False + + num_tokens = 256 + packed_qkv = torch.zeros( + ( + num_tokens, + 3, + block.num_heads, + block.hidden_size_per_attention_head, + ), + dtype=dtype, + device=device, + ) + q = packed_qkv[:, 0].contiguous() + k = packed_qkv[:, 1].contiguous() + v = packed_qkv[:, 2] + cu_seqlens = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + metadata = prepare_vision_attention_metadata(cu_seqlens, device=device) + with torch.inference_mode(): + block.attention_backend_impls["fa4"]( + q, + k, + v, + cu_seqlens=cu_seqlens, + bsz=1, + seq_len=num_tokens, + forward_metadata=metadata, + ) + torch.cuda.synchronize(device) + return True + + def prepare_forward_metadata( + self, + *, + grid_thws: torch.Tensor, + total_tokens: int, + dtype: torch.dtype, + device: torch.device, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + ) -> KimiK3VisionForwardMetadata: + shapes = _resolve_grid_thw_list(grid_thws, grid_thw_list) + rope_freqs_cis = self.rope_2d.get_freqs_cis( + grid_thws=grid_thws, + device=device, + grid_thw_list=shapes, + ) + cumulative_lengths = [0] + for t, h, w in shapes: + cumulative_lengths.append(cumulative_lengths[-1] + t * h * w) + cu_seqlens = torch.tensor(cumulative_lengths, dtype=torch.int32, device=device) + + if self.attention_backend == "flashinfer_cudnn": + attention = prepare_flashinfer_cudnn_vision_attention_metadata( + cu_seqlens, + device=device, + elem_per_token=self.attention_width, + ) + else: + attention = prepare_vision_attention_metadata(cu_seqlens, device=device) + + selected_attention_backend = _resolve_mm_attention_backend( + self.attention_backend, + max_seqlen=attention.max_seqlen, + total_tokens=total_tokens, + device=device, + ) + return KimiK3VisionForwardMetadata( + grid_thw_list=shapes, + segment_bounds=tuple(zip(cumulative_lengths[:-1], cumulative_lengths[1:])), + rope_freqs_cis=rope_freqs_cis, + attention=attention, + use_fused_rope=_can_use_fused_rope_for_shape( + dtype=dtype, + device=device, + freqs_cis=rope_freqs_cis, + ), + selected_attention_backend=selected_attention_backend, + ) + + def forward( + self, + hidden_states: torch.Tensor, + grid_thws: torch.Tensor, + *, + forward_metadata: Optional[KimiK3VisionForwardMetadata] = None, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + ) -> torch.Tensor: + if forward_metadata is None: + forward_metadata = self.prepare_forward_metadata( + grid_thws=grid_thws, + total_tokens=hidden_states.shape[0], + dtype=hidden_states.dtype, + device=hidden_states.device, + grid_thw_list=grid_thw_list, + ) + forward_metadata = replace( + forward_metadata, + use_fused_rope=_can_use_fused_rope( + hidden_states, forward_metadata.rope_freqs_cis + ), + ) + attention = forward_metadata.attention + cu_seqlens = attention.cu_seqlens + rope_freqs_cis = forward_metadata.rope_freqs_cis + + for block in self.blocks: + hidden_states = block( + hidden_states, + cu_seqlens, + forward_metadata.segment_bounds, + rope_freqs_cis, + attention, + forward_metadata.use_fused_rope, + forward_metadata.selected_attention_backend, + ) + + return self.final_layernorm(hidden_states) + + +class KimiK3VisionTower(nn.Module): + def __init__(self, vision_config, **kwargs): + super().__init__() + config = vision_config + self.config = config + self.merge_kernel_size = tuple(config.merge_kernel_size) + self.patch_size = config.patch_size + self.merge_type = config.merge_type + if self.merge_type != "sd2_tpool": + raise NotImplementedError(f"Not support merge_type: {self.merge_type}") + + hidden_size = getattr(config, "vt_hidden_size", None) or config.hidden_size + num_heads = ( + getattr(config, "vt_num_attention_heads", None) + or config.num_attention_heads + ) + num_layers = ( + getattr(config, "vt_num_hidden_layers", None) or config.num_hidden_layers + ) + intermediate_size = ( + getattr(config, "vt_intermediate_size", None) or config.intermediate_size + ) + + self.patch_embed = MoonVision3dPatchEmbed( + out_dim=hidden_size, + patch_size=config.patch_size, + pos_emb_height=config.init_pos_emb_height, + pos_emb_width=config.init_pos_emb_width, + pos_emb_time=config.init_pos_emb_time, + pos_emb_type=config.pos_emb_type, + pos_emb_interpolation_mode=config.pos_emb_interpolation_mode, + patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True), + ) + + activation_func = getattr(config, "activation_func", "gelu_pytorch_tanh") + if activation_func == "gelu_pytorch_tanh": + activation = lambda x: F.gelu(x, approximate="tanh") + elif activation_func == "gelu": + activation = F.gelu + else: + raise NotImplementedError(f"Not support activation_func: {activation_func}") + + self.encoder = MoonViT3dEncoder( + hidden_dim=hidden_size, + num_layers=num_layers, + block_cfg={ + "num_heads": num_heads, + "hidden_dim": hidden_size, + "qkv_hidden_size": getattr(config, "qkv_hidden_size", None), + "mlp_dim": intermediate_size, + "norm_type": getattr(config, "norm_type", "layernorm"), + "activation": activation, + "attn_bias": getattr(config, "attn_bias", True), + "linear_bias": getattr(config, "linear_bias", True), + }, + ) + self.cuda_graph_runner = None + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.patch_embed.proj.weight.device + + def precompile_fused_rope(self) -> bool: + return self.encoder.precompile_fused_rope(self.dtype, self.device) + + def precompile_attention_backend(self) -> bool: + return self.encoder.precompile_attention_backend(self.dtype, self.device) + + def prepare_forward_metadata( + self, + grid_thws: torch.Tensor, + *, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + total_tokens: int, + dtype: Optional[torch.dtype] = None, + ) -> KimiK3VisionForwardMetadata: + metadata = self.encoder.prepare_forward_metadata( + grid_thws=grid_thws, + total_tokens=total_tokens, + dtype=self.dtype if dtype is None else dtype, + device=self.device, + grid_thw_list=grid_thw_list, + ) + return replace( + metadata, + position_embeddings=self.patch_embed.pos_emb.position_embeddings( + grid_thws, grid_thw_list=metadata.grid_thw_list + ), + ) + + def forward( + self, + pixel_values: torch.Tensor, + grid_thws: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + *, + grid_hw: Optional[torch.Tensor] = None, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + forward_metadata: Optional[KimiK3VisionForwardMetadata] = None, + ) -> List[torch.Tensor]: + # run_dp_sharded_mrope_vision_model calls rope_2d towers with + # grid_hw=/max_seqlen= keywords (#30878); K3 grids are (t, h, w) and the + # encoder derives its own varlen metadata, so max_seqlen is unused. + if grid_thws is None: + grid_thws = grid_hw + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, grid_thws.shape + if ( + forward_metadata is None + and pixel_values.is_cuda + and envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get() + ): + if grid_thw_list is None: + grid_thw_list = _resolve_grid_thw_list(grid_thws) + else: + grid_thw_list = _resolve_grid_thw_list(grid_thws, grid_thw_list) + if self.cuda_graph_runner is None: + from sglang.srt.multimodal.kimi_k3_vit_cuda_graph_runner import ( + KimiK3ViTCudaGraphRunner, + ) + + self.cuda_graph_runner = KimiK3ViTCudaGraphRunner( + self, + capacity=(envs.SGLANG_KIMI_K3_VIT_CUDA_GRAPH_CACHE_CAPACITY.get()), + min_hits=envs.SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MIN_HITS.get(), + max_seqlen=(envs.SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MAX_SEQLEN.get()), + ) + return self.cuda_graph_runner.run(pixel_values, grid_thws, grid_thw_list) + return self._forward_eager( + pixel_values, + grid_thws, + grid_thw_list=grid_thw_list, + forward_metadata=forward_metadata, + ) + + def _forward_eager( + self, + pixel_values: torch.Tensor, + grid_thws: torch.Tensor, + *, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, + forward_metadata: Optional[KimiK3VisionForwardMetadata] = None, + ) -> List[torch.Tensor]: + if forward_metadata is not None: + grid_thw_list = forward_metadata.grid_thw_list + hidden_states = self.patch_embed( + pixel_values, + grid_thws, + grid_thw_list=grid_thw_list, + position_embeddings=( + None + if forward_metadata is None + else forward_metadata.position_embeddings + ), + ) + hidden_states = self.encoder( + hidden_states, + grid_thws, + forward_metadata=forward_metadata, + grid_thw_list=grid_thw_list, + ) + return tpool_patch_merger( + hidden_states, + grid_thws, + merge_kernel_size=self.merge_kernel_size, + grid_thw_list=grid_thw_list, + ) + + +class KimiK3MultiModalProjector(nn.Module): + """PatchMergerMLPV2: bias-free two-layer MLP over merged patches with a + post RMSNorm; K3 has no pre-norm, unlike the K2.5 projector.""" + + def __init__(self, vision_config): + super().__init__() + config = vision_config + mm_hidden_size = ( + getattr(config, "mm_hidden_size", None) or config.vt_hidden_size + ) + merge_h, merge_w = config.merge_kernel_size + self.hidden_size = mm_hidden_size * merge_h * merge_w + text_hidden_size = getattr(config, "text_hidden_size", None) or getattr( + config, "hidden_size" + ) + eps = config.projector_ln_eps + + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size, bias=False), + nn.GELU(), + nn.Linear(self.hidden_size, text_hidden_size, bias=False), + ) + self.post_norm = nn.RMSNorm(text_hidden_size, eps=eps) + + def forward( + self, image_features: Union[torch.Tensor, List[torch.Tensor]] + ) -> torch.Tensor: + if isinstance(image_features, (list, tuple)): + x = concat_or_single( + [item.reshape(item.shape[0], -1) for item in image_features] + ) + else: + x = image_features.reshape(image_features.shape[0], -1) + return self.post_norm(self.proj(x)) diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index c2dd7895c..b8fc10636 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -113,6 +113,9 @@ class KimiMoE(nn.Module): layer_id=self.layer_idx, quant_config=quant_config, routed_scaling_factor=self.routed_scaling_factor, + activation=config.hidden_act, + gemm1_alpha=config.activation_situ_beta, + gemm1_clamp_limit=config.activation_situ_linear_beta, prefix=add_prefix("experts", prefix), ) @@ -405,7 +408,10 @@ class KimiDeltaAttention(nn.Module): forget_gate = forget_gate.unflatten( -1, (-1, self.head_dim) ) # [T, H*K] -> [T, H, K] - beta = beta.float().sigmoid() + if not forward_batch.forward_mode.is_target_verify(): + # Only chunk_kda (extend) wants pre-activated beta; the verify + # kernel sigmoids it in-kernel like decode. + beta = beta.float().sigmoid() forget_gate = forget_gate.unsqueeze(0) beta = beta.unsqueeze(0) @@ -866,6 +872,12 @@ class KimiLinearForCausalLM(nn.Module): weight_loader(param, loaded_weight, **kwargs) loaded_params.add(name) + self.post_load_weights() + + def post_load_weights(self): + # Derive the absorbed MLA weights. Also called by the dummy loader + # (`_post_load_weights`), which never runs `load_weights` — keeping the + # derivation only there left `w_kc` None under --load-format dummy. for layer_id in self.config.full_attention_layer_ids: if not self.model.start_layer <= layer_id < self.model.end_layer: continue diff --git a/python/sglang/srt/models/kimi_vl_moonvit.py b/python/sglang/srt/models/kimi_vl_moonvit.py index c4adf0a07..ac4276584 100644 --- a/python/sglang/srt/models/kimi_vl_moonvit.py +++ b/python/sglang/srt/models/kimi_vl_moonvit.py @@ -570,13 +570,21 @@ def tpool_patch_merger( x: torch.Tensor, grid_thws: torch.Tensor, merge_kernel_size: tuple[int, int] = (2, 2), + *, + grid_thw_list: Optional[Sequence[Sequence[int]]] = None, ) -> List[torch.Tensor]: - """Group spatial patches and average only across real video frames.""" + """Group spatial patches and average only across real video frames. + + ``grid_thw_list`` lets a graph-aware tower pass the host grid it already + has instead of paying a device sync for ``grid_thws.tolist()``. + """ d_model = x.size(-1) outputs = [] pre_sum = 0 - for t, h, w in grid_thws.tolist(): + shapes = grid_thws.tolist() if grid_thw_list is None else grid_thw_list + for t, h, w in shapes: + t, h, w = int(t), int(h), int(w) seq = x[pre_sum : pre_sum + t * h * w] kernel_height, kernel_width = merge_kernel_size new_height, new_width = h // kernel_height, w // kernel_width diff --git a/python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py b/python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py new file mode 100644 index 000000000..3002a7578 --- /dev/null +++ b/python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py @@ -0,0 +1,211 @@ +"""Bounded full-tower CUDA graph runner for Kimi K3 vision.""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Hashable, List, Tuple + +import torch + +from sglang.srt.model_executor.runner_utils.pool import ( + get_or_create_global_graph_memory_pool, +) + +if TYPE_CHECKING: + from sglang.srt.models.kimi_k3_vl import ( + GridTHW, + KimiK3VisionForwardMetadata, + KimiK3VisionTower, + ) + +logger = logging.getLogger(__name__) + + +@dataclass +class _CapturedVisionGraph: + graph: torch.cuda.CUDAGraph + input_buffer: torch.Tensor + outputs: Tuple[torch.Tensor, ...] + metadata: KimiK3VisionForwardMetadata + + +class KimiK3ViTCudaGraphRunner: + def __init__( + self, + tower: KimiK3VisionTower, + *, + capacity: int, + min_hits: int, + max_seqlen: int | None = None, + ) -> None: + if capacity <= 0: + raise ValueError(f"capacity must be positive, got {capacity}") + if min_hits <= 0: + raise ValueError(f"min_hits must be positive, got {min_hits}") + if max_seqlen is not None and max_seqlen <= 0: + raise ValueError(f"max_seqlen must be positive, got {max_seqlen}") + self.tower = tower + self.capacity = capacity + self.min_hits = min_hits + self.max_seqlen = max_seqlen + self.graphs: dict[Hashable, _CapturedVisionGraph] = {} + self.seen: OrderedDict[Hashable, int] = OrderedDict() + self.failed_keys: set[Hashable] = set() + self._graph_memory_pool: Any = None + self._capacity_logged = False + self._max_seqlen_logged = False + logger.info( + "Kimi-K3 ViT CUDA graph: capacity=%d min_hits=%d max_seqlen=%s", + capacity, + min_hits, + max_seqlen, + ) + + @staticmethod + def graph_key(grid_thw_list: Tuple[GridTHW, ...]) -> Hashable: + return grid_thw_list + + def _record_hit(self, key: Hashable) -> int: + count = self.seen.pop(key, 0) + 1 + self.seen[key] = count + seen_capacity = max(self.capacity * 8, self.capacity) + if len(self.seen) > seen_capacity: + self.seen.popitem(last=False) + return count + + def _run_eager( + self, + pixel_values: torch.Tensor, + grid_thws: torch.Tensor, + grid_thw_list: Tuple[GridTHW, ...], + ) -> tuple[List[torch.Tensor], KimiK3VisionForwardMetadata]: + metadata = self.tower.prepare_forward_metadata( + grid_thws, + grid_thw_list=grid_thw_list, + total_tokens=pixel_values.shape[0], + dtype=pixel_values.dtype, + ) + outputs = self.tower._forward_eager( + pixel_values, + grid_thws, + grid_thw_list=grid_thw_list, + forward_metadata=metadata, + ) + return outputs, metadata + + def _capture( + self, + key: Hashable, + pixel_values: torch.Tensor, + grid_thws: torch.Tensor, + grid_thw_list: Tuple[GridTHW, ...], + metadata: KimiK3VisionForwardMetadata, + ) -> _CapturedVisionGraph: + allocated_before = torch.cuda.memory_allocated(pixel_values.device) + reserved_before = torch.cuda.memory_reserved(pixel_values.device) + input_buffer = torch.empty_like(pixel_values) + input_buffer.copy_(pixel_values) + graph = torch.cuda.CUDAGraph() + if self._graph_memory_pool is None: + self._graph_memory_pool = get_or_create_global_graph_memory_pool(torch.cuda) + with torch.cuda.graph(graph, pool=self._graph_memory_pool): + outputs = self.tower._forward_eager( + input_buffer, + grid_thws, + grid_thw_list=grid_thw_list, + forward_metadata=metadata, + ) + torch.cuda.synchronize(pixel_values.device) + entry = _CapturedVisionGraph( + graph=graph, + input_buffer=input_buffer, + outputs=tuple(outputs), + metadata=metadata, + ) + position_embeddings = getattr(metadata, "position_embeddings", None) + position_cache_bytes = ( + position_embeddings.numel() * position_embeddings.element_size() + if position_embeddings is not None + else 0 + ) + logger.info( + "Captured Kimi-K3 ViT CUDA graph: key=%s cache=%d/%d " + "allocated_delta_mib=%.1f reserved_delta_mib=%.1f " + "position_cache_mib=%.1f allocated_total_mib=%.1f " + "reserved_total_mib=%.1f", + key, + len(self.graphs) + 1, + self.capacity, + (torch.cuda.memory_allocated(pixel_values.device) - allocated_before) + / 2**20, + (torch.cuda.memory_reserved(pixel_values.device) - reserved_before) / 2**20, + position_cache_bytes / 2**20, + torch.cuda.memory_allocated(pixel_values.device) / 2**20, + torch.cuda.memory_reserved(pixel_values.device) / 2**20, + ) + return entry + + def run( + self, + pixel_values: torch.Tensor, + grid_thws: torch.Tensor, + grid_thw_list: Tuple[GridTHW, ...], + ) -> List[torch.Tensor]: + key = self.graph_key(grid_thw_list) + entry = self.graphs.get(key) + if entry is not None: + entry.input_buffer.copy_(pixel_values) + entry.graph.replay() + return list(entry.outputs) + + max_seqlen = max(t * h * w for t, h, w in grid_thw_list) + if self.max_seqlen is not None and max_seqlen > self.max_seqlen: + if not self._max_seqlen_logged: + logger.info( + "Kimi-K3 ViT CUDA graph uses eager fallback above " + "max_seqlen=%d to avoid low-value captures and HBM use", + self.max_seqlen, + ) + self._max_seqlen_logged = True + outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list) + return outputs + + if key in self.failed_keys or self._record_hit(key) < self.min_hits: + outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list) + return outputs + + if len(self.graphs) >= self.capacity: + if not self._capacity_logged: + logger.warning( + "Kimi-K3 ViT CUDA graph cache reached capacity=%d; " + "new shapes use eager fallback to bound HBM", + self.capacity, + ) + self._capacity_logged = True + outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list) + return outputs + + metadata = self.tower.prepare_forward_metadata( + grid_thws, + grid_thw_list=grid_thw_list, + total_tokens=pixel_values.shape[0], + dtype=pixel_values.dtype, + ) + try: + entry = self._capture(key, pixel_values, grid_thws, grid_thw_list, metadata) + except Exception: + self.failed_keys.add(key) + logger.exception( + "Kimi-K3 ViT CUDA graph capture failed for key=%s; " + "using eager fallback", + key, + ) + outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list) + return outputs + + self.graphs[key] = entry + entry.input_buffer.copy_(pixel_values) + entry.graph.replay() + return list(entry.outputs) diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index 3c4b50fa1..d7548ba20 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/python/sglang/srt/multimodal/mm_utils.py @@ -557,9 +557,11 @@ def run_dp_sharded_mrope_vision_model( grid_thw_list: list, *, rope_type: Literal["rope_3d", "rope_2d", "rope_2d_packed"], + pool_temporal_dimension: bool = False, load_local_pixel_values: Optional[Callable[[list[int]], torch.Tensor]] = None, pixel_values_device: Optional[torch.device] = None, pixel_values_dtype: Optional[torch.dtype] = None, + pass_grid_thw_list: bool = False, ): """Run a vision model with data parallelism (DP) sharding. The function will shard the input image tensor on the @@ -576,6 +578,11 @@ def run_dp_sharded_mrope_vision_model( "rope_2d" for packed 2D rope outputs (e.g., Kimi-VL) "rope_2d_packed" for packed 2D rope outputs that accept ``grid_thws`` positionally (e.g., Kimi-K2.5/K2.7) + pool_temporal_dimension: Whether the vision model pools away the temporal + grid dimension. Its output length is then h * w divided by + the spatial merge area instead of t * h * w divided by it. + pass_grid_thw_list: Forward the existing host grid list to the vision + model so graph-aware towers do not materialize it from a CUDA tensor. Returns: torch.Tensor: Output image embeddings @@ -614,11 +621,13 @@ def run_dp_sharded_mrope_vision_model( device=pixel_values.device if rope_type == "rope_2d" else None, ) if rope_type == "rope_2d": - image_embeds = vision_model( - pixel_values, - grid_hw=grid_thw, - max_seqlen=max(math.prod(grid) for grid in grid_thw_list), - ) + kwargs = { + "grid_hw": grid_thw, + "max_seqlen": max(math.prod(grid) for grid in grid_thw_list), + } + if pass_grid_thw_list: + kwargs["grid_thw_list"] = grid_thw_list + image_embeds = vision_model(pixel_values, **kwargs) # MoonViT returns one tensor per image. The multi-GPU path below # already concatenates these tensors before returning, so keep the # TP=1 DP-encoder path on the same projector-facing contract. @@ -687,7 +696,9 @@ def run_dp_sharded_mrope_vision_model( ) output_tokens_per_image = [ - math.prod(grid) // embed_dim_reduction_factor for grid in grid_thw_list + math.prod(grid[1:] if pool_temporal_dimension else grid) + // embed_dim_reduction_factor + for grid in grid_thw_list ] grouped_output_lengths = [] assignment_offset = 0 @@ -717,11 +728,13 @@ def run_dp_sharded_mrope_vision_model( device=(pixel_values_local.device if rope_type == "rope_2d" else None), ) if rope_type == "rope_2d": - image_embeds_local = vision_model( - pixel_values_local, - grid_hw=local_grid_thw, - max_seqlen=max(math.prod(grid) for grid in local_grid_thw_list), - ) + kwargs = { + "grid_hw": local_grid_thw, + "max_seqlen": max(math.prod(grid) for grid in local_grid_thw_list), + } + if pass_grid_thw_list: + kwargs["grid_thw_list"] = local_grid_thw_list + image_embeds_local = vision_model(pixel_values_local, **kwargs) else: image_embeds_local = vision_model(pixel_values_local, local_grid_thw) if isinstance(image_embeds_local, list): diff --git a/python/sglang/srt/multimodal/processors/kimi_k25.py b/python/sglang/srt/multimodal/processors/kimi_k25.py index 4fb0ab375..34816acb7 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k25.py +++ b/python/sglang/srt/multimodal/processors/kimi_k25.py @@ -1,7 +1,7 @@ import math import re from collections import defaultdict -from typing import Dict, List, Union +from typing import Callable, Dict, List, Optional, Union import numpy as np import torch @@ -194,15 +194,24 @@ def _process_single_image( image_scale: torch.Tensor, image_bias: torch.Tensor, patch_size: int, + to_chw: Callable[[Union[torch.Tensor, Image.Image]], torch.Tensor] = _to_cuda_chw, + post_resize: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, ) -> torch.Tensor: - """Process a single image on GPU: resize -> pad -> normalize -> patchify.""" - image = _to_cuda_chw(image) + """Process a single image on GPU: resize -> pad -> normalize -> patchify. + + ``to_chw`` converts the input to a CUDA CHW tensor (a model may keep an + alpha channel here); ``post_resize`` runs on the resized ``(B, C, H, W)`` + batch before patchify (K3 composites transparent backgrounds there). + """ + image = to_chw(image) new_h, new_w = config["new_height"], config["new_width"] padded_h = new_h + config["pad_height"] padded_w = new_w + config["pad_width"] x = _resize_bicubic_if_needed(image.unsqueeze(0), new_h, new_w) + if post_resize is not None: + x = post_resize(x) return normalize_and_patchify( x, image_scale, image_bias, patch_size, padded_h, padded_w @@ -250,6 +259,8 @@ def _gpu_preprocess_images( image_scale: torch.Tensor, image_bias: torch.Tensor, patch_size: int, + to_chw: Callable[[Union[torch.Tensor, Image.Image]], torch.Tensor] = _to_cuda_chw, + post_resize: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """GPU preprocessing pipeline for a batch of images. @@ -278,21 +289,30 @@ def _gpu_preprocess_images( if len(group) == 1: idx, image, config = group[0] patches = _process_single_image( - image, config, image_scale, image_bias, patch_size + image, + config, + image_scale, + image_bias, + patch_size, + to_chw=to_chw, + post_resize=post_resize, ) all_patches[idx] = patches all_grids[idx] = _grid_thw_from_resize_config(config, patch_size) else: - indexed_images = [(idx, _to_cuda_chw(image)) for idx, image, _ in group] + indexed_images = [(idx, to_chw(image)) for idx, image, _ in group] # One NaViT target group can include several original resolutions. # Batch only source-compatible images, which removes redundant # bicubic launches for common multi-image requests without padding # random-size inputs to a larger source resolution. - batch = torch.cat( - _resize_images_by_source_shape(indexed_images, target_h, target_w), - dim=0, - ) + resized = _resize_images_by_source_shape(indexed_images, target_h, target_w) + if post_resize is not None: + # Before the concat: a hook may change the channel count (K3 + # composites RGBA onto a background and returns RGB), and mixed + # 3/4-channel sources cannot be concatenated first. + resized = [post_resize(part) for part in resized] + batch = torch.cat(resized, dim=0) T = 1 gh, gw = padded_h // patch_size, padded_w // patch_size diff --git a/python/sglang/srt/multimodal/processors/kimi_k3.py b/python/sglang/srt/multimodal/processors/kimi_k3.py new file mode 100644 index 000000000..cf9df0563 --- /dev/null +++ b/python/sglang/srt/multimodal/processors/kimi_k3.py @@ -0,0 +1,438 @@ +"""Kimi K3 multimodal processor. + +GPU image preprocessing dedicated to K3: unlike the K2.5 wrapper it keeps +the alpha channel through the bicubic resize and then composites RGBA +images onto the checkpoint-configured background +(``transparent_bg_config`` with ``transparent_bg_fill_stage == +"after_resize"`` in preprocessor_config.json), instead of dropping alpha +at load time. +""" + +import re +from typing import Dict, List, Union + +import numpy as np +import torch +from PIL import Image + +from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput +from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration +from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor as SGLangBaseProcessor, +) +from sglang.srt.multimodal.processors.base_processor import ( + MultimodalSpecialTokens, +) +from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin +from sglang.srt.multimodal.processors.kimi_k25 import ( + KimiGPUProcessorWrapper, + _get_image_dimensions, + _gpu_preprocess_images, + navit_resize_config, +) +from sglang.srt.utils.cuda_ipc_transport_utils import ( + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, +) + + +def _encode_k3_special_tokens(tokenizer, text: str) -> list[int]: + """Encode K3 control tokens without allowing them to be BPE-split.""" + try: + return list(tokenizer.encode(text, allowed_special="all")) + except TypeError: + # Keep the helper usable with lightweight tokenizer stubs in CPU tests. + return list(tokenizer.encode(text)) + + +def _expand_k3_image_prompt_token_ids( + input_ids: Union[List[int], torch.Tensor], + image_token_id: int, + image_token_counts: List[int], + image_sizes: List[tuple[int, int]], + tokenizer, +) -> torch.Tensor: + """Expand K3 image placeholders into the checkpoint's media contract. + + K3 requires each image feature span to be enclosed by its original uploaded + dimensions. The chat template deliberately emits one ``media_pad`` per + image; after decode, insert the surrounding control tokens and expand that + one placeholder to the NaViT feature count. + """ + if len(image_token_counts) != len(image_sizes): + raise ValueError("Expected one original size for each K3 image.") + + if isinstance(input_ids, torch.Tensor): + input_ids = input_ids.detach().flatten().cpu().numpy() + input_ids = np.asarray(input_ids, dtype=np.int64) + + placeholder_count = np.count_nonzero(input_ids == image_token_id) + if placeholder_count != len(image_token_counts): + raise ValueError( + f"Expected {len(image_token_counts)} image placeholder token(s), " + f"found {placeholder_count}." + ) + + output = [] + image_index = 0 + for token_id in input_ids: + if token_id != image_token_id: + output.append(int(token_id)) + continue + + width, height = image_sizes[image_index] + output.extend( + _encode_k3_special_tokens( + tokenizer, + f"<|media_begin|>image {width}x{height}<|media_content|>", + ) + ) + output.extend([image_token_id] * image_token_counts[image_index]) + output.extend(_encode_k3_special_tokens(tokenizer, "<|media_end|>")) + image_index += 1 + + return torch.tensor(output, dtype=torch.long).unsqueeze(0) + + +def _expand_k3_image_prompt_text( + input_text: str, + image_token: str, + image_token_counts: List[int], + image_sizes: List[tuple[int, int]], +) -> str: + """Render the K3 media framing for the CPU HF-processor fallback.""" + parts = input_text.split(image_token) + if len(parts) - 1 != len(image_token_counts): + raise ValueError( + f"Expected {len(image_token_counts)} image placeholder(s), " + f"found {len(parts) - 1}." + ) + + output = [parts[0]] + for image_token_count, (width, height), suffix in zip( + image_token_counts, image_sizes, parts[1:] + ): + output.extend( + ( + f"<|media_begin|>image {width}x{height}<|media_content|>", + image_token * image_token_count, + "<|media_end|>", + suffix, + ) + ) + return "".join(output) + + +def _k3_to_cuda_chw(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor: + if isinstance(image, Image.Image): + # The checkpoint's fill_transparent_bg_with() returns RGB-mode images + # untouched before it ever inspects the alpha bands, so an RGB image + # carrying a stray "transparency" info key must NOT be promoted to + # RGBA here. + has_alpha = image.mode != "RGB" and ( + "A" in image.getbands() or "transparency" in image.info + ) + arr = np.asarray(image.convert("RGBA" if has_alpha else "RGB")) + return torch.from_numpy(arr).permute(2, 0, 1).cuda() + + image = image.cuda() + if image.dim() == 2: + image = image.unsqueeze(0) + if image.shape[0] == 1: + image = image.repeat(3, 1, 1) + return image + + +def _chessboard_background( + height: int, width: int, cfg: dict, device: torch.device +) -> torch.Tensor: + square = cfg.get("chessboard_square_size", 16) + white = float(cfg.get("chessboard_white_value", 255)) + gray = float(cfg.get("chessboard_gray_value", 200)) + on_top_left = cfg.get("chessboard_square_on_top_left", True) + + ys = torch.arange(height, device=device) // square + xs = torch.arange(width, device=device) // square + parity = (ys.unsqueeze(1) + xs.unsqueeze(0)) % 2 + gray_parity = 1 if on_top_left else 0 + bg = torch.where(parity == gray_parity, gray, white) + return bg.unsqueeze(0).expand(3, height, width) + + +def _fill_transparent_bg(x: torch.Tensor, bg_cfg: Union[dict, None]) -> torch.Tensor: + """Composite a resized (1, 4, H, W) float image in [0, 255] onto the + configured background; 3-channel input passes through.""" + if x.shape[1] == 3: + return x + rgb = x[:, :3] + if bg_cfg is None: + return rgb + + _, _, height, width = x.shape + pattern = bg_cfg.get("pattern", "black") + if pattern == "chessboard": + bg = _chessboard_background(height, width, bg_cfg, x.device) + elif pattern == "white": + bg = torch.full((3, height, width), 255.0, device=x.device) + elif pattern == "black": + bg = torch.zeros(3, height, width, device=x.device) + elif pattern == "gray": + bg = torch.full((3, height, width), 128.0, device=x.device) + else: + raise ValueError(f"Invalid background pattern: {pattern}") + + alpha = (x[:, 3:4] / 255.0).clamp(0.0, 1.0) + # The checkpoint processor casts the composited float result back with + # numpy's astype(np.uint8), which truncates; floor matches that exactly + # (a composite of [0, 255] inputs is always non-negative). + return (alpha * rgb + (1.0 - alpha) * bg).clamp(0.0, 255.0).floor_() + + +class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper): + def __init__(self, *args, transparent_bg_config=None, **kwargs): + super().__init__(*args, **kwargs) + self._transparent_bg_config = transparent_bg_config + + def _prepare_input_ids( + self, input_text, resize_configs, original_input_ids, image_sizes + ): + image_token_counts = [config["num_tokens"] for config in resize_configs] + if original_input_ids is None: + original_input_ids = _encode_k3_special_tokens( + self._hf_processor.tokenizer, input_text + ) + return _expand_k3_image_prompt_token_ids( + original_input_ids, + self._image_token_id, + image_token_counts, + image_sizes, + self._hf_processor.tokenizer, + ) + + def __call__(self, text=None, images=None, **kwargs): + images = images or kwargs.pop("images", None) + original_input_ids = kwargs.pop("sglang_original_input_ids", None) + if images and torch.cuda.is_available(): + return self._gpu_call(text, images, original_input_ids) + return self._cpu_call(text, images, original_input_ids, **kwargs) + + def _gpu_call(self, text, images, original_input_ids=None): + input_text = text[0] if isinstance(text, list) else text + + resize_configs = [] + image_sizes = [] + for image in images: + w, h = _get_image_dimensions(image) + image_sizes.append((w, h)) + resize_configs.append( + navit_resize_config( + w, + h, + self._patch_size, + self._merge_kernel_size, + self._in_patch_limit, + self._patch_limit_on_one_side, + self._fixed_output_tokens, + ) + ) + + input_ids = self._prepare_input_ids( + input_text, resize_configs, original_input_ids, image_sizes + ) + + image_scale, image_bias = self._get_gpu_norm_tensors() + # Shared source-compatible batched pipeline (same as K2.5): RGBA + # inputs land in their own source-shape groups, and the + # transparent-background compositing runs on each resized batch + # before patchify -- identical order to the previous per-image path. + pixel_values, grid_thws = _gpu_preprocess_images( + images, + resize_configs, + image_scale, + image_bias, + self._patch_size, + to_chw=_k3_to_cuda_chw, + post_resize=lambda x: _fill_transparent_bg(x, self._transparent_bg_config), + ) + + return { + "input_ids": input_ids, + "pixel_values": pixel_values, + "image_grid_thw": grid_thws, + } + + def _cpu_call(self, text, images, original_input_ids=None, **kwargs): + """HF fallback with the same K3 media framing as the GPU path.""" + input_text = text[0] if isinstance(text, list) else text + if not images: + return self._hf_processor(text=[input_text], **kwargs) + + image_sizes = [_get_image_dimensions(image) for image in images] + image_token_counts = [ + self._hf_processor.media_processor.media_tokens_calculator( + {"type": "image", "image": image} + ) + for image in images + ] + expanded_text = _expand_k3_image_prompt_text( + input_text, + self._image_token, + image_token_counts, + image_sizes, + ) + kwargs["medias"] = [{"type": "image", "image": image} for image in images] + out = self._hf_processor(text=[expanded_text], **kwargs) + out["input_ids"] = self._prepare_input_ids( + input_text, + [{"num_tokens": count} for count in image_token_counts], + original_input_ids, + image_sizes, + ) + grid_thws = out.pop("grid_thws", None) + if grid_thws is not None: + out["image_grid_thw"] = grid_thws + return out + + +class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): + models = [KimiK3ForConditionalGeneration] + gpu_image_decode = True + prefer_tokenized_input = True + precompute_hash_before_cpu_transfer = True + auto_mm_processor_worker_num = 2 + auto_mm_io_worker_num = 16 + supports_mm_processor_concurrency = True + preserve_processor_input_ids = True + + def __init__(self, hf_config, server_args, _processor, *args, **kwargs): + mm_tokens = MultimodalSpecialTokens( + image_token="<|media_pad|>", + image_token_id=hf_config.media_placeholder_token_id, + image_token_regex=re.compile(r"(?:<\|media_pad\|>)+"), + ).build(_processor) + + media_proc_cfg = _processor.media_processor.media_proc_cfg + + processor = KimiK3GPUProcessorWrapper( + _processor, + image_token=mm_tokens.image_token, + image_token_id=mm_tokens.image_token_id, + patch_size=media_proc_cfg["patch_size"], + merge_kernel_size=media_proc_cfg["merge_kernel_size"], + in_patch_limit=media_proc_cfg["in_patch_limit"], + patch_limit_on_one_side=media_proc_cfg["patch_limit_on_one_side"], + fixed_output_tokens=media_proc_cfg.get("fixed_output_tokens"), + image_mean=media_proc_cfg["image_mean"], + image_std=media_proc_cfg["image_std"], + transparent_bg_config=media_proc_cfg.get("transparent_bg_config"), + ) + super().__init__(hf_config, server_args, processor, *args, **kwargs) + self.mm_tokens = mm_tokens + + async def process_mm_data_async( + self, + image_data: List[Union[str, bytes, Dict]], + input_text, + request_obj, + *args, + **kwargs, + ): + if getattr(request_obj, "video_data", None) or kwargs.get("audio_data"): + raise ValueError("Kimi-K3 supports image input only") + + expected_image_count = len(image_data or []) + placeholder_count = self.count_image_placeholders( + input_text, self.mm_tokens.image_token_id + ) + if placeholder_count is not None: + if placeholder_count != expected_image_count: + raise ValueError( + "Kimi image placeholders must map one-to-one to image data: " + f"expected {expected_image_count}, found {placeholder_count} token(s)" + ) + # Keep structural media tokens distinct from user text that happens to + # spell ``<|media_pad|>``. Decoding the whole prompt and matching the + # resulting string would lose that distinction and could bind an image + # to user-provided text instead of the renderer-inserted token. + base_output = await self.fast_load_mm_data( + prompt=input_text, + image_data=image_data, + multimodal_tokens=self.mm_tokens, + discard_alpha_channel=False, + # Unlike load_mm_data, fast_load_mm_data does not derive + # input_ids from the prompt. Without this the wrapper falls + # back to re-encoding the decoded string, which is the loss of + # the structural/user distinction described above. + input_ids=input_text, + ) + else: + base_output = await self.load_mm_data( + prompt=input_text, + image_data=image_data, + multimodal_tokens=self.mm_tokens, + discard_alpha_channel=False, + ) + + if len(base_output.images) != expected_image_count: + raise ValueError( + "Kimi image placeholders must map one-to-one to image data: " + f"expected {expected_image_count}, loaded {len(base_output.images)}" + ) + + mm_items, input_ids, _ = await self.process_and_combine_mm_data_async( + base_output, + self.mm_tokens, + sglang_original_input_ids=base_output.input_ids, + ) + + # K3's tower is unconditionally image-wise data-parallel (each image + # is consumed by exactly one TP rank), so keep IPC proxies lazy until + # that assignment is known: one tokenizer/scheduler crossing per + # image instead of one per rank. K2.5 gates this on + # --mm-enable-dp-encoder; K3 needs no flag. + if getattr(self, "use_cuda_ipc", False): + for item in mm_items: + item.model_specific_data[DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY] = ( + True + ) + + return MultimodalProcessorOutput( + input_ids=input_ids.tolist(), + mm_items=mm_items, + im_token_id=self.mm_tokens.image_token_id, + ) + + def get_mm_data(self, prompt, embeddings, **kwargs): + img_grid_thw = kwargs.get("img_grid_thw", None) + output = self._build_kimi_mm_data_from_grids( + prompt=prompt, + embeddings=embeddings, + image_token_id=self.mm_tokens.image_token_id, + img_grid_thw=img_grid_thw, + ) + image_sizes = kwargs.get("original_image_sizes") + if image_sizes is None: + return output + + counts = [self._num_image_tokens_from_grid(grid) for grid in img_grid_thw] + if len(image_sizes) != len(counts): + raise ValueError( + "Expected one original image size for each K3 encoder grid." + ) + output.input_ids = ( + _expand_k3_image_prompt_token_ids( + prompt, + self.mm_tokens.image_token_id, + counts, + [tuple(size) for size in image_sizes], + self._tokenizer, + ) + .flatten() + .tolist() + ) + + search_start = 0 + for item, count in zip(output.mm_items, counts): + start = output.input_ids.index(self.mm_tokens.image_token_id, search_start) + item.offsets = [(start, start + count - 1)] + search_start = start + count + return output diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 81186ee65..b94e35eb5 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -78,6 +78,7 @@ from sglang.srt.utils.common import ( is_hip, is_hopper_with_cuda_12_3, is_host_cpu_arm64, + is_mnnvl_fabric_device, is_mps, is_musa, is_no_spec_infer_or_topk_one, @@ -313,7 +314,23 @@ RL_ON_POLICY_TARGET_CHOICES = ["fsdp"] LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"] -ENCODER_TRANSFER_BACKEND_CHOICES = ["zmq_to_scheduler", "zmq_to_tokenizer", "mooncake"] +ENCODER_TRANSFER_BACKEND_CHOICES = [ + "auto", + "zmq_to_scheduler", + "zmq_to_tokenizer", + "mooncake", +] + + +def resolve_encoder_transfer_backend( + backend: str, model_arch: str, tp_size: int +) -> str: + if backend != "auto": + return backend + if model_arch == "KimiK3ForConditionalGeneration" and tp_size > 1: + return "zmq_to_tokenizer" + return "zmq_to_scheduler" + DSA_PREFILL_CP_SPLIT_CHOICES = ["in-seq-split", "round-robin-split"] NSA_PREFILL_CP_SPLIT_CHOICES = DSA_PREFILL_CP_SPLIT_CHOICES # deprecated alias @@ -348,7 +365,14 @@ MAMBA_RADIX_CACHE_STRATEGY_CHOICES = [ MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"] -LINEAR_ATTN_KERNEL_BACKEND_CHOICES = ["triton", "cutedsl", "flashinfer", "flashkda"] +LINEAR_ATTN_KERNEL_BACKEND_CHOICES = [ + "triton", + "cutedsl", + "flashinfer", + "flashkda", + "nvidia_kda", + "ptx_kda", +] # Allow external code to add more choices @@ -1866,7 +1890,12 @@ class ServerArgs: NS("exec.comm"), ] = False enable_symm_mem: A[ - bool, "Enable NCCL symmetric memory for fast collectives.", NS("exec.comm") + bool, + Arg( + help="Enable NCCL symmetric memory for fast collectives.", + resolvable=True, + ), + NS("exec.comm"), ] = False triton_attention_reduce_in_fp32: A[ bool, @@ -2097,6 +2126,7 @@ class ServerArgs: Arg( help="Attention backend for speculative decoding operations (both target verify and draft extend). Can be one of 'prefill' (default) or 'decode'.", choices=["prefill", "decode"], + resolvable=True, ), NS("spec"), ] = "prefill" @@ -2523,6 +2553,14 @@ class ServerArgs: ), NS("exec.mamba"), ] = None + linear_attn_verify_backend: A[ + Optional[str], + Arg( + help="Override the kernel backend for linear attention speculative target-verify. If not set, follows the decode backend (flashinfer decode -> flashinfer verify, otherwise triton). KDA supports triton, nv_cutedsl, and flashinfer verify backends.", + choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES + ["nv_cutedsl"], + ), + NS("exec.mamba"), + ] = None # ReplaySSM buffered output-only linear-attn decode (GDN + KDA): per-slot # ring + periodic flush to cut per-step HBM state traffic. enable_linear_replayssm: A[ @@ -2542,14 +2580,15 @@ class ServerArgs: "Ring-buffer length L for ReplaySSM linear-attn decode. The full recurrent state is flushed to HBM every L decode steps.", NS("exec.mamba"), ] = 16 - # ReplaySSM spec-verify (Part B of RFC #28511): GDN linear-chain target-verify - # via fold-every-commit instead of per-draft full-state snapshots -- the - # verify stores each draft step's raw inputs into a per-slot window and the - # commit replays the accepted prefix into the fp32 checkpoint. GDN only; - # linear-chain (topk <= 1) only. - enable_gdn_replayssm_spec: A[ + # ReplaySSM spec-verify (Part B of RFC #28511): linear-attn target-verify via + # fold-every-commit instead of per-draft full-state snapshots -- the verify + # stores each draft step's raw inputs into a per-slot window and the commit + # replays the accepted prefix into the fp32 checkpoint. GDN sizes the window + # to the draft maximum; KDA folds a (raw v, pre-norm k, gate, beta) ring of + # length --linear-replayssm-cache-len. Linear-chain (topk <= 1) only. + enable_linear_replayssm_spec: A[ bool, - "Enable the ReplaySSM GDN spec-verify (Part B of RFC #28511): fold-every-commit -- a per-slot raw-input window sized to the draft maximum replaces the recurrent verify's per-draft full-state snapshots. GDN only, linear-chain (--speculative-eagle-topk in {None, 1}) only.", + "Enable the ReplaySSM spec-verify: fold-every-commit -- a per-slot raw-input window replaces the recurrent verify's per-draft full-state snapshots. GDN or KDA hybrid linear-attn models, linear-chain (--speculative-eagle-topk in {None, 1}) only.", NS("exec.mamba"), ] = False @@ -2704,7 +2743,10 @@ class ServerArgs: mm_feature_transport: A[ Optional[Literal["cpu", "cuda_ipc"]], "Transport multimodal features through CPU memory or a bounded CUDA IPC pool. " - "The default is CPU transport; CUDA IPC reserves GPU memory on the base GPU.", + "Unset resolves automatically: single-node CUDA deployments (without " + "disaggregation) use cuda_ipc, everything else uses cpu. CUDA IPC reserves " + "SGLANG_MM_FEATURE_CACHE_MB (default 1024 MiB) on the base GPU and falls " + "back to CPU transport per tensor when the pool is full.", NS("mm"), ] = None keep_mm_feature_on_device: A[ @@ -3005,7 +3047,7 @@ class ServerArgs: encoder_transfer_backend: A[ str, Arg( - help="The backend for encoder disaggregation transfer. Default is zmq_to_scheduler.", + help="The backend for encoder disaggregation transfer. Auto selects a model- and TP-aware backend.", choices=ENCODER_TRANSFER_BACKEND_CHOICES, ), NS("disagg"), @@ -3550,8 +3592,6 @@ class ServerArgs: handle_speculative_decoding(self) - # Needs the draft-token count derived just above. - # Validate the CuteDSL A2A token budget now that num_tokens_per_req is final. self._validate_cutedsl_a2a_token_budget() @@ -3753,16 +3793,20 @@ class ServerArgs: ) def _handle_model_source_paths(self): - """Resolve model/tokenizer paths backed by remote object stores.""" - if is_runai_obj_uri(self.model_path): - ObjectStorageModel.download_and_get_path(self.model_path) - - if ( - self.tokenizer_path is not None - and is_runai_obj_uri(self.tokenizer_path) - and self.tokenizer_path != self.model_path + """Prepare metadata for model paths backed by remote object stores.""" + seen_paths = set() + for model_path in ( + self.model_path, + self.tokenizer_path, + self.speculative_draft_model_path, ): - ObjectStorageModel.download_and_get_path(self.tokenizer_path) + if ( + model_path is not None + and model_path not in seen_paths + and is_runai_obj_uri(model_path) + ): + ObjectStorageModel.download_and_get_path(model_path) + seen_paths.add(model_path) def _handle_pd_disaggregation(self): from sglang.srt.arg_groups.pd_disaggregation_hook import ( @@ -3772,9 +3816,6 @@ class ServerArgs: handle_pd_disaggregation(self) def _handle_dcp_validation(self): - # Decode context parallel (DCP) is currently implemented and validated - # only on AMD HIP/ROCm. Reject invalid or unverified configurations - # early instead of letting them fail deeper in model initialization. if self.dcp_size < 1: raise ValueError( "Decode context parallel size (--dcp-size / " @@ -3805,48 +3846,6 @@ class ServerArgs: "communication backend (it removes the head-dim Q all-gather); " f"got --dcp-comm-backend={self.dcp_comm_backend}." ) - if not self.dcp_size > 1: - return - if is_hip(): - return - elif is_cuda(): - if self.speculative_algorithm is not None: - model_arches = self.get_model_config().hf_config.architectures - decode_backend = self.decode_attention_backend or self.attention_backend - kimi_linear_dspark = ( - self.speculative_algorithm == "DSPARK" - and "KimiLinearForCausalLM" in model_arches - and self.speculative_attention_mode == "decode" - and decode_backend in ("tokenspeed_mla", "cutedsl_mla") - ) - if kimi_linear_dspark: - ragged_verify_mode = envs.SGLANG_RAGGED_VERIFY_MODE.get() - if ragged_verify_mode != "static": - raise ValueError( - "Kimi Linear DCP + DSPARK currently requires " - "SGLANG_RAGGED_VERIFY_MODE=static, but got " - f"{ragged_verify_mode!r}." - ) - else: - raise ValueError( - "Decode context parallel (--dcp-size / " - "--decode-context-parallel-size > 1) with speculative " - "decoding on CUDA is supported only for Kimi Linear + " - "DSPARK + --speculative-attention-mode decode + " - "tokenspeed_mla, or experimental cutedsl_mla, but got " - f"architectures={model_arches}, " - f"speculative_algorithm={self.speculative_algorithm!r}, " - "speculative_attention_mode=" - f"{self.speculative_attention_mode!r}, " - f"decode_attention_backend={decode_backend!r}." - ) - else: - raise ValueError( - "Decode context parallel (--dcp-size / " - "--decode-context-parallel-size > 1) is currently only " - f"supported on the AMD HIP platform, but got dcp_size=" - f"{self.dcp_size} on a non-HIP platform." - ) def _handle_load_balance_method(self): if self.disaggregation_mode not in ("null", "prefill", "decode"): @@ -4231,7 +4230,12 @@ class ServerArgs: self.cuda_graph_backend_prefill = Backend.FULL def _handle_cuda_graph_config(self): + from sglang.srt.arg_groups.kimi_k3_hook import disable_kimi_k3_symm_mem + self._parse_cuda_graph_config() + # Reads the resolved per-phase backends; must precede the compat rules + # below and _handle_gpu_memory_settings, which key off enable_symm_mem. + disable_kimi_k3_symm_mem(self) self._apply_cuda_graph_compatibility() self._apply_deepep_adjustments() self._apply_cuda_graph_disaggregation_roles() @@ -5088,6 +5092,18 @@ class ServerArgs: ) validate_declarations(self, self._resolved_overrides) + if model_arch in ( + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", + ): + from sglang.srt.arg_groups.kimi_k3_hook import ( + apply_kimi_k3_linear_attn_defaults, + apply_kimi_k3_spec_backend_defaults, + ) + + apply_kimi_k3_linear_attn_defaults(self) + apply_kimi_k3_spec_backend_defaults(self) + if model_arch in [ "DeepseekV4ForCausalLM", ]: @@ -5582,7 +5598,9 @@ class ServerArgs: "--mamba-radix-cache-strategy extra_buffer." ) algo = (view.speculative_algorithm or "").upper() - assert algo not in ("DFLASH", "DSPARK"), ( + # dspark verifies through prepare_mamba_track_for_verify (lazy plan + # wired); dflash bypasses that hook, so it stays unsupported. + assert algo != "DFLASH", ( f"extra_buffer_lazy unsupported with {view.speculative_algorithm}; " "use --mamba-radix-cache-strategy extra_buffer." ) @@ -6018,6 +6036,10 @@ class ServerArgs: self.linear_attn_decode_backend is None and is_sm100_supported() and self.mamba_ssm_dtype == "bfloat16" + # Stage 4: flashinfer's recurrent_kda compiles the state slot stride + # as a free int64, so it reads the page-major/unified envelope-strided + # state correctly — the unified-memory skip is no longer needed (the + # page-major gate now allows flashinfer for linear-attn decode). ): self.linear_attn_decode_backend = "flashinfer" logger.info( @@ -6059,6 +6081,21 @@ class ServerArgs: f"got {self.mamba_ssm_dtype!r}" ) + verify = self.linear_attn_verify_backend + if verify is None and decode == "flashinfer": + verify = "flashinfer" + if ( + verify == "flashinfer" + and self.mamba_ssm_dtype != "bfloat16" + and is_cuda() + and torch.cuda.get_device_capability()[0] >= 10 + ): + raise ValueError( + "--linear-attn-verify-backend flashinfer on SM100+ requires " + "--mamba-ssm-dtype bfloat16, " + f"got {self.mamba_ssm_dtype!r}" + ) + # SM100+ FlashInfer GDN prefill requires CUDA 13+ (CuTe DSL kernel) # for correctness and best performance. prefill = self.linear_attn_prefill_backend or self.linear_attn_backend @@ -6123,20 +6160,20 @@ class ServerArgs: f"{self.linear_replayssm_cache_len}." ) - # ReplaySSM spec-verify (Part B of #28511): GDN-only, linear-chain target - # verify. Reuses the `linear_replayssm` ring (replayssm_d/k/g + write_pos) - # plus two extra per-slot cursors (cache_base, is_flush) and the chunked - # (I+A)^-1 reconstruction verify kernel. The intra-window interaction uses a - # strictly-lower causal mask, so it is valid ONLY for a linear draft chain - # (speculative_eagle_topk in {None, 1}, i.e. NEXTN / MTP); EAGLE tree verify - # (topk > 1) must fall back to the recurrent verify. GDN-only is enforced at - # runtime (KDA routes through kda_backend, which never enters this path; the - # pool gate also checks `not cache_params.is_kda`). The ring length reuses - # --linear-replayssm-cache-len (no separate flag). - if self.enable_gdn_replayssm_spec: + # ReplaySSM spec-verify (Part B of #28511): linear-chain target verify via + # fold-every-commit -- the verify stores each draft step's raw inputs into + # the per-slot (rawv, rawk, g, beta) window and the commit replays the + # accepted prefix into the fp32 checkpoint. The intra-window interaction + # uses a strictly-lower causal mask, so it is valid ONLY for a linear + # draft chain (speculative_eagle_topk in {None, 1}, i.e. NEXTN / MTP); + # EAGLE tree verify (topk > 1) must fall back to the recurrent verify. + # GDN sizes the window to the draft maximum; KDA (kda_backend) keeps a + # --linear-replayssm-cache-len window and folds via its own fused + # verify ring-write + commit_kda_replayssm_after_verify. + if self.enable_linear_replayssm_spec: if self.speculative_eagle_topk not in (None, 1): raise ValueError( - "--enable-gdn-replayssm-spec requires a linear draft chain " + "--enable-linear-replayssm-spec requires a linear draft chain " "(--speculative-eagle-topk in {None, 1}); the chunked verify " "kernel uses a strictly-lower causal mask and is invalid for " "EAGLE tree verify. Got " @@ -6144,45 +6181,61 @@ class ServerArgs: ) if decode not in ("triton", "flashinfer"): raise ValueError( - "--enable-gdn-replayssm-spec requires the triton or " + "--enable-linear-replayssm-spec requires the triton or " "flashinfer linear-attn decode backend, got " f"--linear-attn-decode-backend={decode!r}." ) - # The auto->extra_buffer strategy resolution is still a declaration - # here, so read it through the resolved view. - view = self._resolved() - if ( - view.disable_radix_cache is False - and view.mamba_radix_cache_strategy == "extra_buffer_lazy" - ): - raise ValueError( - "--enable-gdn-replayssm-spec is not validated with " - "--mamba-radix-cache-strategy extra_buffer_lazy yet; " - "use extra_buffer." - ) + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) + + ragged_mode = read_ragged_verify_mode() + if ragged_mode is not RaggedVerifyMode.STATIC: + # Ragged ring-writes need the KDA fold-every-commit family + # (DSPARK/DFLASH) + the triton verify kernel (nv_cutedsl falls + # back to it for ragged layouts). The GDN ring-write kernels do + # not take the ragged layout and the flashinfer verify kernel + # never writes the ring -> a stale ring would be folded; keep + # refusing those combinations. + _algo = (self.speculative_algorithm or "").upper() + verify = self.linear_attn_verify_backend + if _algo not in ("DSPARK", "DFLASH") or verify not in ( + "triton", + "nv_cutedsl", + ): + raise ValueError( + "--enable-linear-replayssm-spec with " + f"SGLANG_RAGGED_VERIFY_MODE={ragged_mode.value} requires the " + "KDA fold-every-commit family (DSPARK/DFLASH) and a " + "ring-writing verify kernel (--linear-attn-verify-backend " + "triton or nv_cutedsl); got " + f"algorithm={self.speculative_algorithm!r}, " + f"verify={verify!r}. Use SGLANG_RAGGED_VERIFY_MODE=static." + ) if self.disaggregation_mode == "prefill": raise ValueError( - "--enable-gdn-replayssm-spec is not supported on a PD " + "--enable-linear-replayssm-spec is not supported on a PD " "prefill server: the ring is spec-verify-only scratch and " "the prefill server never runs spec verify." ) if self.enable_linear_replayssm: raise ValueError( - "--enable-gdn-replayssm-spec and --enable-linear-replayssm are " + "--enable-linear-replayssm-spec and --enable-linear-replayssm are " "mutually exclusive: they share the ring storage but drive it " "with incompatible cursor protocols (per-decode-forward vs " "per-verify-commit advance)." ) if self.mamba_ssm_dtype is None: logger.info( - "--enable-gdn-replayssm-spec: setting --mamba-ssm-dtype " + "--enable-linear-replayssm-spec: setting --mamba-ssm-dtype " "float32 (the closed-loop exact fold keeps the SSM checkpoint " "bit-identical to the recurrent baseline)." ) self.mamba_ssm_dtype = "float32" elif self.mamba_ssm_dtype != "float32": logger.warning( - "--enable-gdn-replayssm-spec with --mamba-ssm-dtype=%s: the " + "--enable-linear-replayssm-spec with --mamba-ssm-dtype=%s: the " "closed-loop fold re-quantizes the committed state each " "commit/flush (fp32 keeps it bit-exact to the fp32 recurrent " "baseline), so it may drift over long sequences. Validate " @@ -7212,6 +7265,13 @@ class ServerArgs: elif is_remote_url(self.model_path): self.load_format = "remote" + if ( + self.speculative_draft_model_path is not None + and is_runai_obj_uri(self.speculative_draft_model_path) + and self.speculative_draft_load_format is None + ): + self.speculative_draft_load_format = "runai_streamer" + if self.custom_weight_loader is None: self.custom_weight_loader = [] @@ -7361,6 +7421,17 @@ class ServerArgs: # Validate model type for encoder disaggregation hf_config = self.get_model_config().hf_config model_arch = hf_config.architectures[0] + if self.encoder_transfer_backend == "auto": + self.encoder_transfer_backend = resolve_encoder_transfer_backend( + self.encoder_transfer_backend, model_arch, self.tp_size + ) + if self.encoder_only or self.language_only: + logger.info( + "Encoder transfer backend auto-resolved to %s for %s at TP%d.", + self.encoder_transfer_backend, + model_arch, + self.tp_size, + ) if (self.encoder_only or self.language_only) and model_arch not in [ "Qwen2VLForConditionalGeneration", "Qwen3VLForConditionalGeneration", @@ -7374,6 +7445,7 @@ class ServerArgs: "Qwen2_5OmniForConditionalGeneration", "KimiVLForConditionalGeneration", "KimiK25ForConditionalGeneration", + "KimiK3ForConditionalGeneration", "MiMoV2ForCausalLM", ]: raise ValueError( @@ -7523,6 +7595,26 @@ class ServerArgs: "--mm-feature-transport=%s instead.", requested_transport, ) + elif self.encoder_only: + requested_transport = "cpu" + logger.info( + "Multimodal feature transport auto-resolved to cpu for " + "encoder-only serving; encoder outputs use " + "--encoder-transfer-backend instead." + ) + elif is_cuda() and self.nnodes == 1 and self.disaggregation_mode == "null": + # Auto policy: single-node CUDA serving defaults to the bounded + # CUDA-IPC pool. The pool is only allocated when a multimodal + # processor exists, so text-only deployments are unaffected; a + # full pool degrades to CPU transport per tensor. Multi-node + # (IPC handles are intra-node) and PD-disaggregated deployments + # keep CPU transport. + requested_transport = "cuda_ipc" + logger.info( + "Multimodal feature transport auto-resolved to cuda_ipc " + "(single-node CUDA). Pass --mm-feature-transport=cpu to " + "opt out." + ) else: requested_transport = "cpu" elif legacy_ipc_is_set and legacy_ipc_enabled != ( @@ -7535,6 +7627,14 @@ class ServerArgs: int(legacy_ipc_enabled), ) + if self.encoder_only and requested_transport == "cuda_ipc": + logger.warning( + "--mm-feature-transport=cuda_ipc does not control encoder-only " + "output transfer; using cpu for this inactive transport. Select " + "--encoder-transfer-backend for encoder outputs." + ) + requested_transport = "cpu" + if requested_transport == "cuda_ipc": if not is_cuda(): raise ValueError( @@ -7554,6 +7654,16 @@ class ServerArgs: self.base_gpu_id, self.tokenizer_worker_num, ) + logger.info( + "CUDA IPC pool-handle caching is %s. It reuses mappings to the " + "existing bounded pool without reserving another pool; set " + "SGLANG_USE_IPC_POOL_HANDLE_CACHE=0 to disable it.", + ( + "enabled" + if envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.get() + else "disabled" + ), + ) self.mm_feature_transport = requested_transport # The bounded IPC pool owns device residency. Do not retain unpooled @@ -7563,6 +7673,41 @@ class ServerArgs: "1" if requested_transport == "cuda_ipc" else "0" ) + def _handle_custom_all_reduce_v2_multinode(self): + # Custom all-reduce v2's graph zero-copy path uses IPC handles and is + # intra-node only. On MNNVL-fabric devices (GB200/GB300) the eager pull + # path works across nodes via the symm-mem workspace, so opt into the + # multinode mode automatically (a failed fabric rendezvous falls back + # to the legacy path at init). Elsewhere force-disable v2 on + # multi-node so the dispatch falls back to the legacy CustomAllreduce + # path, unless the MNNVL opt-in is set explicitly. + if self.nnodes <= 1 or not envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get(): + return + if ( + not envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.is_set() + and is_mnnvl_fabric_device() + # CustomAllReduceV2 supports world sizes 2..8 only + # (can_use_custom_all_reduce_v2 rejects larger groups); don't + # auto-opt-in a TP16+ launch just to fall back downstream. + and self.tp_size <= 8 + ): + logger.info( + "MNNVL fabric device detected with nnodes=%d: enabling " + "custom all-reduce v2 multinode mode " + "(SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE=1; set it " + "to 0 to opt out).", + self.nnodes, + ) + envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.set("1") + if not envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get(): + if envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.is_set(): + logger.warning( + "Disabling SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 because nnodes=%d " + "(custom all-reduce v2 is intra-node only).", + self.nnodes, + ) + envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.set("0") + def _handle_environment_variables(self): self._handle_multimodal_feature_transport() envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0") @@ -7574,6 +7719,7 @@ class ServerArgs: envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set( "1" if self.enable_deterministic_inference else "0" ) + self._handle_custom_all_reduce_v2_multinode() if self.debug_cuda_graph: if not (is_cuda() or is_hip()): logger.warning( @@ -7815,19 +7961,43 @@ class ServerArgs: f"paged MLA backends); got {sorted(backends)}, allowed " f"{sorted(allowed_full)}. Pass a compatible --attention-backend." ) - # The Mamba state is stored in envelope-strided views; only the - # stride-aware Triton causal-conv / SSM kernels read them correctly. - linear_backends = { - self.linear_attn_backend, - self.linear_attn_decode_backend, - self.linear_attn_prefill_backend, - self.mamba_backend, - } - linear_backends.discard(None) - assert linear_backends <= {"triton"}, ( - "--enable-page-major-kv-layout requires the Triton linear-attention / " - f"Mamba kernels for the strided conv/SSM state; got " - f"{sorted(linear_backends)}. Pass --linear-attn-backend triton and " + # The Mamba/KDA state is stored in envelope-strided views; only + # stride-audited kernels may read it (Stage 4 audit, per slot): + # - decode: triton; flashinfer (recurrent_kda compiles the state slot + # stride as a free int64 — natively strided); cutedsl (KDA fused + # sigmoid-gating update made stride-safe) on KDA-hybrid models only — + # cutedsl_gdn still compiles h0 against a contiguous dummy. + # - prefill: triton; flashkda (wrapper gathers/scatters a contiguous + # per-slot copy, external kernel never sees the pool); cutedsl + # (kernel_h compiles h0/ht with dynamic int64 strides), same + # KDA-only caveat. + # - mamba (mamba2/short-conv state): triton only. + # use_mla_backend() distinguishes the KDA-hybrid family (K3/KimiLinear + # are MLA-hybrid) from GDN models (GQA-hybrid) for the cutedsl caveat. + decode_allowed = {"triton", "flashinfer"} + prefill_allowed = {"triton", "flashkda"} + if self.use_mla_backend(): + decode_allowed.add("cutedsl") + prefill_allowed.add("cutedsl") + resolved_linear_decode = ( + self.linear_attn_decode_backend or self.linear_attn_backend + ) + resolved_linear_prefill = ( + self.linear_attn_prefill_backend or self.linear_attn_backend + ) + assert resolved_linear_decode in decode_allowed | {None}, ( + "--enable-page-major-kv-layout: linear-attention DECODE backend must " + f"be one of {sorted(decode_allowed)} for the strided conv/SSM state; " + f"got {resolved_linear_decode!r}." + ) + assert resolved_linear_prefill in prefill_allowed | {None}, ( + "--enable-page-major-kv-layout: linear-attention PREFILL backend must " + f"be one of {sorted(prefill_allowed)} for the strided conv/SSM state; " + f"got {resolved_linear_prefill!r}." + ) + assert self.mamba_backend in (None, "triton"), ( + "--enable-page-major-kv-layout requires the Triton Mamba kernels for " + f"the strided conv/SSM state; got {self.mamba_backend!r}. Pass " "--mamba-backend triton." ) @@ -8226,6 +8396,13 @@ class ServerArgs: new_flag="--enable-prefill-cp", help="[Deprecated] Use --enable-prefill-cp instead.", ) + parser.add_argument( + "--enable-gdn-replayssm-spec", + dest="enable_linear_replayssm_spec", + action=DeprecatedStoreTrueAction, + new_flag="--enable-linear-replayssm-spec", + help="[Deprecated] Use --enable-linear-replayssm-spec instead.", + ) parser.add_argument( "--enable-prefill-context-parallel", dest="enable_prefill_context_parallel", diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 8ec5420bc..6b2cc6ec4 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -105,6 +105,9 @@ class DFlashVerifyInput(SpecInput): bs = len(req_pool_indices) layout = self.ragged_verify_layout + if layout is not None and layout.bs != bs: + # Graph replay pads the batch to the captured slots; match it. + layout = layout.padded_to_bucket(padded_bs=bs) if layout is None: qo_indptr = torch.arange( diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 30b113167..d7de9bf49 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -13,7 +13,7 @@ from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.sampler import apply_custom_logit_processor from sglang.srt.managers.schedule_batch import Req from sglang.srt.speculative.spec_utils import _sample_simulated_acc_len -from sglang.srt.utils import is_cuda, is_musa +from sglang.srt.utils import is_cuda, is_hip, is_musa DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>" @@ -46,6 +46,15 @@ if is_cuda() or is_musa(): top_k_renorm_prob = None top_p_renorm_prob = None tree_speculative_sampling_target_only = None +elif is_hip(): + from sglang.kernels.ops.sampling.renorm_triton import ( + top_k_renorm_probs_triton as top_k_renorm_prob, + ) + from sglang.kernels.ops.sampling.renorm_triton import ( + top_p_renorm_probs_triton as top_p_renorm_prob, + ) + + _DFLASH_SAMPLING_VERIFY_AVAILABLE = True else: top_k_renorm_prob = None top_p_renorm_prob = None diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index abf643b71..2e1685da1 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -22,13 +22,16 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +# trtllm_mha: decode-only dense-MQA drafts (dspark). DFLASH excludes it +# earlier, at arg resolution (speculative_hook.py) -- its draft path needs +# per-layer DFlash attention -- so it never reaches this gate with it. _SUPPORTED_DRAFT_BACKENDS = ( "flashinfer", "fa3", "fa4", "triton", - "trtllm_mha", "ascend", + "trtllm_mha", ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 962ef5bd8..d5b1ec03b 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -92,8 +92,15 @@ class DSparkWorkerV2(BaseSpecWorker): self._draft_dp_context_enabled = ( server_args.enable_dp_attention and not self._draft_is_moe ) - attn_tp_size = server_args.tp_size // max(server_args.dp_size, 1) - if server_args.enable_dp_attention and self._draft_is_moe and attn_tp_size > 1: + self._is_pd_prefill = server_args.disaggregation_mode == "prefill" + self._decode_graph_allowed = ( + not server_args.disable_cuda_graph and not self._is_pd_prefill + ) + if ( + server_args.enable_dp_attention + and self._draft_is_moe + and ps.attn_tp_size > 1 + ): raise ValueError( "DSpark + dp attention with a DeepSeek-V4 (MoE) draft requires " "attn_tp == 1 (set --dp-size == --tp). attn_tp > 1 corrupts the " @@ -173,7 +180,7 @@ class DSparkWorkerV2(BaseSpecWorker): server_args.enable_dp_attention and not self._draft_is_moe and self._verify_planner.is_compact_mode - and not server_args.disable_cuda_graph + and self._decode_graph_allowed ): raise ValueError( "DSpark dense-draft compact verify under --enable-dp-attention does not " @@ -201,7 +208,7 @@ class DSparkWorkerV2(BaseSpecWorker): self._verify_epilogue = None if ( self._verify_planner.is_compact_mode - and not server_args.disable_cuda_graph + and self._decode_graph_allowed and is_cuda() ): self._verify_epilogue = DsparkVerifyEpilogue( @@ -262,6 +269,9 @@ class DSparkWorkerV2(BaseSpecWorker): simulate_acc_len=self._simulate_acc_len, ) + if self._is_pd_prefill and not self._draft_is_moe: + self.draft_model.prune_to_ctx_kv_injection() + def _resolve_target_embed_tokens(self, target_model): if hasattr(target_model, "get_input_embeddings"): return target_model.get_input_embeddings() @@ -311,7 +321,7 @@ class DSparkWorkerV2(BaseSpecWorker): ) def init_cuda_graphs(self): - capture_decode_cuda_graph = not get_exec().graph.disable_cuda_graph + capture_decode_cuda_graph = self._decode_graph_allowed if is_cuda() and capture_decode_cuda_graph: available_mem = get_available_gpu_memory(self.device, self.gpu_id) if available_mem < 1.0: @@ -420,6 +430,8 @@ class DSparkWorkerV2(BaseSpecWorker): if batch.out_cache_loc is None: raise RuntimeError("DSpark prefill expected out_cache_loc, but got None.") + # Must inject before prefill returns: the scheduler may update radix + # afterward, invalidating out_cache_loc. device = next_token_ids.device ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device) draft_seq_lens = torch.tensor( @@ -436,6 +448,7 @@ class DSparkWorkerV2(BaseSpecWorker): cache_loc=batch.out_cache_loc, positions=positions, ) + # Avoid copying large hidden-state buffers to CPU in overlap scheduling. logits_output.hidden_states = None batch_output.next_draft_input = make_next_draft_input( diff --git a/python/sglang/srt/speculative/dspark_disaggregation.py b/python/sglang/srt/speculative/dspark_disaggregation.py new file mode 100644 index 000000000..569f92c2b --- /dev/null +++ b/python/sglang/srt/speculative/dspark_disaggregation.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.managers.overlap_utils import RelayPayload +from sglang.srt.speculative.dspark_components.dspark_draft import make_next_draft_input + +if TYPE_CHECKING: + from sglang.srt.managers.overlap_utils import FutureMap + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.server_args import ServerArgs + from sglang.srt.speculative.spec_info import SpecInput + + +def build_dspark_disagg_draft_input( + batch: ScheduleBatch, + server_args: ServerArgs, + last_tokens_tensor: torch.Tensor, + future_map: FutureMap, +) -> SpecInput: + spec_info = make_next_draft_input( + bonus_tokens=last_tokens_tensor, + new_seq_lens=batch.seq_lens, + ) + if batch.enable_overlap: + spec_info.future_dsa_topk_indices_available = False + spec_info.future_indices = batch.req_pool_indices + future_map.publish(spec_info.future_indices, batch.seq_lens) + future_map.stash( + spec_info.future_indices, RelayPayload.from_draft_input(spec_info) + ) + return spec_info diff --git a/python/sglang/srt/speculative/ragged_verify.py b/python/sglang/srt/speculative/ragged_verify.py index a2906681c..92542e553 100644 --- a/python/sglang/srt/speculative/ragged_verify.py +++ b/python/sglang/srt/speculative/ragged_verify.py @@ -55,6 +55,9 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True): kv_lens_host: Optional[torch.Tensor] = None max_q_len: Optional[int] = None max_kv_len: Optional[int] = None + # Per-row upper bound (capped padded variant); rows never exceed it, so + # dense [bs, cap] consumers stay in bounds. None = full-coverage variant. + cap: Optional[int] = None def __post_init__(self) -> None: if self.verify_lens_cpu is None: @@ -66,6 +69,11 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True): f"every request must verify the anchor (verify_len >= 1), got " f"{self.verify_lens_cpu}" ) + if self.cap is not None and max(self.verify_lens_cpu) > self.cap: + raise ValueError( + f"capped layout has a row exceeding cap={self.cap}: " + f"{self.verify_lens_cpu}" + ) if self.total_verify_tokens != sum(self.verify_lens_cpu): raise ValueError( f"total_verify_tokens {self.total_verify_tokens} != " @@ -89,6 +97,7 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True): graph_num_tokens: int, verify_lens_cpu: Optional[list[int]] = None, total_verify_tokens: Optional[int] = None, + cap: Optional[int] = None, ) -> RaggedVerifyLayout: from sglang.kernels.ops.speculative.ragged_verify_kernels import ( BuildQoIndptr, @@ -103,6 +112,7 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True): qo_indptr_device=indptr.qo_indptr, verify_lens_cpu=verify_lens_cpu, total_verify_tokens=total_verify_tokens, + cap=cap, ) @classmethod @@ -154,7 +164,16 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True): device=device, ) - def padded_to_bucket(self, *, padded_bs: int) -> RaggedVerifyLayout: + def padded_to_bucket( + self, *, padded_bs: int, cap: Optional[int] = None + ) -> RaggedVerifyLayout: + """Pad to the captured slot count. The full-coverage variant (cap=None) + keeps sum(verify_lens) == graph_num_tokens, so padding rows (or, with + no padding rows, the last real row) can exceed the per-request verify + window. Dense-layout consumers (mamba/KDA) must pass cap=N: rows are + clamped to N and the layout no longer covers the tier's leftover + tokens -- those collapse into the consumer's ghost slot. Real tokens + map identically under both variants.""" from sglang.kernels.ops.speculative.ragged_verify_kernels import ( PaddedToBucket, ) @@ -165,11 +184,14 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True): bs=self.bs, padded_bs=padded_bs, ) + if cap is not None: + padded = torch.clamp(padded, max=cap) return RaggedVerifyLayout._assemble_device( verify_lens=padded, graph_num_tokens=self.graph_num_tokens, - total_verify_tokens=self.graph_num_tokens, + total_verify_tokens=None if cap is not None else self.graph_num_tokens, + cap=cap, ) diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 898c6384e..440bcc0be 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -184,6 +184,14 @@ class SpeculativeAlgorithm(Enum): return build_eagle_disagg_draft_input( batch, server_args, last_tokens_tensor, future_map ) + if self.is_dspark(): + from sglang.srt.speculative.dspark_disaggregation import ( + build_dspark_disagg_draft_input, + ) + + return build_dspark_disagg_draft_input( + batch, server_args, last_tokens_tensor, future_map + ) return None def need_topk(self) -> bool: diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 1c0f3fbb3..2845a2d6c 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -935,10 +935,72 @@ def commit_mamba_states_after_verify( # NOTE: radix mamba prefix-caching (mamba_track / extra_buffer) would need # a device-side force-flush so `temporal` reflects the ring before a # snapshot; not wired for Part B (server_args forbids extra_buffer with - # --enable-gdn-replayssm-spec), so the per-track scatters are intentionally + # --enable-linear-replayssm-spec), so the per-track scatters are intentionally # skipped here. return + # KDA ReplaySSM (fold-every-commit): KDA keeps its own recurrent verify kernel + # for the OUTPUT, so we replay the accepted window into the fp32 checkpoint + # (`temporal`) here on commit -- `temporal` is always the current committed + # state. The draft window's raw inputs were written to the ring during verify + # by the KDA backend. Gate on the fold flag + is_kda (the cursor tensors are + # never allocated under fold, so they cannot serve as the signal). + if ( + mamba_pool is not None + and getattr(mamba_pool, "replayssm_spec_fold", False) + and getattr(mamba_pool, "replayssm_is_kda", False) + ): + if batch.forward_mode.is_idle() or accept_index.numel() == 0: + return + from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( + commit_kda_replayssm_after_verify, + ) + + spec_state = req_pool.get_speculative_mamba2_params_all_layers() + bs = accept_lens.shape[0] + state_batch_indices = req_pool.get_mamba_indices(batch.req_pool_indices) + accept_indices_offset = torch.arange( + 0, + bs * draft_token_num, + step=draft_token_num, + dtype=accept_lens.dtype, + device=accept_lens.device, + ) + req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device) + last_correct_step_indices = ( + accept_index[req_idx, (accept_lens - 1).to(torch.int64)] + - accept_indices_offset + ) + # extra_buffer: the interval-crossing step whose state must snapshot into + # the track ping-pong slot (mirrors the regular commit's + # mamba_steps_to_track); commit_kda_replayssm_spec folds it in one pass, so + # `temporal` stays current and no device-side force-flush is needed. + mamba_track_indices = batch.mamba_track_indices + mamba_steps_to_track = None + if mamba_track_indices is not None: + ti = get_exec().mamba.mamba_track_interval + seq_pre = batch.seq_lens + seq_post = batch.seq_lens + accept_lens + to_track_mask = seq_pre // ti != seq_post // ti + tracking_point = seq_post // ti * ti + to_track_ith = torch.clamp(tracking_point - seq_pre - 1, min=0).to( + torch.int64 + ) + candidate = accept_index[req_idx, to_track_ith] - accept_indices_offset + mamba_steps_to_track = torch.where( + to_track_mask, candidate, torch.full_like(candidate, -1) + ) + commit_kda_replayssm_after_verify( + spec_state=spec_state, + state_batch_indices=state_batch_indices, + accept_lens=accept_lens, # incl. bonus token + last_correct_step_indices=last_correct_step_indices, + mamba_track_indices=mamba_track_indices, + mamba_steps_to_track=mamba_steps_to_track, + null_block_id=-1, # SGLang: valid slots >= 0, padding == -1 + ) + return + attn_backend = model_runner.attn_backend bs = accept_lens.shape[0] diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 1b01e9269..6636d06ab 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -843,6 +843,18 @@ def get_device_name(device_id: int = 0) -> str: return torch.npu.get_device_name(device_id) +@lru_cache(maxsize=1) +def is_mnnvl_fabric_device() -> bool: + """Whether the GPU sits on an MNNVL fabric (cross-node NVLink), keyed on + the device name: the GB200/GB300 superchips. Used to auto-select + fabric-dependent communication paths (NCCL cuMem/MNNVL, custom all-reduce + v2 multinode, DCP fi_a2a).""" + if not (hasattr(torch, "cuda") and torch.cuda.is_available()): + return False + name = (torch.cuda.get_device_name(0) or "").upper() + return any(tag in name for tag in ("GB200", "GB300")) + + @lru_cache(maxsize=1) def is_habana_available() -> bool: return find_spec("habana_frameworks") is not None diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index cd84b2885..3a237069f 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -39,6 +39,7 @@ from sglang.srt.configs import ( InternS2PreviewConfig, JetNemotronConfig, JetVLMConfig, + KimiK3Config, KimiK25Config, KimiLinearConfig, KimiVLConfig, @@ -98,6 +99,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { Step3VLConfig, LongcatFlashConfig, Olmo3Config, + KimiK3Config, KimiLinearConfig, Qwen3NextConfig, FalconH1Config, diff --git a/scripts/ci/cuda/ci_install_kimi_k3.sh b/scripts/ci/cuda/ci_install_kimi_k3.sh new file mode 100755 index 000000000..bb8d7e45a --- /dev/null +++ b/scripts/ci/cuda/ci_install_kimi_k3.sh @@ -0,0 +1,139 @@ +#!/bin/bash +# Install the standard CUDA CI dependencies plus Kimi-K3's FlashInfer assets. +set -euxo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" + +# Source (not bash) so the generic install's Python/venv selection remains +# active while patching the FlashInfer package it just installed. +# shellcheck disable=SC1091 +source "${SCRIPT_DIR}/ci_install_dependency.sh" "$@" + +TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip" +TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192" +TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1" +export SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool" + +install_required_tools() { + local missing=() + local tool + for tool in patch unzip wget; do + command -v "${tool}" >/dev/null 2>&1 || missing+=("${tool}") + done + if [ ${#missing[@]} -eq 0 ]; then + return + fi + + apt-get update || true + apt-get install -y --no-install-recommends "${missing[@]}" +} + +cubin_pool_is_valid() { + [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" ] && + [ "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" -type f -name '*.cubin' | wc -l)" -eq 1696 ] && + [ -f "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/flashinferMetaInfo.h" ] && + [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/local" ] && + [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/overlay/csrc" ] +} + +install_trtllm_gen_moe_cubin_pool() ( + if cubin_pool_is_valid; then + echo "Reusing validated TRT-LLM Gen MoE cubin pool" + return + fi + + local cubin_archive + local cubin_extract_dir + local extracted_pool + cubin_archive="$(mktemp /tmp/trtllm_gen_moe_cubin_pool.XXXXXX.zip)" + cubin_extract_dir="$(mktemp -d /tmp/trtllm_gen_moe_cubin_extract.XXXXXX)" + extracted_pool="${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" + trap 'rm -f "${cubin_archive}"; rm -rf "${cubin_extract_dir}"' EXIT + + wget --no-verbose --output-document="${cubin_archive}" \ + "${TRTLLM_GEN_MOE_CUBIN_URL}" + echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \ + sha256sum --check --strict - + unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" + test "$(find "${extracted_pool}" -type f -name '*.cubin' | wc -l)" -eq 1696 + + rm -rf "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" + mv "${extracted_pool}" "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" + cubin_pool_is_valid +) + +apply_flashinfer_dcp_patch() { + local flashinfer_dcp_patch + local flashinfer_site_packages + flashinfer_dcp_patch="${REPO_ROOT}/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt" + flashinfer_site_packages="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" + + sed '/^diff --git a\/tests\//,$d' "${flashinfer_dcp_patch}" | \ + patch --dry-run --batch --forward --strip=1 \ + --directory="${flashinfer_site_packages}" + sed '/^diff --git a\/tests\//,$d' "${flashinfer_dcp_patch}" | \ + patch --batch --forward --strip=1 \ + --directory="${flashinfer_site_packages}" + + rm -rf /root/.cache/flashinfer /root/.cache/pip + python3 -c 'import inspect; from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla; assert "enable_dcp" in inspect.signature(trtllm_batch_decode_with_kv_cache_mla).parameters' +} + +apply_transformers_symlink_patch() { + # transformers 5.12.1 resolves custom-code symlinks out of the HF snapshot + # and into blobs/, then looks for relative imports by filename in blobs/. + # Mirror the hot patch in docker/rocm.Dockerfile until upstream PR #46618 + # reaches a transformers release. + python3 - <<'PY' +import importlib +import pathlib +import tempfile + +import transformers.dynamic_module_utils as dynamic_module_utils + +marks = ["Path(resolved_module_file).resolve()", "Path(source_file).resolve()"] +path = pathlib.Path(dynamic_module_utils.__file__) +src = path.read_text() +if not any(mark in src for mark in marks): + print("transformers dynamic_module_utils already fixed; no patch needed") +else: + patched = ( + src.replace( + "Path(resolved_module_file).resolve()", "Path(resolved_module_file)" + ).replace("Path(source_file).resolve()", "Path(source_file)") + ) + assert patched != src, "FATAL: transformers symlink patch matched nothing" + path.write_text(patched) + print("patched transformers dynamic_module_utils.py (symlink hash fix)") + +# Exercise the exact HF-cache layout that failed in the B300 Kimi-K3 run. +dynamic_module_utils = importlib.reload(dynamic_module_utils) +with tempfile.TemporaryDirectory() as tmp: + cache = pathlib.Path(tmp) / "models--org--model" + blobs = cache / "blobs" + snapshot = cache / "snapshots" / "revision" + blobs.mkdir(parents=True) + snapshot.mkdir(parents=True) + (blobs / "modeling-blob").write_text("from .media_utils import VALUE\n") + (blobs / "media-blob").write_text("VALUE = 1\n") + (snapshot / "modeling.py").symlink_to("../../blobs/modeling-blob") + (snapshot / "media_utils.py").symlink_to("../../blobs/media-blob") + source_hash = dynamic_module_utils._compute_local_source_files_hash( + snapshot, snapshot / "modeling.py" + ) + assert len(source_hash) == 16, source_hash +PY +} + +install_required_tools +install_trtllm_gen_moe_cubin_pool +apply_flashinfer_dcp_patch +apply_transformers_symlink_patch + +# The install runs in its own shell. Persist the pool path for later workflow +# steps so Kimi-K3 selects the FlashInfer MXFP4 MoE runner instead of failing +# its startup validation. +if [ -n "${GITHUB_ENV:-}" ]; then + echo "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" >> "${GITHUB_ENV}" +fi diff --git a/scripts/ci/runner_configs.yml b/scripts/ci/runner_configs.yml index 605a79d12..dfac4efd9 100644 --- a/scripts/ci/runner_configs.yml +++ b/scripts/ci/runner_configs.yml @@ -19,6 +19,7 @@ _anchors: default_install: &default scripts/ci/cuda/ci_install_dependency.sh deepep_install: &deepep scripts/ci/cuda/ci_install_deepep.sh + kimi_k3_install: &kimi_k3 scripts/ci/cuda/ci_install_kimi_k3.sh runner_configs: 1-gpu-small: { install: *default, artifact_version: v4, install_timeout: "20", runs_on: 1-gpu-5090 } @@ -29,6 +30,7 @@ runner_configs: 4-gpu-h100: { install: *default, artifact_version: v4, install_timeout: "20", runs_on: 4-gpu-h100 } 8-gpu-h200: { install: *default, artifact_version: v4, install_timeout: "20", runs_on: 8-gpu-h200 } 8-gpu-b200: { install: *default, artifact_version: v6, install_timeout: "20", runs_on: 8-gpu-b200 } + 8-gpu-b300: { install: *kimi_k3, artifact_version: v6, install_timeout: "20", runs_on: 8-gpu-b300 } 8-gpu-h20: { install: *deepep, artifact_version: v4, install_timeout: "20", runs_on: 8-gpu-h20, rdma_devices: "mlx5_1,mlx5_2,mlx5_3,mlx5_4" } deepep-4-gpu-h100: { install: *deepep, artifact_version: v4, install_timeout: "20", runs_on: 4-gpu-h100 } deepep-4-gpu-b200: { install: *deepep, artifact_version: v6, install_timeout: "20", runs_on: $b200_runner } diff --git a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index b32cbf1ff..5034eb4b1 100644 --- a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py +++ b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py @@ -168,6 +168,31 @@ class TestBreakableCUDAGraphBasic(CustomTestCase): torch.cuda.synchronize() self.assertTrue(torch.allclose(y, torch.full((4,), 33.0, device=self.device))) + def test_side_stream_join_across_break(self): + """A side-stream producer may be joined after a graph break.""" + x = torch.zeros(4, device=self.device) + y = torch.zeros(4, device=self.device) + stream = torch.cuda.Stream(self.device) + + @self.eager_on_graph(enable=True) + def identity(src): + return src + + graph = self.BreakableCUDAGraph() + capture_stream = torch.cuda.Stream(self.device) + with self.BreakableCUDAGraphCapture(graph, stream=capture_stream): + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + side_output = x + 1.0 + identity(x) + torch.cuda.current_stream().wait_stream(stream) + y.copy_(side_output * 2.0) + + x.fill_(5.0) + graph.replay() + torch.cuda.synchronize() + self.assertTrue(torch.allclose(y, torch.full((4,), 12.0, device=self.device))) + def test_eager_output_is_held_strongly_for_replay_bridge(self): """The replay closure must keep the eager output bridge buffer alive.""" x = torch.zeros(4, device=self.device) diff --git a/test/registered/dcp/test_kimi_linear_dcp4.py b/test/registered/dcp/test_kimi_linear_dcp4.py index 37ab83ef5..227757292 100644 --- a/test/registered/dcp/test_kimi_linear_dcp4.py +++ b/test/registered/dcp/test_kimi_linear_dcp4.py @@ -1,7 +1,7 @@ """Four-Blackwell acceptance coverage for Kimi Linear TokenSpeed MLA DCP. The captured-shape and eager-shape requests deliberately straddle -``--cuda-graph-max-bs-decode=64``. This guards both the regular CUDA graph +``--cuda-graph-max-bs-decode``. This guards both the regular CUDA graph decode path and the full-capacity eager DCP LSE scratch-buffer path. """ @@ -20,9 +20,10 @@ from sglang.test.test_utils import ( popen_launch_server, ) -register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-b200") +register_cuda_ci(est_time=240, stage="extra-b", runner_config="4-gpu-b200") KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" +CUDA_GRAPH_MAX_BS_DECODE = 256 def _has_four_blackwell_gpus() -> bool: @@ -41,12 +42,11 @@ def _has_four_blackwell_gpus() -> bool: class TestKimiLinearDCP4(GSM8KMixin, CustomTestCase): model = KIMI_LINEAR_MODEL base_url = DEFAULT_URL_FOR_TEST - gsm8k_score_threshold = 0.90 + gsm8k_score_threshold = 0.88 gsm8k_num_examples = 200 # Keep accuracy evaluation within the captured decode batch sizes so its - # score is batch-invariant. The separate smoke test still exercises the - # eager path with batch size 65. - gsm8k_num_threads = 4 + # score is batch-invariant. + gsm8k_num_threads = 128 gsm8k_num_shots = 5 @classmethod @@ -74,7 +74,7 @@ class TestKimiLinearDCP4(GSM8KMixin, CustomTestCase): "--dtype", "bfloat16", "--cuda-graph-max-bs-decode", - "64", + str(CUDA_GRAPH_MAX_BS_DECODE), "--cuda-graph-backend-prefill", "disabled", "--mem-fraction-static", @@ -112,12 +112,24 @@ class TestKimiLinearDCP4(GSM8KMixin, CustomTestCase): self.assertTrue(output["text"].strip()) self.assertGreater(output["meta_info"]["completion_tokens"], 0) + def _effective_max_running_requests(self) -> int: + response = requests.get(self.base_url + "/server_info", timeout=30) + response.raise_for_status() + return min( + state["effective_max_running_requests_per_dp"] + for state in response.json()["internal_states"] + ) + def test_decode_cuda_graph_and_eager_batch(self): - # Batch two replays a captured shape; batch 65 is above the configured - # regular CUDA graph maximum and therefore exercises eager decode. self._assert_batch_completes(2) self._assert_batch_completes(2) - self._assert_batch_completes(65) + self.assertGreater( + self._effective_max_running_requests(), + CUDA_GRAPH_MAX_BS_DECODE, + "eager DCP decode is unreachable: concurrency was capped at or " + "below the CUDA graph capture ceiling", + ) + self._assert_batch_completes(CUDA_GRAPH_MAX_BS_DECODE + 1) def test_physical_capacity_sanity(self): response = requests.get(self.base_url + "/server_info", timeout=30) diff --git a/test/registered/dcp/test_kimi_linear_dcp_dspark4.py b/test/registered/dcp/test_kimi_linear_dcp_dspark4.py index 5d6929fb0..bc58043f8 100644 --- a/test/registered/dcp/test_kimi_linear_dcp_dspark4.py +++ b/test/registered/dcp/test_kimi_linear_dcp_dspark4.py @@ -1,3 +1,5 @@ +"""Four-Blackwell Kimi Linear TokenSpeed DCP + DSpark static acceptance test.""" + import json import socket import tempfile @@ -39,6 +41,7 @@ def _has_four_blackwell_gpus() -> bool: def _write_dummy_qwen3_dspark_draft(root: Path) -> str: + """Write a dummy Qwen3 DSpark config with Kimi Linear dimensions.""" draft_dir = root / "qwen3-dspark-kimi-proxy" draft_dir.mkdir() config = { diff --git a/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py b/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py index 3db73910c..f52b1137e 100644 --- a/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py +++ b/test/registered/kernels/ops/gemm/test_cutedsl_bf16_gemm.py @@ -17,7 +17,11 @@ register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b if not torch.cuda.is_available(): pytest.skip("CUDA required", allow_module_level=True) -from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import cutedsl_bf16_gemm # noqa: E402 +from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import ( # noqa: E402 + _K3_TGV_WIN_SHAPES, + cutedsl_bf16_gemm, + use_cutedsl_bf16_gemm, +) N_VALUES = [1024, 2624, 6144] K_VALUES = [2048, 6144] @@ -48,5 +52,31 @@ def test_cutedsl_bf16_gemm(num_tokens, k, n, has_bias): torch.testing.assert_close(out, ref.bfloat16(), rtol=2e-2, atol=2.5) +@pytest.mark.parametrize("n, k", sorted(_K3_TGV_WIN_SHAPES) + [(1024, 2048)]) +def test_empty_batch_not_tgv_eligible(n, k): + """DP-attention idle groups run a 0-token dummy forward to keep the + mlp-sync lockstep; every m == 0 shape must route to cuBLAS.""" + assert not use_cutedsl_bf16_gemm(0, n, k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("has_bias", [False, True]) +def test_cutedsl_bf16_gemm_empty_batch(has_bias): + """Empty input must yield the empty [0, N] output, mirroring F.linear — + launching TGV with a 0-CTA grid fails with CUDA_ERROR_INVALID_VALUE.""" + if is_hip_runtime() or get_jit_cuda_arch().major != 10: + pytest.skip("SM100/SM103 required") + + n, k = 6144, 2048 + x = torch.empty(0, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + bias = torch.randn(n, dtype=torch.bfloat16, device="cuda") if has_bias else None + + out = cutedsl_bf16_gemm(x, weight, bias) + torch.cuda.synchronize() + assert out.shape == (0, n) + assert out.dtype == torch.bfloat16 + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py new file mode 100644 index 000000000..585610b53 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py @@ -0,0 +1,451 @@ +"""Correctness test for the K3 MNNVL fused all-reduce (ar_fusion) kernels. + +Compares the 1shot multicast-push and the in-place low-SM NVLS 2shot pull +(with and without the fused residual) against +NCCL, bit-exact on small-int bf16 inputs; the fused-RMSNorm pull against a +torch reference; the pull tuning knobs (num_blocks, unroll) on sizes whose +shard split is uneven; plus a CUDA-graph capture/replay pass and a mixed +stress loop exercising the push phase double-buffering and the pull +semaphore window cycling. + +Usage:: + + python test/registered/jit/kimi_k3/test_ar_fusion.py # relaunches under torchrun (8 GPUs) +""" + +from __future__ import annotations + +import atexit +import logging +import os + +import pytest +import torch +import torch.distributed as dist + +import sglang.srt.distributed.parallel_state as ps +from sglang.kernels.jit.utils import cache_once, get_ci_test_range +from sglang.kernels.ops.communication.mp import register_comm_cleanup +from sglang.kernels.ops.kimi_k3 import all_reduce +from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main + +register_cuda_ci( + est_time=240, + stage="extra-b", + runner_config="8-gpu-h200", +) + +H = 7168 # Kimi-K3 hidden size; the kernels are tuned/used at multiples of it +NORM_DIM = 3584 # latent width; the norm buffer is [N, NORM_DIM] + [N, 2*NORM_DIM] +MB = 1024 * 1024 + +PUSH_BS = [1, 2, 8, 32, 128] +PULL_BS = [1, 8, 64, 1024, 4096] +PUSH_BS = get_ci_test_range(PUSH_BS, [1, 32, 128]) +PULL_BS = get_ci_test_range(PULL_BS, [1, 64, 4096]) + + +def _precompile(num_gpus): + for ws in num_gpus: + all_reduce._jit_module(ws) + + +@cache_once +def _init_world(): + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="gloo") + ps._WORLD = coord = ps.init_world_group( + ranks=list(range(world_size)), + local_rank=local_rank, + backend="nccl", + ) + atexit.register(dist.destroy_process_group) + logging.disable(logging.INFO) + torch.cuda.set_stream(torch.cuda.Stream()) + return coord.cpu_group + + +@cache_once +def _init_nccl_group(): + _init_world() + local_rank = int(os.environ["LOCAL_RANK"]) + group = dist.new_group(backend="nccl", device_id=torch.device(f"cuda:{local_rank}")) + assert isinstance(group, dist.ProcessGroup) + return group + + +def _symm_alloc_mc(shape, dtype) -> tuple[torch.Tensor, int]: + import torch.distributed._symmetric_memory as torch_symm_mem + + cpu_group = _init_world() + rank = dist.get_rank() + device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + pool = torch_symm_mem.get_mem_pool(device) + with torch.cuda.use_mem_pool(pool): + buf = torch.empty(shape, dtype=dtype, device=device) + hdl = torch_symm_mem.rendezvous(buf, cpu_group.group_name) + assert hdl.multicast_ptr != 0 + mc = hdl.multicast_ptr + (buf.data_ptr() - hdl.buffer_ptrs[rank]) + return buf, mc + + +@cache_once +def _init_comm() -> CustomAllReduceV2: + cpu_group = _init_world() + device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + comm = CustomAllReduceV2( + cpu_group, device, max_pull_size=1 * MB, max_push_size=2 * MB + ) + if comm.disabled or comm.mc_base_ptr == 0: + raise RuntimeError("ar_fusion requires CustomAllReduceV2 with multicast") + all_reduce.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr) + register_comm_cleanup(comm) + return comm + + +@cache_once +def _init_pool_buf() -> tuple[torch.Tensor, int]: + # 1.5x headroom: the norm tests view the buffer as [N, 3584 + 7168] + return _symm_alloc_mc((max(PULL_BS) * H * 3 // 2,), torch.bfloat16) + + +def _device() -> torch.device: + return torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + + +def _int_input(n: int, seed: int, per_rank: bool) -> torch.Tensor: + # small ints are exact in bf16 even after an fp32-accumulated 8-way sum, + # so the comparison against NCCL is bit-exact + rank = dist.get_rank() if per_rank else 0 + g = torch.Generator().manual_seed(seed * 1009 + rank) + return torch.randint(0, 16, (n,), dtype=torch.bfloat16, generator=g).to(_device()) + + +def _nccl_ref(x: torch.Tensor, residual): + ref = x.clone() + dist.all_reduce(ref, group=_init_nccl_group()) + return ref if residual is None else ref + residual + + +def _norm_ref(x_reduced: torch.Tensor, num_norm_rows: int, weight, eps: float): + """allreduce result -> RMSNorm over the first num_norm_rows rows of the + [numel / NORM_DIM, NORM_DIM] row view, in fp32 like the kernels.""" + out = x_reduced.clone() + normed_part = out[: num_norm_rows * NORM_DIM].view(num_norm_rows, NORM_DIM).float() + factor = torch.rsqrt(normed_part.pow(2).mean(-1, keepdim=True) + eps) + normed = (normed_part * factor * weight.float()).to(torch.bfloat16) + out[: num_norm_rows * NORM_DIM] = normed.view(-1) + return out + + +def _assert_norm_close(x: torch.Tensor, ref: torch.Tensor, num_norm_rows: int): + # the non-normed tail is a plain allreduce: bit-exact; the normed prefix + # gets the fp32 norm epilogue: bf16 tolerances + torch.testing.assert_close( + x[num_norm_rows * NORM_DIM :], ref[num_norm_rows * NORM_DIM :], atol=0, rtol=0 + ) + torch.testing.assert_close( + x[: num_norm_rows * NORM_DIM], ref[: num_norm_rows * NORM_DIM] + ) + + +@pytest.mark.parametrize("bs", PUSH_BS) +@pytest.mark.parametrize("use_residual", [False, True]) +@torch.inference_mode() +def test_ar_fusion_push(bs: int, use_residual: bool): + comm = _init_comm() + world = comm.world_size + n = bs * H + x = _int_input(n, bs, per_rank=True) + residual = _int_input(n, bs + 7, per_rank=False) if use_residual else None + ref = _nccl_ref(x, residual) + all_reduce.all_reduce_push_res(world, x, residual, ws_mc_base=comm.mc_base_ptr) + torch.cuda.synchronize() + torch.testing.assert_close(x, ref, atol=0, rtol=0) + + +@pytest.mark.parametrize("bs", PULL_BS) +@pytest.mark.parametrize("use_residual", [False, True]) +@torch.inference_mode() +def test_ar_fusion_pull_2shot(bs: int, use_residual: bool): + comm = _init_comm() + world = comm.world_size + buf, mc = _init_pool_buf() + n = bs * H + x = buf[:n] + x.copy_(_int_input(n, bs + 13, per_rank=True)) + residual = _int_input(n, bs + 17, per_rank=False) if use_residual else None + ref = _nccl_ref(x, residual) + all_reduce.all_reduce_pull_res(world, x, residual, input_mc_ptr=mc) + torch.cuda.synchronize() + torch.testing.assert_close(x, ref, atol=0, rtol=0) + + +@pytest.mark.parametrize("num_blocks", [1, 2, 4, 8]) +@pytest.mark.parametrize("unroll", [4, 8]) +@torch.inference_mode() +def test_ar_fusion_pull_tuning_grid(num_blocks: int, unroll: int): + """Every (num_blocks, unroll) combination must agree with NCCL on a size + whose 16B-vector count is not divisible by the world size (uneven shards) + and whose per-thread range leaves an unrolled-loop tail.""" + _init_comm() + world = dist.get_world_size() + buf, mc = _init_pool_buf() + n = (3 * H + 7) * 8 # 21511 vecs: % 8 ranks != 0, small vs blocks*512*unroll + x = buf[:n] + x.copy_(_int_input(n, num_blocks * 10 + unroll, per_rank=True)) + ref = _nccl_ref(x, None) + all_reduce.all_reduce_pull_res( + world, x, None, input_mc_ptr=mc, num_blocks=num_blocks, unroll=unroll + ) + torch.cuda.synchronize() + torch.testing.assert_close(x, ref, atol=0, rtol=0) + + +@pytest.mark.parametrize("num_tokens", PULL_BS) +@pytest.mark.parametrize("rows_per_token", [3, 1]) # [N|2N] MoE buf / latent-only +@torch.inference_mode() +def test_ar_fusion_pull_norm(num_tokens: int, rows_per_token: int): + _init_comm() + world = dist.get_world_size() + buf, mc = _init_pool_buf() + n = num_tokens * rows_per_token * NORM_DIM + x = buf[:n] + x.copy_(_int_input(n, num_tokens + 23 + rows_per_token, per_rank=True)) + weight = _int_input(NORM_DIM, 29, per_rank=False) + 1 # small positive ints + ref = _norm_ref(_nccl_ref(x, None), num_tokens, weight, eps=1e-6) + all_reduce.all_reduce_pull_norm( + world, x, weight, 1e-6, num_norm_rows=num_tokens, input_mc_ptr=mc + ) + torch.cuda.synchronize() + _assert_norm_close(x, ref, num_tokens) + + +@pytest.mark.parametrize("num_tokens", [1, 8, 24]) +@pytest.mark.parametrize("rows_per_token", [3, 1]) +@torch.inference_mode() +def test_ar_fusion_push_norm(num_tokens: int, rows_per_token: int): + """The push-side norm (small-message regime of the serving dispatch) with + an explicit num_norm_rows, on both the MoE-buffer and latent-only row + layouts.""" + comm = _init_comm() + world = comm.world_size + n = num_tokens * rows_per_token * NORM_DIM + x = _int_input(n, num_tokens + 41 + rows_per_token, per_rank=True) + weight = _int_input(NORM_DIM, 43, per_rank=False) + 1 + ref = _norm_ref(_nccl_ref(x, None), num_tokens, weight, eps=1e-6) + all_reduce.all_reduce_push_norm( + world, x, weight, 1e-6, num_norm_rows=num_tokens, ws_mc_base=comm.mc_base_ptr + ) + torch.cuda.synchronize() + _assert_norm_close(x, ref, num_tokens) + + +FIN_TOPK = 16 + + +def _build_permuted_layout(num_tokens: int, seed: int): + """trtllm-gen permuted gemm2 layout (rows grouped by expert, per-expert + tile padding). Deterministic on CPU: idx/weights are identical on every + rank (TP semantics — same routing), gemm2 values are per-rank.""" + num_experts, tile = 896, 8 + gen = torch.Generator(device="cpu").manual_seed(seed) + topk_ids = torch.stack( + [ + torch.randperm(num_experts, generator=gen)[:FIN_TOPK] + for _ in range(num_tokens) + ] + ) + counts = torch.bincount(topk_ids.flatten(), minlength=num_experts) + padded = (counts + tile - 1) // tile * tile + bases = torch.cumsum(padded, 0) - padded + fill = torch.zeros(num_experts, dtype=torch.long) + idx = torch.empty(num_tokens * FIN_TOPK, dtype=torch.int32) + for i, e in enumerate(topk_ids.flatten().tolist()): + idx[i] = bases[e] + fill[e] + fill[e] += 1 + weights = torch.rand(num_tokens, FIN_TOPK, generator=gen).to(torch.bfloat16) + num_rows = int(padded.sum()) + g = torch.Generator(device="cpu").manual_seed(seed * 31 + dist.get_rank()) + gemm2 = (torch.randn(num_rows, NORM_DIM, generator=g) * 2).to(torch.bfloat16) + dev = _device() + return gemm2.to(dev), idx.to(dev), weights.to(dev) + + +def _finalize_norm_ref(gemm2, idx, weights, norm_w, eps: float) -> torch.Tensor: + """Replicates the fused kernel numerics: fp32 ascending-k local finalize + cast to bf16 (the staged push value), rank-ordered fp32 cross-rank sum, + fp32 RMSNorm. Only the rsqrt may differ from the kernel by ulps.""" + num_tokens = weights.shape[0] + idx2 = idx.view(num_tokens, FIN_TOPK).long() + acc = torch.zeros(num_tokens, NORM_DIM, dtype=torch.float32, device=gemm2.device) + for k in range(FIN_TOPK): + acc += weights[:, k, None].float() * gemm2[idx2[:, k]].float() + local = acc.to(torch.bfloat16) + world = dist.get_world_size() + gathered = [torch.empty_like(local) for _ in range(world)] + dist.all_gather(gathered, local, group=_init_nccl_group()) + total = torch.zeros_like(acc) + for r in range(world): + total += gathered[r].float() + factor = torch.rsqrt(total.square().mean(dim=-1, keepdim=True) + eps) + return (total * factor * norm_w.float()).to(torch.bfloat16) + + +@pytest.mark.parametrize("bs", PUSH_BS) +@torch.inference_mode() +def test_ar_fusion_finalize_push_norm(bs: int): + comm = _init_comm() + world = comm.world_size + eps = 1e-6 + gemm2, idx, weights = _build_permuted_layout(bs, seed=bs + 23) + g = torch.Generator(device="cpu").manual_seed(77) + norm_w = (torch.rand(NORM_DIM, generator=g) + 0.5).to(torch.bfloat16).to(_device()) + ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps) + out = torch.empty(bs, NORM_DIM, dtype=torch.bfloat16, device=_device()) + all_reduce.finalize_all_reduce_push_norm( + world, out, gemm2, idx, weights, norm_w, eps, ws_mc_base=comm.mc_base_ptr + ) + torch.cuda.synchronize() + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +@torch.inference_mode() +def test_ar_fusion_finalize_push_norm_stress(): + """Back-to-back fused calls interleaved with plain pushes exercise the + shared push-workspace phase double-buffering across kernel variants.""" + comm = _init_comm() + world = comm.world_size + eps = 1e-5 + g = torch.Generator(device="cpu").manual_seed(78) + norm_w = (torch.rand(NORM_DIM, generator=g) + 0.5).to(torch.bfloat16).to(_device()) + for it in range(12): + bs = (1, 8, 32)[it % 3] + gemm2, idx, weights = _build_permuted_layout(bs, seed=9000 + it) + ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps) + out = torch.empty(bs, NORM_DIM, dtype=torch.bfloat16, device=_device()) + all_reduce.finalize_all_reduce_push_norm( + world, out, gemm2, idx, weights, norm_w, eps, ws_mc_base=comm.mc_base_ptr + ) + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + x = _int_input(bs * H, 8000 + it, per_rank=True) + ref2 = _nccl_ref(x, None) + all_reduce.all_reduce_push_res(world, x, None, ws_mc_base=comm.mc_base_ptr) + torch.testing.assert_close(x, ref2, atol=0, rtol=0) + + +@pytest.mark.parametrize("num_blocks", [1, 4, 16]) +@pytest.mark.parametrize("unroll", [4, 8]) +@torch.inference_mode() +def test_ar_fusion_pull_norm_tuning_grid(num_blocks: int, unroll: int): + """Every (num_blocks, unroll) combination must agree on a token count + whose row count is not divisible by the world size (uneven row shards) + and not by unroll (partial last row group per block).""" + _init_comm() + world = dist.get_world_size() + buf, mc = _init_pool_buf() + num_tokens = 13 # 39 rows: % 8 ranks != 0, per-rank rows < num_blocks*unroll + n = num_tokens * 3 * NORM_DIM + x = buf[:n] + x.copy_(_int_input(n, 500 + num_blocks * 10 + unroll, per_rank=True)) + weight = _int_input(NORM_DIM, 31, per_rank=False) + 1 + ref = _norm_ref(_nccl_ref(x, None), num_tokens, weight, eps=1e-6) + all_reduce.all_reduce_pull_norm( + world, + x, + weight, + 1e-6, + num_norm_rows=num_tokens, + input_mc_ptr=mc, + num_blocks=num_blocks, + unroll=unroll, + ) + torch.cuda.synchronize() + _assert_norm_close(x, ref, num_tokens) + + +@torch.inference_mode() +def test_ar_fusion_stress_mixed(): + """Back-to-back mixed calls exercise the push phase double-buffering and + the pull semaphore window cycling (with varying grids).""" + comm = _init_comm() + world = comm.world_size + buf, mc = _init_pool_buf() + for it in range(32): + n = (1, 8, 64)[it % 3] * H + num_blocks = (1, 2, 4, 8)[it % 4] + x = _int_input(n, 3000 + it, per_rank=True) + ref = _nccl_ref(x, None) + all_reduce.all_reduce_push_res(world, x, None, ws_mc_base=comm.mc_base_ptr) + torch.testing.assert_close(x, ref, atol=0, rtol=0) + y = buf[:n] + y.copy_(_int_input(n, 4000 + it, per_rank=True)) + ref2 = _nccl_ref(y, None) + all_reduce.all_reduce_pull_res( + world, y, None, input_mc_ptr=mc, num_blocks=num_blocks + ) + torch.testing.assert_close(y, ref2, atol=0, rtol=0) + + +@torch.inference_mode() +def test_ar_fusion_graph_capture(): + comm = _init_comm() + world = comm.world_size + buf, mc = _init_pool_buf() + cpu_group = _init_world() + n = 64 * H + gres = _int_input(n, 99, per_rank=False) + gx = torch.zeros(n, dtype=torch.bfloat16, device=_device()) + # two disjoint regions of the symm buffer, one per captured pull kernel + gy, mc_y = buf[:n], mc + gz, mc_z = buf[n : 2 * n], mc + n * buf.element_size() + + def _run_all(): + all_reduce.all_reduce_push_res(world, gx, gres, ws_mc_base=comm.mc_base_ptr) + all_reduce.all_reduce_pull_res(world, gy, gres, input_mc_ptr=mc_y) + all_reduce.all_reduce_pull_res(world, gz, gres, input_mc_ptr=mc_z) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _run_all() + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + dist.barrier(group=cpu_group) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run_all() + + for it in range(4): + vx = _int_input(n, 5000 + it, per_rank=True) + vy = _int_input(n, 6000 + it, per_rank=True) + vz = _int_input(n, 7000 + it, per_rank=True) + ref_x = _nccl_ref(vx, gres) + ref_y = _nccl_ref(vy, gres) + ref_z = _nccl_ref(vz, gres) + gx.copy_(vx) + gy.copy_(vy) + gz.copy_(vz) + dist.barrier(group=cpu_group) + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(gx, ref_x, atol=0, rtol=0) + torch.testing.assert_close(gy, ref_y, atol=0, rtol=0) + torch.testing.assert_close(gz, ref_z, atol=0, rtol=0) + + +if __name__ == "__main__": + multigpu_pytest_main( + __name__, + __file__, + num_gpus=(8,), + pre_launch_fn=_precompile, + ) diff --git a/test/registered/kernels/ops/moe/test_renorm.py b/test/registered/kernels/ops/moe/test_renorm.py index 7aafcf979..ae66e0ff9 100644 --- a/test/registered/kernels/ops/moe/test_renorm.py +++ b/test/registered/kernels/ops/moe/test_renorm.py @@ -4,12 +4,23 @@ import sys import pytest -import sgl_kernel import torch -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import is_hip +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x") + +if is_hip(): + from sglang.kernels.ops.sampling.renorm_triton import ( + top_k_renorm_probs_triton as top_k_renorm_prob, + ) + from sglang.kernels.ops.sampling.renorm_triton import ( + top_p_renorm_probs_triton as top_p_renorm_prob, + ) +else: + from sgl_kernel import top_k_renorm_prob, top_p_renorm_prob @pytest.mark.parametrize("batch_size", [1, 99, 989]) @@ -37,7 +48,7 @@ def test_top_k_renorm_probs(batch_size, vocab_size, k): dim=-1, keepdim=True ) - renorm_prob = sgl_kernel.top_k_renorm_prob(normalized_prob, k) + renorm_prob = top_k_renorm_prob(normalized_prob, k) for i in range(batch_size): torch.testing.assert_close( renorm_prob_ground_truth[i], @@ -72,7 +83,7 @@ def test_top_p_renorm_probs(batch_size, vocab_size, p): dim=-1, keepdim=True ) - renorm_prob = sgl_kernel.top_p_renorm_prob(normalized_prob, p) + renorm_prob = top_p_renorm_prob(normalized_prob, p) torch.testing.assert_close( renorm_prob_ground_truth, renorm_prob, diff --git a/test/registered/models_e2e/test_kimi_k3_b300.py b/test/registered/models_e2e/test_kimi_k3_b300.py new file mode 100644 index 000000000..9da3bd023 --- /dev/null +++ b/test/registered/models_e2e/test_kimi_k3_b300.py @@ -0,0 +1,117 @@ +"""B300 per-commit CI coverage for Kimi-K3 serving recipes. + +Runs the Low Latency DSPARK recipe and the Balanced DCP/HiCache recipe on +eight B300 GPUs. Each server must preserve basic model quality on GSM8K. +""" + +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + _wait_for_gpu_idle_in_ci, + popen_launch_server, +) + +register_cuda_ci(est_time=1800, stage="base-c", runner_config="8-gpu-b300") + +MODEL_PATH = ( + "/data/radixark/model-cache/hub/models--moonshotai--Kimi-K3/" + "snapshots/9f62e4e9fffbd0a83ddd60e1c209d828994b3569" +) +DSPARK_DRAFT_MODEL = "RadixArk/Kimi-K3-DSpark" +SERVER_LAUNCH_TIMEOUT = 3600 +GPU_IDLE_TIMEOUT = 120 + + +def _stop_server(process): + if process: + kill_process_tree(process.pid) + _wait_for_gpu_idle_in_ci(timeout=GPU_IDLE_TIMEOUT) + + +class TestKimiK3B300LowLatency(GSM8KMixin, CustomTestCase): + """TP8 Low Latency recipe with DSPARK linear ReplaySSM speculation.""" + + gsm8k_score_threshold = 0.95 + gsm8k_num_examples = 200 + + @classmethod + def setUpClass(cls): + cls.model = MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--tp-size", + "8", + "--mem-fraction-static", + "0.85", + "--weight-loader-prefetch-checkpoints", + "--reasoning-parser", + "kimi_k3", + "--tool-call-parser", + "kimi_k3", + "--mamba-full-memory-ratio", + "0.86", + "--speculative-algorithm", + "DSPARK", + "--speculative-draft-model-path", + DSPARK_DRAFT_MODEL, + "--speculative-dspark-block-size", + "7", + ], + ) + + @classmethod + def tearDownClass(cls): + _stop_server(getattr(cls, "process", None)) + + +class TestKimiK3B300Balanced(GSM8KMixin, CustomTestCase): + """TP8/DCP8 Balanced recipe with hierarchical cache.""" + + gsm8k_score_threshold = 0.95 + gsm8k_num_examples = 200 + + @classmethod + def setUpClass(cls): + cls.model = MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--tp-size", + "8", + "--dcp-size", + "8", + "--disable-custom-all-reduce", + "--mem-fraction-static", + "0.85", + "--weight-loader-prefetch-checkpoints", + "--reasoning-parser", + "kimi_k3", + "--tool-call-parser", + "kimi_k3", + "--mamba-full-memory-ratio", + "7.21", + "--enable-hierarchical-cache", + ], + ) + + @classmethod + def tearDownClass(cls): + _stop_server(getattr(cls, "process", None)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/configs/test_model_config_scaling.py b/test/registered/unit/configs/test_model_config_scaling.py index f04e1681f..4d3c0fa7d 100644 --- a/test/registered/unit/configs/test_model_config_scaling.py +++ b/test/registered/unit/configs/test_model_config_scaling.py @@ -9,6 +9,15 @@ register_cpu_ci(est_time=1, suite="base-a-test-cpu") class TestMlaMscaleScaling(CustomTestCase): + def test_ignores_transformers_v5_default_rope_parameters(self): + base_scaling = 1 / math.sqrt(72) + rope_scaling = {"rope_theta": 10000.0, "rope_type": "default"} + + with self.assertNoLogs("sglang.srt.configs.model_config", level="WARNING"): + scaling = compute_mla_mscale_scaling(rope_scaling, base_scaling) + + self.assertEqual(scaling, base_scaling) + def test_respects_disabled_yarn_scaling(self): base_scaling = 1 / math.sqrt(128) rope_scaling = { @@ -34,6 +43,14 @@ class TestMlaMscaleScaling(CustomTestCase): compute_mla_mscale_scaling(rope_scaling, base_scaling), base_scaling ) + def test_applies_legacy_scaling_without_rope_type(self): + base_scaling = 1 / math.sqrt(128) + rope_scaling = {"factor": 128, "mscale_all_dim": 1} + + self.assertGreater( + compute_mla_mscale_scaling(rope_scaling, base_scaling), base_scaling + ) + def test_respects_disabled_native_apply_scale(self): base_scaling = 1 / math.sqrt(128) rope_scaling = { diff --git a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py index 0b82c263e..b871b259d 100644 --- a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py +++ b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py @@ -35,6 +35,7 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase): runner._capture_chunked_prefix = False runner.prefill_backend_name = backend runner.has_mha_companion_layers = backend == Backend.BREAKABLE + runner.mla_pinned_under_bcg = False runner.capture_hidden_mode = CaptureHiddenMode.NULL runner.capture_num_tokens = [4, 16] runner.max_num_tokens = 16 diff --git a/test/registered/unit/disaggregation/test_encoder_health.py b/test/registered/unit/disaggregation/test_encoder_health.py new file mode 100644 index 000000000..b9b64df85 --- /dev/null +++ b/test/registered/unit/disaggregation/test_encoder_health.py @@ -0,0 +1,78 @@ +import asyncio +import sys + +import pytest + +from sglang.srt.disaggregation import encode_server +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _FakeEncoder: + def __init__(self): + self.audio_processor = None + self.image_processor = object() + self.embedding_to_send = {} + self.encode_dispatch_lock = asyncio.Lock() + self.encode_calls = [] + + async def encode(self, **kwargs): + self.encode_calls.append(kwargs) + return 1, 1, 1, None, None + + +def _install_tp_encoder(monkeypatch, encoder): + broadcasts = [] + monkeypatch.setattr(encode_server, "dp_dispatcher", None) + monkeypatch.setattr(encode_server, "encoder", encoder) + monkeypatch.setattr(encode_server, "send_sockets", [object()]) + monkeypatch.setattr( + encode_server, + "sock_send", + lambda socket, payload: broadcasts.append((socket, payload)), + ) + return broadcasts + + +def test_health_encode_waits_for_collective_dispatch_lock(monkeypatch): + async def run_test(): + encoder = _FakeEncoder() + broadcasts = _install_tp_encoder(monkeypatch, encoder) + await encoder.encode_dispatch_lock.acquire() + + task = asyncio.create_task(encode_server.health_generate()) + await asyncio.sleep(0) + assert broadcasts == [] + assert encoder.encode_calls == [] + + encoder.encode_dispatch_lock.release() + response = await task + assert response.status_code == 200 + assert len(broadcasts) == 1 + assert len(encoder.encode_calls) == 1 + + asyncio.run(run_test()) + + +def test_health_encode_rechecks_busy_state_after_waiting(monkeypatch): + async def run_test(): + encoder = _FakeEncoder() + broadcasts = _install_tp_encoder(monkeypatch, encoder) + await encoder.encode_dispatch_lock.acquire() + + task = asyncio.create_task(encode_server.health_generate()) + await asyncio.sleep(0) + encoder.embedding_to_send["real-request"] = object() + encoder.encode_dispatch_lock.release() + + response = await task + assert response.status_code == 200 + assert broadcasts == [] + assert encoder.encode_calls == [] + + asyncio.run(run_test()) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/disaggregation/test_encoder_scheduler.py b/test/registered/unit/disaggregation/test_encoder_scheduler.py new file mode 100644 index 000000000..1430f31b0 --- /dev/null +++ b/test/registered/unit/disaggregation/test_encoder_scheduler.py @@ -0,0 +1,125 @@ +import asyncio +import sys + +import pytest + +from sglang.srt.disaggregation.encode_server import ( + EncoderScheduler, + PendingRequest, + _resolve_encoder_batch_policy, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _pending(modality: str = "image") -> PendingRequest: + return PendingRequest( + {"req_id": f"{modality}-request", "modality": modality}, + asyncio.get_running_loop(), + ) + + +def test_collect_batch_yields_for_concurrent_image_request_without_fixed_wait(): + # The end-to-end coalescing test cannot replace this case: asyncio.gather + # enqueues both requests within one event-loop turn, so it passes even with + # the yield removed. Only a second request enqueued from a separate task + # observes whether _collect_batch yields at all. + async def run_test(): + scheduler = EncoderScheduler( + encoder=None, + send_sockets=[], + max_batch_size=8, + coalesce_same_turn=True, + ) + first = _pending() + second = _pending() + await scheduler.pending_queue.put(first) + + async def enqueue_after_worker_yields(): + await scheduler.pending_queue.put(second) + + producer = asyncio.create_task(enqueue_after_worker_yields()) + batch = await scheduler._collect_batch() + await producer + + assert batch == [first, second] + + asyncio.run(run_test()) + + +def test_collect_batch_respects_max_batch_size(): + async def run_test(): + scheduler = EncoderScheduler( + encoder=None, + send_sockets=[], + max_batch_size=2, + coalesce_same_turn=True, + ) + requests = [_pending() for _ in range(3)] + for request in requests: + await scheduler.pending_queue.put(request) + + assert await scheduler._collect_batch() == requests[:2] + assert scheduler.pending_queue.get_nowait() is requests[2] + + asyncio.run(run_test()) + + +def test_scheduler_coalesces_concurrent_submissions(): + class FakeEncoder: + def __init__(self): + self.encode_dispatch_lock = asyncio.Lock() + self.batches = [] + + async def batch_encode(self, requests, _modality): + self.batches.append([request["req_id"] for request in requests]) + return [(1, 2, 3, None, None) for _ in requests] + + async def run_test(): + encoder = FakeEncoder() + scheduler = EncoderScheduler( + encoder=encoder, + send_sockets=[], + max_batch_size=8, + coalesce_same_turn=True, + ) + scheduler.start() + try: + requests = [ + { + "req_id": f"image-{index}", + "modality": "image", + "mm_items": [object()], + "num_parts": 1, + "part_idx": 0, + } + for index in range(2) + ] + results = await asyncio.gather( + *(scheduler.submit(request) for request in requests) + ) + finally: + await scheduler.stop() + + assert encoder.batches == [["image-0", "image-1"]] + assert results == [(1, 2, 3, None, None)] * 2 + + asyncio.run(run_test()) + + +@pytest.mark.parametrize( + ("model_type", "configured", "explicit", "expected"), + [ + ("kimi_k3", 8, False, (2, True)), + ("kimi_k3", 8, True, (8, True)), + ("kimi_k3", 1, False, (1, True)), + ("qwen3_vl", 8, False, (8, False)), + ], +) +def test_resolve_encoder_batch_policy(model_type, configured, explicit, expected): + assert _resolve_encoder_batch_policy(model_type, configured, explicit) == expected + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py b/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py new file mode 100644 index 000000000..2210d06b5 --- /dev/null +++ b/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py @@ -0,0 +1,401 @@ +import asyncio +import pickle +import sys +import threading +import time +from array import array +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch +import zmq +import zmq.asyncio +from fastapi import HTTPException +from PIL import Image + +from sglang.srt.disaggregation.encode_receiver import ( + EmbeddingData, + MMReceiverHTTP, + MultiModalEmbeddingData, + _select_mm_processor_prompt, +) +from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim +from sglang.srt.managers.schedule_batch import Modality +from sglang.srt.managers.tokenizer_manager import ( + _reject_missing_dispatched_encoder_embedding, +) +from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration +from sglang.srt.runtime_context import get_context +from sglang.srt.server_args import resolve_encoder_transfer_backend +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def test_kimi_k3_encoder_transfer_backend_auto_avoids_tp_fanout(): + assert ( + resolve_encoder_transfer_backend("auto", "KimiK3ForConditionalGeneration", 8) + == "zmq_to_tokenizer" + ) + assert ( + resolve_encoder_transfer_backend("auto", "KimiK3ForConditionalGeneration", 1) + == "zmq_to_scheduler" + ) + assert ( + resolve_encoder_transfer_backend("auto", "Qwen3VLForConditionalGeneration", 8) + == "zmq_to_scheduler" + ) + assert ( + resolve_encoder_transfer_backend( + "zmq_to_scheduler", "KimiK3ForConditionalGeneration", 8 + ) + == "zmq_to_scheduler" + ) + assert ( + resolve_encoder_transfer_backend( + "mooncake", "KimiK3ForConditionalGeneration", 8 + ) + == "mooncake" + ) + + +def test_epd_language_only_rejects_missing_dispatched_embedding(): + server_args = SimpleNamespace( + language_only=True, + encoder_transfer_backend="zmq_to_tokenizer", + ) + request = SimpleNamespace(need_wait_for_mm_inputs=True) + + with pytest.raises(HTTPException) as exc_info: + _reject_missing_dispatched_encoder_embedding(server_args, request, None) + + assert getattr(exc_info.value, "status_code", None) == 503 + + +def test_epd_allows_local_processing_when_request_was_not_dispatched(): + server_args = SimpleNamespace( + language_only=True, + encoder_transfer_backend="zmq_to_tokenizer", + ) + request = SimpleNamespace(need_wait_for_mm_inputs=False) + + _reject_missing_dispatched_encoder_embedding(server_args, request, None) + + +def _encoder(model_type="kimi_k3"): + encoder = MMEncoder.__new__(MMEncoder) + encoder.model_type = model_type + encoder.model_config = SimpleNamespace( + hf_config=SimpleNamespace( + vision_config=SimpleNamespace(merge_kernel_size=(2, 2)) + ) + ) + return encoder + + +def test_kimi_k3_encoder_normalizes_pillow_images_to_media_dicts(): + image = Image.new("RGB", (2, 2)) + encoder = _encoder() + + assert encoder._grid_count_per_leaf( + [image, {"type": "image", "image": [image, image]}], Modality.IMAGE + ) == [1, 2] + + normalized = encoder._normalize_kimi_encoder_images( + [image, {"type": "image", "image": [image, image]}] + ) + assert len(normalized) == 3 + assert all(item["type"] == "image" for item in normalized) + assert all(item["image"] is image for item in normalized) + + +def test_kimi_k3_encoder_passes_media_dicts_to_image_processor(): + image = Image.new("RGB", (3, 2)) + processor_calls = [] + + def image_processor(*, images, **kwargs): + processor_calls.append((images, kwargs)) + return {"pixel_values": torch.ones(1, 3), "grid_thws": [[1, 1, 1]]} + + encoder = _encoder() + encoder.image_processor = image_processor + encoder.vision_config = {"image": {"return_tensors": "pt"}} + encoder._flatten_and_load_images = AsyncMock(return_value=[image]) + encoder.preproc_executor = ThreadPoolExecutor(max_workers=1) + try: + output = asyncio.run(encoder._process_image_items([image], None)) + finally: + encoder.preproc_executor.shutdown() + + assert "pixel_values" in output + assert output["original_image_sizes"] == [[3, 2]] + assert len(processor_calls) == 1 + images, kwargs = processor_calls[0] + assert images[0]["type"] == "image" + assert images[0]["image"] is image + assert kwargs == {"return_tensors": "pt"} + + +def test_kimi_k3_epd_aggregates_original_image_sizes_in_part_order(): + first = EmbeddingData( + req_id="request", + num_parts=2, + part_idx=0, + grid_dim=torch.tensor([[1, 2, 6]]), + modality=Modality.IMAGE, + embedding=torch.ones(3, 4), + original_image_sizes=[[1536, 1024]], + ) + second = EmbeddingData( + req_id="request", + num_parts=2, + part_idx=1, + grid_dim=torch.tensor([[1, 2, 4]]), + modality=Modality.IMAGE, + embedding=torch.ones(2, 4), + original_image_sizes=[[1024, 1536]], + ) + + combined = MultiModalEmbeddingData.from_embedding_data(first, model_type="kimi_k3") + combined.add(second) + + assert combined.ready + assert combined.get_mm_extra_meta()["original_image_sizes"] == [ + [1536, 1024], + [1024, 1536], + ] + + +def test_kimi_k3_encoder_prefers_grid_thws_and_uses_temporal_pool_length(): + grid_thws = torch.tensor([[3, 8, 12]]) + stale_grid = torch.tensor([[1, 2, 2]]) + mm_inputs = {"grid_thws": grid_thws, "image_grid_thw": stale_grid} + + assert _get_mm_grid_dim(mm_inputs, Modality.IMAGE, "kimi_k3") is grid_thws + assert _encoder().get_num_tokens(grid_thws[0], Modality.IMAGE) == 24 + + +def test_kimi_k3_encoder_splits_cross_request_batch_into_single_grid_items(): + encoder = _encoder() + grid_thws = torch.tensor([[1, 2, 2], [2, 2, 4], [1, 4, 2]]) + feature = torch.arange(56, dtype=torch.float32).reshape(28, 2) + embeddings = torch.arange(15, dtype=torch.float32).reshape(5, 3) + captured = {} + + def get_feature_fn(items): + captured["items"] = items + return embeddings + + output = encoder._encode_missing( + feature, + {"pixel_values": feature, "grid_thws": grid_thws}, + indices=[2, 0, 1], + modality=Modality.IMAGE, + get_feature_fn=get_feature_fn, + grid_thw=grid_thws, + keep_on_gpu=True, + ) + + items = captured["items"] + assert len(items) == 3 + expected_feature_slices = [feature[20:28], feature[0:4], feature[4:20]] + expected_grids = [grid_thws[2:3], grid_thws[0:1], grid_thws[1:2]] + for item, expected_feature, expected_grid in zip( + items, expected_feature_slices, expected_grids + ): + torch.testing.assert_close(item.feature, expected_feature) + torch.testing.assert_close(item.model_specific_data["grid_thws"], expected_grid) + + assert [embedding.shape[0] for embedding in output] == [2, 1, 2] + torch.testing.assert_close(torch.cat(output), embeddings) + + +def test_kimi_k3_encoder_only_wrapper_guards_language_tower_hooks(): + model = SimpleNamespace(language_model=None) + + KimiK3ForConditionalGeneration.post_load_weights(model) + with pytest.raises(AttributeError, match="lm_head"): + KimiK3ForConditionalGeneration.lm_head.fget(model) + with pytest.raises(AttributeError, match="DSPARK"): + KimiK3ForConditionalGeneration.set_dspark_layers_to_capture(model, [0]) + + +def test_epd_scheduler_uses_token_ids_for_tokenized_mm_processors(): + recv_req = SimpleNamespace( + input_text="unexpanded prompt", input_ids=array("q", [11, 22, 33]) + ) + + prompt = _select_mm_processor_prompt( + recv_req, SimpleNamespace(prefer_tokenized_input=True) + ) + + assert prompt == [11, 22, 33] + assert isinstance(prompt, list) + assert ( + _select_mm_processor_prompt( + recv_req, SimpleNamespace(prefer_tokenized_input=False) + ) + == "unexpanded prompt" + ) + + +def test_epd_scheduler_routes_many_requests_over_one_receive_socket(): + context = zmq.Context() + receiver = MMReceiverHTTP.__new__(MMReceiverHTTP) + receiver.scheduler_recv_socket = context.socket(zmq.PULL) + port = receiver.scheduler_recv_socket.bind_to_random_port("tcp://127.0.0.1") + received = [] + + class Sink: + def consume_parts(self, parts): + received.append(pickle.loads(parts[0]).req_id) + + receiver.waiting_by_rid = {f"rid-{i}": Sink() for i in range(32)} + sender = context.socket(zmq.PUSH) + try: + sender.connect(f"tcp://127.0.0.1:{port}") + for i in range(32): + mm_data = EmbeddingData( + req_id=f"rid-{i}_local_part_0", + num_parts=1, + part_idx=0, + grid_dim=None, + modality=Modality.IMAGE, + error_msg="probe", + error_code=599, + ) + sender.send_multipart([pickle.dumps(mm_data)]) + + deadline = time.monotonic() + 2 + while len(received) < 32 and time.monotonic() < deadline: + receiver._drain_scheduler_embeddings() + time.sleep(0.01) + assert received == [f"rid-{i}_local_part_0" for i in range(32)] + finally: + sender.close(linger=0) + receiver.scheduler_recv_socket.close(linger=0) + context.term() + + +def test_epd_encoder_reuses_scheduler_zmq_peer(): + async def send_twice(): + context = zmq.asyncio.Context() + receiver = context.socket(zmq.PULL) + port = receiver.bind_to_random_port("tcp://127.0.0.1") + encoder = MMEncoder.__new__(MMEncoder) + config_override = get_context().override_server_args( + encoder_transfer_backend="zmq_to_scheduler" + ) + with config_override as server_args: + encoder.server_args = server_args + encoder.send_timeout = 3 + encoder.context = context + encoder.scheduler_send_sockets = {} + encoder.scheduler_send_locks = {} + mm_data = EmbeddingData( + req_id="test-rid_local_part_0", + num_parts=1, + part_idx=0, + grid_dim=None, + modality=Modality.IMAGE, + error_msg="probe", + error_code=599, + ) + try: + for _ in range(2): + await encoder._send(None, mm_data, url=f"127.0.0.1:{port}") + parts = await asyncio.wait_for(receiver.recv_multipart(), timeout=1) + assert pickle.loads(parts[0]).req_id == mm_data.req_id + assert len(encoder.scheduler_send_sockets) == 1 + finally: + for socket in encoder.scheduler_send_sockets.values(): + socket.close(linger=0) + receiver.close(linger=0) + context.term() + + asyncio.run(send_twice()) + + +def test_epd_encoder_pipelines_zero_copy_sends_per_peer(): + class FakeTracker: + def __init__(self, release): + self.release = release + + def wait(self, timeout): + assert self.release.wait(timeout) + + class FakeSocket: + def __init__(self, release, second_queued): + self.release = release + self.second_queued = second_queued + self.send_count = 0 + + def setsockopt(self, *_args): + pass + + def connect(self, _endpoint): + pass + + def close(self, **_kwargs): + pass + + async def send_multipart(self, _frames, **_kwargs): + self.send_count += 1 + if self.send_count == 2: + self.second_queued.set() + return FakeTracker(self.release) + + class FakeContext: + def __init__(self, socket): + self.socket_instance = socket + + def socket(self, _socket_type): + return self.socket_instance + + async def run_test(): + release = threading.Event() + second_queued = asyncio.Event() + socket = FakeSocket(release, second_queued) + encoder = MMEncoder.__new__(MMEncoder) + config_override = get_context().override_server_args( + encoder_transfer_backend="zmq_to_scheduler" + ) + with config_override as server_args: + encoder.server_args = server_args + encoder.send_timeout = 1 + encoder.context = FakeContext(socket) + encoder.scheduler_send_sockets = {} + encoder.scheduler_send_locks = {} + mm_data = EmbeddingData( + req_id="test-rid_local_part_0", + num_parts=1, + part_idx=0, + grid_dim=None, + modality=Modality.IMAGE, + error_msg="probe", + error_code=599, + ) + + first = asyncio.create_task( + encoder._send(None, mm_data, url="127.0.0.1:12345") + ) + while socket.send_count < 1: + await asyncio.sleep(0) + second = asyncio.create_task( + encoder._send(None, mm_data, url="127.0.0.1:12345") + ) + try: + await asyncio.wait_for(second_queued.wait(), timeout=0.5) + finally: + release.set() + await asyncio.gather(first, second) + + assert socket.send_count == 2 + + asyncio.run(run_test()) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 2496f7def..90d58f91b 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -269,6 +269,29 @@ class ServingChatTestCase(unittest.TestCase): self.assertEqual(adapted.sampling_params["stop"], ["STOP"]) conv_mock.assert_not_called() + def test_kimi_k3_usage_excludes_assistant_generation_stub(self): + self.chat.chat_encoding_spec = "kimi_k3" + ret = [ + { + "text": "Answer", + "meta_info": { + "id": "chatcmpl-kimi-k3-usage", + "prompt_tokens": 2075, + "completion_tokens": 1, + "cached_tokens": 0, + "image_tokens": 2035, + "finish_reason": {"type": "stop", "matched": None}, + "weight_version": "default", + }, + } + ] + + response = self.chat._build_chat_response(self.basic_req, ret, created=123) + + self.assertEqual(response.usage.prompt_tokens, 2072) + self.assertEqual(response.usage.total_tokens, 2073) + self.assertEqual(response.usage.prompt_tokens_details.image_tokens, 2035) + def test_kimi_tool_call_keeps_default_reasoning(self): self.template_manager.reasoning_config = ReasoningToggleConfig( toggle_param="thinking", default_enabled=True diff --git a/test/registered/unit/layers/attention/linear/kernels/test_kda_nvidia.py b/test/registered/unit/layers/attention/linear/kernels/test_kda_nvidia.py new file mode 100644 index 000000000..0859e8c8f --- /dev/null +++ b/test/registered/unit/layers/attention/linear/kernels/test_kda_nvidia.py @@ -0,0 +1,202 @@ +"""Unit tests for the NVIDIA KDA prefill routing/repacking wrapper.""" + +import unittest +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.layers.attention.linear.kernels.kda_nvidia import ( + NvidiaKDAKernel, + _from_nvidia_kda_state_layout, + _to_nvidia_kda_state_layout, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class _RejectTriton: + def extend(self, *args, **kwargs): + raise AssertionError("ordinary prefill unexpectedly fell back to Triton") + + +class TestNvidiaKDAAllPrefillWrapper(CustomTestCase): + def test_state_layout_round_trip(self): + state = torch.arange(2 * 3 * 5 * 7, dtype=torch.float32).view(2, 3, 5, 7) + + nvidia_kda_state = _to_nvidia_kda_state_layout( + state, head_k_dim=7, head_v_dim=5 + ) + + self.assertEqual(tuple(nvidia_kda_state.shape), (2, 3, 7, 5)) + self.assertTrue(nvidia_kda_state.is_contiguous()) + self.assertEqual(nvidia_kda_state[1, 2, 6, 4], state[1, 2, 4, 6]) + + restored = _from_nvidia_kda_state_layout( + nvidia_kda_state, + head_k_dim=7, + head_v_dim=5, + dtype=torch.bfloat16, + ) + + self.assertEqual(tuple(restored.shape), (2, 3, 5, 7)) + self.assertTrue(restored.is_contiguous()) + self.assertTrue(torch.equal(restored, state.to(torch.bfloat16))) + + def test_state_layout_rejects_swapped_contract(self): + with self.assertRaisesRegex(ValueError, "SGLang KDA state"): + _to_nvidia_kda_state_layout( + torch.zeros(1, 2, 7, 5), + head_k_dim=7, + head_v_dim=5, + ) + + def _make_kernel(self): + calls = [] + kernel = NvidiaKDAKernel() + kernel._l2norm = lambda x: x + kernel._triton = _RejectTriton() + + def fake_fwd(q, k, v, g, beta, **kwargs): + calls.append( + { + "q": q.clone(), + "k": k.clone(), + "v": v.clone(), + "g": g.clone(), + "beta": beta.clone(), + "initial_state": kwargs["initial_state"].clone(), + "cu_seqlens": kwargs["cu_seqlens"], + } + ) + return v.clone(), kwargs["initial_state"] + 1.0 + + kernel._fwd = fake_fwd + return kernel, calls + + @staticmethod + def _inputs(seq_lens): + total = sum(seq_lens) + token_values = torch.arange(total * 128, dtype=torch.bfloat16).view( + 1, total, 1, 128 + ) + query_start_loc = torch.tensor( + [0] + list(torch.tensor(seq_lens).cumsum(0).tolist()), dtype=torch.int32 + ) + return { + "q": token_values + 100, + "k": token_values + 200, + "v": token_values + 300, + "g": (token_values + 400).view(1, total, 128), + "beta": torch.arange(total, dtype=torch.float32).view(1, total, 1), + "query_start_loc": query_start_loc, + } + + def test_short_single_sequence_uses_triton(self): + kernel, calls = self._make_kernel() + kernel._triton.extend = Mock(return_value="triton") + x = self._inputs([5]) + states = torch.zeros(3, 1, 128, 128, dtype=torch.bfloat16) + + output = kernel.extend( + x["q"], + x["k"], + x["v"], + x["g"], + x["beta"], + ssm_states=states, + cache_indices=torch.tensor([1], dtype=torch.int32), + query_start_loc=x["query_start_loc"], + extend_seq_lens_cpu=[5], + A_log=torch.zeros(128, dtype=torch.float32), + ) + + self.assertEqual(output, "triton") + self.assertEqual(calls, []) + kernel._triton.extend.assert_called_once() + self.assertTrue(torch.count_nonzero(states).item() == 0) + + def test_packed_multi_sequence_repacking_preserves_order_and_slots(self): + kernel, calls = self._make_kernel() + seq_lens = [2, 3, 1] + x = self._inputs(seq_lens) + states = torch.arange(5 * 128 * 128, dtype=torch.bfloat16).view(5, 1, 128, 128) + states_before = states.clone() + slots = torch.tensor([2, 0, 4], dtype=torch.int32) + + output = kernel.extend( + x["q"], + x["k"], + x["v"], + x["g"], + x["beta"], + ssm_states=states, + cache_indices=slots, + query_start_loc=x["query_start_loc"], + extend_seq_lens_cpu=seq_lens, + A_log=torch.zeros(128, dtype=torch.float32), + ) + + self.assertEqual(len(calls), 1) + call = calls[0] + self.assertEqual(tuple(call["q"].shape), (3, 2048, 1, 128)) + self.assertIsNone(call["cu_seqlens"]) + self.assertTrue(torch.equal(output, x["v"])) + + start = 0 + for row, length in enumerate(seq_lens): + end = start + length + self.assertTrue(torch.equal(call["v"][row, :length], x["v"][0, start:end])) + self.assertTrue(torch.count_nonzero(call["q"][row, length:]).item() == 0) + self.assertTrue(torch.count_nonzero(call["k"][row, length:]).item() == 0) + self.assertTrue(torch.count_nonzero(call["v"][row, length:]).item() == 0) + self.assertTrue(torch.count_nonzero(call["beta"][row, length:]).item() == 0) + self.assertTrue(torch.all(call["g"][row, length:] == -1000)) + start = end + + for slot in slots.tolist(): + self.assertTrue(torch.equal(states[slot], states_before[slot] + 1)) + untouched = {0, 1, 2, 3, 4} - set(slots.tolist()) + for slot in untouched: + self.assertTrue(torch.equal(states[slot], states_before[slot])) + + def test_non_fp32_beta_falls_back_to_triton(self): + kernel, calls = self._make_kernel() + x = self._inputs([2, 3]) + x["beta"] = x["beta"].bfloat16() + states = torch.zeros(3, 1, 128, 128, dtype=torch.bfloat16) + + with self.assertRaisesRegex( + AssertionError, "ordinary prefill unexpectedly fell back to Triton" + ): + kernel.extend( + x["q"], + x["k"], + x["v"], + x["g"], + x["beta"], + ssm_states=states, + cache_indices=torch.tensor([1, 2], dtype=torch.int32), + query_start_loc=x["query_start_loc"], + extend_seq_lens_cpu=[2, 3], + A_log=torch.zeros(128, dtype=torch.float32), + ) + self.assertEqual(calls, []) + + def test_supports_only_datacenter_blackwell(self): + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=(10, 0)), + ): + self.assertTrue(NvidiaKDAKernel().supports_prefill) + + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=(12, 0)), + ): + self.assertFalse(NvidiaKDAKernel().supports_prefill) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/attention/test_linear_attn_config.py b/test/registered/unit/layers/attention/test_linear_attn_config.py index 41d78a788..d0af273b4 100644 --- a/test/registered/unit/layers/attention/test_linear_attn_config.py +++ b/test/registered/unit/layers/attention/test_linear_attn_config.py @@ -21,18 +21,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") class TestLinearAttnConfig(CustomTestCase): def setUp(self): - saved = ( - linear_utils.LINEAR_ATTN_DECODE_BACKEND, - linear_utils.LINEAR_ATTN_PREFILL_BACKEND, - ) - - def restore(): - ( - linear_utils.LINEAR_ATTN_DECODE_BACKEND, - linear_utils.LINEAR_ATTN_PREFILL_BACKEND, - ) = saved - - self.addCleanup(restore) + self.addCleanup(linear_utils._BACKENDS.update, linear_utils._BACKENDS.copy()) def _init(self, prefill_default=None, **fields): args = ServerArgs(model_path="dummy") @@ -40,8 +29,8 @@ class TestLinearAttnConfig(CustomTestCase): setattr(args, key, value) initialize_linear_attn_config(args, prefill_default) return ( - linear_utils.LINEAR_ATTN_PREFILL_BACKEND, - linear_utils.LINEAR_ATTN_DECODE_BACKEND, + linear_utils.get_linear_attn_prefill_backend(), + linear_utils.get_linear_attn_decode_backend(), ) def test_default_applies_when_the_flag_is_unset(self): diff --git a/test/registered/unit/layers/attention/test_vision_max_seqlen.py b/test/registered/unit/layers/attention/test_vision_max_seqlen.py index 27e9e39f2..84c5a6e74 100644 --- a/test/registered/unit/layers/attention/test_vision_max_seqlen.py +++ b/test/registered/unit/layers/attention/test_vision_max_seqlen.py @@ -1,9 +1,11 @@ import sys +import pytest import torch from torch import nn from sglang.srt.layers.attention import vision +from sglang.srt.models import kimi_k25 from sglang.srt.models.kimi_k25 import MoonViT3dEncoder, MoonViTEncoderLayer from sglang.test.ci.ci_register import register_cpu_ci @@ -166,6 +168,7 @@ def test_kimi_moonvit_precomputes_sequence_lengths_once(): encoder = MoonViT3dEncoder.__new__(MoonViT3dEncoder) nn.Module.__init__(encoder) encoder.rope_2d = CapturingRope() + encoder.use_fused_rope = False encoder.blocks = nn.ModuleList([CapturingBlock()]) encoder.final_layernorm = nn.Identity() @@ -179,6 +182,59 @@ def test_kimi_moonvit_precomputes_sequence_lengths_once(): assert recorded["max_seqlen"] == 4 +def test_kimi_moonvit_prepares_cuda_rope_inputs_once(): + recorded = {} + + class CapturingRope: + def get_freqs_cis(self, grid_thws, device): + real = torch.arange(14, dtype=torch.float32, device=device).view(7, 2) + return torch.complex(real, real + 1) + + class CapturingBlock(nn.Module): + def forward( + self, + hidden_states, + cu_seqlens, + max_seqlen, + rope_freqs_cis, + **kwargs, + ): + recorded["rope_freqs_cis"] = rope_freqs_cis + return hidden_states + + encoder = MoonViT3dEncoder.__new__(MoonViT3dEncoder) + nn.Module.__init__(encoder) + encoder.rope_2d = CapturingRope() + encoder.use_fused_rope = True + encoder.blocks = nn.ModuleList([CapturingBlock()]) + encoder.final_layernorm = nn.Identity() + + # The fused path is gated on the q/k dtype; fp32 stays on the portable one. + hidden_states = torch.ones(7, 4, dtype=torch.bfloat16) + encoder(hidden_states, torch.tensor([[1, 1, 7]], dtype=torch.int32)) + + cos_sin_cache, positions = recorded["rope_freqs_cis"] + assert cos_sin_cache.shape == (7, 4) + assert torch.equal(cos_sin_cache[:, :2] + 1, cos_sin_cache[:, 2:]) + assert torch.equal(positions, torch.arange(7)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_kimi_moonvit_fused_rope_matches_portable_path(): + torch.manual_seed(0) + q = torch.randn(256, 4, 72, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + angles = torch.randn(256, 36, device="cuda", dtype=torch.float32) + freqs_cis = torch.polar(torch.ones_like(angles), angles) + + q_ref, k_ref = kimi_k25.apply_rope(q.clone(), k.clone(), freqs_cis) + prepared_rope = kimi_k25.prepare_fused_qk_complex_rope_inplace(freqs_cis) + q_fused, k_fused = kimi_k25.apply_rope(q.clone(), k.clone(), prepared_rope) + + torch.testing.assert_close(q_fused, q_ref, rtol=0.01, atol=0.01) + torch.testing.assert_close(k_fused, k_ref, rtol=0.01, atol=0.01) + + if __name__ == "__main__": import pytest diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py index 2fd5cf637..4d1c0ae7e 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py @@ -48,6 +48,8 @@ from flashinfer.fused_moe import ( ) from flashinfer.fused_moe.core import ActivationType +from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig + GROUP_SIZE = 32 # MXFP4 block size @@ -58,6 +60,12 @@ class _MockLayer: path (``get_tp_group`` etc.). """ + def __init__(self): + # The SM90 weight-processing path reads the runner config for the + # gate/up row layout (``gate_up_interleaved``) and the activation. A + # real ``FusedMoE`` always carries one, so the stand-in does too. + self.moe_runner_config = MoeRunnerConfig() + class _MockTopKOutput: def __init__(self, weights, ids): diff --git a/test/registered/unit/managers/test_mm_hashes.py b/test/registered/unit/managers/test_mm_hashes.py index 85fb500ec..648f06a32 100644 --- a/test/registered/unit/managers/test_mm_hashes.py +++ b/test/registered/unit/managers/test_mm_hashes.py @@ -77,6 +77,15 @@ class TestMmHashesContract(CustomTestCase): b.set_pad_value() self.assertNotEqual(a.pad_value, b.pad_value) + def test_set_hash_updates_an_existing_pad_value(self): + item = MultimodalDataItem(modality=Modality.IMAGE, hash=0xAAAA) + item.set_pad_value() + + item.set_hash(0xBBBB) + + self.assertEqual(item.hash, 0xBBBB) + self.assertEqual(item.pad_value, _compute_pad_value(0xBBBB)) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py b/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py new file mode 100644 index 000000000..7860e8791 --- /dev/null +++ b/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py @@ -0,0 +1,233 @@ +"""FlashKDA prefill wrapper vs envelope-strided Mamba state pools (CPU). + +Derived property under test: ``FlashKDAKernel`` (the wrapper around the +external, contiguous-only ``flash_kda`` CUTLASS kernel) touches the SSM state +pool ONLY through torch advanced indexing — a gather into a contiguous local +copy before the kernel and a scatter write-back after. Advanced indexing is +layout-agnostic, so the wrapper works unchanged on the envelope-strided +temporal views used by --enable-page-major-kv-layout / --enable-unified-memory +(slot pitch == the multi-layer entry envelope, NOT H*V*K). This insulation is +the justification for allowing prefill=flashkda under the page-major backend +gate without ever teaching the external kernel about strides. + +What turns this red: any "optimization" that hands the pool view to +``flash_kda.fwd`` directly, replaces the gather with a ``.view()`` / pointer +reshape that assumes the contiguous slot pitch, or drops the scatter +write-back. On a contiguous pool such a change is invisible; on the strided +pool it mis-addresses state exactly like the chunk_delta_h hardcoded-pitch bug +(GSM8K 0.17). + +Runs on CPU — the external kernel is replaced by a stub; only the pool access +pattern (the code under test) executes. + + python -m pytest test/registered/unit/mem_cache/test_flashkda_strided_state_access.py -v +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") + +import sys +import types +import unittest + +import torch + +from sglang.srt.layers.attention.linear.kernels.kda_flashkda import FlashKDAKernel +from sglang.srt.mem_cache.layout.page_major import ( + build_page_major_mamba_views, + mamba_entry_bytes, +) + +_DEV = "cpu" + +# Tiny KDA-like geometry (multi-layer so the envelope slot pitch != H*V*K). +_LAYERS = 3 +_LAYER_UNDER_TEST = 1 +_H = 2 +_K = 4 +_V = 4 +_SLOTS = 8 +_CONV_SHAPES = ( + (3, 8), +) # KDA conv layout [kernel-1, dim]; bf16 region pads the envelope +_CONV_DTYPE = torch.bfloat16 +_TEMPORAL_DTYPE = torch.float32 +# FlashKDA fused-path window: per-seq len must be in [chunk_size, max_seq_len]. +_SEQ_LEN = 128 + + +def _make_strided_temporal_views(): + """Envelope-strided conv/temporal views, as UnifiedMambaPool / the + page-major MambaPool serve them ((num_layers, max_slots, *inner)).""" + entry = mamba_entry_bytes( + layer_num=_LAYERS, + conv_state_shapes=_CONV_SHAPES, + conv_dtype=_CONV_DTYPE, + temporal_state_shape=(_H, _V, _K), + temporal_dtype=_TEMPORAL_DTYPE, + ) + raw = torch.zeros(_SLOTS * entry, dtype=torch.uint8, device=_DEV) + conv_views, temporal = build_page_major_mamba_views( + raw, + layer_num=_LAYERS, + conv_state_shapes=_CONV_SHAPES, + conv_dtype=_CONV_DTYPE, + temporal_state_shape=(_H, _V, _K), + temporal_dtype=_TEMPORAL_DTYPE, + max_slots=_SLOTS, + ) + return conv_views, temporal + + +class _FakeFlashKDA: + """Stand-in for the external ``flash_kda`` module. Records what the wrapper + hands it and applies a deterministic state update so the write-back is + checkable: final = 2 * initial + 1.""" + + def __init__(self): + self.calls = 0 + self.initial_state_was_contiguous = None + self.initial_state_copy = None + + def fwd( + self, + q, + k, + v, + g, + beta, + scale, + out_buf, + A_log, + dt_bias, + lower_bound, + *, + initial_state, + final_state, + cu_seqlens, + ): + self.calls += 1 + self.initial_state_was_contiguous = initial_state.is_contiguous() + self.initial_state_copy = initial_state.clone() + final_state.copy_(initial_state * 2.0 + 1.0) + out_buf.fill_(0.25) + + +class TestFlashKDAStridedStateAccess(unittest.TestCase): + def setUp(self): + self._saved_module = sys.modules.get("flash_kda") + self.fake = _FakeFlashKDA() + mod = types.ModuleType("flash_kda") + mod.fwd = self.fake.fwd + sys.modules["flash_kda"] = mod + + def tearDown(self): + if self._saved_module is None: + sys.modules.pop("flash_kda", None) + else: + sys.modules["flash_kda"] = self._saved_module + + def _run_extend(self, ssm_states, cache_indices): + num_seqs = cache_indices.numel() + packed = num_seqs * _SEQ_LEN + torch.manual_seed(0) + q = torch.randn(1, packed, _H, _K, dtype=torch.bfloat16) + k = torch.randn(1, packed, _H, _K, dtype=torch.bfloat16) + v = torch.randn(1, packed, _H, _V, dtype=torch.bfloat16) + g = torch.randn(1, packed, _H, _K, dtype=torch.bfloat16) + beta = torch.rand(1, packed, _H, dtype=torch.bfloat16) * 0.8 + 0.1 + query_start_loc = torch.arange(0, packed + 1, _SEQ_LEN, dtype=torch.int32) + return FlashKDAKernel().extend( + q, + k, + v, + g, + beta, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + A_log=torch.randn(1, 1, _H, 1), + dt_bias=torch.randn(_H * _K), + lower_bound=-10.0, # safe gate => fused path (no triton fallback) + extend_seq_lens_cpu=[_SEQ_LEN] * num_seqs, + ) + + def test_gather_kernel_scatter_on_envelope_strided_pool(self): + conv_views, temporal = _make_strided_temporal_views() + ssm_states = temporal[_LAYER_UNDER_TEST] # what mamba2_layer_cache serves + + # Precondition of the property: the pool really is envelope-strided. + self.assertNotEqual( + ssm_states.stride(0), + _H * _V * _K, + "test setup no longer produces a strided pool; the property below " + "would be vacuous", + ) + + # Distinct value per (layer, slot); sentinel in the conv regions that + # interleave the temporal regions inside each slot envelope. + seed = ( + torch.arange(_LAYERS, dtype=torch.float32)[:, None] * 100.0 + + torch.arange(_SLOTS, dtype=torch.float32)[None, :] + ) + temporal[:] = seed.view(_LAYERS, _SLOTS, 1, 1, 1) + 1.0 + for cv in conv_views: + cv.fill_(3.0) + temporal_before = temporal.clone() + conv_before = [cv.clone() for cv in conv_views] + + cache_indices = torch.tensor([5, 2], dtype=torch.int32) + out = self._run_extend(ssm_states, cache_indices) + + # Routing: the fused path ran exactly once (a silent re-route to the + # triton fallback would make every assertion below vacuous). + self.assertEqual(self.fake.calls, 1) + self.assertEqual(tuple(out.shape), (1, 2 * _SEQ_LEN, _H, _V)) + + # Gather: the external kernel must receive a CONTIGUOUS copy whose rows + # are the addressed slots of the strided pool. + self.assertTrue(self.fake.initial_state_was_contiguous) + self.assertTrue( + torch.equal( + self.fake.initial_state_copy, + temporal_before[_LAYER_UNDER_TEST][cache_indices.long()], + ), + "gather mis-addressed the envelope-strided slots", + ) + + # Scatter: the committed state lands in exactly the addressed slots. + expected = temporal_before[_LAYER_UNDER_TEST][cache_indices.long()] * 2.0 + 1.0 + self.assertTrue( + torch.equal(ssm_states[cache_indices.long()], expected), + "write-back mis-addressed the envelope-strided slots", + ) + + # Isolation: untouched slots of this layer, ALL slots of the other + # layers, and the interleaved conv regions are byte-identical. A + # contiguous-pitch (H*V*K) access pattern would corrupt these. + touched = torch.zeros(_SLOTS, dtype=torch.bool) + touched[cache_indices.long()] = True + self.assertTrue( + torch.equal( + ssm_states[~touched], + temporal_before[_LAYER_UNDER_TEST][~touched], + ), + "write-back leaked into unaddressed slots", + ) + for layer in range(_LAYERS): + if layer == _LAYER_UNDER_TEST: + continue + self.assertTrue( + torch.equal(temporal[layer], temporal_before[layer]), + f"write-back leaked into layer {layer}'s envelope region", + ) + for cv, before in zip(conv_views, conv_before): + self.assertTrue( + torch.equal(cv, before), + "write-back leaked into the conv region of the slot envelope", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py b/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py index 6376126f5..6bcf78c4b 100644 --- a/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py +++ b/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py @@ -3,14 +3,18 @@ The memory solver charges this on top of mamba_cache_per_req so num_slots is not over-provisioned (the ring is allocated per slot but is NOT part of the state cache cost). Pins the arithmetic against hand-computed byte counts for the -fold window (raw v / pre-norm k / g / beta). If the MambaPool allocation -changes shape, update both together. +fold window (raw v / pre-norm k / g / beta) across both gate layouts: GDN +per-head scalar g vs KDA per-K vector g (KDA also keeps the chunked d/k rings +under spec, see MambaPool). If the MambaPool allocation changes shape, update +both the allocation and this expectation together. """ import pytest import torch from sglang.srt.configs.mamba_utils import ( + KimiLinearCacheParams, + KimiLinearStateShape, Mamba2CacheParams, Mamba2StateDType, Mamba2StateShape, @@ -23,14 +27,22 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") # temporal = (hv=4, v_dim=8, k_dim=8), num_k_heads_per_tp = 4, record_len = 8, # 2 layers. conv bf16 (2B), fp32 gate/beta (4B). Ring tensors (per slot, per # layer): -# rawv hv*RL*v_dim, rawk h_k*RL*k_dim -> conv dtype -# g hv*RL -> fp32 -# beta hv*RL -> fp32 +# rawv hv*RL*v_dim, rawk h_k*RL*k_dim -> conv dtype +# g hv*RL (GDN) / hv*RL*k_dim (KDA) -> fp32 +# beta hv*RL -> fp32 +# d/k like rawv/rawk -> conv dtype (KDA only) DTYPE = Mamba2StateDType(conv=torch.bfloat16, temporal=torch.float32) RL = 8 LAYERS = [0, 1] +def _kda_params(): + shape = KimiLinearStateShape.create( + tp_world_size=1, num_heads=4, head_dim=8, num_k_heads=4, head_k_dim=8 + ) + return KimiLinearCacheParams(shape=shape, dtype=DTYPE, layers=LAYERS) + + def _gdn_params(): # Only shape.temporal and shape.num_k_heads_per_tp are read here; the rest # are dummy (the accounting does not depend on them). @@ -57,8 +69,17 @@ class TestReplaySSMRingAccounting(CustomTestCase): 1280 * len(LAYERS), ) + def test_kda_fold(self): + # rawv 512 + rawk 512 + g(per-K, 4*8*8*4) 1024 + beta 128 + # + d 512 + k 512 (KDA keeps the chunked rings under spec) = 3200 + self.assertEqual( + _kda_params().replayssm_ring_bytes_per_req(record_len=RL), + 3200 * len(LAYERS), + ) + def test_zero_len_ring(self): self.assertEqual(_gdn_params().replayssm_ring_bytes_per_req(record_len=0), 0) + self.assertEqual(_kda_params().replayssm_ring_bytes_per_req(record_len=0), 0) if __name__ == "__main__": diff --git a/test/registered/unit/mem_cache/test_unified_mamba_views.py b/test/registered/unit/mem_cache/test_unified_mamba_views.py index e234ab574..b7170b06b 100644 --- a/test/registered/unit/mem_cache/test_unified_mamba_views.py +++ b/test/registered/unit/mem_cache/test_unified_mamba_views.py @@ -37,7 +37,9 @@ These tests prove the views: (the shape `MambaPool.State.conv[i]` / `.temporal` expose); - reject a deliberately mis-aligned spec via the alignment assert. -Skipped on CPU — these views back GPU kernels and we mirror the GPU path. +The round-trip class is skipped on CPU — those views back GPU kernels and we +mirror the GPU path. ``TestKDAFlashInferEnvelopeStateContract`` is pure stride +arithmetic and runs everywhere. python -m pytest test/registered/unit/mem_cache/test_shared_mamba_views.py -v """ @@ -288,5 +290,174 @@ class TestUnifiedMambaViews(unittest.TestCase): self._fill_and_roundtrip(pool, spec) +def _k3_kda_mamba_geometry(heads_per_rank: int) -> dict: + """Kimi K3 KDA per-rank state geometry: 69 KDA layers, K = V = 128, + conv width 4 (=> 3 cached tokens), conv row ``(kernel-1, q+k+v dim)`` + in the KimiLinear layout (``KimiLinearStateShape.create`` with + num_k_heads == num_heads, head_k_dim == head_dim — see + ``models/kimi_linear.py``), temporal/SSM state ``(HV, V, K)``. + ``heads_per_rank`` = 96 total KDA heads / attn_tp (12 at the TP8 + deployment shape, cf. ``kernels/ops/attention/kda_fused_decode.py``).""" + h = heads_per_rank + return dict( + layer_num=69, + conv_state_shapes=((3, 3 * h * 128),), + conv_dtype=torch.bfloat16, + temporal_state_shape=(h, 128, 128), + # FlashInfer recurrent_kda requires a bf16 state pool (the server-args + # gate enforces --mamba-ssm-dtype bfloat16 for flashinfer decode). + temporal_dtype=torch.bfloat16, + ) + + +class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase): + """Derived property: the envelope-strided KDA temporal view (unified memory + / page-major layout) must satisfy the state contract of FlashInfer + ``recurrent_kda`` (pinned ``flashinfer_python==0.6.14``), because the KDA + flashinfer decode wrapper (``linear/kernels/kda_flashinfer.py``) passes the + committed per-layer pool view straight into the kernel (in-place state + update on the cu_seqlens path — no gather/scatter copy around the call). + + The kernel compiles its state argument as a CuTe fake tensor of shape + ``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)`` + and ``assumed_align=32`` (flashinfer ``kda_kernels/recurrent_kda.py``), so + a per-layer pool view is only readable by the kernel when: + + * its inner strides are exactly compact ``(V*K, K, 1)``; + * its slot stride — the per-slot envelope pitch, NOT ``HV*V*K`` — is a + multiple of 16 elements (32 bytes at bf16); + * its base byte offset is 32-byte aligned (for every layer). + + Any envelope-layout change that breaks one of these (per-slot padding that + is not a 32 B multiple, a conv-shape change misaligning the temporal + region, a transposed/padded temporal inner layout) would silently + mis-address every KDA state read/write on SM100 flashinfer decode; this + test turns such a diff red without a GPU. + """ + + # 32 B: recurrent_kda's assumed_align AND its slot-stride divisibility + # (16 elements * 2 B bf16). External-source literal from flashinfer + # kda_kernels/recurrent_kda.py (S_batch = cute.sym_int64(divisibility=16), + # make_fake_tensor(..., assumed_align=32)). + _KERNEL_ALIGN_BYTES = 32 + + @staticmethod + def _build_tp8_views(): + """Real TP8 K3 KDA envelope views on CPU (2 slots suffice — the + per-slot geometry is slot-count independent).""" + from sglang.srt.mem_cache.layout.page_major import ( + build_page_major_mamba_views, + mamba_entry_bytes, + ) + + geom = _k3_kda_mamba_geometry(12) # 96 heads / TP8 + entry_bytes = mamba_entry_bytes(**geom) + max_slots = 2 + raw = torch.empty(max_slots * entry_bytes, dtype=torch.uint8, device="cpu") + _, temporal_view = build_page_major_mamba_views( + raw, max_slots=max_slots, **geom + ) + return geom, entry_bytes, temporal_view + + def test_k3_tp8_envelope_view_matches_recurrent_kda_contract(self): + """Check every per-layer temporal view against the kernel contract.""" + geom, entry_bytes, temporal_view = self._build_tp8_views() + + itemsize = temporal_view.element_size() + _, v, k = geom["temporal_state_shape"] + for layer in (0, geom["layer_num"] - 1): + view = temporal_view[layer] # [slots, HV, V, K], what decode() gets + self.assertEqual( + view.stride()[1:], + (v * k, k, 1), + "temporal inner strides must stay compact (V*K, K, 1): " + "recurrent_kda compiles them as constants", + ) + self.assertEqual( + view.stride(0), + entry_bytes // itemsize, + "slot stride must be the envelope pitch (entry_bytes)", + ) + self.assertEqual( + view.stride(0) % (self._KERNEL_ALIGN_BYTES // itemsize), + 0, + "slot stride must satisfy recurrent_kda's " + "sym_int64(divisibility=16) — 16 elements = 32 B at bf16", + ) + self.assertEqual( + (view.storage_offset() * itemsize) % self._KERNEL_ALIGN_BYTES, + 0, + f"layer {layer} temporal view base is not 32 B aligned " + "(recurrent_kda assumed_align=32)", + ) + + def test_k3_entry_and_temporal_offset_32B_multiples_across_tp(self): + """The two byte quantities that feed the contract above — the per-slot + envelope pitch and the temporal region's offset inside the envelope + (= all-layers conv region, temporal comes last) — must be 32 B + multiples for every plausible attn-TP shard of K3's 96 KDA heads.""" + import math + + from sglang.srt.mem_cache.layout.page_major import mamba_entry_bytes + + for heads_per_rank in (96, 48, 24, 12): # attn_tp 1 / 2 / 4 / 8 + geom = _k3_kda_mamba_geometry(heads_per_rank) + entry_bytes = mamba_entry_bytes(**geom) + conv_region_bytes = ( + geom["layer_num"] + * math.prod(geom["conv_state_shapes"][0]) + * geom["conv_dtype"].itemsize + ) + self.assertEqual( + entry_bytes % self._KERNEL_ALIGN_BYTES, + 0, + f"tp shard h={heads_per_rank}: envelope pitch {entry_bytes} B " + "breaks recurrent_kda's slot-stride divisibility", + ) + self.assertEqual( + conv_region_bytes % self._KERNEL_ALIGN_BYTES, + 0, + f"tp shard h={heads_per_rank}: temporal region offset " + f"{conv_region_bytes} B breaks assumed_align=32", + ) + + def test_wrapper_state_contract_check_matches_layout(self): + """The KDA flashinfer decode wrapper enforces this same contract at + runtime (``FlashInferKDAKernel._check_state_stride_contract``, called + once per pool view before handing the pool to ``recurrent_kda``). A + regression in that check would only surface on SM100 hardware, so pin + its accept/reject behavior here: it must ACCEPT exactly what the + layouts produce — the envelope-strided per-layer view and a plain + contiguous pool — and REJECT views the kernel would silently + mis-address (wrong inner strides; a slot stride off the divisibility).""" + import types + + from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import ( + FlashInferKDAKernel, + ) + + check = FlashInferKDAKernel._check_state_stride_contract + + def run(view): + # Fresh stub per call: the real kernel caches approvals by id(). + check(types.SimpleNamespace(_state_contract_ok=set()), view) + + _, _, temporal_view = self._build_tp8_views() + envelope = temporal_view[0] # what forward_decode hands to the kernel + run(envelope) # must not raise + + contiguous = torch.empty(2, 12, 128, 128, dtype=torch.bfloat16) + run(contiguous) # locally-allocated pools must keep working + + with self.assertRaises(ValueError): + run(envelope.transpose(-1, -2)) # inner strides not compact + + # Slot stride 196616 elements: envelope-like but % 16 != 0. + flat = torch.empty(2 * 196616, dtype=torch.bfloat16) + misaligned = flat.as_strided((2, 12, 128, 128), (196616, 16384, 128, 1)) + with self.assertRaises(ValueError): + run(misaligned) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/models/test_kimi_k25.py b/test/registered/unit/models/test_kimi_k25.py index 51ed152bc..6f5f57ad7 100644 --- a/test/registered/unit/models/test_kimi_k25.py +++ b/test/registered/unit/models/test_kimi_k25.py @@ -1,7 +1,8 @@ """CPU coverage for Kimi-K2.5/K2.7 encoder-DP wiring.""" +import asyncio from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, Mock, patch import numpy as np import pytest @@ -24,8 +25,15 @@ from sglang.srt.models.kimi_vl_moonvit import tpool_patch_merger from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin +from sglang.srt.multimodal.processors.kimi_k3 import ( + KimiK3GPUProcessorWrapper, + KimiK3ImageProcessor, + _expand_k3_image_prompt_text, + _expand_k3_image_prompt_token_ids, +) from sglang.srt.multimodal.processors.kimi_k25 import ( KimiGPUProcessorWrapper, + KimiK2_5VLImageProcessor, _ensure_chw_rgb, _expand_image_token_ids, _resize_bicubic_if_needed, @@ -516,5 +524,222 @@ def test_kimi_lazy_ipc_feature_acknowledges_all_tp_consumers(): proxy.reconstruct_on_target_device.assert_called_once_with(0, consumer_count=8) +class _Tokenizer: + def encode(self, text, allowed_special=None): + tokens = { + "<|media_begin|>image 1536x1024<|media_content|>": [10, 11], + "<|media_begin|>image 1024x1536<|media_content|>": [12, 13], + "<|media_end|>": [14], + } + return tokens.get(text, []) + + +class _HFProcessor: + def __init__(self): + self.tokenizer = _Tokenizer() + self.image_processor = SimpleNamespace() + self.media_processor = SimpleNamespace( + media_proc_cfg={ + "patch_size": 14, + "merge_kernel_size": 2, + "in_patch_limit": 16384, + "patch_limit_on_one_side": 256, + "fixed_output_tokens": None, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + "transparent_bg_config": None, + } + ) + + +@pytest.mark.parametrize( + ("processor_cls", "wrapper_cls"), + [ + (KimiK2_5VLImageProcessor, KimiGPUProcessorWrapper), + (KimiK3ImageProcessor, KimiK3GPUProcessorWrapper), + ], +) +def test_kimi_processor_workers_clone_the_gpu_wrapper(processor_cls, wrapper_cls): + server_args = SimpleNamespace( + mm_feature_transport="cpu", + disable_fast_image_processor=False, + skip_tokenizer_init=False, + mm_process_config={}, + mm_io_worker_num=0, + mm_processor_worker_num=0, + tokenizer_worker_num=1, + base_gpu_id=0, + ) + processor = processor_cls( + hf_config=SimpleNamespace(media_placeholder_token_id=42), + server_args=server_args, + _processor=_HFProcessor(), + transport_mode=None, + ) + try: + worker_processor = asyncio.run( + processor.mm_processor_executor.run(lambda *, processor: processor) + ) + assert isinstance(processor._processor, wrapper_cls) + assert isinstance(worker_processor, wrapper_cls) + assert worker_processor is not processor._processor + finally: + processor.mm_processor_executor.shutdown() + processor.io_executor.shutdown() + processor.cpu_executor.shutdown() + + +def test_kimi_k3_expands_image_placeholders_with_original_dimensions(): + actual = _expand_k3_image_prompt_token_ids( + [1, 99, 2, 99, 3], + 99, + [3, 2], + [(1536, 1024), (1024, 1536)], + _Tokenizer(), + ) + + assert actual.tolist() == [[1, 10, 11, 99, 99, 99, 14, 2, 12, 13, 99, 99, 14, 3]] + + +def test_kimi_k3_cpu_prompt_uses_the_same_media_contract(): + actual = _expand_k3_image_prompt_text( + "before<|media_pad|>between<|media_pad|>after", + "<|media_pad|>", + [3, 2], + [(1536, 1024), (1024, 1536)], + ) + + assert actual == ( + "before<|media_begin|>image 1536x1024<|media_content|>" + "<|media_pad|><|media_pad|><|media_pad|><|media_end|>between" + "<|media_begin|>image 1024x1536<|media_content|>" + "<|media_pad|><|media_pad|><|media_end|>after" + ) + + +def test_kimi_k3_epd_rebuild_uses_the_same_media_contract(): + processor = object.__new__(KimiK3ImageProcessor) + processor.hf_config = SimpleNamespace( + vision_config=SimpleNamespace(merge_kernel_size=(2, 2)) + ) + processor.mm_tokens = SimpleNamespace(image_token_id=99) + processor._tokenizer = _Tokenizer() + embeddings = {Modality.IMAGE: torch.arange(20, dtype=torch.float32).reshape(5, 4)} + + output = processor.get_mm_data( + [1, 99, 2, 99, 3], + embeddings, + img_grid_thw=torch.tensor([[1, 2, 6], [1, 2, 4]]), + original_image_sizes=[[1536, 1024], [1024, 1536]], + ) + + assert output.input_ids == [ + 1, + 10, + 11, + 99, + 99, + 99, + 14, + 2, + 12, + 13, + 99, + 99, + 14, + 3, + ] + assert [item.offsets for item in output.mm_items] == [[(3, 5)], [(10, 11)]] + torch.testing.assert_close( + output.mm_items[0].precomputed_embeddings, embeddings[Modality.IMAGE][:3] + ) + torch.testing.assert_close( + output.mm_items[1].precomputed_embeddings, embeddings[Modality.IMAGE][3:] + ) + + +def test_kimi_k3_rejects_silently_dropped_images(): + processor = object.__new__(KimiK3ImageProcessor) + processor.mm_tokens = Mock() + processor.load_mm_data = AsyncMock(return_value=SimpleNamespace(images=[object()])) + + with pytest.raises(ValueError, match="expected 2, loaded 1"): + asyncio.run( + processor.process_mm_data_async( + image_data=["image-1", "image-2"], + input_text="<|media_pad|><|media_pad|>", + request_obj=SimpleNamespace(video_data=None), + ) + ) + + +def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries(): + processor = object.__new__(KimiK3ImageProcessor) + processor.mm_tokens = SimpleNamespace(image_token_id=99) + processor.fast_load_mm_data = AsyncMock( + return_value=SimpleNamespace( + images=[object(), object()], input_ids=[1, 99, 2, 99, 3] + ) + ) + processor.load_mm_data = AsyncMock() + processor.process_and_combine_mm_data_async = AsyncMock( + return_value=([], torch.tensor([[1, 2]]), None) + ) + + asyncio.run( + processor.process_mm_data_async( + image_data=["image-1", "image-2"], + input_text=[1, 99, 2, 99, 3], + request_obj=SimpleNamespace(video_data=None), + ) + ) + + processor.fast_load_mm_data.assert_awaited_once() + processor.load_mm_data.assert_not_awaited() + + +def test_kimi_k3_rejects_tokenized_placeholder_mismatch(): + processor = object.__new__(KimiK3ImageProcessor) + processor.mm_tokens = SimpleNamespace(image_token_id=99) + processor.fast_load_mm_data = AsyncMock() + processor.load_mm_data = AsyncMock() + + with pytest.raises(ValueError, match=r"expected 2, found 1 token\(s\)"): + asyncio.run( + processor.process_mm_data_async( + image_data=["image-1", "image-2"], + input_text=torch.tensor([[1, 99, 2]]), + request_obj=SimpleNamespace(video_data=None), + ) + ) + + processor.fast_load_mm_data.assert_not_awaited() + processor.load_mm_data.assert_not_awaited() + + +@pytest.mark.parametrize( + ("request_obj", "extra_kwargs"), + [ + (SimpleNamespace(video_data=["video"]), {}), + (SimpleNamespace(video_data=None), {"audio_data": ["audio"]}), + ], +) +def test_kimi_k3_rejects_unsupported_modalities(request_obj, extra_kwargs): + processor = object.__new__(KimiK3ImageProcessor) + processor.mm_tokens = Mock() + processor.load_mm_data = AsyncMock() + + with pytest.raises(ValueError, match="supports image input only"): + asyncio.run( + processor.process_mm_data_async( + image_data=[], + input_text="prompt", + request_obj=request_obj, + **extra_kwargs, + ) + ) + processor.load_mm_data.assert_not_awaited() + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/models/test_kimi_k25_mm_projection.py b/test/registered/unit/models/test_kimi_k25_mm_projection.py new file mode 100644 index 000000000..dfe39d9b6 --- /dev/null +++ b/test/registered/unit/models/test_kimi_k25_mm_projection.py @@ -0,0 +1,41 @@ +"""CPU coverage for the Kimi vision-projector packing fast path.""" + +import pytest +import torch +import torch.nn as nn + +from sglang.srt.models.kimi_k25 import mm_projection_auto +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _FlattenProjector(nn.Module): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states.flatten(start_dim=1) + + +def test_mm_projection_auto_packs_variable_image_outputs_once(): + outputs = [ + torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3), + torch.arange(3 * 4 * 3, dtype=torch.float32).reshape(3, 4, 3), + ] + expected = torch.cat([output.flatten(start_dim=1) for output in outputs], dim=0) + + actual = mm_projection_auto(_FlattenProjector(), outputs) + + torch.testing.assert_close(actual, expected) + assert actual.shape == (5, 12) + + +def test_mm_projection_auto_single_item_avoids_cat_copy(): + output = torch.randn(5, 4, 3) + + actual = mm_projection_auto(_FlattenProjector(), [output]) + + torch.testing.assert_close(actual, output.flatten(start_dim=1)) + assert actual.data_ptr() == output.data_ptr() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/models/test_kimi_k3_bfa_overlap.py b/test/registered/unit/models/test_kimi_k3_bfa_overlap.py new file mode 100644 index 000000000..128fc28c2 --- /dev/null +++ b/test/registered/unit/models/test_kimi_k3_bfa_overlap.py @@ -0,0 +1,102 @@ +"""KDA bfa side-stream overlap: forward_qkvbfg_fused must produce outputs +bit-identical to the serial path, both eager and under CUDA graph +capture/replay (the overlap only engages in capture mode).""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.models.kimi_k3 import KimiK3DeltaAttention +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=90, stage="base-b", runner_config="1-gpu-large") + +_H = 7168 +_QKVG = 6144 # q,k,v,g slices per rank at TP8 +_N_FA = 128 +_N_B = 12 +_BFA_W_ROWS = 144 # [f_a | b] padded to 8 rows like _merge_bfa_weights + + +def _make_owner(with_stream: bool): + gen = torch.Generator(device="cuda").manual_seed(0) + + def _randn(*shape): + return ( + torch.randn(*shape, generator=gen, device="cuda", dtype=torch.float32) + .mul(0.05) + .to(torch.bfloat16) + ) + + qkvg_w = _randn(_QKVG, _H) + + def fused_qkvg_proj(x): + return torch.nn.functional.linear(x, qkvg_w), None + + owner = SimpleNamespace( + use_full_rank_gate=True, + _bfa_w=_randn(_BFA_W_ROWS, _H).contiguous(), + _bfa_fa_size=_N_FA, + _bfa_b_size=_N_B, + f_b_proj=SimpleNamespace(weight=_randn(1536, _N_FA).contiguous()), + fused_qkvg_proj=fused_qkvg_proj, + split_sizes=[3 * 1536, 1536], + _bfa_alt_stream=torch.cuda.Stream() if with_stream else None, + _bfa_bs_limit=128 if with_stream else 0, + ) + return owner + + +def _run(owner, x): + out = KimiK3DeltaAttention.forward_qkvbfg_fused(owner, x) + return [t.clone() for t in out] + + +class TestKimiK3BfaOverlap(CustomTestCase): + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA is not available") + + def test_capture_replay_matches_serial(self): + torch.manual_seed(0) + for T in (1, 4, 12): + with self.subTest(T=T): + x = ( + torch.randn(T, _H, device="cuda", dtype=torch.float32) + .mul(0.05) + .to(torch.bfloat16) + ) + serial = _run(_make_owner(with_stream=False), x) + + owner = _make_owner(with_stream=True) + with patch( + "sglang.srt.models.kimi_k3.get_is_capture_mode", + return_value=True, + ): + # warm up allocations/JIT outside capture + _ = _run(owner, x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = KimiK3DeltaAttention.forward_qkvbfg_fused(owner, x) + graph.replay() + torch.cuda.synchronize() + # note: owners share the same seeded weights + for got, ref, name in zip( + captured, serial, ("qkv", "beta", "forget_gate", "g") + ): + self.assertTrue(torch.equal(got, ref), f"T={T} {name} mismatch") + + def test_eager_stream_branch_not_taken(self): + x = torch.randn(3, _H, device="cuda", dtype=torch.bfloat16) + serial = _run(_make_owner(with_stream=False), x) + overlap = _run(_make_owner(with_stream=True), x) # capture mode False + for got, ref in zip(overlap, serial): + self.assertTrue(torch.equal(got, ref)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_kimi_k3_vision.py b/test/registered/unit/models/test_kimi_k3_vision.py new file mode 100644 index 000000000..02418ce84 --- /dev/null +++ b/test/registered/unit/models/test_kimi_k3_vision.py @@ -0,0 +1,524 @@ +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from sglang.srt.layers.attention.vision import ( + prepare_flashinfer_cudnn_vision_attention_metadata, +) +from sglang.srt.models import kimi_k3_vl +from sglang.srt.models.kimi_k3_vl import ( + KimiK3VisionTower, + MoonViT3dEncoder, + _resolve_mm_attention_backend, + interpolate_pos_emb, + sdpa_varlen_attention, +) +from sglang.srt.multimodal import kimi_k3_vit_cuda_graph_runner +from sglang.srt.multimodal.kimi_k3_vit_cuda_graph_runner import ( + KimiK3ViTCudaGraphRunner, +) +from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model +from sglang.srt.runtime_context import get_parallel +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +@pytest.mark.parametrize( + ( + "configured_backend", + "device_type", + "capability", + "max_seqlen", + "total_tokens", + "fa4_available", + "expected", + ), + [ + ("sdpa", "cuda", (10, 3), 8192, 8192, True, "sdpa"), + ("auto", "cpu", None, 1024, 1024, True, "sdpa"), + ("auto", "cuda", (10, 0), 1024, 1024, True, "sdpa"), + ("auto", "cuda", (10, 3), 1536, 1536, True, "triton_attn"), + ("auto", "cuda", (10, 3), 1600, 1600, True, "fa4"), + ("auto", "cuda", (10, 3), 1024, 4096, True, "fa4"), + ("auto", "cuda", (10, 3), 1536, 1536, False, "triton_attn"), + ("auto", "cuda", (10, 3), 1600, 1600, False, "sdpa"), + ], +) +def test_kimi_k3_resolves_shape_aware_attention_backend( + monkeypatch, + configured_backend, + device_type, + capability, + max_seqlen, + total_tokens, + fa4_available, + expected, +): + if capability is not None: + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: capability) + + actual = _resolve_mm_attention_backend( + configured_backend, + max_seqlen=max_seqlen, + total_tokens=total_tokens, + device=torch.device(device_type), + fa4_available=fa4_available, + ) + + assert actual == expected + + +def test_kimi_k3_skips_attention_precompile_on_cpu(): + encoder = MoonViT3dEncoder( + hidden_dim=8, + num_layers=1, + block_cfg={ + "num_heads": 1, + "hidden_dim": 8, + "qkv_hidden_size": 8, + "mlp_dim": 16, + "norm_type": "rmsnorm", + "activation": F.gelu, + "attn_bias": False, + "linear_bias": False, + }, + ) + + assert not encoder.precompile_attention_backend(torch.bfloat16, torch.device("cpu")) + + +def test_kimi_k3_sdpa_reuses_prepared_segment_bounds(): + class SeqlensThatMustNotSync: + def tolist(self): + raise AssertionError("prepared segment bounds must avoid tensor.tolist()") + + q = torch.randn(4, 1, 8) + k = torch.randn(4, 1, 8) + v = torch.randn(4, 1, 8) + bounds = ((0, 2), (2, 4)) + expected = sdpa_varlen_attention(q, k, v, torch.tensor([0, 2, 4])) + actual = sdpa_varlen_attention( + q, + k, + v, + SeqlensThatMustNotSync(), + segment_bounds=bounds, + ) + + assert torch.equal(actual, expected) + + +def test_kimi_k3_vision_tower_reuses_prepared_forward_metadata(monkeypatch): + config = SimpleNamespace( + patch_size=2, + init_pos_emb_height=2, + init_pos_emb_width=2, + init_pos_emb_time=1, + pos_emb_type="divided_fixed", + pos_emb_interpolation_mode="bilinear", + patch_embed_proj_bias=False, + merge_kernel_size=(1, 1), + merge_type="sd2_tpool", + vt_hidden_size=8, + vt_num_attention_heads=1, + vt_num_hidden_layers=0, + num_hidden_layers=0, + vt_intermediate_size=16, + qkv_hidden_size=8, + norm_type="rmsnorm", + activation_func="gelu_pytorch_tanh", + attn_bias=False, + linear_bias=False, + ) + tower = KimiK3VisionTower(config).eval() + pixel_values = torch.randn(4, 3, 2, 2) + grid_thws = torch.tensor([[1, 2, 2]]) + grid_thw_list = ((1, 2, 2),) + reference = tower(pixel_values, grid_thws) + metadata = tower.prepare_forward_metadata( + grid_thws, + grid_thw_list=grid_thw_list, + total_tokens=pixel_values.shape[0], + dtype=pixel_values.dtype, + ) + assert metadata.position_embeddings is not None + + def fail_reprepare(**_): + raise AssertionError("forward metadata must be reused") + + def fail_position_recompute(*_args, **_kwargs): + raise AssertionError("prepared position embeddings must be reused") + + monkeypatch.setattr(tower.encoder, "prepare_forward_metadata", fail_reprepare) + monkeypatch.setattr( + tower.patch_embed.pos_emb, + "position_embeddings", + fail_position_recompute, + ) + actual = tower( + pixel_values, + grid_thws, + grid_thw_list=grid_thw_list, + forward_metadata=metadata, + ) + + assert len(actual) == len(reference) == 1 + assert torch.equal(actual[0], reference[0]) + + +def test_kimi_k3_dp_helper_passes_host_grid_list_to_capable_tower(): + class RecordingTower: + def __call__( + self, + pixel_values, + *, + grid_hw, + max_seqlen, + grid_thw_list, + ): + self.grid_hw = grid_hw + self.max_seqlen = max_seqlen + self.grid_thw_list = grid_thw_list + return [pixel_values.unsqueeze(0)] + + tower = RecordingTower() + pixels = torch.randn(4, 2) + grids = [[1, 2, 2]] + with get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0): + output = run_dp_sharded_mrope_vision_model( + tower, + pixels, + grids, + rope_type="rope_2d", + pass_grid_thw_list=True, + ) + + assert tower.grid_hw.device == pixels.device + assert tower.max_seqlen == 4 + assert tower.grid_thw_list is grids + assert torch.equal(output, pixels.unsqueeze(0)) + + +def test_kimi_k3_vit_graph_runner_bounds_shape_observations(): + runner = KimiK3ViTCudaGraphRunner(object(), capacity=2, min_hits=2) + + for index in range(20): + runner._record_hit(((1, index + 1, 2),)) + + assert len(runner.seen) == 16 + assert ((1, 1, 2),) not in runner.seen + assert ((1, 20, 2),) in runner.seen + + +def test_kimi_k3_vit_graph_runner_skips_eager_on_capture(monkeypatch): + class Tower: + def prepare_forward_metadata(self, *_args, **_kwargs): + return object() + + runner = KimiK3ViTCudaGraphRunner(Tower(), capacity=1, min_hits=1) + pixels = torch.randn(4, 2) + grids = torch.tensor([[1, 2, 2]]) + replayed = [] + + def fail_eager(*_args, **_kwargs): + raise AssertionError("the capture request must not run eager first") + + def capture(_key, pixel_values, _grid_thws, _grid_thw_list, metadata): + return SimpleNamespace( + graph=SimpleNamespace(replay=lambda: replayed.append(True)), + input_buffer=torch.empty_like(pixel_values), + outputs=(pixel_values.clone(),), + metadata=metadata, + ) + + monkeypatch.setattr(runner, "_run_eager", fail_eager) + monkeypatch.setattr(runner, "_capture", capture) + outputs = runner.run(pixels, grids, ((1, 2, 2),)) + + assert replayed == [True] + assert torch.equal(outputs[0], pixels) + + +def test_kimi_k3_vit_graph_runner_uses_eager_above_max_seqlen(monkeypatch): + runner = KimiK3ViTCudaGraphRunner(object(), capacity=1, min_hits=1, max_seqlen=4) + pixels = torch.randn(8, 2) + grids = torch.tensor([[1, 2, 4]]) + eager_calls = [] + + def run_eager(pixel_values, *_args, **_kwargs): + eager_calls.append(True) + return [pixel_values.clone()], object() + + monkeypatch.setattr(runner, "_run_eager", run_eager) + monkeypatch.setattr( + runner, + "_capture", + lambda *_args, **_kwargs: pytest.fail("large shapes must not be captured"), + ) + + for _ in range(3): + outputs = runner.run(pixels, grids, ((1, 2, 4),)) + + assert eager_calls == [True, True, True] + assert torch.equal(outputs[0], pixels) + assert not runner.seen + assert not runner.graphs + + +def test_kimi_k3_vit_graph_runner_reuses_global_graph_pool(monkeypatch): + class Tower: + def _forward_eager(self, pixel_values, *_args, **_kwargs): + return [pixel_values.clone()] + + pool = object() + graph = object() + pool_creations = [] + capture_pools = [] + + def get_pool(device_module): + pool_creations.append(device_module) + return pool + + def graph_context(captured_graph, *, pool): + assert captured_graph is graph + capture_pools.append(pool) + return nullcontext() + + monkeypatch.setattr( + kimi_k3_vit_cuda_graph_runner, + "get_or_create_global_graph_memory_pool", + get_pool, + ) + monkeypatch.setattr(torch.cuda, "CUDAGraph", lambda: graph) + monkeypatch.setattr(torch.cuda, "graph", graph_context) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda _device: 0) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda _device: 0) + monkeypatch.setattr(torch.cuda, "synchronize", lambda _device: None) + + runner = KimiK3ViTCudaGraphRunner(Tower(), capacity=2, min_hits=1) + pixels = torch.randn(4, 2) + grids = torch.tensor([[1, 2, 2]]) + metadata = object() + runner._capture(((1, 2, 2),), pixels, grids, ((1, 2, 2),), metadata) + runner._capture(((1, 1, 4),), pixels, grids, ((1, 1, 4),), metadata) + + assert pool_creations == [torch.cuda] + assert capture_pools == [pool, pool] + + +@pytest.mark.parametrize( + ("capacity", "min_hits", "max_seqlen"), + [(0, 1, None), (1, 0, None), (1, 1, 0)], +) +def test_kimi_k3_vit_graph_runner_rejects_invalid_limits( + capacity, min_hits, max_seqlen +): + with pytest.raises(ValueError): + KimiK3ViTCudaGraphRunner( + object(), + capacity=capacity, + min_hits=min_hits, + max_seqlen=max_seqlen, + ) + + +def test_kimi_k3_position_interpolation_uses_contiguous_chw(monkeypatch): + weight = torch.randn(4, 5, 8) + output_size = (3, 7) + expected = ( + F.interpolate( + weight.permute(2, 0, 1).unsqueeze(0), + size=output_size, + mode="bilinear", + ) + .squeeze(0) + .permute(1, 2, 0) + .flatten(end_dim=1) + ) + original_interpolate = F.interpolate + captured = {} + + def capture_layout(input_tensor, *args, **kwargs): + captured["is_contiguous"] = input_tensor.is_contiguous() + captured["stride"] = input_tensor.stride() + return original_interpolate(input_tensor, *args, **kwargs) + + monkeypatch.setattr(kimi_k3_vl.F, "interpolate", capture_layout) + actual = interpolate_pos_emb(weight, "bilinear", output_size) + + assert captured == { + "is_contiguous": True, + "stride": (160, 20, 5, 1), + } + assert torch.equal(actual, expected) + + +def test_kimi_k3_prepares_shared_attention_metadata_once(monkeypatch): + metadata_ids = [] + values_are_contiguous = [] + + class FakeAttention(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + + def forward(self, q, k, v, *, forward_metadata, **kwargs): + metadata_ids.append(id(forward_metadata)) + values_are_contiguous.append(v.is_contiguous()) + return q + + monkeypatch.setattr( + kimi_k3_vl, + "get_server_args", + lambda: SimpleNamespace(mm_attention_backend="flashinfer_cudnn"), + ) + monkeypatch.setitem(kimi_k3_vl.QKV_BACKEND_IMPL, "flashinfer_cudnn", FakeAttention) + + encoder = MoonViT3dEncoder( + hidden_dim=8, + num_layers=2, + block_cfg={ + "num_heads": 1, + "hidden_dim": 8, + "qkv_hidden_size": 8, + "mlp_dim": 16, + "norm_type": "rmsnorm", + "activation": F.gelu, + "attn_bias": False, + "linear_bias": False, + }, + ) + output = encoder(torch.randn(4, 8), torch.tensor([[1, 2, 2]])) + + assert output.shape == (4, 8) + assert len(metadata_ids) == 2 + assert len(set(metadata_ids)) == 1 + assert values_are_contiguous == [True, True] + + +def test_flashinfer_cudnn_metadata_uses_bucketed_element_indptrs(): + metadata = prepare_flashinfer_cudnn_vision_attention_metadata( + torch.tensor([0, 480, 1200], dtype=torch.int32), + device=torch.device("cpu"), + elem_per_token=1536, + ) + + expected_indptr = torch.tensor( + [0, 480 * 1536, 1200 * 1536] + [1200 * 1536] * 6, + dtype=torch.int32, + ) + assert torch.equal(metadata.packed_indptrs, expected_indptr.repeat(3)) + assert metadata.sequence_lengths.flatten().tolist() == [480, 720] + [0] * 6 + assert metadata.flashinfer_max_seqlen == 4096 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) + + +class _K3TowerStub: + device = torch.device("cpu") + merge_kernel_size = (2, 2) + + def __init__(self): + self.config = SimpleNamespace(hidden_size=2) + self.patch_embed = SimpleNamespace( + proj=SimpleNamespace(weight=torch.empty(1, dtype=torch.float32)) + ) + + +def test_kimi_k3_encoder_dp_defers_feature_materialization(monkeypatch): + """K3 vision is image-wise DP: the DP runner must receive lazy features + (pixel_values=None + a loader), and the loader must materialize exactly + the requested images on the owner rank with the tower dtype.""" + from unittest.mock import patch as mock_patch + + from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem + from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration + + model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration) + torch.nn.Module.__init__(model) + model.use_data_parallel = True + model.vision_tower = _K3TowerStub() + model.mm_projector = lambda image_embeds: image_embeds + + items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + offsets=[(0, 1)], + feature=torch.randn(4, 2, dtype=torch.float64), + model_specific_data={"grid_thws": torch.tensor([[1, 2, 2]])}, + ), + MultimodalDataItem( + modality=Modality.IMAGE, + offsets=[(1, 2)], + feature=torch.randn(4, 2, dtype=torch.float64), + model_specific_data={"grid_thws": torch.tensor([[1, 2, 2]])}, + ), + ] + sharded_embeddings = torch.randn(2, 2) + + with mock_patch( + "sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model", + return_value=sharded_embeddings, + ) as run_dp, mock_patch( + "sglang.srt.models.kimi_k3.get_server_args", + return_value=SimpleNamespace(tp_size=1), + ), mock_patch( + "sglang.srt.models.kimi_k3.get_parallel", + return_value=SimpleNamespace(attn_tp_size=1), + ): + output = model.get_image_feature(items) + # exercise the loader inside the patch scope: it reads server args + loader_in_scope = run_dp.call_args.kwargs["load_local_pixel_values"] + local = loader_in_scope([1]) + both = loader_in_scope([0, 1]) + + assert output is sharded_embeddings + tower, pixel_values, grid_thws = run_dp.call_args.args + assert tower is model.vision_tower + assert pixel_values is None + assert grid_thws == [[1, 2, 2], [1, 2, 2]] + assert run_dp.call_args.kwargs["rope_type"] == "rope_2d" + assert run_dp.call_args.kwargs["pass_grid_thw_list"] is True + assert run_dp.call_args.kwargs["pool_temporal_dimension"] is True + loader = run_dp.call_args.kwargs["load_local_pixel_values"] + assert callable(loader) + + # Owner-rank materialization: only the requested image, tower dtype. + assert local.shape == (4, 2) + assert local.dtype == torch.float32 + assert torch.equal(local, items[1].feature.to(torch.float32)) + + assert both.shape == (8, 2) + assert torch.equal( + both, + torch.cat([items[0].feature, items[1].feature]).to(torch.float32), + ) + + +def test_kimi_k3_rejects_aggregated_items(): + """One item must carry exactly one logical image: the DP owner + assignment and the bounded CUDA-IPC lease accounting are per-item, so + aggregated encoder inputs must be split upstream (EPD encode server).""" + from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem + from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration + + model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration) + torch.nn.Module.__init__(model) + model.use_data_parallel = True + model.vision_tower = _K3TowerStub() + model.mm_projector = lambda image_embeds: image_embeds + + aggregated = MultimodalDataItem( + modality=Modality.IMAGE, + offsets=[(0, 2)], + feature=torch.arange(12 * 2, dtype=torch.float64).reshape(12, 2), + model_specific_data={"grid_thws": torch.tensor([[1, 2, 2], [1, 2, 4]])}, + ) + + with pytest.raises(ValueError, match="one vision grid per MultimodalDataItem"): + model.get_image_feature([aggregated]) diff --git a/test/registered/unit/multimodal/test_kimi_k3_gpu_preprocess.py b/test/registered/unit/multimodal/test_kimi_k3_gpu_preprocess.py new file mode 100644 index 000000000..1b978fd9e --- /dev/null +++ b/test/registered/unit/multimodal/test_kimi_k3_gpu_preprocess.py @@ -0,0 +1,114 @@ +"""K3 GPU preprocess: shared batched pipeline hook contract (CPU parts).""" + +import sys + +import numpy as np +import pytest +import torch +from PIL import Image + +from sglang.srt.multimodal.processors.kimi_k3 import _fill_transparent_bg +from sglang.srt.multimodal.processors.kimi_k25 import _resize_bicubic_if_needed +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _natural_image(height: int, width: int) -> np.ndarray: + """Deterministic natural-image-like content: gradients, hard edges, + and high-frequency texture (the aliasing-sensitive case).""" + yy, xx = np.mgrid[0:height, 0:width].astype(np.float32) + base = ( + 127 + + 60 * np.sin(2 * np.pi * xx / (width / 7.3)) + + 50 * np.cos(2 * np.pi * yy / (height / 5.1)) + ) + edges = 255.0 * ((xx // 9 + yy // 7) % 2) + tex = 30.0 * np.sin(xx * 12.9898 + yy * 78.233) + img = np.clip(0.55 * base + 0.30 * edges + 0.15 * (127 + tex), 0, 255) + return np.stack( + [img, np.roll(img, 13, axis=0), np.roll(img, 29, axis=1)], axis=-1 + ).astype(np.uint8) + + +def test_resize_matches_pil_bicubic_golden(): + """The GPU resize must reproduce the checkpoint processor's + PIL.Image.resize(..., BICUBIC) downscale: PIL antialiases (kernel support + scales with the ratio) and returns uint8. Without antialias=True the + difference on textured content reaches tens of pixel levels.""" + arr = _natural_image(1200, 1600) + for target_w, target_h in ((800, 600), (1120, 840)): + golden = np.asarray( + Image.fromarray(arr).resize((target_w, target_h), Image.Resampling.BICUBIC) + ).astype(np.float32) + + x = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) + ours = ( + _resize_bicubic_if_needed(x, target_h, target_w) + .squeeze(0) + .permute(1, 2, 0) + .numpy() + ) + + diff = np.abs(ours - golden) + # Integer pixel domain: everything within 1 level, most pixels exact. + assert diff.max() <= 1.0, f"max |diff|={diff.max()} at {target_w}x{target_h}" + assert (diff == 0).mean() > 0.7, f"bitwise ratio={(diff == 0).mean():.3f}" + + +def test_fill_transparent_bg_matches_checkpoint_composite(): + """Composite + truncation must match the checkpoint's numpy reference: + alpha * rgb + (1 - alpha) * chessboard, then astype(np.uint8).""" + cfg = { + "pattern": "chessboard", + "chessboard_square_size": 8, + "chessboard_square_on_top_left": True, + "chessboard_white_value": 255, + "chessboard_gray_value": 180, + } + rgba = _natural_image(32, 40) + alpha = ((np.mgrid[0:32, 0:40][0] * 6) % 256).astype(np.uint8) + img = np.concatenate([rgba, alpha[..., None]], axis=-1) + + # Checkpoint reference (media_utils.fill_transparent_bg_with). + bg = np.ones((32, 40, 3), dtype=np.uint8) * 255 + for y in range(0, 32, 8): + for x0 in range(0, 40, 8): + if (y // 8 + x0 // 8) % 2 == 1: + bg[y : y + 8, x0 : x0 + 8] = 180 + a3 = np.stack([alpha.astype(np.float32) / 255.0] * 3, axis=2) + golden = (a3 * img[:, :, :3] + (1 - a3) * bg).astype(np.uint8) + + x = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0) + ours = _fill_transparent_bg(x, cfg).squeeze(0).permute(1, 2, 0).numpy() + assert np.array_equal(ours, golden.astype(np.float32)) + + +def test_fill_transparent_bg_batch_matches_per_image(): + """Compositing a batch must be bitwise identical to per-image calls + (the batched pipeline applies it to whole resize groups).""" + torch.manual_seed(0) + batch = torch.rand(3, 4, 8, 6) * 255.0 + cfg = {"pattern": "chessboard", "chessboard_square_size": 2} + + batched = _fill_transparent_bg(batch, cfg) + per_image = torch.cat( + [_fill_transparent_bg(batch[i : i + 1], cfg) for i in range(batch.shape[0])] + ) + assert torch.equal(batched, per_image) + + +def test_fill_transparent_bg_rgb_passthrough_batch(): + batch = torch.rand(2, 3, 4, 4) * 255.0 + assert _fill_transparent_bg(batch, {"pattern": "white"}) is batch + + +def test_fill_transparent_bg_no_config_drops_alpha(): + batch = torch.rand(2, 4, 4, 4) * 255.0 + out = _fill_transparent_bg(batch, None) + assert out.shape == (2, 3, 4, 4) + assert torch.equal(out, batch[:, :3]) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/server_args/test_dcp_config.py b/test/registered/unit/server_args/test_dcp_config.py deleted file mode 100644 index dc3a8ebf7..000000000 --- a/test/registered/unit/server_args/test_dcp_config.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Unit tests for DCP (Decode Context Parallelism) server args configuration. - -Covers the ``--dcp-comm-backend`` field ({ag_rs, a2a, fi_a2a}) and its -validation in ``ServerArgs._handle_dcp_validation``: - - a2a / fi_a2a require --dcp-size > 1 - - fi_a2a requires a CUDA platform (the authoritative MNNVL fabric probe runs - later, at model-runner init) - - dcp>1 requires CUDA or HIP (base behavior from the merged DCP PR) - -Tests construct with safe defaults (dcp_size=1) then mutate the fields and call -``_handle_dcp_validation`` directly, so construction never trips the platform -gate; is_cuda / is_hip are patched per-test to pin the platform deterministically -(these are CPU-CI tests, where the real is_cuda() is False). -""" - -import dataclasses -import unittest -from unittest.mock import patch - -from sglang.srt.server_args import ServerArgs -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") - -_mock_device = patch("sglang.srt.server_args.get_device", return_value="cuda") -_mock_device.start() - - -class TestDCPFieldDefaults(CustomTestCase): - """Verify DCP-related dataclass fields exist with correct defaults.""" - - def test_dcp_size_field_exists(self): - fields = {f.name for f in dataclasses.fields(ServerArgs)} - self.assertIn("dcp_size", fields) - - def test_dcp_comm_backend_field_exists(self): - fields = {f.name for f in dataclasses.fields(ServerArgs)} - self.assertIn("dcp_comm_backend", fields) - - def test_dcp_size_default(self): - self.assertEqual(ServerArgs.dcp_size, 1) - - def test_dcp_comm_backend_default(self): - self.assertEqual(ServerArgs.dcp_comm_backend, "ag_rs") - - -class TestDCPCommBackendValidation(CustomTestCase): - """Verify ``_handle_dcp_validation`` accepts/rejects the right combos.""" - - @staticmethod - def _make_args(dcp_size, dcp_comm_backend): - # Construct with safe defaults (dcp_size=1) so __post_init__ never trips - # the dcp>1 platform gate, then set the fields under test. - args = ServerArgs(model_path="dummy") - args.dcp_size = dcp_size - args.dcp_comm_backend = dcp_comm_backend - return args - - def test_a2a_requires_dcp_size_gt_1(self): - args = self._make_args(dcp_size=1, dcp_comm_backend="a2a") - with self.assertRaises(ValueError): - args._handle_dcp_validation() - - def test_fi_a2a_requires_dcp_size_gt_1(self): - args = self._make_args(dcp_size=1, dcp_comm_backend="fi_a2a") - with self.assertRaises(ValueError): - args._handle_dcp_validation() - - @patch("sglang.srt.server_args.is_hip", return_value=False) - @patch("sglang.srt.server_args.is_cuda", return_value=True) - def test_a2a_with_dcp_size_2_on_cuda_passes(self, *_): - args = self._make_args(dcp_size=2, dcp_comm_backend="a2a") - args._handle_dcp_validation() # no raise - self.assertEqual(args.dcp_comm_backend, "a2a") - - @patch("sglang.srt.server_args.is_hip", return_value=False) - @patch("sglang.srt.server_args.is_cuda", return_value=True) - def test_fi_a2a_with_dcp_size_2_on_cuda_passes_server_args(self, *_): - # server_args accepts fi_a2a on CUDA; the MNNVL fabric probe is deferred - # to model-runner init (init_fi_a2a_workspace). - args = self._make_args(dcp_size=2, dcp_comm_backend="fi_a2a") - args._handle_dcp_validation() # no raise - self.assertEqual(args.dcp_comm_backend, "fi_a2a") - - @patch("sglang.srt.server_args.is_hip", return_value=False) - @patch("sglang.srt.server_args.is_cuda", return_value=False) - def test_fi_a2a_on_non_cuda_raises(self, *_): - args = self._make_args(dcp_size=2, dcp_comm_backend="fi_a2a") - with self.assertRaises(ValueError): - args._handle_dcp_validation() - - @patch("sglang.srt.server_args.is_hip", return_value=False) - @patch("sglang.srt.server_args.is_cuda", return_value=True) - def test_ag_rs_with_dcp_size_8_on_cuda_passes(self, *_): - args = self._make_args(dcp_size=8, dcp_comm_backend="ag_rs") - args._handle_dcp_validation() # no raise - self.assertEqual(args.dcp_size, 8) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/server_args/test_mnnvl_auto_inference.py b/test/registered/unit/server_args/test_mnnvl_auto_inference.py new file mode 100644 index 000000000..cd0afda09 --- /dev/null +++ b/test/registered/unit/server_args/test_mnnvl_auto_inference.py @@ -0,0 +1,89 @@ +"""Unit tests for the MNNVL auto-inference gate. + +The TP8 best-throughput launch used to require exporting +``SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE=1`` by hand. It is now +capability-inferred; these cases pin the negative-branch contracts so a +refactor cannot silently turn the predicate into always-true (engaging fabric +paths on non-fabric clusters) or drop the explicit-off override. +""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.environ import envs +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +_HANDLE = ServerArgs._handle_custom_all_reduce_v2_multinode + + +def _cleared(*fields): + """Context helper: run with the given env fields unset, restore after.""" + import contextlib + import os + + @contextlib.contextmanager + def ctx(): + backup = {f.name: os.environ.pop(f.name, None) for f in fields} + try: + yield + finally: + for name, val in backup.items(): + if val is None: + os.environ.pop(name, None) + else: + os.environ[name] = val + + return ctx() + + +class TestCaV2MultinodeAuto(CustomTestCase): + def test_fabric_multinode_auto_enables(self): + """GB200/GB300 + nnodes>1 + unset opt-in -> multinode mode on, v2 kept.""" + with _cleared( + envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE, + envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2, + ), patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True): + _HANDLE(SimpleNamespace(nnodes=2, tp_size=8)) + self.assertTrue(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get()) + self.assertTrue(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get()) + + def test_non_fabric_multinode_still_disables_v2(self): + """Non-fabric multi-node keeps the legacy force-disable (the predicate + must not degrade to always-true).""" + with _cleared( + envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE, + envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2, + ), patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=False): + _HANDLE(SimpleNamespace(nnodes=2, tp_size=8)) + self.assertFalse(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get()) + self.assertFalse(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get()) + + def test_explicit_off_wins_over_fabric(self): + """SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE=0 on a fabric device + must still force-disable v2 (explicit off beats auto-detection).""" + with _cleared(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2), patch( + "sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True + ), envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.override("0"): + _HANDLE(SimpleNamespace(nnodes=2, tp_size=8)) + self.assertFalse(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get()) + self.assertFalse(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get()) + + def test_tp16_not_auto_opted_in(self): + """CustomAllReduceV2 supports world sizes 2..8 only; a TP16 fabric + launch must not auto-set the multinode opt-in (it would log + 'enabling' and then silently fall back downstream).""" + with _cleared( + envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE, + envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2, + ), patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True): + _HANDLE(SimpleNamespace(nnodes=2, tp_size=16)) + self.assertFalse(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.is_set()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_dflash_overlap_hostsync.py b/test/registered/unit/spec/test_dflash_overlap_hostsync.py index 83ad8df38..d3a574211 100644 --- a/test/registered/unit/spec/test_dflash_overlap_hostsync.py +++ b/test/registered/unit/spec/test_dflash_overlap_hostsync.py @@ -220,14 +220,14 @@ class TestRebuildCompactDraftReqToToken(CustomTestCase): class TestHybridNeedsCpuSeqLens(CustomTestCase): - def _make(self, prefill_flag, decode_flag): + def _make(self, prefill_flag, decode_flag, spec_mode="decode"): from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend def backend(flag): return SimpleNamespace(needs_cpu_seq_lens=flag) runner = SimpleNamespace( - server_args=SimpleNamespace(speculative_attention_mode="decode"), + server_args=SimpleNamespace(speculative_attention_mode=spec_mode), kv_cache_dtype=torch.bfloat16, token_to_kv_pool=None, req_to_token_pool=None, @@ -236,9 +236,15 @@ class TestHybridNeedsCpuSeqLens(CustomTestCase): return HybridAttnBackend(runner, backend(prefill_flag), backend(decode_flag)) def test_delegation(self): + # Only backends serving the spec decode loop count: decode always, + # prefill only when speculative_attention_mode routes verify to it. self.assertFalse(self._make(False, False).needs_cpu_seq_lens) - self.assertTrue(self._make(True, False).needs_cpu_seq_lens) + self.assertFalse(self._make(True, False).needs_cpu_seq_lens) self.assertTrue(self._make(False, True).needs_cpu_seq_lens) + self.assertTrue(self._make(True, False, spec_mode="prefill").needs_cpu_seq_lens) + self.assertFalse( + self._make(False, False, spec_mode="prefill").needs_cpu_seq_lens + ) class TestFilterBatchHostIndices(CustomTestCase): diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 3c3e2630e..c3241e5c5 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -90,6 +90,8 @@ class TestModelOverridableWhitelist(CustomTestCase): "fp8_gemm_runner_backend", "disable_custom_all_reduce", "enable_aiter_allreduce_fusion", + "enable_symm_mem", + "speculative_attention_mode", } ), ) diff --git a/test/registered/vlm/test_vision_openai_server_a.py b/test/registered/vlm/test_vision_openai_server_a.py index 00d546f23..c0f7b7cf4 100644 --- a/test/registered/vlm/test_vision_openai_server_a.py +++ b/test/registered/vlm/test_vision_openai_server_a.py @@ -149,7 +149,8 @@ class TestKimiVLServer(ImageOpenAITestMixin): extra_args = [ "--context-length=8192", "--dtype=bfloat16", - "--mem-fraction-static=0.40", + # Weights alone need ~0.39; 0.40 left <0.001 headroom and flaked at load. + "--mem-fraction-static=0.42", ] def test_video_images_chat_completion(self): diff --git a/test/run_suite.py b/test/run_suite.py index de4143366..cab88424f 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -82,6 +82,7 @@ PER_COMMIT_SUITES = { "base-c-test-8-gpu-h20", "base-c-test-8-gpu-h200", "base-c-test-8-gpu-b200", + "base-c-test-8-gpu-b300", "base-c-test-deepep-4-gpu-h100", "base-c-test-deepep-4-gpu-b200", "base-c-test-deepep-8-gpu-h200",