[Example] Add GQA Sink Backward BHSD with robustness improvements by wch0810 · Pull Request #1 · wch0810/tilelang-ascend

wch0810 · 2026-08-20T02:47:14Z

Summary

Add GQA Sink Backward BHSD example with robustness improvements and 30-run cold-start stability verification.

Files Added

examples/attention_sink/example_gqa_sink_bwd_bhsd/example_gqa_sink_bwd_bhsd.py (1179 lines)
- 9-kernel backward: 5-kernel bwd split (k1-k5) + fwd + preprocess + postprocess + dsink
- No-scope Developer mode (no T.Scope/flag/barrier_all, pure AUTO_CV_SYNC + AUTO_SYNC)
- Compensated GEMM for fp16 precision recovery
- Online softmax with attention sink + sliding window mask
examples/attention_sink/example_gqa_sink_bwd_bhsd/test_gqa_sink_bwd_bhsd.py (1809 lines)
- L0(7) + L1(8) + L2(5) + Boundary(4) + do_bench + msprof
- 169-line precision standard (fp16: atol=6.10e-5, rtol=1.95e-3, max_abs_limit=0.1, required_ratio=0.99)

Robustness Improvements

stderr capture: _capture_bisheng_stderr(exc) — captures full C++ stderr on compile failure
/tmp cleanup hook: _cleanup_tmp_compilation_files() — cleans /tmp/tmp*.so after each test case (only .so, not .cpp, to avoid race condition)
compile retry: _run_with_retry(fn, max_retries=2) — retries on transient compile errors (precision failures propagate immediately)

msprof Bug Fixes

Fix per_kernel parameter shadowed by local dict variable
Fix BWD main kernel misidentification (was selecting pure Vector k2, now analyzes all 5 bwd kernels)
Replace single-kernel analysis with per-BWD-kernel bottleneck table (k1-k5)
Add explicit launch→kernel label mapping (was relying on execution order)

Verification Results

Precision: L0 7/7 + L1 8/8 all PRECISION_PASS (matches precision-standard.md double-gate)
Performance: bwd main 6151us vs GPU 14287us (56.9% faster)
Stability: 30 cold-start runs, 0 retry, 0 race condition, 450 cases all PASS
msprof: Per-BWD-kernel bottleneck table works (k2/k4 sync-bound, k1/k3/k5 mixed)

Test Commands

# Precision
PYTHONPATH=/mnt/workspace/cann/tilelang-github2/tilelang-ascend python test_gqa_sink_bwd_bhsd.py --level all

# Performance
PYTHONPATH=/mnt/workspace/cann/tilelang-github2/tilelang-ascend python test_gqa_sink_bwd_bhsd.py --level bench --per-kernel
PYTHONPATH=/mnt/workspace/cann/tilelang-github2/tilelang-ascend python test_gqa_sink_bwd_bhsd.py --level msprof

Known Limitations (non-blocking)

l2_shape_d64: D=64 silently accepted (BOUNDARY_WARN)
boundary_sink_inf: sink with ±inf causes precision divergence (BOUNDARY_WARN)
boundary_dbound: Q/K at ±32000 extreme values (BOUNDARY_WARN)
Expert 95% target (4464us) not reached (gap +1687us, needs Stage 3 optimization)

… 30-run cold-start stability verification - Add example_gqa_sink_bwd_bhsd.py: 9-kernel backward (5-kernel bwd split + fwd + preprocess + postprocess + dsink), no-scope Developer mode, Compensated GEMM for fp16 precision recovery, online softmax with attention sink + sliding window mask - Add test_gqa_sink_bwd_bhsd.py: L0(7) + L1(8) + L2(5) + Boundary(4) + do_bench + msprof, 169-line precision standard (fp16 atol=6.10e-5, rtol=1.95e-3, max_abs_limit=0.1, required_ratio=0.99) - Robustness improvements: stderr capture (_capture_bisheng_stderr), /tmp cleanup hook (_cleanup_tmp_compilation_files, .so only not .cpp to avoid race condition), compile retry (_run_with_retry max_retries=2) - msprof bug fixes: per_kernel shadow bug, BWD main kernel misidentification, single-kernel analysis replaced with per-BWD-kernel bottleneck table, explicit launch->kernel label mapping - Verified: bwd main 6151us vs GPU 14287us (56.9% faster), L0 7/7 + L1 8/8 all PRECISION_PASS, 30 cold-start runs 0 retry 0 race condition

…dispatch - Reduce 4 kernels to 1: flashattn_mla_decode (merge #2/tile-ai#3, delete #1/tile-ai#4) - workspace_3 dtype: fp16 -> fp32 (unified for all dim/block_N configs, no host dispatch) - Remove num_split>1 two-phase path (Phase 1 + Phase 2 deleted) - Remove ws3_fp16 host dispatch branching (single kernel for all configs) - run_mla_decode signature: drop num_split parameter - Test: remove 2 num_split>1 L0 cases (23 -> 21 cases) - Precision: 21/21 PASS (golden ratio=1.0000 max_abs=1.221e-04, unchanged) - Performance: 3877us (gap 1.28x, +6.4% vs 4-kernel 3644us, fp32 GM traffic) - ruff format + ruff check passed (line-length=140)

…8.5% latency) Responds to hedi515 review "需要这么多kernel吗" on PR tile-ai#1638. Merges 9 bwd kernels into 6 via combineCV auto-sync (Developer mode, no T.Scope/flag/barrier_all). Key finding: combineCV FetchWorkspaceName (ascend_combinecv.cc:190-202) requires intra-kernel CV transfer buffers named with "workspace" substring to trigger auto set_flag/wait_flag. Merge plan (9→6 kernels): - flashattn_fwd (unchanged) - flashattn_bwd_preprocess = preprocess + dsink (pure Vector) - flashattn_bwd_qk_softmax = k1 (QK^T Cube) + k2 (softmax Vector) - flashattn_bwd_dv_dp_ds = k3 (dV+dP Cube) + k4 (dS Vector) - flashattn_bwd_dk_dq = k5 (dK+dQ, unchanged) - flashattn_bwd_postprocess = dK + dV cast (dual-output) intra-kernel CV transfers (workspace_s, workspace_dp) named with "workspace" substring to trigger combineCV auto-sync; inter-kernel workspace (ws_p, ws_ds, etc.) kept as ws_* for host synchronize. Kernel names aligned to GPU reference style (flashattn_bwd_preprocess/ postprocess/qk_softmax/dv_dp_ds/dk_dq), removing k1_k2/k3_k4/k5 internal numbering that exposed merge history. Precision: L0 7/7 + L1 8/8 PASS (fp16 atol=6.10e-5, rtol=1.95e-3, 169-line standard, check_precision double-gate). Precision identical to baseline (ratio/max_abs match). Performance: bwd main 6152→5020us (-18.5% vs baseline, gap +556us to Expert 95% target 4464us). Per-kernel (msprof): qk_softmax 1593us, dv_dp_ds 1986us, dk_dq 1460us. L2 improvements: - Add explicit assert dim_qk % 128 == 0 (l2_shape_d64 now correctly rejected with clear error instead of silent compile failure) - l2_shape_n192 now valid (merged preprocess uses blk=32, 192%32==0), updated test expectation from BOUNDARY_WARN to BOUNDARY_PASS Stage 3 attempt (not adopted): Plan D merged k3_k4+k5 into 5 kernels, precision PASS but -16.2% regression (V→C cross_flag sync serializes CV overlap, +812us sync overhead outweighs 192MB L2 cache savings). Rolled back to 6-kernel.

- F841: remove unused dO tensor (case only runs fwd to verify N=192 is valid) - F541: remove extraneous f-string prefix (no placeholders) ruff check + format: all pass.

