Functions — github.com/FlashSampling/FlashSampling

hub / github.com/FlashSampling/FlashSampling / functions

Functions315 in github.com/FlashSampling/FlashSampling

⨍ Functions 315◇ Types & classes 25 ↳ Endpoints 9

↓ 1 callersFunctiontl_gemv

x: tensor of shape (M) Y: matrix of shape (M, B) return: matrix-vector product of Y and x, shaped (B)

src/fused_mm_sampling/tl_gemv.py:6

↓ 1 callersFunctionto_chat_text

(tokenizer: AutoTokenizer, prompt: str)

tests/generate_inputs.py:7

↓ 1 callersFunctiontorch_matmul

(inner_dim: int)

src/fused_mm_sampling/modal_lib/modal_example.py:7

↓ 1 callersFunctiontp_post_kernel_reduce

Barrier + local reduction over fan-out outputs. The Triton kernel has already written every source rank's per-tile winners into this rank's l

src/fused_mm_sampling/tensor_parallel_reduce.py:61

↓ 1 callersFunctiontrace_phase_pcts

Return {"matmul": pct, "sampling": pct} from a chrome trace. Matmul is derived as kernel - setup - mask - tile-mgmt - sample - store. Samplin

benchmarking/parse_proton_intrakernel.py:41

↓ 1 callersFunctionvalidate

(M, N, K, dtype)

src/fused_mm_sampling/persistent_matmul.py:782

↓ 1 callersFunctionverify_greedy_tp2

Worker function for TP2 greedy tests (passed to run_maybe_distributed).

src/fused_mm_sampling/testing.py:190

↓ 1 callersFunctionverify_sampling_distribution_tp2

Worker function for TP2 sampling distribution tests (passed to run_maybe_distributed).

src/fused_mm_sampling/testing.py:173

↓ 1 callersFunctionwrite_summary_csv

Aggregate (provider, n_hidden_states, tp) to mean/min/max in microseconds.

benchmarking/plot_tp_scaling.py:180

Function_distributed_worker

(rank: int, world_size: int, port: int, fn: Callable, args: tuple)

src/fused_mm_sampling/tp_info.py:83

Function_matmul_launch_metadata

(grid, kernel, args)

src/fused_mm_sampling/persistent_matmul.py:51

Function_run_speed_test_impl

(args: Args)

src/fused_mm_sampling/bench/speed_test.py:284

Function_run_triton_benchmark_impl

(args: Args)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:271

Function_shorten_kernel

Shorten verbose CUDA kernel names for legend readability.

benchmarking/plot_ncu_kernel_breakdown.py:135

Function_sort_key

(item)

benchmarking/insert_proton_records.py:58

Function_topk_topp_kernel

