moe: add planned EXL3 Trellis path for fused Blackwell MoE#49
moe: add planned EXL3 Trellis path for fused Blackwell MoE#49brandonmmusic-max wants to merge 9 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds Trellis3.0 decoding and native Trellis256 weight preparation, extends W4A16 GEMM and full-rotation MoE execution, exposes a lazy-loaded operator API, and updates compile-cache payloads and manifest validation for device-aware formats. ChangesTrellis MoE and cache integration
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant TrellisMoEAPI
participant WeightPreparation
participant W4A16Kernels
participant TopKSum
Caller->>TrellisMoEAPI: prepare_weights, plan, bind
TrellisMoEAPI->>WeightPreparation: validate and wrap native Trellis weights
WeightPreparation-->>TrellisMoEAPI: TrellisMoEWeights
TrellisMoEAPI->>W4A16Kernels: run full-rotation fused MoE
W4A16Kernels->>TopKSum: activation output and route metadata
TopKSum-->>TrellisMoEAPI: accumulated output
Possibly related PRs
🚥 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e3ff9087-de31-44d4-ba04-a0a09dbdaff1
📒 Files selected for processing (10)
b12x/cute/compiler.pyb12x/cute/fp4.pyb12x/moe/fused/w4a16/__init__.pyb12x/moe/fused/w4a16/host.pyb12x/moe/fused/w4a16/kernel.pyb12x/moe/fused/w4a16/prepare.pytests/test_cutlass_runtime_patches.pyvalidation/cutlass_migration/acceptance/corpus/ptx_capture.pyvalidation/cutlass_migration/acceptance/e2e/contract.pyvalidation/cutlass_migration/evidence/kernel_resources.py
40f85fc to
1937274
Compare
brandonmmusic-max
changed the title
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
sparkinfer/moe/trellis_moe/api.py (1)
41-52: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
__all__is not isort-sorted (RUF022). Sort to keep lint green.🤖 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/trellis_moe/api.py` around lines 41 - 52, Sort the exported names in the module-level __all__ list alphabetically to satisfy RUF022, preserving every existing export and its spelling.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.
Inline comments:
In `@sparkinfer/moe/_shared/kernels/w4a16/kernel.py`:
- Around line 3481-3525: Change the trellis256 B-source offset calculations in
both branches of the shown loading loop to use Int64 for all offset arithmetic,
including expert, tile, projection, and local terms. Ensure b_src_i32 is widened
consistently before forming the final source address, matching the sibling
weight-loading path and preventing valid large-shape offsets from overflowing
Int32.
In `@sparkinfer/moe/trellis_moe/__init__.py`:
- Line 69: Add triton to the project dependencies in pyproject.toml, matching
the existing dependency declaration conventions. Keep the requires=("triton",)
metadata in the trellis_moe module unchanged so support detection and runtime
requirements remain aligned.
---
Nitpick comments:
In `@sparkinfer/moe/trellis_moe/api.py`:
- Around line 41-52: Sort the exported names in the module-level __all__ list
alphabetically to satisfy RUF022, preserving every existing export and its
spelling.
🪄 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: 619928cf-145f-4ffe-a35a-4ddb9590c87e
📒 Files selected for processing (17)
sparkinfer/__init__.pysparkinfer/_lib/compiler.pysparkinfer/_lib/intrinsics.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
- w4a16 kernel: compute the trellis256 B-source offset in Int64 in both the proj and non-proj branches. The expert-stride term overflows Int32 for large-but-valid MoE shapes (e.g. E=256, H=8192, bits=6 reaches ~3.21e9 > 2^31-1), which would fetch wrong or out-of-bounds B weights. Matches the sibling modelopt-native staging path that already builds this offset with Int64(...). Value-identical for in-range offsets, so existing shapes/tests are unaffected. - pyproject: declare the `triton` dependency. trellis_moe META.requires it and is_supported() gates on it, so a clean install lacking triton would silently skip the op. - trellis_moe.api: sort __all__ (RUF022). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- w4a16 kernel: compute the trellis256 B-source offset in Int64 in both the proj and non-proj branches. The expert-stride term overflows Int32 for large-but-valid MoE shapes (e.g. E=256, H=8192, bits=6 reaches ~3.21e9 > 2^31-1), which would fetch wrong or out-of-bounds B weights. Matches the sibling modelopt-native staging path that already builds this offset with Int64(...). Value-identical for in-range offsets, so existing shapes/tests are unaffected. - pyproject: declare the `triton` dependency. trellis_moe META.requires it and is_supported() gates on it, so a clean install lacking triton would silently skip the op. - trellis_moe.api: sort __all__ (RUF022). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
fd1d7b5 to
879ca0a
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/moe/_shared/kernels/w4a16/host.py`:
- Around line 395-412: Change the full-rotation activation buffers
rotation_a_gate and rotation_a_up to allocate torch.bfloat16, and update the
corresponding full-rotation scratch validation and kernel ABI expectations used
by run_w4a16_moe() to require BF16 consistently. Preserve the existing
conditional allocation behavior.
🪄 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: 6f5a8e43-01e2-4b48-aaec-40ef680de1e3
📒 Files selected for processing (18)
pyproject.tomlsparkinfer/__init__.pysparkinfer/_lib/compiler.pysparkinfer/_lib/intrinsics.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
🚧 Files skipped from review as they are similar to previous changes (13)
- sparkinfer/moe/_shared/kernels/w4a16/init.py
- sparkinfer/init.py
- sparkinfer/moe/init.py
- sparkinfer/moe/trellis_moe/init.py
- validation/cutlass_migration/acceptance/e2e/contract.py
- sparkinfer/moe/trellis_moe/api.py
- validation/cutlass_migration/evidence/kernel_resources.py
- sparkinfer/moe/_shared/kernels/w4a16/route_pack.py
- validation/cutlass_migration/acceptance/corpus/ptx_capture.py
- tests/moe/test_trellis_moe.py
- sparkinfer/moe/trellis_moe/_impl.py
- sparkinfer/_lib/intrinsics.py
- sparkinfer/moe/_shared/kernels/w4a16/kernel.py
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).
|
Great job here man! before it lands, #163 looks like it's stacked on an earlier snapshot oof my pr 139. The three shared files it carries —
— are from before the MTP + tool-call fixes I pushed to #139 today. Concretely, #163's
Could you rebase onto the current #139 head ( |
Builds on the rank-sliced EXL3 (ExLlamaV3 trellis) support added in local-inference-lab#139 by @brandonmusic, paired with the trellis3_t256 planned MoE from local-inference-lab/sparkinfer#49. No change to that work's weight loading, formats, or decode path; this commit only extends its batch dispatch. The vllm-project#139 backend serves every batch with m > 32 through the eager exllamav3_ext parity path: fp16-in/fp32-out staging (~25 GB extra HBM traffic per 4096-token step at GLM-5.2 geometry), python chunk loops, and one expert re-stream per 128-row chunk. That covers 100% of prefill, so prefill throughput is capped at ~2.0-2.3k tok/s while the same weights decode at full speed through the planned Trellis window. Build a second trellis_moe plan at max_tokens=max_num_batched_tokens with block_size_m=64 (env VLLM_EXL3_PREFILL_BLOCK_M, allowed {8,16,32,48,64}) beside the existing decode plan (32/8), and dispatch three ways in _apply_rank_sliced: * m in [min, max] -> decode plan (unchanged) * max < m <= capacity -> prefill plan (sparkinfer bind already accepts any tokens in [1, max_tokens]; zero kernel changes needed) * m < min -> parity path, whose persistent staging (xh/out32/token_sorted/weight_sorted) shrinks from capacity rows to one chunk while the prefill plan is live (~110 MiB/GPU returned) VLLM_EXL3_PREFILL_TRELLIS=0 restores the exact single-plan parity behavior (one-variable serving A/B lever). The parity branch now raises during CUDA graph capture instead of silently recording eager ext calls. The prefill arena (~1.05 GiB/GPU at capacity 3072) allocates during the profile pass, so vLLM's memory profiler sees it and the KV pool auto-shrinks instead of risking request-time OOM. Matched A/B serving brandonmusic/GLM-5.2-EXL3-TR3-3.0bpw on the published verdictai/glm52-exl3-sparkinfer runtime image (4x RTX PRO 6000 Blackwell, TP4/DCP4/MTP3, identical boot geometry, only the kill-switch flipped, fresh JIT cache both arms; llm_decode_bench prefill-only, exact /tokenize, C1, 60 s/context): ctx parity trellis delta 8k 2,287 3,742 +63.6% 32k 2,218 3,557 +60.4% 64k 2,078 3,381 +62.7% 128k 1,920 3,032 +57.9% Decode C1 unchanged-to-better; 57k-token needle retrieval exact.
Builds on the rank-sliced EXL3 (ExLlamaV3 trellis) support added in local-inference-lab#139 by @brandonmusic, paired with the trellis3_t256 planned MoE from local-inference-lab/sparkinfer#49. No change to that work's weight loading, formats, or decode path; this commit only extends its batch dispatch. The vllm-project#139 backend serves every batch with m > 32 through the eager exllamav3_ext parity path: fp16-in/fp32-out staging (~25 GB extra HBM traffic per 4096-token step at GLM-5.2 geometry), python chunk loops, and one expert re-stream per 128-row chunk. That covers 100% of prefill, so prefill throughput is capped at ~2.0-2.3k tok/s while the same weights decode at full speed through the planned Trellis window. Build a second trellis_moe plan at max_tokens=max_num_batched_tokens with block_size_m=64 (env VLLM_EXL3_PREFILL_BLOCK_M, allowed {8,16,32,48,64}) beside the existing decode plan (32/8), and dispatch three ways in _apply_rank_sliced: * m in [min, max] -> decode plan (unchanged) * max < m <= capacity -> prefill plan (sparkinfer bind already accepts any tokens in [1, max_tokens]; zero kernel changes needed) * m < min -> parity path, whose persistent staging (xh/out32/token_sorted/weight_sorted) shrinks from capacity rows to one chunk while the prefill plan is live (~110 MiB/GPU returned) VLLM_EXL3_PREFILL_TRELLIS=0 restores the exact single-plan parity behavior (one-variable serving A/B lever). The parity branch now raises during CUDA graph capture instead of silently recording eager ext calls. The prefill arena (~1.05 GiB/GPU at capacity 3072) allocates during the profile pass, so vLLM's memory profiler sees it and the KV pool auto-shrinks instead of risking request-time OOM. Matched A/B serving brandonmusic/GLM-5.2-EXL3-TR3-3.0bpw on the published verdictai/glm52-exl3-sparkinfer runtime image (4x RTX PRO 6000 Blackwell, TP4/DCP4/MTP3, identical boot geometry, only the kill-switch flipped, fresh JIT cache both arms; llm_decode_bench prefill-only, exact /tokenize, C1, 60 s/context): ctx parity trellis delta 8k 2,287 3,742 +63.6% 32k 2,218 3,557 +60.4% 64k 2,078 3,381 +62.7% 128k 1,920 3,032 +57.9% Decode C1 unchanged-to-better; 57k-token needle retrieval exact. (cherry picked from commit d57d775)
- w4a16 kernel: compute the trellis256 B-source offset in Int64 in both the proj and non-proj branches. The expert-stride term overflows Int32 for large-but-valid MoE shapes (e.g. E=256, H=8192, bits=6 reaches ~3.21e9 > 2^31-1), which would fetch wrong or out-of-bounds B weights. Matches the sibling modelopt-native staging path that already builds this offset with Int64(...). Value-identical for in-range offsets, so existing shapes/tests are unaffected. - pyproject: declare the `triton` dependency. trellis_moe META.requires it and is_supported() gates on it, so a clean install lacking triton would silently skip the op. - trellis_moe.api: sort __all__ (RUF022). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
879ca0a to
5dd6f15
Compare
Rebased onto SparkInfer
|
| Test | v20 | Prior base |
|---|---|---|
| MTP-3 code generation | 94–100% valid — reported ~50% syntax-error symptom eliminated | ~50% |
| LAVD r10/c5 | 10/10 (0 fail) | 8/10 (2 fail) |
Tested image: verdictai/glm52-exl3-sparkinfer:v20-gg6722c1d-si1a88b38-cu132-sm120a.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
sparkinfer/moe/_shared/kernels/w4a16/prepare.py (1)
1564-1565: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: use iterable unpacking instead of tuple concatenation.
Ruff flags these two lines (RUF005). Purely stylistic; functionally equivalent.
♻️ Suggested tweak
- 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 appending their final dimension instead of tuple concatenation, preserving the existing shapes and values.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.
Inline comments:
In `@sparkinfer/_lib/compiler.py`:
- Around line 2047-2065: Update _semantic_compile_manifest_payload to include
cache_payload[4] as device_arch in semantic_payload for both explicit and
non-explicit cache formats, preserving all existing index mappings. Update
_build_compile_manifest to expose the same device_arch value alongside toolchain
so manifest semantic identity remains device-aware and consistent across
builders.
---
Nitpick comments:
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 appending their final dimension
instead of tuple concatenation, preserving the existing shapes and values.
🪄 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: 3224cf74-0a61-4973-a96d-cd82260c5596
📒 Files selected for processing (19)
pyproject.tomlsparkinfer/__init__.pysparkinfer/_lib/compiler.pysparkinfer/_lib/intrinsics.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/attention/test_fused_indexer.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
🚧 Files skipped from review as they are similar to previous changes (16)
- sparkinfer/moe/_shared/kernels/w4a16/init.py
- pyproject.toml
- sparkinfer/init.py
- tests/_lib/test_compile_cache.py
- sparkinfer/moe/trellis_moe/init.py
- validation/cutlass_migration/acceptance/e2e/contract.py
- sparkinfer/moe/trellis_moe/api.py
- sparkinfer/moe/init.py
- sparkinfer/moe/_shared/kernels/w4a16/route_pack.py
- validation/cutlass_migration/acceptance/corpus/ptx_capture.py
- tests/moe/test_trellis_moe.py
- validation/cutlass_migration/evidence/kernel_resources.py
- sparkinfer/moe/_shared/kernels/w4a16/host.py
- sparkinfer/moe/trellis_moe/_impl.py
- sparkinfer/_lib/intrinsics.py
- sparkinfer/moe/_shared/kernels/w4a16/kernel.py
|
@coderabbitai review Fixes are in — |
🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
ast-grep outline sparkinfer/_lib/compiler.py --items all --type function --match 'device_arch|semantic_compile_manifest_payload|build_compile_manifest|compile_disk_cache_payload|static_compile_cache_context'
ast-grep outline tests/_lib/test_compile_cache.py --items all --type function --match 'device_arch|disk_cache_key'
printf '\n--- compiler references ---\n'
rg -n -C 5 '_device_arch_key|_static_compile_cache_context|_compile_disk_cache_payload|_semantic_compile_manifest_payload|_build_compile_manifest|semantic' sparkinfer/_lib/compiler.py
printf '\n--- focused tests ---\n'
rg -n -C 8 'test_disk_cache_key_includes_device_arch|test_device_arch_key_retries_after_unavailable|device_arch' tests/_lib/test_compile_cache.pyLength of output: 13287
The focused tests also cover payload inclusion and retry behavior when CUDA is initially unavailable. Re-reviewing the PR now. ✅ Action performedReview finished.
|
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 #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).
Validation — this PR's fused Trellis MoE kernels drive a tr3 MTP draft layerConfirming the fused Trellis MoE kernels from this PR ( One runtime knob: the MTP-N draft issues small-m GEMMs (m=1…N), so the Trellis decode window must include them — set Image: |
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Brandon Music <brandon.m.music@gmail.com>
- w4a16 kernel: compute the trellis256 B-source offset in Int64 in both the proj and non-proj branches. The expert-stride term overflows Int32 for large-but-valid MoE shapes (e.g. E=256, H=8192, bits=6 reaches ~3.21e9 > 2^31-1), which would fetch wrong or out-of-bounds B weights. Matches the sibling modelopt-native staging path that already builds this offset with Int64(...). Value-identical for in-range offsets, so existing shapes/tests are unaffected. - pyproject: declare the `triton` dependency. trellis_moe META.requires it and is_supported() gates on it, so a clean install lacking triton would silently skip the op. - trellis_moe.api: sort __all__ (RUF022). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
device_arch (cache_payload index 4) was read into the disk-cache key but dropped from _semantic_compile_manifest_payload, so two payloads differing only by GPU architecture hashed to the same semantic identity. Add it to the common semantic dict (index 4 in both v6_explicit_spec and v3 formats). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
63bcedf to
419cb1e
Compare
Builds on the rank-sliced EXL3 (ExLlamaV3 trellis) support added in local-inference-lab#139 by @brandonmusic, paired with the trellis3_t256 planned MoE from local-inference-lab/sparkinfer#49. No change to that work's weight loading, formats, or decode path; this commit only extends its batch dispatch. The vllm-project#139 backend serves every batch with m > 32 through the eager exllamav3_ext parity path: fp16-in/fp32-out staging (~25 GB extra HBM traffic per 4096-token step at GLM-5.2 geometry), python chunk loops, and one expert re-stream per 128-row chunk. That covers 100% of prefill, so prefill throughput is capped at ~2.0-2.3k tok/s while the same weights decode at full speed through the planned Trellis window. Build a second trellis_moe plan at max_tokens=max_num_batched_tokens with block_size_m=64 (env VLLM_EXL3_PREFILL_BLOCK_M, allowed {8,16,32,48,64}) beside the existing decode plan (32/8), and dispatch three ways in _apply_rank_sliced: * m in [min, max] -> decode plan (unchanged) * max < m <= capacity -> prefill plan (sparkinfer bind already accepts any tokens in [1, max_tokens]; zero kernel changes needed) * m < min -> parity path, whose persistent staging (xh/out32/token_sorted/weight_sorted) shrinks from capacity rows to one chunk while the prefill plan is live (~110 MiB/GPU returned) VLLM_EXL3_PREFILL_TRELLIS=0 restores the exact single-plan parity behavior (one-variable serving A/B lever). The parity branch now raises during CUDA graph capture instead of silently recording eager ext calls. The prefill arena (~1.05 GiB/GPU at capacity 3072) allocates during the profile pass, so vLLM's memory profiler sees it and the KV pool auto-shrinks instead of risking request-time OOM. Matched A/B serving brandonmusic/GLM-5.2-EXL3-TR3-3.0bpw on the published verdictai/glm52-exl3-sparkinfer runtime image (4x RTX PRO 6000 Blackwell, TP4/DCP4/MTP3, identical boot geometry, only the kill-switch flipped, fresh JIT cache both arms; llm_decode_bench prefill-only, exact /tokenize, C1, 60 s/context): ctx parity trellis delta 8k 2,287 3,742 +63.6% 32k 2,218 3,557 +60.4% 64k 2,078 3,381 +62.7% 128k 1,920 3,032 +57.9% Decode C1 unchanged-to-better; 57k-token needle retrieval exact. (cherry picked from commit d57d775)
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
sparkinfer/moe/_shared/kernels/w4a16/prepare.py (1)
1855-1882: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDevice comparison here is looser than the
gate_suh/up_suhcheck.Line 1861 tests
resolved_device.index is Noneinstead of the supplied tensor's index, so an indexlessresolved_device(e.g.torch.device("cuda")from the synthetic path) accepts adummy_scaleon any CUDA device. Aligning with the strict comparison used at Line 1839 would keep the zero-copy binding contract uniform.🤖 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 1855 - 1882, Update the dummy_scale device validation in the surrounding preparation function to use the same strict device-type-and-index comparison as the gate_suh/up_suh validation, including when resolved_device.index is None. Do not allow an indexless resolved_device to accept dummy_scale from an arbitrary device; preserve the existing dtype, shape, contiguity, and alignment checks.sparkinfer/_lib/compiler.py (1)
1465-1475: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winArch key is captured from whichever device is current at first call.
get_device_capability()/get_device_name()with no argument read the current device, and the result is memoized process-wide. If the key is computed beforetorch.cuda.set_device(...)(or on a heterogeneous-GPU node), the on-disk cache can be keyed by a different device than the one actually compiled for. Consider binding to the explicitly selected device and, optionally, keying per device index.♻️ Suggested hardening
- major, minor = torch.cuda.get_device_capability() - name = torch.cuda.get_device_name() + index = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(index) + name = torch.cuda.get_device_name(index)As per coding guidelines, "Migration acceptance must be GPU-only and exercise the real production launch path on explicitly selected physical GPUs".
🤖 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 1465 - 1475, Update the device-architecture key computation around torch.cuda.get_device_capability and torch.cuda.get_device_name to use the explicitly selected physical GPU rather than the implicit current device, and include the device index in the memoization/cache key when devices may differ. Preserve the existing unavailable fallback, and ensure migration acceptance exercises the production launch path on explicitly selected GPUs only.Source: Coding guidelines
tests/attention/test_fused_indexer.py (1)
549-575: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMake the high-page boundary an executable invariant.
Add
assert pid_lo * record_bytes >= (1 << 31)near the fixture constants. Otherwise, a future stride or page-ID change could leave this test passing without exercising the intended Int32-offset boundary.As per coding guidelines, paged-pool tests must include live pages beyond the
2^31 / strideboundary.Suggested guard
pid_lo, pages_used = 2000, 64 pool_pages = pid_lo + pages_used + 4 need = pool_pages * record_bytes + (1 << 29) + assert pid_lo * record_bytes >= (1 << 31)🤖 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/attention/test_fused_indexer.py` around lines 549 - 575, Add an executable assertion near the fixture constants requiring pid_lo * record_bytes to be at least 1 << 31. Keep the existing high-page setup unchanged so the paged-pool test continues exercising live pages beyond the Int32 offset boundary.Source: Coding guidelines
🤖 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 `@tests/attention/test_fused_indexer.py`:
- Around line 589-605: Update the validation around _golden_topk to return
ordered golden indices alongside golden values. For each row, use the returned
idx[row] to gather the corresponding expected golden values, then compare those
values and indices pairwise against the implementation output; remove the
independent sorted-value and set-only assertions.
---
Nitpick comments:
In `@sparkinfer/_lib/compiler.py`:
- Around line 1465-1475: Update the device-architecture key computation around
torch.cuda.get_device_capability and torch.cuda.get_device_name to use the
explicitly selected physical GPU rather than the implicit current device, and
include the device index in the memoization/cache key when devices may differ.
Preserve the existing unavailable fallback, and ensure migration acceptance
exercises the production launch path on explicitly selected GPUs only.
In `@sparkinfer/moe/_shared/kernels/w4a16/prepare.py`:
- Around line 1855-1882: Update the dummy_scale device validation in the
surrounding preparation function to use the same strict device-type-and-index
comparison as the gate_suh/up_suh validation, including when
resolved_device.index is None. Do not allow an indexless resolved_device to
accept dummy_scale from an arbitrary device; preserve the existing dtype, shape,
contiguity, and alignment checks.
In `@tests/attention/test_fused_indexer.py`:
- Around line 549-575: Add an executable assertion near the fixture constants
requiring pid_lo * record_bytes to be at least 1 << 31. Keep the existing
high-page setup unchanged so the paged-pool test continues exercising live pages
beyond the Int32 offset boundary.
🪄 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: be21eac0-1684-4c5f-8446-eedb7b38aef9
📒 Files selected for processing (19)
pyproject.tomlsparkinfer/__init__.pysparkinfer/_lib/compiler.pysparkinfer/_lib/intrinsics.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/attention/test_fused_indexer.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
🚧 Files skipped from review as they are similar to previous changes (16)
- sparkinfer/init.py
- sparkinfer/moe/trellis_moe/api.py
- pyproject.toml
- sparkinfer/moe/_shared/kernels/w4a16/route_pack.py
- sparkinfer/moe/trellis_moe/init.py
- validation/cutlass_migration/acceptance/e2e/contract.py
- sparkinfer/moe/_shared/kernels/w4a16/init.py
- tests/_lib/test_compile_cache.py
- sparkinfer/moe/init.py
- validation/cutlass_migration/acceptance/corpus/ptx_capture.py
- validation/cutlass_migration/evidence/kernel_resources.py
- sparkinfer/moe/_shared/kernels/w4a16/host.py
- tests/moe/test_trellis_moe.py
- sparkinfer/moe/trellis_moe/_impl.py
- sparkinfer/_lib/intrinsics.py
- sparkinfer/moe/_shared/kernels/w4a16/kernel.py
The direct-K high-page-id test compared the sorted top-k values and the returned index set independently, so a regression that emitted the correct values against the wrong indices would still pass. _golden_topk grows an opt-in return_scores flag (existing call sites keep the two-tuple contract) and the test now gathers the golden score for each index the kernel actually returned and compares it position by position. Gathering by the returned index instead of asserting an index order keeps the check robust to ties at the top-k boundary.
…est invariant compiler: read the arch cache key from the explicitly selected device. torch.cuda.get_device_capability()/get_device_name() with no argument read the implicitly current device and the result is memoized process-wide, so a lookup made before a worker selects its GPU pins the key to the wrong card. Observable on mixed rigs: Max-Q and non-Max-Q RTX PRO 6000 boards share compute capability 12.0 but report different device names. prepare: normalise an indexless resolved_device to a concrete ordinal in the synthesize path, then validate dummy_scale with the same strict comparison used for the suh/svh tables. Previously the two disagreed -- suh/svh used strict equality (so cuda:0 != cuda rejected a correctly placed tensor) while dummy_scale used an index-optional comparison that accepted a tensor from any ordinal. Normalising once makes both strict and correct. tests: assert the direct-K case actually crosses the Int32 page-offset boundary. The i64 path it exists to cover is implied only by the constants, so retuning record_bytes or pid_lo downward would leave a passing test that no longer exercises the boundary.
|
@coderabbitai All three nitpicks addressed in
I deliberately did not add the device index to the cache key itself. All ranks on a homogeneous-capability rig should keep sharing compiled artifacts; keying on the ordinal would fragment the JIT cache per rank and multiply cold-start compile time for no correctness gain.
Instead the synthesize path now pins an indexless
All three verified by syntax check and the boundary arithmetic above. The CUDA suite has not been run against this head from CI. |
Runtime bumped to v24 (target/draft scratch isolation)The validated runtime referenced in this PR is now: It supersedes the earlier v23 image, which predates the rank-sliced runtime scoping fix. The pins in the description are updated accordingly. Why: the EXL3 rank-sliced runtime cache was keyed only on device/dtype/shape/topk/planner settings, and the cached entry owns mutable Trellis/prefill scratch. The MTP-78 draft layer matches a target MoE layer on every one of those components, so it was reusing the target's scratch. The key is now scoped to the owning quant config, giving each model its own runtime. Measured effect (the fix is visible in memory, not in the log -- the planner line uses
The 119,552-token KV reduction corresponds to ~1056 MiB against a 1054.2 MiB arena, which is direct evidence that a second runtime is allocated. The reduced KV capacity is the intended cost of the isolation, and it is also why the fix is scoped per model rather than per layer -- per-layer ownership would have needed tens of GiB per rank across 75+ layers. Quality on v24: Estonia c2 x5 = 5/5 pass. LAVD c5 across two independent 5-run samples = 4 EXACT / 5 NEAR / 1 FAIL over 10 runs (sample 1: 2E/3N/1F, sample 2: 2E/3N/0F, matching the pre-fix image exactly). The single miss was a wrong ledger total at 7,114 completion tokens against a 24,576 cap, so answer quality rather than truncation, and within this profile's historical spread. Reporting the aggregate including the failure rather than only the clean sample. |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
git show --no-ext-diff --format='commit %H%n%s' --stat 72d7434
printf '\n--- relevant declarations ---\n'
ast-grep outline sparkinfer/_lib/compiler.py --items all --type function --match 'device_arch'
ast-grep outline sparkinfer/moe/_shared/kernels/w4a16/prepare.py --items all --type function --match 'prepare_trellis256_moe_weights|prepare_trellis256_dense_weight'
ast-grep outline tests/attention/test_fused_indexer.py --items all --type function --match 'high_page_id'
printf '\n--- compiler device-arch implementation ---\n'
rg -n -C 12 '_device_arch_key|current_device|get_device_capability|get_device_name' sparkinfer/_lib/compiler.py
printf '\n--- trellis preparation device normalization/validation ---\n'
rg -n -C 10 'resolved_device|dummy_scale|gate_suh|up_suh|svh|cuda' sparkinfer/moe/_shared/kernels/w4a16/prepare.py
printf '\n--- high-page boundary assertion ---\n'
rg -n -C 12 'pid_lo|record_bytes|1 << 31|high_page_id_i64' tests/attention/test_fused_indexer.py
printf '\n--- targeted commit diff ---\n'
git show --no-ext-diff --format= -- \
72d7434 \
sparkinfer/_lib/compiler.py \
sparkinfer/moe/_shared/kernels/w4a16/prepare.py \
tests/attention/test_fused_indexer.pyLength of output: 31151
One correction on To guarantee the intended behavior on a mixed-GPU process, the architecture key needs to be derived from an explicit target/compile device propagated to this code path (and cached per architecture identity), or cache initialization must be deferred until after device selection. As written, this finding remains unresolved. |
The previous change was not a fix. torch resolves get_device_capability() and get_device_name() against torch.cuda.current_device() when called without an argument, so passing that same ordinal explicitly is equivalent and left the process-wide _DEVICE_ARCH_KEY free to freeze whichever GPU happened to be current on the first call. It also broke test_device_arch_key_retries_after_unavailable, whose torch stubs take no device argument. Replace the single global with a per-ordinal dict so each entry is bound to the device it was measured on, and thread the ordinal through _static_compile_cache_context, which is lru_cached on compile_callable and would otherwise re-freeze the identity at that layer regardless. The returned key still omits the ordinal: two GPUs with the same capability and name produce the same key and keep sharing compiled artifacts, so a homogeneous rig does not fragment its JIT cache. Only the lookup is per ordinal. Adds test_device_arch_key_is_memoized_per_ordinal, which models a mixed rig (two boards sharing compute capability 12.0 but reporting different names) and asserts the second board does not inherit the first board's identity. Restores the retry test against the per-ordinal cache.
|
@coderabbitai You're right, and thanks for not letting that one go. My previous change was not a fix — It was also worse than a no-op: it broke Fixed properly in Per-ordinal memoization. Second cache layer, which your comment led me to. Deliberate non-change: the returned key still omits the ordinal. Two GPUs with the same capability and name produce the same key and keep sharing compiled artifacts; keying on the ordinal would fragment the JIT cache per rank and multiply cold-start compiles for no correctness gain on a homogeneous rig. Only the lookup is per ordinal. Tests: added
|
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
git show --no-ext-diff --format='commit %H%n%s' --stat d4438d4
printf '\n--- compiler function map ---\n'
ast-grep outline sparkinfer/_lib/compiler.py --items all --type function --match 'device_arch|static_compile_cache_context|compile_disk_cache_payload'
printf '\n--- compiler implementation and call sites ---\n'
rg -n -C 10 '_DEVICE_ARCH_KEYS|_DEVICE_ARCH_KEY|_current_device_ordinal|_device_arch_key|_static_compile_cache_context' sparkinfer/_lib/compiler.py
printf '\n--- targeted test map ---\n'
ast-grep outline tests/_lib/test_compile_cache.py --items all --type function --match 'device_arch|disk_cache_key'
printf '\n--- targeted test implementations ---\n'
rg -n -C 12 'test_device_arch_key_retries_after_unavailable|test_device_arch_key_is_memoized_per_ordinal|_DEVICE_ARCH_KEYS|device_arch' tests/_lib/test_compile_cache.py
printf '\n--- focused commit diff ---\n'
git show --no-ext-diff --format= d4438d4 -- \
sparkinfer/_lib/compiler.py \
tests/_lib/test_compile_cache.pyLength of output: 19918
The returned semantic key appropriately remains based on architecture identity rather than ordinal, retaining artifact sharing across equivalent GPUs. The focused tests cover unavailable-probe retry, mixed-device identities, explicit ordinal lookup, and simultaneous cache entries for both devices. ✏️ Learnings added
If you found this review helpful, would you consider giving us a shout-out on X? Thank you for using CodeRabbit! |
Runtime bumped to v26 — supersedes the v24 note aboveThe description pins have been updated to this image. v24 and the intermediate v25 are superseded — please disregard both. What changed since v24. v24 carried the per-model rank-sliced runtime scoping (target/draft scratch isolation). v25 additionally carried a change to the compile-cache architecture key that CodeRabbit correctly identified as a no-op: v26 carries the real fix:
This is not academic on the validation rig: its four boards report two different device names (Max-Q and non-Max-Q) while sharing compute capability 12.0. Validation on v26 (4x RTX PRO 6000 SM120a, TP4/DCP4, MTP-3, NVFP4 DS-MLA KV):
GPU KV cache 998,656 tokens (1.90x at 524,288). That is ~1 GiB below the pre-isolation 1,115,904 (2.13x), which is the intended cost of giving the draft model its own Trellis/prefill arena. Unit tests: CI status unchanged: |
1 participant
Summary
Adds the EXL3 Trellis compute path and a stable planned MoE API to Sparkinfer.
It consumes native EXL3 MCG-codebook tensors without repacking or offline
requantization and supports fused routed-expert execution on Blackwell.
This PR includes:
sparkinfer.moe.trellis_moeAPI used by vLLMThe production consumer is
local-inference-lab/vllm#139.
Rebase and release pins
masteratc39b8062ba450c030e669d898a026d10980c9470d4438d490691f79022fdfc8149e1c5f161d1544526c4bfdd3ff2be0433e6fe07e0c3be535f5bb318The exact validated image is:
Finalized Gilded Gnosis v20 common base (2026-07-25)
This branch is rebased onto its own base branch for review and merge. It has also
been built and validated as an EXL3 source layer on the finalized v20 common
base. Per the v20 release note the base packages no EXL3/Trellis loader or
kernels, so EXL3 is enabled by rebasing the EXL3 layer onto the pinned stack --
not by pointing the
voipmonitor/vllmtag at an EXL3 checkpoint.voipmonitor/vllm:gilded-gnosis-v20-vllm5517197-sibe0edca-fi801d57a-cu132-20260725551719766029e78824a30d97ae6ac63917405b5fbe0edcaae6f5d284bb29a82325aba7a0ead6960fmodels/deepseek_v2.pyandv1/attention/backends/mla/indexer.pydiffer between the previous and thefinalized v20 tree. Both were re-derived from the v20-final versions with the
EXL3 edits replayed on top, so the new DCP/indexer/prefill work is preserved
rather than overwritten by the overlay.
Validated image:
Runtime validation, 4x RTX PRO 6000 Blackwell (SM120a), TP4/DCP4, MTP-3, NVFP4 DS-MLA KV:
v0.11.2.dev280+gilded.gnosis.v20.vllm5517197.sibe0edca.fi801d57a.cu132.20260725vLLM is using nccl==2.30.4Trellis m=1..32 block_m=8, prefill trellis block_m=64 arena=1054.2MiB capacity=3072 chunk=128 topk=8DCP_*=auto): query split 1,full CKV gather 1, top-k owner merge 1, indexer shards 0, CKV prefetch depth 1,
and
DCP_PREFILL_WORKSPACE=1(VLLM_DCP_PROJECT_BEFORE_MERGE=1+VLLM_B12X_MLA_DCP_GATHER_IN_WORKSPACE=1)Preallocated 30.8 MiB for 2 persistent CKV execution lane(s)andUsing native CKV layer prefetch with depth=1 and 2 workspace slotsthan the pre-fix image because the draft model now gets its own rank-sliced
runtime arena instead of sharing the target's -- the intended cost of the
target/draft scratch isolation.
Earlier images in this series measured LAVD at 2E/3N/0F (pre-fix) and, on one
5-run sample of the first scoped build, 2E/2N/1F; a second sample of that build
returned 2E/3N/0F. The single failure there was a wrong ledger total at 7,114
completion tokens against a 24,576 cap -- an answer-quality miss rather than
truncation, and within this profile's run-to-run spread. This build shows no
failures on either suite.
Note: the base's
Loading safetensors using InstantTensor loaderassertion doesnot appear on this path -- the EXL3 checkpoint is loaded by the EXL3 rank-sliced
loader, not the base InstantTensor path. All other boot assertions hold.
Validation
The kernel ladder was validated on RTX PRO 6000 Blackwell (SM120a), including
bitrate-specific reconstruction, batch-shape stability, and CUDA graph replay.
Focused exact-image results:
End-to-end validation used the 332.19 GB GLM-5.2 rank-sliced EXL3 checkpoint on
TP4 + DCP4 A2A with MTP1 greedy, NVFP4 DeepSeek-MLA KV cache, and async
scheduling disabled. All 81 shards loaded, CUDA graphs captured through C8,
LAVD passed 20/20 across two full runs, and Estonia passed 5/5.
Exact-image performance:
Notes and limits
alignment and is performance-validated on SM120a.
the vLLM consumer owns its parity fallback.
chunked by the consumer according to its configured capacity.
remaining release constraint is in vLLM speculative/CKV lifetime management,
not this Trellis compute API.
AI assistance disclosure
Implementation, debugging, and documentation were developed with Codex and
Claude Code assistance. The author selected the design, reviewed the changes,
and ran the reported hardware validation.
Summary by CodeRabbit