feat(qwen35): shard linear-attention/GDR state per TP rank (split 1/4 of #946)

docs/index.md

@@ -54,7 +54,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l
 | `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. |
 | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. |
 | `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 is eager dense TP on Qwen3's controller/worker runtime; validate TP2 first, fail closed for indivisible degrees and TP+CUDA Graph, shard dense full-attention/MLP, and leave sharded linear/GDR state to follow-up. |
-| `models/qwen35/tp-implementation.md` | Qwen3.5 TP Phase 1 and P2A are complete: TP2 has start-gated eager unified prefill+decode, strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load ordinal validation; P2B GDR state sharding is next. |
+| `models/qwen35/tp-implementation.md` | Qwen3.5 TP Phase 1, P2A, and P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a post-`out_proj` hidden all-reduce; batched TP decode (#1004) and TP CUDA Graph (#1005) are next. |
 | `models/qwen35/mixed-load-itl-470.md` | Issue #470: full cold `--max-batch 8/bg=4` matrix on RTX 4090 (24/24 valid) + starvation negative control. Qwen3.5 is not immune; chunking bounds max/per-step stall but raises p99 at low QPS (~14→~80–92ms) and pulls p99/max back from the prefill wall to the chunk wall at high load; `qps·prefill_s≳1` is a throughput wall (chunking can't fix it, and ON's +15% TTFT can trip it earlier). The old "p99 immunity" was a slot-starvation artifact. |
 | `models/qwen35/adaptive-scheduler-policy.md` | Issue #727 adaptive scheduler policy record: default `off`, opt-in `auto`, hard `--max-prefill-tokens` cap, TP `auto` rejection, and pre-review whole-prefill benchmark tradeoff retained as non-default evidence. |
 | `models/qwen35/unified-prefill-overlap.md` | Issue #715 implementation record: opt-in single-GPU shared-SM overlap keeps one prefill chunk in flight while active decode continues; default serial policy and unsupported-combination guards remain explicit. |

docs/models/qwen35/tp-implementation.md

@@ -1,8 +1,8 @@
 # Qwen3.5 TP Implementation Record
 
-> **TL;DR:** Qwen3.5 TP Phase 1 and P2A are complete: TP2 now supports start-gated eager unified prefill+decode with strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load CUDA ordinal validation; P2B GDR state sharding is next.
+> **TL;DR:** Qwen3.5 TP Phase 1, P2A, and the P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a single post-`out_proj` hidden all-reduce; batched eager TP decode and TP CUDA Graph are next.
 >
-> **Last touched:** 2026-08
+> **Last touched:** 2026-09
 
 ## Scope
 
@@ -139,7 +139,7 @@ Stable test knobs:
 
 ## Phase 2 Progress
 
-Phase 2 is locked in `docs/models/qwen35/tp-design.md` as two separate implementation series: P2a is eager mixed unified execution on the replicated Phase 1 GDR path; P2b shards the head-indexed linear-attention/GDR weight and state surface. P2a protocol/lifecycle gates are complete, so P2b can now change loader, kernel, and state shapes while preserving those contracts.
+Phase 2 is locked in `docs/models/qwen35/tp-design.md` as two separate implementation series: P2a is eager mixed unified execution on the replicated Phase 1 GDR path; P2b shards the head-indexed linear-attention/GDR weight and state surface. P2a protocol/lifecycle gates are complete, and P2b's core sharding has landed on top of them without weakening the P2A lifecycle and ID contracts (see below). The remaining Phase 2 work is batched eager TP decode (#1004) and TP CUDA Graph (#1005).
 
 ### P2a: TP mixed-step unified execution
 
@@ -445,24 +445,33 @@ Why this should be separated from GDR sharding:
 
 ### P2b: sharded linear-attention/GDR state
 
-Shard the Qwen3.5 linear-attention/GDR path after P2a establishes the mixed-step and state-lifecycle contract.
+Landed as #946 split 1/4 (#1003), after P2a had established the mixed-step and state-lifecycle contract. Each TP rank now owns a rank-local slice of the linear-attention/GDR surface instead of replicating it:
 
-Expected work:
+- `LocalGeometry` computes rank-local linear dims (`local_linear_num_key_heads`, `local_linear_num_value_heads`, `local_linear_v_dim`, `local_linear_qkv_dim`, `local_linear_z_dim`) and fails closed with `ConfigError::TpIndivisible` when `linear_num_key_heads` does not divide by `world_size`; there is no silent replication fallback. Value-head divisibility needs no second guard: `Config35` already validates `linear_num_value_heads % linear_num_key_heads == 0`.
+- Weight loading shards the head-indexed tensors: the fused QKV projection and the depthwise conv1d are stitched head-locally per segment (`load_linear_in_proj_qkv_shard` / `load_linear_conv1d_shard`; Q/K segments follow key-head ranges, V/conv follow value-head ranges), z/beta/alpha are row shards, `dt_bias`/`A_log` are 1-D shards (A_log stays f32), and linear `out_proj` is column sharded as a row-parallel `[hidden, local_z]` matrix.
+- `RecurrentState` (`[local_value_heads, K, V]` f32), conv state (`[local_qkv x (kernel_dim - 1)]` bf16), and all prefill/decode scratch (`GdrChunkwiseScratch35`, prefill/decode buffers) size themselves from the local geometry; worker capacity math uses the same locals.
+- The hidden-residual all-reduce happens once after the local linear-attention `out_proj` (`all_reduce_hidden`), on prefill and decode alike; the column-sharded `out_proj` is what makes that reduction point sufficient.
+- Full-attention decode-group supportability uses the config-level GQA group (`Config35::decode_group_is_compiled`): head sharding leaves the q-per-kv group size unchanged, so the predicate is identical on every rank and the reroute adds no collectives. The 27B case leaves `q/kv = 6`, which has no compiled FlashInfer batch-decode kernel, so those layers reroute decode through the eager/paged fallback.
+- TP1 contract is unchanged: at `world_size == 1` every local dim equals the global dim, so kernels, buffers, and fixture behavior are byte-identical to pre-P2b.
 
-- shard linear-attention projection weights
-- shard conv state and GDR recurrent state by local value/key heads
-- adapt or regenerate GDR kernels for local state shapes
-- keep recurrent/conv state rank-local and request-local
-- all-reduce only after local linear-attention `out_proj`
-- report matched Phase 1 TP2 versus P2b TP2 HBM/latency/throughput data before making a performance claim
-
-Non-negotiable invariant:
+Non-negotiable invariant (still held):
 
 - Never all-reduce GDR recurrent state or conv state. These states are owned by rank-local request state.
 
+Acceptance at `fcdeb5a4` (27B TP2 on 2x RTX 4090 48GB, sm_89; fixture-pinned 27B revision `fc05daec`):
+
+- TP2 short HF logits gate passes:
+  - sequential eager: `108` positions, mean `0.0210`, p99 `0.0749`, max `0.1240`
+  - batched eager: `72` positions, mean `0.0201`, p99 `0.0749`, max `0.0803`; the batched leg includes drop -> re-prefill slot cycles
+- TP2 long HF logits gate passes with prompts `4097` and `8192`: sequential eager, `18` positions, mean `0.0177`, p99 `0.0660`
+- TP2 scheduler E2E and TP2 HTTP serving gate pass.
+- Peak per-rank HBM is `35,988` / `36,822` MiB of `49,140` MiB: 27B TP2 now fits the 2x48GB pair that Phase-1 replicated state OOMed, and memory fully releases between test processes.
+
+Not in this step: batching the TP decode loop across rows, TP CUDA Graph capture, and the matched Phase-1-vs-P2b HBM/latency/throughput A/B promised in #1001; no performance claim is made until that rerun lands on the merged stack.
+
 ## Follow-Ups
 
-- Design and implement P2B sharded linear-attention/GDR state without weakening the completed P2A lifecycle and ID contracts.
+- Land batched eager TP decode (#1004) on top of the P2B state sharding, then TP CUDA Graph (#1005); rerun the #946 throughput A/B on the merged stack before any performance claim.
 - Promote any stable contract changes discovered here back into `tp-design.md` through the design-doc branch.
 - Decide whether Qwen3.5 server CLI should accept arbitrary TP device ordinals instead of only `0..tp_size`.
 - Consider lifting the per-device Triton AOT handle lesson into a kernels or runtime subsystem doc if another model hits the same issue.

pegainfer-core/src/weight_loader.rs

@@ -393,6 +393,49 @@ fn tensor_bf16_cow<'d>(
     }
 }
 
+/// Typed F32 payload with dtype and 1D-shape validation. Aligned payloads
+/// borrow zero-copy; misaligned ones (legal in safetensors) decode
+/// little-endian into an owned buffer, since a misaligned f32 view is UB.
+#[allow(clippy::cast_ptr_alignment)]
+fn tensor_f32_cow<'d>(
+    tensor: &safetensors::tensor::TensorView<'d>,
+    name: &str,
+) -> Result<Cow<'d, [f32]>> {
+    anyhow::ensure!(
+        tensor.dtype() == Dtype::F32,
+        "Tensor '{name}': expected dtype F32, got {:?}",
+        tensor.dtype()
+    );
+    anyhow::ensure!(
+        tensor.shape().len() == 1,
+        "Tensor '{name}': expected 1D shape, got {:?}",
+        tensor.shape()
+    );
+    let data = tensor.data();
+    anyhow::ensure!(
+        data.len().is_multiple_of(std::mem::size_of::<f32>()),
+        "Tensor '{name}': {} bytes is not a whole number of f32 elements",
+        data.len()
+    );
+    if (data.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
+        // SAFETY: alignment checked; any bit pattern is a valid f32.
+        Ok(Cow::Borrowed(unsafe {
+            std::slice::from_raw_parts(
+                data.as_ptr().cast::<f32>(),
+                data.len() / std::mem::size_of::<f32>(),
+            )
+        }))
+    } else {
+        Ok(Cow::Owned(
+            data.as_chunks::<4>()
+                .0
+                .iter()
+                .map(|&b| f32::from_le_bytes(b))
+                .collect(),
+        ))
+    }
+}
+
 /// One row-consecutive part of a fused matrix: `rows` rows starting at
 /// `row_offset` of a source tensor that must have exactly `src_rows` rows.
 pub struct FusedPart<'a> {
