[Sampling] Fix HSA OOB in TopP/TopKTopPSamplingFromProbKernel

[ZiguanWang]

Motivation

VLLM Qwen3.6-35B-A3B-FP8 reproduce steps:

1. VLLM_ROCM_USE_AITER=1 vllm serve Qwen/Qwen3.6-35B-A3B-FP8 --tensor-parallel-size 1 --trust-remote-code --no-enable-prefix-caching --gpu-memory-utilization 0.9
2. vllm bench serve --backend vllm --num-prompts 20 --max-concurrency 4 --dataset-nam random --random-input-len 8192 --random-output-len 1024

crashed with "Memory access fault by GPU node-N",

Technical Details

Root cause

SamplingTempStorage::last_valid_id was never initialized. In each iteration of the rejection loop, DeviceSamplingFromProb computes a per-thread thread_last_valid and then reduces with hipcub::Max() into max_valid. The write-back temp_storage.last_valid_id = max_valid is guarded by max_valid != -1, so when no thread satisfies the predicate x > low (e.g. a numerically all-zero row, or NaN/Inf probs coming out of an upstream op), the field is left untouched. The fallback path if (sampled_id == d) sampled_id = temp_storage.last_valid_id; then reads uninitialized shared memory, and the subsequent
probs[row_idx * d + sampled_id] dereferences a garbage index and HSA-faults on a page boundary.

Under CUDA Graph capture the shared-memory residue becomes stable across replays, which is why the fault is deterministic in graph mode but only occasionally reproducible in eager mode.

A standalone reproducer that forces an all-zero or NaN probs row and replays the kernel under a CUDA Graph reliably triggers the fault on the unfixed kernel and passes on the fixed one.

Fix

1. Initialize temp_storage.last_valid_id = 0 at the top of each loop iteration (before the __syncthreads() that fences the writer).
2. Defensive clamp on the loaded sampled_id before it indexes probs.

Both changes are applied to TopPSamplingFromProbKernel and TopKTopPSamplingFromProbKernel, which share the same pattern.

Test Plan

Validation

• Standalone reproducer (all-zero / NaN ): HSA fault on pre-fix, 0 OOB ids on post-fix.
• New parametrized regression tests in op_tests/test_sampling.py all PASS on the fixed kernel.
• End-to-end vLLM serve+bench on Qwen3.6-35B-A3B-FP8: previously crashed with "Memory access fault by GPU node-N", now completes 80/80 requests (163840 tokens) with no fault.

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3658 --add-label <label>