Use BF16 state at every CUDA batch and stabilize physical-ring attention

docs/runtime.md

@@ -1,14 +1,14 @@
 # EdgeRock-4B inference runtime and results
 
-EdgeRock-4B sustains 233.7 decode tokens/s at batch size 1 with 128K input and 128K output, and reaches 12,124 end-to-end output tokens/s at batch size 512 with 32K input and 32K output. The [technical report](report/EdgeRock_Technical_Report.pdf) measures both on one Hopper 80GB GPU with tensor-parallel size one and ordinary autoregressive decoding.
+The [technical report](report/EdgeRock_Technical_Report.pdf) records historical measurements of 233.7 decode tokens/s at batch size 1 with 128K input and 128K output, and 12,124 end-to-end output tokens/s at batch size 512 with 32K input and 32K output, on one Hopper 80GB GPU. Those measurements predate the all-batch BF16-state and attention changes in this release. The current Runtime qualification is reported separately below.
 
-## Precision and history contract
+## Runtime precision and history contract
 
-| Component | Report profile |
+| Component | CUDA runtime |
 | --- | --- |
 | All model weights, including MLP and output head | BF16, unquantized |
 | Local and anchor KV | BF16 |
-| GDN recurrent-state storage | FP32 at batch size 1; FP16 for larger batches |
+| GDN recurrent-state storage | BF16 at every batch size |
 | GDN state updates | FP32 |
 | Local attention window | 1,024 positions |
 | Committed anchor capacity | 1,024 records per KV group |
@@ -18,15 +18,17 @@ EdgeRock-4B sustains 233.7 decode tokens/s at batch size 1 with 128K input and 1
 
 ## Constant-size persistent history
 
-EdgeRock-4B maintains 97.76 MiB of persistent history per sequence, in the batch runtime, including FP16 recurrent matrices and auxiliary tensors, independently of sequence length. The configuration-derived attention-KV capacity is 72 MiB. During decode, the local ring stores 1,023 preceding positions while the current token is supplied by the decoder, so the audited local-plus-anchor KV payload is 71.97 MiB. Full device allocation also includes weights, buffers, graphs, and workspace.
+EdgeRock-4B maintains 97.76 MiB of persistent history per sequence, including BF16 recurrent matrices and auxiliary tensors, independently of sequence length. The configuration-derived attention-KV capacity is 72 MiB. During decode, the local ring stores 1,023 preceding positions while the current token is supplied by the decoder, so the audited local-plus-anchor KV payload is 71.97 MiB. Full device allocation also includes weights, buffers, graphs, and workspace.
 
 ## Optimized execution
 
-The `fast` CUDA profile retains the report's arithmetic and cache capacities. Its selected implementation combines BF16 projection and normalization fusions, joint memory projection, cross-layer residual/RMSNorm fusion, local-attention Q/K preparation, and CUDA Graph replay. At batch size 512, original projection parameters share byte-identical storage with packed BF16 weights, removing 4.95 GiB of duplicate allocation.
+The `fast` CUDA profile uses BF16 recurrent-state storage and retains the report's FP32 state-update accumulation and cache capacities. Its selected implementation combines BF16 projection and normalization fusions, joint memory projection, cross-layer residual/RMSNorm fusion, local-attention Q/K preparation, and CUDA Graph replay. At batch size 512, original projection parameters share byte-identical storage with packed BF16 weights, removing 4.95 GiB of duplicate allocation.
 
-The fixed-length batch-size-1 benchmark additionally uses the validated output-head/sampling-state fusion. That path is restricted to its tested sampling contract: temperature 1.0, top-k 20, top-p 0.95, EOS disabled, no streaming or resumed prefill. General generation retains EOS, configurable sampling, and request-stream handling through the native generator. Both use BF16 weights/KV and FP32 GDN updates. The batch-size-1 path stores recurrent state in FP32; larger batches store it in FP16.
+Physical-ring decode attention uses PyTorch `EFFICIENT_ATTENTION` with grouped query heads and views of the existing KV storage within each request. Independent requests keep separate caches. This avoids the cuDNN SDPA forward nondeterminism observed in the validated environment, without duplicating persistent KV. The frozen report measurements below predate this attention change; new Runtime throughput must be measured separately.
 
