cute ops: probe nvvm.fmax signature instead of sniffing CUDA_VERSION (fixes engine init on PyPI DSL 4.6.0)#70
Merged
lukealonso merged 1 commit intoJul 23, 2026
Conversation
Author
|
Companion vLLM prefill PR measured on this stack: local-inference-lab/vllm#163 |
Contributor
📝 WalkthroughWalkthroughChangesThe PR adds a full-rotation EXL3 Trellis MoE API and W4A16 execution path, including native weight preparation, Trellis dequantization, routing, scratch binding, CUDA graph-safe execution, and GPU correctness tests. It also adds device-aware compile cache formats, updates cache validators, adds Triton, and makes NVVM Trellis MoE
Compile cache format
NVVM compatibility
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant trellis_moe
participant W4A16Kernels
participant ScratchArena
participant Output
Caller->>trellis_moe: prepare_weights
trellis_moe->>W4A16Kernels: plan and compile Trellis kernels
Caller->>ScratchArena: provide scratch storage
trellis_moe->>ScratchArena: bind typed workspace views
Caller->>trellis_moe: run
trellis_moe->>W4A16Kernels: dispatch fused MoE and top-k sum
W4A16Kernels->>Output: write rotated routed output
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Unblocks building the trellis3_t256 stack from local-inference-lab/ b12x#49 (@brandonmusic) on the public PyPI nvidia-cutlass-dsl 4.6.0 wheel, which local-inference-lab#49 requires (>= 4.6.0) but currently cannot run on. The DSL wheels disagree about both the fmax ABI and what cutlass.CUDA_VERSION means: the PyPI 4.6.0 wheel reports its bundled CTK as DSLCudaVersion(12, 9) while exposing the unified nvvm.fmax(a, b, *, c=None, ...) signature, so the existing "major == 12 and minor == 9" branch calls the removed nvvm.fmax(T.f32(), a, b) form and every attention JIT fails at engine init with "fmax() takes 2 positional arguments but 3 ... were given". The CTK-13.3 DSL build used by the published runtime image sidesteps the sniff by reporting 13.3; PyPI users hit it every time. Probe the actual signature once at import (first parameter named "a" or "b" means the unified form) and pick the call shape from that. Verified on PyPI 4.5.2 (legacy result-type form preserved), PyPI 4.6.0 (unified; attention JIT compiles clean), and the CTK-13.3 build (unified; same branch it already takes).
davidsyoung
force-pushed
the
fmax-dsl46-signature-probe
branch
from
July 22, 2026 12:00
5ea5503 to
0611af8
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/_lib/test_compile_cache.py (1)
56-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the v6 explicit-spec payload branch.
The new test exercises the non-explicit (
v3, 9-field)_compile_disk_cache_payloadpath withdevice_archat index 4, but there's no test for the explicit-spec (v6_explicit_spec, 11-field) branch that also shifted indices in this PR (_compile_disk_cache_payloadwithcompile_specset,_compile_cache_payload_log_value,_is_explicit_spec_payload,_build_compile_manifest).Want me to draft a test that constructs a
KernelCompileSpec, calls_compile_disk_cache_payload(..., compile_spec=spec), and assertspayload[0] == "sparkinfer_cute_compile_cache_v6_explicit_spec"withpayload[4]holding the device-arch tuple?🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/_lib/test_compile_cache.py` around lines 56 - 79, Add a test alongside test_disk_cache_key_includes_device_arch that constructs a KernelCompileSpec, calls _compile_disk_cache_payload with compile_spec=spec, and verifies the explicit-spec payload uses the "sparkinfer_cute_compile_cache_v6_explicit_spec" marker at payload[0] and stores the device-architecture tuple at payload[4].sparkinfer/moe/_shared/kernels/w4a16/kernel.py (1)
6494-6528: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffDuplicated
_had128_quadacrossW4A16FusedMoeKernelandW4A16TopKSumKernel.This blockwise-128 Walsh–Hadamard helper (lines 6494–6528) is byte-identical to
W4A16FusedMoeKernel._had128_quad(lines 5540–5578), including the0.088388347648scale. Consider hoisting it to a shared@cute.jitfree function to keep the two rotation paths from drifting.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/moe/_shared/kernels/w4a16/kernel.py` around lines 6494 - 6528, Hoist the duplicated _had128_quad implementation out of W4A16FusedMoeKernel and W4A16TopKSumKernel into one shared `@cute.jit` free function, then update both rotation paths to call it. Preserve the existing butterfly computation, lane handling, return tuple, and 0.088388347648 scaling exactly.sparkinfer/moe/_shared/kernels/w4a16/prepare.py (1)
1564-1565: 📐 Maintainability & Code Quality | 🔵 TrivialPrefer unpacking over tuple concatenation (RUF005).
Minor readability/idiom nit flagged by Ruff.
♻️ Proposed change
- expected_i16_shape = expected_prefix_shape + (16 * trellis_bits,) - expected_i32_shape = expected_prefix_shape + (8 * trellis_bits,) + expected_i16_shape = (*expected_prefix_shape, 16 * trellis_bits) + expected_i32_shape = (*expected_prefix_shape, 8 * trellis_bits)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sparkinfer/moe/_shared/kernels/w4a16/prepare.py` around lines 1564 - 1565, Update the expected_i16_shape and expected_i32_shape assignments to use iterable unpacking for expected_prefix_shape instead of tuple concatenation, preserving the existing trailing dimensions and resulting shapes.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@sparkinfer/moe/_shared/kernels/w4a16/kernel.py`:
- Around line 6494-6528: Hoist the duplicated _had128_quad implementation out of
W4A16FusedMoeKernel and W4A16TopKSumKernel into one shared `@cute.jit` free
function, then update both rotation paths to call it. Preserve the existing
butterfly computation, lane handling, return tuple, and 0.088388347648 scaling
exactly.
In `@sparkinfer/moe/_shared/kernels/w4a16/prepare.py`:
- Around line 1564-1565: Update the expected_i16_shape and expected_i32_shape
assignments to use iterable unpacking for expected_prefix_shape instead of tuple
concatenation, preserving the existing trailing dimensions and resulting shapes.
In `@tests/_lib/test_compile_cache.py`:
- Around line 56-79: Add a test alongside
test_disk_cache_key_includes_device_arch that constructs a KernelCompileSpec,
calls _compile_disk_cache_payload with compile_spec=spec, and verifies the
explicit-spec payload uses the "sparkinfer_cute_compile_cache_v6_explicit_spec"
marker at payload[0] and stores the device-architecture tuple at payload[4].
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 45ab270f-78a1-4aa0-8a42-a14d85aa94fa
📒 Files selected for processing (19)
pyproject.tomlsparkinfer/__init__.pysparkinfer/_lib/compiler.pysparkinfer/_lib/intrinsics.pysparkinfer/attention/_shared/cute/ops.pysparkinfer/moe/__init__.pysparkinfer/moe/_shared/kernels/w4a16/__init__.pysparkinfer/moe/_shared/kernels/w4a16/host.pysparkinfer/moe/_shared/kernels/w4a16/kernel.pysparkinfer/moe/_shared/kernels/w4a16/prepare.pysparkinfer/moe/_shared/kernels/w4a16/route_pack.pysparkinfer/moe/trellis_moe/__init__.pysparkinfer/moe/trellis_moe/_impl.pysparkinfer/moe/trellis_moe/api.pytests/_lib/test_compile_cache.pytests/moe/test_trellis_moe.pyvalidation/cutlass_migration/acceptance/corpus/ptx_capture.pyvalidation/cutlass_migration/acceptance/e2e/contract.pyvalidation/cutlass_migration/evidence/kernel_resources.py
davidsyoung
changed the title
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
2 participants
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.
Standalone fix on
master(single commit, +21/−15 insparkinfer/attention/_shared/cute/ops.py). Also unblocks #49, which requires nvidia-cutlass-dsl ≥ 4.6.0 — the stack currently cannot run on the only public 4.6.0 wheel.The DSL wheels disagree about both the fmax ABI and what
cutlass.CUDA_VERSIONmeans:nvvm.fmax(T.f32(), a, b, …)(a, b, *, c=None, …)TypeError: fmax() takes 2 positional arguments but 3 … were given— every attention JIT fails at engine initSince
fmaxsits under every attention softmax reduction (fmax_reduce, indexer quad-max, paged/extend running max), PyPI-4.6.0 users can't boot at all.Fix: probe
inspect.signature(nvvm.fmax)once at import (first parameter nameda/b⇒ unified form) and pick the call shape from the actual ABI instead of a version proxy. Verified on all three builds above — legacy form preserved on 4.5.2, attention JIT compiles clean on PyPI 4.6.0, and the 13.3 build takes the same branch it already does. No numerical or behavioral change.