…6→4 kernels) Responds to ShareableXue review requesting single-kernel bwd architecture. Merges rev0's 3 bwd main kernels (flashattn_bwd_qk_softmax + flashattn_bwd_dv_dp_ds + flashattn_bwd_dk_dq) into 1 flashattn_bwd single kernel via Developer mode + combineCV (4 True, zero T.Scope/flag/barrier_all). Architecture: per-iteration C-V-C-V-C fusion in single T.serial loop. 5 GEMMs + softmax + dS all in one kernel, combineCV auto-separates Cube/Vector code by buffer storage scope. Key technical points: - 6 intra-kernel workspace_* buffers (workspace_s/dp/p/p_delta/ds/ds_delta) named with 'workspace' substring to trigger combineCV FetchWorkspaceName (ascend_combinecv.cc:190-202) auto set_flag/wait_flag sync - P_fp32 UB retention: P_fp32 retained in s_ub (UB), used directly in dS compute, NOT written to GM — eliminates workspace_p_fp32 - CompGEMM dual-read solved: workspace_p (main GEMM2) + workspace_p_delta (corr GEMM2corr), each strictly 1 write + 1 read per iteration - Workspace dimensions [bwd_block_num, block_M, block_N] (no kv_iter dim): per-iteration produce+consume, not cross-iteration Results vs rev0 (6-kernel): - bwd main kernel count: 3 -> 1 (aligns with GPU reference flashattn_bwd) - total kernel count: 6 -> 4 (fwd + preprocess + bwd + postprocess) - bwd inter-kernel host synchronize: 2 -> 0 (combineCV auto-sync) - bwd inter-kernel GM workspace: 449MB -> 0 (all intra-kernel) Precision: L0 7/7 + L1 8/8 ALL PRECISION_PASS (fp16 169-line standard, identical to rev0). Coverage 21/21 PASS. ruff check + format passed. Performance: bwd main 5556us (vs rev0 5020us, +10.7% regression from V->C sync serializing CV overlap — accepted as architecture fusion cost).

…2 kernels) Responds to ShareableXue review: integrate bwd into one kernel, keep tiled intermediates on chip, eliminate host sync between stages. Architecture (2 kernels, Developer + combineCV, 0 T.Scope/flag/barrier_all): 1. flashattn_fwd: Forward (online softmax + sink + window) -> O, lse 2. flashattn_bwd: Single kernel: Phase0(Delta) + KV-loop(5 GEMM + softmax + dS) + Phase5(dSink). dK/dV fp32 atomic_add + host .half(). Key changes from previous rev (4→2 kernels): - Merge preprocess into bwd Phase 0 (Delta = sum(O*dO)) + Phase 5 (dSink) - Eliminate postprocess kernel: dK/dV cast via host .half() (same pattern as varlen rev3, atomic_add cross-Q-block accumulation requires fp32 + host sync) - Eliminate all 6 GM workspace buffers: alloc_shared/fragment for on-chip CV transfer (T.copy(fragment, shared) — compiler auto-inserts GM relay via AUTO_CV_SYNC, no explicit workspace_* tensors) - P_fp32 retained in UB (s_ub) for dS compute — no GM roundtrip - Compensated GEMM via L0C accumulation (init=True main + init=False corr) - threads=1, no vid Results: - bwd kernel count: 3 → 1 (single flashattn_bwd) - total kernel count: 4 → 2 (fwd + bwd) - bwd launches: 2 → 1 - bwd inter-kernel GM workspace: 480MB → 0 - bwd inter-kernel host sync: 2 → 1 (atomic_add completion, cannot eliminate) - Precision: L0 7/7 + L1 8/8 ALL PASS (fp16 169-line standard, identical to previous revs). Smoke test checks all 6 outputs (fwd_O/Delta/dQ/dK/dV/dSinks). - Performance: bwd 5127us (vs Expert 4699us: -9.1%, vs GPU 14287us: +64.1%) Remaining host sync (1x) is atomic_add cross-Q-block accumulation — dK/dV receive contributions from multiple Q-blocks via fp32 atomic_add; cast to fp16 must wait for all blocks. fp16 atomic_add attempted but precision fails (accumulation rounding error scales with Q-block count). This is a hardware constraint, not an implementation choice.

…e bench/msprof - Move ref_fwd/ref_bwd golden functions from example to test file - Simplify example __main__ to single smoke test case (shape + finiteness check) - Remove bench and msprof from test file (precision-only, -816 lines total) - Convert all comments to English, remove version tags (P1/P2/rev3 etc.) - Keep 2-kernel architecture: flashattn_fwd + flashattn_bwd (single kernel bwd)

wch0810 · 2026-08-25T07:10:54Z

