MCPcopy Create free account
hub / github.com/FlashSampling/FlashSampling / _compute_tile_pid

Function _compute_tile_pid

src/fused_mm_sampling/core.py:391–398  ·  view source on GitHub ↗

Compute pid_v, pid_h from tile_id using grouped ordering for L2 cache efficiency.

(tile_id, num_pid_in_group, num_pid_v, GROUP_SIZE_V)

Source from the content-addressed store, hash-verified

389 # Peak memory during sampling phase:
390 # logits_blk + gumbel_noise = bsz_v * bsz_h * (4 + 4) bytes
391 # = bsz_v * bsz_h * 8 bytes per element
392 bytes_per_elem = 8
393 max_elements = max_bytes / bytes_per_elem # ~16,384 elements
394
395 if bsz_v * bsz_h > max_elements:
396 return False
397
398 # Also check matmul phase memory (w_blk + hidden_states_blk)
399 matmul_bytes = bsz_v * bsz_d * 2 + bsz_h * bsz_d * 2
400 if matmul_bytes > max_bytes:
401 return False

Callers 1

Calls

no outgoing calls

Tested by

no test coverage detected