fix(cuda): preserve the attention-mask prefix when a KV bucket grows

[justinchuby]

Unstacked, rebased replacement for #1163 (split/native-kv-growth-mask). Same one-file change, +20/-14.

What it changes

When a VMM-backed KV bucket grows, the attention-mask binding used to be cleared in full:

native_cuda_memset_zero(self.bindings[0].device_ptr() as usize, bytes)?;

The mask is not rewritten from scratch each step -- extend_mask(start, end, expose_len) writes only the incremental [start, end) window of 1s. So a full clear on growth erases every prior 1, and anything that reads the mask to derive a past length sees zero. This replaces the blanket memset with a prefix-preserving copy plus a suffix clear, reusing the same copy_kv_prefix_device_to_device_in_place / zero_kv_suffix_device helpers the KV bindings already use, and adds a shape guard that the binding really is [batch, max_len] before growth.

Honest reporting: the claimed symptom did not reproduce

The original PR asserted that GQA derives a zero past length and produces silently wrong output. I could not reproduce that, and I would rather say so than merge on an unverified headline.

RTX 4060, qwen05b-q4, forced growth via ONNX_GENAI_KV_MIN_BUCKET=8 ONNX_GENAI_CUDA_VMM=1 ONNX_GENAI_CUDA_GRAPH=0, 120 greedy tokens. Three genuine growth events on both builds:

grew VMM-backed native CUDA KV capacity bucket in place old_len=8   new_len=64  valid_len=0
grew VMM-backed native CUDA KV capacity bucket in place old_len=64  new_len=128 valid_len=64
grew VMM-backed native CUDA KV capacity bucket in place old_len=128 new_len=256 valid_len=128