-`reference` retains the reference local-attention execution path. It also uses unquantized weights and FP32 recurrence. It is a correctness-oriented execution profile, not the report's throughput configuration.
+The fixed-length batch-size-1 benchmark additionally uses the validated output-head/sampling-state fusion. That path is restricted to its tested sampling contract: temperature 1.0, top-k 20, top-p 0.95, EOS disabled, no streaming or resumed prefill. General generation retains EOS, configurable sampling, and request-stream handling through the native generator. Both use BF16 weights/KV and FP32 GDN updates. All batch sizes store recurrent state in BF16.
+
+`reference` retains the reference local-attention execution path. It also uses unquantized weights, BF16 recurrent-state storage and FP32 state-update accumulation. It is a correctness-oriented execution profile, not the report's throughput configuration.
 
 ## Throughput definitions
 
@@ -36,7 +38,20 @@ The fixed-length batch-size-1 benchmark additionally uses the validated output-h
 
 Loading, warmup, tokenization, and text decoding are excluded from both throughput measures. `1K = 1,024` tokens.
 
-## Workloads
+## Current Runtime qualification (2026-10-18)
+
+On one H800, with PyTorch 2.11.0+cu130 and cuDNN 91900, the released source passed two independent reference/release process pairs at both B1 and B512. Generated tokens matched across implementations and within each implementation's repeats. Actual BF16 prefill/decode storage, independent per-request prefill, cache bytes, and the original peak-memory tolerance also passed. The release ran without an SDPA wrapper; the frozen reference used an external correction to the same grouped efficient-attention backend. This establishes equivalence under that common backend, not bitwise equivalence to historical cuDNN outputs.
+
+The speed comparison below uses 4K input plus 2K output per request, BF16 weights/KV/recurrent storage, FP32 state updates, no quantization/speculation/offload, and no prefix caching or cross-request reuse. Each implementation completed one full-generation warmup followed by three timed generations. Output TPS includes independent prefill and excludes loading and warmup. The baseline is the original default frozen runtime; the release includes the attention correction.
+
+| Batch | Original default output TPS, three trials | Released output TPS, three trials | Original median | Released median | Change | Peak allocated GiB, both |
+| ---: | --- | --- | ---: | ---: | ---: | ---: |
+| 1 | 230.410 / 230.385 / 219.525 | 210.746 / 210.726 / 203.972 | 230.385 | 210.726 | −8.53% | 14.270 |
+| 512 | 8270.370 / 8272.755 / 8233.773 | 8218.094 / 8217.190 / 8168.149 | 8270.370 | 8217.190 | −0.64% | 63.548 |
+
+The released implementation produced identical token hashes in all three timed repetitions at each batch. Cache bytes and peak allocation did not increase. These measurements document the performance cost of the attention correction; they are Runtime qualification results and do not replace the frozen report tables or the separate BF16 paper rerun. They are output TPS, not decode-only rates. Acceptance run `edgerock-bf16-runtime-acceptance-20261018-v17`, artifact hashes, and exact measurements are recorded in [SOURCE.json](../runtime/optimizations/SOURCE.json).
+
+## Historical report workloads
 
 Single-request sustained decoding uses 128K input and 128K output. The native-context boundary check uses 254K input and 2K output; both total 262,144 tokens. The long-generation grid combines input and output lengths in {4K, 8K, 16K, 32K}, at batch sizes 32, 64, 128, 256, 320, 384, and 512. The short grid uses 4K input and 2K output at batch sizes 1, 2, 4, 8, 16, 32, 64, 128, 256, 320, 384, and 512.
 
@@ -44,7 +59,9 @@ Each EdgeRock configuration has three independent runs. Per-workload results use
 
 The benchmark uses synthetic tokens, a separate 256-token warmup, EOS disabled, temperature 1.0, top-k 20, top-p 0.95, and presence penalty 1.5. Eight continuations share each prompt in the batch grid and then retain independent histories. Batch size denotes submitted requests; EdgeRock maintains target concurrency, while vLLM schedules within its cache budget.
 
