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
[Kernel] Pack topk id/weights triton kernel (#42527)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
May 18, 2026 at 18:04 UTC
2267f70070bdee8057b4afae69cba9b847add587
2 files changed
+54
-8
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+1
@@ -343,6 +343,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
343
# the TRTLLM C++ kernel supports.
344
"routing_method_type": RoutingMethodType.Renormalize,
345
"do_finalize": True,
346
+ "enable_pdl": True,
347
"output": output,
348
"tune_max_num_tokens": max(self.max_capture_size, 1),
349
}
vllm/model_executor/layers/fused_moe/utils.py
+53
-8
@@ -376,17 +376,62 @@ def disable_inplace() -> bool:
376
return is_torch_equal_or_newer("2.9")
377
378
379
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
379
+@triton.jit
380
+def _pack_topk_ids_weights_kernel(
381
+ topk_ids_ptr,
382
+ topk_weights_ptr,
383
+ output_ptr,
384
+ n_elements,
385
+ BLOCK_SIZE: tl.constexpr,
386
+ USE_GDC: tl.constexpr,
387
+ launch_pdl: tl.constexpr, # triton metadata
388
+):
389
+ pid = tl.program_id(axis=0)
390
+ offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
391
+ mask = offsets < n_elements
392
+ if USE_GDC:
393
+ tl.extra.cuda.gdc_launch_dependents()
394
+ tl.extra.cuda.gdc_wait()
395
+ expert_id = tl.load(topk_ids_ptr + offsets, mask=mask, other=0).to(tl.int32)
396
+ expert_id_shifted = expert_id << 16
397
+
398
+ weight = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0)
399
+ weight_bf16 = weight.to(tl.bfloat16)
400
+ weight_int16 = weight_bf16.to(tl.int16, bitcast=True)
401
+
402
+ weight_int32 = weight_int16.to(tl.int32) & 0xFFFF
403
+
404
+ packed = expert_id_shifted | weight_int32
405
+ tl.store(output_ptr + offsets, packed, mask=mask)
406
+
407
+
408
def trtllm_moe_pack_topk_ids_weights(
381
- topk_ids: torch.Tensor, topk_weights: torch.Tensor
409
+ topk_ids: torch.Tensor,
410
+ topk_weights: torch.Tensor,
411
+ block_size: int = 1024,
412
) -> torch.Tensor:
383
- """
384
- Pack topk_ids and topk_weights into a single int32 tensor.
385
- Format: (expert_id << 16) | weight_bf16.view(int16)
386
- """
387
- return (topk_ids.to(torch.int32) << 16) | topk_weights.to(torch.bfloat16).view(
388
- torch.int16
413
+ assert topk_ids.shape == topk_weights.shape
414
+ assert topk_ids.is_contiguous() and topk_weights.is_contiguous()
415
+
416
+ original_shape = topk_ids.shape
417
+ ids_flat = topk_ids.reshape(-1)
418
+ weights_flat = topk_weights.reshape(-1)
419
+
420
+ n_elements = ids_flat.numel()
421
+ output = torch.empty(n_elements, dtype=torch.int32, device=topk_ids.device)
422
+
423
+ use_gdc = current_platform.is_cuda() and current_platform.has_device_capability(90)
424
+ grid = (triton.cdiv(n_elements, block_size),)
425
+ _pack_topk_ids_weights_kernel[grid](
426
+ ids_flat,
427
+ weights_flat,
428
+ output,
429
+ n_elements,
430
+ BLOCK_SIZE=block_size,
431
+ USE_GDC=use_gdc,
432
+ launch_pdl=use_gdc,
433
)
434
+ return output.reshape(original_shape)
435
436
437
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)