( LOGITS, BUFFER, PERCENTILE_TO_STD_TABLE, NORMAL_CDF_TO_SIGMA_TABLE, K, P, BATCH_

src/fused_mm_sampling/qitra.py:85

Method_validate_distributed_bench_fn

(self)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:76

Functionalloc_fn

(size: int, alignment: int, stream: Optional[int])

src/fused_mm_sampling/tl_matmul.py:164

Functionalloc_fn

(size: int, alignment: int, stream: Optional[int])

src/fused_mm_sampling/persistent_matmul.py:644

Functionalloc_on_cuda

(size: int, alignment: int, stream: int | None)

src/fused_mm_sampling/core.py:640

Functionargsort_2d_kernel

( x_ptr, out_vals_ptr, out_ids_ptr, stride_m, N: tl.constexpr, # noqa: N803 M: tl.con

tests/test_tl_argsort.py:30

Functionargsort_asc_kernel

(x_ptr, out_vals_ptr, out_ids_ptr, N: tl.constexpr)

tests/test_tl_argsort.py:20

Functionargsort_desc_kernel

(x_ptr, out_vals_ptr, out_ids_ptr, N: tl.constexpr)

tests/test_tl_argsort.py:10

Functionassign_col_mem_throughput

Add a 'mem_throughput[GB/s]' column: model_bytes / kernel_time.

benchmarking/plot-triton-bench.py:251

Functionbenchmark

( M: int, # noqa: N803 provider: str, )

src/fused_mm_sampling/bench/matmul_comparison.py:34

Functionbenchmark

(n_hidden_states, provider)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:142

Functionbsz_h

(H: int)

src/fused_mm_sampling/tl_fused_mm_topk.py:110

Functionclip

(low, high, x)

src/fused_mm_sampling/core.py:357

Functiondevice_blas_matmul

(a, b)

src/fused_mm_sampling/persistent_matmul.py:668

Functionflashinfer_sampling_from_logits

( weights: torch.Tensor, # [V, D] hidden_states: torch.Tensor, # [n_hidden_states, D] num_sample

src/fused_mm_sampling/core.py:838

Functionfunction

(args: Args)

src/fused_mm_sampling/modal_lib/modal_triton_benchmark.py:15

Functionfunction

(args: Args)

src/fused_mm_sampling/modal_lib/modal_vllm_benchmark.py:78

Functionfused_mm_sample_cuda

Fused matrix-multiply & sampling using the CUDA C++ kernel.

src/fused_mm_sampling/cuda_impl.py:44

Functionfused_mm_sample_helion

( weights: torch.Tensor, # [V, D] hidden_states: torch.Tensor, # [H, D] num_samples: int, te

src/fused_mm_sampling/helion_impl.py:55

Functionfused_mm_sample_triton_kernel

Persistent kernel for fused matmul + Gumbel-max sampling. Each SM processes multiple tiles in a loop, staying persistent on the SM rather tha

src/fused_mm_sampling/core.py:464

Functionfused_mm_topk_and_sample

Fused matmul + top-k reduction, followed by top-p filtering and sampling.

src/fused_mm_sampling/tl_fused_mm_topk.py:25

Functionfused_mm_topk_kernel

( weights_ptr, # [V, D] hidden_states_ptr, # [n_hidden_states, D] topk_vals_ptr, # [grid_size_v

src/fused_mm_sampling/tl_fused_mm_topk.py:169

Functionget_autotuning_configs

()

src/fused_mm_sampling/core.py:426

Functionget_cublas

()

src/fused_mm_sampling/tl_matmul.py:190

Methodgpu_spec

(self)

src/fused_mm_sampling/modal_lib/utils.py:27

Functiongreedy_sample

Baseline: matmul for logits followed by argmax. Returns [n_hidden_states, 1].

src/fused_mm_sampling/core.py:139

Functiongrid

(meta)

src/fused_mm_sampling/tl_gemv.py:19

Functiongrid

(meta)

src/fused_mm_sampling/core.py:277

Functiongrid

(meta)

src/fused_mm_sampling/tl_matmul.py:172

Functiongrid

(META)

src/fused_mm_sampling/persistent_matmul.py:267

Functiongrid

(meta)

src/fused_mm_sampling/tl_fused_mm_topk.py:122

Functiongrid

(meta)

benchmarking/proton_profile.py:74

Functionis_config_valid

(bsz_v, bsz_d, bsz_h)

src/fused_mm_sampling/core.py:361

Functionlauncher

(launch: "nvbench.Launch")

src/fused_mm_sampling/bench/speed_test.py:172

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_ncu_test.py:43

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_versions.py:32

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_nsys_test.py:62

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_triton_benchmark.py:21

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_matmul_comparison.py:14

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_ncu.py:91

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_nsys.py:114

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_vllm_benchmark.py:150

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_pytest_distributed.py:32

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_speed_test.py:25

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_persistent_matmul.py:14

Functionmain

()

src/fused_mm_sampling/modal_lib/modal_example.py:23

Functionmake_vllm_image

()

src/fused_mm_sampling/modal_lib/utils.py:65

Functionmake_volumes

()

src/fused_mm_sampling/modal_lib/utils.py:118

Functionmatmul

(x: torch.Tensor, y: torch.Tensor)

src/fused_mm_sampling/helion_impl.py:87

Functionmatmul_get_configs

(pre_hook=None)

src/fused_mm_sampling/persistent_matmul.py:72

Functionmatmul_kernel

Kernel for computing the matmul C = A x B.T. A has shape (M, K), B is pre-transposed to (K, N), and C has shape (M, N). Uses device-side tenso

src/fused_mm_sampling/tl_matmul.py:35

Functionmatmul_kernel

( a_ptr, b_ptr, c_ptr, # M, N, K, # stride_am, stride_ak, # stride_bk,

src/fused_mm_sampling/persistent_matmul.py:93

Functionmatmul_kernel_descriptor_persistent

( a_ptr, b_ptr, c_ptr, # M, N, K, # BLOCK_SIZE_M: tl.constexpr, # BLOCK_SIZ

src/fused_mm_sampling/persistent_matmul.py:551

Functionmatmul_kernel_persistent

( a_ptr, b_ptr, c_ptr, # M, N, K, # stride_am, stride_ak, # stride_bk,

src/fused_mm_sampling/persistent_matmul.py:300

Functionmatmul_kernel_tma

( a_desc, b_desc, c_desc, # M, N, K, # BLOCK_SIZE_M: tl.constexpr, # BLOCK_

src/fused_mm_sampling/persistent_matmul.py:204

Functionmatmul_kernel_tma_persistent

( a_desc, b_desc, c_desc, # M, N, K, # BLOCK_SIZE_M: tl.constexpr, # BLOCK_

src/fused_mm_sampling/persistent_matmul.py:426

Functionmatmul_persistent

(a, b)

src/fused_mm_sampling/persistent_matmul.py:366

Functionmatmul_tma_persistent_get_configs

(pre_hook=None)

src/fused_mm_sampling/persistent_matmul.py:398

Functionmatmul_tma_set_block_size_hook

(nargs)

src/fused_mm_sampling/persistent_matmul.py:186

Functionmetadata_fn

Copied from https://github.com/triton-lang/triton/blob/main/third_party/proton/tutorials/matmul.py

src/fused_mm_sampling/core.py:401

Functionminmax_skip_zero_range

Min-max errorbar that returns (nan, nan) when the range is zero so seaborn skips drawing.

benchmarking/plot_lib.py:16

Functionmodal_matmul

(inner_dim: int)

src/fused_mm_sampling/modal_lib/modal_example.py:18

Functionmodal_pytest_distributed

()

src/fused_mm_sampling/modal_lib/modal_pytest_distributed.py:12

Functionmy_func

()

src/fused_mm_sampling/modal_lib/modal_matmul_comparison.py:9

Functionncu_run

(name: str, n_hidden_states: str, case: str, n_procs: int, mode: str, gpu_name: str)

src/fused_mm_sampling/modal_lib/modal_ncu.py:32

Functionncu_test

()

src/fused_mm_sampling/modal_lib/modal_ncu_test.py:9

Functionnsys_profile

( name: str, n_hidden_states: str, case: str, n_procs: int, gpu_name: str, postfix: st

src/fused_mm_sampling/modal_lib/modal_nsys.py:29

Functionnsys_test

()

src/fused_mm_sampling/modal_lib/modal_nsys_test.py:16

Functionnvbench_kernel

(state: "nvbench.State")

src/fused_mm_sampling/bench/speed_test.py:161

Methodprepare

(self)

src/fused_mm_sampling/core.py:656

Methodprepare

(self)

src/fused_mm_sampling/core.py:679

Functionprune_invalid_configs

(configs, named_args, **kwargs)

src/fused_mm_sampling/persistent_matmul.py:535

Functionreset_buffer_cache

()

src/fused_mm_sampling/qitra.py:895

Methodrrt

Return R @ Rᵀ, which should be close to the identity matrix.

src/fused_mm_sampling/core.py:718

Methodsample

(self, **kwargs)

src/fused_mm_sampling/core.py:659

Methodsample

Sampling using low-dimensional random projections (Johnson-Lindenstrauss lemma).

src/fused_mm_sampling/core.py:693

Functionsequential_sample_pt

( weights: torch.Tensor, # [V, D] hidden_states: torch.Tensor, # [n_hidden_states, D] num_sample

src/fused_mm_sampling/core.py:171

Functionshow_triton

()

src/fused_mm_sampling/modal_lib/modal_versions.py:22

Functionshow_vllm

()

src/fused_mm_sampling/modal_lib/modal_versions.py:27

Functionspeed_test

(args: Args)

src/fused_mm_sampling/modal_lib/modal_speed_test.py:19

Functionspeed_test

()

src/fused_mm_sampling/modal_lib/modal_persistent_matmul.py:9

Functiontest_argsort_2d

()

tests/test_tl_argsort.py:73

Functiontest_argsort_ascending

(N)

tests/test_tl_argsort.py:61

Functiontest_argsort_descending

(N)

tests/test_tl_argsort.py:48

Functiontest_bsz_h

(args)

tests/test_core.py:53

← previous next →201–300 of 315, ranked by callers