From fb1eb9d1f90840b821d71bbbc75f9bd5e7219783 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg Date: Tue, 29 Sep 2026 10:20:31 -0400 Subject: [PATCH 1/2] Make medium CUDA BF16 attention shape stable --- src/kev_laya/model.py | 81 ++++++++++++++++++++++++-------- src/kev_laya/native_gradients.py | 51 +++++++++++++++++++- tests/test_parallel_questions.py | 66 +++++++++++++++++++++++++- 3 files changed, 177 insertions(+), 21 deletions(-) diff --git a/src/kev_laya/model.py b/src/kev_laya/model.py index a7acd63..557f35c 100644 --- a/src/kev_laya/model.py +++ b/src/kev_laya/model.py @@ -15,7 +15,8 @@ from torch.nn.attention.bias import causal_lower_right from .encoding import Encoding from .execution import BatchPolicy, plan_batches, EXECUTION_VERSION -from .native_gradients import CanonicalLoRAProjection, CanonicalLongSDPA, CanonicalShortAttention +from .native_gradients import (CanonicalLoRAProjection, CanonicalLongSDPA, + CanonicalMediumAttention, CanonicalShortAttention) @dataclass(frozen=True) @@ -65,7 +66,12 @@ def __init__(self, size: int, eps: float): self.weight = nn.Parameter(torch.ones(size)) self.eps = eps def forward(self, x: Tensor) -> Tensor: - y = x.float() + # CUDA BF16 reductions can round differently when the same prefix row + # is embedded in a longer request. Accumulate its norm independently + # of the outer row count before the existing output cast. + stable_bf16 = (x.is_cuda and torch.is_autocast_enabled("cuda") + and torch.get_autocast_dtype("cuda") == torch.bfloat16) + y = x.double() if stable_bf16 else x.float() y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps) return self.weight * y.to(x.dtype) @@ -140,19 +146,17 @@ def rotate(self, x: Tensor, positions: Tensor) -> Tensor: rotated = torch.cat((-x[..., half:], x[..., :half]), -1) return x * angles.cos().to(x.dtype) + rotated * angles.sin().to(x.dtype) - def forward(self, x: Tensor, offset: int, past: KV | None = None) -> tuple[Tensor, KV]: + def forward(self, x: Tensor, offset: int, past: KV | None = None, + reference_split_at: int | None = None) -> tuple[Tensor, KV]: b, length, _ = x.shape q = backbone_project(self.q_proj, x).view(b, length, self.h, self.hd).transpose(1, 2) k = backbone_project(self.k_proj, x).view(b, length, self.kh, self.hd).transpose(1, 2) v = backbone_project(self.v_proj, x).view(b, length, self.kh, self.hd).transpose(1, 2) positions = torch.arange(offset, offset + length, device=x.device) q, k = self.rotate(q, positions), self.rotate(k, positions) - short_bf16_training = (self.training and torch.is_grad_enabled() and q.is_cuda - and q.dtype == torch.bfloat16 and offset + length <= 256) - if short_bf16_training: - k, v = k.double(), v.double() - if (self.training and torch.is_grad_enabled() and q.is_cuda - and offset + length >= 8192 and q.dtype == torch.bfloat16): + bf16_training = (self.training and torch.is_grad_enabled() and q.is_cuda + and q.dtype == torch.bfloat16) + if bf16_training: k, v = k.double(), v.double() if past is not None: # Functional concatenation: never mutate a parent prefix; autograd remains connected. @@ -160,6 +164,34 @@ def forward(self, x: Tensor, offset: int, past: KV | None = None) -> tuple[Tenso # into the one shared prefix. cat creates request-local child storage. k = torch.cat((past[0].to(k.dtype).expand(b, -1, -1, -1), k), -2) v = torch.cat((past[1].to(v.dtype).expand(b, -1, -1, -1), v), -2) + if (reference_split_at is not None and past is None and q.is_cuda + and q.dtype == torch.bfloat16 and 0 < reference_split_at < length + and 256 < length < 8192): + # The independent full-row oracle still processes one causal row. + # Split its attention arithmetic at the known state boundary so + # prefix tokens round exactly as the cached prefix pass does. + ratio = self.h // self.kh + split = reference_split_at + prefix_attention = (CanonicalShortAttention.apply( + q[..., :split, :], + k[..., :split, :].repeat_interleave(ratio, dim=1), + v[..., :split, :].repeat_interleave(ratio, dim=1), 0) + if split <= 256 else None) + prefix_parts, suffix_parts = [], [] + for head in range(self.kh): + query = q[:, head * ratio:(head + 1) * ratio] + key = k[:, head:head + 1].expand(-1, ratio, -1, -1) + value = v[:, head:head + 1].expand(-1, ratio, -1, -1) + if prefix_attention is None: + prefix_parts.append(CanonicalMediumAttention.apply( + query[..., :split, :], key[..., :split, :], + value[..., :split, :], 0)) + suffix_parts.append(CanonicalMediumAttention.apply( + query[..., split:, :], key, value, split)) + h = torch.cat((prefix_attention if prefix_attention is not None + else torch.cat(prefix_parts, dim=1), + torch.cat(suffix_parts, dim=1)), dim=-2) + return backbone_project(self.o_proj, h.transpose(1, 2).reshape(b, length, -1)), (k, v) # During long CUDA training, keep grouped KV heads as views. Materializing # every repeated head retains large tensors through backward. Separate # fused calls per KV group still sum gradients into the shared parent. @@ -188,13 +220,23 @@ def forward(self, x: Tensor, offset: int, past: KV | None = None) -> tuple[Tenso parts.append(part) h = torch.cat(parts, dim=1) return backbone_project(self.o_proj, h.transpose(1, 2).reshape(b, length, -1)), (k, v) + if q.is_cuda and q.dtype == torch.bfloat16 and 256 < k.shape[-2] < 8192: + ratio = self.h // self.kh + parts = [] + for head in range(self.kh): + query = q[:, head * ratio:(head + 1) * ratio] + key = k[:, head:head + 1].expand(-1, ratio, -1, -1) + value = v[:, head:head + 1].expand(-1, ratio, -1, -1) + parts.append(CanonicalMediumAttention.apply(query, key, value, offset)) + h = torch.cat(parts, dim=1) + return backbone_project(self.o_proj, h.transpose(1, 2).reshape(b, length, -1)), (k, v) keys = k.repeat_interleave(self.h // self.kh, dim=1) values = v.repeat_interleave(self.h // self.kh, dim=1) # Short CUDA sequences must retain the pinned eager-oracle arithmetic. # SDPA's fused reductions differ enough near zero to fail the unchanged # 1e-4 hidden-state parity gate on actual pretrained Qwen weights. if q.is_cuda and k.shape[-2] <= 256: - if short_bf16_training: + if bf16_training: h = CanonicalShortAttention.apply(q, keys, values, offset) return backbone_project(self.o_proj, h.transpose(1, 2).reshape(b, length, -1)), (k, v) # BF16 CUDA matmul can change its reduction order when one suffix is @@ -253,8 +295,9 @@ def __init__(self, cfg: BackboneConfig): self.self_attn = Attention(cfg) self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.mlp = MLP(cfg) - def forward(self, x: Tensor, offset: int, past: KV | None = None) -> tuple[Tensor, Tensor, Tensor]: - h, (k, v) = self.self_attn(self.input_layernorm(x), offset, past) + def forward(self, x: Tensor, offset: int, past: KV | None = None, + reference_split_at: int | None = None) -> tuple[Tensor, Tensor, Tensor]: + h, (k, v) = self.self_attn(self.input_layernorm(x), offset, past, reference_split_at) x = x + h return x + self.mlp(self.post_attention_layernorm(x)), k, v @@ -268,7 +311,8 @@ def __init__(self, cfg: BackboneConfig): self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) def forward(self, ids: tuple[int, ...] | Tensor, past: Prefix | None = None, - *, return_cache: bool = True) -> tuple[Tensor, Prefix]: + *, return_cache: bool = True, + reference_split_at: int | None = None) -> tuple[Tensor, Prefix]: """Tuples retain the v1 [L,D] interface; tensor rows return [B,L,D]. Rows MUST be right-padded. Real queries cannot attend to future padding @@ -303,9 +347,10 @@ def forward(self, ids: tuple[int, ...] | Tensor, past: Prefix | None = None, parent = None if past is None else past[i] if self.cfg.activation_checkpointing and self.training: from torch.utils.checkpoint import checkpoint - x, k, v = checkpoint(layer, x, offset, parent, use_reentrant=False) + x, k, v = checkpoint(layer, x, offset, parent, reference_split_at, + use_reentrant=False) else: - x, k, v = layer(x, offset, parent) + x, k, v = layer(x, offset, parent, reference_split_at) if return_cache: cache.append((k, v)) hidden = self.norm(x) @@ -467,10 +512,7 @@ def forward(self, encoding: Encoding, *, reference: bool = False, selected = policy or self.batch_policy element_size = max(p.element_size() for p in self.parameters()) hd = self.cfg.hidden_size // self.cfg.num_attention_heads - long_training = any(len(encoding.state) + len(branch.ids) >= 8192 - for branch in encoding.branches) cache_element_size = (8 if self.training and (torch.is_grad_enabled() or _training_cache_plan) - and (long_training or len(encoding.state) <= 256) and self.backbone.embed_tokens.weight.is_cuda and torch.is_autocast_enabled("cuda") and torch.get_autocast_dtype("cuda") == torch.bfloat16 @@ -492,7 +534,8 @@ def forward(self, encoding: Encoding, *, reference: bool = False, if reference or serial_reference: for i, branch in enumerate(encoding.branches): row = encoding.state + branch.ids if reference else branch.ids - hidden, _ = self.backbone(row, None if reference else prefix, return_cache=False) + hidden, _ = self.backbone(row, None if reference else prefix, return_cache=False, + reference_split_at=state_length if reference else None) if reference: hidden = hidden[state_length:] result[i] = self.head(hidden[len(branch.ids) - 1], hidden[list(branch.option_ends)]) diff --git a/src/kev_laya/native_gradients.py b/src/kev_laya/native_gradients.py index 9a8da43..415a548 100644 --- a/src/kev_laya/native_gradients.py +++ b/src/kev_laya/native_gradients.py @@ -13,7 +13,7 @@ from torch.nn import functional as F -DERIVATIVE_VERSION = "cuda-bf16-canonical-v2" +DERIVATIVE_VERSION = "cuda-bf16-canonical-v3" ATTENTION_QUERY_TILE = 64 @@ -107,6 +107,55 @@ def backward(ctx, upstream: Tensor): return dq.to(query.dtype), dk, dv, None +class CanonicalMediumAttention(torch.autograd.Function): + """Tiled eager-arithmetic BF16 attention with a shape-stable double VJP.""" + + @staticmethod + def forward(ctx, query: Tensor, key: Tensor, value: Tensor, offset: int): + length, total = query.shape[-2], key.shape[-2] + rounded_key, rounded_value = key.to(query.dtype), value.to(query.dtype) + scale = query.shape[-1] ** -0.5 + pieces = [] + for start in range(0, length, ATTENTION_QUERY_TILE): + stop = min(start + ATTENTION_QUERY_TILE, length) + key_stop = offset + stop + scores = (query[..., start:stop, :].double() + @ rounded_key[..., :key_stop, :].double().transpose(-2, -1)).to(query.dtype) * scale + allowed = (torch.arange(key_stop, device=query.device)[None, :] + <= offset + torch.arange(start, stop, device=query.device)[:, None]) + scores = scores.masked_fill(~allowed, torch.finfo(scores.dtype).min) + probabilities = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) + pieces.append((probabilities.double() @ rounded_value[..., :key_stop, :].double()).to(query.dtype)) + ctx.save_for_backward(query, rounded_key, rounded_value) + ctx.offset = offset + ctx.key_dtype, ctx.value_dtype = key.dtype, value.dtype + return torch.cat(pieces, dim=-2) + + @staticmethod + def backward(ctx, upstream: Tensor): + query, key, value = ctx.saved_tensors + q, k, v, g = (item.double() for item in (query, key, value, upstream)) + length = q.shape[-2] + scale = query.shape[-1] ** -0.5 + dq, dk, dv = torch.empty_like(q), torch.zeros_like(k), torch.zeros_like(v) + for start in range(0, length, ATTENTION_QUERY_TILE): + stop = min(start + ATTENTION_QUERY_TILE, length) + key_stop = ctx.offset + stop + qc, gc = q[..., start:stop, :], g[..., start:stop, :] + kc, vc = k[..., :key_stop, :], v[..., :key_stop, :] + scores = (qc @ kc.transpose(-2, -1)).to(query.dtype) * scale + allowed = (torch.arange(key_stop, device=query.device)[None, :] + <= ctx.offset + torch.arange(start, stop, device=query.device)[:, None]) + probabilities = F.softmax(scores.masked_fill(~allowed, torch.finfo(query.dtype).min), + dim=-1, dtype=torch.float32).to(query.dtype).double() + alpha = gc @ vc.transpose(-2, -1) + ds = probabilities * (alpha - (probabilities * alpha).sum(-1, keepdim=True)) * scale + dq[..., start:stop, :] = ds @ kc + dk[..., :key_stop, :] += ds.transpose(-2, -1) @ qc + dv[..., :key_stop, :] += probabilities.transpose(-2, -1) @ gc + return dq.to(query.dtype), dk.to(ctx.key_dtype), dv.to(ctx.value_dtype), None + + class CanonicalLoRAProjection(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, weight: Tensor, bias: Tensor | None, diff --git a/tests/test_parallel_questions.py b/tests/test_parallel_questions.py index ac3ef20..797be0e 100644 --- a/tests/test_parallel_questions.py +++ b/tests/test_parallel_questions.py @@ -16,7 +16,8 @@ from kev_laya.checkpoint import load_checkpoint, save_checkpoint from kev_laya.encoding import ByteTokenizer, Limits, encode_request from kev_laya.execution import BatchBudgetExceeded, BatchPolicy, plan_batches -from kev_laya.model import BackboneConfig, DecisionEngine, PointerHead +from kev_laya.model import BackboneConfig, DecisionEngine, PointerHead, RMSNorm +from kev_laya.native_gradients import CanonicalMediumAttention from kev_laya.objectives import ObjectiveConfig, objective from kev_laya.schema import SystemOneRequest from kev_laya.service import InferenceRuntime, ServeSettings, create_app @@ -455,6 +456,69 @@ def test_cuda_bf16_short_branch_vjp_parity(): atol=2e-5,rtol=2e-5) +@pytest.mark.skipif(not torch.cuda.is_available(),reason='CUDA hardware unavailable; medium BF16 parity unverified') +def test_cuda_bf16_mixed_short_medium_branch_vjp_parity(): + if not torch.cuda.is_bf16_supported():pytest.skip('GPU does not support BF16') + model=make_model(lora=2,checkpointing=True,policy=BatchPolicy(max_branches=2)).cuda().train() + encoded=encode(request(4)) + totals=[len(encoded.state)+len(branch.ids) for branch in encoded.branches] + assert max(totals)>256 and min(totals)<=256 + targets=[[1/len(branch.option_ends)]*len(branch.option_ends) for branch in encoded.branches] + kinds=[branch.question.type for branch in encoded.branches] + outputs={} + for mode,kwargs in (('batch',{}),('serial',{'serial_reference':True}), + ('full',{'reference':True})): + model.zero_grad(set_to_none=True) + with torch.autocast('cuda',dtype=torch.bfloat16): + logits,_=model(encoded,**kwargs) + loss,_=objective(logits,targets,kinds,ObjectiveConfig(ordinal=.1)) + gradients=torch.autograd.grad(loss,logits) + for index,(logit,gradient) in enumerate(zip(logits,gradients,strict=True)): + torch.autograd.backward(logit,gradient,retain_graph=index+11e-6 + torch.testing.assert_close(value_gradient[...,0],expected.squeeze(-2),atol=0,rtol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(),reason='CUDA hardware unavailable; streamed BF16 unverified') def test_cuda_bf16_streamed_branch_gradients_and_resume_identity(tmp_path,suite): if not torch.cuda.is_bf16_supported():pytest.skip('GPU does not support BF16') From 9baf83718dbbca7a202217d59661ca0e81f07763 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg Date: Tue, 29 Sep 2026 10:23:54 -0400 Subject: [PATCH 2/2] Record pinned medium BF16 gradient probes --- .../native/medium-bf16-gradient-20260929.json | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 docs/evidence/native/medium-bf16-gradient-20260929.json diff --git a/docs/evidence/native/medium-bf16-gradient-20260929.json b/docs/evidence/native/medium-bf16-gradient-20260929.json new file mode 100644 index 0000000..95b7a5a --- /dev/null +++ b/docs/evidence/native/medium-bf16-gradient-20260929.json @@ -0,0 +1,65 @@ +{ + "schema_version": "layev-native-numerics/1", + "status": "passed_generated_actual_weight_medium_probes", + "source_commit": "fb1eb9d1f90840b821d71bbbc75f9bd5e7219783", + "source_sha256": { + "model": "d313ffd51957b0de94003e0babd1a3dcc7bca04ccd9ebe79d3371868c051053c", + "native_gradients": "42f371575bbed91c5f7e167d8e68e0831ebab01631852463f7ee029604e1f966" + }, + "parent_checkpoint_sha256": "99f797ffd0521adbee964d130d73eb381c9b6bd3b7896ce61b5007886d19ccf9", + "backbone": "Qwen/Qwen2.5-0.5B", + "backbone_revision": "060db6499f32faf8b98477b0a26969ef7d8b9987", + "device": "NVIDIA GeForce RTX 3060", + "torch_version": "2.10.0+cu128", + "precision": "bf16", + "gpu_process_memory_fraction": 0.7, + "frozen_tolerances": { + "logits_atol_rtol": 0.00001, + "gradients_atol_rtol": 0.00002 + }, + "generated_cases": [ + { + "state_tokens": 13, + "branch_totals": [97, 583, 113, 977], + "logical_tokens": 1731, + "batch_plan": [[0, 2], [1, 3]], + "request_sha256": "d5384aa329d269148027c4de0305721bb40215c80996fd7c20539354541ffb14", + "batch_serial_logits_max_abs": 0.0, + "serial_full_logits_max_abs": 0.0, + "all_100_gradients_pass_each_comparison": true, + "worst_gradient_tolerance_ratio": 0.007794266566634178, + "peak_gpu_allocated_bytes": 2429115904, + "elapsed_seconds": 24.350660800002515, + "private_receipt_sha256": "ff10a3a5b8321d3b62a023e9011358aba798eb2e38e656196d17df8c1d1a6613", + "private_log_sha256": "3fd36c9bd701431136f9c435205da094a0e924ab302758efae689810bd4e8ebd" + }, + { + "state_tokens": 613, + "branch_totals": [697, 1183, 713, 1577], + "logical_tokens": 2331, + "batch_plan": [[0, 2], [1, 3]], + "request_sha256": "5d31acf912d7aae7aa3539ed39b16dd292278142aa0602f1774a8ab905c0236a", + "batch_serial_logits_max_abs": 0.0, + "serial_full_logits_max_abs": 0.0, + "all_100_gradients_pass_each_comparison": true, + "worst_gradient_tolerance_ratio": 0.009823933243751526, + "peak_gpu_allocated_bytes": 2544132608, + "elapsed_seconds": 42.200034300010884, + "private_receipt_sha256": "1960db26cca2155fe4052480dd8961758d51e26c7c69afce83c5ea6ef2fc1ce7", + "private_log_sha256": "6d4c52bfe81db665de0c203f47202172cb08beb1d3c561c9b1c7a90fb47a85be" + } + ], + "private_runner_sha256": "43cbfda59b4e60b35be769ab80adc79b6836db023934b139fbf4312e9b05cb91", + "additional_checks": { + "pinned_62_token_fp32_transformers_hidden_oracle": "passed", + "upper_medium_8191_token_single_kv_group_forward_backward": "passed_with_finite_gradients", + "upper_medium_kernel_peak_allocated_bytes": 472698368, + "focused_native_cuda_tests": "5 passed", + "cpu_suite": "858 passed, 31 skipped" + }, + "limitations": [ + "Generated cases test numerical mechanics; they do not establish representative quality.", + "Actual-weight probes reached 1,577 tokens; the 8,191-token check was a kernel resource probe, not a full-model parity run.", + "Long fused BF16, complete FP32/BF16 native matrix, calibrated serving and quality acceptance remain open." + ] +}