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 3, 11:45 UTC
[Bugfix] moe lora align kernel grid (#40131)
Signed-off-by: TheDuyIT <nduy250299@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Signed-off-by: dtnguyen <dtnguyen@nvidia.com> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Nguyễn Thế Duy committed
May 18, 2026 at 14:17 UTC
e3aeee5ff8bf7e89fea231d2a965701248eb43c0
2 files changed
+263
-9
csrc/moe/moe_align_sum_kernels.cu
+46
-9
@@ -390,7 +390,13 @@ __global__ void moe_lora_align_block_size_kernel(
390
int32_t* __restrict__ token_mask, bool has_expert_map) {
391
int lora_idx = blockIdx.x / 2;
392
int lora_id = lora_ids[lora_idx];
393
- if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
393
+ // Output buffers are indexed by lora_id (in [0, max_loras)). The grid
394
+ // iterates one extra slot to accommodate the "-1" entry that
395
+ // active_lora_ids may hold in position 0 for mixed base + LoRA batches;
396
+ // guard against any other unexpected lora_id >= max_loras to avoid
397
+ // out-of-bounds writes. This mirrors the `lora_id >= max_loras` guard in
398
+ // the Triton _fused_moe_lora_kernel.
399
+ if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
400
return;
401
}
402
@@ -420,10 +426,21 @@ __global__ void lora_count_and_sort_expert_tokens_kernel(
426
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
427
int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
428
int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask,
423
- int32_t* lora_ids, bool has_expert_map) {
429
+ int32_t max_loras, int32_t* lora_ids, int32_t* adapter_enabled,
430
+ bool has_expert_map) {
431
int lora_idx = blockIdx.x;
432
int lora_id = lora_ids[lora_idx];
426
- if (lora_id == -1) {
433
+ // Same guard rationale as moe_lora_align_block_size_kernel. Additionally
434
+ // skip disabled adapter slots: moe_lora_align_block_size_kernel early-returns
435
+ // for them and leaves token_mask[lora_id, :] uninitialized (token_mask is
436
+ // allocated with torch::empty), so running the sort loop here would traverse
437
+ // garbage mask bits and pollute this slot's rows of sorted_token_ids and
438
+ // cumsum_buffer. Downstream consumers already skip disabled slots, so the
439
+ // pollution is dormant today, but the check keeps behavior symmetric with
440
+ // the other two align kernels and avoids O(numel) wasted work per disabled
441
+ // slot. Short-circuit evaluation ensures adapter_enabled is only indexed
442
+ // after lora_id is confirmed to be in [0, max_loras).
443
+ if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
444
return;
445
}
446
@@ -446,7 +463,8 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
463
int32_t* token_mask, bool has_expert_map) {
464
int lora_idx = blockIdx.x;
465
int lora_id = lora_ids[lora_idx];
449
- if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
466
+ // Same guard rationale as moe_lora_align_block_size_kernel.
467
+ if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
468
return;
469
}
470
@@ -698,7 +716,15 @@ void moe_lora_align_block_size(
716
scalar_t, fill_threads>;
717
AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
718
(void*)kernel, shared_mem));
701
- kernel<<<max_loras, blockDim, shared_mem, stream>>>(
719
+ // Grid size is (max_loras + 1) because active_lora_ids has length
720
+ // max_loras + 1: sorted-unique values of token_lora_mapping, which
721
+ // can include -1 (base-model tokens) in addition to up to max_loras
722
+ // real LoRA slots. Using max_loras would drop the real LoRA slot
723
+ // when -1 is present at position 0 and leave output buffers
724
+ // uninitialized, causing illegal memory accesses in downstream
725
+ // MoE-LoRA kernels. This mirrors the fix made for the Triton
726
+ // _fused_moe_lora_kernel grid in vllm-project/vllm#32277.
727
+ kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>(
728
topk_ids.data_ptr<scalar_t>(),
729
token_lora_mapping.data_ptr<int32_t>(), block_size,
730
expert_map.data_ptr<int32_t>(), num_experts, max_loras,
@@ -722,10 +748,17 @@ void moe_lora_align_block_size(
748
auto align_kernel =
749
vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;
750
725
- // launch two threadblocks for each lora
751
+ // Launch two threadblocks per LoRA slot, across max_loras + 1 slots
752
+ // to cover the extra "-1" (base-model tokens) entry that
753
+ // active_lora_ids may contain in addition to up to max_loras real
754
+ // LoRA slots. Using max_loras would drop the real LoRA slot when -1
755
+ // occupies position 0 and leave the output buffers uninitialized,
756
+ // causing illegal memory accesses downstream. Mirrors the grid fix
757
+ // applied to _fused_moe_lora_kernel in vllm-project/vllm#32277.
758
// blockIdx.x % 2 == 0: counting experts and aligning
759
// blockIdx.x % 2 == 1: filling sorted_token_ids
728
- align_kernel<<<max_loras * 2, blockDim, shared_mem_size, stream>>>(
760
+ align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size,
761
+ stream>>>(
762
topk_ids.data_ptr<scalar_t>(),
763
token_lora_mapping.data_ptr<int32_t>(), block_size,
764
expert_map.data_ptr<int32_t>(), num_experts, max_loras,
@@ -744,7 +777,10 @@ void moe_lora_align_block_size(
777
const int max_blocks = 65535;
778
const int actual_blocks = std::min(num_blocks, max_blocks);
779
747
- dim3 gridDims(max_loras, actual_blocks);
780
+ // Same rationale as align_kernel above: iterate over max_loras + 1
781
+ // slots so the sort kernel processes the real LoRA slot even when
782
+ // active_lora_ids has -1 at position 0.
783
+ dim3 gridDims(max_loras + 1, actual_blocks);
784
auto sort_kernel =
785
vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;
786
@@ -753,7 +789,8 @@ void moe_lora_align_block_size(
789
sorted_token_ids.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
790
expert_map.data_ptr<int32_t>(), topk_ids.numel(), num_experts,
791
max_num_tokens_padded, topk_num, token_mask.data_ptr<int32_t>(),
756
- lora_ids.data_ptr<int32_t>(), has_expert_map);
792
+ max_loras, lora_ids.data_ptr<int32_t>(),
793
+ adapter_enabled.data_ptr<int32_t>(), has_expert_map);
794
}
795
});
796
}
\ No newline at end of file
tests/lora/test_moe_lora_align_sum.py
+217
@@ -104,5 +104,222 @@ def test_moe_lora_align_block_size(
104
assert torch.all(topk_ids.view(-1)[indices] == expert_id)
105
106
107
+# Sentinel values for the regression tests below. Distinctive out-of-domain
108
+# ints so that "kernel never wrote this slot" is directly observable: the
109
+# kernel only ever writes a real expert id in [0, num_experts) or -1
110
+# (expert_ids), a token index or the `numel` padding value (sorted_token_ids),
111
+# and a block-aligned cumsum count (num_tokens_post_pad).
112
+SENTINEL_EXPERT = -2
113
+SENTINEL_TOKEN = -7
114
+SENTINEL_NPAD = -13
115
+
116
+
117
+def _build_and_run_align(
118
+ *,
119
+ num_lora_tokens,
120
+ num_base_tokens,
121
+ max_loras,
122
+ num_experts=64,
123
+ topk_num=6,
124
+ block_size=16,
125
+ lora_ids_override=None,
126
+ disabled_slots=(),
127
+ seed=1,
128
+):
129
+ """Build inputs the way ``LoRAKernelMeta.prepare_tensors`` does, run
130
+ ``moe_lora_align_block_size``, and return a dict of result tensors plus
131
+ derived sizes. Output buffers are pre-filled with ``SENTINEL_*`` so
132
+ callers can assert which slots the kernel did / did not touch.
133
+
134
+ Tokens are assigned to LoRA slot 0 (first ``num_lora_tokens``) then -1
135
+ (remaining ``num_base_tokens``), matching the "mixed base + 1 LoRA"
136
+ shape used to repro vllm-project/vllm#32235.
137
+
138
+ ``lora_ids_override``: optional 1-D int tensor of length ``max_loras+1``
139
+ used verbatim. Default mirrors ``prepare_tensors`` (sorted-unique into
140
+ the head, -1 tail).
141
+ ``disabled_slots``: iterable of slot indices to clear in ``adapter_enabled``.
142
+ """
143
+ random.seed(seed)
144
+ num_tokens = num_lora_tokens + num_base_tokens
145
+ assert num_tokens > 0, "test requires at least one token"
146
+
147
+ topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
148
+ token_lora_mapping = torch.empty((num_tokens,), dtype=torch.int32)
149
+ for i in range(num_tokens):
150
+ pool = list(range(num_experts))
151
+ random.shuffle(pool)
152
+ for j in range(topk_num):
153
+ topk_ids[i, j] = pool[j]
154
+ token_lora_mapping[i] = 0 if i < num_lora_tokens else -1
155
+ topk_ids = topk_ids.to(DEVICE_TYPE)
156
+ token_lora_mapping = token_lora_mapping.to(DEVICE_TYPE)
157
+
158
+ max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
159
+ max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
160
+ if topk_ids.numel() < num_experts:
161
+ max_num_tokens_padded = topk_ids.numel() * block_size
162
+ max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
163
+
164
+ if lora_ids_override is None:
165
+ lora_ids = torch.full(
166
+ (max_loras + 1,), -1, dtype=torch.int32, device=DEVICE_TYPE
167
+ )
168
+ unique_ids = torch.unique(token_lora_mapping, sorted=True)
169
+ lora_ids[: unique_ids.numel()] = unique_ids.to(torch.int32)
170
+ else:
171
+ assert lora_ids_override.numel() == max_loras + 1
172
+ lora_ids = lora_ids_override.to(dtype=torch.int32, device=DEVICE_TYPE)
173
+
174
+ adapter_enabled = torch.ones(
175
+ (max_loras + 1,), dtype=torch.int32, device=DEVICE_TYPE
176
+ )
177
+ for slot in disabled_slots:
178
+ adapter_enabled[slot] = 0
179
+
180
+ sorted_token_ids = torch.full(
181
+ (max_loras * max_num_tokens_padded,),
182
+ SENTINEL_TOKEN,
183
+ dtype=torch.int32,
184
+ device=DEVICE_TYPE,
185
+ )
186
+ expert_ids = torch.full(
187
+ (max_loras * max_num_m_blocks,),
188
+ SENTINEL_EXPERT,
189
+ dtype=torch.int32,
190
+ device=DEVICE_TYPE,
191
+ )
192
+ num_tokens_post_pad = torch.full(
193
+ (max_loras,), SENTINEL_NPAD, dtype=torch.int32, device=DEVICE_TYPE
194
+ )
195
+
196
+ ops.moe_lora_align_block_size(
197
+ topk_ids,
198
+ token_lora_mapping,
199
+ num_experts,
200
+ block_size,
201
+ max_loras,
202
+ max_num_tokens_padded,
203
+ max_num_m_blocks,
204
+ sorted_token_ids,
205
+ expert_ids,
206
+ num_tokens_post_pad,
207
+ adapter_enabled,
208
+ lora_ids,
209
+ )
210
+
211
+ return {
212
+ "lora_ids": lora_ids,
213
+ "sorted_token_ids": sorted_token_ids,
214
+ "expert_ids": expert_ids,
215
+ "num_tokens_post_pad": num_tokens_post_pad,
216
+ "max_num_tokens_padded": max_num_tokens_padded,
217
+ "block_size": block_size,
218
+ "max_loras": max_loras,
219
+ }
220
+
221
+
222
+@pytest.mark.parametrize(
223
+ "max_loras",
224
+ [
225
+ 1,
226
+ 2,
227
+ ],
228
+)
229
+def test_moe_lora_align_block_size_mixed_base_and_lora(max_loras):
230
+ """Regression test for issue #32235: real LoRA slot must not be skipped
231
+ when ``active_lora_ids`` has -1 at position 0."""
232
+ out = _build_and_run_align(
233
+ num_lora_tokens=8, num_base_tokens=8, max_loras=max_loras
234
+ )
235
+
236
+ # Sanity check on the layout being tested.
237
+ assert out["lora_ids"][0].item() == -1, (
238
+ "prepare_tensors layout mismatch: -1 expected at position 0 for mixed batch"
239
+ )
240
+
241
+ real_slot = 0
242
+ post_pad = out["num_tokens_post_pad"][real_slot].item()
243
+ assert post_pad != SENTINEL_NPAD, (
244
+ f"num_tokens_post_pad[{real_slot}] was never written by the kernel; "
245
+ "the align kernel skipped the real LoRA slot."
246
+ )
247
+ assert (
248
+ 0 < post_pad <= out["max_num_tokens_padded"]
249
+ and post_pad % out["block_size"] == 0
250
+ ), f"num_tokens_post_pad[{real_slot}]={post_pad} is not a valid block-aligned count"
251
+
252
+ expert_row = out["expert_ids"].view(max_loras, -1)[real_slot]
253
+ assert (expert_row != SENTINEL_EXPERT).all(), (
254
+ f"expert_ids row for slot {real_slot} has unwritten sentinel entries; "
255
+ "the align kernel skipped the real LoRA slot."
256
+ )
257
+
258
+ sorted_row = out["sorted_token_ids"].view(max_loras, -1)[real_slot]
259
+ assert (sorted_row != SENTINEL_TOKEN).all(), (
260
+ f"sorted_token_ids row for slot {real_slot} has unwritten sentinel "
261
+ "entries; the align kernel skipped the real LoRA slot."
262
+ )
263
+
264
+
265
+def test_moe_lora_align_block_size_disabled_adapter_untouched():
266
+ """Disabled-adapter slot rows must remain untouched by all three align
267
+ kernels. Pins the invariant protected by the ``adapter_enabled`` guard
268
+ in ``lora_count_and_sort_expert_tokens_kernel``: without it the sort
269
+ kernel reads uninitialized ``token_mask`` values for disabled slots and
270
+ pollutes ``sorted_token_ids`` / ``cumsum_buffer``."""
271
+ max_loras = 1
272
+ out = _build_and_run_align(
273
+ num_lora_tokens=16,
274
+ num_base_tokens=0,
275
+ max_loras=max_loras,
276
+ disabled_slots=(0,),
277
+ )
278
+ # Sanity: slot 0 IS present in active_lora_ids (otherwise we would only
279
+ # exercise the lora_id == -1 / >= max_loras guards).
280
+ assert (out["lora_ids"] == 0).any().item()
281
+
282
+ assert out["num_tokens_post_pad"][0].item() == SENTINEL_NPAD, (
283
+ "num_tokens_post_pad[0] was modified for a disabled adapter slot."
284
+ )
285
+ expert_row = out["expert_ids"].view(max_loras, -1)[0]
286
+ assert (expert_row == SENTINEL_EXPERT).all(), (
287
+ "expert_ids row for disabled slot 0 was partially written."
288
+ )
289
+ # Row specifically protected by the sort-kernel adapter_enabled guard.
290
+ sorted_row = out["sorted_token_ids"].view(max_loras, -1)[0]
291
+ assert (sorted_row == SENTINEL_TOKEN).all(), (
292
+ "sorted_token_ids row for disabled slot 0 was polluted by the sort "
293
+ "kernel; lora_count_and_sort_expert_tokens_kernel must skip "
294
+ "adapter_enabled == 0 slots."
295
+ )
296
+
297
+
298
+def test_moe_lora_align_block_size_lora_id_oob_guard():
299
+ """Regression test for the ``lora_id >= max_loras`` guard.
300
+
301
+ Production ``LoRAKernelMeta.prepare_tensors`` pre-fills the tail of
302
+ ``active_lora_ids`` with -1, so the existing ``lora_id == -1`` check
303
+ covers the extra slot. This test bypasses that invariant and injects
304
+ an out-of-range value (5 with max_loras=1) at the tail to verify the
305
+ explicit guard prevents OOB reads against ``adapter_enabled`` and
306
+ OOB writes against the max_loras-sized output buffers. Without the
307
+ guard, an illegal-memory-access would surface on the next CUDA sync.
308
+ """
309
+ max_loras = 1
310
+ lora_ids_override = torch.tensor([0, 5], dtype=torch.int32)
311
+ out = _build_and_run_align(
312
+ num_lora_tokens=16,
313
+ num_base_tokens=0,
314
+ max_loras=max_loras,
315
+ lora_ids_override=lora_ids_override,
316
+ )
317
+ # The .item() call below syncs and would surface any async
318
+ # illegal-memory-access from the OOB iteration.
319
+ assert out["num_tokens_post_pad"][0].item() != SENTINEL_NPAD, (
320
+ "real LoRA slot 0 was skipped by the align kernel"
321
+ )
322
+
323
+
324
if __name__ == "__main__":
325
pytest.main([__file__])