Skip to content

fix(jit): separate TGV SM100 target module identities - #6

Draft
heiheiha798 wants to merge 4 commits into
mainfrom
b300/tgv-sm100-module-identity
Draft

heiheiha798 wants to merge 4 commits into
mainfrom
b300/tgv-sm100-module-identity

Conversation

@heiheiha798

@heiheiha798 heiheiha798 commented Aug 20, 2026

Copy link
Copy Markdown
Owner

📌 Description

This change gives the TGV BF16/FP16 SM100a and SM100f builds distinct JIT/AOT identities:

  • tgv_gemm_{bf16|fp16}_sm100a
  • tgv_gemm_{bf16|fp16}_sm100f

The generated-source directories carry the same target suffix. Kernel sources, arithmetic, architecture flags, public API, runner behavior, autotune keys, and tactic space are unchanged.

The collision matters on B300 because tgv_gemm_sm100 requests the SM100f family, while multi-architecture AOT generation asks for SM100a first and deduplicates specs by name. On the baseline, the retained same-named SM100a image shadows the requested image and fails on CC 10.3 with cudaErrorNoKernelImageForDevice.

The focused regression test covers all four dtype/target specs, their target-specific generated paths and exact architecture flags, and the real AOT inventory after first-wins name deduplication. The existing trace correctness test now includes SM103 after the rebuilt SM100f image passed on B300.

No performance improvement is claimed. The candidate compiles the same SM100f TGV sources and is parity with the baseline JIT path under the bounded regression gate.

🔍 Related Issues

Validation

Baseline: 5366177a074e27df7db527f5b744c77dfd748484

Candidate: 2438fa19d55c31380c7459898e0a0fe6d4b487c3

Environment: NVIDIA B300 SXM6 AC (CC 10.3, 148 SMs), driver 580.126.09, CUDA/nvcc 13.0/13.0.88, Python 3.12.3, PyTorch 2.9.1+cu130, FlashInfer 0.6.18. GPU compilation and validation ran through Slurm.

Baseline Job 4871 reproduced the collision in two fresh JIT-disabled processes: both selected the staged SM100a artifact and terminated with CUDA error 209 (no kernel image is available for execution on the device). The wrapper itself ended nonzero because its original classifier expected a catchable Python status instead of the native CUDA 209 process exit; the two raw process exits are the gate evidence.

Candidate Job 4875 completed successfully:

  • focused ruff, format, and Python syntax checks passed;
  • the focused JIT/AOT identity inventory test passed;
  • all four target artifacts had distinct names and the expected SM100a/SM100f flags;
  • the JIT-disabled B300 AOT matrix passed 16/16 across BF16/FP16, bias, PDL, and caller-output variants;
  • tests/gemm/test_tgv_gemm.py passed 90/90;
  • the SM103 trace correctness test passed 2/2;
  • fresh-cache BF16 and FP16 JIT cases passed with cosine similarities 0.99609375 and 0.99951171875.

The current head adds only two review-requested test corrections after Job 4875. At the current head, the focused JIT-spec regression passed in a no-device CPU process, and focused ruff, format, Python syntax, pre-commit, and git diff --check checks passed.

Performance regression check

This is an availability fix, not a speedup. Job 4877 ran the unchanged 26-case BF16 TGV benchmark suite in one B300 allocation using isolated caches and paired process order. Twenty-four rows were below the repeatable 2% threshold. Two very short rows triggered a bounded diagnostic because each initially showed the same 0.0002016 ms median delta.

Job 4926 then used a single neutral exact-bucket tune, fixed tactics, ten fresh timed processes, balanced B,C / C,B / B,C / C,B / B,C ordering, CUDA graphs, 30 aggregate warmups, and 100 observations of 100 launches each:

Shape Baseline p10/p50/p90 (ms) Candidate p10/p50/p90 (ms) Paired ratio p10/p50/p90 Pairs >2% Result
M32 N7168 K2048, no bias, tactic 8 0.00476512 / 0.00476832 / 0.00477021 0.00476426 / 0.00476880 / 0.00477123 0.99907 / 1.00007 / 1.00105 0/5 parity
M1 N2880 K1024, bias, tactic 0 0.00323430 / 0.00323472 / 0.00323914 0.00323478 / 0.00323808 / 0.00324096 0.99947 / 0.99995 / 1.00190 0/5 parity

CUPTI was unavailable, and the benchmark helper could not provide closure-driven cold-L2 timing because the closure exposed no tensor arguments. The diagnostic is therefore warm-L2 CUDA-event evidence and is used only to resolve the two short-kernel signals, not as an absolute-latency replacement for the unchanged suite.

🚀 Pull Request Checklist

✅ Pre-commit Checks

  • I have installed pre-commit and its hooks in the source environment.
  • Full pre-commit run --all-files was not run in the approved environment.
  • Focused pre-commit passed for all three changed Python files (file hygiene, mypy, ruff check, and ruff format).
  • git diff --check passed against the current upstream base.

🧪 Tests

  • Tests were added or updated for module identity, AOT inventory, and SM103 trace correctness.
  • Focused tests and the B300 AOT/JIT correctness matrix passed as described above.
  • The whole repository test suite was not run.

Reviewer Notes

  • Renaming deliberately invalidates old name-addressed TGV caches; affected users may see a one-time JIT rebuild.
  • Multi-architecture AOT packages retain two additional SM100f TGV modules that were previously dropped, modestly increasing build time and package size.
  • Physical GPU validation used one 148-SM B300/x86_64 node. SM100a generation and flags were covered statically and by artifact build, but no physical B200 runtime was available.
  • The current head's final two commits only adjust tests; the runtime implementation is the one validated in Job 4875.
  • CUPTI, closure-driven cold-L2 timing for the unchanged suite, full pre-commit, and the whole-repository test suite remain unavailable or unrun as disclosed above.

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