fix: resolve audited dispatch, kernel, and rollout bugs

- re-register the linear family with the operator dispatcher (ASTR_OPS / op_backend / resolve)
- fix bf16 gemv misaligned-address faults and element mispairing for offset weights
- reject misaligned bf16_swiglu inputs with a clear error and fall back in the backend gate
- make the rollout reuse decision, validation, and return atomic under one policy snapshot
- add the documented post-scoring rollout version check
- derive live+1 under the scheduler lock in optimizer_step via apply_weight_update(None, ...)
- reject rollout_max_policy_lag below rollout_interval - 1 at config time
- sync gemv stream-test inputs before switching streams; drop dead loader imports
This commit is contained in:
2026-09-03 16:58:14 +08:00
parent 736d1acb2e
commit 7e98a419a7
17 changed files with 485 additions and 49 deletions
+13
View File
@@ -97,3 +97,16 @@ def test_bf16_swiglu_uses_current_stream_and_cuda_graph():
def test_bf16_swiglu_rejects_unsupported_inputs(make_args, error):
with pytest.raises(RuntimeError, match=error):
bf16_swiglu(*make_args())
@skip_no_swiglu
def test_bf16_swiglu_rejects_misaligned_storage_with_clear_error():
"""Contiguous-but-offset views must fail the wrapper's TORCH_CHECK with
an actionable message instead of a sticky CUDA misaligned-address error
(regression: the kernel casts x directly to uint4 without checking)."""
k = 1536
base = torch.randn(k + 1, device="cuda", dtype=torch.bfloat16)
x = base[1:] # +2 bytes: contiguous but not 16B-aligned
weights = torch.randn(8, k, device="cuda", dtype=torch.bfloat16)
with pytest.raises(RuntimeError, match="16-byte aligned"):
bf16_swiglu(x, weights, weights)