算子名	主仓路径（tilelang主仓）https://github.com/tile-ai/tilelang	性能指标（估算）H20
GQA Forward (VarLen)	examples/flash_attention/example_gqa_fwd_varlen.py	batch=8,heads=64,groups=16,q_seqlen=2048,k_seqlen=2048,dim=128，fp16，性能11429us
GQA Backward	examples/flash_attention/example_gqa_bwd.py	batch=8,h=32,n_ctx=1024,d_head_qk=192,d_head_v=128,groups=16，fp16，性能4469us
MHA Sink Forward	examples/attention_sink/example_mha_sink_fwd_bhsd.py	batch=8,heads=32,seq_q=4096,seq_kv=4096,dim=128，fp16，22859us
MHA Sink Backward	examples/attention_sink/example_mha_sink_bwd_bhsd.py	batch=1,h=64,n_ctx=4096,d_head=128，fp16，性能14287us
GQA Sink Forward VarLen	examples/attention_sink/example_gqa_sink_fwd_varlen.py	batch=8,heads=64,groups=16,q_seqlen=2048,k_seqlen=2048,dim=128，fp16，性能11429us
GQA Sink Backward	examples/attention_sink/example_gqa_sink_bwd_bhsd.py	batch=1,h=64,n_ctx=4096,d_head=128,groups=8，fp16，性能14287us
GQA Sink Backward VarLen	examples/attention_sink/example_gqa_sink_bwd_varlen.py	batch=8,heads=64,groups=16,q_seqlen=2048,k_seqlen=2048,dim=128，fp16，性能28574us
GQA Decode VarLen Mask	examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py	batch=8,heads=32,heads_kv=8,max_cache_seqlen=8192,dim=128,dim_v=128,sparse_ratio=0.8，fp16,性能20us
Fused MOE	examples/fusedmoe/example_fusedmoe_tilelang.py	d_hidden=7168,d_expert=2048,n_routed_experts=8,n_shared_experts=1,n_experts_per_token=4,batch_size=1,seq_len=8192，fp16，性能37503us
MHC Pre	examples/deepseek_mhc/example_mhc_pre.py	n=2048,hidden_size=2560,hc_mult=4，fp16，性能27us
MHC Post	examples/deepseek_mhc/example_mhc_post.py	n=4096,h=2560,hc_mult=4，fp16，性能108us
MHC Backward	examples/deepseek_mhc/example_mhc_bwd.py	seqlen=65536,n_stream=8，fp16，性能32us
MLA Decode	examples/deepseek_mla/example_mla_decode.py	batch=132,heads=128,kv_heads=1,kv_ctx=8192,dim=512,pe_dim=64，fp16，性能3131us
MLA Decode Paged	examples/deepseek_mla/example_mla_decode_paged.py	batch=128,h_q=128,h_kv=1,cache_seqlen=8192,d=576,dv=512，fp16，性能3036us
MLA Decode Persistent	examples/deepseek_mla/example_mla_decode_persistent.py	batch=128,heads=128,kv_heads=1,kv_ctx=8192,dim=512,pe_dim=64，fp16，性能3036us
NSA Forward	examples/deepseek_nsa/example_tilelang_nsa_fwd.py	B=2,SEQ_LEN=64,H=1,HQ=16,D=32,S=1,block_size=32，fp16，性能20us
NSA Forward VarLen	examples/deepseek_nsa/example_tilelang_nsa_fwd_varlen.py	N=2,C_SEQ_LEN=64,H=1,HQ=16,D=64,S=1,block_size=32，fp16，性能20us
NSA Backward	examples/deepseek_nsa/example_tilelang_nsa_bwd.py	B=1,T=32,H=1,HQ=16,D=32,S=1,block_size=32，fp16，性能20us
NSA Decode	examples/deepseek_nsa/example_tilelang_nsa_decode.py	B=2,SEQ_LEN=64,H=1,HQ=16,D=16,S=1,block_size=32,SEQ_LEN_Q=1，fp16，性能20us
Sparse MLA Forward	examples/deepseek_v32/sparse_mla_fwd.py	B=1,S=4096,SKV=8192,H=128,HKV=1,DQK=576,DV=512,topk=2048，fp16，性能24288us
Sparse MLA Backward	examples/deepseek_v32/sparse_mla_bwd.py
Sparse MLA Forward Pipelined	examples/deepseek_v32/sparse_mla_fwd_pipelined.py
Sparse MLA Forward SeeSaw	examples/deepseek_v32/sparse_mla_fwd_seesaw.py
Chunk Delta Backward	examples/gdn/example_chunk_delta_bwd.py
Chunk Delta H	examples/gdn/example_chunk_delta_h.py
Chunk O Forward	examples/gdn/example_chunk_o.py
Chunk O Backward	examples/gdn/example_chunk_o_bwd.py
Chunk Scaled Dot KKT	examples/gdn/example_chunk_scaled_dot_kkt.py
WY Fast	examples/gdn/example_wy_fast.py
WY Fast Backward Split	examples/gdn/example_wy_fast_bwd_split.py
Chunk Delta Backward	examples/kda/chunk_delta_bwd.py
Chunk Delta H Forward	examples/kda/chunk_delta_h_fwd.py
Chunk Inter Solve Fused	examples/kda/chunk_inter_solve_fused.py
Chunk Intra Token Parallel	examples/kda/chunk_intra_token_parallel.py
Chunk O	examples/kda/chunk_o.py
Chunk Backward (DQKWG)	examples/kda/chunk_bwd_dqkwg.py
Chunk Backward (DV)	examples/kda/chunk_bwd_dv.py
Chunk Backward (GLA dA)	examples/kda/chunk_bwd_gla_dA.py
Chunk Backward Intra	examples/kda/chunk_bwd_intra.py
WY Fast	examples/kda/wy_fast.py
WY Fast Backward	examples/kda/wy_fast_bwd.py
FLA Chunk Delta	examples/kda/FLA_KDA/fla_chunk_delta.py
FLA Chunk Inter	examples/kda/FLA_KDA/fla_chunk_inter.py
FLA Chunk Intra	examples/kda/FLA_KDA/fla_chunk_intra.py
FLA Chunk Intra Token Parallel	examples/kda/FLA_KDA/fla_chunk_intra_token_parallel.py
FLA Chunk O	examples/kda/FLA_KDA/fla_chunk_o.py
FLA WY Fast	examples/kda/FLA_KDA/fla_wy_fast.py
DSA Sparse MLA Forward	examples/dsa_sparse_finetune/sparse_mla_fwd.py
DSA Sparse MLA Backward	examples/dsa_sparse_finetune/sparse_mla_bwd.py
DSA Indexer TopK ReduceSum	examples/dsa_sparse_finetune/indexer_topk_reducesum.py
DSA Indexer Backward	examples/dsa_sparse_finetune/indexer_bwd.py
DSA Sparse MLA TopK ReduceSum	examples/dsa_sparse_finetune/sparse_mla_topk_reducesum.py

