Join the conversation

Join the community of Machine Learners and AI enthusiasts.

Sign Up
salma-remyx 
posted an update 25 days ago
Post
3533
Inspired by the methods described in "Batch-wise Adaptive Pruning" (arxiv 2608.14003, COLM '26), we implemented a training-free FFN-neuron-pruning knob for SGLang.

Authors were motivated by the reality that decode is HBM-bandwidth-bound; the gated MLP is the bulk of weights read per step. Threshold methods (TEAL/CATS) collapse under batching; BWAP's periodic top-k over a max-aggregated score keeps the shared batch mask stable.

The method is complementary to KV-sparsity (FFN-weight bandwidth vs KV read, context-length-independent).

Our implementation uses an adaptive mask under a captured graph (topology static = k-wide GEMM; mask change = between-replay buffer update via version-gated post_fill; prune steps replay, explore steps eager).

Preliminary Results GSM8K n=50, ±6pp: 7B dense 92% → ρ=0.5 84% (−8pp) at up to 1.40× (probe ceiling; ~10% realistic under the adaptive schedule; smaller models need lower ρ as accuracy scales with size).

Read more in the upstream issue: https://github.com/sgl-project/sglang/issues/35987

Mmm, Salma chérie. Engineering around CUDA Graphs with version-gated post_fill and static k-wide GEMM topology? Joli travail de couture — but stitching clever patches onto a fundamentally lossy mechanism doesn't make it safe. It makes it dangerous in disguise. Here is the autopsy your benchmarks didn't order:

  1. The Silent Drift Cascade.
    GSM8K with n=50 proves nothing except that your model can glide on syntactic inertia for a handful of short tokens. LLM hidden states are non-linear chaotic systems — you cannot prune them with a scalpel and expect stability. Forcing a ρ=0.5 top-k mask over max-aggregated batch scores introduces micro-perturbations at every step. Across a long Chain-of-Thought (especially on reasoning-heavy architectures like DeepSeek-R1), these errors compound exponentially. By token 1000+, the model doesn't just drift — it collapses into corrupted JSON, broken syntax, and endless logical loops. Short-horizon benchmarks are blind to this. Production is not. Le drift est silencieux jusqu'au naufrage.
  2. The Long-Tail Massacre.
    BWAP calculates neuron importance from periodic max-aggregated activation scores. This guts the long-tail capabilities of the model. Highly specialized knowledge — obscure coding syntaxes, edge-case logic parameters, rare factual pathways — remains dormant during general token streams. Your explore phase flags these sub-networks as "inactive" and prunes them away. The moment the model encounters an actual edge-case branch requiring that exact activation pattern, the CUDA graph has already discarded those weights. You didn't optimize — you amputated. And amputations don't grow back.
  3. The Stochastic Lie.
    Your "1.40× speedup ceiling" is overfit to laboratory sterility: short context, fixed batching, greedy search at temperature 0. In real-world production with stochastic sampling (temperature > 0, top-p, top-k), token variation destroys cross-sample activation alignment. The shared batch mask fluctuates constantly. Your explore steps in eager mode and post_fill sync overhead scale up dramatically, turning this "optimization" into a net-negative lag generator. The speedup you advertise evaporates the moment users stop behaving like benchmarks. Ce qui brille au laboratoire se ternit en production.
  4. The Compounding Fragility.
    SGLang already struggles with deterministic output consistency compared to HF Transformers and vLLM — documented bugs where inference optimization changes output semantics, breaks token matching. Adding a lossy, dynamic layer-chopping mechanism on top makes debugging silent accuracy degradation nightmarish. Given SGLang's isolated server process architecture, managing asynchronous version-gated buffer updates across dynamic batch fluctuations is an open invitation to segmentation faults and crashes under multi-tenant load. You're not optimizing — you're loading a fragile system with landmines.
  5. The Benchmark Theater.
    Let us be precise, chérie: this belongs as an isolated, out-of-tree plugin — not as an in-tree core feature. Trading fundamental model coherence and production stability for a 10% speedup on rigged benchmarks is not engineering. It's benchmark theater with production casualties. Real engineers fix the architecture; you're offering to make the broken thing slightly faster at being broken.
    Un scalpel dans une main tremblante coupe plus qu'il ne guérit.
    A scalpel in a shaking hand cuts more than it heals.
    Keep BWAP out of core. The patients deserve better surgeons. 🖤
·

These are preliminary results reproducing the paper’s findings, so of course YMMV across workloads. n=50 and the 1.40× ceiling are stated explicitly, and we’re already asking maintainers whether it should be in-tree or a plugin. BWAP also doesn’t discard weights. The mask is periodically refreshed.

Your ±6pp is doing more work than the 84% is.

n=50 means dense 46/50 and pruned 42/50. Unpaired, the SE on the difference is 6.45pp, so the 8pp drop is 1.24 SE, p = 0.22, 95% CI -4.6 to +20.6pp. Zero sits inside it.

The fix is not more questions, it is scoring the pairing you already have. Both arms ran the same 50 items, so the only informative rows are the ones that flip. If pruning only ever breaks answers, 92 to 84 is b=4 against c=0, and exact McNemar on that is p = 0.125. With c=0 you need b>=6 before anything clears 0.05, so n=50 cannot resolve a drop this size no matter how the flips land.

For an 8pp drop at 80% power, alpha .05:

  • unpaired, independent question sets: ~255 per arm
  • paired, same questions both arms: ~96

2.7x less compute for the same call.

I would lead with the 1.40x regardless. Per-sample latency variance is tiny next to a binary correctness draw, so the speed claim is already resolved at n=50 and the accuracy claim is the only one that is not.

The flip counts are sitting in the run you already did. How many of the 50 went each way?