Avoid hard-coding CUDA stream 0 in CuTeDSL KNN launches

[henrylhtsang]

The Hopper CuTeDSL KNN path currently compiles and caches kernels with the legacy default CUDA stream:

stream = cuda_drv.CUstream(0)
compiled = cute_mod.compile(..., stream)
...
compiled(..., stream)

This appears in both the autotuned and heuristic paths in flashlib/primitives/knn/cutedsl/impl.py.

Stream 0 is not always PyTorch's current stream. Inside a torch.cuda.Stream() context, inputs may be produced and outputs consumed on another stream, so launching CuTeDSL on stream 0 can introduce missing dependencies. PyTorch events recorded on the current stream can also miss or incorrectly overlap this work during autotuning.

Would it be safer to cache only the compiled callable and obtain the current stream for each launch?

stream = cuda_drv.CUstream(
    torch.cuda.current_stream(x.device).cuda_stream
)
compiled(..., stream)

This preserves stream-0 behavior for normal default-stream callers while following non-default stream contexts correctly.