[Bugfix][Kimi-K3] Enforce packed rows and op availability in AttnRes dispatch (#50567)
Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
namgyu-youn committed
Aug 4, 2026 at 09:13 UTC
41ba11b8413cefc8b8f0c917d01594cbb26fdfa7
2 files changed
+11
-1
csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu
+7
@@ -914,6 +914,13 @@ void kimi_k3_attn_res(torch::stable::Tensor& prefix,
914
cudaDeviceProp const* properties = get_device_prop();
915
STD_TORCH_CHECK(properties->major == 10,
916
"Kimi K3 AttnRes requires the SM100 family");
917
+ int64_t const hidden_size = prefix.size(1);
918
+ STD_TORCH_CHECK(prefix.stride(0) == hidden_size &&
919
+ delta.stride(0) == hidden_size &&
920
+ output.stride(0) == hidden_size,
921
+ "Kimi K3 AttnRes requires densely packed rows; got strides ",
922
+ prefix.stride(0), ", ", delta.stride(0), ", ",
923
+ output.stride(0), " for hidden_size ", hidden_size);
924
925
using namespace sm100::fwd_prod_v2;
926
// Two-source chunks and two resident CTAs are beneficial once setup is
vllm/models/kimi_k3/nvidia/ops/attn_res.py
+4
-1
@@ -187,7 +187,9 @@ def attn_res(
187
assert qk_weight.stride(-1) == 1
188
assert output_norm_weight is None or output_norm_weight.stride(-1) == 1
189
# The in-tree NVIDIA kernel covers the common fused-add + output-norm path;
190
- # Triton handles block boundaries and final pre-norm output.
190
+ # Triton handles block boundaries and final pre-norm output. The native op
191
+ # is only compiled for SM100 under CUDA >= 13, so a device check alone is
192
+ # not enough to know it exists.
193
if (
194
hidden_size == 7168
195
and delta is not None
@@ -195,6 +197,7 @@ def attn_res(
197
and num_blocks > 0
198
and block_write_idx < 0
199
and current_platform.is_device_capability_family(100)
200
+ and hasattr(torch.ops._C, "kimi_k3_attn_res")
201
):
202
return ops.kimi_k3_attn_res(
203
prefix,