EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Lint]Style: Convert vllm-ascend/ to ruff format(Batch #12) (#6177)

Browse Source 
### What this PR does / why we need it?
**Scope of Changes**:
| File Path |
| :--- |
| `vllm_ascend/ops/triton/activation/swiglu_quant.py` |
| `vllm_ascend/ops/triton/batch_invariant/matmul.py` |
| `vllm_ascend/ops/triton/batch_invariant/mean.py` |
| `vllm_ascend/ops/triton/batch_invariant/rmsnorm.py` |
| `vllm_ascend/ops/triton/fla/chunk.py` |
| `vllm_ascend/ops/triton/fla/chunk_delta_h.py` |
| `vllm_ascend/ops/triton/fla/chunk_o.py` |
| `vllm_ascend/ops/triton/fla/chunk_scaled_dot_kkt.py` |
| `vllm_ascend/ops/triton/fla/cumsum.py` |
| `vllm_ascend/ops/triton/fla/fused_qkvzba_split_reshape.py` |
| `vllm_ascend/ops/triton/fla/l2norm.py` |
| `vllm_ascend/ops/triton/fla/layernorm_guard.py` |
| `vllm_ascend/ops/triton/fla/sigmoid_gating.py` |
| `vllm_ascend/ops/triton/fla/solve_tril.py` |
| `vllm_ascend/ops/triton/fla/utils.py` |
| `vllm_ascend/ops/triton/fla/wy_fast.py` |
| `vllm_ascend/ops/triton/fused_gdn_gating.py` |
| `vllm_ascend/ops/triton/layernorm_gated.py` |
| `vllm_ascend/ops/triton/linearnorm/split_qkv_rmsnorm_rope.py` |
| `vllm_ascend/ops/triton/mamba/causal_conv1d.py` |
| `vllm_ascend/ops/triton/reject_sample.py` |
| `vllm_ascend/ops/triton/rope.py` |
| `vllm_ascend/ops/triton/spec_decode/utils.py` |
| `vllm_ascend/ops/triton/triton_utils.py` |

### Does this PR introduce _any_ user-facing change?

### How was this patch tested?

- vLLM version: v0.14.0
- vLLM main:
d68209402d

Signed-off-by: MrZ20 <2609716663@qq.com>
This commit is contained in:
SILONG ZENG
2026-01-23 14:59:19 +08:00
committed by GitHub
parent 193acc2c19
commit 78af0c30a3
25 changed files with 760 additions and 996 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
vllm_ascend/ops/triton/fla/fused_qkvzba_split_reshape.py
View File

@@ -31,29 +31,30 @@ def fused_qkvzba_split_reshape_cat_kernel(
BA_DIM_T: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK * 2
QKV_DIM_T: tl.constexpr = HEAD_QK * 2 + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V
q_end: tl.constexpr = HEAD_QK
blk_q_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T +
i_qk * QKVZ_DIM_T + tl.arange(0, q_end))
blk_q_ptr = mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(0, q_end)
k_end: tl.constexpr = q_end + HEAD_QK
blk_k_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T +
i_qk * QKVZ_DIM_T + tl.arange(q_end, k_end))
blk_k_ptr = mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(q_end, k_end)
v_end: tl.constexpr = k_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V
blk_v_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T +
i_qk * QKVZ_DIM_T + tl.arange(k_end, v_end))
blk_v_ptr = mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(k_end, v_end)
z_end: tl.constexpr = v_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V
blk_z_ptr = (mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T +
i_qk * QKVZ_DIM_T + tl.arange(v_end, z_end))
blk_q_st_ptr = (mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T +
i_qk * HEAD_QK + tl.arange(0, HEAD_QK))
blk_k_st_ptr = (mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T +
NUM_HEADS_QK * HEAD_QK + i_qk * HEAD_QK +
tl.arange(0, HEAD_QK))
blk_v_st_ptr = (mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T +
NUM_HEADS_QK * HEAD_QK * 2 +
i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK +
tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK))
blk_z_st_ptr = (z + i_bs * NUM_HEADS_V * HEAD_V +
i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK +
tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK))
blk_z_ptr = mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(v_end, z_end)
blk_q_st_ptr = mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
blk_k_st_ptr = (
mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T + NUM_HEADS_QK * HEAD_QK + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
)
blk_v_st_ptr = (
mixed_qkv
+ i_bs * NUM_HEADS_QK * QKV_DIM_T
+ NUM_HEADS_QK * HEAD_QK * 2
+ i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK
+ tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK)
)
blk_z_st_ptr = (
z
+ i_bs * NUM_HEADS_V * HEAD_V
+ i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK
+ tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK)
)
tl.store(blk_q_st_ptr, tl.load(blk_q_ptr))
tl.store(blk_k_st_ptr, tl.load(blk_k_ptr))
tl.store(blk_v_st_ptr, tl.load(blk_v_ptr))
@@ -66,8 +67,7 @@ def fused_qkvzba_split_reshape_cat_kernel(
tl.store(blk_b_st_ptr, tl.load(blk_b_ptr))
for i in tl.static_range(b_end, a_end):
blk_a_ptr = mixed_ba + i_bs * NUM_HEADS_QK * BA_DIM_T + i_qk * BA_DIM_T + i
blk_a_st_ptr = (a + i_bs * NUM_HEADS_V +
i_qk * NUM_HEADS_V // NUM_HEADS_QK + (i - b_end))
blk_a_st_ptr = a + i_bs * NUM_HEADS_V + i_qk * NUM_HEADS_V // NUM_HEADS_QK + (i - b_end)
tl.store(blk_a_st_ptr, tl.load(blk_a_ptr))