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

[Fix] concat_mla_absorb_q_kernel fails for long inputs (#12453)

Browse Source
This commit is contained in:
bingps
2025-11-02 11:52:06 -08:00
committed by GitHub
parent 66fb9b1307
commit 15ed27d7d4
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/elementwise/concat_mla.cu
+6 -6
View File
@@ -18,11 +18,11 @@ __global__ void concat_mla_k_kernel(
const nv_bfloat16* __restrict__ k_nope, const nv_bfloat16* __restrict__ k_nope,
const nv_bfloat16* __restrict__ k_rope, const nv_bfloat16* __restrict__ k_rope,
const int num_tokens, const int num_tokens,
const int k_stride_0, const int64_t k_stride_0,
const int k_stride_1, const int k_stride_1,
const int k_nope_stride_0, const int64_t k_nope_stride_0,
const int k_nope_stride_1, const int k_nope_stride_1,
const int k_rope_stride_0) { const int64_t k_rope_stride_0) {
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int token_id = flat_warp_id / NUM_HEAD_CHUNKS; const int token_id = flat_warp_id / NUM_HEAD_CHUNKS;
const int head_chunk_id = flat_warp_id % NUM_HEAD_CHUNKS; const int head_chunk_id = flat_warp_id % NUM_HEAD_CHUNKS;
@@ -126,11 +126,11 @@ __global__ void concat_mla_absorb_q_kernel(
nv_bfloat16* out, nv_bfloat16* out,
const int num_items, const int num_items,
const int dim_1, const int dim_1,
const int a_stride_0, const int64_t a_stride_0,
const int a_stride_1, const int a_stride_1,
const int b_stride_0, const int64_t b_stride_0,
const int b_stride_1, const int b_stride_1,
const int out_stride_0, const int64_t out_stride_0,
const int out_stride_1) { const int out_stride_1) {
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane_id = get_lane_id(); const int lane_id = get_lane_id();