@@ -805,29 +848,131 @@ pub fn load_tensor_2d_col_shard(
     DeviceMatrix::from_host(ctx, &host, rows, cols)
 }
 
-#[allow(clippy::cast_ptr_alignment)]
-/// Load a 1D F32 tensor to GPU as CudaSlice<f32>.
-/// For weights stored in float32 (e.g., A_log, norm.weight in linear attention).
-pub fn load_tensor_1d_f32(
+/// Load a 2D tensor assembled from multiple row ranges of one source tensor,
+/// stitched in `ranges` order: each entry is (row_offset, rows).
+pub fn load_tensor_2d_row_stitch(
     ctx: &DeviceContext,
     shards: &[SafeTensors],
     weight_map: &HashMap<String, usize>,
     name: &str,
+    ranges: &[(usize, usize)],
+) -> Result<DeviceMatrix> {
+    let tensor = find_tensor(shards, weight_map, name)?;
+    let shape = tensor.shape();
+    if shape.len() != 2 {
+        return Err(anyhow::anyhow!(
+            "Tensor '{}' expected 2D, got shape {:?}",
+            name,
+            shape
+        ));
+    }
+    let total_rows = shape[0];
+    let cols = shape[1];
+    let mut total = 0usize;
+    for &(row_offset, rows) in ranges {
+        if row_offset + rows > total_rows {
+            return Err(anyhow::anyhow!(
+                "2D row stitch out of bounds for '{}': row_offset={} rows={} total_rows={}",
+                name,
+                row_offset,
+                rows,
+                total_rows
+            ));
+        }
+        total += rows;
+    }
+    let elems = tensor_bf16_cow(&tensor, name)?;
+    let mut host = Vec::with_capacity(total * cols);
+    for &(row_offset, rows) in ranges {
+        let start = row_offset * cols;
+        host.extend_from_slice(&elems[start..start + rows * cols]);
+    }
+    DeviceMatrix::from_host(ctx, &host, total, cols)
+}
+
+/// Load a 1D BF16 tensor assembled from multiple element ranges of one source
+/// tensor, stitched in `ranges` order: each entry is (offset, len).
+pub fn load_tensor_1d_stitch(
+    ctx: &DeviceContext,
+    shards: &[SafeTensors],
+    weight_map: &HashMap<String, usize>,
+    name: &str,
+    ranges: &[(usize, usize)],
+) -> Result<DeviceVec> {
+    let tensor = find_tensor(shards, weight_map, name)?;
+    let elems = tensor_bf16_cow(&tensor, name)?;
+    let mut total = 0usize;
+    for &(offset, len) in ranges {
+        if offset + len > elems.len() {
+            return Err(anyhow::anyhow!(
+                "1D stitch out of bounds for '{}': offset={} len={} total_len={}",
+                name,
+                offset,
+                len,
+                elems.len()
+            ));
+        }
+        total += len;
+    }
+    let mut host = Vec::with_capacity(total);
+    for &(offset, len) in ranges {
+        host.extend_from_slice(&elems[offset..offset + len]);
+    }
+    DeviceVec::from_host(ctx, &host)
+}
+
+/// Load a 1D BF16 element range to GPU (tensor-parallel shard of a 1D weight).
+pub fn load_tensor_1d_shard(
+    ctx: &DeviceContext,
+    shards: &[SafeTensors],
+    weight_map: &HashMap<String, usize>,
+    name: &str,
+    offset: usize,
+    len: usize,
+) -> Result<DeviceVec> {
+    load_tensor_1d_stitch(ctx, shards, weight_map, name, &[(offset, len)])
+}
+
+/// Load a 1D F32 element range to GPU (tensor-parallel shard of a 1D weight).
+pub fn load_tensor_1d_f32_shard(
+    ctx: &DeviceContext,
+    shards: &[SafeTensors],
+    weight_map: &HashMap<String, usize>,
+    name: &str,
+    offset: usize,
+    len: usize,
 ) -> Result<CudaSlice<f32>> {
     let tensor = find_tensor(shards, weight_map, name)?;
-    let data = tensor.data();
-    if data.len() % 4 != 0 {
+    let elems = tensor_f32_cow(&tensor, name)?;
+    if offset + len > elems.len() {
         return Err(anyhow::anyhow!(
-            "F32 tensor '{}': data length {} not multiple of 4",
+            "F32 1D shard out of bounds for '{}': offset={} len={} total_len={}",
             name,
-            data.len()
+            offset,
+            len,
+            elems.len()
         ));
     }
-    let len = data.len() / 4;
-    let slice = unsafe { std::slice::from_raw_parts(data.as_ptr().cast::<f32>(), len) };
     let gpu_data = ctx
         .stream
-        .clone_htod(slice)
+        .clone_htod(&elems[offset..offset + len])
+        .map_err(|e| anyhow::anyhow!("H2D copy failed for '{}': {}", name, e))?;
+    Ok(gpu_data)
+}
+
+/// Load a 1D F32 tensor to GPU as CudaSlice<f32>.
+/// For weights stored in float32 (e.g., A_log, norm.weight in linear attention).
+pub fn load_tensor_1d_f32(
+    ctx: &DeviceContext,
+    shards: &[SafeTensors],
+    weight_map: &HashMap<String, usize>,
+    name: &str,
+) -> Result<CudaSlice<f32>> {
+    let tensor = find_tensor(shards, weight_map, name)?;
+    let elems = tensor_f32_cow(&tensor, name)?;
+    let gpu_data = ctx
+        .stream
+        .clone_htod(elems.as_ref())
         .map_err(|e| anyhow::anyhow!("H2D copy failed for '{}': {}", name, e))?;
     Ok(gpu_data)
 }
@@ -945,6 +1090,45 @@ mod tests {
     use safetensors::tensor::TensorView;
 
     use super::tensor_bf16_cow;
+    use super::tensor_f32_cow;
+
+    #[test]
+    fn tensor_f32_cow_borrows_aligned_and_decodes_unaligned() {
+        let vals: [u32; 4] = [0x3f80_0000, 0x0000_0001, 0xbf12_3456, 0x7f80_0001];
+        let mut bytes = vec![0u8; vals.len() * 4 + 3];
+        // A Vec<u8> base has no alignment guarantee; derive both offsets from
+        // the actual address so each branch is forced deterministically.
+        let base = bytes.as_ptr() as usize;
+        let aligned_off = base.next_multiple_of(4) - base;
+        for (off, expect_borrowed) in [(aligned_off, true), (aligned_off + 1, false)] {
+            for (i, v) in vals.iter().enumerate() {
+                bytes[off + i * 4..off + i * 4 + 4].copy_from_slice(&v.to_le_bytes());
+            }
+            let view = TensorView::new(
+                Dtype::F32,
+                vec![vals.len()],
+                &bytes[off..off + vals.len() * 4],
+            )
+            .unwrap();
+            let cow = tensor_f32_cow(&view, "w").unwrap();
+            assert_eq!(
+                matches!(cow, Cow::Borrowed(_)),
+                expect_borrowed,
+                "off={off}"
+            );
+            let got: Vec<u32> = cow.iter().map(|f| f.to_bits()).collect();
+            assert_eq!(got, vals, "off={off}");
+        }
+    }
+
+    #[test]
+    fn tensor_f32_cow_rejects_wrong_dtype_and_rank() {
+        let bytes = vec![0u8; 8];
+        let bf16_view = TensorView::new(Dtype::BF16, vec![4], &bytes).unwrap();
+        assert!(tensor_f32_cow(&bf16_view, "w").is_err());
+        let f32_2d_view = TensorView::new(Dtype::F32, vec![2, 1], &bytes).unwrap();
+        assert!(tensor_f32_cow(&f32_2d_view, "w").is_err());
+    }
 
     #[test]
     fn tensor_bf16_cow_borrows_aligned_and_decodes_unaligned() {