[Perf][Moe]improve cutlass_moe_fp4 performance by using apply_router_weight_on_i… (#19493)
This commit is contained in:
@@ -491,10 +491,12 @@ def cutlass_moe_fp4(
|
||||
params.to_gemm2_args(),
|
||||
)
|
||||
del int_fp4, int_blockscale
|
||||
c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size))
|
||||
c2 = c2.view(m_a, num_topk, params.hidden_size)
|
||||
|
||||
if no_combine:
|
||||
c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size))
|
||||
c2 = c2.view(m_a, num_topk, params.hidden_size)
|
||||
return c2.to(out_dtype)
|
||||
if not apply_router_weight_on_input:
|
||||
c2 = c2 * topk_weights.view(m_a, num_topk, 1).to(out_dtype)
|
||||
return c2.sum(dim=1).to(out_dtype)
|
||||
output = torch.empty((m_a, k_a), device=device, dtype=out_dtype)
|
||||
weights = topk_weights.to(out_dtype) if not apply_router_weight_on_input else None
|
||||
apply_shuffle_mul_sum(c2, output, c_map, weights)
|
||||
return output
|
||||
|
||||
Reference in New Issue
Block a user