feat(kernels): port standalone Kimi K3 kernels (#32890)
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
co-authored by
Claude Opus 5
hnyls2002
zhangxiaolei
Claude Opus 5
hnyls2002
zhangxiaolei
parent
a1344fad4e
commit
fb207b72b0
@@ -0,0 +1,215 @@
|
||||
"""Kimi-K3 fused KDA decode must match the existing unfused decode chain.
|
||||
|
||||
The fused kernel replaces:
|
||||
|
||||
causal_conv1d_update -> kda_packed_decode -> sigmoid-gated RMSNorm
|
||||
|
||||
This file covers the local head layouts used by Kimi-K3 TP8/TP16/TP32:
|
||||
H = 12/6/3. The H=6 and H=3 cases are the branches added by the fixed-head
|
||||
dispatch in ``kda_fused_decode.cuh``.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.attention import kda_fused_decode
|
||||
from sglang.kernels.ops.attention.fla.fused_norm_gate import rms_norm_gated
|
||||
from sglang.kernels.ops.attention.fla.fused_recurrent import (
|
||||
fused_recurrent_kda_packed_decode,
|
||||
)
|
||||
from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
_HEAD_DIM = 128
|
||||
_CONV_STATE_W = 3
|
||||
_SLOTS = 8
|
||||
_BATCH = 4
|
||||
|
||||
|
||||
def _randn(shape, dtype, generator, scale=1.0):
|
||||
return (torch.randn(shape, device="cuda", generator=generator) * scale).to(dtype)
|
||||
|
||||
|
||||
def _make_case(heads: int, seed: int):
|
||||
generator = torch.Generator(device="cuda").manual_seed(seed)
|
||||
seg = heads * _HEAD_DIM
|
||||
conv_dim = 3 * seg
|
||||
|
||||
# Keep magnitudes moderate so fp32 state updates stay in a stable range.
|
||||
mixed_qkv = _randn((_BATCH, conv_dim), torch.bfloat16, generator, scale=0.2)
|
||||
a = _randn((_BATCH, seg), torch.bfloat16, generator, scale=0.2)
|
||||
b = _randn((_BATCH, heads), torch.bfloat16, generator, scale=0.2)
|
||||
onorm_g = _randn((_BATCH, seg), torch.bfloat16, generator, scale=0.2)
|
||||
|
||||
conv_states = _randn(
|
||||
(_SLOTS, _CONV_STATE_W, conv_dim), torch.bfloat16, generator, scale=0.2
|
||||
)
|
||||
ssm_states = _randn(
|
||||
(_SLOTS, heads, _HEAD_DIM, _HEAD_DIM), torch.float32, generator, scale=0.02
|
||||
)
|
||||
cache_indices = torch.arange(_BATCH, device="cuda", dtype=torch.int32)
|
||||
|
||||
conv_weights = _randn((conv_dim, 4), torch.float32, generator, scale=0.1)
|
||||
conv_bias = _randn((conv_dim,), torch.float32, generator, scale=0.05)
|
||||
a_log = _randn((heads,), torch.float32, generator, scale=0.1)
|
||||
dt_bias = _randn((seg,), torch.float32, generator, scale=0.1)
|
||||
onorm_weight = _randn((_HEAD_DIM,), torch.float32, generator, scale=0.1) + 1.0
|
||||
|
||||
return (
|
||||
mixed_qkv,
|
||||
a,
|
||||
b,
|
||||
onorm_g,
|
||||
conv_states,
|
||||
ssm_states,
|
||||
cache_indices,
|
||||
conv_weights,
|
||||
conv_bias,
|
||||
a_log,
|
||||
dt_bias,
|
||||
onorm_weight,
|
||||
)
|
||||
|
||||
|
||||
def _run_unfused_reference(
|
||||
mixed_qkv,
|
||||
a,
|
||||
b,
|
||||
onorm_g,
|
||||
conv_states,
|
||||
ssm_states,
|
||||
cache_indices,
|
||||
conv_weights,
|
||||
conv_bias,
|
||||
a_log,
|
||||
dt_bias,
|
||||
onorm_weight,
|
||||
):
|
||||
heads = ssm_states.shape[-3]
|
||||
qkv = causal_conv1d_update(
|
||||
mixed_qkv,
|
||||
conv_states.transpose(-1, -2),
|
||||
conv_weights,
|
||||
conv_bias,
|
||||
activation="silu",
|
||||
conv_state_indices=cache_indices,
|
||||
)
|
||||
out = torch.empty(
|
||||
(_BATCH, 1, heads, _HEAD_DIM), dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
out, _ = fused_recurrent_kda_packed_decode(
|
||||
qkv,
|
||||
a,
|
||||
b,
|
||||
a_log,
|
||||
dt_bias,
|
||||
_HEAD_DIM**-0.5,
|
||||
ssm_states,
|
||||
out,
|
||||
cache_indices,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
ref = rms_norm_gated(
|
||||
out,
|
||||
onorm_g.view(1, _BATCH, heads, _HEAD_DIM),
|
||||
onorm_weight,
|
||||
None,
|
||||
activation="sigmoid",
|
||||
eps=1e-6,
|
||||
)
|
||||
return ref.transpose(0, 1).contiguous()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
||||
@pytest.mark.parametrize(
|
||||
"heads,tp_size",
|
||||
[
|
||||
pytest.param(3, 32, id="tp32_h3"),
|
||||
pytest.param(6, 16, id="tp16_h6"),
|
||||
pytest.param(12, 8, id="tp8_h12"),
|
||||
],
|
||||
)
|
||||
def test_kda_fused_decode_matches_unfused_chain(heads: int, tp_size: int):
|
||||
(
|
||||
mixed_qkv,
|
||||
a,
|
||||
b,
|
||||
onorm_g,
|
||||
conv_states,
|
||||
ssm_states,
|
||||
cache_indices,
|
||||
conv_weights,
|
||||
conv_bias,
|
||||
a_log,
|
||||
dt_bias,
|
||||
onorm_weight,
|
||||
) = _make_case(heads=heads, seed=20260731 + tp_size)
|
||||
|
||||
conv_ref = conv_states.clone()
|
||||
conv_fused = conv_states.clone()
|
||||
state_ref = ssm_states.clone()
|
||||
state_fused = ssm_states.clone()
|
||||
|
||||
w_q_t, w_k_t, w_v_t = [
|
||||
weight.t().contiguous()
|
||||
for weight in conv_weights.split(heads * _HEAD_DIM, dim=0)
|
||||
]
|
||||
|
||||
assert kda_fused_decode.covered(
|
||||
mixed_qkv,
|
||||
a,
|
||||
b,
|
||||
conv_fused,
|
||||
state_fused,
|
||||
cache_indices,
|
||||
onorm_g,
|
||||
)
|
||||
|
||||
ref = _run_unfused_reference(
|
||||
mixed_qkv.clone(),
|
||||
a,
|
||||
b,
|
||||
onorm_g,
|
||||
conv_ref,
|
||||
state_ref,
|
||||
cache_indices,
|
||||
conv_weights,
|
||||
conv_bias,
|
||||
a_log,
|
||||
dt_bias,
|
||||
onorm_weight,
|
||||
)
|
||||
fused = kda_fused_decode.kda_fused_decode(
|
||||
mixed_qkv.clone(),
|
||||
a,
|
||||
b,
|
||||
conv_fused,
|
||||
w_q_t,
|
||||
w_k_t,
|
||||
w_v_t,
|
||||
conv_bias,
|
||||
a_log,
|
||||
dt_bias,
|
||||
onorm_g,
|
||||
onorm_weight,
|
||||
state_fused,
|
||||
cache_indices,
|
||||
scale=_HEAD_DIM**-0.5,
|
||||
onorm_eps=1e-6,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# JIT log breadcrumb for PR/CI evidence that the fused fixed-head branch ran.
|
||||
print(f"K3 fused KDA decode test used fused path: TP{tp_size}, H={heads}")
|
||||
|
||||
torch.testing.assert_close(fused, ref, rtol=2e-2, atol=2e-2)
|
||||
torch.testing.assert_close(state_fused, state_ref, rtol=2e-2, atol=2e-2)
|
||||
torch.testing.assert_close(conv_fused, conv_ref, rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,166 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.kernels.ops.attention.fla.kda import chunk_kda
|
||||
from sglang.kernels.ops.attention.linear.kda_nvidia_prefill import (
|
||||
chunk_kda_fwd as nvidia_chunk_kda_fwd,
|
||||
)
|
||||
from sglang.kernels.ops.attention.linear.kda_ptx_prefill import (
|
||||
chunk_kda_fwd as ptx_chunk_kda_fwd,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=180, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
register_cuda_ci(est_time=180, stage="base-c", runner_config="4-gpu-gb300")
|
||||
|
||||
|
||||
def _inputs(seed, seq_len=128):
|
||||
generator = torch.Generator(device="cuda").manual_seed(seed)
|
||||
batch_size, num_heads, head_dim = 1, 2, 128
|
||||
shape = (batch_size, seq_len, num_heads, head_dim)
|
||||
q = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||
v = (
|
||||
0.1
|
||||
* torch.randn(
|
||||
shape,
|
||||
generator=generator,
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
).to(torch.bfloat16)
|
||||
gate = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||
beta_logits = torch.randn(
|
||||
shape[:-1],
|
||||
generator=generator,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
a_log = torch.randn(
|
||||
num_heads, generator=generator, device="cuda", dtype=torch.float32
|
||||
)
|
||||
dt_bias = torch.randn(
|
||||
num_heads * head_dim,
|
||||
generator=generator,
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
state = torch.zeros(
|
||||
batch_size,
|
||||
num_heads,
|
||||
head_dim,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
return q, k, v, gate, beta_logits, a_log, dt_bias, state
|
||||
|
||||
|
||||
def _reference(q, k, v, gate, beta, a_log, dt_bias, state, fused_qk_norm):
|
||||
return chunk_kda(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=gate,
|
||||
beta=beta,
|
||||
scale=q.shape[-1] ** -0.5,
|
||||
initial_state=state,
|
||||
initial_state_indices=torch.arange(
|
||||
q.shape[0], device="cuda", dtype=torch.int32
|
||||
),
|
||||
use_qk_l2norm_in_kernel=fused_qk_norm,
|
||||
A_log=a_log,
|
||||
dt_bias=dt_bias,
|
||||
lower_bound=-5.0,
|
||||
)
|
||||
|
||||
|
||||
class TestKdaPrefill(CustomTestCase):
|
||||
@torch.inference_mode()
|
||||
def test_nvidia_prefill(self):
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
|
||||
self.skipTest("NVIDIA KDA prefill requires datacenter Blackwell")
|
||||
q, k, v, gate, beta_logits, a_log, dt_bias, state = _inputs(0)
|
||||
q = F.normalize(q.float(), dim=-1).to(torch.bfloat16)
|
||||
k = F.normalize(k.float(), dim=-1).to(torch.bfloat16)
|
||||
beta = torch.sigmoid(beta_logits.float()).to(torch.bfloat16)
|
||||
actual, actual_state = nvidia_chunk_kda_fwd(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=gate,
|
||||
beta=beta,
|
||||
scale=q.shape[-1] ** -0.5,
|
||||
initial_state=state.transpose(-1, -2).contiguous(),
|
||||
output_final_state=True,
|
||||
safe_gate=True,
|
||||
lower_bound=-5.0,
|
||||
use_gate_in_kernel=True,
|
||||
A_log=a_log,
|
||||
dt_bias=dt_bias,
|
||||
)[:2]
|
||||
expected = _reference(
|
||||
q, k, v, gate, beta, a_log, dt_bias, state, fused_qk_norm=False
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
actual.float(), expected.float(), rtol=2e-2, atol=3e-2
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
actual_state.transpose(-1, -2),
|
||||
state,
|
||||
rtol=2e-2,
|
||||
atol=3e-2,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_ptx_prefill(self):
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (
|
||||
10,
|
||||
3,
|
||||
):
|
||||
self.skipTest("PTX KDA prefill requires GB300")
|
||||
q, k, v, gate, beta_logits, a_log, dt_bias, state = _inputs(1)
|
||||
actual, actual_state = ptx_chunk_kda_fwd(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=gate,
|
||||
beta=beta_logits,
|
||||
scale=q.shape[-1] ** -0.5,
|
||||
initial_state=state.transpose(-1, -2).contiguous(),
|
||||
output_final_state=True,
|
||||
safe_gate=True,
|
||||
lower_bound=-5.0,
|
||||
use_gate_in_kernel=True,
|
||||
A_log=a_log,
|
||||
dt_bias=dt_bias,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
use_beta_sigmoid_in_kernel=True,
|
||||
)[:2]
|
||||
expected = _reference(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
gate,
|
||||
torch.sigmoid(beta_logits.float()).to(torch.bfloat16),
|
||||
a_log,
|
||||
dt_bias,
|
||||
state,
|
||||
fused_qk_norm=True,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
actual.float(), expected.float(), rtol=2e-2, atol=3e-2
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
actual_state.transpose(-1, -2),
|
||||
state,
|
||||
rtol=2e-2,
|
||||
atol=3e-2,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,359 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
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
|
||||
from sglang.kernels.ops.communication.mp import register_comm_cleanup
|
||||
from sglang.kernels.ops.kimi_k3 import (
|
||||
all_reduce,
|
||||
attn_res,
|
||||
gemm_ag,
|
||||
gemm_ar,
|
||||
sp_collective,
|
||||
)
|
||||
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="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
register_cuda_ci(est_time=480, suite="nightly-8-gpu-b200", nightly=True)
|
||||
|
||||
_HIDDEN_SIZE = 7168
|
||||
_GEMM_AR_K_TOTAL = 12288
|
||||
_GEMM_AG_WORLD_SIZE = 8
|
||||
_MB = 1024 * 1024
|
||||
_SP_TUNING = sp_collective.Tuning(num_blocks=1, block_size=256)
|
||||
|
||||
|
||||
def _device():
|
||||
return torch.device("cuda", int(os.environ["LOCAL_RANK"]))
|
||||
|
||||
|
||||
def _require_sm100():
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0):
|
||||
pytest.skip("Kimi K3 collectives require SM100+")
|
||||
|
||||
|
||||
@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)
|
||||
cpu_group = coord.cpu_group
|
||||
assert isinstance(cpu_group, dist.ProcessGroup)
|
||||
nccl_group = dist.new_group(backend="nccl", device_id=_device())
|
||||
return cpu_group, nccl_group
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_comm():
|
||||
cpu_group, _ = _init_world()
|
||||
comm = CustomAllReduceV2(
|
||||
cpu_group,
|
||||
_device(),
|
||||
max_pull_size=4 * _MB,
|
||||
max_push_size=4 * _MB,
|
||||
)
|
||||
if comm.disabled or comm.mc_base_ptr == 0:
|
||||
raise RuntimeError("Kimi K3 collectives require multicast symmetric memory")
|
||||
all_reduce.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
|
||||
sp_collective.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
|
||||
attn_res.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
|
||||
register_comm_cleanup(comm)
|
||||
return comm
|
||||
|
||||
|
||||
@cache_once
|
||||
def _init_gemm_ar():
|
||||
cpu_group, _ = _init_world()
|
||||
world_size = dist.get_world_size()
|
||||
gemm_ar.init(
|
||||
world_size=world_size,
|
||||
rank=dist.get_rank(),
|
||||
group=cpu_group,
|
||||
k=_GEMM_AR_K_TOTAL // world_size,
|
||||
)
|
||||
|
||||
|
||||
def _symmetric_tensor(shape):
|
||||
from torch._C._distributed_c10d import _SymmetricMemory
|
||||
|
||||
cpu_group, _ = _init_world()
|
||||
tensor = _SymmetricMemory.empty_strided_p2p(
|
||||
shape,
|
||||
torch.empty(shape).stride(),
|
||||
torch.bfloat16,
|
||||
_device(),
|
||||
cpu_group.group_name,
|
||||
)
|
||||
handle = _SymmetricMemory.rendezvous(tensor)
|
||||
rank = dist.get_rank()
|
||||
multicast_ptr = (
|
||||
int(handle.multicast_ptr) + tensor.data_ptr() - int(handle.buffer_ptrs[rank])
|
||||
)
|
||||
if multicast_ptr == 0:
|
||||
raise RuntimeError("symmetric tensor has no multicast mapping")
|
||||
return tensor, handle, multicast_ptr
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_all_reduce_push():
|
||||
_require_sm100()
|
||||
comm = _init_comm()
|
||||
rank = dist.get_rank()
|
||||
generator = torch.Generator().manual_seed(10 + rank)
|
||||
x = torch.randint(
|
||||
0,
|
||||
16,
|
||||
(_HIDDEN_SIZE,),
|
||||
generator=generator,
|
||||
dtype=torch.bfloat16,
|
||||
).to(_device())
|
||||
residual = (
|
||||
torch.arange(_HIDDEN_SIZE, dtype=torch.int32, device=_device())
|
||||
.remainder_(7)
|
||||
.to(torch.bfloat16)
|
||||
)
|
||||
expected = x.clone()
|
||||
_, nccl_group = _init_world()
|
||||
dist.all_reduce(expected, group=nccl_group)
|
||||
expected += residual
|
||||
|
||||
all_reduce.all_reduce_push_res(
|
||||
comm.world_size,
|
||||
x,
|
||||
residual,
|
||||
ws_mc_base=comm.mc_base_ptr,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(x, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_sequence_parallel_collectives():
|
||||
_require_sm100()
|
||||
comm = _init_comm()
|
||||
rank, world_size = dist.get_rank(), comm.world_size
|
||||
local_tokens = 2
|
||||
generator = torch.Generator(device="cuda").manual_seed(20 + rank)
|
||||
reduce_input = torch.randn(
|
||||
world_size * local_tokens,
|
||||
_HIDDEN_SIZE,
|
||||
generator=generator,
|
||||
device=_device(),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
residual = torch.randn(
|
||||
local_tokens,
|
||||
_HIDDEN_SIZE,
|
||||
generator=torch.Generator(device="cuda").manual_seed(21),
|
||||
device=_device(),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
expected_reduce = reduce_input.float()
|
||||
_, nccl_group = _init_world()
|
||||
dist.all_reduce(expected_reduce, group=nccl_group)
|
||||
lo = rank * local_tokens
|
||||
expected_reduce = (expected_reduce[lo : lo + local_tokens] + residual.float()).to(
|
||||
torch.bfloat16
|
||||
)
|
||||
reduce_output = torch.empty_like(expected_reduce)
|
||||
sp_collective.reduce_scatter_res(
|
||||
world_size,
|
||||
reduce_input,
|
||||
reduce_output,
|
||||
residual,
|
||||
tuning=_SP_TUNING,
|
||||
)
|
||||
|
||||
gather_input = torch.randn(
|
||||
local_tokens,
|
||||
_HIDDEN_SIZE,
|
||||
generator=generator,
|
||||
device=_device(),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
expected_gather = torch.empty(
|
||||
world_size * local_tokens,
|
||||
_HIDDEN_SIZE,
|
||||
device=_device(),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
dist.all_gather_into_tensor(
|
||||
expected_gather,
|
||||
gather_input,
|
||||
group=nccl_group,
|
||||
)
|
||||
gather_output = torch.empty_like(expected_gather)
|
||||
sp_collective.all_gather(
|
||||
world_size,
|
||||
gather_input,
|
||||
gather_output,
|
||||
ws_mc_base=comm.mc_base_ptr,
|
||||
tuning=_SP_TUNING,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(reduce_output, expected_reduce, rtol=2e-2, atol=3e-2)
|
||||
torch.testing.assert_close(gather_output, expected_gather, rtol=0, atol=0)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_gemm_all_gather():
|
||||
_require_sm100()
|
||||
if int(os.environ["WORLD_SIZE"]) != _GEMM_AG_WORLD_SIZE:
|
||||
pytest.skip("Kimi K3 gemm_ag is compiled for TP8")
|
||||
comm = _init_comm()
|
||||
generator = torch.Generator().manual_seed(30)
|
||||
x = (
|
||||
(torch.randn(1, gemm_ag.K, generator=generator) * 0.05)
|
||||
.to(torch.bfloat16)
|
||||
.to(_device())
|
||||
)
|
||||
weight = (
|
||||
(torch.randn(gemm_ag.N, gemm_ag.K, generator=generator) * 0.05)
|
||||
.to(torch.bfloat16)
|
||||
.to(_device())
|
||||
)
|
||||
bias = torch.randn(1, gemm_ag.N, generator=generator).to(
|
||||
device=_device(), dtype=torch.bfloat16
|
||||
)
|
||||
output = torch.empty(1, gemm_ag.N, device=_device(), dtype=torch.bfloat16)
|
||||
expected = (x.float() @ weight.float().t() + bias.float()).to(torch.bfloat16)
|
||||
|
||||
gemm_ag.gemm_ag_up_proj(
|
||||
comm.world_size,
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
None,
|
||||
output,
|
||||
ws_mc_base=comm.mc_base_ptr,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(output, expected, rtol=3e-2, atol=3e-2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_gemm_all_reduce():
|
||||
_require_sm100()
|
||||
_init_gemm_ar()
|
||||
rank, world_size = dist.get_rank(), dist.get_world_size()
|
||||
local_k = _GEMM_AR_K_TOTAL // world_size
|
||||
generator = torch.Generator().manual_seed(40 + rank)
|
||||
x = torch.randn(1, local_k, generator=generator).to(
|
||||
device=_device(), dtype=torch.bfloat16
|
||||
)
|
||||
weight = torch.randn(gemm_ar.N, local_k, generator=generator).to(
|
||||
device=_device(), dtype=torch.bfloat16
|
||||
)
|
||||
expected = (x.float() @ weight.float().t()).to(torch.bfloat16).float()
|
||||
_, nccl_group = _init_world()
|
||||
dist.all_reduce(expected, group=nccl_group)
|
||||
|
||||
output = gemm_ar.o_proj_gemm_ar(x, weight)
|
||||
torch.cuda.synchronize()
|
||||
bad = ((output.float() - expected).abs() > 0.05 + 0.02 * expected.abs()).sum()
|
||||
assert bad.item() <= output.numel() / 1000
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_attention_residual_direct_all_gather():
|
||||
_require_sm100()
|
||||
comm = _init_comm()
|
||||
rank, local_tokens, num_bank_rows = dist.get_rank(), 2, 3
|
||||
generator = torch.Generator(device="cuda").manual_seed(50 + rank)
|
||||
prefix = torch.randn(
|
||||
local_tokens,
|
||||
_HIDDEN_SIZE,
|
||||
generator=generator,
|
||||
device=_device(),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
bank = torch.randn(
|
||||
local_tokens,
|
||||
num_bank_rows + 1,
|
||||
_HIDDEN_SIZE,
|
||||
generator=generator,
|
||||
device=_device(),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
combine_weight = torch.linspace(
|
||||
-0.01, 0.01, _HIDDEN_SIZE, device=_device(), dtype=torch.bfloat16
|
||||
)
|
||||
output_weight = torch.linspace(
|
||||
1.25, 0.75, _HIDDEN_SIZE, device=_device(), dtype=torch.bfloat16
|
||||
)
|
||||
local_reference = torch.empty_like(prefix)
|
||||
attn_res.attn_res_fused_tma(
|
||||
prefix,
|
||||
bank.clone(),
|
||||
combine_weight,
|
||||
output_weight,
|
||||
local_reference,
|
||||
num_bank_rows,
|
||||
1e-6,
|
||||
)
|
||||
full_reference = torch.empty(
|
||||
comm.world_size * local_tokens,
|
||||
_HIDDEN_SIZE,
|
||||
device=_device(),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
_, nccl_group = _init_world()
|
||||
dist.all_gather_into_tensor(
|
||||
full_reference,
|
||||
local_reference,
|
||||
group=nccl_group,
|
||||
)
|
||||
|
||||
output, handle, multicast_ptr = _symmetric_tensor(tuple(full_reference.shape))
|
||||
attn_res.attn_res_fused_direct_ag(
|
||||
comm.world_size,
|
||||
prefix,
|
||||
bank,
|
||||
combine_weight,
|
||||
output_weight,
|
||||
output,
|
||||
num_bank_rows,
|
||||
1e-6,
|
||||
output_mc_ptr=multicast_ptr,
|
||||
max_blocks=4,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(output, full_reference, rtol=2e-2, atol=3e-2)
|
||||
assert handle is not None
|
||||
|
||||
|
||||
def _precompile(num_gpus):
|
||||
for world_size in num_gpus:
|
||||
all_reduce._jit_module(world_size)
|
||||
sp_collective._jit_module(world_size)
|
||||
gemm_ar._jit_module(_GEMM_AR_K_TOTAL // world_size, world_size)
|
||||
if _GEMM_AG_WORLD_SIZE in num_gpus:
|
||||
gemm_ag._jit_module()
|
||||
attn_res._jit_fused_tma_module(4, 1, 200)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
multigpu_pytest_main(
|
||||
__name__,
|
||||
__file__,
|
||||
num_gpus=(4, 8),
|
||||
pre_launch_fn=_precompile,
|
||||
)
|
||||
@@ -0,0 +1,453 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import (
|
||||
commit_kda_replayssm_spec,
|
||||
)
|
||||
from sglang.kernels.ops.kimi_k3 import (
|
||||
situ_and_mul,
|
||||
situ_and_mul_masked_post_quant,
|
||||
)
|
||||
from sglang.kernels.ops.kimi_k3.attn_res import attn_res_fused_tma
|
||||
from sglang.kernels.ops.kimi_k3.kda_decode_mtp import (
|
||||
fused_kda_decode_mtp_dspark,
|
||||
)
|
||||
from sglang.kernels.ops.kimi_k3.mla_output_gate import (
|
||||
covered,
|
||||
kimi_k3_mla_output_gate,
|
||||
)
|
||||
from sglang.kernels.ops.moe.moe_front import (
|
||||
NUM_EXPERTS,
|
||||
TOPK,
|
||||
fused_front,
|
||||
)
|
||||
from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate
|
||||
from sglang.srt.utils import get_device_sm
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
|
||||
_HIDDEN_SIZE = 7168
|
||||
_GROUP_SIZE = 128
|
||||
_BETA = 4.0
|
||||
_LINEAR_BETA = 25.0
|
||||
|
||||
|
||||
def _situ_reference(gate_up):
|
||||
gate, up = gate_up.chunk(2, dim=-1)
|
||||
gate = gate.float()
|
||||
up = up.float()
|
||||
return (
|
||||
_BETA
|
||||
* torch.tanh(gate / _BETA)
|
||||
* torch.sigmoid(gate)
|
||||
* _LINEAR_BETA
|
||||
* torch.tanh(up / _LINEAR_BETA)
|
||||
)
|
||||
|
||||
|
||||
def _unpack_ue8m0_scales(packed, num_groups):
|
||||
num_experts, groups_per_word, num_tokens = packed.shape
|
||||
exponents = packed.contiguous().view(torch.uint8)
|
||||
exponents = exponents.view(num_experts, groups_per_word, num_tokens, 4)
|
||||
exponents = exponents.permute(0, 2, 1, 3).reshape(
|
||||
num_experts, num_tokens, num_groups
|
||||
)
|
||||
return torch.exp2(exponents.float() - 127.0)
|
||||
|
||||
|
||||
class TestKimiK3ComputeKernels(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA is not available")
|
||||
if get_device_sm() < 100:
|
||||
raise unittest.SkipTest("Kimi K3 compute kernels require SM100a+")
|
||||
|
||||
def test_attn_residual_and_prefix_write(self):
|
||||
generator = torch.Generator(device="cuda").manual_seed(0)
|
||||
|
||||
def randn(*shape):
|
||||
return torch.randn(*shape, generator=generator, device="cuda")
|
||||
|
||||
num_tokens, num_bank_rows, num_valid_bank_rows = 5, 8, 5
|
||||
prefix = randn(num_tokens, _HIDDEN_SIZE).to(torch.bfloat16)
|
||||
bank = randn(num_tokens, num_bank_rows, _HIDDEN_SIZE).to(torch.bfloat16)
|
||||
combine_weight = (randn(_HIDDEN_SIZE) * _HIDDEN_SIZE**-0.5).to(torch.bfloat16)
|
||||
output_weight = (1 + 0.1 * randn(_HIDDEN_SIZE)).to(torch.bfloat16)
|
||||
output = torch.empty_like(prefix)
|
||||
|
||||
rows = torch.cat(
|
||||
[
|
||||
bank[:, :num_valid_bank_rows].float(),
|
||||
prefix.unsqueeze(1).float(),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
rms = torch.rsqrt(rows.square().mean(-1) + 1e-6)
|
||||
scores = (rows * combine_weight.float()).sum(-1) * rms
|
||||
mixed = (torch.softmax(scores, dim=-1).unsqueeze(-1) * rows).sum(1)
|
||||
expected = (
|
||||
mixed
|
||||
* torch.rsqrt(mixed.square().mean(-1, keepdim=True) + 1e-6)
|
||||
* output_weight.float()
|
||||
)
|
||||
|
||||
attn_res_fused_tma(
|
||||
prefix,
|
||||
bank,
|
||||
combine_weight,
|
||||
output_weight,
|
||||
output,
|
||||
num_valid_bank_rows,
|
||||
1e-6,
|
||||
write_prefix=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output.float(), expected, rtol=2e-2, atol=4e-2)
|
||||
self.assertTrue(torch.equal(bank[:, num_valid_bank_rows], prefix))
|
||||
|
||||
def test_mla_output_gate(self):
|
||||
generator = torch.Generator(device="cuda").manual_seed(1)
|
||||
shape = (5, 12, 128)
|
||||
x = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16)
|
||||
gate = torch.randn(
|
||||
shape, generator=generator, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
self.assertTrue(covered(x, gate))
|
||||
expected = x * torch.sigmoid(gate).to(torch.bfloat16)
|
||||
self.assertTrue(torch.equal(kimi_k3_mla_output_gate(x, gate), expected))
|
||||
|
||||
def test_situ_and_mul(self):
|
||||
generator = torch.Generator(device="cuda").manual_seed(2)
|
||||
hidden_size = 1024
|
||||
storage = torch.randn(
|
||||
(7, 2 * hidden_size + 16),
|
||||
generator=generator,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
gate_up = storage[:, : 2 * hidden_size]
|
||||
output = torch.empty(
|
||||
(gate_up.shape[0], hidden_size),
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
returned = situ_and_mul(gate_up, output, beta=_BETA, linear_beta=_LINEAR_BETA)
|
||||
|
||||
self.assertIs(returned, output)
|
||||
torch.testing.assert_close(
|
||||
returned.float(),
|
||||
_situ_reference(gate_up).to(torch.bfloat16).float(),
|
||||
rtol=2e-2,
|
||||
atol=4e-2,
|
||||
)
|
||||
|
||||
def test_situ_mul_quant(self):
|
||||
torch.cuda.manual_seed_all(3)
|
||||
num_experts, num_tokens, hidden_size, topk = 8, 32, 1024, 16
|
||||
gate_up = (
|
||||
torch.randn(
|
||||
num_experts,
|
||||
num_tokens,
|
||||
2 * hidden_size,
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
* 2.0
|
||||
).to(torch.bfloat16)
|
||||
masked_m = torch.randint(
|
||||
0,
|
||||
num_tokens + 1,
|
||||
(num_experts,),
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
)
|
||||
masked_m[0] = 0
|
||||
masked_m[-1] = num_tokens
|
||||
|
||||
output = torch.full(
|
||||
(num_experts, num_tokens, hidden_size),
|
||||
0x7F,
|
||||
device="cuda",
|
||||
dtype=torch.uint8,
|
||||
).view(torch.float8_e4m3fn)
|
||||
num_groups = hidden_size // _GROUP_SIZE
|
||||
output_scale = torch.zeros(
|
||||
(num_experts, num_groups // 4, num_tokens),
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
situ_and_mul_masked_post_quant(
|
||||
input=gate_up,
|
||||
output=output,
|
||||
output_scale=output_scale,
|
||||
quant_group_size=_GROUP_SIZE,
|
||||
masked_m=masked_m,
|
||||
beta=_BETA,
|
||||
linear_beta=_LINEAR_BETA,
|
||||
scale_ue8m0=True,
|
||||
topk=topk,
|
||||
transposed=True,
|
||||
)
|
||||
|
||||
scales = _unpack_ue8m0_scales(output_scale, num_groups)
|
||||
expanded_scales = scales.repeat_interleave(_GROUP_SIZE, dim=-1)
|
||||
dequantized = output.float() * expanded_scales
|
||||
expected = _situ_reference(gate_up)
|
||||
error_bound = expanded_scales * 17.0
|
||||
raw_output = output.view(torch.uint8)
|
||||
for expert in range(num_experts):
|
||||
valid_tokens = int(masked_m[expert].item())
|
||||
self.assertTrue(
|
||||
bool(
|
||||
(
|
||||
(
|
||||
dequantized[expert, :valid_tokens]
|
||||
- expected[expert, :valid_tokens]
|
||||
).abs()
|
||||
<= error_bound[expert, :valid_tokens]
|
||||
).all()
|
||||
)
|
||||
)
|
||||
self.assertTrue(bool((raw_output[expert, valid_tokens:] == 0x7F).all()))
|
||||
|
||||
def test_moe_front(self):
|
||||
torch.manual_seed(4)
|
||||
num_tokens, latent_dim = 1, 128
|
||||
hidden = (
|
||||
torch.randn(
|
||||
num_tokens,
|
||||
_HIDDEN_SIZE,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
/ 32
|
||||
)
|
||||
weight = (
|
||||
torch.randn(
|
||||
NUM_EXPERTS + latent_dim,
|
||||
_HIDDEN_SIZE,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
/ 32
|
||||
)
|
||||
bias = torch.randn(NUM_EXPERTS, device="cuda")
|
||||
|
||||
weights, ids, routed = fused_front(
|
||||
hidden,
|
||||
weight,
|
||||
bias,
|
||||
latent_dim,
|
||||
renormalize=True,
|
||||
routed_scaling_factor=2.5,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
merged = torch.mm(hidden, weight.t(), out_dtype=torch.float32)
|
||||
ref_weights, ref_ids = moe_fused_gate(
|
||||
merged[:, :NUM_EXPERTS],
|
||||
bias,
|
||||
topk=TOPK,
|
||||
scoring_func="sigmoid",
|
||||
renormalize=True,
|
||||
routed_scaling_factor=2.5,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
order = ids.argsort(dim=-1)
|
||||
ref_order = ref_ids.argsort(dim=-1)
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
ids.gather(1, order),
|
||||
ref_ids.to(torch.int32).gather(1, ref_order),
|
||||
)
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
weights.gather(1, order),
|
||||
ref_weights.gather(1, ref_order),
|
||||
rtol=1e-6,
|
||||
atol=0,
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
routed,
|
||||
merged[:, NUM_EXPERTS:].to(torch.bfloat16),
|
||||
)
|
||||
)
|
||||
|
||||
def test_mtp_replayssm_ring(self):
|
||||
num_requests, num_heads, num_spec, key_dim = 2, 2, 2, 128
|
||||
num_tokens = num_requests * (1 + num_spec)
|
||||
num_slots, ring_size, conv_width = num_requests + 2, 16, 4
|
||||
|
||||
def run(cache_ring):
|
||||
torch.manual_seed(5)
|
||||
x_q = torch.randn(
|
||||
1,
|
||||
num_tokens,
|
||||
num_heads,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
x_k = torch.randn_like(x_q)
|
||||
x_v = torch.randn_like(x_q)
|
||||
gate = torch.randn_like(x_q)
|
||||
beta = torch.randn(
|
||||
1,
|
||||
num_tokens,
|
||||
num_heads,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
conv_weight = [
|
||||
torch.randn(
|
||||
num_heads * key_dim,
|
||||
conv_width,
|
||||
device="cuda",
|
||||
)
|
||||
* 0.1
|
||||
for _ in range(3)
|
||||
]
|
||||
conv_state = [
|
||||
torch.randn(
|
||||
num_slots,
|
||||
num_heads * key_dim,
|
||||
conv_width - 1,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
for _ in range(3)
|
||||
]
|
||||
slots = torch.arange(1, num_requests + 1, device="cuda", dtype=torch.int32)
|
||||
scratch = torch.arange(num_requests, device="cuda", dtype=torch.int32)
|
||||
state = torch.randn(
|
||||
num_slots,
|
||||
num_heads,
|
||||
key_dim,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
intermediate_conv = torch.zeros(
|
||||
num_requests,
|
||||
1 + num_spec,
|
||||
num_heads * key_dim,
|
||||
conv_width - 1,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kwargs = dict(
|
||||
x_q=x_q,
|
||||
x_k=x_k,
|
||||
x_v=x_v,
|
||||
w_q=conv_weight[0],
|
||||
w_k=conv_weight[1],
|
||||
w_v=conv_weight[2],
|
||||
cs_q=conv_state[0],
|
||||
cs_k=conv_state[1],
|
||||
cs_v=conv_state[2],
|
||||
g=gate,
|
||||
beta=beta,
|
||||
A_log=torch.randn(num_heads, device="cuda"),
|
||||
dt_bias=torch.randn(num_heads * key_dim, device="cuda"),
|
||||
recurrent_state=state,
|
||||
intermediate_state_indices=scratch,
|
||||
intermediate_conv_q=intermediate_conv.clone(),
|
||||
intermediate_conv_k=intermediate_conv.clone(),
|
||||
intermediate_conv_v=intermediate_conv.clone(),
|
||||
ssm_state_indices=slots,
|
||||
cu_seqlens=torch.arange(
|
||||
0,
|
||||
num_tokens + 1,
|
||||
1 + num_spec,
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
),
|
||||
lower_bound=-5.0,
|
||||
)
|
||||
if not cache_ring:
|
||||
intermediate = torch.zeros(
|
||||
num_requests,
|
||||
1 + num_spec,
|
||||
num_heads,
|
||||
key_dim,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
output = fused_kda_decode_mtp_dspark(
|
||||
intermediate_ssm=intermediate,
|
||||
**kwargs,
|
||||
)
|
||||
return output, intermediate, slots, scratch
|
||||
|
||||
raw_v = torch.zeros(
|
||||
num_slots,
|
||||
num_heads,
|
||||
ring_size,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
raw_k = torch.zeros_like(raw_v)
|
||||
ring_gate = torch.zeros(
|
||||
num_slots,
|
||||
num_heads,
|
||||
ring_size,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
ring_beta = torch.zeros(
|
||||
num_slots,
|
||||
num_heads,
|
||||
ring_size,
|
||||
device="cuda",
|
||||
)
|
||||
output = fused_kda_decode_mtp_dspark(
|
||||
intermediate_ssm=None,
|
||||
replayssm_rawv=raw_v,
|
||||
replayssm_rawk=raw_k,
|
||||
replayssm_g=ring_gate,
|
||||
replayssm_beta=ring_beta,
|
||||
**kwargs,
|
||||
)
|
||||
return (
|
||||
output,
|
||||
state,
|
||||
slots,
|
||||
(raw_v, raw_k, ring_gate, ring_beta),
|
||||
)
|
||||
|
||||
baseline, intermediate, slots, scratch = run(cache_ring=False)
|
||||
ring_output, checkpoint, ring_slots, rings = run(cache_ring=True)
|
||||
self.assertTrue(torch.equal(ring_output, baseline))
|
||||
commit_kda_replayssm_spec(
|
||||
checkpoint,
|
||||
*rings,
|
||||
ring_slots,
|
||||
torch.full(
|
||||
(num_requests,),
|
||||
1 + num_spec,
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
),
|
||||
max_cache_len=ring_size,
|
||||
num_k_heads=num_heads,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
null_block_id=-1,
|
||||
)
|
||||
for request in range(num_requests):
|
||||
expected = intermediate[scratch[request], num_spec]
|
||||
actual = checkpoint[slots[request]]
|
||||
relative_error = (
|
||||
actual - expected
|
||||
).abs().max() / expected.abs().max().clamp_min(1e-6)
|
||||
self.assertLess(relative_error.item(), 2e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,456 @@
|
||||
"""Representative parity coverage for the lightweight Kimi-K3 prerequisites."""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.attention.concat_mla import concat_mla_absorb_q
|
||||
from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
|
||||
fused_sigmoid_gating_delta_rule_update,
|
||||
)
|
||||
from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import (
|
||||
commit_kda_replayssm_spec,
|
||||
)
|
||||
from sglang.kernels.ops.attention.set_mla_kv_concat_q import (
|
||||
can_use_set_mla_kv_concat_q,
|
||||
can_use_set_mla_kv_concat_q_fp8,
|
||||
set_mla_kv_concat_q,
|
||||
set_mla_kv_concat_q_fp8,
|
||||
)
|
||||
from sglang.kernels.ops.attention.utils import concat_mla_absorb_q_general
|
||||
from sglang.kernels.ops.attention.vision_rope import (
|
||||
apply_fused_qk_complex_rope,
|
||||
)
|
||||
from sglang.kernels.ops.elementwise import add3
|
||||
from sglang.kernels.ops.gemm.tiny_gemm import (
|
||||
tiny_k_gemm_bf16,
|
||||
tiny_n_gemm_bf16,
|
||||
)
|
||||
from sglang.kernels.ops.kvcache.set_mla_kv_buffer import set_mla_kv_buffer
|
||||
from sglang.kernels.ops.mm.process.image import (
|
||||
_normalize_and_patchify_torch,
|
||||
normalize_and_patchify,
|
||||
)
|
||||
from sglang.kernels.ops.moe import moe_route_quant_fused
|
||||
from sglang.kernels.ops.moe.moe_route_radix import route_radix
|
||||
from sglang.kernels.ops.moe.moe_topk_sum import moe_topk_sum
|
||||
from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds
|
||||
from sglang.kernels.ops.quantization.per_token_group_quant import (
|
||||
per_token_group_quant,
|
||||
)
|
||||
from sglang.kernels.ops.sampling.top_p_renorm_triton import (
|
||||
top_p_renorm_probs_triton,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
NUM_EXPERTS = 896
|
||||
TOPK = 16
|
||||
NOPE_DIM = 512
|
||||
ROPE_DIM = 64
|
||||
MLA_DIM = NOPE_DIM + ROPE_DIM
|
||||
MLA_PAGES = 256
|
||||
|
||||
|
||||
def _route_oracle(
|
||||
scores, bias, topk, renormalize, routed_scaling_factor, apply_scale, sorted
|
||||
):
|
||||
"""Pure-torch fp32 reference for route_radix.
|
||||
|
||||
Deliberately independent of moe_fused_gate: that entry dispatches back to
|
||||
route_radix whenever scoring is sigmoid with no shared experts, no expert
|
||||
groups and no softcapping (moe_fused_gate.py, the covered() fast path), which
|
||||
is exactly the configuration under test.
|
||||
|
||||
Contract, from route_radix.cuh: bias participates in RANKING only and the
|
||||
emitted weight stays bias-free; NaN is floored so it can never win; ties go to
|
||||
the lower expert id; renormalize divides by the winners' sum (guarded to 1 when
|
||||
that sum is non-positive) and only then is routed scaling applied; sorted=True
|
||||
emits (biased desc, id asc) while sorted=False emits ascending expert id.
|
||||
"""
|
||||
s = torch.sigmoid(scores.float())
|
||||
biased = s + bias.float()
|
||||
biased = torch.where(torch.isnan(biased), torch.full_like(biased, -1e30), biased)
|
||||
# stable + descending: equal biased values keep ascending-id order
|
||||
ranked = torch.argsort(biased, dim=-1, descending=True, stable=True)[:, :topk]
|
||||
w = s.gather(1, ranked)
|
||||
total = w.sum(-1, keepdim=True)
|
||||
norm = torch.where(total > 0, total, torch.ones_like(total))
|
||||
if renormalize:
|
||||
w = w / norm
|
||||
if apply_scale:
|
||||
w = w * routed_scaling_factor
|
||||
if sorted:
|
||||
return w, ranked.to(torch.int32)
|
||||
by_id = ranked.argsort(dim=-1)
|
||||
return w.gather(1, by_id), ranked.gather(1, by_id).to(torch.int32)
|
||||
|
||||
|
||||
def _make_mla_inputs(batch_size, num_heads, seed):
|
||||
generator = torch.Generator(device="cuda").manual_seed(seed)
|
||||
|
||||
def randn(*shape):
|
||||
return (
|
||||
torch.randn(*shape, generator=generator, device="cuda", dtype=torch.float32)
|
||||
.mul(0.1)
|
||||
.to(torch.bfloat16)
|
||||
)
|
||||
|
||||
pool = randn(MLA_PAGES, MLA_DIM)
|
||||
latent = randn(batch_size, MLA_DIM)
|
||||
query = randn(batch_size, num_heads, MLA_DIM)
|
||||
loc = torch.randperm(MLA_PAGES, generator=generator, device="cuda")[:batch_size].to(
|
||||
torch.int64
|
||||
)
|
||||
return (
|
||||
pool,
|
||||
loc,
|
||||
latent[:, :NOPE_DIM],
|
||||
latent[:, NOPE_DIM:],
|
||||
query[..., :NOPE_DIM],
|
||||
query[..., NOPE_DIM:],
|
||||
)
|
||||
|
||||
|
||||
class TestKimiK3PrerequisiteOps(CustomTestCase):
|
||||
def test_mla_scatter_concat_bf16_and_fp8(self):
|
||||
batch_size, num_heads = 64, 8
|
||||
pool, loc, k_nope, k_rope, q_nope, q_rope = _make_mla_inputs(
|
||||
batch_size, num_heads, seed=0
|
||||
)
|
||||
|
||||
if not can_use_set_mla_kv_concat_q(NOPE_DIM * 2, ROPE_DIM * 2):
|
||||
self.skipTest("fused MLA scatter+concat requires SM90+")
|
||||
pool_ref = pool.clone()
|
||||
query = set_mla_kv_concat_q(pool, loc, k_nope, k_rope, q_nope, q_rope)
|
||||
set_mla_kv_buffer(pool_ref, loc, k_nope, k_rope)
|
||||
query_ref = concat_mla_absorb_q(q_nope, q_rope)
|
||||
self.assertTrue(torch.equal(pool, pool_ref))
|
||||
self.assertTrue(torch.equal(query, query_ref))
|
||||
|
||||
if not can_use_set_mla_kv_concat_q_fp8():
|
||||
self.skipTest("fused FP8 MLA scatter+concat requires SM90+")
|
||||
fp8_pool = torch.zeros(
|
||||
MLA_PAGES, MLA_DIM, device="cuda", dtype=torch.float8_e4m3fn
|
||||
)
|
||||
fp8_ref = fp8_pool.clone()
|
||||
fp8_query = set_mla_kv_concat_q_fp8(
|
||||
fp8_pool, loc, k_nope, k_rope, q_nope, q_rope
|
||||
)
|
||||
row = torch.cat([k_nope, k_rope], dim=-1).to(torch.float8_e4m3fn)
|
||||
fp8_ref[loc] = row
|
||||
fp8_query_ref = concat_mla_absorb_q_general(q_nope, q_rope).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(fp8_pool.view(torch.uint8), fp8_ref.view(torch.uint8))
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
fp8_query.view(torch.uint8),
|
||||
fp8_query_ref.view(torch.uint8),
|
||||
)
|
||||
)
|
||||
|
||||
def test_replayssm_ring_fold(self):
|
||||
batch_size, num_steps = 8, 4
|
||||
num_value_heads, num_key_heads = 8, 2
|
||||
key_dim = value_dim = 128
|
||||
ring_size = 16
|
||||
torch.manual_seed(6)
|
||||
|
||||
q = torch.randn(
|
||||
batch_size,
|
||||
num_steps,
|
||||
num_key_heads,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn(
|
||||
batch_size,
|
||||
num_steps,
|
||||
num_value_heads,
|
||||
value_dim,
|
||||
device="cuda",
|
||||
)
|
||||
a = torch.randn(
|
||||
batch_size,
|
||||
num_steps,
|
||||
num_value_heads,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
b = torch.randn(batch_size, num_steps, num_value_heads, device="cuda")
|
||||
a_log = torch.randn(num_value_heads, device="cuda")
|
||||
dt_bias = torch.randn(num_value_heads, key_dim, device="cuda")
|
||||
slots = torch.arange(1, batch_size + 1, device="cuda", dtype=torch.int32)
|
||||
slots[-1] = -1
|
||||
num_slots = batch_size + 1
|
||||
state = torch.randn(
|
||||
num_slots,
|
||||
num_value_heads,
|
||||
value_dim,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
intermediate = torch.zeros(
|
||||
num_slots,
|
||||
num_steps,
|
||||
num_value_heads,
|
||||
value_dim,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
raw_v = torch.zeros(
|
||||
num_slots,
|
||||
num_value_heads,
|
||||
ring_size,
|
||||
value_dim,
|
||||
device="cuda",
|
||||
)
|
||||
raw_k = torch.zeros(
|
||||
num_slots,
|
||||
num_key_heads,
|
||||
ring_size,
|
||||
key_dim,
|
||||
device="cuda",
|
||||
)
|
||||
gate = torch.zeros_like(raw_v)
|
||||
beta = torch.zeros(
|
||||
num_slots,
|
||||
num_value_heads,
|
||||
ring_size,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=a_log,
|
||||
a=a,
|
||||
dt_bias=dt_bias,
|
||||
softplus_beta=1.0,
|
||||
softplus_threshold=20.0,
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
b=b,
|
||||
initial_state_source=state,
|
||||
initial_state_indices=slots,
|
||||
scale=key_dim**-0.5,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
is_kda=True,
|
||||
lower_bound=-5.0,
|
||||
disable_state_update=True,
|
||||
intermediate_states_buffer=intermediate,
|
||||
intermediate_state_indices=slots,
|
||||
cache_steps=num_steps,
|
||||
cache_ring=True,
|
||||
replayssm_rawv=raw_v,
|
||||
replayssm_rawk=raw_k,
|
||||
replayssm_g=gate,
|
||||
replayssm_beta=beta,
|
||||
)
|
||||
checkpoint = state.clone()
|
||||
commit_kda_replayssm_spec(
|
||||
checkpoint,
|
||||
raw_v,
|
||||
raw_k,
|
||||
gate,
|
||||
beta,
|
||||
slots,
|
||||
torch.full((batch_size,), num_steps, device="cuda", dtype=torch.int32),
|
||||
max_cache_len=ring_size,
|
||||
num_k_heads=num_key_heads,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
null_block_id=-1,
|
||||
)
|
||||
for slot in slots[:-1].tolist():
|
||||
expected = intermediate[slot, num_steps - 1]
|
||||
actual = checkpoint[slot]
|
||||
relative_error = (
|
||||
actual - expected
|
||||
).abs().max() / expected.abs().max().clamp_min(1e-6)
|
||||
self.assertLess(relative_error.item(), 1e-3)
|
||||
|
||||
def test_add3_bit_exact(self):
|
||||
torch.manual_seed(0)
|
||||
tensors = [
|
||||
torch.randn(9, 112, device="cuda", dtype=torch.bfloat16) for _ in range(3)
|
||||
]
|
||||
actual = add3.add3(*tensors, prefetch_bc=True)
|
||||
expected = (tensors[0] + tensors[1]) + tensors[2]
|
||||
self.assertTrue(torch.equal(actual, expected))
|
||||
|
||||
def test_moe_auxiliary_kernels(self):
|
||||
x = torch.randn(2, TOPK, 7168, device="cuda", dtype=torch.bfloat16)
|
||||
out = torch.empty(2, 7168, device="cuda", dtype=torch.bfloat16)
|
||||
self.assertIs(moe_topk_sum(x, out), out)
|
||||
self.assertTrue(torch.equal(out, x.float().sum(1).to(torch.bfloat16)))
|
||||
|
||||
def test_moe_route_and_quant(self):
|
||||
torch.manual_seed(1)
|
||||
scores = torch.randn(8, NUM_EXPERTS, device="cuda", dtype=torch.bfloat16)
|
||||
bias = torch.randn(NUM_EXPERTS, device="cuda", dtype=torch.float32)
|
||||
args = (scores, bias, TOPK, True, 2.5, True)
|
||||
weights, ids = route_radix(*args, sorted=True)
|
||||
# Oracle, NOT moe_fused_gate: for this exact configuration (sigmoid, no
|
||||
# shared experts, no expert groups, no softcapping) moe_fused_gate
|
||||
# dispatches straight back to route_radix, so using it as the reference
|
||||
# compares the kernel with itself and cannot see a selection, tie-break,
|
||||
# NaN, renormalize or scaling error.
|
||||
ref_weights, ref_ids = _route_oracle(*args, sorted=True)
|
||||
self.assertTrue(torch.equal(ids, ref_ids))
|
||||
# rtol is not 1e-6: the kernel computes sigmoid with __fdividef/__expf,
|
||||
# whose last bits differ from torch's. The old self-comparison could
|
||||
# afford atol=0; a real oracle cannot.
|
||||
torch.testing.assert_close(weights, ref_weights, rtol=1e-5, atol=1e-6)
|
||||
|
||||
if not moe_route_quant_fused.available():
|
||||
self.skipTest("fused route+quant kernel unavailable")
|
||||
hidden = torch.randn(8, 3584, device="cuda", dtype=torch.bfloat16)
|
||||
ref_weights, ref_ids = route_radix(*args, sorted=False)
|
||||
ref_packed = PackTopkIds.execute(ref_ids, ref_weights)
|
||||
ref_q, ref_scale = per_token_group_quant(
|
||||
hidden, group_size=32, scale_ue8m0=True
|
||||
)
|
||||
actual = moe_route_quant_fused.route_quant_fused(
|
||||
scores,
|
||||
bias,
|
||||
hidden,
|
||||
TOPK,
|
||||
renormalize=True,
|
||||
routed_scaling_factor=2.5,
|
||||
apply_scale=True,
|
||||
)
|
||||
weights, ids, packed, quantized, scale = actual
|
||||
self.assertTrue(torch.equal(ids, ref_ids))
|
||||
self.assertTrue(
|
||||
torch.equal(weights.view(torch.int32), ref_weights.view(torch.int32))
|
||||
)
|
||||
self.assertTrue(torch.equal(packed, ref_packed))
|
||||
self.assertTrue(
|
||||
torch.equal(quantized.view(torch.uint8), ref_q.view(torch.uint8))
|
||||
)
|
||||
torch.testing.assert_close(scale, ref_scale, rtol=0, atol=0)
|
||||
|
||||
def test_route_radix_ties_and_nan(self):
|
||||
"""The cases the self-comparison could not see.
|
||||
|
||||
Exact ties: many experts share one biased value, so the winner set is only
|
||||
determined by the lowest-id rule. NaN: floored, so a NaN expert must never
|
||||
be selected while enough finite ones exist. Both run with renormalize and
|
||||
scaling on and off, since those are applied in a fixed order.
|
||||
"""
|
||||
bias = torch.zeros(NUM_EXPERTS, device="cuda", dtype=torch.float32)
|
||||
|
||||
tied = torch.full((4, NUM_EXPERTS), 0.25, device="cuda", dtype=torch.bfloat16)
|
||||
# a handful of strict winners above the tied plateau, the rest exactly equal
|
||||
tied[:, 300] = 2.0
|
||||
tied[:, 7] = 2.0
|
||||
tied[:, 800] = 1.5
|
||||
|
||||
nan_scores = torch.randn(4, NUM_EXPERTS, device="cuda", dtype=torch.bfloat16)
|
||||
nan_scores[:, 100] = float("nan")
|
||||
nan_scores[:, 500] = float("nan")
|
||||
# make the NaN experts the ones that would otherwise win outright
|
||||
nan_scores[:, 101] = 5.0
|
||||
|
||||
for name, scores in (("ties", tied), ("nan", nan_scores)):
|
||||
for renormalize in (False, True):
|
||||
for apply_scale in (False, True):
|
||||
for sorted_ in (False, True):
|
||||
args = (scores, bias, TOPK, renormalize, 2.5, apply_scale)
|
||||
ids = route_radix(*args, sorted=sorted_)[1]
|
||||
ref_ids = _route_oracle(*args, sorted=sorted_)[1]
|
||||
tag = (
|
||||
f"{name} renorm={renormalize} "
|
||||
f"scale={apply_scale} sorted={sorted_}"
|
||||
)
|
||||
self.assertTrue(torch.equal(ids, ref_ids), msg=tag)
|
||||
if name == "nan":
|
||||
self.assertFalse(
|
||||
bool(((ids == 100) | (ids == 500)).any()),
|
||||
msg=f"{tag}: a NaN expert was selected",
|
||||
)
|
||||
|
||||
def test_tiny_gemm_variants(self):
|
||||
torch.manual_seed(2)
|
||||
x = torch.randn(2, 7168, device="cuda", dtype=torch.bfloat16) / 8
|
||||
weight = torch.randn(144, 7168, device="cuda", dtype=torch.bfloat16) / 8
|
||||
actual = tiny_n_gemm_bf16(x, weight, out_dtype=torch.float32)
|
||||
torch.testing.assert_close(
|
||||
actual.double(), x.double() @ weight.double().t(), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
x = torch.randn(7, 128, device="cuda", dtype=torch.bfloat16) / 4
|
||||
weight = torch.randn(1536, 128, device="cuda", dtype=torch.bfloat16) / 4
|
||||
actual = tiny_k_gemm_bf16(x, weight)
|
||||
torch.testing.assert_close(
|
||||
actual.double(), x.double() @ weight.double().t(), rtol=2e-2, atol=2e-2
|
||||
)
|
||||
|
||||
def test_top_p_renorm(self):
|
||||
torch.manual_seed(3)
|
||||
probs = torch.randn(3, 1024, device="cuda").softmax(-1)
|
||||
top_p = torch.tensor([0.5, 0.8, 0.95], device="cuda")
|
||||
sorted_probs = probs.sort(-1).values
|
||||
cutoff = torch.searchsorted(
|
||||
sorted_probs.cumsum(-1), (1 - top_p).unsqueeze(1)
|
||||
).squeeze(1)
|
||||
cutoff.clamp_(max=probs.shape[1] - 1)
|
||||
pivot = sorted_probs.gather(1, cutoff[:, None])
|
||||
expected = torch.where(probs >= pivot, probs, 0)
|
||||
expected /= expected.sum(-1, keepdim=True)
|
||||
torch.testing.assert_close(
|
||||
top_p_renorm_probs_triton(probs, top_p),
|
||||
expected,
|
||||
rtol=2e-6,
|
||||
atol=1e-8,
|
||||
)
|
||||
|
||||
def test_vision_rope(self):
|
||||
torch.manual_seed(4)
|
||||
qkv = torch.randn(480, 3, 12, 128, device="cuda", dtype=torch.bfloat16)
|
||||
q, k, _ = qkv.unbind(1)
|
||||
angles = torch.randn(480, 64, device="cuda")
|
||||
freqs = torch.polar(torch.ones_like(angles), angles)
|
||||
freqs_expanded = freqs.unsqueeze(-2)
|
||||
|
||||
def reference(x):
|
||||
value = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
|
||||
return torch.view_as_real(value * freqs_expanded).flatten(-2).type_as(x)
|
||||
|
||||
actual_q, actual_k = apply_fused_qk_complex_rope(q, k, freqs)
|
||||
atol = 2 * torch.finfo(torch.bfloat16).eps
|
||||
torch.testing.assert_close(actual_q, reference(q), rtol=0, atol=atol)
|
||||
torch.testing.assert_close(actual_k, reference(k), rtol=0, atol=atol)
|
||||
|
||||
def test_normalize_and_patchify(self):
|
||||
torch.manual_seed(5)
|
||||
image = torch.randn(2, 3, 17, 19, device="cuda")
|
||||
scale = torch.randn(1, 3, 1, 1, device="cuda")
|
||||
bias = torch.randn(1, 3, 1, 1, device="cuda")
|
||||
args = (image, scale, bias, 4, 20, 20)
|
||||
actual = normalize_and_patchify(
|
||||
args[0],
|
||||
args[1],
|
||||
args[2],
|
||||
patch_size=args[3],
|
||||
padded_height=args[4],
|
||||
padded_width=args[5],
|
||||
)
|
||||
expected = _normalize_and_patchify_torch(
|
||||
args[0],
|
||||
args[1],
|
||||
args[2],
|
||||
patch_size=args[3],
|
||||
padded_height=args[4],
|
||||
padded_width=args[5],
|
||||
)
|
||||
torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user