fix(moe): reject unresolved capture launches - #108
Closed
voipmonitor wants to merge 5 commits into
Closed
Conversation
Contributor
📝 WalkthroughWalkthroughChangesCompiler and W4A16 execution
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant W4A16MoE
participant CaptureDetection
participant FusedCompiler
participant LaunchCache
W4A16MoE->>CaptureDetection: Check selected stream capture
CaptureDetection-->>W4A16MoE: Return capture status
W4A16MoE->>FusedCompiler: Compile with cache requirement
FusedCompiler->>LaunchCache: Resolve cached launch
LaunchCache-->>FusedCompiler: Return launch or cache miss
FusedCompiler-->>W4A16MoE: Return launch or capture error
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@sparkinfer/_lib/compiler.py`:
- Around line 2651-2659: Preserve the existing raise_if_kernel_resolution_frozen
check before cute.compile, and add a final check immediately before
_call_cute_compile(). Also perform the check before loading the cached object in
both disk-cache paths, covering _compile_disk_cache_payload and
_load_cute_compile_from_disk so no compile or cache-load operation proceeds
after kernel resolution is frozen.
In `@sparkinfer/moe/_shared/kernels/dynamic.py`:
- Around line 825-829: Update the w4a8_split_materialized selection and launch
flow to require a supported gated activation, such as self.is_gated or
activation in {"silu", "situ"}. Ensure non-gated activations like relu2 cannot
enter self.materialized_phase1_kernel or W4A8MaterializedPhase1Kernel, while
preserving the existing valid monolithic path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 635e579d-34b8-41d5-891b-24bbde74703f
📒 Files selected for processing (8)
sparkinfer/_lib/compiler.pysparkinfer/moe/_shared/kernels/dynamic.pysparkinfer/moe/_shared/kernels/w4a16/host.pysparkinfer/moe/_shared/kernels/w4a16/kernel.pysparkinfer/moe/_shared/kernels/w4a16/route_pack.pysparkinfer/moe/fused_moe/_impl.pytests/_lib/test_compile_cache.pytests/moe/test_w4a16_e2e.py
Comment on lines
+2651
to
+2659
| from sparkinfer._lib.runtime_control import ( | ||
| raise_if_kernel_resolution_frozen, | ||
| ) | ||
|
|
||
| raise_if_kernel_resolution_frozen( | ||
| "cute.compile", | ||
| target=func, | ||
| cache_key=compile_spec if compile_spec is not None else memory_cache_key, | ||
| ) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Locate compiler.py and related runtime_control symbols"
git ls-files | rg '(^|/)compiler\.py$|runtime_control' || true
echo
echo "Relevant sections in sparkinfer/_lib/compiler.py"
sed -n '2600,2795p' sparkinfer/_ib/compiler.py 2>/dev/null || sed -n '2600,2795p' sparkinfer/_lib/compiler.py
echo
echo "Search for raise_if_kernel_resolution_frozen and freeze_kernel_resolution"
rg -n "raise_if_kernel_resolution_frozen|freeze_kernel_resolution|kernel_resolution_frozen" .Repository: local-inference-lab/sparkinfer
Length of output: 18869
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "runtime_control implementation"
cat -n sparkinfer/_lib/runtime_control.py
echo
echo "compile cache tests around freeze checks"
sed -n '280,365p' tests/_lib/test_compile_cache.py
echo
echo "nearby _compile_disk_cache_payload / _load_cute_compile_from_disk definitions"
rg -n "def _compile_disk_cache_payload|def _load_cute_compile_from_disk|def _call_cute_compile|def _spec_memo" sparkinfer/_lib/compiler.py
sed -n '2350,2600p' sparkinfer/_lib/compiler.py | cat -nRepository: local-inference-lab/sparkinfer
Length of output: 17375
Keep the freeze check right before cute.compile.
raise_if_kernel_resolution_frozen() acquires its lock and returns, then compile() can call _compile_disk_cache_payload(), _load_cute_compile_from_disk(), or _call_cute_compile() while another thread invokes freeze_kernel_resolution(). Add a final check immediately before _call_cute_compile() and load the cached object in both disk-cache paths to maintain the frozen-state guarantee.
🤖 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/_lib/compiler.py` around lines 2651 - 2659, Preserve the existing
raise_if_kernel_resolution_frozen check before cute.compile, and add a final
check immediately before _call_cute_compile(). Also perform the check before
loading the cached object in both disk-cache paths, covering
_compile_disk_cache_payload and _load_cute_compile_from_disk so no compile or
cache-load operation proceeds after kernel resolution is frozen.
Comment on lines
+825
to
+829
| # This helper is gated-only and is never launched unless the split | ||
| # materialized path is active. Use a valid inert specialization for | ||
| # non-split activations (notably ReLU2) instead of rejecting them | ||
| # during otherwise valid monolithic-kernel construction. | ||
| activation=self.activation if self.w4a8_split_materialized else "silu", |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 12 \
'class W4A8MaterializedPhase1Kernel|w4a8_split_materialized|external_materialized_fc1|materialized_phase1_kernel|is_gated' \
sparkinfer/moe/_shared/kernels/dynamic.pyRepository: local-inference-lab/sparkinfer
Length of output: 41771
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -a 'w4a8_phase1\.(py|cu|cuh)$' . | sed 's#^\./##'Repository: local-inference-lab/sparkinfer
Length of output: 233
🏁 Script executed:
#!/bin/bash
set -euo pipefail
path="sparkinfer/moe/_shared/kernels/w4a8_phase1.py"
wc -l "$path"
ast-grep outline "$path" --view expanded || true
sed -n '1,260p' "$path" | cat -nRepository: local-inference-lab/sparkinfer
Length of output: 12497
🏁 Script executed:
#!/bin/bash
set -euo pipefail
path="sparkinfer/moe/_shared/kernels/dynamic.py"
python3 - <<'PY'
import ast
from pathlib import Path
p = Path("sparkinfer/moe/_shared/kernels/dynamic.py")
tree = ast.parse(p.read_text())
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == "DynamicMLERepeatPathKernel":
constructor = None
for item in node.body:
if isinstance(item, ast.FunctionDef) and item.name == "__init__":
constructor = item
break
print("constructor_found", bool(constructor))
if constructor:
for stmt in constructor.body:
if isinstance(stmt, ast.Assign):
targets = [t.id for t in stmt.targets if isinstance(t, ast.Name)]
if "w4a8_split_materialized" in targets or "materialized_phase1_kernel" in targets:
print(ast.get_source_segment(p, stmt))
print("\nactivation helper contract:")
activations = Path("sparkinfer/moe/_shared/kernels/activations.py")
act_tree = ast.parse(activations.read_text())
for node in act_tree.body:
if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in {"is_gated_moe_activation", "normalize_moe_activation"}:
print(ast.get_source_segment(activations, node))
PY
echo
printf '--- w4a8_phase1 activation constructor gate ---\n'
sed -n '81,110p' sparkinfer/moe/_shared/kernels/w4a8_phase1.py | cat -n
echo
printf '--- dynamic split-materialized condition ---\n'
sed -n '791,830p' sparkinfer/moe/_shared/kernels/dynamic.py | cat -nRepository: local-inference-lab/sparkinfer
Length of output: 287
🏁 Script executed:
#!/bin/bash
set -euo pipefail
path="sparkinfer/moe/_shared/kernels/dynamic.py"
python3 - <<'PY'
import ast
from pathlib import Path
tree = ast.parse(Path(path).read_text())
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == "DynamicMLERepeatPathKernel":
constructor = None
for item in node.body:
if isinstance(item, ast.FunctionDef) and item.name == "__init__":
constructor = item
break
print("constructor_found", bool(constructor))
if constructor:
for stmt in constructor.body:
if isinstance(stmt, ast.Assign):
targets = [t.id for t in stmt.targets if isinstance(t, ast.Name)]
if "w4a8_split_materialized" in targets or "materialized_phase1_kernel" in targets:
print(ast.get_source_segment(Path(path), stmt))
print("\nactivation helper contract:")
for act in ("activations.py", "w4a8_phase1.py"):
p = Path(filename := f"sparkinfer/moe/_shared/kernels/{act}")
act_tree = ast.parse(p.read_text())
for node in act_tree.body:
if isinstance(node, (ast.FunctionDef, ast.ClassDef)):
if node.name == "is_gated_moe_activation":
print(f"\n--- {act}: {node.name} ---")
print(ast.get_source_segment(p, node))
elif act == "w4a8_phase1.py" and node.name == "W4A8MaterializedPhase1Kernel":
print(f"\n--- {act}: {node.name}.activation validation ---")
for child in node.body:
if isinstance(child, ast.FunctionDef) and child.name == "__init__":
print(ast.get_source_segment(p, child))
PY
printf '\n--- dynamic split-materialized condition ---\n'
sed -n '791,830p' sparkinfer/moe/_shared/kernels/dynamic.py | cat -nRepository: local-inference-lab/sparkinfer
Length of output: 279
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("sparkinfer/moe/_shared/kernels/dynamic.py")
tree = ast.parse(path.read_text())
constructor = None
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == "DynamicMLERepeatPathKernel":
for item in node.body:
if isinstance(item, ast.FunctionDef) and item.name == "__init__":
constructor = item
break
if constructor:
break
print("constructor_found", bool(constructor))
if constructor:
for stmt in constructor.body:
if isinstance(stmt, ast.Assign):
targets = [t.id for t in stmt.targets if isinstance(t, ast.Name)]
if "w4a8_split_materialized" in targets or "materialized_phase1_kernel" in targets:
print(ast.get_source_segment(path, stmt))
print("\nactivation helper contract:")
act_path = Path("sparkinfer/moe/_shared/kernels/activations.py")
act_tree = ast.parse(act_path.read_text())
for node in act_tree.body:
if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name == "is_gated_moe_activation":
print(f"\n--- activations.py: {node.name} ---")
print(ast.get_source_segment(act_path, node))
phase1_path = Path("sparkinfer/moe/_shared/kernels/w4a8_phase1.py")
phase1_tree = ast.parse(phase1_path.read_text())
for node in phase1_tree.body:
if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name == "W4A8MaterializedPhase1Kernel":
print(f"\n--- w4a8_phase1.py: {node.name}.activation validation ---")
for child in node.body:
if isinstance(child, ast.FunctionDef) and child.name == "__init__":
print(ast.get_source_segment(phase1_path, child))
print("\n--- dynamic split-materialized condition ---\n")
import sys
for i, line in enumerate(phase1_path.read_text().splitlines()[:0] or [], 1):
pass
print(f"python ok")
PY
printf '\n--- dynamic split-materialized condition ---\n'
sed -n '791,830p' sparkinfer/moe/_shared/kernels/dynamic.py | cat -nRepository: local-inference-lab/sparkinfer
Length of output: 601
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,80p' sparkinfer/moe/_shared/kernels/activations.py | cat -n
echo
echo '--- w4a8_phase1.py activation validation ---'
sed -n '81,107p' sparkinfer/moe/_shared/kernels/w4a8_phase1.py | cat -n
echo
echo '--- dynamic split materialization condition ---'
sed -n '791,830p' sparkinfer/moe/_shared/kernels/dynamic.py | cat -n
echo
echo '--- search split materialized path launch ---'
rg -n -C 10 'cutlass.const_expr\(self.external_materialized_fc1\)|self.materialized_phase1_kernel\(|w4a8_split_materialized' sparkinfer/moe/_shared/kernels/dynamic.pyRepository: local-inference-lab/sparkinfer
Length of output: 15730
Reject non-gated activations for split materialized FC1.
w4a8_split_materialized is enabled from w4a8_repacked + materialize_intermediate + M64/M128, without checking self.is_gated. At launch, self.materialized_phase1_kernel(...) is called, and W4A8MaterializedPhase1Kernel itself rejects activations outside {"silu", "situ"}, so activation="relu2" can fail with invalid materialized phase-1 inputs. Gate the split-materialized path by activation or add the non-gated phase-1 contract.
🤖 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/dynamic.py` around lines 825 - 829, Update the
w4a8_split_materialized selection and launch flow to require a supported gated
activation, such as self.is_gated or activation in {"silu", "situ"}. Ensure
non-gated activations like relu2 cannot enter self.materialized_phase1_kernel or
W4A8MaterializedPhase1Kernel, while preserving the existing valid monolithic
path.
Contributor
Author
|
Superseded by #110. Although #108 and #109 were individually mergeable against |
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.
Summary
Reject unresolved W4A16 launch resolution before compilation, persistent-cache lookup, filesystem/module loading, or CUDA module loading when the selected launch stream is being captured.
This is the current-
mastersuccessor to #102. It contains only the capture-resolution commits after the obsolete stacked base.Change
cuStreamIsCapturing, while keeping the current-stream fast path.The failure is early and actionable: the operator must run eager warmup for the missing token count before capture. Non-capture behavior, launch ABI, kernel selection, and numerics are unchanged.
Verification
Focused suites cover memory hits, disk-hit rejection, selected-stream capture, generic capacity buckets, exact Trellis capacity, and capture-time planned prewarm. Changed-file Ruff checks and
git diff --checkpass.No merge is performed by this PR creation.
Summary by CodeRabbit
New Features
Bug Fixes