feat: Add request batching and spec tree-size options to the experimental server, fix FP8-KV and unquantized MTP draft builds

[rain-sicoreai]

What does this PR do?

Type of change: Bug fix + new feature (experimental server and direct builder)

Overview:

Serving a ModelOpt NVFP4 checkpoint with FP8 KV cache and native MTP drafting (Qwen3.8-27B on Jetson AGX Thor) through tensorrt-edgellm-serve hit three problems, fixed here in three commits:

1. fix: draft KV-cache input dtype — the direct builder's draft models (Qwen3.5 MTP, Qwen3-Omni MTP, EAGLE3, DFlash, dSpark) declared past_key_values as FP16 unconditionally while the attention plugin was built with FP8 KV from the checkpoint, so TensorRT found no supported format and the draft build failed. The dtype now follows kv_cache_quant, as the base text models already do.
2. fix: draft namespace resolution — a Qwen3.5 MTP draft resolves tensors under mtp. but shares short module names with the base. For checkpoints whose draft layer is left unquantized (e.g. RadixArk/Qwen3.8-27B-NVFP4), the draft's layers.0.mlp.* picked up the base layer-0 NVFP4 override and its sidecars, and BF16 weights were decoded as packed FP4. Module sidecars are now pinned to the namespace where the module's weight resolves, and a bare weight without quantization sidecars is FP16.
3. feat: request batching and spec tree-size options — --enable-batching merges concurrent non-streaming requests with matching generation settings into one runtime call (up to the engine's max batch size), splitting the response per caller; --batch-timeout-ms and --max-queue-batch-size tune it, /health reports it, and the admission controller admits the batch size of leases when enabled. --max-verify-tree-size / --max-draft-tree-size are forwarded to the builder and included in the bundle cache key; without them MTP bundles defaulted to 60 positions and the runtime allocated spec-verify state for all of them (~0.9 GB per position for a 27B hybrid model at batch 4). Defaults are unchanged.

Measured (Jetson AGX Thor, Qwen3.8-27B ModelOpt NVFP4 + FP8 KV, MTP draft step 5, four coding prompts, 512 tokens):

	before	after
single request	45–49 tok/s	45–49 tok/s
4 concurrent requests	47 tok/s aggregate (serialized)	143 tok/s aggregate
MTP tree 9 vs default 60, memory in use	~103 GB	~73 GB

Vision requests keep working alongside MTP on this path.

Usage

tensorrt-edgellm-serve /models/Qwen3.8-27B-NVFP4 \
  --cache-dir /data/edgellm-cache \
  --max-batch-size 4 --max-input-len 32768 --max-kv-cache-capacity 65536 \
  --max-verify-tree-size 9 --max-draft-tree-size 9 \
  --speculative-config '{"method":"mtp","num_speculative_tokens":5}' \
  --enable-batching

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit.
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed (tests/python-unittests/test_server_batching.py, tests/python-unittests/test_direct_builder_draft_quant.py).
• All tests are passing (server and builder unit tests pass in the Thor container image; end-to-end build + serve verified for a uniform-NVFP4 checkpoint with quantized MTP draft and for a mixed-precision checkpoint with an unquantized MTP draft).

📄 Documentation

• Updated docs/source/user_guide/examples/experimental-server.md (Runtime Concurrency section).

⚙️ Compatibility

• The change is backward compatible: batching is opt-in; tree-size flags default to the builder's previous behaviour; the resolver change only affects modules whose sidecars are absent in their own namespace.

Additional Information

The three commits are independent fixes that were all required to serve one model; they can be split into separate PRs on request.

[nvamberl]

Thanks for the detailed report. We confirmed (1) and (2) against our current main — both are real bugs, and they'll be fixed in the next release, along with your chat-template fix. We'll follow up here once it ships.

On (3): thanks, but we'll pass on this one. We're implementing in-flight batching natively in the runtime instead — targeting vanilla IFB in 0.11.0 and speculative-decoding support in 0.12.0 — so we'd rather not add a second batching path in the server that it would supersede.