[WIP] test(kernel): verify kernel selection cross-platform and drop expected_kernel_name#659
Open
antiagainst wants to merge 4 commits into
Open
[WIP] test(kernel): verify kernel selection cross-platform and drop expected_kernel_name#659antiagainst wants to merge 4 commits into
antiagainst wants to merge 4 commits into
Conversation
FA3 registers only on Hopper hosts and FA4 only on SM100 hosts, but their declared capabilities had no max_arch_version, so the specs claimed to support every newer NVIDIA arch. Real hosts never noticed because the import guard hides the spec, but any selection run against a platform other than the host (Platform.override, cross-platform golden tests) would let FA3/FA4 outrank kernels that actually support that arch. Declare the ceilings the import guards already imply so capability filtering is truthful on its own. Signed-off-by: Lei Zhang <antiagainst@gmail.com>
Each golden case previously ran only when the host GPU matched its arch predicate, so an NVIDIA CI machine never checked the AMD golden selections and vice versa. Run every case on every host by injecting the case's target platform with Platform.override (the pattern the cdna4 attn_merge_state test already used), parametrized over conftest fixture platforms mapped from the case's arch tag. A kernel whose registration is import-guarded on a missing optional backend package skips; a host that matches the case natively still hard-fails if the kernel is not registered, preserving the old guarantee on matching CI. Also add a b300 (sm103) fixture so the blackwell-sm103 cases actually run (renamed from blackwell-non-sm100: B300 is the only non-sm100 Blackwell part, and it is what the fixture tests), add an mxfp4 dense GEMM golden case (triton_mm_mxfp4 on cdna4) matching the runtime Mxfp4LinearMethod call, and skip the argmax cases without CUDA since argmax only dispatches through selection for CUDA tensors. Signed-off-by: Lei Zhang <antiagainst@gmail.com>
The AST scanner keyed on expected_kernel_name literals at runtime call sites, and its hand-maintained MoE entries duplicated trait combinations that the runtime now builds through tokenspeed_kernel.moe_plan. Every selection it asserted is covered by the cross-platform golden cases in test_kernel_api_selection.py, which exercise the public APIs end to end (signature and trait construction included) instead of calling select_kernel directly: - gemm.mm nvfp4 (cublaslt_mm_nvfp4): blackwell-plus golden case - gemm.mm mxfp4 (triton_mm_mxfp4): cdna4 golden case - the six MoE apply trait combos: the corresponding moe_plan/moe_apply golden cases Signed-off-by: Lei Zhang <antiagainst@gmail.com>
The hint only logged a warning when auto-selection disagreed with a name hard-coded at the call site, and it pointed runtime code in the wrong direction: call sites should rely on kernel auto-selection, with per-arch expectations recorded once in the golden selection tests and explicit pins going through the override mechanism. Debugging needs are covered by TOKENSPEED_KERNEL_VERBOSE and explain_selection. Remove the parameter from select_kernel and mm, the two runtime dense call sites (nvfp4 keeps override for its sanctioned kernel pin), and the nvfp4 fused-GEMM numerics test. Signed-off-by: Lei Zhang <antiagainst@gmail.com>
antiagainst
force-pushed
the
lei/kernel-selection-golden
branch
from
c039609 to
edbd606
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: edbd606d48
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Comment on lines
+1495
to
+1503
| if expected_spec is None: | ||
| # Registrations are import-guarded on optional backend packages, so a | ||
| # missing spec is only a failure when this host should run the case | ||
| # natively. | ||
| assert not case.matches(host_platform), ( | ||
| f"{case.expected!r} is not registered on " | ||
| f"{host_platform.device_name} ({host_platform.arch_version})" | ||
| ) | ||
| pytest.skip(f"{case.expected!r} is not registered (optional backend missing)") |
There was a problem hiding this comment.
On non-native CI hosts, this skip path hides the new cross-platform golden cases instead of exercising them. The autouse _kernel_registry fixture reloads registration modules before this test calls Platform.override(platform), and several in-tree kernels are registered only when current_platform() matches the real host architecture; for example, an H100 host never registers FA4/B200 kernels, so expected_spec is None and the B200 case is skipped because case.matches(host_platform) is false. To get the intended fixture-platform coverage, reload the registrations after overriding to the target platform (or only skip truly missing optional backend packages).
Useful? React with 👍 / 👎.
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
1 participant
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.
expected_kernel_namewas an early-enablement safeguard: call siteshard-coded the expected kernel and
test_runtime_callsite_selection.pychecked those literals against fixture platforms. Move the expectations into
the kernel package's golden selection tests and remove the mechanism.
test_kernel_api_selection.pygolden cases on every host viaPlatform.overridewith fixture platforms, instead of only on matchingGPUs. New
b300_platformfixture activates the never-executedblackwell-non-sm100 cases; missing optional backends skip.
max_arch_versionon FA3/FA4 registrations. Their import guardshid overly-broad capabilities, which cross-platform selection exposed
(FA3 outranked
flashinfer_trtllm_*on B300). No change on real hosts.test_runtime_callsite_selection.py; its coverage is subsumed by thegolden cases, including a new
triton_mm_mxfp4dense GEMM case.expected_kernel_namefromselect_kernel,mm, and call sites.Explicit pins keep using
override=; debugging keepsTOKENSPEED_KERNEL_VERBOSE/explain_selection.