wch0810 · 2026-08-27T06:19:25Z

DSA Sparse MLA Forward examples/dsa_sparse_finetune/sparse_mla_fwd.py B=1,S=4096,H=128,HKV=1,DQK=576,DV=512,topk=1024，fp16，性能12144us
DSA Sparse MLA Backward examples/dsa_sparse_finetune/sparse_mla_bwd.py B=1,S=2048,H=64,HKV=1,DQKV=576,DV=512,topk=512，fp16，性能3840us
DSA Indexer TopK ReduceSum examples/dsa_sparse_finetune/indexer_topk_reducesum.py B=1,S=2048,H=64,D=128,topk=64，fp16，性能714us
DSA Indexer Backward examples/dsa_sparse_finetune/indexer_bwd.py B=1,S=2048,H=16,D=128,topk=64，fp16，性能40us
DSA Sparse MLA TopK ReduceSum examples/dsa_sparse_finetune/sparse_mla_topk_reducesum.py B=1,S=2048,H=16,D=512,tail_D=64,topk=128，fp16，性能110us

wch0810 · 2026-08-27T06:22:26Z

B=1,S=4096,SKV=8192,H=128,HKV=1,DQK=576,DV=512,topk=2048，fp16，性能24288us
B=1,S=4096,SKV=8192,H=64,HKV=1,DQKV=576,DV=512,topk=2048，fp16，性能30715us
B=1,S=4096,SKV=8192,H=128,HKV=1,DQK=576,DV=512,topk=2048，fp16，性能22560us
B=2,S=4096,SKV=8192,H=128,HKV=1,DQK=576,DV=512,topk=2048，fp16，性能42100us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，fp16，性能1300us
B=1,S=32768,H=32,DK=128,DV=128,chunk_size=64，fp16，性能810us
B=1,S=32768,H=32,DK=128,DV=128,chunk_size=64，fp16，性能750us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，fp16，性能550us
B=1,S=32768,H=32,DK=128,chunk_size=64，fp16，性能210us
B=1,S=32768,H=32,DK=128,DV=128,chunk_size=64，fp16，性能370us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，fp16，性能400us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能700us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能570us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16，fp16，性能190us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16，fp16，性能185us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能460us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp32，性能1700us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp32，性能200us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能125us
B=1,S=8192,H=8,DK=128,chunk_size=64，fp32，性能80us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能185us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，fp32，性能500us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能1270us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp32，性能1700us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16，fp16，性能270us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16，fp16，性能185us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能660us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能685us
B=1,S=4096,H=128,HKV=1,DQK=576,DV=512,topk=1024，fp16，性能12144us
B=1,S=2048,H=64,HKV=1,DQKV=576,DV=512,topk=512，fp16，性能3840us
B=1,S=2048,H=64,D=128,topk=64，fp16，性能714us
B=1,S=2048,H=16,D=128,topk=64，fp16，性能40us
B=1,S=2048,H=16,D=512,tail_D=64,topk=128，fp16，性能110us

