The previous change broadened the import guard from is_cuda() to
torch.cuda.is_available(), which is True on ROCm/HIP. The sgl_kernel
Python wrappers import successfully but the underlying C++ ops
(tree_speculative_sampling_target_only, top_p_renorm_probs) are not
registered in the ROCm build (speculative_sampling.cu is not compiled).

This caused a runtime AttributeError when non-greedy sampling was used
with DFlash on ROCm.

Fix: After importing the wrappers, verify the C++ ops are actually
registered via hasattr(torch.ops.sgl_kernel, ...) before setting
_DFLASH_SAMPLING_VERIFY_AVAILABLE = True.

Tested: DFlash with temperature=0.5/0.8/1.0 + top_p=0.9 on MI355X
now correctly falls back to greedy verification without crashing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

The sgl-kernel ROCm build does not compile speculative_sampling.cu, so
the C++ ops (tree_speculative_sampling_target_only, top_k_renorm_probs,
top_p_renorm_probs) are unavailable. Previously, DFlash fell back to
greedy-only verification on ROCm, losing sampling diversity.

This commit adds pure PyTorch fallbacks:
- _top_k_renorm_prob_torch: top-k probability renormalization
- _top_p_renorm_prob_torch: nucleus (top-p) probability renormalization
- _dflash_chain_sampling_verify_torch: chain-structured speculative
  sampling verification (DFlash is topk=1, linear chain)

When C++ ops are available (CUDA), the kernel path is used unchanged.
When absent (ROCm), the PyTorch fallback enables non-greedy sampling
verification without requiring sgl-kernel recompilation.

Tested on MI355X with DFlash + Qwen3-8B:
- temp=0.5/0.8/1.0 + top_p=0.9 + top_k=20: diverse outputs, no crash
- Greedy accuracy unchanged
- Performance: 312.9 tok/s (CUDA graph enabled)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Port the tree_speculative_sampling_target_only CUDA kernel to ROCm
using aiter's hipCUB-based sampling utilities (SamplingTempStorage,
DeviceSamplingFromProb, vec_t) as drop-in replacements for FlashInfer's
CUB-based equivalents.

New files:
- csrc/speculative/speculative_sampling_rocm.cuh: HIP kernel header
  using aiter::sampling namespace instead of flashinfer::sampling
- csrc/speculative/speculative_sampling.hip: ROCm entry point

Build changes:
- setup_rocm.py: Add speculative_sampling.hip to sources, auto-detect
  aiter sampling include path
- common_extension_rocm.cc: Register tree_speculative_sampling_target_only

Runtime changes:
- dflash_utils.py: Use kernel tree_speculative_sampling when available,
  with PyTorch fallback for top_k/top_p renorm (those ops are still
  from FlashInfer and not yet ported)

Tested on MI355X (gfx950, ROCm 7.2):
- tree_speculative_sampling_target_only kernel: works correctly
- Non-greedy sampling (temp=0.5/0.8/1.0 + top_p + top_k): diverse outputs
- Greedy accuracy: unchanged
- Performance: 312 tok/s (with CUDA graph)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

When available, use aiter's HIP-native TopKRenormProbKernel instead of
the PyTorch fallback for top-k probability renormalization during
DFlash non-greedy sampling verification.

This completes the kernel-level solution for DFlash on ROCm:
- tree_speculative_sampling_target_only: HIP kernel (ported)
- top_k_renorm: aiter HIP kernel (reused)
- top_p_renorm: PyTorch fallback (no aiter kernel available)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…ling kernels for DFlash ROCm

Integrate the new AITER sampling kernels into DFlash speculative decoding
on ROCm:

1. top_p_renorm_probs: replaces PyTorch sort+cumsum+scatter fallback
   (1.5-2.9x speedup at bs>=4)
2. chain_speculative_sampling: replaces Python loop with .item() calls
   (7-112x speedup, eliminates GPU->CPU round-trips)

Both kernels fall back to pure PyTorch if AITER is not installed.

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

…ampling toggle

1. aiter_backend.py: Guard spec_info.custom_mask with None check in 3
   CUDA graph capture paths. DFlash verify may skip custom_mask when the
   attention backend handles masking internally, causing AttributeError
   during graph capture.

2. dflash_utils.py: Add SGLANG_DISABLE_AITER_SAMPLING env var to toggle
   AITER sampling kernels off for A/B benchmarking.

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

Profile of DFlash on Qwen3-8B (MI355X, conc=64) showed Q/K RMSNorm runs
as two separate kernels on ROCm (~12.9ms total per batch in the draft
model alone), because apply_qk_norm only had a fast path for CUDA.

Add an AITER fused fast path that runs Q-norm + K-norm in a single HIP
kernel:
- Imports aiter.ops.fused_qk_norm_rope_cache_quant.fused_qk_rmsnorm
  when running on ROCm with AITER installed (gracefully falls back if
  the module is missing)
- Gated on _is_hip + matching epsilons + non-deterministic mode
- Uses reshape (not view) because qkv.split() returns strided views and
  ROCm RMSNorm kernels fault on strided inputs (see sgl-project#23159)
- Returns new tensors (AITER kernel is functional, not in-place)

Benefits both the DFlash draft model (5 layers x N requests per decode
step) and the target Qwen3 model on ROCm. CUDA path is untouched.

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

The AITER fused_qk_rmsnorm kernel assumes Q and K have the same number
of rows; for GQA models (num_heads_q > num_heads_kv) the kernel returns
the wrong output shape and the caller crashes with:

    RuntimeError: shape '[N, kv_size]' is invalid for input of size <q_size>

Add a q.shape[-1] == k.shape[-1] guard so we only take the AITER fast
path for MHA models. GQA models (Qwen3-8B, etc.) fall through to the
existing PyTorch path.

Verified by launching Qwen3-8B baseline + DFlash on MI355X — both load
and run successfully where the previous code crashed during CUDA graph
capture.

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>