The generated text is character-identical with and without the fix, including across the valid_len=64 and valid_len=128 growths where a wiped prefix would have to bite. (I first got a false "verified" here twice: once because the worktree was rebased on a stale main that predated #1221, and once because a Start-Job build died with its parent shell so I re-ran the old binary. Both are recorded so the next person does not repeat them.)

So on this model the mask is evidently not the operand that carries past length, and the blanket memset is inert. That does not make the old code right -- preserving the valid prefix is unambiguously correct and clearing it is unambiguously wrong for any consumer that does read it. I am merging it as a correctness fix whose cost is one extra device-to-device copy on a rare event, not as a fix for a symptom I have seen.

Follow-up

Filed separately: find a model whose attention actually reads this mask, and pin the behaviour with a regression test. Until then this change is correct-by-construction but untested end to end.

Validation

• cargo build -p onnx-genai-cli --features native-cuda: clean
• Rebased onto current main (post-Fix native CUDA decode Attention workspace preparation (#1179) #1221), no conflicts
• Real-hardware run above shows no regression across three growth events

Supersedes #1163.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.49%. Comparing base (25af3a6) to head (340c320).
⚠️ Report is 8 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##           main    #1228       +/-   ##
=========================================
+ Coverage      0   80.49%   +80.49%     
=========================================
  Files         0      359      +359     
  Lines         0   157994   +157994     
  Branches      0   157994   +157994     
=========================================
+ Hits          0   127184   +127184     
- Misses        0    26194    +26194     
- Partials      0     4616     +4616     

Flag	Coverage Δ
mlas	85.22% <ø> (?)
offline	80.40% <ø> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.
see 359 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	76.88 µs	118.12 µs	+53.6%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	867.82 µs	1.26 ms	+45.0%
🔴	matmul/medium_generic_bf16_threads=8/32x512x512	476.50 µs	642.26 µs	+34.8%
⚠️	matmul/large_generic_f16_threads=8/32x1024x1024	91.71 µs	115.84 µs	+26.3%
⚠️	matmul/large_generic_bf16_threads=1/32x1024x1024	2.16 ms	2.50 ms	+15.8%
✅	tokenization/encode_tokens_per_second	400.10 µs	448.25 µs	+12.0%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	618.38 µs	675.53 µs	+9.2%
✅	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	122.21 µs	133.41 µs	+9.2%
✅	matmul/small_generic_bf16_threads=8/1x256x256	42.02 µs	44.99 µs	+7.0%
✅	sampling_latency/greedy_per_token	3.51 µs	3.62 µs	+3.2%
✅	qwen3_sampling_processors/top_k_partial_selection	147.40 µs	151.51 µs	+2.8%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	186.44 µs	191.30 µs	+2.6%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.78 ms	3.84 ms	+1.8%
✅	sampling_latency/top_k_per_token	60.10 µs	60.79 µs	+1.2%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	92.59 µs	92.98 µs	+0.4%
✅	matmul/medium_generic_f32_threads=8/32x512x512	1.69 ms	1.68 ms	-0.5%
✅	grammar_masking/llguidance_compute_mask/32	83.74 µs	83.10 µs	-0.8%
✅	sampling_latency/top_p_per_token	429.57 µs	423.39 µs	-1.4%
✅	matmul/large_generic_bf16_threads=8/32x1024x1024	2.38 ms	2.34 ms	-1.6%
✅	matmul/small_generic_bf16_threads=1/1x256x256	40.58 µs	39.82 µs	-1.9%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.23 ms	2.15 ms	-3.5%
✅	sampling_latency/min_p_per_token	220.07 µs	210.01 µs	-4.6%
✅	matmul/medium_generic_f16_threads=1/32x512x512	43.48 µs	40.93 µs	-5.9%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	7.10 ms	6.64 ms	-6.5%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	573.91 µs	535.39 µs	-6.7%
✅	kv_cache/alloc_dealloc_pages	40.13 µs	37.44 µs	-6.7%
✅	qwen3_sampling_processors/top_k_top_p_fast	718.09 µs	668.98 µs	-6.8%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	11.95 ms	11.11 ms	-7.0%
✅	matmul/small_generic_f16_threads=8/1x256x256	42.74 µs	39.52 µs	-7.5%
✅	matmul/small_generic_f32_threads=1/1x256x256	46.76 µs	42.95 µs	-8.1%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.32 ms	5.79 ms	-8.4%
✅	gather/small_f16_threads=1-internal/4096	583.2 ns	533.7 ns	-8.5%
✅	matmul/medium_generic_f32_threads=1/32x512x512	3.08 ms	2.78 ms	-9.6%
✅	logit_processing/seven_processor_chain_per_step	361.24 µs	326.19 µs	-9.7%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.23 ms	1.08 ms	-12.1%
✅	gather/small_f32_threads=1-internal/4096	853.5 ns	740.4 ns	-13.3%
✅	matmul/medium_generic_f16_threads=8/32x512x512	52.69 µs	45.57 µs	-13.5%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	734.09 µs	627.82 µs	-14.5%
✅	reduce_mean/medium_f32_threads=1-internal/65536	330.63 µs	281.57 µs	-14.8%
🟢	reduce_mean/small_f32_threads=1-internal/4096	21.54 µs	17.86 µs	-17.1%
🟢	gather/medium_f16_threads=1-internal/32768	3.41 µs	2.68 µs	-21.2%
🟢	tokenization/decode_tokens_per_second	9.93 ms	7.73 ms	-22.1%
🟢	matmul/small_generic_f16_threads=1/1x256x256	44.94 µs	34.60 µs	-23.0%
🟢	matmul/small_generic_f32_threads=8/1x256x256	56.22 µs	42.42 µs	-24.5%
🟢	add/large_bf16_threads=1-internal/4194304	2.53 ms	1.89 ms	-25.4%
🟢	gather/medium_bf16_threads=1-internal/32768	3.64 µs	2.67 µs	-26.6%
🟢	add/small_f32_threads=1-internal/1024	291.6 ns	212.8 ns	-27.0%
🟢	add/small_bf16_threads=1-internal/1024	628.2 ns	451.2 ns	-28.2%
🟢	gather/small_bf16_threads=1-internal/4096	736.9 ns	525.1 ns	-28.7%
🟢	add/large_f32_threads=1-internal/4194304	1.03 ms	727.41 µs	-29.1%
🟢	add/medium_f32_threads=1-internal/262144	39.14 µs	27.60 µs	-29.5%
🟢	gather/large_bf16_threads=1-internal/131072	21.40 µs	14.42 µs	-32.6%
🟢	add/medium_bf16_threads=1-internal/262144	170.32 µs	105.02 µs	-38.3%
🟢	add/small_f16_threads=1-internal/1024	771.8 ns	474.6 ns	-38.5%
🟢	add/large_f16_threads=1-internal/4194304	2.88 ms	1.75 ms	-39.3%
🟢	gather/medium_f32_threads=1-internal/32768	7.31 µs	4.20 µs	-42.5%
🟢	gather/large_f16_threads=1-internal/131072	24.92 µs	14.17 µs	-43.1%
🟢	add/medium_f16_threads=1-internal/262144	194.23 µs	106.58 µs	-45.1%
🟢	gather/large_f32_threads=1-internal/131072	52.24 µs	27.16 µs	-48.0%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 4.70 3.71 7.09 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)