A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 23:15 UTC
Adding SplitK in fused_moe_lora kernel (#27818)
Signed-off-by: Yu Gong <yu3.gong@gmail.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
yugong333 committed
Oct 31, 2025 at 21:55 UTC
29de3cdee4dd7f805931b459398b15c3b5f7057c
1 file changed
+9
-5
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+9
-5
@@ -88,14 +88,17 @@ def _fused_moe_lora_kernel(
88
grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)
89
90
# calculate pid_m,pid_n
91
+ pid_sk = pid % SPLIT_K
92
+ pid_m_n = pid // SPLIT_K
93
num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
94
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
95
+
96
num_pid_in_group = GROUP_SIZE_M * num_pid_n
94
- group_id = pid // num_pid_in_group
97
+ group_id = pid_m_n // num_pid_in_group
98
first_pid_m = group_id * GROUP_SIZE_M
99
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
97
- pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
98
- pid_n = (pid % num_pid_in_group) // group_size_m
100
+ pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
101
+ pid_n = (pid_m_n % num_pid_in_group) // group_size_m
102
103
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_idx)
104
if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
@@ -113,7 +116,7 @@ def _fused_moe_lora_kernel(
116
cur_c_ptr = c_ptr + (slice_id % num_slice_c) * slice_c_size
117
118
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
116
- offs_k = tl.arange(0, BLOCK_SIZE_K)
119
+ offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
120
121
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
122
token_ind = stride_tl * lora_idx + offs_token_id
@@ -131,7 +134,8 @@ def _fused_moe_lora_kernel(
134
cur_b_ptr
135
+ lora_idx * stride_bl
136
+ expert_id * stride_be
134
- + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
137
+ + offs_k[:, None] * stride_bk
138
+ + offs_bn[None, :] * stride_bn
139
)
140
141
# accumulator