Skip to content

[TEST] Cover Gluon Blackwell tcgen05 scaled ops#460

Draft
Jokeren wants to merge 1 commit into
mainfrom
test/gluon-blackwell-tcgen05-scaled-ops
Draft

[TEST] Cover Gluon Blackwell tcgen05 scaled ops#460
Jokeren wants to merge 1 commit into
mainfrom
test/gluon-blackwell-tcgen05-scaled-ops

Conversation

@Jokeren

@Jokeren Jokeren commented Jun 20, 2026

Copy link
Copy Markdown
Member

Summary

Add CPU-result coverage for Blackwell tcgen05_mma_scaled simulation through Gluon.

This adds a focused end-to-end test that:

  • stores A/B tiles in shared memory
  • stores uint8 scale groups in tensor memory scale layouts
  • runs blackwell.tcgen05_mma_scaled with e4m3 format constexpr operands
  • checks the simulated CPU result against a PyTorch scaled matmul reference

The branch also teaches the Gluon CPU runner to recover constexpr annotations from wrapped Gluon JIT functions and to wrap string constexpr arguments as tl.constexpr, which current Blackwell scaled MMA builtins require for format operands.

Validation

  • PYTHONPATH=/mnt/keren/triton-viz pytest tests/end_to_end/test_gluon_blackwell_tcgen05_scaled_ops.py -q
  • PYTHONPATH=/mnt/keren/triton-viz pytest tests/end_to_end/test_gluon_blackwell_tcgen05_scaled_ops.py tests/end_to_end/test_gluon_blackwell_tcgen05_ops.py tests/end_to_end/test_gluon_blackwell_tensor_memory_ops.py tests/end_to_end/test_gluon_wgmma_ops.py tests/end_to_end/test_gluon_tma_im2col_ops.py tests/end_to_end/test_gluon_blackwell_tma_ops.py tests/end_to_end/test_gluon_tma_ops.py tests/end_to_end/test_gluon_async_copy_ops.py tests/end_to_end/test_gluon_core_ops.py tests/end_to_end/test_gluon.py tests/end_to_end/test_core.py tests/unit/test_adapters.py -q
  • pre-commit run --files triton_viz/core/simulation/gluon.py tests/end_to_end/test_gluon_blackwell_tcgen05_scaled_ops.py

@Jokeren Jokeren force-pushed the test/gluon-blackwell-tcgen05-ops branch from 6049438 to b1422ef Compare June 26, 2026 14:48
Base automatically changed from test/gluon-blackwell-tcgen05-ops to main June 26, 2026 16:14
@Jokeren Jokeren force-pushed the test/gluon-blackwell-tcgen05-scaled-ops branch from 9f12302 to 59256be Compare June 26, 2026 16:18
@github-actions

github-actions Bot commented Jun 26, 2026

Copy link
Copy Markdown

Performance Benchmark

Benchmark main (min) PR (min) Change Samples
gemm 0.105s 0.106s +0.3% 20 / 20
gemm_oob 0.117s 0.117s +0.1% 20 / 20
indirect_load 0.022s 0.022s -1.3% 20 / 20
nested_loop 0.234s 0.236s +1.0% 20 / 20
block_pointer_loop_advance 0.205s 0.205s -0.2% 20 / 20
liger_jsd 0.137s 0.138s +0.6% 20 / 20
flaggems_layernorm 0.398s 0.399s +0.4% 20 / 20
swiglu 0.172s 0.172s +0.2% 20 / 20
cross_entropy 0.984s 0.990s +0.6% 20 / 20
fused_linear_jsd 0.210s 0.210s +0.2% 20 / 20
Total 2.584s 2.596s +0.4% N/A

Iterations: 1 warmup + 20 measured
Samples are shown as main / PR; long pytest benchmarks may use fewer samples.

@Jokeren Jokeren force-pushed the test/gluon-blackwell-tcgen05-scaled-ops branch from 59256be to b05edb7 Compare June 26, 2026 17:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant