[LoRA] Guard TMA down path for LoRA hooks (#31608)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user