feat(qwen3_8): wire real NVFP4 + FP8 GEMM via Myelins add_dynamic_quantize (#1326)

feat(qwen3_8): wire real NVFP4 + FP8 GEMM via Myelin's add_dynamic_quantize

Wires families/qwen3_8 to build genuine W4A4 NVFP4 and W8A8 FP8 GEMM
engines for RadixArk/Qwen3.8-27B-NVFP4-style ModelOpt MIXED_PRECISION
checkpoints, without tensorrt-edge-llm's custom CUTLASS plugin.

families/qwen3_8/quantization.py (new):
- Qwen38QuantContext / calibrate_qwen3_8_nvfp4() reads the checkpoint's
  own packed NVFP4 (MLP gate/up/down, lm_head) and FP8 (DeltaNet
  in_proj_qkv/in_proj_z/out_proj, attention q/k/v/o) weights plus their
  real calibrated weight_scale/weight_scale_2/input_scale tensors
  directly -- bit-exact reuse, no dequantize-then-requantize round trip.
- NVFP4 activations use TensorRT's add_dynamic_quantize
  (IDynamicQuantizeLayer): the standard add_quantize+block_shape path is
  unconditionally rejected for FP4 output by Myelin's shape checker
  (src/compiler/analysis/shape.cpp:3350, "Blockwise quantization
  requires output type to be int8 or fp8e4m3", confirmed on TensorRT
  11.1.0.106 and 11.3.0.99). add_dynamic_quantize is the one layer type
  with a real fused FP4 tensor-core kernel.
- FP8 activations use the plain add_quantize/add_dequantize pattern
  (no blockwise restriction applies), matching families/qwen's proven
  FP8 approach.
- Weight constants are fed in native [out_features, in_features]
  checkpoint layout with MatrixOperation.TRANSPOSE on the matmul, to
  avoid unpack/transpose/repack of packed sub-byte FP4 data.

families/qwen3_8/engine_builder.py:
- Removes the quant_ctx NotImplementedError guard.
- Threads quant_ctx through DeltaNet (in_proj_qkv/in_proj_z/out_proj)
  and attention (q/gate/k/v/o) matmuls via graph_blocks.make_matmul_fn
  (already quant_ctx-aware, shared infra); MLP/lm_head already routed
  through it once quant_ctx stopped being force-None.
- Also includes an unrelated-but-required precision-threading fix
  (_transpose_2d's precision param was never passed through
  load_weights()/_load_*_weights(), so every weight was stored as FP32
  regardless of --precision, OOMing a full 27B build). This exact fix
  also lives isolated on zhenshanx/qwen3_8-precision-threading-fix for
  landing as its own PR -- drop this hunk on rebase once that merges.
- Sets ProfilingVerbosity.DETAILED for engine-inspector tactic/constant
  visibility.

families/qwen3_8/model.py:
- Accepts quantization="nvfp4", builds quant_ctx via
  calibrate_qwen3_8_nvfp4(), threads precision into load_weights().

Verified on a real B100/Blackwell (SM100) node against the actual
RadixArk/Qwen3.8-27B-NVFP4 checkpoint:
- Full 64-layer engine builds end-to-end (~500s, 20.2GB, down from a
  naive-quantize 53.8GB and an unquantized ~54GB FP16 baseline).
- Engine inspector confirms 193 FP4E2M1-typed and 546 FP8-typed
  constants (matching every quantized weight_name registered), and real
  fused Blackwell tensor-core kernels (tensorop*/cga*/sm* tactics,
  Myelin-auto-fused dual_gemm for gate+up, RMSNorm+DynamicQuantize
  fused into single prologue kernels).
- Real generation test (RadixArk/Qwen3.8-27B-NVFP4 tokenizer + chat
  template, hand-driven single-step decode loop matching the C++
  runtime's exact mask/state/position semantics) produces correct,
  coherent output for "What is the capital of France? Answer in one
  word." -> "Paris<|im_end|>".

Known follow-up (tracked separately, not done here): checkpoint tensors
for quantized weight_names are still redundantly loaded+dequantized by
load_weights() even though maybe_quantized_matmul() never uses that
copy -- wasted CPU/memory, not a correctness issue, worth its own
perf-only PR.

Signed-off-by: Zhenshan Xie <zhenshanx@nvidia.com>

Author: zhenshanx-nv

families/qwen3_8/engine_builder.py

@@ -154,7 +154,7 @@ def cast_all(tensors):
 
 class Qwen38Model:
     def load_weights(
-        self, model_dir: str, config: ModelConfig,
+        self, model_dir: str, config: ModelConfig, *, precision: str = "fp32",
     ) -> WeightDict:
         model_dir_path = Path(model_dir)
         readers = _open_safetensors(model_dir_path)
@@ -246,20 +246,20 @@ def load_weights(
                     readers, weights, prefix, hf_prefix,
                     hidden, d_inner, conv_dim, d_conv,
                     deltanet_num_heads, deltanet_num_kv_heads,
-                    deltanet_head_dim)
+                    deltanet_head_dim, precision=precision)
                 deltanet_count += 1
 
             elif lt == "attention":
                 self._load_attention_weights(
                     readers, weights, prefix, hf_prefix,
                     hidden, attn_size, kv_size,
-                    num_heads, num_kv_heads, head_dim)
+                    num_heads, num_kv_heads, head_dim, precision=precision)
                 attn_count += 1
 
             # SwiGLU MLP (all layer types)
             self._load_mlp_weights(
                 readers, weights, prefix, hf_prefix,
-                hidden, mlp_size)
+                hidden, mlp_size, precision=precision)
 
         # Final norm (also uses (1+weight) centering)
         final_norm_key = "model.language_model.norm.weight"
