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

This commit is contained in:
Yanbin Jiang
2026-07-21 09:47:17 -07:00
committed by GitHub
parent 604b1507d2
commit 57e5846b90
@@ -473,6 +473,12 @@ def _fused_moe_kernel_sequence(
topk = topk_ids.shape[1]
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 = (
min(num_tokens * topk, E + 1) * (config["BLOCK_SIZE_M"] - 1)
if down_moe_use_tma