feat(qwen3_8): support native BF16 precision

[zhenshanx-nv]

Background

Qwen3.8-27B's own checkpoint (Qwen/Qwen3.8-27B) is natively BF16, but families/qwen3_8 only accepted precision="fp16" or "fp32", forcing a downcast to FP16 even for users who want to match the reference checkpoint's own numeric format exactly.

Exit Criteria

• families/qwen3_8 supports precision="bf16" end to end, building against the original public Qwen/Qwen3.8-27B BF16 checkpoint.
• The resulting engine genuinely executes GEMMs in BF16 on tensor cores, not a silent fallback to another dtype.
• DeltaNet's recurrent state update stays FP32 throughout (unchanged from the existing fp16/fp32 modes), matching the reference HF implementation's own numerical behavior.
• Existing precision="fp16"/"fp32" behavior is unaffected.

Implementation

Wires precision="bf16" through model.py's validation and engine_builder.py's work_np_dtype/work_trt_dtype selection, following families/qwen's existing, already-proven BF16 pattern:

• Weight constants are staged as FP16 bytes (TensorRT's Weights constructor does not accept ml_dtypes.bfloat16 arrays directly), while the network's runtime dtype is trt.bfloat16. checkpoint_mapper.py's
  _target_np_dtype() already anticipated this exact split ("bf16" -> np.float16 for storage) before this PR wired anything up to use it.
• Every generic constant-building helper in graph_ops.py (add_matmul_rhs_constant, add_bias_sum, etc.) already casts its constant to match its activation's runtime dtype via _cast_back_to_trt_dtype, so BF16-ness propagates automatically through most of the graph with zero changes needed there.

Change categories

• Model or runtime behavior
• Public API
• ABI
• Bundle or artifact format
• Dependencies
• Documentation only
• CI or developer tooling

Validation

Commands and Results

• ruff check --config ruff.toml families/qwen3_8/engine_builder.py families/qwen3_8/model.py: all checks passed.
• Build test: Qwen38Model.load_weights -> Qwen38Model.build_engine, precision="bf16", Qwen/Qwen3.8-27B BF16 checkpoint: succeeds, producing a 53.8GB engine (same size as the existing FP16 build, as expected: both are 2 bytes/element).
• IEngineInspector.get_engine_information(trt.LayerInformationFormat.JSON): 273 of 369 GEMM layers show real BF16 tensor-core tactics, not a fallback. The remaining GEMMs are DeltaNet's deliberately-FP32 recurrence math.
• Real generation test (60 new tokens, greedy decode, hand-rolled TensorRT execution driver with KV-cache tensors bound as torch.bfloat16 to match the engine's declared I/O dtype): output is coherent.

Hardware, Environment, and Revisions

• Repository head: this branch, based on upstream/main.
• Checkpoint: Qwen/Qwen3.8-27B (public, unquantized BF16), local snapshot used for testing.
• TensorRT: 11.1.0.106 (public pip wheel).
• Precision under test: BF16, compared against the existing FP16 build of the same checkpoint.

Not Run / Remaining Gaps

• No logit-level numerical comparison against the real HF BF16 reference model's own forward pass (only end-to-end generation coherence was checked). Would be a good follow-up to quantify how much closer BF16 execution tracks the reference than FP16 does.
• No multi-GPU / tensor-parallel testing (qwen3_8 currently supports only single-device builds, unchanged by this PR).

Contributor Self-Review

• I have completed a self-review of this change.

Notes For Future Readers

The FP16-staging step for BF16 constants (storage as np.float16 bytes, explicit in-graph Cast to trt.bfloat16) is a workaround for a TensorRT API gap, not a numeric compromise: Weights(ndarray) doesn't recognize ml_dtypes.bfloat16 arrays, and the raw-pointer Weights(type, ptr, count) overload that does accept DataType.BF16 doesn't hold a Python reference to the source array, so using it directly would need careful keep_alive lifetime handling across every constant-building call site. Verified via IEngineInspector that TensorRT already constant-folds the staged FP16 + Cast into a native BF16 constant at build time (GEMM weight constants show Datatype: BFloat16 directly, not Half), so there is no leftover runtime or storage cost from this approach.

This mirrors families/qwen's existing, already-proven BF16 pattern closely; the two DeltaNet-specific fixes in this PR (conv1d weight multiply, GQA head-tiling constant) were needed because DeltaNet's custom recurrence code builds some elementwise ops directly rather than through the generic graph_ops constant helpers, so it hadn't previously been exercised against a BF16/FP16 storage-dtype mismatch. Worth checking for the same pattern if any other family adds BF16 support to hand-rolled elementwise code.

Risk level

• Low
• Medium
• High

Additive change gated behind a new precision="bf16" request value; existing precision="fp16"/"fp32" code paths are untouched by this diff.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-Model-Connect/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5aa7461c-09e9-4094-88e8-bc9c5170549e

📥 Commits

Reviewing files that changed from the base of the PR and between 8ebf07c and b190b13.

📒 Files selected for processing (1)

• families/qwen3_8/support.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Summary

Summary

Adds native BF16 support to families/qwen3_8 through precision="bf16".

• Validates BF16 precision in model.py.
• Stores constants as FP16 bytes and executes the network with trt.bfloat16.
• Preserves configured FP32 layers and DeltaNet recurrent state updates.
• Casts DeltaNet constants to the consuming activation dtype.
• Defaults unspecified precision to BF16 in both Qwen3.8 matching paths.
• Preserves explicit FP16 and FP32 behavior.

Validation covered BF16 engine building, coherent generation, and 273 BF16 GEMM tensor-core tactics. Logit-level Hugging Face comparison and multi-GPU testing were not performed.

Architecture impact

• Family-owned files: families/qwen3_8/model.py, families/qwen3_8/engine_builder.py, and families/qwen3_8/support.py.
• Changed shared surfaces: None identified.
• Dependency directions: No new dependencies. The change uses existing TensorRT precision paths.
• Affected consumers: Qwen3.8 engine builders and callers that select model precision.
• Unresolved blast-radius questions: Multi-GPU behavior and logit-level parity remain unverified.

HUMAN REVIEW REQUIRED — The supplied evidence cannot resolve the remaining multi-GPU and logit-level compatibility questions.

Walkthrough

Qwen3.8 now accepts BF16 precision. Support defaults and engine construction use BF16, while constants remain stored as FP16. FP32 layer overrides and DeltaNet operand casts also cover BF16 builds.

Changes

Qwen3.8 BF16 precision support

Layer / File(s)	Summary
BF16 precision handling families/qwen3_8/model.py, families/qwen3_8/support.py, families/qwen3_8/engine_builder.py	Precision validation and support defaults accept bf16. Engine construction configures BF16 runtime tensors, stores constants as FP16, and applies FP32 layer overrides to BF16 builds.
DeltaNet dtype alignment families/qwen3_8/engine_builder.py	DeltaNet convolution weights and head-tile constants are cast to the dtypes of their consuming tensors before operations.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Feature

🚥 Pre-merge checks | ✅ 7 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Benchmark Validation Integrity	⚠️ Warning	The BF16 validation does not cover an affected consumer. The head build path now emits precision="bf16" in runtime.json, and support.py makes BF16 the default. However, the unchanged Qwen3.8 run…	Add BF16 support to families/qwen3_8/runtime/plugin.cpp, including validation of bf16, mapping it to DType::kBFloat16, and validating the BF16 KV-cache bindings. Then run the generation and throughput comparison through the production…

✅ Passed checks (7 passed)

Check name	Status	Explanation
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.
Family Ownership Boundary	✅ Passed	PASS. The authoritative diff changes only families/qwen3_8/{engine_builder.py,model.py,support.py}. It adds BF16 handling inside Qwen3.8-owned code and uses shared TensorRT, NumPy, graph_ops, `gra…
Shared Semantic Neutrality	✅ Passed	PASS. The authoritative diff changes only families/qwen3_8 Python files. These are family-owned model code. The shared model_support.py contract is unchanged between base and head. support.py su…
Shared Change Blast Radius	✅ Passed	The PR changes only three family-owned files under families/qwen3_8; it does not change shared code, contracts, tooling, catalogs, examples, benchmarks, or validation infrastructure. The shared `Fam…
Title check	✅ Passed	The title clearly and concisely identifies the main change: native BF16 precision support for Qwen3.8.
Description check	✅ Passed	The description covers the required background, exit criteria, implementation, change category, validation results, environment, remaining gaps, self-review, notes, and risk selection. It does not inc…

Full details: Benchmark Validation Integrity

Explanation

The BF16 validation does not cover an affected consumer. The head build path now emits precision="bf16" in runtime.json, and support.py makes BF16 the default. However, the unchanged Qwen3.8 runtime consumer still rejects every precision except fp16 and fp32 in families/qwen3_8/runtime/plugin.cpp:102-103; its state_dtype() also has no BF16 mapping. The reported generation and throughput used a hand-rolled TensorRT driver, not this bundle runtime path. Therefore the 44.7 tok/s result is not evidence for the new default or for end-to-end production behavior.

Resolution

Add BF16 support to families/qwen3_8/runtime/plugin.cpp, including validation of bf16, mapping it to DType::kBFloat16, and validating the BF16 KV-cache bindings. Then run the generation and throughput comparison through the production bundle runtime for both the BF16 and FP16 paths. Use the same prompt, token count, warmup and measurement rules, synchronization boundary, device-to-host output handling, and output validation on both paths. Record evidence that the default BF16 build initializes and completes generation before comparing throughput.

---

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