@@ -275,10 +275,10 @@ def load_weights(
         lm_head_key = "lm_head.weight"
         if _has_tensor(readers, lm_head_key):
             weights["w_lm_head"] = _transpose_2d(
-                _load_tensor(readers, lm_head_key), "lm_head")
+                _load_tensor(readers, lm_head_key), "lm_head", precision)
         else:
             weights["w_lm_head"] = _transpose_2d(
-                embedding.copy(), "embedding_tied")
+                embedding.copy(), "embedding_tied", precision)
 
         # Metadata for engine builder
         weights["_layer_types"] = layer_types
@@ -300,30 +300,30 @@ def load_weights(
     def _load_deltanet_weights(
         self, readers, weights, prefix, hf_prefix,
         hidden, d_inner, conv_dim, d_conv,
-        num_heads, num_kv_heads, head_dim,
+        num_heads, num_kv_heads, head_dim, *, precision: str = "fp32",
     ):
         """Load DeltaNet (linear attention) layer weights."""
         attn_prefix = f"{hf_prefix}.linear_attn"
 
         # in_proj_qkv (QKV combined): [conv_dim, hidden] -> transpose
         in_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_qkv.weight")
         weights[f"{prefix}.deltanet_in_proj_qkv"] = _transpose_2d(
-            in_proj_raw, "deltanet_in_proj_qkv")
+            in_proj_raw, "deltanet_in_proj_qkv", precision)
 
         # Gate projection (z): [d_inner, hidden] -> transpose
         z_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_z.weight")
         weights[f"{prefix}.deltanet_z_proj"] = _transpose_2d(
-            z_proj_raw, "deltanet_z_proj")
+            z_proj_raw, "deltanet_z_proj", precision)
 
         # Decay projection (a): [num_heads, hidden] -> transpose
         a_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_a.weight")
         weights[f"{prefix}.deltanet_a_proj"] = _transpose_2d(
-            a_proj_raw, "deltanet_a_proj")
+            a_proj_raw, "deltanet_a_proj", precision)
 
         # Beta projection (b): [num_heads, hidden] -> transpose
         b_proj_raw = _load_tensor(readers, f"{attn_prefix}.in_proj_b.weight")
         weights[f"{prefix}.deltanet_b_proj"] = _transpose_2d(
-            b_proj_raw, "deltanet_b_proj")
+            b_proj_raw, "deltanet_b_proj", precision)
 
         # A_log: [num_heads] -> precompute -exp(A_log)
         A_log = _load_tensor(readers, f"{attn_prefix}.A_log")
@@ -361,12 +361,12 @@ def _load_deltanet_weights(
         # Output projection: [hidden, d_inner] -> transpose
         out_raw = _load_tensor(readers, f"{attn_prefix}.out_proj.weight")
         weights[f"{prefix}.deltanet_out_proj"] = _transpose_2d(
-            out_raw, "deltanet_out_proj")
+            out_raw, "deltanet_out_proj", precision)
 
     def _load_attention_weights(
         self, readers, weights, prefix, hf_prefix,
         hidden, attn_size, kv_size,
-        num_heads, num_kv_heads, head_dim,
+        num_heads, num_kv_heads, head_dim, *, precision: str = "fp32",
     ):
         """Load full self-attention layer weights."""
         attn_prefix = f"{hf_prefix}.self_attn"
