[Kimi] Support kimi-k3 (#32541)

Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: Mick <mickjagger19@icloud.com> Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Co-authored-by: Ke Bao <ispobaoke@gmail.com> Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com> Co-authored-by: Chunan Zeng <zcnrex@gmail.com> Co-authored-by: Khoa Pham <khoa.pham@radixark.ai> Co-authored-by: Ziyi Xu <ziyi.xu@radixark.ai> Co-authored-by: Zijie Xia <37504505+zijiexia@users.noreply.github.com> Co-authored-by: Yuwei An <ayw.sirius19@gmail.com> Co-authored-by: zhangxiaohao <1024393531@qq.com> Co-authored-by: Yangmin Li <yangminl@nvidia.com> Co-authored-by: Julien Lin <jullin@nvidia.com> Co-authored-by: Hao Phan <htphan@nvidia.com> Co-authored-by: Thomas Wang <1am9trash@gmail.com> Co-authored-by: RolaoDenthu <xinyisong0111@gmail.com> Co-authored-by: pigeonsoup <32922982+pigeonsoup@users.noreply.github.com> Co-authored-by: HaiShaw <hixiao@gmail.com> Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com> Co-authored-by: Lee Nau <lee.nau@gmail.com> Co-authored-by: HMING <126185151+Hearum@users.noreply.github.com> Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: Byron Hsu <byronhsu1230@gmail.com> Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Thomas Wang <thomawan@amd.com> Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai> Co-authored-by: BBuf <xiaoyu.zhang@radixark.ai> Co-authored-by: Hanming Lu <hanminglu@meta.com> Co-authored-by: Xinyi Song <xinyis10@illinois.edu>
2026-08-04 13:22:49 -07:00
parent 0753663b8e
commit abddb1c7e9
139 changed files with 15414 additions and 911 deletions
@@ -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)
@@ -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)
@@ -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 = {
@@ -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"]))
@@ -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,
+    )
@@ -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,
@@ -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()
@@ -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 = {
@@ -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
@@ -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"]))
@@ -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"]))
@@ -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"]))
@@ -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
@@ -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()
@@ -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):
@@ -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

@@ -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):
@@ -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()
@@ -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()
@@ -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__":
@@ -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()
@@ -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"]))
@@ -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"]))
@@ -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()
@@ -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])
@@ -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"]))
@@ -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()
@@ -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()
@@ -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):
@@ -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",
                }
            ),
        )
@@ -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):
@@ -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",