Repository navigation
perf(moe): full-codeword lookup for prepared K2 Trellis decode - #436
Conversation
Preserve source-coordinate expert extents, retained two-bit lookup tables and declared row capacities. Keep block-codec epilogues and pipeline tuning intact. Optional FP32 atomic prefill reduction does not guarantee bitwise summation order. Validation: Host eligibility checks pass for packed, ModelOpt, EXL3 and three block codecs. Syntax and whitespace checks pass. GPU validation of the canonical-base composition remains pending. Co-authored-by: Martin Vit <martin@voipmonitor.org> Co-authored-by: Codex <noreply@openai.com>
Normalize uniform EXL3 slots into TrellisSource and TrellisWeights. Preserve codewords, FP16 scale boundaries, global Hadamard coordinates and bounded CUDA staging. Route the compatibility loader through common preparation while retaining asymmetric-pair support. Native config callers keep their schema and projection-tiered path. CPU layout and adapter validation: 98 tests passed. GPU preparation, graph replay, performance and full-model fidelity remain unqualified in this environment. Assisted-by: OpenAI Codex
…packing Remove the full-codeword shared LUT and stable per-expert route ordering from the preparation interface change. Retain exact route capacities, prepared launch ownership and opt-in FP32 prefill reduction. Bump the MoE query schema for the reduced control set and scope the release fragment to PR 433. Validation: selected host suites report 181 passed and 28 failures; the same 28 missing-decode_config failures reproduce at 7fca2aa. Ten LUT-specific cases move with the LUT change. Changed Python sources parse and git diff --check passes. GPU execution and serving qualification were not run. Assisted-by: OpenAI Codex
Restore the 64 KiB shared codeword table, immutable selection control, residency checks and decoding tests as an independently reviewable change. Preserve the original auto/compact/full behavior; qualification remains pending. Document the instruction-scheduling difference that a matched compact-decoder benchmark must control. Validation: adapter, config, schema and variant-selection suites: 170 passed. Python source matches the extracted implementation except for the query schema version; syntax and whitespace checks pass. CUDA execution, graph replay and performance were not run. Source-commit: 7fca2aa Assisted-by: OpenAI Codex
|
Warning Review limit reachedNext included review available in 35 minutes. View limit detailsLimit details: You’ve used all 2 included reviews currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Repository: local-inference-lab/b12x/.coderabbit.yaml Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (9)
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 |
Kimi-K3 TP9 qualificationThe extracted full-codeword decoder is qualified for the profile below and restores a measured serving regression after switching to the canonical common-Trellis API. Same independently encoded two-bit checkpoint, nine RTX PRO 6000 Blackwell GPUs, TP9/DCP9, Red Hat DSpark K5, BF16 activation arithmetic, FP8 KV, scheduler chunk 4096, fixed 183-token coding prompt, 4096 generated tokens, three runs:
All nine output hashes match: A component A/B uses actual checkpoint layer-1 weights, hidden width 3584, 896 experts, top-16 routing, six BF16 input rows and CUDA Graph replay. Compact/full medians are 72.308/51.811 µs for a 384-channel TP extent and 68.002/49.280 µs for a 256-channel extent. Outputs are bit-identical for live row counts 1/5/6 and five routing-sharing patterns; replay allocations are stable. These synthetic routes are not claimed to reproduce a model's routing distribution. Validation commands: python -m pytest tests/moe/test_trellis_direct_lut_decode.py -q
python -m pytest tests/moe/test_fused_moe_variant_selection.py -q9 GPU decoder cases and 39 policy tests pass, including exhaustive codewords and randomized crossing windows. The comparisons qualify the implementation as a whole; they do not isolate table size from instruction scheduling or prove an equivalently rescheduled compact intrinsic cannot match it. Recommendation: review this PR as a serving dependency for the qualified K2 profile. Keep #437 independent. Exact source/image identities, scripts and raw receipts are tracked in local-inference-lab/vllm#906 under AI assistance: OpenAI Codex. No additional quantization or FP8 activation path was introduced. |
voipmonitor
changed the title
2 participants
Behavior
Eligible prepared uniform K2 Trellis kernels with intermediate Hadamard transforms stage a 64 KiB full-codeword table per CTA. Direct byte lookups replace the compact decoder's integer permutation without changing decoded weights, cooperative residency or output arithmetic.
B12X_TRELLIS_DECODE_TABLE=auto|compact|fullis captured in the immutable preparation query.autoselects the full table for eligible capacities up to 16 tokens;compactretains the comparison implementation. Ineligible explicitfullrequests fail instead of changing occupancy silently.Implemented. Qualified: Kimi-K3 uniform two-bit EXL3, TP9/DCP9, Red Hat DSpark K5 on RTX PRO 6000 Blackwell. Other serving profiles have not been qualified by this evidence.
Measured benefit and correctness
Matched 183-token coding prompt, 4096 generated tokens, temperature zero, seed 1, three runs on identical GPUs/settings:
All output-text hashes, token counts and acceptance counters match. The implementation restores decode throughput without changing the generated workload. Prefill remains approximately 2.6k/2.6k/2.5k tok/s at 8k/32k/64k. Stable route packing (#437) is not part of either canonical arm.
The qualification comment records exact output identity and the real-weight component A/B. At six rows, the full decoder reduces whole-MoE graph latency from 72.308 to 51.811 microseconds for a 384-channel extent and from 68.002 to 49.280 microseconds for a 256-channel extent. Eager live counts and graph route-pattern changes produce bit-identical outputs.
Scope and alternatives
Targets
masterafter merged #433. #287 implements a related mechanism on the legacy execution API, and #300 concerns staging. Do not merge those alternatives automatically with this prepared-capacity implementation.The full-table intrinsic also exposes eight independent loads and uses different byte packing. The measurements qualify that implementation as a whole. They do not isolate table size from instruction scheduling or prove that an equivalently rescheduled compact decoder cannot match its speed.
Validation
9 GPU decoder cases and 39 policy cases passed. Coverage includes all 65,536 codewords, random cross-word windows, immutable query controls, eligibility and cooperative-residency guards. The integration includes three full-model decode runs and repeated prefill measurements; vLLM #906 records the source-locked composition and its dependencies.
AI assistance: OpenAI Codex. Original contributor attribution is preserved. Maintainer review remains required.