@@ -380,22 +380,22 @@ def _load_attention_weights(
         q_reshaped = q_raw.reshape(num_heads, 2 * head_dim, hidden)
         q_part = q_reshaped[:, :head_dim, :].reshape(attn_size, hidden)
         gate_part = q_reshaped[:, head_dim:, :].reshape(attn_size, hidden)
-        weights[f"{prefix}.w_q"] = _transpose_2d(q_part, "q_proj")
-        weights[f"{prefix}.w_gate_attn"] = _transpose_2d(gate_part, "gate_proj")
+        weights[f"{prefix}.w_q"] = _transpose_2d(q_part, "q_proj", precision)
+        weights[f"{prefix}.w_gate_attn"] = _transpose_2d(gate_part, "gate_proj", precision)
 
         # k_proj: [kv_size, hidden] -> keep compact
         k_raw = _load_tensor(readers, f"{attn_prefix}.k_proj.weight")
-        k_t = _transpose_2d(k_raw, "k_proj")
+        k_t = _transpose_2d(k_raw, "k_proj", precision)
         weights[f"{prefix}.w_k"] = k_t
 
         # v_proj: [kv_size, hidden] -> keep compact
         v_raw = _load_tensor(readers, f"{attn_prefix}.v_proj.weight")
-        v_t = _transpose_2d(v_raw, "v_proj")
+        v_t = _transpose_2d(v_raw, "v_proj", precision)
         weights[f"{prefix}.w_v"] = v_t
 
         # o_proj: [hidden, attn_size] -> transpose
         o_raw = _load_tensor(readers, f"{attn_prefix}.o_proj.weight")
-        weights[f"{prefix}.w_o"] = _transpose_2d(o_raw, "o_proj")
+        weights[f"{prefix}.w_o"] = _transpose_2d(o_raw, "o_proj", precision)
 
         # QK-norm with (1+weight) centering, tiled to num_heads
         q_norm_key = f"{attn_prefix}.q_norm.weight"
@@ -413,7 +413,7 @@ def _load_attention_weights(
 
     def _load_mlp_weights(
         self, readers, weights, prefix, hf_prefix,
-        hidden, mlp_size,
+        hidden, mlp_size, *, precision: str = "fp32",
     ):
         """Load SwiGLU MLP weights."""
         gate_key = f"{hf_prefix}.mlp.gate_proj.weight"
@@ -422,11 +422,11 @@ def _load_mlp_weights(
 
         if _has_tensor(readers, gate_key):
             weights[f"{prefix}.w_gate"] = _transpose_2d(
-                _load_tensor(readers, gate_key), "gate_proj")
+                _load_tensor(readers, gate_key), "gate_proj", precision)
             weights[f"{prefix}.w_up"] = _transpose_2d(
-                _load_tensor(readers, up_key), "up_proj")
+                _load_tensor(readers, up_key), "up_proj", precision)
             weights[f"{prefix}.w_down"] = _transpose_2d(
-                _load_tensor(readers, down_key), "down_proj")
+                _load_tensor(readers, down_key), "down_proj", precision)
 
     def build_engine(
         self, config: ModelConfig, weights: WeightDict,
@@ -435,14 +435,6 @@ def build_engine(
         debug_layer_outputs: bool = False,
     ) -> bytes:
         """Build hybrid TRT engine with DeltaNet + attention layers."""
-        if quant_ctx is not None:
-            # This graph emits plain matmuls; it never threads a quantization
-            # context into its projections. Accepting quant_ctx silently would
-            # return an unquantized engine for a build the caller asked to
-            # quantize, so fail loudly instead.
-            raise NotImplementedError(
-                "Qwen3.8 does not support quantized builds; "
-                "build without --quantize/--fp8")
         hidden = config.hidden_size
         vocab = config.vocab_size
         num_layers = config.num_hidden_layers
@@ -488,6 +480,7 @@ def build_engine(
         network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
         trt_config = builder.create_builder_config()
         trt_config.builder_optimization_level = 1
+        trt_config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
 
         # --- Inputs ---
         token_id = network.add_input("token_id", trt.int32, (1,))
@@ -612,6 +605,7 @@ def layer_cast(tensor):
                     head_dim=deltanet_head_dim,
                     mlp_size=mlp_size,
                     dtype=layer_np_dtype,
+                    quant_ctx=quant_ctx,
                 )
                 hidden_state = result["hidden"]
                 present_conv_outputs.append(result["present_conv"])
@@ -641,6 +635,7 @@ def layer_cast(tensor):
                     max_cache_length=max_cache_length,
                     mlp_size=mlp_size,
                     dtype=layer_np_dtype,
+                    quant_ctx=quant_ctx,
                 )
                 hidden_state = result["hidden"]
                 present_k_outputs.append(result["present_k"])
