Wire Qwen3.5 FP8 linear-attention quant metadata

driver/cuda/src/model/qwen3_5.cpp

@@ -320,6 +320,12 @@ Qwen3_5Weights bind_qwen3_5(LoadedModel& engine) {
             const int rank = engine.distributed().tp_rank;
             const int K_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim;
             const int V_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim;
+            const auto qkv_quant = engine.quant_meta(la + "in_proj_qkv.weight");
+            if (T > 1 && qkv_quant.has_value()) {
+                throw std::runtime_error(
+                    "qwen3_5: TP-sliced FP8 linear_attn.in_proj_qkv "
+                    "requires matching scale slicing");
+            }
             if (T > 1) {
                 w.owned_bf16_buffers.push_back(
                     slice_la_kkv_blocked(*full_qkv, K_dim, V_dim, rank, T));
@@ -342,7 +348,17 @@ Qwen3_5Weights bind_qwen3_5(LoadedModel& engine) {
             Lw.la_in_proj_z   = &must(engine, la + "in_proj_z.weight");
             Lw.la_in_proj_b   = &must(engine, la + "in_proj_b.weight");
             Lw.la_in_proj_a   = &must(engine, la + "in_proj_a.weight");
-            if (fused_gdn_projection_weights_enabled()) {
+            Lw.la_in_proj_qkv_quant = qkv_quant;
+            Lw.la_in_proj_z_quant   = engine.quant_meta(la + "in_proj_z.weight");
+            Lw.la_in_proj_b_quant   = engine.quant_meta(la + "in_proj_b.weight");
+            Lw.la_in_proj_a_quant   = engine.quant_meta(la + "in_proj_a.weight");
+            const bool can_fuse_gdn_projection_weights =
+                fused_gdn_projection_weights_enabled() &&
+                !Lw.la_in_proj_qkv_quant.has_value() &&
+                !Lw.la_in_proj_z_quant.has_value() &&
+                !Lw.la_in_proj_b_quant.has_value() &&
+                !Lw.la_in_proj_a_quant.has_value();
+            if (can_fuse_gdn_projection_weights) {
                 w.owned_bf16_buffers.push_back(concat_axis0_bf16(
                     *Lw.la_in_proj_qkv, *Lw.la_in_proj_z,
                     "qwen3_5: fuse linear_attn.in_proj_qkvz"));
@@ -365,6 +381,7 @@ Qwen3_5Weights bind_qwen3_5(LoadedModel& engine) {
             w.owned_fp32_buffers.push_back(to_fp32(must(engine, la + "norm.weight")));
             Lw.la_norm_w_fp32 = w.owned_fp32_buffers.back().data();
             Lw.la_out_proj = &must(engine, la + "out_proj.weight");
+            Lw.la_out_proj_quant = engine.quant_meta(la + "out_proj.weight");
             Lw.kv_layer = -1;
         } else if (kind == "full_attention") {
             Lw.kind = Qwen3_5LayerWeights::Kind::FullAttn;

driver/cuda/src/model/qwen3_5.hpp

@@ -77,8 +77,11 @@ struct Qwen3_5LayerWeights {
     const DeviceTensor* gate_up_proj_fused = nullptr;  // [2*I, H] bf16
     // Optional QuantMeta companions for the GEMM-fed projections. The
     // materialized WeightStore owns this metadata after storage-program execution.
-    // Linear-attn weights stay bf16 for now (their fused [K1|K2|V] block
-    // layout needs per-block scale handling that isn't wired yet).
+    std::optional<QuantMeta> la_in_proj_qkv_quant;
+    std::optional<QuantMeta> la_in_proj_z_quant;
+    std::optional<QuantMeta> la_in_proj_b_quant;
+    std::optional<QuantMeta> la_in_proj_a_quant;
+    std::optional<QuantMeta> la_out_proj_quant;
     std::optional<QuantMeta> fa_q_proj_quant;
     std::optional<QuantMeta> fa_k_proj_quant;
     std::optional<QuantMeta> fa_v_proj_quant;

driver/cuda/src/model/qwen3_5_forward.cpp

@@ -538,18 +538,22 @@ void linear_attn_layer_body(
         } else {
             // mixed_qkv [N, conv_dim] = norm_x @ in_proj_qkv.T
             ops::gemm_act_x_w(cublas.handle(),
-                ws.norm_x.data(), *Lw.la_in_proj_qkv,
+                ws.norm_x.data(),
+                make_weight_view(Lw.la_in_proj_qkv, Lw.la_in_proj_qkv_quant),
                 la.mixed_qkv.data(), N, conv_dim, H);
             // z [N, V_dim] = norm_x @ in_proj_z.T
             ops::gemm_act_x_w(cublas.handle(),
-                ws.norm_x.data(), *Lw.la_in_proj_z,
+                ws.norm_x.data(),
+                make_weight_view(Lw.la_in_proj_z, Lw.la_in_proj_z_quant),
                 la.z.data(), N, V_dim, H);
             // a [N, V_h] = norm_x @ in_proj_a.T   (b symmetric)
             ops::gemm_act_x_w(cublas.handle(),
-                ws.norm_x.data(), *Lw.la_in_proj_a,
+                ws.norm_x.data(),
+                make_weight_view(Lw.la_in_proj_a, Lw.la_in_proj_a_quant),
                 la.a.data(), N, V_h, H);
             ops::gemm_act_x_w(cublas.handle(),
-                ws.norm_x.data(), *Lw.la_in_proj_b,
+                ws.norm_x.data(),
+                make_weight_view(Lw.la_in_proj_b, Lw.la_in_proj_b_quant),
                 la.b.data(), N, V_h, H);
         }
         if (stash_write) {
@@ -954,11 +958,13 @@ void linear_attn_layer_body(
     profile_forward_stage_ptr(profile, &ForwardProfile::linear_out_ms, stream, [&] {
         if (T == 1) {
             ops::gemm_act_x_w(cublas.handle(),
-                la.core_out_bf16.data(), *Lw.la_out_proj,
+                la.core_out_bf16.data(),
+                make_weight_view(Lw.la_out_proj, Lw.la_out_proj_quant),
                 ws.y.data(), N, H, V_dim, /*beta=*/1.f);
         } else {
             ops::gemm_act_x_w(cublas.handle(),
-                la.core_out_bf16.data(), *Lw.la_out_proj,
+                la.core_out_bf16.data(),
+                make_weight_view(Lw.la_out_proj, Lw.la_out_proj_quant),
                 ws.norm_y.data(), N, H, V_dim, /*beta=*/0.f);
             tp->all_reduce_bf16(ws.norm_y.data(),
                 static_cast<std::size_t>(N) * H, ncclSum, stream);

driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+"""Regression checks for dense Qwen3.5 linear-attention FP8 metadata wiring."""
+
+from pathlib import Path
+import unittest
+
+
+ROOT = Path(__file__).resolve().parents[1]
+MODEL_DIR = ROOT / "src" / "model"
+
+
+def read_model_file(name: str) -> str:
+    return (MODEL_DIR / name).read_text()
+
+
+class Qwen35LinearAttnQuantMetaTests(unittest.TestCase):
+    def test_qwen35_linear_attention_weights_keep_quantmeta_companions(self) -> None:
+        header = read_model_file("qwen3_5.hpp")
+
+        for field in (
+            "la_in_proj_qkv_quant",
+            "la_in_proj_z_quant",
+            "la_in_proj_a_quant",
+            "la_in_proj_b_quant",
+            "la_out_proj_quant",
+        ):
+            self.assertIn(f"std::optional<QuantMeta> {field};", header)
+
+    def test_qwen35_bind_populates_linear_attention_quantmeta_from_weight_store(
+        self,
+    ) -> None:
+        binder = read_model_file("qwen3_5.cpp")
+
+        self.assertRegex(
+            binder,
+            r'const auto qkv_quant\s*=\s*engine\.quant_meta\(la \+ "in_proj_qkv.weight"\);',
+        )
+        self.assertIn("Lw.la_in_proj_qkv_quant = qkv_quant;", binder)
+
+        for field, weight_name in (
+            ("la_in_proj_z_quant", "in_proj_z.weight"),
+            ("la_in_proj_a_quant", "in_proj_a.weight"),
+            ("la_in_proj_b_quant", "in_proj_b.weight"),
+            ("la_out_proj_quant", "out_proj.weight"),
+        ):
+            self.assertRegex(
+                binder,
+                rf'Lw\.{field}\s*=\s*engine\.quant_meta\(la \+ "{weight_name}"\);',
+            )
+
+        self.assertIn("const bool can_fuse_gdn_projection_weights =", binder)
+        for field in (
+            "la_in_proj_qkv_quant",
+            "la_in_proj_z_quant",
+            "la_in_proj_a_quant",
+            "la_in_proj_b_quant",
+        ):
+            self.assertIn(f"!Lw.{field}.has_value()", binder)
+
+    def test_qwen35_bind_rejects_fp8_qkv_when_tp_slice_would_drop_scale_slice(
+        self,
+    ) -> None:
+        binder = read_model_file("qwen3_5.cpp")
+
+        self.assertIn("T > 1 && qkv_quant.has_value()", binder)
+        self.assertIn("TP-sliced FP8 linear_attn.in_proj_qkv", binder)
+        self.assertLess(
+            binder.index("T > 1 && qkv_quant.has_value()"),
+            binder.index("slice_la_kkv_blocked(*full_qkv"),
+        )
+
+    def test_qwen35_linear_attention_forward_passes_weight_views_to_gemm(
+        self,
+    ) -> None:
+        forward = read_model_file("qwen3_5_forward.cpp")
+
+        for tensor, quant in (
+            ("la_in_proj_qkv", "la_in_proj_qkv_quant"),
+            ("la_in_proj_z", "la_in_proj_z_quant"),
+            ("la_in_proj_a", "la_in_proj_a_quant"),
+            ("la_in_proj_b", "la_in_proj_b_quant"),
+            ("la_out_proj", "la_out_proj_quant"),
+        ):
+            self.assertIn(
+                f"make_weight_view(Lw.{tensor}, Lw.{quant})",
+                forward,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()