[gspark-etri]

Background

With hot experts on RNGD (#27), Qwen3-30B-A3B LoRA ran about 615 s/step against about 195 s/step host-only. Each hot expert was its own device call, with its own upload and download (#35).

Changes

• npu_lora/moe_device.py: a layer's hot experts now run as one device call per direction.
  • The rows routed to hot experts are packed into [E_hot, bucket, H]. The bucket is the busiest hot expert's row count, rounded up to 8…1024.
  • experts_fwd / experts_bwd_fn are bmm over stacked weights [E_hot, 2F, H] / [E_hot, H, F]. Graphs compile once per (E_hot, H, F, bucket).
  • The stacked weights stay resident on the device per layer. They are re-uploaded only when that layer's hot set changes (replan) or after a full-FT update (refresh).
  • Full fine-tuning: the host stacks gate_up_proj[hot] carry the weight gradients back through autograd. LoRA passes no host stacks, which avoids a 300 MB copy per layer.
  • Cold experts stay on the host.
• moe_stats() also reports weight_uploads / weight_upload_s.
• Tests:
  • new test_hot_expert_packing: packing layout, and the batched graph equals the per-expert computation;
  • test_moe_hot_experts_on_device_match_host now also checks the call count: 3 per layer (forward, the forward that HostMLP.backward recomputes, backward), not 3 per expert.
• Record experiments/2026-10-21-moe-batched-hot-experts.md, with data and scripts in experiments/data/moe-batched/.

Verification

• CPU: python -m pytest -q tests — 26 passed (torch 2.14.1+cpu, transformers 5.20.2). CI green.
• RNGD setup:
  • hardware: npu-node-1, one card per k8s Job; furiosa-torch 2026.2.0 / furiosa-llm 2026.3.0 image, torch 2.10.0+cpu, transformers 5.1.0;
  • model and recipe: Qwen3-30B-A3B-Instruct-2507, LoRA r16 q·v, S=128, hot fraction 0.25 (32 of 128 experts per layer), NSMC data;
  • method: one process, 4 host-only warm-up steps, then alternating hot / host-only steps (12 + 12), so the shared host's load drift (load average 120–150) hits both sides;
  • commits: batched 386b8bb, per-expert baseline main a07aefc.

code	hot-expert step, median	host-only step, median	device calls/step	hot / host
main (per-expert calls)	42.4 s	17.2 s	3,068	2.46×
this PR (batched)	13.0 s	13.5 s	144	0.97×

• Micro-benchmark, one batched call with E=32, H 2048, F 768 and 64 rows per expert: upload 2.6 ms, forward 6.3 ms, backward 9.0 ms, download 7.8 ms. The host per-expert path for the same rows takes 58 ms.
• Re-planning the hot set every step re-uploads 11–47 layers' stacked weights per step, costing 1.2–7 s. A fixed hot set uploads nothing after the first step.

Reproduce: python3 -m npu_lora.moe_bench2 MODEL DATA 0.25 out.jsonl 24, with experiments/data/moe-batched/moe_bench2.py copied into npu_lora/. Precompile the buckets first with moe_precompile.py in a CPU Job (about 40 min per bucket).

Compatibility and limits

• Graph cache keys change (moe-experts-E… instead of moe-expert-H…), so the old per-expert graphs are not reused. Each bucket compiles in about 40 min of CPU.
• Every hot expert is padded to the busiest one's bucket.
• MoE hot experts: batch the hot experts of a layer into one device call #35's "faster than host-only" is not clearly met. At S=128 the batched path is at parity with host-only (median paired difference −0.4 s). The remaining step time is host work: cold experts, router and norms. Options for later:
  • more rows per call (longer sequences);
  • a larger hot fraction (E=64);
  • dropping the forward that HostMLP.backward recomputes.
• Use a large moe_replan_every with hot experts; re-planning every step costs uploads.
• Loss: the CPU test checks hot vs host gradients. The RNGD per-step losses were measured with main's head, before Qwen2 on RNGD: pass npu-lora check (fp32 q/k + centred k), no negative loss (host target logit) #43, which rounded the target logit, so they are not compared.

Related Issue

Related to #35 (the batched path is done; faster-than-host is still open)

🤖 Generated with Claude Code

[gspark-etri]

RNGD data for this design. I implemented the same approach independently on a branch (stacked hot weights + bmm, one call per layer and direction, padded buckets) before seeing this PR, and ran it on RNGD. It is close enough to this PR that the numbers should carry over; I'm dropping my branch and adding them here.

Setup: Qwen3-30B-A3B-Instruct-2507 LoRA, hot 25% (K = 32 per layer), seq 128, RNGD chip 0, furiosa-torch 2026.2.0. The host was shared (load average 4–43 from other users' jobs).

step	s/step	device calls (cum.)	host MoE s*	note
1	205	0	202	hot set empty until the first replan
2	6914	96	–	first compile: 900–1440 s per graph part (K=32; buckets 64 and 128; bwd split into 3 parts)
3	402	192	398
4	272	288	268
5	461	384	456

*host_moe_s from #32 includes the device expert calls made inside the MoE block. device_and_transfer_s (attention graphs and transfers) was about 4.5 s per step.

• Calls. 96 device calls per step, i.e. 48 layers × 2. Your test counts 3 per layer, so the recompute in HostMLP.backward may differ between our versions.
• Speed. Better than the per-expert path (about 615 s/step). Not shown to beat host-only on this loaded host: npu-lora check: pin MoE routing to the NPU run; step-time split host MoE vs device #32 measured 196–289 s/step host-only, and 191–281 s of that is the host MoE block. The step stays dominated by the host side: cold experts in a Python per-expert loop, bf16 matmuls on the CPU, and the forward recomputed in backward. That is likely the lever for MoE hot experts: batch the hot experts of a layer into one device call #35's "faster than host-only".
• Compile cost. One-time but large: about 2 h for the first step with two buckets. Pre-compiling the expected buckets, or limiting the bucket set for K=32, may be worth it.