[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
@@ -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()