pypto-gym/tests/ops/gemma4_31b_it · CANN/pypto-gym

cann-robotfix(Operator):Adjust pypto tensor directory

文件	最后提交记录	最后更新时间
attn_softmax_golden.py	开源整改--静态扫描 Co-authored-by: s00454010<suntao37@huawei.com> # message auto-generated for no-merge-commit merge: !197 merge master into master 开源整改--静态扫描 Created-by: u010986451 Commit-by: s00454010 Merged-by: cann-robot Description: 开源整改--静态扫描 See merge request: cann/pypto-gym!197	1 个月前
bench_gqa_decode_attn.py	fix(Operator):Adjust pypto tensor directory Co-authored-by: huangyuqian<huangyuqian2@huawei.com> # message auto-generated for no-merge-commit merge: !321 merge master into master fix(Operator):Adjust pypto tensor directory Created-by: huangyuqian Commit-by: huangyuqian Merged-by: cann-robot Description: ## 变更描述 / Description <!-- 本 PR 做了什么，为什么需要 / What does this PR do and why --> ## 改动类型 / Change Type - [ ] Bug 修复 / Bug Fix - [ ] 新功能 / New Feature - [ ] 性能优化 / Performance - [ ] 代码重构 / Refactoring - [ ] 文档更新 / Documentation - [ ] 测试相关 / Test - [ ] 其它 / Other ## 关联 Issue / Related Issues <!-- Closes #000 可自动关闭 / Closes #000 to auto-close --> - Closes # - References # ## 测试信息 / Testing <!-- 简要测试说明或关键结果 / Brief test description or key results --> - [ ] 单元测试通过 / UT passed - [ ] 集成测试通过 / ST passed - [ ] 人工验证通过 / Manual verified ## 检查清单 / Checklist - [ ] 代码符合规范 / Code follows style guide - [ ] 测试添加并通过 / Tests added and passed - [ ] 文档已更新 / Docs updated if needed - [ ] 无硬编码敏感信息 / No secrets hardcoded - [ ] 提交信息符合规范 / Commit message follows convention See merge request: cann/pypto-gym!321	26 天前
gqa_decode_attn_golden.py	开源整改-LICENSE Co-authored-by: s00454010<suntao37@huawei.com> # message auto-generated for no-merge-commit merge: !179 merge dev_bak into master 开源整改-LICENSE Created-by: u010986451 Commit-by: s00454010 Merged-by: cann-robot Description: ## 变更描述 / Description <!-- 本 PR 做了什么，为什么需要 / What does this PR do and why --> ## 改动类型 / Change Type - [ ] Bug 修复 / Bug Fix - [ ] 新功能 / New Feature - [ ] 性能优化 / Performance - [ ] 代码重构 / Refactoring - [ ] 文档更新 / Documentation - [ ] 测试相关 / Test - [ ] 其它 / Other ## 关联 Issue / Related Issues <!-- Closes #000 可自动关闭 / Closes #000 to auto-close --> - Closes # - References # ## 测试信息 / Testing <!-- 简要测试说明或关键结果 / Brief test description or key results --> - [ ] 单元测试通过 / UT passed - [ ] 集成测试通过 / ST passed - [ ] 人工验证通过 / Manual verified ## 检查清单 / Checklist - [ ] 代码符合规范 / Code follows style guide - [ ] 测试添加并通过 / Tests added and passed - [ ] 文档已更新 / Docs updated if needed - [ ] 无硬编码敏感信息 / No secrets hardcoded - [ ] 提交信息符合规范 / Commit message follows convention See merge request: cann/pypto-gym!179	1 个月前
test_attn_softmax.py	fix(Operator):Adjust pypto tensor directory Co-authored-by: huangyuqian<huangyuqian2@huawei.com> # message auto-generated for no-merge-commit merge: !321 merge master into master fix(Operator):Adjust pypto tensor directory Created-by: huangyuqian Commit-by: huangyuqian Merged-by: cann-robot Description: ## 变更描述 / Description <!-- 本 PR 做了什么，为什么需要 / What does this PR do and why --> ## 改动类型 / Change Type - [ ] Bug 修复 / Bug Fix - [ ] 新功能 / New Feature - [ ] 性能优化 / Performance - [ ] 代码重构 / Refactoring - [ ] 文档更新 / Documentation - [ ] 测试相关 / Test - [ ] 其它 / Other ## 关联 Issue / Related Issues <!-- Closes #000 可自动关闭 / Closes #000 to auto-close --> - Closes # - References # ## 测试信息 / Testing <!-- 简要测试说明或关键结果 / Brief test description or key results --> - [ ] 单元测试通过 / UT passed - [ ] 集成测试通过 / ST passed - [ ] 人工验证通过 / Manual verified ## 检查清单 / Checklist - [ ] 代码符合规范 / Code follows style guide - [ ] 测试添加并通过 / Tests added and passed - [ ] 文档已更新 / Docs updated if needed - [ ] 无硬编码敏感信息 / No secrets hardcoded - [ ] 提交信息符合规范 / Commit message follows convention See merge request: cann/pypto-gym!321	26 天前
test_cases.json	feat(gemma4_31b_it): add Gemma-4-31B-it PyPTO attention kernels Co-authored-by: leedongkun30-arch<lee.dongkun30@gmail.com> # message auto-generated for no-merge-commit merge: !81 merge feature/gemma4-31b-it into master feat(gemma4_31b_it): add Gemma-4-31B-it PyPTO attention kernels Created-by: leedongkun30-arch Commit-by: leedongkun30-arch Merged-by: cann-robot Description: ## Description Add PyPTO fused attention kernels for Gemma-4-31B-it inference on Ascend 910B: - attn_softmax: 3-pass tiled softmax (max / sum / normalize), S_TILE=64, FP32 internal, bf16 I/O - gqa_decode_attn: GQA decode attention with KV head averaging (16→4 KV heads, 75% bandwidth reduction), online softmax, S2_TILE=64, sliding window support - modeling_gemma4.py: HF source copy with PyPTO kernel hooks - ask_Gemma-4-31B-it.py / bench_Gemma-4-31B-it.sh: E2E inference script and benchmark ### Model architecture \| Parameter \| Value \| \|---\|---\| \| hidden_size \| 5376 \| \| intermediate_size \| 21504 \| \| num_attention_heads / num_key_value_heads \| 32 / 16 (GQA ratio=2) \| \| head_dim \| 256 (sliding) / 512 (global) \| \| num_hidden_layers \| 60 \| \| sliding_window \| 1024 \| ### GQA kernel: KV head averaging The GQA decode kernel averages groups of 4 adjacent KV heads (16→4), reducing KV memory bandwidth by 75%. Query heads are reshaped from `[Nq=32, D=256]` to `[Nkv=4, GROUPS=8, D=256]`. This lossy compression trades minimal accuracy for significant bandwidth savings at large sequence lengths. ### Integration approach The modified `modeling_gemma4.py` checks `sys.modules.get("gemma4_pto_kernels")`. When `USE_PTO_SOFTMAX` or `USE_PTO_GQA` is `True`, the corresponding operations dispatch to PyPTO kernels. The ask script injects the kernel switch via `types.ModuleType` before transformers import. The GQA wrapper receives 16 KV heads from the model and performs head averaging internally. ### E2E benchmark results (Ascend 910B, BF16, 9900 input tokens, 100 output tokens) \| Mode \| Prefill (s) \| Decode Time (s) \| Throughput (tok/s) \| \|---\|---\|---\|---\| \| Baseline \| 6.387 \| 20.635 \| 4.8 \| \| PyPTO \| 6.385 \| 21.337 \| 4.7 \| ### GQA kernel microbenchmark (Lossy Nkv=4, S2_TILE=64) \| Skv \| Lossy (ms) \| Lossless (ms) \| Eager (ms) \| Lossy vs Lossless \| \|---:\|---:\|---:\|---:\|---:\| \| 1,024 \| 0.799 \| 0.690 \| 0.248 \| 0.86x \| \| 2,048 \| 0.943 \| 0.969 \| 0.279 \| 1.03x \| \| 4,096 \| 1.129 \| 1.531 \| 0.666 \| 1.36x \| \| 8,192 \| 1.852 \| 2.696 \| 1.508 \| 1.46x \| \| 16,384 \| 3.421 \| 4.902 \| 3.125 \| 1.43x \| \| 32,768 \| 6.076 \| 9.326 \| 6.163 \| 1.54x \| The lossy kernel achieves up to 1.54x speedup over the lossless kernel at Skv=32768 via 75% KV bandwidth reduction. With S2_TILE=128 (doubling tile size), the kernel also beats eager attention at Skv>=8192 (up to 1.41x). ### Softmax kernel microbenchmark \| S (KV length) \| PyPTO (ms) \| Eager (ms) \| Speedup \| \|---:\|---:\|---:\|---:\| \| 128 \| 0.333 \| 0.079 \| 0.24x \| \| 512 \| 0.355 \| 0.079 \| 0.22x \| \| 1,024 \| 0.412 \| 0.079 \| 0.19x \| \| 4,096 \| 0.919 \| 0.078 \| 0.08x \| Softmax kernel does not outperform eager at current scales due to PyPTO dispatch overhead dominating the small computation. ## Change Type - [x] New Feature - [x] Performance - [x] Test ## Changed Files \| Path \| Description \| \|---\|---\| \| `src/pypto_gym/ops/pypto_tile/gemma4_31b_it/attn_softmax/attn_softmax_impl.py` \| Tiled softmax kernel \| \| `src/pypto_gym/ops/pypto_tile/gemma4_31b_it/gqa_decode_attn/gqa_decode_attn_impl.py` \| GQA decode attention with KV head averaging (16→4) \| \| `src/pypto_gym/ops/pypto_tile/gemma4_31b_it/__init__.py` \| Kernel module init and feature flags \| \| `src/pypto_gym/ops/pypto_tile/gemma4_31b_it/README.md` \| Kernel documentation \| \| `src/pypto_gym/transformers/gemma4_31b_it/modeling_gemma4.py` \| HF source copy with PyPTO hooks \| \| `src/pypto_gym/transformers/gemma4_31b_it/configuration_gemma4.py` \| HF config copy \| \| `modeling/transformers/gemma4_31b_it/ask_Gemma-4-31B-it.py` \| E2E inference script \| \| `modeling/transformers/gemma4_31b_it/bench_Gemma-4-31B-it.sh` \| Benchmark shell script \| \| `modeling/transformers/gemma4_31b_it/bench_baseline.json` \| Baseline E2E benchmark results \| \| `modeling/transformers/gemma4_31b_it/bench_pypto.json` \| PyPTO E2E benchmark results \| \| `modeling/transformers/gemma4_31b_it/README.md` \| Model integration documentation \| \| `tests/ops/gemma4_31b_it/test_attn_softmax.py` \| Softmax precision test \| \| `tests/ops/gemma4_31b_it/test_gqa_decode_attn.py` \| GQA decode attention test (subprocess-isolated) \| \| `tests/ops/gemma4_31b_it/bench_gqa_decode_attn.py` \| GQA kernel microbenchmark \| \| `tests/ops/gemma4_31b_it/gqa_decode_attn_golden.py` \| Lossy GQA golden reference \| \| `tests/ops/gemma4_31b_it/attn_softmax_golden.py` \| Softmax golden reference \| \| `tests/ops/gemma4_31b_it/test_cases.json` \| Test case configurations \| ## Testing - [x] UT passed - `test_attn_softmax.py`: 4 cases (S=1,128,512,1024), max diff = 0.0 - `test_gqa_decode_attn.py`: 5 cases (global/local, Skv=64~4096), max diff <= 4.88e-04 (lossy kernel vs lossy golden) - [x] E2E verified (Ascend 910B, Gemma-4-31B-it, 9900 input tokens, 100 output tokens) `bash source /usr/local/Ascend/ascend-toolkit/set_env.sh export PTO_TILE_LIB_CODE_PATH=/path/to/pto-isa export TILE_FWK_DEVICE_ID=14 python tests/ops/gemma4_31b_it/test_attn_softmax.py python tests/ops/gemma4_31b_it/test_gqa_decode_attn.py` ## Known Limits - GQA decode kernel supports Sq=1 only (decode phase); prefill falls back to eager attention with KV head averaging - Tested on sliding-window layers (D=256) only; global layers (D=512) use eager path - KV head averaging is lossy (groups of 4 KV heads averaged); max diff vs lossless <= 4.88e-04 - pypto JIT state contamination: different DYNAMIC Skv values in the same process may produce incorrect output; tests use subprocess isolation - Multi-NPU via device_map="auto" is supported for E2E inference ## Checklist - [x] Code follows style guide - [x] Tests added and passed - [x] Docs updated - [x] No secrets hardcoded - [x] Commit message follows convention See merge request: cann/pypto-gym!81	1 个月前
test_gqa_decode_attn.py	fix(Operator):Adjust pypto tensor directory Co-authored-by: huangyuqian<huangyuqian2@huawei.com> # message auto-generated for no-merge-commit merge: !321 merge master into master fix(Operator):Adjust pypto tensor directory Created-by: huangyuqian Commit-by: huangyuqian Merged-by: cann-robot Description: ## 变更描述 / Description <!-- 本 PR 做了什么，为什么需要 / What does this PR do and why --> ## 改动类型 / Change Type - [ ] Bug 修复 / Bug Fix - [ ] 新功能 / New Feature - [ ] 性能优化 / Performance - [ ] 代码重构 / Refactoring - [ ] 文档更新 / Documentation - [ ] 测试相关 / Test - [ ] 其它 / Other ## 关联 Issue / Related Issues <!-- Closes #000 可自动关闭 / Closes #000 to auto-close --> - Closes # - References # ## 测试信息 / Testing <!-- 简要测试说明或关键结果 / Brief test description or key results --> - [ ] 单元测试通过 / UT passed - [ ] 集成测试通过 / ST passed - [ ] 人工验证通过 / Manual verified ## 检查清单 / Checklist - [ ] 代码符合规范 / Code follows style guide - [ ] 测试添加并通过 / Tests added and passed - [ ] 文档已更新 / Docs updated if needed - [ ] 无硬编码敏感信息 / No secrets hardcoded - [ ] 提交信息符合规范 / Commit message follows convention See merge request: cann/pypto-gym!321	26 天前