feat(moe): serve canonical Fruit QSRT atoms

[malaiwah]

Summary

• Serve canonical Fruit qsrt_atoms_v1 geometry (H=1024, I=512, 96 fixed atoms per expert) through the public B12X MoE API.
• Decode per-expert P24/P33 pair modes and physical atom rotation without reconstructing dense expert weights.
• Preserve W4A16 prefill/reference execution and select W4A8 only for bounded decode rows.
• Support per-expert SiLU/SiTU metadata and fail closed on malformed layouts, unsupported activations, mixed gate/up modes, and invalid route maps.
• Keep steady-state storage caller-owned and CUDA-graph replay safe.

Current canonical head: f4064d06f029240040a128388f32c861a535ad68.

Hot path

• FC1 input rotation and MXFP8 quantization use one fixed CUDA kernel rather than separate rotate/quantize launches.
• A route map is materialized once; FC1 and FC2 consume it directly instead of building Python-side old/new expert partitions and concatenating outputs.
• FC1 publishes one route-major output row, eliminating multipart accumulation for independently selected gate/up pair modes.
• FC2 remains route-major and preserves exact top-k weighting and physical-rotation semantics.
• W4A8 decode expands fixed-stride trellis payload directly into E4M3 MMA operands; W4A16 remains the graph-safe prefill/reference and unsupported-shape fallback.

Verification

The current head is the exact B12X source fingerprinted into the final runtime image and sealed publication. The full 2,816-expert Fruit Instruct package loaded every routed and MTP layer under vLLM FULL_AND_PIECEWISE CUDA graphs. Runtime-path evidence records W4A16 prompt/prefill and W4A8 decode graph capture and replay, including the packaged MTP layer.

The sealed matched protocol used one immutable RTX 5090 image, launch order QSRT → BF16 → SIQ, TP1, max_num_seqs=1, identical prompt tokens/settings, and three warmed repetitions per arm:

Arm	Median client-observed generated-token rate	Loader weight memory
BF16	445.04 tok/s	9.730 GiB
SIQ	428.63 tok/s	3.240 GiB
QSRT	425.67 tok/s	3.100 GiB

These rates include serving/request overhead and are not decode-only or general-throughput claims. QSRT used 68.14% less loader weight memory than BF16 and 4.32% less than SIQ; its median rate was 4.35% below BF16 and 0.69% below SIQ under this exact protocol.

Across 5,870 full-vocabulary BF16-reference positions:

Candidate	Mean forward KL	Max forward KL	Top-1	Top-10
SIQ	0.0551245	3.22484	90.05%	99.88%
QSRT	0.0528832	1.27830	90.24%	99.78%

All three arms passed 0/8 focused absolute behavior contracts. The result qualifies storage, loading, bounded runtime execution, and relative fidelity—not assistant quality.

Publication

• Published artifact: malaiwah/GLM-5.2-QSRT-Fruit-Instruct@cba27c73.
• QSRT encoder: local-inference-lab/qsrt#4 at 2113af303f37cedf4b538dcf68eb699d5e31f7df.
• vLLM loader/launcher: serve: load canonical Fruit QSRT atoms vllm#269 at 0429cb4c11ad7a67ec8b40b621eae061a363b6cb.
• Qualified runtime image: sha256:e1b411c3c1ef02e0e8be966de31d727f04683f93d0c61711cb6ae703c9ee8a19.
• Completion marker SHA-256: 41f32656b7cc68e4c1bc6cec8c91964672ff2ff0ff989448a6d5f37e3b7fb9c3.
• Runtime receipt SHA-256: c56f33b5da5e813d83152c5775288c5364c74205a3d70aacf528eafdc327ea78 (B12X distribution 1.2.3).

Scope

Only TP1 physical serving is qualified. TP2 atom ownership is unit-tested but was not physically served on this single-GPU host. The sparse-prefill integration enforces max_num_seqs=1; this PR does not claim broader concurrency, TP>1, long-context qualification, standardized benchmark quality, or general assistant capability.

Depends on local-inference-lab/qsrt#4 and local-inference-lab/vllm#269.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The PR adds SiLU support, shared and per-expert SUH handling, allocation-free W4A8 scratch views, prepared multipart execution, dynamic QSRT geometry, stricter validation, and expanded Trellis MoE tests.

Changes

Trellis MoE updates

