[1/2] Speed up trtllm_mla attention backend (>10% e2e) (#10473)

2025-09-15 11:53:21 -07:00
parent 5c08d7d21d
commit 3b25dc127a
6 changed files with 119 additions and 3 deletions
@@ -23,6 +23,7 @@ from sgl_kernel.cutlass_moe import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_
 from sgl_kernel.elementwise import (
    FusedSetKVBufferArg,
    apply_rope_with_cos_sin_cache_inplace,
+    concat_mla_absorb_q,
    concat_mla_k,
    copy_to_gpu_no_ce,
    downcast_fp8,