Skip to content

perf(moe): skip inactive block-fp8 decode routes - #665

Open
taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:perf/fp8-block-inactive-slots
Open

taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:perf/fp8-block-inactive-slots

Conversation

@taking-lying-flat

Copy link
Copy Markdown
Contributor

With hybrid block-FP8 decode, CPU-assigned routes currently reach the GPU as slot 0 with routing weight 0. The gate/up GEMM still reads and computes that expert, and the down GEMM only multiplies by zero after its K-loop. With two GPU routes out of top-k 8, six additional routes therefore run both GEMMs.

Handle negative slots in the shared decode kernel by writing zero to the route's output tile and returning before reading activations, expert weights, scales or routing weights. Both gate/up and down use this kernel. Writing the outputs is necessary because the following activation/reduction reads those buffers, including on CUDA Graph replay. Then enable supports_inactive_slots so hybrid preserves its -1 sentinels.

The launch shapes stay fixed, and the activation still processes the zeroed buffer. Follow-up to #662.

Validation

Tested against main 95ac86f21567bc244c7a9d91972838fbf049b4ed on an NVIDIA RTX A1000 Laptop GPU (4 GiB, sm_86), Intel Core i7-12800HX, Ubuntu 26.04.1, NVIDIA driver 595.91.07, PyTorch 2.12.1+cu130 and Triton 3.7.1. This exercises the FP8 compatibility path. The existing test environment's PyTorch version is below main's currently declared >=2.14 requirement.

  • 19 passed: the three existing block-FP8 reference tests plus 16 local regression cases. These cover both GEMMs, bf16/fp16/fp32 outputs, noncontiguous strides, partial N/K tiles, padded scales, five activations, empty expert banks with all routes inactive, and hybrid dispatch with a CPU-executor stub
  • CUDA Graph replay: switch the same graph between all-active, mixed and all-inactive routes while changing inputs. Poison the gate/up, activation and down buffers with NaN before replay; inactive routes become exactly zero and outputs match the old clamped-slot/zero-weight path exactly
  • Compute Sanitizer memcheck: 10 targeted cases passed, 0 errors
  • Native FP8 gate/up and down kernels compile for sm_89 and sm_90; compilation only, not execution on those GPUs
  • ruff check --select E9,F63,F7,F82, py_compile and git diff --check passed

Regression and benchmark scripts were kept outside the repository; this PR contains only the two implementation files. Local commands:

python -m pytest tests/kernels/test_fp8_blockscale_moe.py /tmp/freetoken-fp8-inactive-validation/test_inactive_slots.py -q --tb=short
compute-sanitizer --tool memcheck --error-exitcode 99 --log-file /tmp/freetoken-fp8-inactive-validation/memcheck.log python -m pytest /tmp/freetoken-fp8-inactive-validation/test_inactive_slots.py -k 'both_gemms or all_inactive or cuda_graph' -q --tb=short
python /tmp/freetoken-fp8-inactive-validation/bench_inactive_slots.py

GPU decode A/B

Synthetic block-quantized expert banks, batch size 1, top-k 8, two resident expert banks, bf16 activations and padded block scales. No model checkpoint was loaded. The baseline maps inactive slots to 0 and zeroes their routing weights; the new path keeps negative slots. Times cover gate/up, activation, down and reduction under CUDA Graphs. Each value is the median of five 100 ms measurements, alternating baseline/new order.

Hidden / intermediate Active GPU routes Main This PR Speedup
2048 / 768 8 / 8 179.218 us 179.020 us 1.001x
2048 / 768 4 / 8 180.292 us 102.818 us 1.754x
2048 / 768 2 / 8 181.389 us 60.589 us 2.994x
2048 / 768 0 / 8 173.964 us 6.748 us 25.782x
4096 / 1024 8 / 8 472.466 us 472.730 us 0.999x
4096 / 1024 4 / 8 462.024 us 255.034 us 1.812x
4096 / 1024 2 / 8 476.249 us 148.278 us 3.212x
4096 / 1024 0 / 8 459.185 us 8.819 us 52.070x

These are local GPU decode measurements on the sm_86 compatibility path. End-to-end model tokens/s, the complete CPU/GPU overlap path and native-FP8 GPU runtime performance have not been measured.

KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 12, 2026
… decode routes

Assisted-by: Claude Opus 5.5
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 12, 2026
Follow-up to upstream PR FlashML-org#665: hybrid decode passes CPU-assigned routes as slot -1, which must add nothing even with a nonzero weight; fails before FlashML-org#665.

Assisted-by: Claude Opus 5.5
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 12, 2026
FlashML-org#665 adopted with a negative-slot test, FlashML-org#378 deferred again (ROCm, now ready); FlashML-org#399 landed upstream via FlashML-org#662.

Assisted-by: Claude Opus 5.5

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant