[VLM][LLM] Optimize fused_moe triton kernel tma (#18782)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import functools
 import os
+from collections import OrderedDict
 from typing import Any, Dict, List, Optional
 
 import torch
@@ -610,6 +611,67 @@ def fused_moe_kernel(
     tl.store(c_ptrs, accumulator, mask=c_mask)
 
 
+# -----------------------------------------------------------------------------
+# TMA allocator: set once per process (avoid per-call triton.set_allocator)
+# -----------------------------------------------------------------------------
+_TMA_ALLOCATOR_SET = False
+
+
+def _set_triton_tma_allocator():
+    """TMA descriptors require a global allocator; set it once to avoid per-call overhead."""
+    global _TMA_ALLOCATOR_SET
+    if _TMA_ALLOCATOR_SET:
+        return
+
+    # TMA descriptors require a global memory allocation
+    def alloc_fn(size: int, alignment: int, stream: Optional[int]):
+        # NOTE: keep this allocation on CUDA device
+        return torch.empty(size, device="cuda", dtype=torch.int8)
+
+    triton.set_allocator(alloc_fn)
+    _TMA_ALLOCATOR_SET = True
+
+
+# --- B TensorDescriptor cache (LRU) ---
+_B_DESC_CACHE_MAX = 64
+_B_DESC_CACHE: "OrderedDict[tuple, TensorDescriptor]" = OrderedDict()
+
+
+def _get_b_tma_desc_cached(B: torch.Tensor, block_n: int, block_k: int):
+    """
+    Cache TensorDescriptor for constant weight B.
+    Keyed by storage ptr + shape/stride/dtype + tile shape.
+    """
+    key = (
+        int(B.data_ptr()),
+        tuple(B.shape),
+        tuple(B.stride()),
+        str(B.dtype),
+        int(block_n),
+        int(block_k),
+    )
+
+    desc = _B_DESC_CACHE.get(key, None)
+    if desc is not None:
+        _B_DESC_CACHE.move_to_end(key)
+        return desc
+
+    # Create outside lock to reduce lock hold time (ok if duplicated rarely)
+    desc = TensorDescriptor(
+        B,
+        B.shape,
+        B.stride(),
+        [1, block_n, block_k],
+    )
+
+    _B_DESC_CACHE[key] = desc
+    _B_DESC_CACHE.move_to_end(key)
+    if len(_B_DESC_CACHE) > _B_DESC_CACHE_MAX:
+        _B_DESC_CACHE.popitem(last=False)
+
+    return desc
+
+
 def invoke_fused_moe_kernel(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -754,11 +816,8 @@ def invoke_fused_moe_kernel(
 
     else:
         if a_use_tma or b_use_tma:
-            # TMA descriptors require a global memory allocation
+            _set_triton_tma_allocator()
-            def alloc_fn(size: int, alignment: int, stream: Optional[int]):
-                return torch.empty(size, device="cuda", dtype=torch.int8)
 
-            triton.set_allocator(alloc_fn)
         if a_use_tma:
             a_desc = TensorDescriptor(
                 A, A.shape, A.stride(), [config["BLOCK_SIZE_M"], config["BLOCK_SIZE_K"]]
@@ -766,11 +825,11 @@ def invoke_fused_moe_kernel(
         else:
             a_desc = None
         if b_use_tma:
-            b_desc = TensorDescriptor(
+            # B is constant weights -> cache descriptor
+            b_desc = _get_b_tma_desc_cached(
                 B,
-                B.shape,
+                config["BLOCK_SIZE_N"],
-                B.stride(),
+                config["BLOCK_SIZE_K"],
-                [1, config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"]],
             )
         else:
             b_desc = None