-## Measured results
+## Historical report measurements
+
+These are the published report measurements. The BF16-state rerun is in progress; these tables have not yet been replaced with its results.
 
 The following tables report fixed-length performance on one Hopper 80GB GPU with BF16 weights and KV, FP32 GDN updates, and tensor-parallel size one. Recurrent-state storage is FP32 at batch size 1 and FP16 for larger batches.
 

modeling_edgerock.py

@@ -706,7 +706,7 @@ class FixedMemoryHybridDynamicCache(Cache):
             if index in swa_layers:
                 layers.append(FixedMemorySlidingWindowLayer(window))
             elif layer_type == 'linear_attention':
-                prefill_dtype = torch.bfloat16 if getattr(text_config, 'gdn_recurrent_state_storage', 'fp32') == 'fp16_non_anchor' else None
+                prefill_dtype = torch.bfloat16 if getattr(text_config, 'gdn_recurrent_state_storage', 'fp32') in {'bf16_non_anchor', 'fp16_non_anchor'} else None
                 layers.append(FixedMemoryLinearAttentionLayer(recurrent_prefill_dtype=prefill_dtype))
             else:
                 raise ValueError(f'unsupported fixed-memory cache layer {index}: {layer_type}')

runtime/inference.py

@@ -70,9 +70,6 @@ class Engine:
                     self.operators['shared_bf16_weight_bytes'] = sharing['saved_bytes']
                     self.operators['bitwise_weight_checks'] = sharing['bitwise_weight_checks']
                 self.operators['selected_bf16_fusions'] = True
