fix(jit): separate TGV SM100 target module identities - #6
Draft
heiheiha798 wants to merge 4 commits into
Draft
heiheiha798 wants to merge 4 commits into
heiheiha798 wants to merge 4 commits into
Conversation
7 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
This change gives the TGV BF16/FP16 SM100a and SM100f builds distinct JIT/AOT identities:
tgv_gemm_{bf16|fp16}_sm100atgv_gemm_{bf16|fp16}_sm100fThe 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_sm100requests 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 withcudaErrorNoKernelImageForDevice.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:
5366177a074e27df7db527f5b744c77dfd748484Candidate:
2438fa19d55c31380c7459898e0a0fe6d4b487c3Environment: 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:
tests/gemm/test_tgv_gemm.pypassed 90/90;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 --checkchecks 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,Cordering, CUDA graphs, 30 aggregate warmups, and 100 observations of 100 launches each: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
pre-commitand its hooks in the source environment.pre-commit run --all-fileswas not run in the approved environment.git diff --checkpassed against the current upstream base.🧪 Tests
Reviewer Notes