wch0810 · 2026-08-31T01:17:49Z

B=1,S=4096,SKV=8192,H=128,HKV=1,DQK=576,DV=512,topk=2048，bf16，性能 24288us
B=1,S=4096,SKV=8192,H=64,HKV=1,DQKV=576,DV=512,topk=2048，bf16，性能 30715us
B=1,S=4096,SKV=8192,H=128,HKV=1,DQK=576,DV=512,topk=2048，bf16，性能 22560us
B=2,S=4096,SKV=8192,H=128,HKV=1,DQK=576,DV=512,topk=2048，bf16，性能 42100us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，bf16，性能 1300us
B=1,S=32768,H=32,DK=128,DV=128,chunk_size=64，bf16，性能 810us
B=1,S=32768,H=32,DK=128,DV=128,chunk_size=64，bf16，性能 750us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，bf16，性能 550us
B=1,S=32768,H=32,DK=128,chunk_size=64，bf16，性能 210us
B=1,S=32768,H=32,DK=128,DV=128,chunk_size=64，bf16，性能 370us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，bf16，性能 400us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，bf16，性能 700us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp16，性能 570us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16，bf16，性能 190us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16，bf16，性能 185us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，bf16，性能 460us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp32，性能 1700us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp32，性能 200us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，bf16，性能 125us
B=1,S=8192,H=8,DK=128,chunk_size=64，fp32，性能 80us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，bf16，性能 185us
B=1,S=32768,H=8,DK=128,DV=128,chunk_size=64，fp32，性能 500us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fwd:fp16,bwd:bf16，性能 1270us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，fp32，性能 1700us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16（fwd）+B=1,S=8192,H=8（bwd），fwd:bf16,bwd:fp32，性能 270us
B=1,S=8192,H=64,DK=128,chunk_size=64,sub_chunk_size=16，bf16，性能 185us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64，bf16，性能 660us
B=1,S=8192,H=64,DK=128,DV=128,chunk_size=64（fwd）+B=1,S=32768,H=8（bwd），fwd:bf16,bwd:fp32，性能 685us
B=1,S=4096,H=128,HKV=1,DQK=576,DV=512,topk=1024，bf16，性能 12144us
B=1,S=2048,H=64,HKV=1,DQKV=576,DV=512,topk=512，bf16，性能 3840us
B=1,S=2048,H=64,D=128,topk=64，bf16，性能 714us
B=1,S=2048,H=16,D=128,topk=64，bf16，性能 40us
B=1,S=2048,H=16,D=512,tail_D=64,topk=128，bf16，性能 110us

