已合并
flash_mla_with_kvcache优化性能&&代码规范修改 #180
yangyibin创建于 18 天前
flash_mla_with_kvcache优化性能&&代码规范修改 #180
已合并
共 4 个文件变更+21-16
| @@ -137,9 +137,9 @@ flash_mla_with_kvcache( | |||
| 137 | mask_mode=0, | 137 | mask_mode=0, |
| 138 | max_seqlen_q=-1, | 138 | max_seqlen_q=-1, |
| 139 | max_seqlen_kv=-1, | 139 | max_seqlen_kv=-1, |
| 140 | - layout_q="BSND", | 140 | + layout_q="TND", |
| 141 | layout_kv="PA_BBND", | 141 | layout_kv="PA_BBND", |
| 142 | - layout_out="BSND", | 142 | + layout_out="NTD", |
| 143 | return_softmax_lse=False, | 143 | return_softmax_lse=False, |
| 144 | ) | 144 | ) |
| 145 | ``` | 145 | ``` |
| @@ -169,7 +169,7 @@ flash_mla_with_kvcache( | |||
| 169 | - `attention_out`：shape 为 `[N, T, 512]`，dtype 与 Q 一致。 | 169 | - `attention_out`：shape 为 `[N, T, 512]`，dtype 与 Q 一致。 |
| 170 | - `softmax_lse`：启用时 shape 为 `[N, T]`、dtype 为 `float32`；关闭时返回空的 `float32` Tensor。 | 170 | - `softmax_lse`：启用时 shape 为 `[N, T]`、dtype 为 `float32`；关闭时返回空的 `float32` Tensor。 |
| 171 | 171 | ||
| 172 | -注意：虽然函数签名保留了 `BSND` 默认值，当前 checker 开放的主算子组合为 `layout_q="TND"` 和 `layout_out="NTD"`，调用时应显式传入。 | 172 | +注意：默认值即 checker 当前开放的主算子组合（`layout_q="TND"`、`layout_out="NTD"`）；传入 `BSND` 等其他布局仍会被 checker 拦截。 |
| 173 | 173 | ||
| 174 | ## 性能对比 | 174 | ## 性能对比 |
| 175 | 175 | ||
| @@ -252,11 +252,9 @@ def _check_supported_features( | |||
| 252 | """Check limits imposed by the currently implemented feature set.""" | 252 | """Check limits imposed by the currently implemented feature set.""" |
| 253 | layout_q = _normalize_layout(layout_q, SUPPORTED_Q_LAYOUTS, "layout_q") | 253 | layout_q = _normalize_layout(layout_q, SUPPORTED_Q_LAYOUTS, "layout_q") |
| 254 | layout_kv = _normalize_layout(layout_kv, SUPPORTED_KV_LAYOUTS, "layout_kv") | 254 | layout_kv = _normalize_layout(layout_kv, SUPPORTED_KV_LAYOUTS, "layout_kv") |
| 255 | - layout_out = "BSND" if layout_out is None else _normalize_layout( | 255 | + layout_out = "NTD" if layout_out is None else _normalize_layout( |
| 256 | layout_out, ("NTD",), "layout_out" | 256 | layout_out, ("NTD",), "layout_out" |
| 257 | ) | 257 | ) |
| 258 | - if layout_out != "NTD": | ||
| 259 | - raise ValueError("Layout_out must be NTD") | ||
| 260 | if mask_mode not in SUPPORTED_MASK_MODES: | 258 | if mask_mode not in SUPPORTED_MASK_MODES: |
| 261 | raise ValueError( | 259 | raise ValueError( |
| 262 | f"Mask_mode must be {NO_MASK_MODE} or {CAUSAL_MASK_MODE}; " | 260 | f"Mask_mode must be {NO_MASK_MODE} or {CAUSAL_MASK_MODE}; " |
| @@ -342,6 +340,7 @@ def check_flash_mla_inputs( | |||
| 342 | return_softmax_lse, | 340 | return_softmax_lse, |
| 343 | ): | 341 | ): |
| 344 | """Validate inputs by category and derive host launch metadata.""" | 342 | """Validate inputs by category and derive host launch metadata.""" |
| 343 | + layout_q = "TND" if layout_q is None else layout_q | ||
| 345 | _check_required_inputs( | 344 | _check_required_inputs( |
| 346 | metadata=metadata, | 345 | metadata=metadata, |
| 347 | block_table=block_table, | 346 | block_table=block_table, |
| @@ -1659,16 +1659,21 @@ class SplitKvReducer: | |||
| 1659 | partial_output_tile = self.partial_output_handoff.consume() | 1659 | partial_output_tile = self.partial_output_handoff.consume() |
| 1660 | with vf(mode="simd"): | 1660 | with vf(mode="simd"): |
| 1661 | full_lane_mask, _ = rr.update_mask(VL_T, elem_bits=32) | 1661 | full_lane_mask, _ = rr.update_mask(VL_T, elem_bits=32) |
| 1662 | - for row in range(active_rows): | 1662 | + if split_idx > 0: |
| 1663 | - weight_offset = split_idx * FD_WINDOW_ROWS + row | 1663 | + for row in range(active_rows): |
| 1664 | - split_weight = rr.vload_brc(self.split_weight_buffer, weight_offset) | 1664 | + split_weight = rr.vload_brc(self.split_weight_buffer, split_idx * FD_WINDOW_ROWS + row) |
| 1665 | - for col in tuple(range(0, TILE_D, VL_T)): | 1665 | + for col in tuple(range(0, TILE_D, VL_T)): |
| 1666 | - value = rr.vload(partial_output_tile, (row + window_row_offset) * TILE_D + col) | 1666 | + value = rr.vload(partial_output_tile, (row + window_row_offset) * TILE_D + col) |
| 1667 | - result = rr.vmul(value, split_weight, mask=full_lane_mask) | ||
| 1668 | - if split_idx > 0: | ||
| 1669 | old = rr.vload(self.output_accumulator, row * TILE_D + col) | 1667 | old = rr.vload(self.output_accumulator, row * TILE_D + col) |
| 1670 | result = rr.vmadd(value, split_weight, old, mask=full_lane_mask) | 1668 | result = rr.vmadd(value, split_weight, old, mask=full_lane_mask) |
| 1671 | - rr.vstore(self.output_accumulator, row * TILE_D + col, result, full_lane_mask) | 1669 | + rr.vstore(self.output_accumulator, row * TILE_D + col, result, full_lane_mask) |
| 1670 | + else: | ||
| 1671 | + for row in range(active_rows): | ||
| 1672 | + split_weight = rr.vload_brc(self.split_weight_buffer, split_idx * FD_WINDOW_ROWS + row) | ||
| 1673 | + for col in tuple(range(0, TILE_D, VL_T)): | ||
| 1674 | + value = rr.vload(partial_output_tile, (row + window_row_offset) * TILE_D + col) | ||
| 1675 | + result = rr.vmul(value, split_weight, mask=full_lane_mask) | ||
| 1676 | + rr.vstore(self.output_accumulator, row * TILE_D + col, result, full_lane_mask) | ||
| 1672 | rr.vmem_bar("vst_vld") | 1677 | rr.vmem_bar("vst_vld") |
| 1673 | 1678 | ||
| 1674 | def cast_combined_output_aligned(self): | 1679 | def cast_combined_output_aligned(self): |
| @@ -3160,9 +3165,9 @@ def flash_mla_with_kvcache( | |||
| 3160 | mask_mode=0, | 3165 | mask_mode=0, |
| 3161 | max_seqlen_q=-1, | 3166 | max_seqlen_q=-1, |
| 3162 | max_seqlen_kv=-1, | 3167 | max_seqlen_kv=-1, |
| 3163 | - layout_q="BSND", | 3168 | + layout_q="TND", |
| 3164 | layout_kv="PA_BBND", | 3169 | layout_kv="PA_BBND", |
| 3165 | - layout_out="BSND", | 3170 | + layout_out="NTD", |
| 3166 | return_softmax_lse=False, | 3171 | return_softmax_lse=False, |
| 3167 | ): | 3172 | ): |
| 3168 | """Launch Flash MLA with direct Q/O/LSE storage and no fold GM buffers.""" | 3173 | """Launch Flash MLA with direct Q/O/LSE storage and no fold GM buffers.""" |
| @@ -3213,6 +3218,7 @@ def flash_mla_with_kvcache( | |||
| 3213 | ) | 3218 | ) |
| 3214 | softmax_lse = lse_storage.transpose(0, 1) | 3219 | softmax_lse = lse_storage.transpose(0, 1) |
| 3215 | elif layout_q == "BSND": | 3220 | elif layout_q == "BSND": |
| 3221 | + seqlen_q = q.shape[1] # BSND 的 q 为 [B, S, N, D] | ||
| 3216 | lse_storage = torch.empty( | 3222 | lse_storage = torch.empty( |
| 3217 | (batch_size, seqlen_q, num_heads_q), | 3223 | (batch_size, seqlen_q, num_heads_q), |
| 3218 | dtype=torch.float32, | 3224 | dtype=torch.float32, |