[K3 Perf] Enable DSV3 GEMM for inner-contiguous and row-strided tensors, 12%~81% kernel performance improvement (#54565)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 2, 2026 at 15:22 UTC
a56654d6de060495ff2db3b1d9ff0b187084d1a9
3 files changed
+48
-35
csrc/libtorch_stable/dsv3_fused_a_gemm.cu
+35
-27
@@ -326,7 +326,7 @@ struct GmemLoaderB {
326
#endif
327
}
328
329
- __device__ void issue_mainloop() {
329
+ __device__ void issue_mainloop(int row_stride) {
330
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
331
cudaGridDependencySynchronize();
332
#pragma unroll 1
@@ -350,7 +350,7 @@ struct GmemLoaderB {
350
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
351
int n_idx = linear_idx / tile_k;
352
int k_idx = linear_idx % tile_k;
353
- int gmem_offset = n_idx * gemm_k + k_project(k_idx);
353
+ int gmem_offset = n_idx * row_stride + k_project(k_idx);
354
bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
355
ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
356
}
@@ -585,7 +585,8 @@ struct MmaComputer {
585
template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n,
586
int tile_k, int stage_cnt>
587
__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
588
- bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
588
+ bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n,
589
+ int input_row_stride) {
590
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
591
constexpr int load_thread_cnt = 128;
592
constexpr int compute_thread_cnt = 128;
@@ -621,7 +622,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
622
int cta_m_idx = tile_m * blockIdx.x;
623
int cta_n_idx = tile_n * blockIdx.y;
624
bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
624
- bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
625
+ bf16_t const* gmem_b_local = mat_b + cta_n_idx * input_row_stride;
626
bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
627
628
int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
@@ -645,7 +646,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
646
GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(
647
gmem_b_local, smem_b, smem_barrier, gemm_n);
648
b_loader.prepare();
648
- b_loader.issue_mainloop();
649
+ b_loader.issue_mainloop(input_row_stride);
650
} else {
651
MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
652
gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
@@ -660,7 +661,8 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
661
template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256,
662
int kTileM = 16>
663
void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
663
- cudaStream_t const stream, bool enable_pdl) {
664
+ int input_row_stride, cudaStream_t const stream,
665
+ bool enable_pdl) {
666
constexpr int gemm_m = kHdOut;
667
int const gemm_n = num_tokens;
668
constexpr int gemm_k = kHdIn;
@@ -703,19 +705,19 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
705
cudaLaunchKernelEx(&config,
706
fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
707
tile_n, tile_k, stage_cnt>,
706
- output, mat_a, mat_b, gemm_n);
708
+ output, mat_a, mat_b, gemm_n, input_row_stride);
709
}
710
711
template <typename T, int kHdIn, int kHdOut, int kTileK = 256, int kTileM = 16>
712
void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b,
711
- int num_tokens, cudaStream_t const stream,
712
- bool enable_pdl) {
713
+ int num_tokens, int input_row_stride,
714
+ cudaStream_t const stream, bool enable_pdl) {
715
if (num_tokens <= 8) {
716
invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK, kTileM>(
715
- output, mat_a, mat_b, num_tokens, stream, enable_pdl);
717
+ output, mat_a, mat_b, num_tokens, input_row_stride, stream, enable_pdl);
718
} else {
719
invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK, kTileM>(
718
- output, mat_a, mat_b, num_tokens, stream, enable_pdl);
720
+ output, mat_a, mat_b, num_tokens, input_row_stride, stream, enable_pdl);
721
}
722
}
723
@@ -740,11 +742,9 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
742
mat_a.get_device_index() == output.get_device_index(),
743
"mat_a, mat_b, and output must be on the same device");
744
743
- // The kernels index global memory with raw pointers and packed strides, so
744
- // reject any padded or transposed view rather than reading out of bounds.
745
STD_TORCH_CHECK(
746
- mat_a.stride(0) == hd_in && mat_a.stride(1) == 1,
747
- "mat_a must be a packed row-major [num_tokens, hd_in] tensor");
746
+ mat_a.stride(1) == 1 && (num_tokens == 1 || mat_a.stride(0) % 8 == 0),
747
+ "mat_a must have unit inner stride and 16-byte aligned rows");
748
STD_TORCH_CHECK(
749
output.stride(0) == hd_out && output.stride(1) == 1,
750
"output must be a packed row-major [num_tokens, hd_out] tensor");
@@ -770,12 +770,14 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
770
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
771
auto const* mat_b_ptr =
772
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
773
-
774
-#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \
775
- if (hd_in == HD_IN && hd_out == HD_OUT) { \
776
- invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \
777
- output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \
778
- return; \
773
+ int const input_row_stride = mat_a.stride(0);
774
+
775
+#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \
776
+ if (hd_in == HD_IN && hd_out == HD_OUT) { \
777
+ invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \
778
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, input_row_stride, \
779
+ stream, enable_pdl); \
780
+ return; \
781
}
782
783
// Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners
@@ -795,7 +797,8 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
797
798
if (hd_in == 6144 && hd_out == 2624) {
799
invokeFusedAGemmForTokens<__nv_bfloat16, 6144, 2624, 256, 32>(
798
- output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
800
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, input_row_stride, stream,
801
+ enable_pdl);
802
return;
803
}
804
DISPATCH_DSV3_SHAPE(2048, 2048)
@@ -821,30 +824,35 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
824
825
if (hd_in == 128 && hd_out == 1536) {
826
invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>(
824
- output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
827
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, input_row_stride, stream,
828
+ enable_pdl);
829
return;
830
}
831
if (hd_in == 128 && hd_out == 3072) {
832
invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>(
829
- output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
833
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, input_row_stride, stream,
834
+ enable_pdl);
835
return;
836
}
837
// TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple
838
// of 256, so both need the 128 tile_k.
839
if (hd_in == 128 && hd_out == 768) {
840
invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>(
836
- output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
841
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, input_row_stride, stream,
842
+ enable_pdl);
843
return;
844
}
845
if (hd_in == 384 && hd_out == 7168) {
846
invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>(
841
- output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
847
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, input_row_stride, stream,
848
+ enable_pdl);
849
return;
850
}
851
#ifdef VLLM_K3_BENCH_SHAPES
852
if (hd_in == 4224 && hd_out == 7168) {
853
invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>(
847
- output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
854
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, input_row_stride, stream,
855
+ enable_pdl);
856
return;
857
}
858
#endif
tests/kernels/test_bf16_skinny_gemm.py
+6
-6
@@ -996,20 +996,20 @@ def test_glm_dsv3_selected_shapes(
996
assert cosine > 0.999
997
998
999
-def test_nonpacked_single_token_dsv3_falls_back() -> None:
999
+@pytest.mark.parametrize("num_tokens", [1, 4])
1000
+def test_kda_f_b_nonpacked_dsv3_dispatches(num_tokens: int) -> None:
1001
_require_sm103_and_dsv3()
1002
n, k = 1536, 128
1002
- storage = torch.randn(1, k + 16, dtype=torch.bfloat16, device="cuda")
1003
- x = storage[:, :k]
1003
+ storage = torch.randn(num_tokens, 6288, dtype=torch.bfloat16, device="cuda")
1004
+ x = storage[:, 6144 : 6144 + k]
1005
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
1006
spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
1007
method = k3_gemm.KimiK3LowLatencyLinearMethod(
1008
k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec)
1009
)
1010
1010
- assert x.is_contiguous()
1011
- assert x.stride() == (k + 16, 1)
1012
- assert not k3_gemm._runtime_ok(x, weight) # strict guard rejects the view
1011
+ assert x.stride() == (6288, 1)
1012
+ assert k3_gemm._runtime_ok(x, weight)
1013
output = method.apply(SimpleNamespace(weight=weight), x)
1014
1015
reference = torch.nn.functional.linear(x, weight)
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+7
-2
@@ -663,7 +663,9 @@ def _is_packed_row_major(tensor: torch.Tensor) -> bool:
663
664
def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
665
return (
666
- _is_packed_row_major(x)
666
+ x.dim() == 2
667
+ and x.stride(1) == 1
668
+ and (x.shape[0] == 1 or x.stride(0) % 8 == 0)
669
and _is_packed_row_major(weight)
670
and x.dtype == torch.bfloat16
671
and weight.dtype == torch.bfloat16
@@ -676,7 +678,8 @@ def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
678
679
def _residual_ok(x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor) -> bool:
680
return (
679
- residual.dim() == 2
681
+ _is_packed_row_major(x)
682
+ and residual.dim() == 2
683
and residual.dtype == torch.bfloat16
684
and residual.is_cuda
685
and residual.device == x.device
@@ -693,6 +696,8 @@ def _run_plan(
696
return None
697
backend, config = entry
698
if backend == "cute":
699
+ if x.shape[0] != 1 and not _is_packed_row_major(x):
700
+ return None
701
if not shape_dynamic_skinny_gemm.is_available():
702
return None
703
return shape_dynamic_skinny_gemm(x, weight, config, None)