Layer / File(s)	Summary
Activation selection and metadata b12x/moe/_shared/kernels/trellis_w4a8_transform.py, b12x/moe/_shared/kernels/w4a16/prepare.py, tests/gemm/test_trellis_linear.py	Activation kernels support silu and situ. Compilation keys include activation. Prepared weights and test descriptors store activation and SUH metadata.
Dynamic QSRT preparation b12x/moe/_shared/kernels/w4a16/prepare.py, b12x/moe/fused_moe/_impl.py, benchmarks/benchmark_trellis_pair_moe_tp12.py	QSRT payload geometry and rotation placement use supplied atom slots, expert counts, and rotation multipliers. Public preparation validates and forwards this metadata.
W4A8 scratch and multipart execution b12x/moe/_shared/kernels/trellis_w4a8.py	Scratch allocation and prefix views support shared and per-expert SUH layouts. Prepared-input and multipart APIs validate routes, tensor layouts, metadata, and accumulation storage.
Integration and reference coverage tests/moe/test_fused_moe_trellis.py	Tests cover MXFP8 references, dynamic QSRT payloads, hidden sizes 256 and 1024, both activation modes, SUH layouts, multipart accumulation, CUDA graph replay, and malformed payloads.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant WeightPreparation
  participant PreparedInput
  participant W4A8Runtime
  participant ActivationKernel
  participant Reduction
  Caller->>WeightPreparation: prepare activation and QSRT metadata
  Caller->>PreparedInput: validate inputs, routes, and SUH layout
  PreparedInput->>W4A8Runtime: provide prepared operands and scratch views
  W4A8Runtime->>ActivationKernel: apply silu or situ
  ActivationKernel->>Reduction: provide activated expert outputs
  Reduction->>Caller: return accumulated MoE output

Loading

Possibly related PRs

• local-inference-lab/b12x#117: Adds related shared/per-expert SUH handling and runtime expert/route geometry in another Trellis MoE kernel implementation.

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the primary change: adding support for serving canonical Fruit QSRT atoms in the MoE preparation path.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

b12x/moe/_shared/kernels/trellis_w4a8.py (1)

262-288: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Validate prepared.down_svh before FC2 launch.

_validate_prepared_part checks gate_suh and up_suh, but run_trellis_w4a8_moe_prepared passes prepared.down_svh to _w4a16_topk_sum_launch_flat after FC1, activation rotation, and FC2 routing already ran. A missing or malformed down_svh raises an AttributeError and leaves scratch partially written. Add down_svh to the _validate_prepared_part loop and include its row count in the shared/per-expert row check, since preps require the same (1, H) or (num_experts, H) shape.

🤖 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 `@b12x/moe/_shared/kernels/trellis_w4a8.py` around lines 262 - 288, The
_validate_prepared_part validation currently omits prepared.down_svh. Add
down_svh to the tensor validation loop and include its row count in the
shared_suh/per-expert expected-row check, requiring it to match gate_suh and
up_suh as a contiguous fp16 tensor on the target device with shape (1|E,
hidden_size) before FC2 launch.

🤖 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.

Outside diff comments:
In `@b12x/moe/_shared/kernels/trellis_w4a8.py`:
- Around line 262-288: The _validate_prepared_part validation currently omits
prepared.down_svh. Add down_svh to the tensor validation loop and include its
row count in the shared_suh/per-expert expected-row check, requiring it to match
gate_suh and up_suh as a contiguous fp16 tensor on the target device with shape
(1|E, hidden_size) before FC2 launch.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f2003ae0-333b-49fa-a1e6-17a0dc8df2b9

📥 Commits

Reviewing files that changed from the base of the PR and between 84df431 and 89876a5.

📒 Files selected for processing (2)

• b12x/moe/_shared/kernels/trellis_w4a8.py
• tests/moe/test_fused_moe_trellis.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/moe/test_fused_moe_trellis.py

[malaiwah]

Rechecked immutable runtime head f4064d06f029240040a128388f32c861a535ad68 against current B12X master ce7f6227.

Master is two commits ahead (fix(mla) and topology-scoped PCIe all-reduce). The merge is conflict-free, and the affected GEMM/MoE/QSRT suite passes: 194 passed.

Those upstream changes do not require a Fruit runtime or package-pin change. The sealed runtime image, QSRT package, and vLLM integration authenticate f4064d06; please preserve that commit with merge-commit semantics rather than squashing or rebasing it away.