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

feat(kernels): port standalone Kimi K3 kernels (#32890)

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: zhangxiaolei <zhangxiaolei.666@bytedance.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-01 13:26:47 +08:00
committed by GitHub
co-authored by Claude Opus 5 hnyls2002 zhangxiaolei
parent a1344fad4e
commit fb207b72b0
84 changed files with 26142 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py
+230
View File
@@ -0,0 +1,230 @@
"""NV KDA fused-decode kernel gate + slot addressing vs envelope-strided SSM
pools (CPU).
Derived property under test: the fully-fused KDA decode kernel
(``kda_fused_decode``) addresses the ssm/temporal state pool by the slot pitch
the pool actually reports (``ssm_states.stride(0)``), NOT the dense ``HV*V*K``
pitch. Under ``--enable-unified-memory`` / ``--enable-page-major-kv-layout`` the
per-layer temporal view is envelope-strided: one slot pitches across ALL layers
(56,171,520 B on K3), so a hardcoded ``slot*HV*V*K`` offset mis-addresses every
slot > 0 (the exact chunk_delta_h hardcoded-pitch bug pattern, GSM8K 0.17).
Two things are pinned here (both CPU-checkable without the CUDA kernel):
1. ``covered()`` ACCEPTS the envelope-strided view. Pre-fix the gate did
``ssm_states.view(-1, HV, V, K).is_contiguous()``, which is False on a
non-dense slot pitch, so decode silently dropped to the slower unfused
chain. Reverting to that ``.view(...)`` gate turns test (1) red. The gate
still REJECTS a view whose inner ``[HV, V, K]`` is non-contiguous (the
one contract the float4 state loads rely on).
2. The kernel's reconstructed slot-addressing formula
``base + slot*stride(0) + i_hv*V*K + v*K + k`` resolves to the exact same
storage element as torch's native ``ssm_states[slot, i_hv, v, k]`` on the
strided pool, while the pre-fix dense-pitch formula
``base + slot*(HV*V*K) + ...`` resolves ELSEWHERE for every slot > 0.
Hardcoding the dense pitch back into the ``.cuh`` turns test (2) red.
Runs on CPU: only the Python gate and the (dtype/stride-only) addressing
arithmetic execute — no CUDA kernel is launched.
python -m pytest test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py -v
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
import unittest
import torch
from sglang.kernels.ops.attention.kda_fused_decode import (
_CONV_STATE_W,
covered,
)
from sglang.srt.mem_cache.layout.page_major import (
build_page_major_mamba_views,
mamba_entry_bytes,
)
_DEV = "cpu"
# The kernel is compiled for the K3 KDA decode regime; covered() enforces these
# supported local head counts. Multi-layer + several slots so the envelope slot
# pitch differs from the dense H*V*K pitch.
_KDA_HEADS = (3, 6, 12)
_V = 128
_K = 128
_LAYERS = 3
_LAYER_UNDER_TEST = 1
_SLOTS = 6
_CONV_SHAPES = ((3, 8),) # tiny bf16 conv region interleaves the temporal region
_CONV_DTYPE = torch.bfloat16
_TEMPORAL_DTYPE = torch.float32
def _seg(heads: int) -> int:
return heads * _V
def _conv_dim(heads: int) -> int:
return 3 * _seg(heads)
def _make_strided_temporal_view(heads: int):
"""Envelope-strided temporal (SSM) view as UnifiedMambaPool / the
page-major MambaPool serve it to the KDA backend: shape
``(num_layers, max_slots, H, V, K)`` with the slot pitch spanning ALL
layers' state, not H*V*K."""
entry = mamba_entry_bytes(
layer_num=_LAYERS,
conv_state_shapes=_CONV_SHAPES,
conv_dtype=_CONV_DTYPE,
temporal_state_shape=(heads, _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=(heads, _V, _K),
temporal_dtype=_TEMPORAL_DTYPE,
max_slots=_SLOTS,
)
return raw, temporal
def _make_covered_side_args(batch: int, heads: int):
"""The non-ssm covered() arguments, in the exact K3 shapes/dtypes so the
gate turns solely on the ssm_states view under test."""
bf16 = torch.bfloat16
seg = _seg(heads)
conv_dim = _conv_dim(heads)
mixed_qkv = torch.zeros((batch, conv_dim), dtype=bf16, device=_DEV)
a = torch.zeros((batch, seg), dtype=bf16, device=_DEV)
b = torch.zeros((batch, heads), dtype=bf16, device=_DEV)
onorm_g = torch.zeros((batch, seg), dtype=bf16, device=_DEV)
conv_states = torch.zeros(
(_SLOTS, _CONV_STATE_W, conv_dim), dtype=bf16, device=_DEV
)
cache_indices = torch.zeros((batch,), dtype=torch.int32, device=_DEV)
return mixed_qkv, a, b, conv_states, onorm_g, cache_indices
def _addressing_samples(heads: int):
return [
(0, 0, 0, 0),
(5, heads - 1, 127, 127),
(3, heads // 2, 64, 100),
(1, 0, 0, 1),
(2, min(heads - 1, 2), 3, 7),
]
class TestKdaFusedDecodeStridedState(unittest.TestCase):
def test_covered_accepts_envelope_strided_and_rejects_noncontiguous_inner(self):
for heads in _KDA_HEADS:
with self.subTest(kda_heads=heads):
_raw, temporal = _make_strided_temporal_view(heads)
ssm = temporal[_LAYER_UNDER_TEST] # what mamba2_layer_cache serves
# Precondition: the pool really is envelope-strided (else the
# property below would be vacuous — a dense pool passes the old
# gate too).
self.assertNotEqual(
ssm.stride(0),
heads * _V * _K,
"test setup no longer produces a strided pool",
)
# Inner [H, V, K] IS contiguous — the contract the kernel's
# float4 state loads rely on and all covered() must still require.
self.assertEqual(
(ssm.stride(-1), ssm.stride(-2), ssm.stride(-3)),
(1, _K, _V * _K),
)
(
mixed_qkv,
a,
b,
conv_states,
onorm_g,
cache_indices,
) = _make_covered_side_args(batch=2, heads=heads)
# (1) Accept the envelope-strided view — pre-fix
# .view(...).is_contiguous() would reject this and drop decode
# to the unfused chain.
self.assertTrue(
covered(mixed_qkv, a, b, conv_states, ssm, cache_indices, onorm_g),
"covered() rejected the envelope-strided ssm pool (fused decode "
"would silently fall back to the unfused chain)",
)
# Still reject a view whose inner dims are NOT contiguous: the
# kernel cannot float4-load a transposed [.., V, K] state.
ssm_bad = ssm.transpose(-1, -2) # stride(-1) == K, not 1
self.assertFalse(
covered(
mixed_qkv,
a,
b,
conv_states,
ssm_bad,
cache_indices,
onorm_g,
),
"covered() must reject a non-inner-contiguous ssm view",
)
def test_slot_formula_resolves_correct_element_and_dense_pitch_misaddresses(self):
for heads in _KDA_HEADS:
with self.subTest(kda_heads=heads):
raw, temporal = _make_strided_temporal_view(heads)
ssm = temporal[_LAYER_UNDER_TEST]
# Distinct value per storage element so an offset that lands
# elsewhere reads a provably different value.
raw_fp32 = raw.view(torch.float32)
raw_fp32.copy_(torch.arange(raw_fp32.numel(), dtype=torch.float32))
base = ssm.storage_offset()
# == state.stride(0) the wrapper passes the kernel.
slot_stride = ssm.stride(0)
dense_pitch = heads * _V * _K # the pre-fix hardcoded slot pitch
for slot, i_hv, v, k in _addressing_samples(heads):
intra = (
i_hv * (_V * _K) + v * _K + k
) # kernel's hardcoded intra-slot offset
kernel_off = base + slot * slot_stride + intra
# (2a) The kernel formula names exactly the element torch
# indexing names — proves slot*stride(0) + intra addresses
# the intended slot.
self.assertEqual(
raw_fp32[kernel_off].item(),
ssm[slot, i_hv, v, k].item(),
f"kernel slot formula mis-addressed "
f"(slot={slot}, i_hv={i_hv}, heads={heads})",
)
# (2b) The pre-fix dense-pitch formula lands on a DIFFERENT
# element (a different layer's envelope region) for every
# slot > 0.
dense_off = base + slot * dense_pitch + intra
if slot > 0:
self.assertNotEqual(
raw_fp32[dense_off].item(),
ssm[slot, i_hv, v, k].item(),
f"dense-pitch formula happened to match at "
f"slot={slot}, heads={heads}; the fix would not "
"be load-bearing",
)
if __name__ == "__main__":
unittest.main()