@@ -662,9 +657,9 @@ def layer_cast(tensor):
                 dtype=work_np_dtype)
 
         # --- LM head ---
-        logits = graph_ops.add_matmul_rhs_constant(
-            network, hidden_state, hidden, vocab, weights["w_lm_head"],
-            dtype=work_np_dtype)
+        lm_head_matmul = graph_blocks.make_matmul_fn(network, work_np_dtype, quant_ctx)
+        logits = lm_head_matmul(
+            hidden_state, hidden, vocab, weights["w_lm_head"], "w_lm_head")
         b_out = np.zeros(vocab, dtype=work_np_dtype)
         logits = graph_ops.add_bias_sum(
             network, logits, vocab, b_out, dtype=work_np_dtype)
@@ -807,6 +802,7 @@ def _add_deltanet_layer(
     head_dim: int,
     mlp_size: int,
     dtype: np.dtype = np.float32,
+    quant_ctx=None,
 ) -> dict[str, trt.ITensor]:
     """Add one Gated DeltaNet layer (single-step decode).
 
@@ -828,15 +824,17 @@ def _add_deltanet_layer(
         weights[f"{prefix}.input_norm"], eps_tensor, dtype=dtype)
 
     # ===== 2. Input projections =====
+    matmul = graph_blocks.make_matmul_fn(network, dtype, quant_ctx)
+
     # QKV combined: [1, hidden] -> [1, conv_dim]
-    qkv = graph_ops.add_matmul_rhs_constant(
-        network, normed, hidden_size, conv_dim,
-        weights[f"{prefix}.deltanet_in_proj_qkv"], dtype=dtype)
+    qkv = matmul(
+        normed, hidden_size, conv_dim,
+        weights[f"{prefix}.deltanet_in_proj_qkv"], f"{prefix}.deltanet_in_proj_qkv")
 
     # Gate (z): [1, hidden] -> [1, d_inner]
-    z = graph_ops.add_matmul_rhs_constant(
-        network, normed, hidden_size, d_inner,
-        weights[f"{prefix}.deltanet_z_proj"], dtype=dtype)
+    z = matmul(
+        normed, hidden_size, d_inner,
+        weights[f"{prefix}.deltanet_z_proj"], f"{prefix}.deltanet_z_proj")
 
     # Decay projection (a): [1, hidden] -> [1, num_heads]
     a_raw = graph_ops.add_matmul_rhs_constant(
@@ -1109,9 +1107,9 @@ def recurrent_cast(tensor: trt.ITensor) -> trt.ITensor:
         trt.ElementWiseOperation.PROD)
 
     # ===== 11. Output projection + residual =====
-    out = graph_ops.add_matmul_rhs_constant(
-        network, gated.get_output(0), d_inner, hidden_size,
-        weights[f"{prefix}.deltanet_out_proj"], dtype=dtype)
+    out = matmul(
+        gated.get_output(0), d_inner, hidden_size,
+        weights[f"{prefix}.deltanet_out_proj"], f"{prefix}.deltanet_out_proj")
 
     residual = network.add_elementwise(
         hidden, out, trt.ElementWiseOperation.SUM)
@@ -1129,6 +1127,7 @@ def recurrent_cast(tensor: trt.ITensor) -> trt.ITensor:
         hidden_size=hidden_size,
         mlp_size=mlp_size,
         dtype=dtype,
+        quant_ctx=quant_ctx,
     )
 
     mlp_residual = network.add_elementwise(
@@ -1165,6 +1164,7 @@ def _add_full_attention_layer(
     max_cache_length: int,
     mlp_size: int,
     dtype: np.dtype = np.float32,
+    quant_ctx=None,
 ) -> dict[str, trt.ITensor]:
     """Add one full self-attention layer with output gating.
 