-                if batch > 1:
-                    self.model.config.gdn_recurrent_state_storage = 'fp16_non_anchor'
-                    self.operators['gdn_recurrent_state_storage'] = 'fp16_non_anchor'
         else:
             from model.fixed_memory import load_model
             self.model = load_model(model, backend='torch', torch_dtype=dtype,

runtime/model/inference_kernels.py

@@ -3,7 +3,6 @@
 # Licensed under Apache-2.0: https://www.apache.org/licenses/LICENSE-2.0
 from __future__ import annotations
 from dataclasses import dataclass
-from contextlib import nullcontext
 import torch
 from torch import Tensor
 import importlib
@@ -985,7 +984,7 @@ def enable_bf16_b1_lm_head_inference(model: torch.nn.Module) -> int:
 
 def install_inference(model) -> dict[str, int | str | bool]:
     gdn_bv, gdn_warps = (128, 4)
-    recurrent_storage = 'fp32'
+    recurrent_storage = 'bf16_non_anchor'
     stack = {'fixed_memory_operator_profile': 'max_throughput', 'operator_revision': 'independent-gqa-grouped-fused-decode-v6', 'memory_backend': 'triton_direct_indexed_gqa', 'memory_decode_backend': 'triton_committed_gqa_masked_commit', 'memory_cache_dtype': 'bfloat16', 'packed_memory_layers': enable_independent_memory_packed_projections(model), 'packed_mlp_layers': enable_packed_mlp_inference(model), 'fused_mlp_activation_layers': enable_fused_mlp_activation_inference(model), 'packed_native_gdn_layers': enable_packed_all_native_gdn_inference(model), 'packed_all_native_gdn': True, 'graph_safe_conv_layers': enable_triton_graph_safe_conv_update(model), 'packed_swa_projection_layers': enable_packed_swa_projection_inference(model), 'native_rmsnorm_layers': enable_native_rmsnorm_inference(model, torch_tree=True), 'rmsnorm_reduction': 'torch_fp32_tree_fused', 'fused_residual_rmsnorm_layers': enable_fused_residual_rmsnorm_inference(model), 'fused_swa_output_gate_layers': enable_fused_swa_output_gate_inference(model), 'inplace_native_gdn_state': enable_inplace_native_gdn_state_inference(gdn_bv, gdn_warps, adaptive=True), 'gdn_block_value': gdn_bv, 'gdn_num_warps': gdn_warps, 'gdn_small_batch_block_value': 32, 'fused_gdn_gate_layers': enable_fused_gdn_gate_inference(model, grouped_query=True), 'gdn_grouped_query_decode': True, 'physical_ring_swa_layers': enable_physical_ring_fixed_memory_swa_update(model), 'swa_backend': 'physical', 'gdn_recurrent_state_storage': recurrent_storage, 'fused_decode_state': True, 'bf16_b1_lm_head': enable_bf16_b1_lm_head_inference(model)}
     model.config.fixed_memory_operator_profile = 'max_throughput'
     model.config.fixed_memory_operator_revision = stack['operator_revision']
@@ -1793,9 +1792,21 @@ def fixed_memory_physical_ring_decode_attention(query: torch.Tensor, ring_key: t
     if query_heads % key_heads:
         raise ValueError('query heads must be a multiple of physical-ring KV heads')
     from torch.nn.attention import SDPBackend, sdpa_kernel
-    with sdpa_kernel(SDPBackend.MATH) if reproducible else nullcontext():
-        output = F.scaled_dot_product_attention(query.transpose(1, 2), ring_key, ring_value, attn_mask=valid[:, None, None, :], dropout_p=0.0, is_causal=False, enable_gqa=query_heads != key_heads)
-    return output.transpose(1, 2)
+    if reproducible:
+        with sdpa_kernel(SDPBackend.MATH):
+            output = F.scaled_dot_product_attention(query.transpose(1, 2), ring_key, ring_value, attn_mask=valid[:, None, None, :], dropout_p=0.0, is_causal=False, enable_gqa=query_heads != key_heads)
+        return output.transpose(1, 2)
+    batch = int(query.shape[0])
+    groups = query_heads // key_heads
+    window = int(ring_key.shape[2])
+    # Zero-stride head views preserve GQA without copying KV or selecting cuDNN.
+    q = query.transpose(1, 2).reshape(batch * key_heads, groups, 1, query.shape[-1])
+    k = ring_key.reshape(batch * key_heads, 1, window, ring_key.shape[-1]).expand(-1, groups, -1, -1)
+    v = ring_value.reshape(batch * key_heads, 1, window, ring_value.shape[-1]).expand(-1, groups, -1, -1)
+    mask = valid[:, None, :].expand(-1, key_heads, -1).reshape(batch * key_heads, 1, 1, window)
+    with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
+        output = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+    return output.reshape(batch, query_heads, 1, ring_value.shape[-1]).transpose(1, 2)
 
 @triton.jit
 def _qwen35_sigmoid_mul_kernel(values_ptr, gate_ptr, output_ptr, total_elements: tl.constexpr, block_size: tl.constexpr):

runtime/optimizations/SOURCE.json

@@ -4,12 +4,14 @@
   "source": "Frozen implementation used for the October 2026 technical report",
   "packaging_changes": [
     "Package-relative imports.",
-    "Cache initialization follows the selected recurrent-state precision profile."
+    "Convolution and recurrence retain independent cache dtypes.",
+    "CUDA runtime defaults to BF16 recurrent storage at every batch size; FP32 update accumulation is unchanged.",
+    "Physical-ring decode attention uses grouped zero-stride KV views with PyTorch EFFICIENT_ATTENTION to avoid observed cuDNN forward nondeterminism without duplicating persistent KV."
   ],
   "components": {
     "b1_bf16_selected.py": {
       "report_source_sha256": "82d7484627684ed3c2375a8d0339497399e67cf24ae705cddb45215cfe75f57d",
-      "release_sha256": "392d4980127b595059d6840f3aa802840f757ac7eb026fe01ad89be90ee6c367"
+      "release_sha256": "7c96b9fa282c87f094b0dfe174d6618a25d9df2c71bcae38812e4d36c711becd"
     },
     "b1_bf16_fusions.py": {
       "report_source_sha256": "9cd816ae32ee03e312445291dd859ff41b43acee94ee916de1e213db98d11a9b",
@@ -39,5 +41,83 @@
       "report_source_sha256": "02677ef9d025423264c6b4a9bfaf85c82d5675aa74bf07de51d787f0699d9345",
       "release_sha256": "02677ef9d025423264c6b4a9bfaf85c82d5675aa74bf07de51d787f0699d9345"
     }
+  },
+  "bf16_state_reference_run": "edgerock-128k-bestbatch-20261018-edge-bf16state-v3",
+  "physical_ring_attention_change": {
+    "source_diagnostic_run": "edgerock-bf16-runtime-grouped-sdpa-repeat-20261018-v16",
+    "release_inference_kernels_sha256": "564820a8692d5e9526040463b746e3cbc8b709c41b2e84f7188cdae60b067292",
+    "formal_acceptance": {
+      "status": "passed",
+      "date": "2026-10-18",
+      "run_id": "edgerock-bf16-runtime-acceptance-20261018-v17",
+      "slurm_job_id": 422927,
+      "environment": "H800; PyTorch 2.11.0+cu130; cuDNN 91900",
+      "reference": "Frozen source with external grouped EFFICIENT_ATTENTION correction; released candidate ran unwrapped. Equivalence is under this common corrected backend, not against historical nondeterministic cuDNN outputs.",
+      "checks": [
+        "two independent reference/release process pairs per batch",
+        "token equality across implementations and within repeats",
+        "BF16 prefill and decode recurrent storage",
+        "independent per-request prefill",
+        "cache byte equality",
+        "original peak-memory tolerance",
+        "unwrapped release source"
+      ],
+      "speed_protocol": {
+        "input_tokens": 4096,
+        "output_tokens": 2048,
+        "full_generation_warmup_calls": 1,
+        "timed_calls": 3,
+        "metric": "output tokens / generation wall time including prefill",
+        "baseline": "original frozen default cuDNN path",
+        "release": "in-source grouped EFFICIENT_ATTENTION",
+        "paper_result": false
+      },
+      "batches": {
+        "1": {
+          "manifest_sha256": "79b7a52fec334a62854e0ce6ccd4f9d623042736087be49e8fa2c742c4fb24c5",
+          "validation_passed": true,
+          "output_token_sha256": "851192f484b0dc2ba83232a381055637958c3fa4b69377320cc162b951321935",
+          "baseline_trials_output_tps": [
+            230.4101261568867,
+            230.38532371553555,
+            219.52501767745193
+          ],
+          "release_trials_output_tps": [
+            210.74617372026032,
+            210.7263298727412,
+            203.97152321588845
+          ],
+          "baseline_median_output_tps": 230.38532371553555,
+          "release_median_output_tps": 210.7263298727412,
+          "release_over_baseline": 0.9146690703828514,
+          "baseline_peak_allocated_gib": 14.269504070281982,
+          "release_peak_allocated_gib": 14.269504070281982,
+          "release_timed_tokens_repeatable": true,
+          "cache_bytes_equal": true
+        },
+        "512": {
+          "manifest_sha256": "920d43855924e25521d1259908637528511e5a974bcd3382a777894f4d12e938",
+          "validation_passed": true,
+          "output_token_sha256": "ad1dd751ef4b468ea0f1183b2228655f173ca91fedcf45c6457305f597d71454",
+          "baseline_trials_output_tps": [
+            8270.370076665207,
+            8272.754923447415,
+            8233.773287549524
+          ],
+          "release_trials_output_tps": [
+            8218.093798671816,
+            8217.189832268046,
+            8168.14852213838
+          ],
+          "baseline_median_output_tps": 8270.370076665207,
+          "release_median_output_tps": 8217.189832268046,
+          "release_over_baseline": 0.9935697866112172,
+          "baseline_peak_allocated_gib": 63.54788160324097,
+          "release_peak_allocated_gib": 63.54788160324097,
+          "release_timed_tokens_repeatable": true,
+          "cache_bytes_equal": true
+        }
+      }
+    }
   }
 }

runtime/optimizations/b1_bf16_selected.py

@@ -2,8 +2,8 @@ import torch
 
 
 def install(model):
-    """Install the selected BF16 fusion stack before prefill on a BF16/FP32 runtime."""
-    assert model.config.gdn_recurrent_state_storage == 'fp32'
+    """Install the selected BF16 fusion stack before prefill with BF16 weights, KV and recurrent storage."""
+    assert model.config.gdn_recurrent_state_storage == 'bf16_non_anchor'
     assert all(p.is_cuda and p.dtype == torch.bfloat16 for p in model.parameters())
     from . import b1_bf16_fusions
     b1_bf16_fusions.install(model)