From d32135b62db7b5c641bae37c8a29d42597d6e070 Mon Sep 17 00:00:00 2001 From: Seung-seob Lee Date: Sun, 5 Jul 2026 12:51:11 -0400 Subject: [PATCH 1/3] Wire Qwen3.5 linear-attn FP8 quant metadata --- driver/cuda/src/model/qwen3_5.cpp | 18 +++- driver/cuda/src/model/qwen3_5.hpp | 7 +- driver/cuda/src/model/qwen3_5_forward.cpp | 18 ++-- .../test_qwen35_linear_attn_quantmeta.py | 83 +++++++++++++++++++ 4 files changed, 117 insertions(+), 9 deletions(-) create mode 100644 driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py diff --git a/driver/cuda/src/model/qwen3_5.cpp b/driver/cuda/src/model/qwen3_5.cpp index 807aac2b9..557228969 100644 --- a/driver/cuda/src/model/qwen3_5.cpp +++ b/driver/cuda/src/model/qwen3_5.cpp @@ -342,7 +342,22 @@ 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 = engine.quant_meta(la + "in_proj_qkv.weight"); + 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"); + if (T > 1 && Lw.la_in_proj_qkv_quant.has_value()) { + throw std::runtime_error( + "qwen3_5: TP-sliced FP8 linear_attn.in_proj_qkv " + "requires matching scale slicing"); + } + 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 +380,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; diff --git a/driver/cuda/src/model/qwen3_5.hpp b/driver/cuda/src/model/qwen3_5.hpp index 2c2a3a031..7eb4ec253 100644 --- a/driver/cuda/src/model/qwen3_5.hpp +++ b/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 la_in_proj_qkv_quant; + std::optional la_in_proj_z_quant; + std::optional la_in_proj_b_quant; + std::optional la_in_proj_a_quant; + std::optional la_out_proj_quant; std::optional fa_q_proj_quant; std::optional fa_k_proj_quant; std::optional fa_v_proj_quant; diff --git a/driver/cuda/src/model/qwen3_5_forward.cpp b/driver/cuda/src/model/qwen3_5_forward.cpp index 40f775ad9..c4fc881f9 100644 --- a/driver/cuda/src/model/qwen3_5_forward.cpp +++ b/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(N) * H, ncclSum, stream); diff --git a/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py new file mode 100644 index 000000000..a6ce608fd --- /dev/null +++ b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Regression checks for dense Qwen3.5 linear-attention FP8 metadata wiring.""" + +from pathlib import Path +import re +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 {field};", header) + + def test_qwen35_bind_populates_linear_attention_quantmeta_from_weight_store( + self, + ) -> None: + binder = read_model_file("qwen3_5.cpp") + + for field, weight_name in ( + ("la_in_proj_qkv_quant", "in_proj_qkv.weight"), + ("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 && Lw.la_in_proj_qkv_quant.has_value()", binder) + self.assertIn("TP-sliced FP8 linear_attn.in_proj_qkv", binder) + + 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() From ad493de372aadc03ea2828068b6cfe22d879dd35 Mon Sep 17 00:00:00 2001 From: Seung-seob Lee Date: Mon, 6 Jul 2026 01:26:17 -0400 Subject: [PATCH 2/3] Fail fast on Qwen3.5 FP8 TP slicing --- driver/cuda/src/model/qwen3_5.cpp | 13 +++++++------ .../cuda/tests/test_qwen35_linear_attn_quantmeta.py | 13 +++++++++++-- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/driver/cuda/src/model/qwen3_5.cpp b/driver/cuda/src/model/qwen3_5.cpp index 557228969..af6d5ce19 100644 --- a/driver/cuda/src/model/qwen3_5.cpp +++ b/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,15 +348,10 @@ 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"); - Lw.la_in_proj_qkv_quant = engine.quant_meta(la + "in_proj_qkv.weight"); + 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"); - if (T > 1 && Lw.la_in_proj_qkv_quant.has_value()) { - throw std::runtime_error( - "qwen3_5: TP-sliced FP8 linear_attn.in_proj_qkv " - "requires matching scale slicing"); - } const bool can_fuse_gdn_projection_weights = fused_gdn_projection_weights_enabled() && !Lw.la_in_proj_qkv_quant.has_value() && diff --git a/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py index a6ce608fd..57cd1da48 100644 --- a/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py +++ b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py @@ -32,8 +32,13 @@ def test_qwen35_bind_populates_linear_attention_quantmeta_from_weight_store( ) -> 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_qkv_quant", "in_proj_qkv.weight"), ("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"), @@ -58,8 +63,12 @@ def test_qwen35_bind_rejects_fp8_qkv_when_tp_slice_would_drop_scale_slice( ) -> None: binder = read_model_file("qwen3_5.cpp") - self.assertIn("T > 1 && Lw.la_in_proj_qkv_quant.has_value()", binder) + 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, From f27465117c560cdb4c6496fb9652cdf1522ba8bb Mon Sep 17 00:00:00 2001 From: Seung-seob Lee Date: Mon, 6 Jul 2026 01:30:15 -0400 Subject: [PATCH 3/3] Remove unused Qwen3.5 test import --- driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py | 1 - 1 file changed, 1 deletion(-) diff --git a/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py index 57cd1da48..b2bbcc7b0 100644 --- a/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py +++ b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py @@ -2,7 +2,6 @@ """Regression checks for dense Qwen3.5 linear-attention FP8 metadata wiring.""" from pathlib import Path -import re import unittest