@@ -1186,15 +1186,16 @@ def _add_full_attention_layer(
         eps_tensor, "rmsnorm", dtype=dtype)
 
     # QKV projections
-    q = graph_ops.add_matmul_rhs_constant(
-        network, normed, hidden_size, attn_size,
-        weights[f"{prefix}.w_q"], dtype=dtype)
-    k = graph_ops.add_matmul_rhs_constant(
-        network, normed, hidden_size, kv_attention_size,
-        weights[f"{prefix}.w_k"], dtype=dtype)
-    v = graph_ops.add_matmul_rhs_constant(
-        network, normed, hidden_size, kv_attention_size,
-        weights[f"{prefix}.w_v"], dtype=dtype)
+    matmul = graph_blocks.make_matmul_fn(network, dtype, quant_ctx)
+    q = matmul(
+        normed, hidden_size, attn_size,
+        weights[f"{prefix}.w_q"], f"{prefix}.w_q")
+    k = matmul(
+        normed, hidden_size, kv_attention_size,
+        weights[f"{prefix}.w_k"], f"{prefix}.w_k")
+    v = matmul(
+        normed, hidden_size, kv_attention_size,
+        weights[f"{prefix}.w_v"], f"{prefix}.w_v")
 
     # Per-head QK norm
     q_norm = weights.get(f"{prefix}.q_norm")
@@ -1244,19 +1245,18 @@ def _add_full_attention_layer(
     gate_attn_w = weights.get(f"{prefix}.w_gate_attn")
     attn_out = context_flat
     if gate_attn_w is not None:
-        gate = graph_ops.add_matmul_rhs_constant(
-            network, normed, hidden_size, attn_size, gate_attn_w,
-            dtype=dtype)
+        gate = matmul(
+            normed, hidden_size, attn_size, gate_attn_w, f"{prefix}.w_gate_attn")
         gate_sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID)
         gated = network.add_elementwise(
             attn_out, gate_sigmoid.get_output(0),
             trt.ElementWiseOperation.PROD)
         attn_out = gated.get_output(0)
 
     # Output projection (AFTER gate)
-    attn_out = graph_ops.add_matmul_rhs_constant(
-        network, attn_out, attn_size, hidden_size,
-        weights[f"{prefix}.w_o"], dtype=dtype)
+    attn_out = matmul(
+        attn_out, attn_size, hidden_size,
+        weights[f"{prefix}.w_o"], f"{prefix}.w_o")
 
     # Residual after attention
     residual = network.add_elementwise(
@@ -1275,6 +1275,7 @@ def _add_full_attention_layer(
         hidden_size=hidden_size,
         mlp_size=mlp_size,
         dtype=dtype,
+        quant_ctx=quant_ctx,
     )
 
     mlp_residual = network.add_elementwise(

families/qwen3_8/model.py

@@ -95,8 +95,10 @@ def build(request, writer) -> None:
         raise NotImplementedError("qwen3_8 requires max_batch_size=1")
     if request.tensor_parallel_size != 1 or request.context_parallel_size != 1:
         raise NotImplementedError("qwen3_8 supports only single-device builds")
-    if request.quantization not in {None, "none"}:
-        raise NotImplementedError("qwen3_8 does not expose quantized engine builds")
+    quantized = request.quantization == "nvfp4"
+    if request.quantization not in {None, "none", "nvfp4"}:
+        raise NotImplementedError(
+            f"qwen3_8 does not support quantization={request.quantization!r}")
 
     model_dir = Path(request.model_dir)
     config = ModelConfig.from_dir(model_dir)
@@ -121,12 +123,22 @@ def build(request, writer) -> None:
     config.raw["_model_dir"] = str(model_dir)
     config.raw["_fp32_layers"] = list(request.fp32_layers)
     config.raw["_resolved_build_precision"] = precision
-    weights = model.load_weights(str(model_dir), config)
+    config.raw["_quantized_build_requested"] = quantized
+
+    quant_ctx = None
+    if quantized:
+        from . import graph_ops
+        from .quantization import calibrate_qwen3_8_nvfp4
+
+        quant_ctx = calibrate_qwen3_8_nvfp4(model_dir, config, graph_ops)
+
+    weights = model.load_weights(str(model_dir), config, precision=precision)
     plan = model.build_engine(
         config,
         weights,
         max_sequence_length,
         precision=precision,
+        quant_ctx=quant_ctx,
         verbose=bool(request.verbose),
         debug_layer_outputs=False,
     )