iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kimi K3][Kernel] Support DS conv-state layout in fused KDA decode kernel (#53396)

Browse Source 
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
This commit is contained in:
Canlin Guo
2026-08-26 23:54:50 -07:00
committed by GitHub
parent 833fa6233c
commit aa640684cc
5 changed files with 213 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/kernels/benchmark_kimi_k3_kda_decode.py
+25 -8
View File
@@ -81,7 +81,7 @@ def _bench_graph(fn, repeats: int = NUM_KDA_LAYERS) -> float:
class Inputs:
def __init__(self, num_tokens: int, num_heads: int) -> None:
def __init__(self, num_tokens: int, num_heads: int, conv_layout: str) -> None:
torch.manual_seed(0)
device = "cuda"
dim = num_heads * HEAD_DIM
@@ -98,9 +98,15 @@ class Inputs:
for i in range(3)
]
)
self.conv_state = torch.randn(
num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE
)
if conv_layout == "DS":
self.conv_state_t = torch.randn(
num_slots, 3 * dim, CONV_WIDTH - 1, device=device, dtype=DTYPE
)
else:
conv_state = torch.randn(
num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE
)
self.conv_state_t = conv_state.transpose(-1, -2)
self.recurrent_state = torch.randn(
num_slots, num_heads, HEAD_DIM, HEAD_DIM, device=device, dtype=torch.float32
)
@@ -120,7 +126,6 @@ class Inputs:
self.state_indices = torch.arange(
1, num_tokens + 1, device=device, dtype=torch.int32
)
self.conv_state_t = self.conv_state.transpose(-1, -2)
self.out = torch.empty(
1, num_tokens, num_heads, HEAD_DIM, device=device, dtype=DTYPE
)
@@ -237,6 +242,12 @@ def main() -> None:
default=[12],
help="KDA heads per rank (96 total / TP size)",
)
parser.add_argument(
"--layout",
choices=("SD", "DS"),
default="SD",
help="physical conv-state cache layout",
)
args = parser.parse_args()
if not hasattr(torch.ops._C, "fused_kda_decode"):
@@ -245,7 +256,10 @@ def main() -> None:
props = torch.cuda.get_device_properties(0)
bench = _bench_graph if args.graph else _bench
mode = "cuda-graph replay" if args.graph else "eager dispatch"
print(f"device: {props.name} ({props.gcnArchName}) timing: {mode}")
print(
f"device: {props.name} ({props.gcnArchName}) timing: {mode} "
f"layout: {args.layout}"
)
print(
f"{'heads':>6} {'tokens':>7} {'eager-norm':>11} {'triton-norm':>12} "
f"{'fused us':>9} {'speedup':>8} {'state TB/s':>11} {'saved ms/step':>14}"
@@ -253,7 +267,10 @@ def main() -> None:
for num_heads in args.heads:
for num_tokens in args.tokens:
if args.graph and args.layers > 1:
layers = [Inputs(num_tokens, num_heads) for _ in range(args.layers)]
layers = [
Inputs(num_tokens, num_heads, args.layout)
for _ in range(args.layers)
]
eager_ms = _bench_graph_layers(
[functools.partial(triton_chain, i, False) for i in layers]
)
@@ -267,7 +284,7 @@ def main() -> None:
del layers
torch.accelerator.empty_cache()
else:
inp = Inputs(num_tokens, num_heads)
inp = Inputs(num_tokens, num_heads, args.layout)
eager_ms = bench(functools.partial(triton_chain, inp, False))
triton_ms = bench(functools.partial(triton_chain, inp, True))
fused_ms = bench(functools.partial(fused, inp))
csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu
+91 -54
View File
@@ -261,15 +261,16 @@ __device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y,
return {scratch[0], scratch[1]};
}
template <bool kApplyOnorm, bool kUseStaticDecodeLayout = false,
int kFixedHeads = 0, int kFixedValueHeads = 0,
bool kUseHeadGrid = false, bool kAccumulateOnormSumsq = false,
bool kUseActiveQkReduction = false, bool kUseCacheGlobalStore = false,
bool kComputeOutputBeforeStore = false, bool kSkipWarpSync = false,
bool kPreloadOnormParams = false,
bool kPrefetchNextStateChunk = false,
bool kUseActiveOnormReduction = false, bool kUpdateConvState = false,
bool kUseLowerBound = false, bool kApplyBetaSigmoid = true>
template <
bool kApplyOnorm, bool kUseStaticDecodeLayout = false, int kFixedHeads = 0,
int kFixedValueHeads = 0, bool kUseHeadGrid = false,
bool kAccumulateOnormSumsq = false, bool kUseActiveQkReduction = false,
bool kUseCacheGlobalStore = false, bool kComputeOutputBeforeStore = false,
bool kSkipWarpSync = false, bool kPreloadOnormParams = false,
bool kPrefetchNextStateChunk = false, bool kUseActiveOnormReduction = false,
bool kUpdateConvState = false, bool kUseLowerBound = false,
bool kApplyBetaSigmoid = true, int kConvStateChannelStride = 1,
int kConvStateTapStride = 3 * kFixedHeads * kDimK>
__global__
__launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
const __nv_bfloat16* __restrict__ x_q,
@@ -329,7 +330,6 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
}
constexpr int kLocalDim = kFixedHeads * kDimK;
constexpr int kPackedDim = 3 * kLocalDim;
const int hk_off = i_h * kDimK;
const int hv_off = i_hv * kDimV;
constexpr int hv_count = kFixedValueHeads;
@@ -368,8 +368,10 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
__nv_bfloat16 k_shift1 = __float2bfloat16(0.0f);
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const __nv_bfloat16 q_state = cs_q_for_slot[hk + w * kPackedDim];
const __nv_bfloat16 k_state = cs_k_for_slot[hk + w * kPackedDim];
const int cs_idx =
hk * kConvStateChannelStride + w * kConvStateTapStride;
const __nv_bfloat16 q_state = cs_q_for_slot[cs_idx];
const __nv_bfloat16 k_state = cs_k_for_slot[cs_idx];
q_acc += __bfloat162float(q_state) *
conv_weight_load<kLocalDim>(w_q_t, hk, w);
k_acc += __bfloat162float(k_state) *
@@ -389,12 +391,16 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
k_acc += __bfloat162float(k_new) *
conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
cs_q_for_slot[hk] = q_shift0;
cs_q_for_slot[hk + kPackedDim] = q_shift1;
cs_q_for_slot[hk + 2 * kPackedDim] = q_new;
cs_k_for_slot[hk] = k_shift0;
cs_k_for_slot[hk + kPackedDim] = k_shift1;
cs_k_for_slot[hk + 2 * kPackedDim] = k_new;
cs_q_for_slot[hk * kConvStateChannelStride] = q_shift0;
cs_q_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] =
q_shift1;
cs_q_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] =
q_new;
cs_k_for_slot[hk * kConvStateChannelStride] = k_shift0;
cs_k_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] =
k_shift1;
cs_k_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] =
k_new;
s_q[k] = silu_fast(q_acc);
s_k[k] = silu_fast(k_acc);
@@ -418,7 +424,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk);
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const int cs_idx = hk + w * kPackedDim;
const int cs_idx =
hk * kConvStateChannelStride + w * kConvStateTapStride;
q_acc += bf16_load(cs_q_for_slot, cs_idx) *
conv_weight_load<kLocalDim>(w_q_t, hk, w);
k_acc += bf16_load(cs_k_for_slot, cs_idx) *
@@ -453,7 +460,9 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
__nv_bfloat16 v_shift1 = __float2bfloat16(0.0f);
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const __nv_bfloat16 v_state = cs_v_for_slot[hvv + w * kPackedDim];
const int cs_idx =
hvv * kConvStateChannelStride + w * kConvStateTapStride;
const __nv_bfloat16 v_state = cs_v_for_slot[cs_idx];
v_acc += __bfloat162float(v_state) *
conv_weight_load<kLocalDim>(w_v_t, hvv, w);
if (w == 1) {
@@ -465,9 +474,11 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
const __nv_bfloat16 v_new = x_v[xv_idx];
v_acc += __bfloat162float(v_new) *
conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1);
cs_v_for_slot[hvv] = v_shift0;
cs_v_for_slot[hvv + kPackedDim] = v_shift1;
cs_v_for_slot[hvv + 2 * kPackedDim] = v_new;
cs_v_for_slot[hvv * kConvStateChannelStride] = v_shift0;
cs_v_for_slot[hvv * kConvStateChannelStride + kConvStateTapStride] =
v_shift1;
cs_v_for_slot[hvv * kConvStateChannelStride + 2 * kConvStateTapStride] =
v_new;
s_v[v] = silu_fast(v_acc);
if constexpr (kApplyOnorm && kPreloadOnormParams) {
@@ -484,7 +495,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv);
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const int cs_idx = hvv + w * kPackedDim;
const int cs_idx =
hvv * kConvStateChannelStride + w * kConvStateTapStride;
v_acc += bf16_load(cs_v_for_slot, cs_idx) *
conv_weight_load<kLocalDim>(w_v_t, hvv, w);
}
@@ -723,7 +735,8 @@ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
}
template <int kHeads, bool kApplyOnorm, bool kUpdateConvState,
bool kUseLowerBound, bool kApplyBetaSigmoid>
bool kUseLowerBound, bool kApplyBetaSigmoid,
int kConvStateChannelStride, int kConvStateTapStride>
void launch_kda_decode_many_heads_raw(
const void* x_q, const void* x_k, const void* x_v, const void* w_q_t,
const void* w_k_t, const void* w_v_t, const void* bias_q,
@@ -736,7 +749,7 @@ void launch_kda_decode_many_heads_raw(
auto kernel = &kda_decode_fusion_many_heads_kernel<
kApplyOnorm, true, kHeads, kHeads, true, false, false, false, false,
false, true, true, true, kUpdateConvState, kUseLowerBound,
kApplyBetaSigmoid>;
kApplyBetaSigmoid, kConvStateChannelStride, kConvStateTapStride>;
cudaLaunchConfig_t config{};
config.gridDim = dim3(B, kHeads);
config.blockDim = dim3(kThreads);
@@ -776,26 +789,39 @@ void launch_kda_decode_many_heads_selected(
const float* a_log, const void* g, const float* dt_bias, const void* beta,
const void* onorm_g, const float* onorm_weight,
const int* ssm_state_indices, const int* cu_seqlens, float* state,
void* out, int B, int H, int HV, bool update_conv_cache, float lower_bound,
float scale, float onorm_eps, KdaDecodeStrides strides,
cudaStream_t stream) {
#define LAUNCH_KDA_DECODE(NUM_HEADS) \
do { \
if (update_conv_cache) { \
launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, true, \
kUseLowerBound, kApplyBetaSigmoid>( \
x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
scale, onorm_eps, strides, stream); \
} else { \
launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, false, \
kUseLowerBound, kApplyBetaSigmoid>( \
x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
scale, onorm_eps, strides, stream); \
} \
void* out, int B, int H, int HV, bool update_conv_cache,
bool conv_state_is_ds, float lower_bound, float scale, float onorm_eps,
KdaDecodeStrides strides, cudaStream_t stream) {
#define LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, CS_CHANNEL_STRIDE, CS_TAP_STRIDE) \
do { \
if (update_conv_cache) { \
launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, true, \
kUseLowerBound, kApplyBetaSigmoid, \
CS_CHANNEL_STRIDE, CS_TAP_STRIDE>( \
x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
scale, onorm_eps, strides, stream); \
} else { \
launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, false, \
kUseLowerBound, kApplyBetaSigmoid, \
CS_CHANNEL_STRIDE, CS_TAP_STRIDE>( \
x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
scale, onorm_eps, strides, stream); \
} \
} while (false)
// SD packs (tap, channel) per slot while DS packs (channel, tap); both are
// baked into the kernel as compile-time strides.
#define LAUNCH_KDA_DECODE(NUM_HEADS) \
do { \
if (conv_state_is_ds) { \
LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, kConvStateWidth, 1); \
} else { \
LAUNCH_KDA_DECODE_LAYOUT(NUM_HEADS, 1, 3 * (NUM_HEADS) * kDimK); \
} \
} while (false)
switch (H) {
@@ -816,6 +842,7 @@ void launch_kda_decode_many_heads_selected(
}
#undef LAUNCH_KDA_DECODE
#undef LAUNCH_KDA_DECODE_LAYOUT
}
struct KdaDecodeLaunchParams {
@@ -845,6 +872,7 @@ struct KdaDecodeLaunchParams {
int H;
int HV;
bool update_conv_cache;
bool conv_state_is_ds;
float lower_bound;
float scale;
float onorm_eps;
@@ -859,8 +887,8 @@ void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) {
p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k,
p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta,
p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state,
p.out, p.B, p.H, p.HV, p.update_conv_cache, p.lower_bound, p.scale,
p.onorm_eps, p.strides, p.stream);
p.out, p.B, p.H, p.HV, p.update_conv_cache, p.conv_state_is_ds,
p.lower_bound, p.scale, p.onorm_eps, p.strides, p.stream);
}
template <bool kApplyOnorm, bool kUseLowerBound>
@@ -908,6 +936,8 @@ extern "C" void launch_kda_decode_many_heads_cuda(
cudaStream_t stream) {
const KdaDecodeStrides strides{raw_strides[0], raw_strides[1], raw_strides[2],
raw_strides[3], raw_strides[4]};
const bool conv_state_is_ds =
raw_strides[5] == kConvStateWidth && raw_strides[6] == 1;
const KdaDecodeLaunchParams params{x_q,
x_k,
x_v,
@@ -934,6 +964,7 @@ extern "C" void launch_kda_decode_many_heads_cuda(
H,
HV,
update_conv_cache,
conv_state_is_ds,
lower_bound,
scale,
onorm_eps,
@@ -1027,10 +1058,15 @@ void fused_kda_decode(
"out must have shape [1, B, H, 128]");
STD_TORCH_CHECK(x.stride(1) == 1,
"x must be contiguous in its channel dimension");
STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
conv_state.stride(1) == 1 &&
conv_state.stride(2) == 3 * dim,
"conv_state must use the SD cache layout");
// Both conv-state cache layouts are supported: SD stores (state_len, dim)
// per block while DS stores (dim, state_len). The kernel only needs the
// per-plane strides.
STD_TORCH_CHECK(
conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
((conv_state.stride(1) == 1 && conv_state.stride(2) == 3 * dim) ||
(conv_state.stride(1) == kConvWidth - 1 &&
conv_state.stride(2) == 1)),
"conv_state must use the SD or DS cache layout");
STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim &&
state.stride(1) == kHeadDim * kHeadDim &&
state.stride(2) == kHeadDim && state.stride(3) == 1,
@@ -1096,9 +1132,10 @@ void fused_kda_decode(
int64_t const conv_segment_bytes =
dim * conv_state.stride(1) * sizeof(__nv_bfloat16);
int64_t const bias_segment_bytes = dim * sizeof(float);
std::array<int64_t, 5> const strides{
std::array<int64_t, 7> const strides{
x.stride(0), raw_beta.stride(1), output_gate_row_stride,
conv_state.stride(0), state.stride(0),
conv_state.stride(0), state.stride(0), conv_state.stride(1),
conv_state.stride(2),
};
bool const use_lower_bound = lower_bound.has_value();
float const lower_bound_value =
csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu
+58 -28
View File
@@ -147,7 +147,8 @@ __device__ __forceinline__ Sum3 block_reduce_sum3(float x, float y, float z,
scratch[3 * kGroups + 2]};
}
template <bool kApplyOnorm, int kFixedHeads, bool kUpdateConvState,
template <bool kApplyOnorm, int kConvStateChannelStride,
int kConvStateTapStride, int kFixedHeads, bool kUpdateConvState,
bool kUseLowerBound, bool kApplyBetaSigmoid>
__global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
const bf16_t* __restrict__ x_q, const bf16_t* __restrict__ x_k,
@@ -186,7 +187,6 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
}
return;
}
constexpr int kPackedDim = 3 * kLocalDim;
const int hk_off = i_h * kDimK;
const int hv_off = i_hv * kDimV;
float* const state_for_slot = state + slot * strides.state_slot;
@@ -242,8 +242,10 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
bf16_t k_shift1 = __float2bfloat16(0.0f);
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const bf16_t q_state = cs_q_for_slot[hk + w * kPackedDim];
const bf16_t k_state = cs_k_for_slot[hk + w * kPackedDim];
const int cs_idx =
hk * kConvStateChannelStride + w * kConvStateTapStride;
const bf16_t q_state = cs_q_for_slot[cs_idx];
const bf16_t k_state = cs_k_for_slot[cs_idx];
q_acc += __bfloat162float(q_state) *
conv_weight_load<kLocalDim>(w_q_t, hk, w);
k_acc += __bfloat162float(k_state) *
@@ -264,16 +266,21 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
k_acc += __bfloat162float(k_new) *
conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
cs_q_for_slot[hk] = q_shift0;
cs_q_for_slot[hk + kPackedDim] = q_shift1;
cs_q_for_slot[hk + 2 * kPackedDim] = q_new;
cs_k_for_slot[hk] = k_shift0;
cs_k_for_slot[hk + kPackedDim] = k_shift1;
cs_k_for_slot[hk + 2 * kPackedDim] = k_new;
cs_q_for_slot[hk * kConvStateChannelStride] = q_shift0;
cs_q_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] =
q_shift1;
cs_q_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] =
q_new;
cs_k_for_slot[hk * kConvStateChannelStride] = k_shift0;
cs_k_for_slot[hk * kConvStateChannelStride + kConvStateTapStride] =
k_shift1;
cs_k_for_slot[hk * kConvStateChannelStride + 2 * kConvStateTapStride] =
k_new;
} else {
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const int cs_idx = hk + w * kPackedDim;
const int cs_idx =
hk * kConvStateChannelStride + w * kConvStateTapStride;
q_acc += bf16_load(cs_q_for_slot, cs_idx) *
conv_weight_load<kLocalDim>(w_q_t, hk, w);
k_acc += bf16_load(cs_k_for_slot, cs_idx) *
@@ -307,7 +314,9 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
bf16_t v_shift1 = __float2bfloat16(0.0f);
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const bf16_t v_state = cs_v_for_slot[hvv + w * kPackedDim];
const int cs_idx =
hvv * kConvStateChannelStride + w * kConvStateTapStride;
const bf16_t v_state = cs_v_for_slot[cs_idx];
v_acc += __bfloat162float(v_state) *
conv_weight_load<kLocalDim>(w_v_t, hvv, w);
if (w == 1) {
@@ -320,13 +329,16 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
const bf16_t v_new = x_v[xv_idx];
v_acc += __bfloat162float(v_new) *
conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1);
cs_v_for_slot[hvv] = v_shift0;
cs_v_for_slot[hvv + kPackedDim] = v_shift1;
cs_v_for_slot[hvv + 2 * kPackedDim] = v_new;
cs_v_for_slot[hvv * kConvStateChannelStride] = v_shift0;
cs_v_for_slot[hvv * kConvStateChannelStride + kConvStateTapStride] =
v_shift1;
cs_v_for_slot[hvv * kConvStateChannelStride + 2 * kConvStateTapStride] =
v_new;
} else {
#pragma unroll
for (int w = 0; w < kConvStateWidth; ++w) {
const int cs_idx = hvv + w * kPackedDim;
const int cs_idx =
hvv * kConvStateChannelStride + w * kConvStateTapStride;
v_acc += bf16_load(cs_v_for_slot, cs_idx) *
conv_weight_load<kLocalDim>(w_v_t, hvv, w);
}
@@ -461,7 +473,8 @@ __global__ __launch_bounds__(kThreads) void kda_decode_fusion_kernel(
}
template <int kHeads, bool kApplyOnorm, bool kUpdateConvState,
bool kUseLowerBound, bool kApplyBetaSigmoid>
bool kUseLowerBound, bool kApplyBetaSigmoid,
int kConvStateChannelStride, int kConvStateTapStride>
void launch_kda_decode_raw(
const void* x_q, const void* x_k, const void* x_v, const void* w_q_t,
const void* w_k_t, const void* w_v_t, const void* bias_q,
@@ -471,7 +484,8 @@ void launch_kda_decode_raw(
const int* ssm_state_indices, float* state, void* out, int B,
float lower_bound, float scale, float onorm_eps, KdaDecodeStrides strides,
hipStream_t stream) {
kda_decode_fusion_kernel<kApplyOnorm, kHeads, kUpdateConvState,
kda_decode_fusion_kernel<kApplyOnorm, kConvStateChannelStride,
kConvStateTapStride, kHeads, kUpdateConvState,
kUseLowerBound, kApplyBetaSigmoid>
<<<dim3(kHeads, B), dim3(kThreads), 0, stream>>>(
reinterpret_cast<const bf16_t*>(x_q),
@@ -517,6 +531,7 @@ struct KdaDecodeLaunchParams {
int B;
int H;
bool update_conv_cache;
bool conv_state_is_ds;
float lower_bound;
float scale;
float onorm_eps;
@@ -527,19 +542,29 @@ struct KdaDecodeLaunchParams {
template <int kHeads, bool kApplyOnorm, bool kUseLowerBound,
bool kApplyBetaSigmoid>
void dispatch_kda_decode_conv(const KdaDecodeLaunchParams& p) {
#define LAUNCH_KDA_DECODE(UPDATE_CONV) \
launch_kda_decode_raw<kHeads, kApplyOnorm, UPDATE_CONV, kUseLowerBound, \
kApplyBetaSigmoid>( \
p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k, \
p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, \
p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.state, p.out, p.B, \
#define LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, CS_CHANNEL_STRIDE, \
CS_TAP_STRIDE) \
launch_kda_decode_raw<kHeads, kApplyOnorm, UPDATE_CONV, kUseLowerBound, \
kApplyBetaSigmoid, CS_CHANNEL_STRIDE, CS_TAP_STRIDE>( \
p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k, \
p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, \
p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.state, p.out, p.B, \
p.lower_bound, p.scale, p.onorm_eps, p.strides, p.stream)
// SD packs (tap, channel) per slot while DS packs (channel, tap); both are
// baked into the kernel as compile-time strides.
#define LAUNCH_KDA_DECODE(UPDATE_CONV) \
if (p.conv_state_is_ds) { \
LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, kConvStateWidth, 1); \
} else { \
LAUNCH_KDA_DECODE_LAYOUT(UPDATE_CONV, 1, 3 * kHeads * kDimK); \
}
if (p.update_conv_cache) {
LAUNCH_KDA_DECODE(true);
} else {
LAUNCH_KDA_DECODE(false);
}
#undef LAUNCH_KDA_DECODE
#undef LAUNCH_KDA_DECODE_LAYOUT
}
template <bool kApplyOnorm, bool kUseLowerBound, bool kApplyBetaSigmoid>
@@ -686,10 +711,12 @@ void fused_kda_decode(
"out must have shape [1, B, H, 128]");
STD_TORCH_CHECK(x.stride(1) == 1,
"x must be contiguous in its channel dimension");
STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
conv_state.stride(1) == 1 &&
conv_state.stride(2) == 3 * dim,
"conv_state must use the SD cache layout");
STD_TORCH_CHECK(
conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
((conv_state.stride(1) == 1 && conv_state.stride(2) == 3 * dim) ||
(conv_state.stride(1) == kConvWidth - 1 &&
conv_state.stride(2) == 1)),
"conv_state must use the SD or DS cache layout");
STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim &&
state.stride(1) == kHeadDim * kHeadDim &&
state.stride(2) == kHeadDim && state.stride(3) == 1,
@@ -761,6 +788,8 @@ void fused_kda_decode(
torch::stable::accelerator::DeviceGuard const device_guard(
x.get_device_index());
bool const conv_state_is_ds =
conv_state.stride(1) == kConvStateWidth && conv_state.stride(2) == 1;
KdaDecodeLaunchParams const params{
x_ptr,
x_ptr + segment_bytes,
@@ -786,6 +815,7 @@ void fused_kda_decode(
batch_size,
static_cast<int>(num_heads),
true,
conv_state_is_ds,
lower_bound_value,
0.08838834764831845f,
static_cast<float>(norm_eps),
tests/models/kimi_k3/test_kda.py
+37 -15
View File
@@ -895,13 +895,18 @@ def test_kda_recoverssm_verify_and_group_commit(
@pytest.mark.parametrize(
("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"),
("num_heads", "num_seqs", "lower_bound", "fuse_output_norm", "conv_layout"),
[
(12, 1, -5.0, True),
(12, 4, None, False),
(24, 4, None, False),
(48, 1, -5.0, True),
(96, 1, -5.0, True),
(12, 1, -5.0, True, "SD"),
(12, 4, None, False, "SD"),
(24, 4, None, False, "SD"),
(48, 1, -5.0, True, "SD"),
(96, 1, -5.0, True, "SD"),
(12, 1, -5.0, True, "DS"),
(12, 4, None, False, "DS"),
(24, 4, None, False, "DS"),
(48, 1, -5.0, True, "DS"),
(96, 1, -5.0, True, "DS"),
],
)
@torch.inference_mode()
@@ -910,6 +915,7 @@ def test_fused_kda_decode_correctness(
num_seqs: int,
lower_bound: float | None,
fuse_output_norm: bool,
conv_layout: str,
):
D, W = 128, 4
if not is_fused_kda_decode_supported(
@@ -921,7 +927,7 @@ def test_fused_kda_decode_correctness(
conv_state_dtype=torch.bfloat16,
):
pytest.skip("Fused KDA decode is not supported on this platform")
torch.manual_seed(967 + num_heads + num_seqs)
torch.manual_seed(967 + num_heads + num_seqs + (conv_layout == "DS"))
dim = num_heads * D
slots = num_seqs + 2
packed_x_storage = torch.randn(
@@ -929,13 +935,25 @@ def test_fused_kda_decode_correctness(
)
packed_x = packed_x_storage[:, : 3 * dim]
weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE)
conv_seed = 0.1 * torch.randn(
slots,
W - 1,
3 * dim,
dtype=torch.bfloat16,
device=DEVICE,
).transpose(1, 2)
if conv_layout == "DS":
# DS cache layout: per slot the taps are innermost
# (stride (W-1, 1)), matching VLLM_SSM_CONV_STATE_LAYOUT=DS.
conv_seed = 0.1 * torch.randn(
slots,
3 * dim,
W - 1,
dtype=torch.bfloat16,
device=DEVICE,
)
else:
# SD cache layout: per slot the channels are innermost.
conv_seed = 0.1 * torch.randn(
slots,
W - 1,
3 * dim,
dtype=torch.bfloat16,
device=DEVICE,
).transpose(1, 2)
raw_g = torch.randn(
1,
num_seqs,
@@ -1017,7 +1035,11 @@ def test_fused_kda_decode_correctness(
conv_actual = torch.as_strided(
cache_storage.view(torch.bfloat16),
size=(slots, 3 * dim, W - 1),
stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim),
stride=(
page_bytes // torch.bfloat16.itemsize,
(W - 1) if conv_layout == "DS" else 1,
1 if conv_layout == "DS" else 3 * dim,
),
)
state_actual = torch.as_strided(
cache_storage.view(torch.float32),
vllm/models/kimi_k3/nvidia/kda.py
+2 -1
View File
@@ -153,6 +153,8 @@ def is_fused_kda_decode_supported(
input_dtype: torch.dtype,
conv_state_dtype: torch.dtype,
) -> bool:
# The fused kernel handles both conv-state cache layouts (SD and DS); the
# inner strides are selected from the tensor at launch time.
if (
num_heads not in (12, 24, 48, 96)
or head_dim != 128
@@ -160,7 +162,6 @@ def is_fused_kda_decode_supported(
or num_spec != 0
or input_dtype != torch.bfloat16
or conv_state_dtype != torch.bfloat16
or is_conv_state_dim_first()
or not hasattr(torch.ops._C, "fused_kda_decode")
):
return False