…kernel) (tile-ai#1647) * [Example] Add MLA Decode Persistent kernel with robustness asserts and num_split>1 test coverage - Developer mode (AUTO_CV_COMBINE/SYNC/AUTO_SYNC/MEMORY_PLANNING=True), zero T.Scope/flag/barrier_all - 36 input validation asserts: dtype/device/ndim/shape/value-range/phase2-heads%4/block_H-even - Fractal 16 alignment asserts (source: src/tl_templates/ascend/common.h:1243 roundUp16) - L0B/L0A ping-pong slot budget asserts (source: common.h:1226, block_N*128*2 <= kL0Budget) - Contiguous check for kv/k_pe before zero-copy view - 2 num_split>1 L0 test cases (two-phase path first tested, ratio=0.9998/1.0000) - Fix run_bench core_num hardcoded 20 -> dynamic _get_core_num() - 23/23 test cases PASS (L0:7 + L1:10 + L2:2 + Boundary:4), golden ratio=1.0000 max_abs=1.221e-04 - Precision standard: fp16 atol=2^-14, rtol=2^-9, max_abs=1e-1, required_ratio=0.99 - ruff format + ruff check passed (line-length=140) * [Refactor] Reduce to single kernel with fp32 workspace_3 for unified dispatch - Reduce 4 kernels to 1: flashattn_mla_decode (merge #2/tile-ai#3, delete #1/tile-ai#4) - workspace_3 dtype: fp16 -> fp32 (unified for all dim/block_N configs, no host dispatch) - Remove num_split>1 two-phase path (Phase 1 + Phase 2 deleted) - Remove ws3_fp16 host dispatch branching (single kernel for all configs) - run_mla_decode signature: drop num_split parameter - Test: remove 2 num_split>1 L0 cases (23 -> 21 cases) - Precision: 21/21 PASS (golden ratio=1.0000 max_abs=1.221e-04, unchanged) - Performance: 3877us (gap 1.28x, +6.4% vs 4-kernel 3644us, fp32 GM traffic) - ruff format + ruff check passed (line-length=140) * [Refactor] CI cleanup: single-case smoke test, golden in test file, remove bench/msprof - example file: simplify __main__ to single-case smoke test (CI requirement) - example file: move ref_mla_decode golden to test file - example file: remove unused imports (torch.nn.functional, einops) - example file: convert 6 Chinese comments to English - test file: add ref_mla_decode golden function - test file: remove run_bench, run_msprof, _MSPROF_SCRIPT_TEMPLATE - test file: remove bench/msprof from --level choices - test file: precision tests only (L0/L1/L2/Boundary/all) - Precision: 21/21 PASS (golden ratio=1.0000 max_abs=1.221e-04, unchanged) - ruff format + ruff check passed (line-length=140) - Files: 1138 -> 987 lines (-151) * [Refactor] Move to examples_experiment/deepseek_mla/example_mla_decode_persistent - Move example_mla_decode_persistent.py and test_mla_decode_persistent.py - Path: examples/ -> examples_experiment/deepseek_mla/ - No code changes (pure rename) * [Fix] L2 test gating + cleanup naming remnants - test_l2() returns bool, merged into blocking exit code (was: non-blocking, ignored by main) - Assert specific exception types: AssertionError for dtype/shape (was: catch-all Exception) - BOUNDARY_WARN -> BOUNDARY_FAIL for silently-accepted illegal inputs - Fix docstring: L0/L1/L2 block (was: L0/L1 block only) - Rename phase1_mod -> kernel_mod (no phase1 concept in single-kernel design) - Rename l0_num_split_1 -> l0_default_config (num_split removed) - Remove D-PARAM-num_split tag (num_split no longer a parameter) - Update L0/L1 config comments (remove num_split reference) - Precision: 21/21 PASS (golden ratio=1.0000 max_abs=1.221e-04, unchanged) - ruff format + ruff check passed

wch0810 force-pushed the feat/gqa_sink_bwd_bhsd_robustness_rev4 branch from 3b95f45 to efa27b4 Compare August 20, 2026 03:02

wch0810 added 5 commits August 21, 2026 22:01

[Test] Fix ruff F841/F541 in l2_shape_n192 case

86f00e3

- F841: remove unused dO tensor (case only runs fwd to verify N=192 is valid) - F541: remove extraneous f-string prefix (no placeholders) ruff check + format: all pass.

Navigation Menu

[Example] Add GQA Sink Backward BHSD with robustness improvements - #1

[Example] Add GQA Sink Backward BHSD with robustness improvements#1
wch0810 wants to merge 6 commits into
ascendc_ptofrom
feat/gqa_sink_bwd_bhsd_robustness_rev4

wch0810 commented Aug 20, 2026

Uh oh!

wch0810 commented Aug 25, 2026

Uh oh!

wch0810 commented Aug 27, 2026

Uh oh!

wch0810 commented Aug 27, 2026

Uh oh!

wch0810 commented Aug 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

wch0810 commented Aug 20, 2026

Summary

Files Added

Robustness Improvements

msprof Bug Fixes

Verification Results

Test Commands

Known Limitations (non-blocking)

Uh oh!

wch0810 commented Aug 25, 2026

Uh oh!

wch0810 commented Aug 27, 2026

Uh oh!

wch0810 commented Aug 27, 2026

Uh oh!

wch0810 commented Aug 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant