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

[LoRA] Guard TMA down path for LoRA hooks (#31608)

Browse Source
This commit is contained in:
Yanbin Jiang
2026-07-21 09:47:17 -07:00
committed by GitHub
parent 604b1507d2
commit 57e5846b90
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+6
View File
@@ -473,6 +473,12 @@ def _fused_moe_kernel_sequence(
topk = topk_ids.shape[1] topk = topk_ids.shape[1]
compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16 compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
# LoRA hooks consume and update route-major intermediate buffers. The TMA
# down path keeps those buffers in expert-sorted, block-padded order, which
# is incompatible with the hook contract.
if hooks and (hooks.after_gate_up is not None or hooks.after_down is not None):
down_moe_use_tma = False
padded_tokens = ( padded_tokens = (
min(num_tokens * topk, E + 1) * (config["BLOCK_SIZE_M"] - 1) min(num_tokens * topk, E + 1) * (config["BLOCK_SIZE_M"] - 1)
if down_moe_use_tma if down_moe_use_tma